[Mlir-commits] [mlir] d343cdd - [mlir][vector] Fix bug when swapping scf.for and vector warp op

[Thomas Raoux]

Author: Thomas Raoux
Date: 2022-06-24T19:13:02Z
New Revision: d343cdd50992ea3137146bebf32e538b73359144

URL: https://github.com/llvm/llvm-project/commit/d343cdd50992ea3137146bebf32e538b73359144
DIFF: https://github.com/llvm/llvm-project/commit/d343cdd50992ea3137146bebf32e538b73359144.diff

LOG: [mlir][vector] Fix bug when swapping scf.for and vector warp op

When creating a scf.for without argument a scf.yield is automatically
created. Make sure we don't create a second one.

Differential Revision: https://reviews.llvm.org/D128405

Added: 
    

Modified: 
    mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
    mlir/test/Dialect/Vector/vector-warp-distribute.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
index c71c31654105..e8602af2a9e5 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
@@ -717,7 +717,8 @@ struct WarpOpScfForOp : public OpRewritePattern<WarpExecuteOnLane0Op> {
     rewriter.setInsertionPoint(innerWarp.getBody(), innerWarp.getBody()->end());
     rewriter.create<vector::YieldOp>(innerWarp.getLoc(), yieldOperands);
     rewriter.setInsertionPointAfter(innerWarp);
-    rewriter.create<scf::YieldOp>(forOp.getLoc(), innerWarp.getResults());
+    if (!innerWarp.getResults().empty())
+      rewriter.create<scf::YieldOp>(forOp.getLoc(), innerWarp.getResults());
     rewriter.eraseOp(forOp);
     // Replace the warpOp result coming from the original ForOp.
     for (const auto &res : llvm::enumerate(resultIdx)) {

diff  --git a/mlir/test/Dialect/Vector/vector-warp-distribute.mlir b/mlir/test/Dialect/Vector/vector-warp-distribute.mlir
index a2c988cae33e..084a935ce23e 100644
--- a/mlir/test/Dialect/Vector/vector-warp-distribute.mlir
+++ b/mlir/test/Dialect/Vector/vector-warp-distribute.mlir
@@ -387,6 +387,26 @@ func.func @warp_scf_for_swap(%arg0: index) {
 
 // -----
 
+// CHECK-PROP-LABEL:   func @warp_scf_for_swap_no_yield(
+// CHECK-PROP:           scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} {
+// CHECK-PROP-NEXT:        vector.warp_execute_on_lane_0(%{{.*}})[32] {
+// CHECK-PROP-NEXT:          "some_op"() : () -> ()
+// CHECK-PROP-NEXT:        }
+// CHECK-PROP-NEXT:      }
+func.func @warp_scf_for_swap_no_yield(%arg0: index) {
+  %c128 = arith.constant 128 : index
+  %c1 = arith.constant 1 : index
+  %c0 = arith.constant 0 : index
+  vector.warp_execute_on_lane_0(%arg0)[32] {
+    scf.for %arg3 = %c0 to %c128 step %c1 {
+      "some_op"() : () -> ()
+    }
+  }
+  return
+}
+
+// -----
+
 #map = affine_map<()[s0] -> (s0 * 4)>
 #map1 = affine_map<()[s0] -> (s0 * 128 + 128)>
 #map2 = affine_map<()[s0] -> (s0 * 4 + 128)>