huge refactor for high RewritePatterns usage and less ad-hoc cpp code

remove Spatial many ops in favor of tensor ops like in pim
2026-05-12 10:35:44 +02:00
parent feaff820e1
commit 909c4acfdd
84 changed files with 4048 additions and 3310 deletions
@@ -0,0 +1,136 @@
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+
+#include "src/Accelerators/PIM/Conversion/SpatialToPim/ChannelLoweringPatterns.hpp"
+#include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace {
+
+static int32_t toPimCoreId(int32_t spatialCoreId) { return spatialCoreId; }
+
+struct ChannelSendLowering : OpRewritePattern<spatial::SpatChannelSendOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatChannelSendOp op, PatternRewriter& rewriter) const override {
+    pim::PimSendOp::create(rewriter,
+                           op.getLoc(),
+                           op.getInput(),
+                           getTensorSizeInBytesAttr(rewriter, op.getInput()),
+                           rewriter.getI32IntegerAttr(toPimCoreId(op.getTargetCoreId())));
+    rewriter.eraseOp(op);
+    return success();
+  }
+};
+
+struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatChannelReceiveOp op, PatternRewriter& rewriter) const override {
+    if (op->use_empty()) {
+      rewriter.eraseOp(op);
+      return success();
+    }
+    auto outputType = cast<ShapedType>(op.getResult().getType());
+    Value outputBuffer =
+      tensor::EmptyOp::create(rewriter, op.getLoc(), outputType.getShape(), outputType.getElementType()).getResult();
+    Value received = pim::PimReceiveOp::create(rewriter,
+                                               op.getLoc(),
+                                               op.getResult().getType(),
+                                               outputBuffer,
+                                               getTensorSizeInBytesAttr(rewriter, op.getResult()),
+                                               rewriter.getI32IntegerAttr(toPimCoreId(op.getSourceCoreId())))
+                       .getOutput();
+    rewriter.replaceOp(op, received);
+    return success();
+  }
+};
+
+struct ChannelSendTensorLowering : OpRewritePattern<spatial::SpatChannelSendTensorOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatChannelSendTensorOp op, PatternRewriter& rewriter) const override {
+    SmallVector<int32_t> targetCoreIds;
+    targetCoreIds.reserve(op.getTargetCoreIds().size());
+    for (int32_t targetCoreId : op.getTargetCoreIds())
+      targetCoreIds.push_back(toPimCoreId(targetCoreId));
+    pim::PimSendTensorOp::create(rewriter, op.getLoc(), op.getInput(), rewriter.getDenseI32ArrayAttr(targetCoreIds));
+    rewriter.eraseOp(op);
+    return success();
+  }
+};
+
+struct ChannelReceiveTensorLowering : OpRewritePattern<spatial::SpatChannelReceiveTensorOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatChannelReceiveTensorOp op, PatternRewriter& rewriter) const override {
+    SmallVector<int32_t> sourceCoreIds;
+    sourceCoreIds.reserve(op.getSourceCoreIds().size());
+    for (int32_t sourceCoreId : op.getSourceCoreIds())
+      sourceCoreIds.push_back(toPimCoreId(sourceCoreId));
+    auto outputType = cast<ShapedType>(op.getOutput().getType());
+    Value outputBuffer =
+      tensor::EmptyOp::create(rewriter, op.getLoc(), outputType.getShape(), outputType.getElementType()).getResult();
+    Value received =
+      pim::PimReceiveTensorOp::create(
+        rewriter, op.getLoc(), op.getOutput().getType(), outputBuffer, rewriter.getDenseI32ArrayAttr(sourceCoreIds))
+        .getOutput();
+    rewriter.replaceOp(op, received);
+    return success();
+  }
+};
+
+struct ExtractRowsLowering : OpRewritePattern<spatial::SpatExtractRowsOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatExtractRowsOp op, PatternRewriter& rewriter) const override {
+    auto inputType = cast<RankedTensorType>(op.getInput().getType());
+    SmallVector<Value> replacements;
+    replacements.reserve(op.getNumResults());
+    for (auto [rowIndex, output] : llvm::enumerate(op.getOutputs())) {
+      auto outputType = cast<RankedTensorType>(output.getType());
+      SmallVector<OpFoldResult> offsets = {
+        rewriter.getIndexAttr(static_cast<int64_t>(rowIndex) * outputType.getDimSize(0)), rewriter.getIndexAttr(0)};
+      SmallVector<OpFoldResult> sizes = {rewriter.getIndexAttr(outputType.getDimSize(0)),
+                                         rewriter.getIndexAttr(inputType.getDimSize(1))};
+      SmallVector<OpFoldResult> strides = {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
+      replacements.push_back(
+        tensor::ExtractSliceOp::create(rewriter, op.getLoc(), outputType, op.getInput(), offsets, sizes, strides)
+          .getResult());
+    }
+    rewriter.replaceOp(op, replacements);
+    return success();
+  }
+};
+
+struct ConcatLowering : OpRewritePattern<spatial::SpatConcatOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatConcatOp op, PatternRewriter& rewriter) const override {
+    auto outputType = cast<ShapedType>(op.getOutput().getType());
+    Value outputBuffer =
+      tensor::EmptyOp::create(rewriter, op.getLoc(), outputType.getShape(), outputType.getElementType()).getResult();
+    Value concatenated =
+      pim::PimConcatOp::create(
+        rewriter, op.getLoc(), op.getOutput().getType(), op.getAxisAttr(), op.getInputs(), outputBuffer)
+        .getOutput();
+    rewriter.replaceOp(op, concatenated);
+    return success();
+  }
+};
+
+} // namespace
+
+void populateChannelLoweringPatterns(RewritePatternSet& patterns) {
+  patterns.add<ChannelSendLowering,
+               ChannelReceiveLowering,
+               ChannelSendTensorLowering,
+               ChannelReceiveTensorLowering,
+               ExtractRowsLowering,
+               ConcatLowering>(patterns.getContext());
+}
+
+} // namespace onnx_mlir