Transpose and Refactor of Patterns

2026-05-29 13:23:31 +02:00
parent 82b44a6387
commit 832bd7f1f7
37 changed files with 285 additions and 153 deletions
@@ -0,0 +1,92 @@
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+
+#include "src/Accelerators/PIM/Conversion/SpatialToPim/Patterns.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 {
+
+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()), 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()),
+                                               op.getSourceCoreId())
+                       .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, ExtractRowsLowering, ConcatLowering>(patterns.getContext());
+}
+
+} // namespace onnx_mlir