[mlir][vector] Prevent duplicating operations during vector distribute

[ThomasRaoux]

Summary

We should distribute ops that have other uses than the yield op as this
would duplicate those ops.

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

 
 /// Helper to know if an op can be hoisted out of the region.
 static bool canBeHoisted(Operation *op,
                          function_ref<bool(Value)> definedOutside) {
   return llvm::all_of(op->getOperands(), definedOutside) &&
          isMemoryEffectFree(op) && op->getNumRegions() == 0;
 }
 
-/// Return a value yielded by `warpOp` which statifies the filter lamdba
-/// condition and is not dead.
+/// Return a value yielded by `warpOp` with no other uses which statifies the
+/// filter lamdba condition and is not dead.
 static OpOperand *getWarpResult(WarpExecuteOnLane0Op warpOp,
                                 const std::function<bool(Operation *)> &fn) {
   auto yield = cast<vector::YieldOp>(
       warpOp.getBodyRegion().getBlocks().begin()->getTerminator());
   for (OpOperand &yieldOperand : yield->getOpOperands()) {
     Value yieldValues = yieldOperand.get();
     Operation *definedOp = yieldValues.getDefiningOp();
-    if (definedOp && fn(definedOp)) {
+    if (definedOp && definedOp->hasOneUse() && fn(definedOp)) {
       if (!warpOp.getResult(yieldOperand.getOperandNumber()).use_empty())
         return &yieldOperand;
     }
   }
   return {};
 }
 
 // Clones `op` into a new operation that takes `operands` and returns

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

func.func @vector_insert_2d_broadcast(%laneid: index) -> (vector<4x96xf32>) {
   %r = vector.warp_execute_on_lane_0(%laneid)[32] -> (vector<4x96xf32>) {
     %0 = "some_def"() : () -> (vector<4x96xf32>)
     %s = "another_def"() : () -> (vector<96xf32>)
     %1 = vector.insert %s, %0[2] : vector<96xf32> into vector<4x96xf32>
     vector.yield %1 : vector<4x96xf32>
   }
   return %r : vector<4x96xf32>
 }
+// -----
+
+// Verify that we don't duplicate the reduction.
+// CHECK-PROP-LABEL: func @vector_reduction_no_duplicate(
+//  CHECK-PROP-SAME:     %[[laneid:.*]]: index)
+//       CHECK-PROP:   %[[warp_op:.*]] = vector.warp_execute_on_lane_0(%[[laneid]])[32] -> (f32) {
+//       CHECK-PROP:     vector.reduction
+//       CHECK-PROP:     vector.yield %{{.*}} : f32
+//       CHECK-PROP:   }
+//  CHECK-PROP-NEXT:   return %{{.*}} : f32
+func.func @vector_reduction_no_duplicate(%laneid: index) -> (f32) {
+  %r = vector.warp_execute_on_lane_0(%laneid)[32] -> (f32) {
+    %0 = "some_def"() : () -> (vector<32xf32>)
+    %1 = vector.reduction <add>, %0 : vector<32xf32> into f32
+    "some_blocking_use"(%1) : (f32) -> ()
+    vector.yield %1 : f32
+  }
+  return %r : f32
+}
 No newline at end of file