[mlir][vector] Don't duplicate transfer_read during vector distribution

[ThomasRaoux]

Summary

Only apply the pattern if the transfer_read can be distributed for all
its uses.

[nicolasvasilache]

Thanks!

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

struct WarpOpTransferRead : public OpRewritePattern<WarpExecuteOnLane0Op> {
   using OpRewritePattern<WarpExecuteOnLane0Op>::OpRewritePattern;
   LogicalResult matchAndRewrite(WarpExecuteOnLane0Op warpOp,
                                 PatternRewriter &rewriter) const override {
     OpOperand *operand = getWarpResult(
         warpOp, [](Operation *op) { return isa<vector::TransferReadOp>(op); });
     if (!operand)
       return failure();
     auto read = operand->get().getDefiningOp<vector::TransferReadOp>();
+    // Don't duplicate transfer_read ops when distributing.
+    if (!read.getResult().hasOneUse())
+      return failure();
     unsigned operandIndex = operand->getOperandNumber();
     Value distributedVal = warpOp.getResult(operandIndex);
 
     SmallVector<Value, 4> indices(read.getIndices().begin(),
                                   read.getIndices().end());
     AffineMap map = calculateImplicitMap(read.getResult(), distributedVal);
     AffineMap indexMap = map.compose(read.getPermutationMap());
     OpBuilder::InsertionGuard g(rewriter);

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

func.func @lane_dependent_warp_propagate_read(
   %cst = arith.constant 0.000000e+00 : f32
   %r = vector.warp_execute_on_lane_0(%laneid)[32] -> (vector<1x1xf32>) {
     %2 = vector.transfer_read %src[%c0, %c0], %cst : memref<1x1024xf32>, vector<1x32xf32>
     vector.yield %2 : vector<1x32xf32>
   }
   vector.transfer_write %r, %dest[%c0, %laneid] : vector<1x1xf32>, memref<1x1024xf32>
   return
 }
+
+// -----
+
+// CHECK-PROP:   func @dont_duplicate_read
+func.func @dont_duplicate_read(
+  %laneid: index, %src: memref<1024xf32>) -> vector<1xf32> {
+  %c0 = arith.constant 0 : index
+  %cst = arith.constant 0.000000e+00 : f32
+//       CHECK-PROP:   vector.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<1xf32>) {
+//  CHECK-PROP-NEXT:     vector.transfer_read
+//  CHECK-PROP-NEXT:     "blocking_use"
+//  CHECK-PROP-NEXT:     vector.yield
+  %r = vector.warp_execute_on_lane_0(%laneid)[32] -> (vector<1xf32>) {
+    %2 = vector.transfer_read %src[%c0], %cst : memref<1024xf32>, vector<32xf32>
+    "blocking_use"(%2) : (vector<32xf32>) -> ()
+    vector.yield %2 : vector<32xf32>
+  }
+  return %r : vector<1xf32>
+}