Use miopen.move_pos op in miopen.gridwise_gemm lowering.

llvm · Jun 6, 2020 · da89235 · da89235
1 parent 7e6fa44
commit da89235
Showing 1 changed file with 29 additions and 22 deletions.
diff --git a/mlir/include/mlir/Dialect/MIOpen/LowerMIOpenOps.h b/mlir/include/mlir/Dialect/MIOpen/LowerMIOpenOps.h
@@ -1299,17 +1299,22 @@ struct GridwiseGemmRewritePattern : public OpRewritePattern<miopen::GridwiseGemm
     // LDS barrier.
     lb.create<miopen::LdsBarrierOp>(op.getLoc());
 
-    // TBD add move_pos op.
-    // TBD move blockwiseCopyASrc by (KPerBlock, 0).
-    // TBD move blockwiseCopyBSrc by (KPerBlock, 0).
-    // TBD revive blockwise_copy commands.
-    // // Blockwise copy from global (generic tensor) to register (naive tensor).
-    // auto blockwiseCopyOpAEven = lb.create<miopen::BlockwiseCopyOp>(op.getLoc(), op.getOperand(0), threadAEvenAllocOp);
-    // // Compute block_slice_copy_steps and set in the attribute.
-    // blockwiseCopyOpAEven.setAttr("move_source_slice_window", b.getI32IntegerAttr(KPerBlock));
-    // auto blockwiseCopyOpBEven = lb.create<miopen::BlockwiseCopyOp>(op.getLoc(), op.getOperand(1), threadBEvenAllocOp);
-    // // Compute block_slice_copy_steps and set in the attribute.
-    // blockwiseCopyOpBEven.setAttr("move_source_slice_window", b.getI32IntegerAttr(KPerBlock));
+    auto KPerBlockConstantI32Op =
+        b.create<ConstantIntOp>(op.getLoc(), KPerBlock, b.getIntegerType(32));
+
+    // Blockwise copy from global (generic tensor) to register (naive tensor).
+    lb.create<miopen::MovePosOp>(
+        op.getLoc(), blockwiseCopyASrc,
+        ValueRange{KPerBlockConstantI32Op, zeroConstantI32Op});
+    auto blockwiseCopyOpAEven = lb.create<miopen::BlockwiseCopyOp>(
+        op.getLoc(), op.getOperand(0), threadAEvenAllocOp, blockwiseCopyASrc,
+        blockwiseCopyZero);
+    lb.create<miopen::MovePosOp>(
+        op.getLoc(), blockwiseCopyBSrc,
+        ValueRange{KPerBlockConstantI32Op, zeroConstantI32Op});
+    auto blockwiseCopyOpBEven = lb.create<miopen::BlockwiseCopyOp>(
+        op.getLoc(), op.getOperand(1), threadBEvenAllocOp, blockwiseCopyBSrc,
+        blockwiseCopyZero);
 
     // Emit blockwise GEMM.
     auto blockwiseGemmEvenOp = lb.create<miopen::BlockwiseGemmOp>(
@@ -1326,17 +1331,19 @@ struct GridwiseGemmRewritePattern : public OpRewritePattern<miopen::GridwiseGemm
     // LDS barrier.
     lb.create<miopen::LdsBarrierOp>(op.getLoc());
 
-    // TBD add move_pos op.
-    // TBD move blockwiseCopyASrc by (KPerBlock, 0).
-    // TBD move blockwiseCopyBSrc by (KPerBlock, 0).
-    // TBD revive blockwise_copy commands.
-    // // Blockwise copy from global (generic tensor) to register (naive tensor).
-    // auto blockwiseCopyOpAOdd = lb.create<miopen::BlockwiseCopyOp>(op.getLoc(), op.getOperand(0), threadAOddAllocOp);
-    // // Compute block_slice_copy_steps and set in the attribute.
-    // blockwiseCopyOpAOdd.setAttr("move_source_slice_window", b.getI32IntegerAttr(KPerBlock));
-    // auto blockwiseCopyOpBOdd = lb.create<miopen::BlockwiseCopyOp>(op.getLoc(), op.getOperand(1), threadBOddAllocOp);
-    // // Compute block_slice_copy_steps and set in the attribute.
-    // blockwiseCopyOpBOdd.setAttr("move_source_slice_window", b.getI32IntegerAttr(KPerBlock));
+    // Blockwise copy from global (generic tensor) to register (naive tensor).
+    lb.create<miopen::MovePosOp>(
+        op.getLoc(), blockwiseCopyASrc,
+        ValueRange{KPerBlockConstantI32Op, zeroConstantI32Op});
+    auto blockwiseCopyOpAOdd = lb.create<miopen::BlockwiseCopyOp>(
+        op.getLoc(), op.getOperand(0), threadAOddAllocOp, blockwiseCopyASrc,
+        blockwiseCopyZero);
+    lb.create<miopen::MovePosOp>(
+        op.getLoc(), blockwiseCopyBSrc,
+        ValueRange{KPerBlockConstantI32Op, zeroConstantI32Op});
+    auto blockwiseCopyOpBOdd = lb.create<miopen::BlockwiseCopyOp>(
+        op.getLoc(), op.getOperand(1), threadBOddAllocOp, blockwiseCopyBSrc,
+        blockwiseCopyZero);
 
     // Emit blockwise GEMM.
     auto blockwiseGemmOddOp = lb.create<miopen::BlockwiseGemmOp>(