Skeleton of miopen.threadwise_copy lowering logic.

Move from -convert-miopen-to-gpu to -miopen-lowering-step4.
llvm · Jun 6, 2020 · 40dd6df · 40dd6df
1 parent 5e336db
commit 40dd6df
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 4 deletions.
diff --git a/mlir/include/mlir/Dialect/MIOpen/LowerMIOpenOps.h b/mlir/include/mlir/Dialect/MIOpen/LowerMIOpenOps.h
@@ -1805,7 +1805,7 @@ struct MovePosRewritePattern : public OpRewritePattern<miopen::MovePosOp> {
 };
 
 //===----------------------------------------------------------------------===//
-// threadwiseGemm lowering.
+// ThreadwiseGemm lowering.
 //===----------------------------------------------------------------------===//
 
 struct ThreadwiseGemmRewritePattern
@@ -1856,6 +1856,21 @@ struct ThreadwiseGemmRewritePattern
   }
 };
 
+//===----------------------------------------------------------------------===//
+// ThreadwiseGemm lowering.
+//===----------------------------------------------------------------------===//
+
+struct ThreadwiseCopyRewritePattern
+    : public OpRewritePattern<miopen::ThreadwiseCopyOp> {
+  using OpRewritePattern<miopen::ThreadwiseCopyOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(miopen::ThreadwiseCopyOp op,
+                                PatternRewriter &b) const override {
+    op.erase();
+    return success();
+  }
+};
+
 //===----------------------------------------------------------------------===//
 // Subview lowering.
 //===----------------------------------------------------------------------===//

diff --git a/mlir/lib/Conversion/MIOpenToGPU/MIOpenToGPU.cpp b/mlir/lib/Conversion/MIOpenToGPU/MIOpenToGPU.cpp
@@ -121,9 +121,6 @@ void LowerMIOpenOpsToGPUPass::runOnOperation() {
 
   for (auto module : op.getOps<gpu::GPUModuleOp>()) {
     module.walk([&](gpu::GPUFuncOp func) {
-
-      func.walk([&](miopen::ThreadwiseCopyOp op) { op.erase(); });
-
       func.walk([&](miopen::GpuAllocOp op) {
         auto loc = op.getLoc();
         auto type = op.output().getType().cast<MemRefType>();

diff --git a/mlir/lib/Dialect/MIOpen/Transforms/LowerMIOpenOps.cpp b/mlir/lib/Dialect/MIOpen/Transforms/LowerMIOpenOps.cpp
@@ -93,6 +93,7 @@ void LowerMIOpenOpsStep3Pass::runOnOperation() {
 void LowerMIOpenOpsStep4Pass::runOnOperation() {
   OwningRewritePatternList patterns;
   patterns.insert<ThreadwiseGemmRewritePattern>(&getContext());
+  patterns.insert<ThreadwiseCopyRewritePattern>(&getContext());
 
   populateAffineToStdConversionPatterns(patterns, &getContext());
   populateLoopToStdConversionPatterns(patterns, &getContext());