Cleanup logic getting transfer read warp results

Author: qedawkins

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

@@ -804,29 +804,14 @@ struct WarpOpTransferRead : public OpRewritePattern<WarpExecuteOnLane0Op> {
     // 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>(
-        warpOp.getBodyRegion().getBlocks().begin()->getTerminator());
-    OpOperand *operand;
-    for (OpOperand &yieldOperand : yield->getOpOperands()) {
-      Value yieldValues = yieldOperand.get();
-      Operation *definedOp = yieldValues.getDefiningOp();
-      if (!definedOp)
-        continue;
-      auto maybeRead = dyn_cast<vector::TransferReadOp>(definedOp);
-      if (!maybeRead)
-        continue;
-      if (warpOp.getResult(yieldOperand.getOperandNumber()).use_empty())
-        continue;
+    OpOperand *operand = getWarpResult(warpOp, [](Operation *op) {
       // Don't duplicate transfer_read ops when distributing.
-      if (!maybeRead.getResult().hasOneUse())
-        continue;
-      read = maybeRead;
-      operand = &yieldOperand;
-      break;
-    }
-    if (!read)
+      return isa<vector::TransferReadOp>(op) && op->hasOneUse();
+    });
+    if (!operand)
       return failure();
+    auto read = operand->get().getDefiningOp<vector::TransferReadOp>();
+
     unsigned operandIndex = operand->getOperandNumber();
     Value distributedVal = warpOp.getResult(operandIndex);
 
@@ -933,11 +918,10 @@ struct WarpOpDeadResult : public OpRewritePattern<WarpExecuteOnLane0Op> {
         rewriter, warpOp, newYieldValues, newResultTypes);
 
     // Simplify the new warp op after dropping dead results.
-    auto simplifyFn = [&](Operation *op) {
+    newWarpOp.getBody()->walk([&](Operation *op) {
       if (isOpTriviallyDead(op))
         rewriter.eraseOp(op);
-    };
-    newWarpOp.getBody()->walk(simplifyFn);
+    });
 
     // Replace results of the old warpOp by the new, deduplicated results.
     SmallVector<Value> newValues;