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

[ThomasRaoux]

Summary

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

mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp

LogicalResult matchAndRewrite(WarpExecuteOnLane0Op warpOp,
     SmallVector<Value> yieldOperands;
     for (Value operand : forOp.getBody()->getTerminator()->getOperands())
       yieldOperands.push_back(operand);
     rewriter.eraseOp(forOp.getBody()->getTerminator());
     rewriter.mergeBlocks(forOp.getBody(), innerWarp.getBody(), argMapping);
     rewriter.setInsertionPoint(innerWarp.getBody(), innerWarp.getBody()->end());
     rewriter.create<vector::YieldOp>(innerWarp.getLoc(), yieldOperands);
     rewriter.setInsertionPointAfter(innerWarp);
+    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)) {
       warpOp.getResult(res.value())
           .replaceAllUsesWith(newForOp.getResult(res.index()));
       newForOp->setOperand(res.index() + 3, warpOp.getResult(res.value()));
     }
     return success();

mlir/test/Dialect/Vector/vector-warp-distribute.mlir

func.func @warp_scf_for_swap(%arg0: index) {
   }
   "some_use"(%0#0) : (vector<4xf32>) -> ()
   "some_use"(%0#1) : (vector<4xf32>) -> ()
   return
 }
 
 // -----
 
+// 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)>
 
 // CHECK-PROP-LABEL:   func @warp_scf_for_multiple_yield(
 //       CHECK-PROP:   vector.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<1xf32>) {
 //  CHECK-PROP-NEXT:     "some_def"() : () -> vector<32xf32>
 //  CHECK-PROP-NEXT:     vector.yield %{{.*}} : vector<32xf32>