Refactor ONNXToSpatial Common and diagnostics

src/PIM/Conversion/ONNXToSpatial/Common/ShapeTilingUtils.cpp

@@ -0,0 +1,96 @@
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+
+#include "llvm/ADT/SmallVector.h"
+
+#include "ShapeTilingUtils.hpp"
+#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+SmallVector<Value> sliceTensor(
+  const Value& tensorToSlice, size_t axis, int64_t sliceSize, ConversionPatternRewriter& rewriter, Location loc) {
+  ArrayRef<long> shape = getTensorShape(tensorToSlice);
+  assert("Invalid axis" && axis < shape.size());
+
+  SmallVector<OpFoldResult> strides(shape.size(), rewriter.getIndexAttr(1));
+  SmallVector<OpFoldResult> offsets(shape.size(), rewriter.getIndexAttr(0));
+  SmallVector<OpFoldResult> sizes;
+  sizes.reserve(shape.size());
+  for (const auto size : shape)
+    sizes.push_back(rewriter.getIndexAttr(size));
+  sizes[axis] = rewriter.getIndexAttr(sliceSize);
+
+  long length = shape[axis];
+  auto [numSlices, lastSliceSize] = ceilIntegerDivideWithRemainder(length, sliceSize);
+  SmallVector<Value> slices;
+  slices.reserve(numSlices);
+
+  for (int64_t i = 0; i < numSlices; i++) {
+    offsets[axis] = rewriter.getIndexAttr(i * sliceSize);
+    if (i == numSlices - 1 && lastSliceSize != 0)
+      sizes[axis] = rewriter.getIndexAttr(lastSliceSize);
+
+    Value slice = tensor::ExtractSliceOp::create(rewriter, loc, tensorToSlice, offsets, sizes, strides);
+    slices.push_back(slice);
+  }
+
+  return slices;
+}
+
+SmallVector<Value>
+sliceVector(const Value& vectorToSlice, int64_t sliceSize, ConversionPatternRewriter& rewriter, Location loc) {
+  ArrayRef<long> shape = getTensorShape(vectorToSlice);
+  assert("Not a vector" && isVectorShape(shape));
+  size_t axis = shape[0] != 1 ? 0 : 1;
+  return sliceTensor(vectorToSlice, axis, sliceSize, rewriter, loc);
+}
+
+DenseMap<CoreId, SmallVector<Value>>
+sliceVectorPerCrossbarPerCore(const Value& vectorToSlice, ConversionPatternRewriter& rewriter, Location loc) {
+  SmallVector<Value> slices = sliceVector(vectorToSlice, crossbarSize, rewriter, loc);
+  DenseMap<CoreId, SmallVector<Value>> slicesPerCore;
+  for (size_t sliceId = 0; sliceId < slices.size(); sliceId++) {
+    size_t coreId = sliceId / crossbarCountInCore;
+    slicesPerCore[coreId].push_back(slices[sliceId]);
+  }
+  return slicesPerCore;
+}
+
+DenseMap<HSliceId, DenseMap<CoreId, SmallVector<Value>>> tileMatrix(
+  Value& matrixToTile, int64_t hSliceSize, int64_t vSliceSize, ConversionPatternRewriter& rewriter, Location& loc) {
+  assert("Not a matrix" && isMatrixShape(getTensorShape(matrixToTile)));
+
+  DenseMap<HSliceId, DenseMap<CoreId, SmallVector<Value>>> tiles;
+
+  SmallVector<Value> hSlices = sliceTensor(matrixToTile, 1, hSliceSize, rewriter, loc);
+  size_t numHSlices = hSlices.size();
+  for (size_t hSliceId = 0; hSliceId < numHSlices; hSliceId++) {
+    Value hSlice = hSlices[hSliceId];
+    SmallVector<Value> vSlices = sliceTensor(hSlice, 0, vSliceSize, rewriter, loc);
+    for (size_t vSliceId = 0; vSliceId < vSlices.size(); vSliceId++) {
+      size_t coreId = vSliceId / crossbarCountInCore;
+      Value vSlice = vSlices[vSliceId];
+      tiles[hSliceId][coreId].push_back(vSlice);
+    }
+  }
+  return tiles;
+}
+
+tensor::SplatOp
+broadcastToVector(Value scalarToBroadcast, int64_t length, ConversionPatternRewriter& rewriter, Location loc) {
+  auto oldType = cast<RankedTensorType>(scalarToBroadcast.getType());
+  Type elementType = oldType.getElementType();
+  int64_t shape[2] = {1, length};
+  Type type = oldType.cloneWith(ArrayRef(shape), elementType);
+
+  auto zero = arith::ConstantIndexOp::create(rewriter, loc, 0).getResult();
+  SmallVector<Value> index(oldType.getRank(), zero);
+  auto elementValue = tensor::ExtractOp::create(rewriter, loc, scalarToBroadcast, index).getResult();
+
+  return tensor::SplatOp::create(rewriter, loc, type, elementValue);
+}
+
+} // namespace onnx_mlir