[Mlir-commits] [mlir] [mlir][vector] Fix cases with multiple yielded transfer_read ops (PR #71625)

[Quinn Dawkins]

================
@@ -801,13 +801,31 @@ 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())
+    // Try to find a distributable yielded read. Note that this pattern can
+    // still fail at the end after distribution, in which case this might have
+    // missed another distributable read.
+    vector::TransferReadOp read;
+    auto yield = cast<vector::YieldOp>(
----------------
qedawkins wrote:

I realized I can just make the extra checks a part of the filter function; looks a lot cleaner now. There is still a problem where if either the lane fails to delinearize, or if there are leftover producers in the body of the warp op, it might miss out on propagating a different read, but by that point we've already cloned the warp op and attempted to move things over, so we can't really roll back and try again.

https://github.com/llvm/llvm-project/pull/71625