NiccoloN
32 lines
1.3 KiB
C++
32 lines
1.3 KiB
C++
#pragma once
|
|
|
|
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
|
#include "mlir/IR/ValueRange.h"
|
|
#include "mlir/Transforms/DialectConversion.h"
|
|
|
|
#include "llvm/ADT/SmallVector.h"
|
|
|
|
#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
|
|
|
|
namespace onnx_mlir {
|
|
|
|
/// Slices a statically shaped tensor along one axis into contiguous pieces of
|
|
/// at most `sliceSize` elements.
|
|
llvm::SmallVector<mlir::Value> sliceTensor(const mlir::Value& tensorToSlice,
|
|
size_t axis,
|
|
int64_t sliceSize,
|
|
mlir::PatternRewriter& rewriter,
|
|
mlir::Location loc);
|
|
|
|
llvm::SmallVector<mlir::Value> sliceVector(const mlir::Value& vectorToSlice,
|
|
int64_t sliceSize,
|
|
mlir::PatternRewriter& rewriter,
|
|
mlir::Location loc);
|
|
|
|
/// Partitions one logical vector into per-core crossbar-sized slices using the
|
|
/// current PIM target geometry.
|
|
llvm::DenseMap<CoreId, llvm::SmallVector<mlir::Value>> sliceVectorPerCrossbarPerCore(
|
|
const mlir::Value& vectorToSlice, mlir::PatternRewriter& rewriter, mlir::Location loc);
|
|
|
|
} // namespace onnx_mlir
|