Transpose and Refactor of Patterns

src/PIM/Conversion/SpatialToPim/CMakeLists.txt

@@ -3,15 +3,17 @@ mlir_tablegen(SpatialToPim.hpp.inc -gen-rewriters "-I${ONNX_MLIR_SRC_ROOT}")
 add_public_tablegen_target(SpatialToPimIncGen)
 
 add_pim_library(OMSpatialToPim
+  Patterns.cpp
   SpatialToPimPass.cpp
   BatchCoreLoweringPatterns.cpp
-  ChannelLoweringPatterns.cpp
   Common.cpp
   ComputeLikeRegionUtils.cpp
   CoreLoweringPatterns.cpp
-  GlobalTensorMaterialization.cpp
   ReturnPathNormalization.cpp
-  TensorPackingPatterns.cpp
+  Patterns/ChannelLowering.cpp
+  Patterns/GlobalTensorMaterialization.cpp
+  Patterns/TensorPacking.cpp
+  Patterns/Transpose.cpp
 
   EXCLUDE_FROM_OM_LIBS
 
@@ -19,6 +21,7 @@ add_pim_library(OMSpatialToPim
   SpatialToPimIncGen
 
   LINK_LIBS PUBLIC
+  MLIRLinalgDialect
   MLIRSCFDialect
   MLIRSCFUtils
   MLIRTransformUtils

src/PIM/Conversion/SpatialToPim/ChannelLoweringPatterns.hpp

@@ -1,9 +0,0 @@
-#pragma once
-
-#include "mlir/IR/PatternMatch.h"
-
-namespace onnx_mlir {
-
-void populateChannelLoweringPatterns(mlir::RewritePatternSet& patterns);
-
-} // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp

@@ -1,5 +1,6 @@
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/IR/IRMapping.h"
@@ -24,7 +25,7 @@ static bool isChannelUseChainOp(Operation* op) {
              tensor::ExpandShapeOp,
              tensor::CastOp,
              tosa::ReshapeOp,
-             ONNXTransposeOp,
+             linalg::TransposeOp,
              pim::PimTransposeOp>(op);
 }
 

src/PIM/Conversion/SpatialToPim/GlobalTensorMaterialization.hpp

@@ -1,9 +0,0 @@
-#pragma once
-
-#include "mlir/IR/PatternMatch.h"
-
-namespace onnx_mlir {
-
-void populateGlobalTensorMaterializationPatterns(mlir::RewritePatternSet& patterns);
-
-}

src/PIM/Conversion/SpatialToPim/Patterns.cpp

@@ -0,0 +1,40 @@
+#include "src/Accelerators/PIM/Conversion/SpatialToPim/Patterns.hpp"
+#include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace raptor {
+
+#include "src/Accelerators/PIM/Conversion/SpatialToPim/SpatialToPim.hpp.inc"
+
+} // namespace raptor
+
+void populateInitialPatterns(RewritePatternSet& patterns) {
+  raptor::populateWithGenerated(patterns);
+  populateTransposeLoweringPatterns(patterns);
+}
+
+void populateGlobalTensorMaterializationPatternPhase(RewritePatternSet& patterns) {
+  populateGlobalTensorMaterializationPatterns(patterns);
+}
+
+void populateInitialTensorPackingPatterns(RewritePatternSet& patterns) {
+  populateTensorPackingPatterns(patterns);
+}
+
+void populateCoreBodyPatterns(RewritePatternSet& patterns) {
+  raptor::populateWithGenerated(patterns);
+  populateTransposeLoweringPatterns(patterns);
+}
+
+void populateFinalTensorPackingPatterns(RewritePatternSet& patterns) {
+  populateTensorPackingPatterns(patterns);
+}
+
+void populateCommunicationPatterns(RewritePatternSet& patterns) {
+  populateChannelLoweringPatterns(patterns);
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/Patterns.hpp

@@ -8,6 +8,18 @@
 
 namespace onnx_mlir {
 
+void populateInitialPatterns(mlir::RewritePatternSet& patterns);
+void populateGlobalTensorMaterializationPatternPhase(mlir::RewritePatternSet& patterns);
+void populateInitialTensorPackingPatterns(mlir::RewritePatternSet& patterns);
+void populateCoreBodyPatterns(mlir::RewritePatternSet& patterns);
+void populateFinalTensorPackingPatterns(mlir::RewritePatternSet& patterns);
+void populateCommunicationPatterns(mlir::RewritePatternSet& patterns);
+
+void populateTransposeLoweringPatterns(mlir::RewritePatternSet& patterns);
+void populateChannelLoweringPatterns(mlir::RewritePatternSet& patterns);
+void populateGlobalTensorMaterializationPatterns(mlir::RewritePatternSet& patterns);
+void populateTensorPackingPatterns(mlir::RewritePatternSet& patterns);
+
 mlir::RankedTensorType getPackedTensorType(mlir::RankedTensorType elementType, int64_t count);
 mlir::Value extractPackedChunk(mlir::Value packedValue,
                                mlir::RankedTensorType chunkType,
@@ -20,7 +32,6 @@ mlir::Value createPackedExtractRowsSlice(spatial::SpatExtractRowsOp extractRowsO
                                          mlir::OpBuilder& builder,
                                          mlir::Location loc);
 mlir::Value createPackedExtractSliceTensor(mlir::ValueRange values, mlir::OpBuilder& builder, mlir::Location loc);
-void populateTensorPackingPatterns(mlir::RewritePatternSet& patterns);
 void eraseUnusedTensorPackingOps(mlir::func::FuncOp funcOp, mlir::IRRewriter& rewriter);
 
 } // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/Patterns/ChannelLowering.cpp

@@ -1,6 +1,6 @@
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 
-#include "src/Accelerators/PIM/Conversion/SpatialToPim/ChannelLoweringPatterns.hpp"
+#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"

src/PIM/Conversion/SpatialToPim/Patterns/GlobalTensorMaterialization.cpp

@@ -16,7 +16,7 @@
 
 #include "Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Conversion/SpatialToPim/ComputeLikeRegionUtils.hpp"
-#include "src/Accelerators/PIM/Conversion/SpatialToPim/GlobalTensorMaterialization.hpp"
+#include "src/Accelerators/PIM/Conversion/SpatialToPim/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 using namespace mlir;

src/PIM/Conversion/SpatialToPim/Patterns/TensorPacking.cpp

@@ -1,4 +1,4 @@
-#include "src/Accelerators/PIM/Conversion/SpatialToPim/TensorPackingPatterns.hpp"
+#include "src/Accelerators/PIM/Conversion/SpatialToPim/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 
 using namespace mlir;

src/PIM/Conversion/SpatialToPim/Patterns/Transpose.cpp

@@ -0,0 +1,38 @@
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+
+#include "src/Accelerators/PIM/Conversion/SpatialToPim/Patterns.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace {
+
+struct LinalgTransposeToPim final : OpRewritePattern<linalg::TransposeOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(linalg::TransposeOp transposeOp, PatternRewriter& rewriter) const override {
+    SmallVector<Attribute> permutationAttrs;
+    permutationAttrs.reserve(transposeOp.getPermutation().size());
+    for (int64_t dim : transposeOp.getPermutation())
+      permutationAttrs.push_back(rewriter.getI64IntegerAttr(dim));
+
+    auto permutation = rewriter.getArrayAttr(permutationAttrs);
+    auto pimTranspose = pim::PimTransposeOp::create(rewriter,
+                                                    transposeOp.getLoc(),
+                                                    TypeRange {transposeOp->getResult(0).getType()},
+                                                    transposeOp.getInput(),
+                                                    permutation,
+                                                    transposeOp.getInit());
+    rewriter.replaceOp(transposeOp, pimTranspose.getOutput());
+    return success();
+  }
+};
+
+} // namespace
+
+void populateTransposeLoweringPatterns(RewritePatternSet& patterns) {
+  patterns.add<LinalgTransposeToPim>(patterns.getContext());
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/ReturnPathNormalization.cpp

@@ -1,5 +1,6 @@
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
@@ -40,7 +41,7 @@ static bool isReturnHelperChainOp(Operation* op) {
              tensor::ExpandShapeOp,
              tensor::CastOp,
              tosa::ReshapeOp,
-             ONNXTransposeOp,
+             linalg::TransposeOp,
              pim::PimTransposeOp>(op);
 }
 
@@ -276,11 +277,10 @@ static LogicalResult mapIndicesThroughHelperChain(ArrayRef<int64_t> sourceIndice
       continue;
     }
 
-    if (auto transposeOp = dyn_cast<ONNXTransposeOp>(op)) {
+    if (auto transposeOp = dyn_cast<linalg::TransposeOp>(op)) {
       SmallVector<int64_t> nextIndices(currentIndices.size());
       SmallVector<int64_t> nextShape(currentShape.size());
-      for (auto [destIndex, attr] : llvm::enumerate(transposeOp.getPermAttr().getAsRange<IntegerAttr>())) {
-        int64_t sourceIndex = attr.getInt();
+      for (auto [destIndex, sourceIndex] : llvm::enumerate(transposeOp.getPermutation())) {
         nextIndices[destIndex] = currentIndices[sourceIndex];
         nextShape[destIndex] = currentShape[sourceIndex];
       }

src/PIM/Conversion/SpatialToPim/SpatialToPim.td

@@ -9,12 +9,6 @@ include "src/Accelerators/PIM/Dialect/Spatial/Spatial.td"
 include "src/Accelerators/PIM/Dialect/Pim/Pim.td"
 #endif // OP_BASE
 
-def onnxToPimTranspose : Pat<
-  (ONNXTransposeOp:$srcOpRes $data, $perms),
-  (PimTransposeOp $data, $perms,
-    (NativeCodeCall<"onnx_mlir::getBestOutputTensorFromOperandsOrAllocate($_builder, $0.getDefiningOp())"> $srcOpRes))
->;
-
 def spatToPimVMM : Pat<
   (SpatVMMOp:$srcOpRes $weight, $vector),
   (PimVMMOp $weight, $vector,

src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp

@@ -27,10 +27,8 @@
 
 #include "Common/PimCommon.hpp"
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
-#include "Conversion/SpatialToPim/ChannelLoweringPatterns.hpp"
 #include "Conversion/SpatialToPim/Common.hpp"
-#include "Conversion/SpatialToPim/GlobalTensorMaterialization.hpp"
-#include "Conversion/SpatialToPim/TensorPackingPatterns.hpp"
+#include "Conversion/SpatialToPim/Patterns.hpp"
 #include "Dialect/Pim/PimOps.hpp"
 #include "Dialect/Spatial/SpatialOps.hpp"
 #include "Pass/PIMPasses.h"
@@ -41,11 +39,6 @@ using namespace onnx_mlir;
 using namespace pim;
 
 namespace onnx_mlir {
-namespace raptor {
-
-#include "src/Accelerators/PIM/Conversion/SpatialToPim/SpatialToPim.hpp.inc"
-
-} // namespace raptor
 
 static memref::GlobalOp getOrCreateZeroGlobal(IRRewriter& rewriter, Location loc, RankedTensorType tensorType) {
   auto moduleOp = rewriter.getBlock()->getParentOp()->getParentOfType<ModuleOp>();
@@ -159,7 +152,7 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
                     spatial::SpatExtractRowsOp>();
 
   RewritePatternSet initialPatterns(ctx);
-  populateWithGenerated(initialPatterns);
+  populateInitialPatterns(initialPatterns);
   if (failed(applyPartialConversion(moduleOp, target, std::move(initialPatterns)))) {
     moduleOp.emitError("failed to lower required Spatial ops to the initial PIM form");
     signalPassFailure();
@@ -167,7 +160,7 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
   }
 
   RewritePatternSet globalTensorPatterns(ctx);
-  populateGlobalTensorMaterializationPatterns(globalTensorPatterns);
+  populateGlobalTensorMaterializationPatternPhase(globalTensorPatterns);
   walkAndApplyPatterns(moduleOp, std::move(globalTensorPatterns));
 
   auto returnOp = cast<func::ReturnOp>(funcOp.front().getTerminator());
@@ -197,7 +190,7 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
   }
 
   RewritePatternSet initialTensorPackingPatterns(ctx);
-  populateTensorPackingPatterns(initialTensorPackingPatterns);
+  populateInitialTensorPackingPatterns(initialTensorPackingPatterns);
   walkAndApplyPatterns(funcOp, std::move(initialTensorPackingPatterns));
   eraseUnusedTensorPackingOps(funcOp, rewriter);
 
@@ -214,7 +207,7 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
   }
 
   RewritePatternSet coreBodyPatterns(ctx);
-  populateWithGenerated(coreBodyPatterns);
+  populateCoreBodyPatterns(coreBodyPatterns);
   populateAffineToStdConversionPatterns(coreBodyPatterns);
   FrozenRewritePatternSet frozenCoreBodyPatterns(std::move(coreBodyPatterns));
 
@@ -257,7 +250,7 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
   eraseOpsToRemove();
 
   RewritePatternSet finalTensorPackingPatterns(ctx);
-  populateTensorPackingPatterns(finalTensorPackingPatterns);
+  populateFinalTensorPackingPatterns(finalTensorPackingPatterns);
   walkAndApplyPatterns(funcOp, std::move(finalTensorPackingPatterns));
   eraseUnusedTensorPackingOps(funcOp, rewriter);
 
@@ -277,7 +270,7 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
                                    spatial::SpatExtractRowsOp>();
 
   RewritePatternSet communicationPatterns(ctx);
-  populateChannelLoweringPatterns(communicationPatterns);
+  populateCommunicationPatterns(communicationPatterns);
   if (failed(applyFullConversion(funcOp, communicationTarget, std::move(communicationPatterns)))) {
     funcOp.emitOpError("failed to lower Spatial communication ops to PIM communication ops");
     signalPassFailure();