diff --git a/include/PTO/Transforms/VMIControlFlowSupport.h b/include/PTO/Transforms/VMIControlFlowSupport.h
index 88e037bc36..94ca126949 100644
--- a/include/PTO/Transforms/VMIControlFlowSupport.h
+++ b/include/PTO/Transforms/VMIControlFlowSupport.h
@@ -12,7 +12,9 @@
 #ifndef PTO_TRANSFORMS_VMICONTROLFLOWSUPPORT_H
 #define PTO_TRANSFORMS_VMICONTROLFLOWSUPPORT_H
 
+#include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Value.h"
 #include "mlir/Support/LLVM.h"
 
@@ -37,6 +39,14 @@ class VMIControlFlowSupport {
       scf::WhileOp whileOp, EquivalenceCallback addEquivalent);
 };
 
+/// Return the result types agreed on by every call to \p func. A null type
+/// marks a result position whose callers currently disagree.
+SmallVector<Type> getConsistentCallResultTypes(ModuleOp module,
+                                               func::FuncOp func);
+
+/// Collect the current argument types used to rebuild a function signature.
+SmallVector<Type> getFunctionInputTypes(func::FuncOp func);
+
 } // namespace pto
 } // namespace mlir
 
diff --git a/include/PTO/Transforms/VMILayoutPropagation.h b/include/PTO/Transforms/VMILayoutPropagation.h
index 131160cd74..0c70a37e23 100644
--- a/include/PTO/Transforms/VMILayoutPropagation.h
+++ b/include/PTO/Transforms/VMILayoutPropagation.h
@@ -68,7 +68,7 @@ class VMILayoutPropagator {
   void enqueue(Value value, VMILayoutAttr layout);
   LogicalResult addUseConflict(OpOperand &operand,
                                VMIValueLayoutAssignment &assignment,
-                               VMILayoutAttr layout);
+                               VMILayoutAttr layout) const;
   LogicalResult propagateFact(Value value, VMILayoutAttr layout);
   LogicalResult propagateOperandFact(OpOperand &operand, VMILayoutAttr layout);
   LogicalResult propagateThrough(Operation *op, Value changedValue,
@@ -81,7 +81,7 @@ class VMILayoutPropagator {
                                    RewriterBase &rewriter,
                                    DenseMap<Value, Value> &assignedValues);
   FailureOr<Value> materializeAt(Value source, VMILayoutAttr layout,
-                                 RewriterBase &rewriter, Location loc);
+                                 RewriterBase &rewriter, Location loc) const;
   LogicalResult materializeUseConflict(Value assignedValue,
                                        VMILayoutConflict conflict,
                                        RewriterBase &rewriter);
diff --git a/lib/PTO/Transforms/PTOInferVPTOVecScope.cpp b/lib/PTO/Transforms/PTOInferVPTOVecScope.cpp
index 5b425582fa..d4bd8453d0 100644
--- a/lib/PTO/Transforms/PTOInferVPTOVecScope.cpp
+++ b/lib/PTO/Transforms/PTOInferVPTOVecScope.cpp
@@ -13,6 +13,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "PTO/Transforms/Passes.h"
+#include "Utils.h"
 
 #include "PTO/Support/CodeConstants.h"
 
@@ -306,15 +307,6 @@ computeMovedOpsForResultlessScope(ArrayRef<Operation *> ops) {
   return movedOps;
 }
 
-static Operation *getAncestorInBlock(Operation *op, Block &block) {
-  for (Operation *cur = op; cur; cur = cur->getParentOp()) {
-    if (cur->getBlock() == &block) {
-      return cur;
-    }
-  }
-  return nullptr;
-}
-
 static FailureOr<Operation *>
 cloneVecScopeProducerForUse(
     Value value, Operation *user, Operation *logicalScopeAnchor,
@@ -481,7 +473,7 @@ static LogicalResult rematerializeEscapingValueForUserSegments(
       continue;
     }
 
-    Operation *ancestor = getAncestorInBlock(user, block);
+    Operation *ancestor = pto::getAncestorInBlock(user, &block);
     if (!ancestor) {
       return failure();
     }
@@ -804,7 +796,7 @@ static LogicalResult repairEscapingSubclusters(Block &block,
       ops.push_back(&op);
     }
 
-    auto flush = [&]() -> FailureOr<bool> {
+    auto flush = [&pending, &cache, context]() -> FailureOr<bool> {
       FailureOr<bool> changed =
           fixOneEscapingSubcluster(pending, cache, context);
       pending.clear();
@@ -854,7 +846,7 @@ static LogicalResult inferVecScopesInBlock(Block &block, MLIRContext *context) {
 
   SmallVector<Operation *, mlir::pto::kValue16> pending;
 
-  auto flush = [&]() -> LogicalResult {
+  auto flush = [&pending, context]() -> LogicalResult {
     if (failed(wrapGreedySubclusters(pending, context))) {
       return failure();
     }
diff --git a/lib/PTO/Transforms/PTORemoveIdentityTMov.cpp b/lib/PTO/Transforms/PTORemoveIdentityTMov.cpp
index da811b4c9e..63456f1604 100644
--- a/lib/PTO/Transforms/PTORemoveIdentityTMov.cpp
+++ b/lib/PTO/Transforms/PTORemoveIdentityTMov.cpp
@@ -16,6 +16,7 @@
 #include "PTO/Transforms/InsertSync/PTOIRTranslator.h"
 #include "PTO/Transforms/InsertSync/SyncCommon.h"
 #include "PTO/Transforms/Passes.h"
+#include "Utils.h"
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
@@ -258,15 +259,6 @@ static bool hasSameConcreteAddressRange(const BaseMemInfo *srcInfo,
   return srcRootAddr && dstRootAddr && *srcRootAddr == *dstRootAddr;
 }
 
-static Operation *getAncestorInBlock(Operation *op, Block *block) {
-  for (Operation *cur = op; cur; cur = cur->getParentOp()) {
-    if (cur->getBlock() == block) {
-      return cur;
-    }
-  }
-  return nullptr;
-}
-
 static bool hasUseAfterOp(Value value, Operation *currentOp) {
   Block *block = currentOp->getBlock();
   for (OpOperand &use : value.getUses()) {
@@ -274,7 +266,7 @@ static bool hasUseAfterOp(Value value, Operation *currentOp) {
     if (owner == currentOp) {
       continue;
     }
-    Operation *ancestor = getAncestorInBlock(owner, block);
+    Operation *ancestor = pto::getAncestorInBlock(owner, block);
     if (!ancestor) {
       return true;
     }
@@ -293,7 +285,7 @@ static bool hasLaterUseOfSameAddressRange(
     if (entry.first == op.getSrc()) {
       continue;
     }
-    bool sameRange = llvm::any_of(entry.second, [&](const auto &info) {
+    bool sameRange = llvm::any_of(entry.second, [dstInfo](const auto &info) {
       return hasSameConcreteAddressRange(info.get(), dstInfo);
     });
     if (sameRange && hasUseAfterOp(entry.first, op)) {
@@ -371,7 +363,7 @@ struct PTORemoveIdentityTMovPass
     SmallVector<TMovOp, mlir::pto::kValue16> identityMoves;
     SmallVector<TMovOp, mlir::pto::kValue16> memInfoCandidates;
 
-    func.walk([&](TMovOp op) {
+    func.walk([&identityMoves, &memInfoCandidates](TMovOp op) {
       if (!hasPlainTMovSemantics(op) || !hasCompatibleIdentityTypes(op) ||
           touchesLowPrecisionElement(op)) {
         return;
diff --git a/lib/PTO/Transforms/PTOVerifyTFreePass.cpp b/lib/PTO/Transforms/PTOVerifyTFreePass.cpp
index 21eabcdcce..f6905a8c04 100644
--- a/lib/PTO/Transforms/PTOVerifyTFreePass.cpp
+++ b/lib/PTO/Transforms/PTOVerifyTFreePass.cpp
@@ -13,6 +13,7 @@
 
 #include "PTO/IR/PTO.h"
 #include "PTO/Transforms/Passes.h"
+#include "Utils.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Pass/Pass.h"
 
@@ -43,22 +44,10 @@ static TFreeOp findMatchingTFree(TPopOp tpopOp) {
   return {};
 }
 
-static Operation *getTopLevelAncestorInBlock(Operation *op, Block *block) {
-  Operation *current = op;
-  while (current && current->getBlock() != block) {
-    Region *parentRegion = current->getParentRegion();
-    if (!parentRegion) {
-      return nullptr;
-    }
-    current = parentRegion->getParentOp();
-  }
-  return current;
-}
-
 static bool hasSamePipeTPopInRegion(Operation *op, Value pipeHandle,
                                     TPopOp current) {
   bool found = false;
-  op->walk([&](TPopOp nestedTpop) {
+  op->walk([current, pipeHandle, &found](TPopOp nestedTpop) {
     if (nestedTpop == current) {
       return WalkResult::advance();
     }
@@ -100,7 +89,7 @@ static LogicalResult verifyNoTileUsesAfterTFree(TPopOp tpopOp,
   Block *block = tpopOp->getBlock();
 
   for (OpOperand &use : tile.getUses()) {
-    Operation *topLevelOwner = getTopLevelAncestorInBlock(use.getOwner(), block);
+    Operation *topLevelOwner = pto::getAncestorInBlock(use.getOwner(), block);
     if (!topLevelOwner) {
       return tpopOp.emitOpError(
           "borrowed tile uses must stay in the same parent block as the producing tpop");
@@ -129,7 +118,7 @@ struct PTOVerifyTFreePass
     func::FuncOp funcOp = getOperation();
 
     SmallVector<TPopOp> tpops;
-    funcOp.walk([&](TPopOp op) { tpops.push_back(op); });
+    funcOp.walk([&tpops](TPopOp op) { tpops.push_back(op); });
 
     for (TPopOp tpopOp : tpops) {
       if (!isInsideSectionOrAttributedKernel(tpopOp, funcOp)) {
diff --git a/lib/PTO/Transforms/TileFusion/FusionAnalysis.cpp b/lib/PTO/Transforms/TileFusion/FusionAnalysis.cpp
index d3cb213a99..ba612e6d71 100644
--- a/lib/PTO/Transforms/TileFusion/FusionAnalysis.cpp
+++ b/lib/PTO/Transforms/TileFusion/FusionAnalysis.cpp
@@ -835,10 +835,11 @@ static void recordLastLocalConsumer(std::optional<unsigned> &lastLocalConsumer,
   }
 }
 
-static void finalizeBlockLiveness(
-    Block &block, DenseMap<Operation *, FusionOpKind> &kindByOp,
-    DenseMap<Operation *, unsigned> &computeNodeByOp,
-    SmallVectorImpl<MutableLiveness> &mutableLiveness) {
+static void
+finalizeBlockLiveness(const Block &block,
+                      DenseMap<Operation *, FusionOpKind> &kindByOp,
+                      DenseMap<Operation *, unsigned> &computeNodeByOp,
+                      SmallVectorImpl<MutableLiveness> &mutableLiveness) {
   for (MutableLiveness &state : mutableLiveness) {
     for (OpOperand &use : state.live.value.getUses()) {
       Operation *user = use.getOwner();
@@ -947,7 +948,7 @@ static void recordWriteInstanceUse(Operation *user,
   }
 }
 
-static void finalizeWriteInstanceUse(Block &block,
+static void finalizeWriteInstanceUse(const Block &block,
                                      const MutableLiveness &storageState,
                                      OpOperand &use,
                                      DFGConstructionState &state) {
diff --git a/lib/PTO/Transforms/TileFusion/PTOFusionLoadStoreElision.cpp b/lib/PTO/Transforms/TileFusion/PTOFusionLoadStoreElision.cpp
index 1acf3678e2..289f7997a1 100644
--- a/lib/PTO/Transforms/TileFusion/PTOFusionLoadStoreElision.cpp
+++ b/lib/PTO/Transforms/TileFusion/PTOFusionLoadStoreElision.cpp
@@ -9,6 +9,7 @@
 #include "PTO/Support/CodeConstants.h"
 #include "PTO/IR/PTO.h"
 #include "PTO/Transforms/Passes.h"
+#include "../Utils.h"
 
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
@@ -218,16 +219,8 @@ static Value getCanonicalTrackedValue(Value value) {
   return value;
 }
 
-static Operation *getTopLevelAncestorInBlock(Operation *op, Block *block) {
-  for (Operation *cur = op; cur; cur = cur->getParentOp()) {
-    if (cur->getBlock() == block) {
-      return cur;
-    }
-  }
-  return nullptr;
-}
-
-static Region *getDirectRegionUnderAncestor(Operation *op, Operation *ancestor) {
+static Region *getDirectRegionUnderAncestor(Operation *op,
+                                            const Operation *ancestor) {
   for (Operation *cur = op; cur; cur = cur->getParentOp()) {
     Operation *parent = cur->getParentOp();
     if (parent == ancestor) {
@@ -402,7 +395,7 @@ static void pruneTrackedStoresForLoadBase(SmallVectorImpl<TrackedStore> &stores,
     stores.clear();
     return;
   }
-  llvm::erase_if(stores, [&](const TrackedStore &store) {
+  llvm::erase_if(stores, [base](const TrackedStore &store) {
     return areEquivalentValues(store.base, base);
   });
 }
@@ -415,7 +408,7 @@ static bool isTailStoreUseCompatible(
     return true;
   }
   if (context.regionOp->isProperAncestor(owner)) {
-    Operation *topLevelUser = getTopLevelAncestorInBlock(owner, context.body);
+    Operation *topLevelUser = pto::getAncestorInBlock(owner, context.body);
     if (!topLevelUser) {
       return false;
     }
@@ -429,7 +422,7 @@ static bool isTailStoreUseCompatible(
   }
 
   Operation *topLevelUser =
-      getTopLevelAncestorInBlock(owner, context.parentBlock);
+      pto::getAncestorInBlock(owner, context.parentBlock);
   if (!topLevelUser) {
     return areMutuallyExclusiveByIfRegion(localScopeOp, owner);
   }
@@ -557,7 +550,7 @@ using RegionContextMap =
 
 static RegionContextMap buildRegionContexts(func::FuncOp func) {
   RegionContextMap contexts;
-  func.walk([&](pto::FusionRegionOp fusionRegion) {
+  func.walk([&contexts](pto::FusionRegionOp fusionRegion) {
     std::optional<FusionRegionStoreContext> context =
         buildFusionRegionStoreContext(fusionRegion);
     if (context) {
@@ -570,15 +563,16 @@ static RegionContextMap buildRegionContexts(func::FuncOp func) {
 static void elideFusionRegionBodies(func::FuncOp func,
                                     RegionContextMap &contexts,
                                     bool &changed) {
-  func.walk([&](pto::FusionRegionOp fusionRegion) {
+  func.walk([&contexts, &changed](pto::FusionRegionOp fusionRegion) {
     auto it = contexts.find(fusionRegion.getOperation());
     if (it == contexts.end()) {
       return;
     }
     Block &body = fusionRegion.getBody().front();
     if (isSupportedStraightLineBlock(body)) {
-      changed |=
-          elideLoadStoreRoundTripsInLeafBody(body, &it->second, nullptr);
+      changed =
+          elideLoadStoreRoundTripsInLeafBody(body, &it->second, nullptr) ||
+          changed;
     }
   });
 }
@@ -591,14 +585,15 @@ static void runElisionForLeafBody(Block *body, Operation *scopeOp,
   }
   auto it = contexts.find(fusionRegion.getOperation());
   if (it != contexts.end()) {
-    changed |= elideLoadStoreRoundTripsInLeafBody(*body, &it->second, scopeOp);
+    changed = elideLoadStoreRoundTripsInLeafBody(*body, &it->second, scopeOp) ||
+              changed;
   }
 }
 
 template <typename ScopeOp>
 static void elideVectorScopeBodies(func::FuncOp func, RegionContextMap &contexts,
                                    bool &changed) {
-  func.walk([&](ScopeOp scope) {
+  func.walk([&contexts, &changed](ScopeOp scope) {
     auto fusionRegion = scope->template getParentOfType<pto::FusionRegionOp>();
     if (fusionRegion && isSupportedStraightLineBlock(scope.getBody().front())) {
       runElisionForLeafBody(&scope.getBody().front(), scope, fusionRegion,
@@ -609,7 +604,7 @@ static void elideVectorScopeBodies(func::FuncOp func, RegionContextMap &contexts
 
 static void elideLoopBodies(func::FuncOp func, RegionContextMap &contexts,
                             bool &changed) {
-  func.walk([&](scf::ForOp loop) {
+  func.walk([&contexts, &changed](scf::ForOp loop) {
     if (!isSupportedLoopRoot(loop)) {
       return;
     }
diff --git a/lib/PTO/Transforms/TileFusion/PTOFusionPredicateElision.cpp b/lib/PTO/Transforms/TileFusion/PTOFusionPredicateElision.cpp
index eb92786f53..efd272c17b 100644
--- a/lib/PTO/Transforms/TileFusion/PTOFusionPredicateElision.cpp
+++ b/lib/PTO/Transforms/TileFusion/PTOFusionPredicateElision.cpp
@@ -326,16 +326,19 @@ buildFusionRegionPredicateContext(pto::FusionRegionOp fusionRegion,
   FusionRegionPredicateContext context;
   context.fusionRegion = fusionRegion;
 
-  fusionRegion.walk([&](Operation *op) -> WalkResult {
-    if (op != fusionRegion.getOperation() && isa<pto::FusionRegionOp>(op)) {
-      return WalkResult::skip();
-    }
-
-    if (std::optional<PltCandidate> candidate = buildPltCandidate(op)) {
-      context.pltCandidates.push_back(std::move(*candidate));
-    }
-    return WalkResult::advance();
-  });
+  fusionRegion.walk(
+      [fusionRegion, &context](Operation *op) mutable -> WalkResult {
+        const bool isNestedFusionRegion =
+            op != fusionRegion.getOperation() && isa<pto::FusionRegionOp>(op);
+        if (isNestedFusionRegion) {
+          return WalkResult::skip();
+        }
+
+        if (std::optional<PltCandidate> candidate = buildPltCandidate(op)) {
+          context.pltCandidates.push_back(std::move(*candidate));
+        }
+        return WalkResult::advance();
+      });
 
   populateDominatingCandidateIndices(context.pltCandidates, dominanceInfo);
   return context;
@@ -417,7 +420,8 @@ struct PTOFusionPredicateElisionPass
 
     DominanceInfo &dominanceInfo = getAnalysis<DominanceInfo>();
     SmallVector<FusionRegionPredicateContext, mlir::pto::kValue4> fusionContexts;
-    func.walk([&](pto::FusionRegionOp fusionRegion) {
+    func.walk([&dominanceInfo,
+               &fusionContexts](pto::FusionRegionOp fusionRegion) {
       FusionRegionPredicateContext context =
           buildFusionRegionPredicateContext(fusionRegion, dominanceInfo);
       if (!context.pltCandidates.empty()) {
@@ -427,7 +431,7 @@ struct PTOFusionPredicateElisionPass
 
     bool changed = false;
     for (FusionRegionPredicateContext &context : fusionContexts) {
-      changed |= elideEquivalentPltCandidates(context);
+      changed = elideEquivalentPltCandidates(context) || changed;
     }
 
     if (!changed) {
diff --git a/lib/PTO/Transforms/TileFusion/PTOFusionRegionGen.cpp b/lib/PTO/Transforms/TileFusion/PTOFusionRegionGen.cpp
index 5fac672850..93b09f9fac 100644
--- a/lib/PTO/Transforms/TileFusion/PTOFusionRegionGen.cpp
+++ b/lib/PTO/Transforms/TileFusion/PTOFusionRegionGen.cpp
@@ -10,6 +10,7 @@
 #include "PTO/IR/PTO.h"
 #include "PTO/Transforms/Passes.h"
 #include "PTO/Transforms/TileFusion/FusionAnalysis.h"
+#include "../Utils.h"
 
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/Builders.h"
@@ -154,7 +155,7 @@ collectGroupSpansInBlock(Block &block, SmallVectorImpl<GroupSpan> &spans) {
   return flushGroupSpan(block, spanIndexByGroupId, current, spans);
 }
 
-static bool isNestedInOp(Operation *op, Operation *ancestor) {
+static bool isNestedInOp(Operation *op, const Operation *ancestor) {
   for (Operation *cur = op; cur; cur = cur->getParentOp()) {
     if (cur == ancestor) {
       return true;
@@ -179,18 +180,9 @@ static void appendUniqueValue(SmallVectorImpl<Value> &values,
   }
 }
 
-static Operation *getTopLevelAncestorInBlock(Operation *op, Block *block) {
-  for (Operation *cur = op; cur; cur = cur->getParentOp()) {
-    if (cur->getBlock() == block) {
-      return cur;
-    }
-  }
-  return nullptr;
-}
-
 static bool canReplaceUseWithRegionResult(OpOperand &use, Operation *boundary) {
   Operation *topLevel =
-      getTopLevelAncestorInBlock(use.getOwner(), boundary->getBlock());
+      pto::getAncestorInBlock(use.getOwner(), boundary->getBlock());
   if (!topLevel || topLevel == boundary) {
     return false;
   }
@@ -403,7 +395,8 @@ static void replaceEscapingUsesOutsideRegion(pto::FusionRegionOp fusionRegion,
   for (auto [oldValueRef, newValue] :
        llvm::zip(oldValues, fusionRegion.getOutputs())) {
     Value oldValue = oldValueRef;
-    oldValue.replaceUsesWithIf(newValue, [&](OpOperand &use) {
+    oldValue.replaceUsesWithIf(newValue, [fusionRegion](
+                                             OpOperand &use) mutable {
       return !isNestedInOp(use.getOwner(), fusionRegion.getOperation()) &&
              canReplaceUseWithRegionResult(use, fusionRegion.getOperation());
     });
diff --git a/lib/PTO/Transforms/TileFusion/PTOLowLevelLoopFusion.cpp b/lib/PTO/Transforms/TileFusion/PTOLowLevelLoopFusion.cpp
index 674467f31b..0d143448e8 100644
--- a/lib/PTO/Transforms/TileFusion/PTOLowLevelLoopFusion.cpp
+++ b/lib/PTO/Transforms/TileFusion/PTOLowLevelLoopFusion.cpp
@@ -291,7 +291,7 @@ static LogicalResult collectAliasRelevantRoots(
 }
 
 static bool containsEquivalentRoot(ArrayRef<Value> roots, Value candidate) {
-  return llvm::any_of(roots, [&](Value root) {
+  return llvm::any_of(roots, [candidate](Value root) {
     return areEquivalentValues(root, candidate);
   });
 }
@@ -303,26 +303,21 @@ static bool canMoveAcrossOperations(Operation *movableOp,
                                     ArrayRef<Value> movableRoots,
                                     ArrayRef<Operation *> crossedOps,
                                     StringRef crossedKind,
-                                    llvm::raw_ostream *debugOS) {
+                                    llvm::raw_ostream &debugOS) {
   for (Operation *op : crossedOps) {
     SmallVector<Value, mlir::pto::kValue4> opRoots;
     if (failed(collectAliasRelevantRoots(op, opRoots))) {
-      if (debugOS) {
-        *debugOS << "[op-fusion] reject movable op " << movableOp->getName()
-                 << " at " << movableOp->getLoc() << ": crossed effects of "
-                 << op->getName() << " are not alias-analyzable\n";
-      }
+      debugOS << "[op-fusion] reject movable op " << movableOp->getName()
+              << " at " << movableOp->getLoc() << ": crossed effects of "
+              << op->getName() << " are not alias-analyzable\n";
       return false;
     }
-    if (llvm::any_of(movableRoots, [&](Value root) {
+    if (llvm::any_of(movableRoots, [&opRoots](Value root) {
           return containsEquivalentRoot(opRoots, root);
         })) {
-      if (debugOS) {
-        *debugOS << "[op-fusion] reject movable op " << movableOp->getName()
-                 << " at " << movableOp->getLoc()
-                 << ": touched root may alias a crossed " << crossedKind
-                 << "\n";
-      }
+      debugOS << "[op-fusion] reject movable op " << movableOp->getName()
+              << " at " << movableOp->getLoc()
+              << ": touched root may alias a crossed " << crossedKind << "\n";
       return false;
     }
   }
@@ -341,14 +336,15 @@ static bool canMoveAcrossStages(Operation *movableOp,
     }
     return false;
   }
+  llvm::raw_ostream &diagnosticStream = debugOS ? *debugOS : llvm::nulls();
   for (const StageInfo &crossStage : crossStages) {
     if (!canMoveAcrossOperations(movableOp, roots, crossStage.leafOps,
-                                 "stage memory op", debugOS)) {
+                                 "stage memory op", diagnosticStream)) {
       return false;
     }
     for (const LoopLevelInfo &level : crossStage.levels) {
       if (!canMoveAcrossOperations(movableOp, roots, level.epilogueOps,
-                                   "stage epilogue op", debugOS)) {
+                                   "stage epilogue op", diagnosticStream)) {
         return false;
       }
     }
@@ -451,7 +447,7 @@ static LogicalResult analyzeStage(scf::ForOp outerLoop, StageInfo &stage) {
 static bool appendStage(scf::ForOp loop,
                         SmallVectorImpl<Operation *> &pendingSetup,
                         SmallVectorImpl<StageInfo> &stages,
-                        llvm::raw_ostream *debugOS) {
+                        llvm::raw_ostream &debugOS) {
   StageInfo stage;
   stage.setupOps.assign(pendingSetup.begin(), pendingSetup.end());
   pendingSetup.clear();
@@ -459,10 +455,8 @@ static bool appendStage(scf::ForOp loop,
     stages.push_back(std::move(stage));
     return true;
   }
-  if (debugOS) {
-    *debugOS << "[op-fusion] stop stage run before " << loop.getLoc()
-             << ": next stage analysis failed\n";
-  }
+  debugOS << "[op-fusion] stop stage run before " << loop.getLoc()
+          << ": next stage analysis failed\n";
   return false;
 }
 
@@ -481,9 +475,10 @@ static SmallVector<StageInfo, mlir::pto::kValue8> collectStageRunFrom(scf::ForOp
   stages.push_back(std::move(firstStage));
 
   SmallVector<Operation *, mlir::pto::kValue4> pendingSetup;
+  llvm::raw_ostream &diagnosticStream = debugOS ? *debugOS : llvm::nulls();
   for (Operation *op = firstLoop->getNextNode(); op; op = op->getNextNode()) {
     if (auto nextLoop = dyn_cast<scf::ForOp>(op)) {
-      if (!appendStage(nextLoop, pendingSetup, stages, debugOS)) {
+      if (!appendStage(nextLoop, pendingSetup, stages, diagnosticStream)) {
         break;
       }
       continue;
@@ -735,8 +730,10 @@ struct PTOLowLevelLoopFusionPass
       }
 
       bool changed = false;
-      func.walk([&](pto::FusionRegionOp fusionRegion) {
-        changed |= fuseStageRunsInBlock(fusionRegion.getBody().front(), traceOS);
+      func.walk([traceOS, &changed](pto::FusionRegionOp fusionRegion) {
+        changed =
+            fuseStageRunsInBlock(fusionRegion.getBody().front(), traceOS) ||
+            changed;
       });
       if (changed) {
         ++fusedFuncs;
diff --git a/lib/PTO/Transforms/TileFusion/PTOMarkLastUse.cpp b/lib/PTO/Transforms/TileFusion/PTOMarkLastUse.cpp
index 9a6864a3df..9da7c315df 100644
--- a/lib/PTO/Transforms/TileFusion/PTOMarkLastUse.cpp
+++ b/lib/PTO/Transforms/TileFusion/PTOMarkLastUse.cpp
@@ -177,7 +177,7 @@ collectGroupSpansInBlock(Block &block, SmallVectorImpl<GroupSpan> &spans) {
 }
 
 static bool isSpanLocalLastUseCandidate(Value value, Operation *currentOp,
-                                        Block *block) {
+                                        const Block *block) {
   if (!value) {
     return false;
   }
@@ -197,7 +197,8 @@ static bool isSpanLocalLastUseCandidate(Value value, Operation *currentOp,
   return true;
 }
 
-static bool hasLaterUseAfterSpan(Value value, Operation *spanEnd, Block *block) {
+static bool hasLaterUseAfterSpan(Value value, Operation *spanEnd,
+                                 const Block *block) {
   for (OpOperand &use : value.getUses()) {
     Operation *user = use.getOwner();
     if (user->getBlock() != block) {
diff --git a/lib/PTO/Transforms/TileFusion/PTOOpScheduling.cpp b/lib/PTO/Transforms/TileFusion/PTOOpScheduling.cpp
index 7c9116226f..3a414b75bf 100644
--- a/lib/PTO/Transforms/TileFusion/PTOOpScheduling.cpp
+++ b/lib/PTO/Transforms/TileFusion/PTOOpScheduling.cpp
@@ -130,7 +130,8 @@ static bool hasTileDependency(Operation *opA, Operation *opB) {
          sharesAnyValue(a.tileOutputs, b.tileOutputs);
 }
 
-static bool crossesOperandDefinition(Operation *movingOp, Operation *candidate) {
+static bool crossesOperandDefinition(Operation *movingOp,
+                                     const Operation *candidate) {
   for (Value operand : movingOp->getOperands()) {
     Operation *defOp = operand.getDefiningOp();
     if (defOp == candidate) {
@@ -269,8 +270,9 @@ collectScheduledGroups(Block &block, SmallVectorImpl<ScheduledGroup> &groups) {
   return validateScheduledGroups(groups);
 }
 
-static bool canPrefixMoveLaterAcross(
-    ArrayRef<GroupMember> members, Operation *placement, Operation *barrier) {
+static bool canPrefixMoveLaterAcross(ArrayRef<GroupMember> members,
+                                     const Operation *placement,
+                                     Operation *barrier) {
   for (const GroupMember &prevMember : members) {
     if (!canMoveLaterAcross(prevMember.op, barrier)) {
       return false;
@@ -283,7 +285,7 @@ static bool canPrefixMoveLaterAcross(
 }
 
 static void movePrefixPastBarrier(ArrayRef<GroupMember> members,
-                                  Operation *placement,
+                                  const Operation *placement,
                                   Operation *barrier) {
   Operation *anchor = barrier;
   for (const GroupMember &prevMember : members) {
diff --git a/lib/PTO/Transforms/TileFusion/PTOPrintPreFusionAnalysis.cpp b/lib/PTO/Transforms/TileFusion/PTOPrintPreFusionAnalysis.cpp
index 24e2797dd1..13efc2adcb 100644
--- a/lib/PTO/Transforms/TileFusion/PTOPrintPreFusionAnalysis.cpp
+++ b/lib/PTO/Transforms/TileFusion/PTOPrintPreFusionAnalysis.cpp
@@ -87,7 +87,7 @@ static StringRef stringifyWriteInstanceEscapeClass(
 static void appendIndexList(llvm::raw_ostream &os, ArrayRef<unsigned> values) {
   os << "[";
   for (auto [idx, value] : llvm::enumerate(values)) {
-    if (idx) {
+    if (idx != 0) {
       os << ", ";
     }
     os << value;
@@ -154,12 +154,14 @@ buildValueLabels(Block &block, const pto::FusionBlockAnalysis &analysis) {
 
     if (semanticsOr->kind == pto::FusionOpKind::LocalBoundary) {
       for (Value input : semanticsOr->tileInputs) {
-        if (!labels.count(input)) {
+        const bool isUnlabeledInput = labels.count(input) == 0;
+        if (isUnlabeledInput) {
           labels.try_emplace(input, makeExternalValueLabel(externalOrdinal++));
         }
       }
       for (Value output : semanticsOr->tileOutputs) {
-        if (!labels.count(output)) {
+        const bool isUnlabeledOutput = labels.count(output) == 0;
+        if (isUnlabeledOutput) {
           labels.try_emplace(output, makeBoundaryValueLabel(boundaryOrdinal++));
         }
       }
@@ -167,7 +169,8 @@ buildValueLabels(Block &block, const pto::FusionBlockAnalysis &analysis) {
     }
 
     for (Value input : semanticsOr->tileInputs) {
-      if (!labels.count(input)) {
+      const bool isUnlabeledInput = labels.count(input) == 0;
+      if (isUnlabeledInput) {
         labels.try_emplace(input, makeExternalValueLabel(externalOrdinal++));
       }
     }
@@ -190,14 +193,14 @@ static void printLocalBoundaries(llvm::raw_ostream &os, Block &block,
     os << "    local_boundary[" << boundaryId++ << "] op="
        << semanticsOr->opName << " inputs=[";
     for (auto [idx, input] : llvm::enumerate(semanticsOr->tileInputs)) {
-      if (idx) {
+      if (idx != 0) {
         os << ", ";
       }
       os << valueLabels.lookup(input);
     }
     os << "] outputs=[";
     for (auto [idx, output] : llvm::enumerate(semanticsOr->tileOutputs)) {
-      if (idx) {
+      if (idx != 0) {
         os << ", ";
       }
       os << valueLabels.lookup(output);
@@ -228,11 +231,11 @@ static void printComputeNodes(llvm::raw_ostream &os,
        << " family=" << stringifyComputeFamily(node.semantics.computeFamily)
        << " domain_class=" << node.iterationDomainClass << " inputs=[";
     for (auto [idx, input] : llvm::enumerate(node.semantics.tileInputs)) {
-      os << (idx ? ", " : "") << labels.lookup(input);
+      os << (idx != 0 ? ", " : "") << labels.lookup(input);
     }
     os << "] outputs=[";
     for (auto [idx, output] : llvm::enumerate(node.semantics.tileOutputs)) {
-      os << (idx ? ", " : "") << labels.lookup(output);
+      os << (idx != 0 ? ", " : "") << labels.lookup(output);
     }
     os << "] incoming=";
     appendIndexList(os, node.incomingEdges);
diff --git a/lib/PTO/Transforms/Utils.cpp b/lib/PTO/Transforms/Utils.cpp
index 63fa540a6f..f22a1eb89e 100644
--- a/lib/PTO/Transforms/Utils.cpp
+++ b/lib/PTO/Transforms/Utils.cpp
@@ -165,6 +165,16 @@ func::ReturnOp getAssumedUniqueReturnOp(func::FuncOp funcOp) {
   return returnOp;
 }
 
+Operation *getAncestorInBlock(Operation *op, const Block *block) {
+  for (Operation *current = op; current; current = current->getParentOp()) {
+    bool isInBlock = current->getBlock() == block;
+    if (isInBlock) {
+      return current;
+    }
+  }
+  return nullptr;
+}
+
 Value peelUnrealized(Value value) {
   if (auto castOp = value.getDefiningOp<UnrealizedConversionCastOp>()) {
     return castOp.getOperand(0);
@@ -828,14 +838,16 @@ static bool rangesOverlap(const SemanticRange &lhs, const SemanticRange &rhs) {
 
 LogicalResult verifySemanticNoAliasRanges(func::FuncOp func) {
   LogicalResult result = success();
-  func.walk([&](Operation *op) {
-    if (failed(result))
+  func.walk([&result](Operation *op) {
+    if (failed(result)) {
       return;
+    }
     for (auto [lhs, rhs] : getSemanticNoAliasPairs(op)) {
       auto lhsRange = resolveSemanticRange(lhs);
       auto rhsRange = resolveSemanticRange(rhs);
-      if (!lhsRange || !rhsRange || !rangesOverlap(*lhsRange, *rhsRange))
+      if (!lhsRange || !rhsRange || !rangesOverlap(*lhsRange, *rhsRange)) {
         continue;
+      }
       op->emitError("PlanMemory semantic no-alias violation: operand byte ranges overlap");
       result = failure();
       return;
@@ -998,7 +1010,9 @@ bool isLocalBuffer(std::optional<AddressSpaceAttr> memorySpaceAttr) {
   if (memorySpaceAttr.value().getAddressSpace() == pto::AddressSpace::GM) {
     return false;
   }
-  if (LocalBufferSpace.count(memorySpaceAttr.value().getAddressSpace())) {
+  const bool isLocalBufferSpace =
+      LocalBufferSpace.count(memorySpaceAttr.value().getAddressSpace()) != 0;
+  if (isLocalBufferSpace) {
     return true;
   }
   llvm_unreachable("Currently only support (UB | L1 | L0C) allocation");
diff --git a/lib/PTO/Transforms/Utils.h b/lib/PTO/Transforms/Utils.h
index ad5c5eb796..4555a7c125 100644
--- a/lib/PTO/Transforms/Utils.h
+++ b/lib/PTO/Transforms/Utils.h
@@ -23,13 +23,9 @@
 #include "mlir/Support/LLVM.h"
 
 #include "llvm/ADT/SmallVector.h"
-#include "llvm/Support/Debug.h"
-
-#include <cassert>
 #include <optional>
 #include <queue>
 #include <set>
-#include <type_traits>
 
 namespace mlir {
 namespace pto {
@@ -66,6 +62,7 @@ namespace pto {
   uint64_t AlignUp(uint64_t lhs, uint64_t rhs);
   LoopLikeOpInterface getParentLoop(Value val);
   ModuleOp getTopLevelModuleOp(Operation *op);
+  Operation *getAncestorInBlock(Operation *op, const Block *block);
   void setBaseMemRefTypeScope(Value val, AddressSpaceAttr targetMemScope);
   BaseMemRefType getBaseMemRefTypeWithNewScope(BaseMemRefType type,
                                              AddressSpaceAttr targetMemScope);
diff --git a/lib/PTO/Transforms/VMIControlFlowSupport.cpp b/lib/PTO/Transforms/VMIControlFlowSupport.cpp
index 4bde54f584..ce0642da05 100644
--- a/lib/PTO/Transforms/VMIControlFlowSupport.cpp
+++ b/lib/PTO/Transforms/VMIControlFlowSupport.cpp
@@ -84,3 +84,40 @@ LogicalResult VMIControlFlowSupport::addWhileConstraints(
   }
   return success();
 }
+
+SmallVector<Type> mlir::pto::getConsistentCallResultTypes(ModuleOp module,
+                                                          func::FuncOp func) {
+  SmallVector<Type> resultTypes;
+  bool found = false;
+  module.walk([func, &resultTypes, &found](func::CallOp call) mutable {
+    bool callsFunction = call.getCallee() == func.getSymName();
+    if (!callsFunction) {
+      return;
+    }
+    if (!found) {
+      resultTypes.assign(call.getResultTypes().begin(),
+                         call.getResultTypes().end());
+      found = true;
+      return;
+    }
+    bool hasMatchingArity = resultTypes.size() == call.getNumResults();
+    if (!hasMatchingArity) {
+      return;
+    }
+    for (auto [index, type] : llvm::enumerate(call.getResultTypes())) {
+      if (resultTypes[index] != type) {
+        resultTypes[index] = {};
+      }
+    }
+  });
+  return found ? resultTypes : SmallVector<Type>{};
+}
+
+SmallVector<Type> mlir::pto::getFunctionInputTypes(func::FuncOp func) {
+  SmallVector<Type> inputs;
+  inputs.reserve(func.getNumArguments());
+  for (BlockArgument argument : func.getArguments()) {
+    inputs.push_back(argument.getType());
+  }
+  return inputs;
+}
diff --git a/lib/PTO/Transforms/VMILayoutAssignment.cpp b/lib/PTO/Transforms/VMILayoutAssignment.cpp
index 186e0b7eff..92f10d1d2e 100644
--- a/lib/PTO/Transforms/VMILayoutAssignment.cpp
+++ b/lib/PTO/Transforms/VMILayoutAssignment.cpp
@@ -179,7 +179,7 @@ struct LayoutSolver {
     return maskNodes[id].parent;
   }
 
-  LogicalResult unite(Value lhs, Value rhs, Operation *op) {
+  LogicalResult unite(Value lhs, Value rhs, const Operation *op) {
     (void)op;
     addDataValue(lhs);
     addDataValue(rhs);
@@ -288,11 +288,11 @@ struct LayoutSolver {
     return success();
   }
 
-  VMILayoutAttr getContiguousLayout() {
+  VMILayoutAttr getContiguousLayout() const {
     return VMILayoutAttr::getContiguous(ctx);
   }
 
-  DataLayoutSeedPhase getCastSeedPhase(const VMICastLayoutFact &fact) {
+  DataLayoutSeedPhase getCastSeedPhase(const VMICastLayoutFact &fact) const {
     if (fact.priority == VMICastLayoutPriority::High) {
       return DataLayoutSeedPhase::CompactCast;
     }
@@ -302,7 +302,7 @@ struct LayoutSolver {
     return DataLayoutSeedPhase::Cast;
   }
 
-  VMILayoutAttr getPreferredDenseStoreLayout(VMIVRegType type) {
+  VMILayoutAttr getPreferredDenseStoreLayout(VMIVRegType type) const {
     VMILayoutSupport supports;
     FailureOr<VMIStoreLayoutFact> fact =
         supports.getPreferredStoreLayoutFact(type);
@@ -323,12 +323,12 @@ struct LayoutSolver {
 
   FailureOr<VMIMaskedStoreLayoutFact>
   getPreferredDenseMaskedStoreLayout(VMIVRegType valueType,
-                                     VMIMaskType maskType) {
+                                     VMIMaskType maskType) const {
     VMILayoutSupport supports;
     return supports.getPreferredMaskedStoreLayoutFact(valueType, maskType);
   }
 
-  VMILayoutAttr getGroupSlotsLayout(int64_t numGroups) {
+  VMILayoutAttr getGroupSlotsLayout(int64_t numGroups) const {
     return VMILayoutAttr::getGroupSlots(ctx, numGroups);
   }
 
@@ -362,9 +362,9 @@ struct LayoutSolver {
     return getContiguousLayout();
   }
 
-  DataLayoutSeedPhase getGroupReduceUseSeedPhase(VMIVRegType sourceType,
-                                                 int64_t numGroups,
-                                                 VMIGroupReduceLayoutFact fact) {
+  DataLayoutSeedPhase
+  getGroupReduceUseSeedPhase(VMIVRegType sourceType, int64_t numGroups,
+                             VMIGroupReduceLayoutFact fact) const {
     if (!fact.sourceLayout || !fact.sourceLayout.isContiguous() ||
         fact.sourceLayout.getLaneStride() != 1) {
       return DataLayoutSeedPhase::Reduce;
@@ -381,7 +381,7 @@ struct LayoutSolver {
     return DataLayoutSeedPhase::Reduce;
   }
 
-  VMILayoutAttr getPreferredGroupSlotLoadLayout(VMIGroupSlotLoadOp op) {
+  VMILayoutAttr getPreferredGroupSlotLoadLayout(VMIGroupSlotLoadOp op) const {
     auto type = cast<VMIVRegType>(op.getResult().getType());
     int64_t numGroups = op.getNumGroupsAttr().getInt();
     if (VMILayoutAttr existing = type.getLayoutAttr()) {
@@ -398,7 +398,7 @@ struct LayoutSolver {
   }
 
   VMILayoutAttr
-  getPreferredGroupBroadcastLoadLayout(VMIGroupBroadcastLoadOp op) {
+  getPreferredGroupBroadcastLoadLayout(VMIGroupBroadcastLoadOp op) const {
     auto type = cast<VMIVRegType>(op.getResult().getType());
     if (VMILayoutAttr existing = type.getLayoutAttr()) {
       return existing;
@@ -445,7 +445,7 @@ struct LayoutSolver {
   }
 
   VMILayoutAttr
-  getPreferredGroupBroadcastResultLayout(VMIGroupBroadcastOp op) {
+  getPreferredGroupBroadcastResultLayout(VMIGroupBroadcastOp op) const {
     auto type = cast<VMIVRegType>(op.getResult().getType());
     if (VMILayoutAttr existing = type.getLayoutAttr()) {
       return existing;
@@ -503,7 +503,7 @@ struct LayoutSolver {
     return getContiguousLayout();
   }
 
-  LogicalResult validateGroupLoadLayoutPlan(VMIGroupLoadOp op) {
+  LogicalResult validateGroupLoadLayoutPlan(VMIGroupLoadOp op) const {
     auto type = cast<VMIVRegType>(op.getResult().getType());
     if (type.getLayoutAttr()) {
       return success();
@@ -588,7 +588,7 @@ struct LayoutSolver {
   }
 
   LogicalResult collect() {
-    module.walk([&](Operation *op) {
+    module.walk([this](Operation *op) {
       for (Value result : op->getResults()) {
         addDataValue(result);
         addMaskValue(result);
@@ -1326,21 +1326,25 @@ struct LayoutSolver {
   }
 
   LogicalResult addExecuteRegionConstraints(scf::ExecuteRegionOp executeOp) {
-    WalkResult result = executeOp.getRegion().walk([&](scf::YieldOp yieldOp) {
-      if (yieldOp->getParentOp() != executeOp.getOperation()) {
-        return WalkResult::advance();
-      }
-      if (failed(
-              addYieldConstraints(executeOp->getResults(), yieldOp, executeOp))) {
-        return WalkResult::interrupt();
-      }
-      return WalkResult::advance();
-    });
+    WalkResult result = executeOp.getRegion().walk(
+        [this, executeOp](scf::YieldOp yieldOp) mutable {
+          const bool belongsToExecuteRegion =
+              yieldOp->getParentOp() == executeOp.getOperation();
+          if (!belongsToExecuteRegion) {
+            return WalkResult::advance();
+          }
+          if (failed(addYieldConstraints(executeOp->getResults(), yieldOp,
+                                         executeOp))) {
+            return WalkResult::interrupt();
+          }
+          return WalkResult::advance();
+        });
     return failure(result.wasInterrupted());
   }
 
   LogicalResult addIndexSwitchConstraints(scf::IndexSwitchOp indexSwitchOp) {
-    auto addBlockTerminator = [&](Block &block) -> LogicalResult {
+    auto addBlockTerminator =
+        [this, indexSwitchOp](Block &block) mutable -> LogicalResult {
       auto yieldOp = dyn_cast<scf::YieldOp>(block.getTerminator());
       if (!yieldOp) {
         return success();
@@ -1361,14 +1365,14 @@ struct LayoutSolver {
 
   LogicalResult addWhileConstraints(scf::WhileOp whileOp) {
     return VMIControlFlowSupport::addWhileConstraints(
-        whileOp, [&](Value lhs, Value rhs, Operation *op) {
+        whileOp, [this](Value lhs, Value rhs, Operation *op) {
           return uniteEquivalentValues(lhs, rhs, op);
         });
   }
 
   LogicalResult addForConstraints(scf::ForOp forOp) {
     return VMIControlFlowSupport::addForConstraints(
-        forOp, [&](Value lhs, Value rhs, Operation *op) {
+        forOp, [this](Value lhs, Value rhs, Operation *op) {
           return uniteEquivalentValues(lhs, rhs, op);
         });
   }
@@ -1415,12 +1419,12 @@ struct LayoutSolver {
     return success();
   }
 
-  bool hasVMIValueTypes(Operation *op) {
+  bool hasVMIValueTypes(Operation *op) const {
     return llvm::any_of(op->getOperandTypes(), containsVMIType) ||
            llvm::any_of(op->getResultTypes(), containsVMIType);
   }
 
-  bool hasVMIFunctionType(func::FuncOp func) {
+  bool hasVMIFunctionType(func::FuncOp func) const {
     FunctionType type = func.getFunctionType();
     return llvm::any_of(type.getInputs(), containsVMIType) ||
            llvm::any_of(type.getResults(), containsVMIType);
@@ -1447,7 +1451,8 @@ struct LayoutSolver {
     }
 
     SmallVector<func::ReturnOp> returns;
-    callee.walk([&](func::ReturnOp returnOp) { returns.push_back(returnOp); });
+    callee.walk(
+        [&returns](func::ReturnOp returnOp) { returns.push_back(returnOp); });
     for (func::ReturnOp returnOp : returns) {
       for (auto [index, result] : llvm::enumerate(callOp.getResults())) {
         if (index >= returnOp.getNumOperands()) {
@@ -1471,7 +1476,8 @@ struct LayoutSolver {
   }
 
   FailureOr<Value> materializeLayoutValue(Value value, Type targetType,
-                                          Location loc, OpBuilder &builder) {
+                                          Location loc,
+                                          OpBuilder &builder) const {
     if (value.getType() == targetType) {
       return value;
     }
@@ -1502,31 +1508,6 @@ struct LayoutSolver {
     return failure();
   }
 
-  SmallVector<Type> getCallResultTypes(func::FuncOp func) {
-    SmallVector<Type> resultTypes;
-    bool found = false;
-    module.walk([&](func::CallOp call) {
-      if (call.getCallee() != func.getSymName()) {
-        return;
-      }
-      if (!found) {
-        resultTypes.assign(call.getResultTypes().begin(),
-                           call.getResultTypes().end());
-        found = true;
-        return;
-      }
-      if (resultTypes.size() != call.getNumResults()) {
-        return;
-      }
-      for (auto [index, type] : llvm::enumerate(call.getResultTypes())) {
-        if (index < resultTypes.size() && resultTypes[index] != type) {
-          resultTypes[index] = {};
-        }
-      }
-    });
-    return found ? resultTypes : SmallVector<Type>{};
-  }
-
   LogicalResult materializeCallOperands(IRRewriter &rewriter) {
     WalkResult result = module.walk([this, &rewriter](func::CallOp call) {
       auto callee = SymbolTable::lookupNearestSymbolFrom<func::FuncOp>(
@@ -1581,7 +1562,8 @@ struct LayoutSolver {
 
   LogicalResult materializeFunctionReturns(IRRewriter &rewriter) {
     WalkResult result = module.walk([this, &rewriter](func::FuncOp func) {
-      SmallVector<Type> resultTypes = getCallResultTypes(func);
+      SmallVector<Type> resultTypes =
+          getConsistentCallResultTypes(module, func);
       if (resultTypes.empty()) {
         return WalkResult::advance();
       }
@@ -1633,11 +1615,11 @@ struct LayoutSolver {
     return success();
   }
 
-  bool hasRequestedLayout(VMILayoutPropagator &propagator, Value value) {
+  bool hasRequestedLayout(VMILayoutPropagator &propagator, Value value) const {
     return static_cast<bool>(propagator.getRequestedLayout(value));
   }
 
-  bool hasLayoutAssignment(VMILayoutPropagator &propagator, Value value) {
+  bool hasLayoutAssignment(VMILayoutPropagator &propagator, Value value) const {
     return propagator.lookup(value) != nullptr;
   }
 
@@ -1804,20 +1786,16 @@ struct LayoutSolver {
   }
 
   void rewriteFunctionType() {
-    module.walk([&](func::FuncOp func) {
+    module.walk([this](func::FuncOp func) {
       if (func.empty()) {
         return;
       }
 
-      SmallVector<Type> inputs;
-      inputs.reserve(func.getNumArguments());
-      for (BlockArgument arg : func.getArguments()) {
-        inputs.push_back(arg.getType());
-      }
-
+      SmallVector<Type> inputs = getFunctionInputTypes(func);
       SmallVector<Type> results;
       auto it = firstReturnOperandsByFunc.find(func);
-      SmallVector<Type> callResultTypes = getCallResultTypes(func);
+      SmallVector<Type> callResultTypes =
+          getConsistentCallResultTypes(module, func);
       if (!callResultTypes.empty()) {
         for (Type type : callResultTypes) {
           results.push_back(type);
diff --git a/lib/PTO/Transforms/VMILayoutPropagation.cpp b/lib/PTO/Transforms/VMILayoutPropagation.cpp
index 0558eb5aa5..b3139badd5 100644
--- a/lib/PTO/Transforms/VMILayoutPropagation.cpp
+++ b/lib/PTO/Transforms/VMILayoutPropagation.cpp
@@ -72,7 +72,7 @@ static bool hasAmbiguousTransferTargets(ArrayRef<VMILayoutFact> facts) {
 }
 
 static bool relationContainsOperandLayout(const VMILayoutRelation &relation,
-                                          OpOperand &operand,
+                                          const OpOperand &operand,
                                           VMILayoutAttr layout) {
   for (const VMILayoutFact &fact : relation.facts) {
     if (fact.operand == &operand) {
@@ -585,22 +585,22 @@ class VMIGroupReduceTransfer final : public VMILayoutTransfer {
     [[maybe_unused]] auto [op, changedValue, changedLayout, propagator,
                            changedOperand] = request;
     if (auto reduce = dyn_cast<VMIGroupReduceAddFOp>(op)) {
-      return queryReduce(reduce, changedValue, changedLayout, changedOperand);
+      return queryReduce(reduce, changedValue, changedLayout);
     }
     if (auto reduce = dyn_cast<VMIGroupReduceMaxFOp>(op)) {
-      return queryReduce(reduce, changedValue, changedLayout, changedOperand);
+      return queryReduce(reduce, changedValue, changedLayout);
     }
     if (auto reduce = dyn_cast<VMIGroupReduceMinFOp>(op)) {
-      return queryReduce(reduce, changedValue, changedLayout, changedOperand);
+      return queryReduce(reduce, changedValue, changedLayout);
     }
     if (auto reduce = dyn_cast<VMIGroupReduceAddIOp>(op)) {
-      return queryReduce(reduce, changedValue, changedLayout, changedOperand);
+      return queryReduce(reduce, changedValue, changedLayout);
     }
     if (auto reduce = dyn_cast<VMIGroupReduceMaxIOp>(op)) {
-      return queryReduce(reduce, changedValue, changedLayout, changedOperand);
+      return queryReduce(reduce, changedValue, changedLayout);
     }
     if (auto reduce = dyn_cast<VMIGroupReduceMinIOp>(op)) {
-      return queryReduce(reduce, changedValue, changedLayout, changedOperand);
+      return queryReduce(reduce, changedValue, changedLayout);
     }
     return failure();
   }
@@ -608,8 +608,8 @@ class VMIGroupReduceTransfer final : public VMILayoutTransfer {
 private:
   template <typename OpTy>
   FailureOr<SmallVector<VMILayoutRelation, mlir::pto::kValue4>>
-  queryReduce(OpTy reduce, Value changedValue, VMILayoutAttr changedLayout,
-              OpOperand *changedOperand) const {
+  queryReduce(OpTy reduce, Value changedValue,
+              VMILayoutAttr changedLayout) const {
     auto sourceType = dyn_cast<VMIVRegType>(reduce.getSource().getType());
     auto resultType = dyn_cast<VMIVRegType>(reduce.getResult().getType());
     if (!sourceType || !resultType) {
@@ -1189,7 +1189,7 @@ void VMILayoutPropagator::addEquivalentValues(Value lhs, Value rhs) {
     return;
   }
 
-  auto addEdge = [&](Value from, Value to) {
+  auto addEdge = [this](Value from, Value to) {
     SmallVector<Value, 2> &values = equivalentValues[from];
     if (!llvm::is_contained(values, to)) {
       values.push_back(to);
@@ -1214,7 +1214,7 @@ void VMILayoutPropagator::enqueue(Value value, VMILayoutAttr layout) {
 LogicalResult
 VMILayoutPropagator::addUseConflict(OpOperand &operand,
                                     VMIValueLayoutAssignment &assignment,
-                                    VMILayoutAttr layout) {
+                                    VMILayoutAttr layout) const {
   for (VMILayoutConflict &conflict : assignment.conflicts) {
     if (conflict.operand != &operand) {
       continue;
@@ -1467,7 +1467,7 @@ bool VMILayoutPropagator::isTypeRewriteable(Value value) const {
 FailureOr<Value> VMILayoutPropagator::materializeAt(Value source,
                                                     VMILayoutAttr layout,
                                                     RewriterBase &rewriter,
-                                                    Location loc) {
+                                                    Location loc) const {
   VMILayoutAttr sourceLayout = getCurrentLayout(source);
   if (!sourceLayout) {
     return failure();
diff --git a/lib/PTO/Transforms/VMILayoutRematerialize.cpp b/lib/PTO/Transforms/VMILayoutRematerialize.cpp
index 4a6bb091b4..5fdc47aa65 100644
--- a/lib/PTO/Transforms/VMILayoutRematerialize.cpp
+++ b/lib/PTO/Transforms/VMILayoutRematerialize.cpp
@@ -98,7 +98,8 @@ static std::optional<Value> rematerializeBinaryDataOp(Operation *op,
                                                       VMIVRegType resultType,
                                                       Location loc,
                                                       OpBuilder &builder) {
-  auto rebuild = [&](auto typedOp) -> std::optional<Value> {
+  auto rebuild = [resultType, loc,
+                  &builder](auto typedOp) -> std::optional<Value> {
     auto lhsType = dyn_cast<VMIVRegType>(typedOp.getLhs().getType());
     auto rhsType = dyn_cast<VMIVRegType>(typedOp.getRhs().getType());
     if (!lhsType || !rhsType) {
@@ -130,7 +131,8 @@ static std::optional<Value> rematerializeUnaryDataOp(Operation *op,
                                                      VMIVRegType resultType,
                                                      Location loc,
                                                      OpBuilder &builder) {
-  auto rebuild = [&](auto typedOp) -> std::optional<Value> {
+  auto rebuild = [resultType, loc,
+                  &builder](auto typedOp) -> std::optional<Value> {
     auto sourceType = dyn_cast<VMIVRegType>(typedOp.getSource().getType());
     if (!sourceType) {
       return std::nullopt;
@@ -381,7 +383,7 @@ struct VMILayoutRematerializePass
     while (changed) {
       changed = false;
       SmallVector<Operation *> helpers;
-      module.walk([&](Operation *op) {
+      module.walk([&helpers](Operation *op) {
         if (isa<VMIEnsureLayoutOp, VMIEnsureMaskLayoutOp,
                 VMIEnsureMaskGranularityOp, VMITruncIOp>(op)) {
           helpers.push_back(op);
@@ -394,22 +396,22 @@ struct VMILayoutRematerializePass
         }
 
         if (auto ensure = dyn_cast<VMIEnsureLayoutOp>(op)) {
-          changed |= tryReplaceDataEnsure(ensure);
+          changed = tryReplaceDataEnsure(ensure) || changed;
           continue;
         }
 
         if (auto ensure = dyn_cast<VMIEnsureMaskLayoutOp>(op)) {
-          changed |= tryReplaceMaskEnsure(ensure);
+          changed = tryReplaceMaskEnsure(ensure) || changed;
           continue;
         }
 
         if (auto ensure = dyn_cast<VMIEnsureMaskGranularityOp>(op)) {
-          changed |= tryReplaceMaskEnsure(ensure);
+          changed = tryReplaceMaskEnsure(ensure) || changed;
           continue;
         }
 
         if (auto trunc = dyn_cast<VMITruncIOp>(op)) {
-          changed |= tryRematerializeTruncIThroughSourceEnsure(trunc);
+          changed = tryRematerializeTruncIThroughSourceEnsure(trunc) || changed;
         }
       }
     }
diff --git a/lib/PTO/Transforms/VMILayoutSinkMaterialization.cpp b/lib/PTO/Transforms/VMILayoutSinkMaterialization.cpp
index e1c3323374..91fcf65358 100644
--- a/lib/PTO/Transforms/VMILayoutSinkMaterialization.cpp
+++ b/lib/PTO/Transforms/VMILayoutSinkMaterialization.cpp
@@ -685,7 +685,7 @@ struct VMILayoutSinkMaterializationPass
   void runOnOperation() override {
     ModuleOp module = getOperation();
     SmallVector<Operation *> candidates;
-    module.walk([&](Operation *op) {
+    module.walk([&candidates](Operation *op) {
       if (isSinkCandidate(op)) {
         candidates.push_back(op);
       }
diff --git a/lib/PTO/Transforms/VMILayoutSupport.cpp b/lib/PTO/Transforms/VMILayoutSupport.cpp
index b227cca7cd..8148c088b3 100644
--- a/lib/PTO/Transforms/VMILayoutSupport.cpp
+++ b/lib/PTO/Transforms/VMILayoutSupport.cpp
@@ -49,6 +49,33 @@ namespace {
 
 constexpr int64_t kLayoutBlockBitWidth = 256;
 
+template <typename ResultT>
+static ResultT failWithReason(std::string *reason, const Twine &message) {
+  if (reason) {
+    *reason = message.str();
+  }
+  return failure();
+}
+
+static FailureOr<VMIGroupBroadcastLoadDirectFact>
+failGroupBroadcastLoadDirect(std::string *reason, const Twine &message) {
+  return failWithReason<FailureOr<VMIGroupBroadcastLoadDirectFact>>(reason,
+                                                                    message);
+}
+
+static std::optional<LogicalResult>
+getEnsureLayoutEarlyExit(VMILayoutAttr sourceLayout,
+                         VMILayoutAttr resultLayout, std::string *reason) {
+  if (!sourceLayout || !resultLayout) {
+    return failWithReason<LogicalResult>(
+        reason, "requires assigned source/result layouts");
+  }
+  if (sourceLayout == resultLayout) {
+    return success();
+  }
+  return std::nullopt;
+}
+
 static llvm::cl::opt<bool> preferLaneStrideNarrowing(
     "vmi-prefer-lane-stride-narrowing",
     llvm::cl::desc(
@@ -63,6 +90,18 @@ static llvm::cl::opt<bool> preferLaneStrideNarrowing(
 
 #include "VMILayoutSupportQueryHelpers.inc"
 
+static VMIGroupBroadcastLoadDirectFact materializeGroupBroadcastLoadDirectFact(
+    const GroupBroadcastLoadDirectPattern &pattern,
+    const GroupBroadcastLoadQuery &query, VMILayoutAttr resultLayout,
+    unsigned elementBits) {
+  return VMIGroupBroadcastLoadDirectFact{
+      pattern.kind,
+      VMIGroupBroadcastLoadLayoutFact{
+          getGroupBlockClassFromPattern(pattern.block), resultLayout,
+          query.key.groupSize, query.key.lanesPerPart, query.key.vcgBlockElems,
+          static_cast<int64_t>(elementBits)}};
+}
+
 //===----------------------------------------------------------------------===//
 // Query implementations
 //===----------------------------------------------------------------------===//
@@ -70,7 +109,7 @@ static llvm::cl::opt<bool> preferLaneStrideNarrowing(
 FailureOr<VMIVselrLayoutFact>
 VMILayoutSupport::getPreferredVselrLayoutFact(
     VMIVselrOp op, std::string *reason) const {
-  auto fail = [&](const Twine &message) -> FailureOr<VMIVselrLayoutFact> {
+  auto fail = [reason](const Twine &message) -> FailureOr<VMIVselrLayoutFact> {
     if (reason) {
       *reason = message.str();
     }
@@ -92,7 +131,7 @@ VMILayoutSupport::getPreferredVselrLayoutFact(
 FailureOr<VMIVselrLayoutFact>
 VMILayoutSupport::getVselrLayoutFact(VMIVselrOp op,
                                      std::string *reason) const {
-  auto fail = [&](const Twine &message) -> FailureOr<VMIVselrLayoutFact> {
+  auto fail = [reason](const Twine &message) -> FailureOr<VMIVselrLayoutFact> {
     if (reason) {
       *reason = message.str();
     }
@@ -136,7 +175,8 @@ FailureOr<VMIGroupReduceLayoutFact>
 VMILayoutSupport::getPreferredGroupReduceLayoutFact(VMIVRegType sourceType,
                                                     int64_t numGroups,
                                                     std::string *reason) const {
-  auto fail = [&](const Twine &message) -> FailureOr<VMIGroupReduceLayoutFact> {
+  auto fail =
+      [reason](const Twine &message) -> FailureOr<VMIGroupReduceLayoutFact> {
     if (reason) {
       *reason = message.str();
     }
@@ -168,7 +208,8 @@ FailureOr<VMIGroupReduceLayoutFact>
 VMILayoutSupport::getGroupReduceLayoutFactForLayouts(
     VMIVRegType sourceType, VMIMaskType maskType, VMIVRegType resultType,
     int64_t numGroups, std::string *reason) const {
-  auto fail = [&](const Twine &message) -> FailureOr<VMIGroupReduceLayoutFact> {
+  auto fail =
+      [reason](const Twine &message) -> FailureOr<VMIGroupReduceLayoutFact> {
     if (reason) {
       *reason = message.str();
     }
@@ -210,7 +251,7 @@ FailureOr<SmallVector<VMIGroupReduceLayoutFact, mlir::pto::kValue4>>
 VMILayoutSupport::getGroupReduceLayoutFactsForLayout(
     VMIVRegType sourceType, int64_t numGroups, VMIGroupReduceLayoutPort port,
     VMILayoutAttr layout, std::string *reason) const {
-  auto fail = [&](const Twine &message)
+  auto fail = [reason](const Twine &message)
       -> FailureOr<SmallVector<VMIGroupReduceLayoutFact, 4>> {
     if (reason) {
       *reason = message.str();
@@ -266,7 +307,7 @@ VMILayoutSupport::getGroupBroadcastLayoutFactForLayouts(
     VMIVRegType sourceType, VMIVRegType resultType, int64_t numGroups,
     std::string *reason) const {
   auto fail =
-      [&](const Twine &message) -> FailureOr<VMIGroupBroadcastLayoutFact> {
+      [reason](const Twine &message) -> FailureOr<VMIGroupBroadcastLayoutFact> {
     if (reason) {
       *reason = message.str();
     }
@@ -308,7 +349,7 @@ VMILayoutSupport::getGroupBroadcastLayoutFactsForLayout(
     VMIVRegType sourceType, VMIVRegType resultType, int64_t numGroups,
     VMIGroupBroadcastLayoutPort port, VMILayoutAttr layout,
     std::string *reason) const {
-  auto fail = [&](const Twine &message)
+  auto fail = [reason](const Twine &message)
       -> FailureOr<SmallVector<VMIGroupBroadcastLayoutFact, 4>> {
     if (reason) {
       *reason = message.str();
@@ -372,7 +413,8 @@ VMILayoutSupport::getGroupBroadcastLoadLayoutFact(VMIVRegType resultType,
                                                   int64_t numGroups,
                                                   std::string *reason) const {
   auto fail =
-      [&](const Twine &message) -> FailureOr<VMIGroupBroadcastLoadLayoutFact> {
+      [reason](
+          const Twine &message) -> FailureOr<VMIGroupBroadcastLoadLayoutFact> {
     if (reason) {
       *reason = message.str();
     }
@@ -430,22 +472,17 @@ FailureOr<VMIGroupBroadcastLoadDirectFact>
 VMILayoutSupport::getGroupBroadcastLoadDirectFact(
     VMIVRegType resultType, Type sourceType, Value sourceGroupStride,
     int64_t numGroups, std::string *reason) const {
-  auto fail =
-      [&](const Twine &message) -> FailureOr<VMIGroupBroadcastLoadDirectFact> {
-    if (reason) {
-      *reason = message.str();
-    }
-    return failure();
-  };
-
   if (!isa<PtrType>(sourceType)) {
-    return fail("group_broadcast_load direct lowering requires !pto.ptr source");
+    return failGroupBroadcastLoadDirect(
+        reason,
+        "group_broadcast_load direct lowering requires !pto.ptr source");
   }
 
   unsigned elementBits =
       pto::getPTOStorageElemBitWidth(resultType.getElementType());
   if (elementBits == 0) {
-    return fail("group_broadcast_load requires known element bit width");
+    return failGroupBroadcastLoadDirect(
+        reason, "group_broadcast_load requires known element bit width");
   }
   std::optional<int64_t> stride = getConstantIndexValue(sourceGroupStride);
 
@@ -471,19 +508,13 @@ VMILayoutSupport::getGroupBroadcastLoadDirectFact(
     if (existing && existing != resultLayout) {
       continue;
     }
-    return VMIGroupBroadcastLoadDirectFact{
-        pattern.kind,
-        VMIGroupBroadcastLoadLayoutFact{
-            getGroupBlockClassFromPattern(pattern.block),
-            resultLayout,
-            query.key.groupSize,
-            query.key.lanesPerPart,
-            query.key.vcgBlockElems,
-            static_cast<int64_t>(elementBits)}};
+    return materializeGroupBroadcastLoadDirectFact(pattern, query, resultLayout,
+                                                   elementBits);
   }
 
-  return fail("group_broadcast_load has no preferred direct lowering layout "
-              "table row");
+  return failGroupBroadcastLoadDirect(
+      reason,
+      "group_broadcast_load has no preferred direct lowering layout table row");
 }
 
 static std::pair<int64_t, int64_t> getCastElementBits(VMIVRegType sourceType,
@@ -722,7 +753,7 @@ VMILayoutSupport::getCastLayoutFactsForLayout(VMIVRegType sourceType,
                                               VMICastLayoutPort port,
                                               VMILayoutAttr layout,
                                               std::string *reason) const {
-  auto fail = [&](const Twine &message)
+  auto fail = [reason](const Twine &message)
       -> FailureOr<SmallVector<VMICastLayoutFact, 4>> {
     if (reason) {
       *reason = message.str();
@@ -775,7 +806,7 @@ VMILayoutSupport::getCastLayoutFactsForLayout(VMIVRegType sourceType,
 static FailureOr<VMICastLayoutFact>
 getUniqueCastLayoutFact(FailureOr<SmallVector<VMICastLayoutFact, mlir::pto::kValue4>> facts,
                         std::string *reason) {
-  auto fail = [&](const Twine &message) -> FailureOr<VMICastLayoutFact> {
+  auto fail = [reason](const Twine &message) -> FailureOr<VMICastLayoutFact> {
     if (reason) {
       *reason = message.str();
     }
@@ -816,7 +847,7 @@ FailureOr<VMICastLayoutFact> VMILayoutSupport::getCastLayoutFactForResultLayout(
 FailureOr<VMICastLayoutFact> VMILayoutSupport::getCastLayoutFactForLayouts(
     VMIVRegType sourceType, VMIVRegType resultType, VMILayoutAttr sourceLayout,
     VMILayoutAttr resultLayout, std::string *reason) const {
-  auto fail = [&](const Twine &message) -> FailureOr<VMICastLayoutFact> {
+  auto fail = [reason](const Twine &message) -> FailureOr<VMICastLayoutFact> {
     if (reason) {
       *reason = message.str();
     }
@@ -856,7 +887,8 @@ struct MaskGranularityCastQuery {
 static FailureOr<MaskGranularityCastQuery> buildMaskGranularityCastQuery(
     VMIMaskType sourceType, VMIMaskType resultType, VMILayoutAttr layout,
     std::string *reason) {
-  auto fail = [&](const Twine &message) -> FailureOr<MaskGranularityCastQuery> {
+  auto fail =
+      [reason](const Twine &message) -> FailureOr<MaskGranularityCastQuery> {
     if (reason) {
       *reason = message.str();
     }
@@ -895,7 +927,7 @@ FailureOr<SmallVector<VMIMaskGranularityCastLayoutFact, mlir::pto::kValue4>>
 VMILayoutSupport::getMaskGranularityCastLayoutFactsForLayout(
     VMIMaskType sourceType, VMIMaskType resultType, VMICastLayoutPort port,
     VMILayoutAttr layout, std::string *reason) const {
-  auto fail = [&](const Twine &message)
+  auto fail = [reason](const Twine &message)
       -> FailureOr<SmallVector<VMIMaskGranularityCastLayoutFact, 4>> {
     if (reason) {
       *reason = message.str();
@@ -952,7 +984,8 @@ VMILayoutSupport::getMaskGranularityCastLayoutFactForLayouts(
     VMIMaskType sourceType, VMIMaskType resultType, VMILayoutAttr sourceLayout,
     VMILayoutAttr resultLayout, std::string *reason) const {
   auto fail =
-      [&](const Twine &message) -> FailureOr<VMIMaskGranularityCastLayoutFact> {
+      [reason](
+          const Twine &message) -> FailureOr<VMIMaskGranularityCastLayoutFact> {
     if (reason) {
       *reason = message.str();
     }
@@ -999,7 +1032,8 @@ FailureOr<VMILayoutAttr> VMILayoutSupport::getWidenSourceLayoutForResultLayout(
 static FailureOr<VMIInterleaveLayoutFact> getPreferredInterleaveLayoutFactImpl(
     ArrayRef<InterleaveLayoutPattern> patterns, VMIVRegType valueType,
     std::string *reason) {
-  auto fail = [&](const Twine &message) -> FailureOr<VMIInterleaveLayoutFact> {
+  auto fail =
+      [reason](const Twine &message) -> FailureOr<VMIInterleaveLayoutFact> {
     if (reason) {
       *reason = message.str();
     }
@@ -1028,7 +1062,7 @@ getInterleaveLayoutFactsForLayoutImpl(
     ArrayRef<InterleaveLayoutPattern> patterns, VMIVRegType valueType,
     VMIInterleaveLayoutPort port, VMILayoutAttr layout,
     std::string *reason) {
-  auto fail = [&](const Twine &message)
+  auto fail = [reason](const Twine &message)
       -> FailureOr<SmallVector<VMIInterleaveLayoutFact, 4>> {
     if (reason) {
       *reason = message.str();
@@ -1153,7 +1187,8 @@ static bool matchesInterleaveLayouts(const VMIInterleaveLayoutFact &fact,
 static FailureOr<VMIInterleaveLayoutFact> getInterleaveLayoutFactForLayoutsImpl(
     ArrayRef<InterleaveLayoutPattern> patterns, InterleaveTypes types,
     std::string *reason) {
-  auto fail = [&](const Twine &message) -> FailureOr<VMIInterleaveLayoutFact> {
+  auto fail =
+      [reason](const Twine &message) -> FailureOr<VMIInterleaveLayoutFact> {
     if (reason) {
       *reason = message.str();
     }
@@ -1243,7 +1278,7 @@ VMILayoutSupport::getVdintlvLayoutFactForLayouts(
 FailureOr<VMILoadLayoutFact>
 VMILayoutSupport::getLoadLayoutFact(VMIVRegType resultType,
                                     std::string *reason) const {
-  auto fail = [&](const Twine &message) -> FailureOr<VMILoadLayoutFact> {
+  auto fail = [reason](const Twine &message) -> FailureOr<VMILoadLayoutFact> {
     if (reason) {
       *reason = message.str();
     }
@@ -1276,7 +1311,8 @@ FailureOr<VMIDeinterleaveLoadLayoutFact>
 VMILayoutSupport::getPreferredDeinterleaveLoadLayoutFact(
     VMIVRegType valueType, std::string *reason) const {
   auto fail =
-      [&](const Twine &message) -> FailureOr<VMIDeinterleaveLoadLayoutFact> {
+      [reason](
+          const Twine &message) -> FailureOr<VMIDeinterleaveLoadLayoutFact> {
     if (reason) {
       *reason = message.str();
     }
@@ -1299,7 +1335,7 @@ FailureOr<SmallVector<VMIDeinterleaveLoadLayoutFact, mlir::pto::kValue4>>
 VMILayoutSupport::getDeinterleaveLoadLayoutFactsForLayout(
     VMIVRegType valueType, VMIDeinterleaveLoadLayoutPort port,
     VMILayoutAttr layout, std::string *reason) const {
-  auto fail = [&](const Twine &message)
+  auto fail = [reason](const Twine &message)
       -> FailureOr<SmallVector<VMIDeinterleaveLoadLayoutFact, 4>> {
     if (reason) {
       *reason = message.str();
@@ -1337,7 +1373,8 @@ FailureOr<VMIDeinterleaveLoadLayoutFact>
 VMILayoutSupport::getDeinterleaveLoadLayoutFactForLayouts(
     VMIVRegType lowType, VMIVRegType highType, std::string *reason) const {
   auto fail =
-      [&](const Twine &message) -> FailureOr<VMIDeinterleaveLoadLayoutFact> {
+      [reason](
+          const Twine &message) -> FailureOr<VMIDeinterleaveLoadLayoutFact> {
     if (reason) {
       *reason = message.str();
     }
@@ -1373,7 +1410,7 @@ VMILayoutSupport::getDeinterleaveLoadLayoutFactForLayouts(
 FailureOr<VMIStoreLayoutFact>
 VMILayoutSupport::getStoreLayoutFact(VMIVRegType valueType,
                                      std::string *reason) const {
-  auto fail = [&](const Twine &message) -> FailureOr<VMIStoreLayoutFact> {
+  auto fail = [reason](const Twine &message) -> FailureOr<VMIStoreLayoutFact> {
     if (reason) {
       *reason = message.str();
     }
@@ -1405,7 +1442,7 @@ VMILayoutSupport::getStoreLayoutFact(VMIVRegType valueType,
 FailureOr<VMIStoreLayoutFact>
 VMILayoutSupport::getPreferredStoreLayoutFact(VMIVRegType valueType,
                                               std::string *reason) const {
-  auto fail = [&](const Twine &message) -> FailureOr<VMIStoreLayoutFact> {
+  auto fail = [reason](const Twine &message) -> FailureOr<VMIStoreLayoutFact> {
     if (reason) {
       *reason = message.str();
     }
@@ -1441,7 +1478,8 @@ VMILayoutSupport::getPreferredStoreLayoutFact(VMIVRegType valueType,
 
 FailureOr<VMIMaskedStoreLayoutFact> VMILayoutSupport::getMaskedStoreLayoutFact(
     VMIVRegType valueType, VMIMaskType maskType, std::string *reason) const {
-  auto fail = [&](const Twine &message) -> FailureOr<VMIMaskedStoreLayoutFact> {
+  auto fail =
+      [reason](const Twine &message) -> FailureOr<VMIMaskedStoreLayoutFact> {
     if (reason) {
       *reason = message.str();
     }
@@ -1482,7 +1520,7 @@ FailureOr<VMIMaskedStoreLayoutFact>
 VMILayoutSupport::getPreferredMaskedStoreLayoutFact(
     VMIVRegType valueType, VMIMaskType maskType, std::string *reason) const {
   auto fail =
-      [&](const Twine &message) -> FailureOr<VMIMaskedStoreLayoutFact> {
+      [reason](const Twine &message) -> FailureOr<VMIMaskedStoreLayoutFact> {
     if (reason) {
       *reason = message.str();
     }
@@ -1532,7 +1570,8 @@ VMILayoutSupport::getPreferredMaskedStoreLayoutFact(
 FailureOr<VMIMaskedLoadLayoutFact> VMILayoutSupport::getMaskedLoadLayoutFact(
     VMIVRegType resultType, VMIMaskType maskType, VMIVRegType passthruType,
     std::string *reason) const {
-  auto fail = [&](const Twine &message) -> FailureOr<VMIMaskedLoadLayoutFact> {
+  auto fail =
+      [reason](const Twine &message) -> FailureOr<VMIMaskedLoadLayoutFact> {
     if (reason) {
       *reason = message.str();
     }
@@ -1575,17 +1614,9 @@ static LogicalResult matchEnsureLayoutPattern(VMIVRegType sourceType,
                                               VMILayoutAttr sourceLayout,
                                               VMILayoutAttr resultLayout,
                                               std::string *reason) {
-  auto fail = [&](const Twine &message) -> LogicalResult {
-    if (reason) {
-      *reason = message.str();
-    }
-    return failure();
-  };
-  if (!sourceLayout || !resultLayout) {
-    return fail("requires assigned source/result layouts");
-  }
-  if (sourceLayout == resultLayout) {
-    return success();
+  if (std::optional<LogicalResult> earlyExit =
+          getEnsureLayoutEarlyExit(sourceLayout, resultLayout, reason)) {
+    return *earlyExit;
   }
 
   int64_t numGroups =
@@ -1613,8 +1644,9 @@ static LogicalResult matchEnsureLayoutPattern(VMIVRegType sourceType,
     return success();
   }
 
-  return fail("source/result layouts do not match a supported ensure_layout "
-              "table row");
+  return failWithReason<LogicalResult>(
+      reason,
+      "source/result layouts do not match a supported ensure_layout table row");
 }
 
 static LogicalResult matchEnsureMaskLayoutPattern(VMIMaskType sourceType,
@@ -1622,17 +1654,9 @@ static LogicalResult matchEnsureMaskLayoutPattern(VMIMaskType sourceType,
                                                   VMILayoutAttr sourceLayout,
                                                   VMILayoutAttr resultLayout,
                                                   std::string *reason) {
-  auto fail = [&](const Twine &message) -> LogicalResult {
-    if (reason) {
-      *reason = message.str();
-    }
-    return failure();
-  };
-  if (!sourceLayout || !resultLayout) {
-    return fail("requires assigned source/result layouts");
-  }
-  if (sourceLayout == resultLayout) {
-    return success();
+  if (std::optional<LogicalResult> earlyExit =
+          getEnsureLayoutEarlyExit(sourceLayout, resultLayout, reason)) {
+    return *earlyExit;
   }
 
   for (const EnsureMaskLayoutPattern &pattern : kEnsureMaskLayoutPatterns) {
@@ -1655,8 +1679,10 @@ static LogicalResult matchEnsureMaskLayoutPattern(VMIMaskType sourceType,
     return success();
   }
 
-  return fail("source/result mask layouts do not match a supported "
-              "ensure_mask_layout table row");
+  return failWithReason<LogicalResult>(
+      reason,
+      "source/result mask layouts do not match a supported ensure_mask_layout "
+      "table row");
 }
 
 FailureOr<VMIEnsureLayoutFact> VMILayoutSupport::getEnsureLayoutFact(
@@ -1683,7 +1709,8 @@ FailureOr<VMIEnsureMaskLayoutFact> VMILayoutSupport::getEnsureMaskLayoutFact(
 
 FailureOr<VMIGroupSlotLayoutFact> VMILayoutSupport::getGroupSlotLoadLayoutFact(
     VMIVRegType resultType, int64_t numGroups, std::string *reason) const {
-  auto fail = [&](const Twine &message) -> FailureOr<VMIGroupSlotLayoutFact> {
+  auto fail =
+      [reason](const Twine &message) -> FailureOr<VMIGroupSlotLayoutFact> {
     if (reason) {
       *reason = message.str();
     }
@@ -1714,7 +1741,8 @@ VMILayoutSupport::getGroupLoadLayoutFact(VMIGroupLoadOp op,
 FailureOr<VMIGroupLoadLayoutFact> VMILayoutSupport::getGroupLoadLayoutFact(
     VMIVRegType resultType, Value rowStride, int64_t numGroups,
     std::string *reason) const {
-  auto fail = [&](const Twine &message) -> FailureOr<VMIGroupLoadLayoutFact> {
+  auto fail =
+      [reason](const Twine &message) -> FailureOr<VMIGroupLoadLayoutFact> {
     if (reason) {
       *reason = message.str();
     }
@@ -1765,7 +1793,8 @@ FailureOr<VMIGroupLoadLayoutFact> VMILayoutSupport::getGroupLoadLayoutFact(
 
 FailureOr<VMIGroupSlotLayoutFact> VMILayoutSupport::getGroupStoreLayoutFact(
     VMIVRegType valueType, int64_t numGroups, std::string *reason) const {
-  auto fail = [&](const Twine &message) -> FailureOr<VMIGroupSlotLayoutFact> {
+  auto fail =
+      [reason](const Twine &message) -> FailureOr<VMIGroupSlotLayoutFact> {
     if (reason) {
       *reason = message.str();
     }
@@ -1785,7 +1814,8 @@ FailureOr<VMIGroupSlotLayoutFact> VMILayoutSupport::getGroupStoreLayoutFact(
 
 FailureOr<VMIGroupStoreLayoutFact> VMILayoutSupport::getGroupStoreLayoutFact(
     VMIGroupStoreOp op, VMIVRegType valueType, std::string *reason) const {
-  auto fail = [&](const Twine &message) -> FailureOr<VMIGroupStoreLayoutFact> {
+  auto fail =
+      [reason](const Twine &message) -> FailureOr<VMIGroupStoreLayoutFact> {
     if (reason) {
       *reason = message.str();
     }
@@ -1841,7 +1871,7 @@ FailureOr<SmallVector<VMIGroupStoreLayoutFact, mlir::pto::kValue4>>
 VMILayoutSupport::getGroupStoreLayoutFactsForLayout(
     VMIGroupStoreOp op, VMIVRegType valueType, VMILayoutAttr layout,
     std::string *reason) const {
-  auto fail = [&](const Twine &message)
+  auto fail = [reason](const Twine &message)
       -> FailureOr<SmallVector<VMIGroupStoreLayoutFact, 4>> {
     if (reason) {
       *reason = message.str();
@@ -1885,7 +1915,8 @@ VMILayoutSupport::getGroupStoreLayoutFactsForLayout(
 FailureOr<VMIGroupStoreLayoutFact>
 VMILayoutSupport::getPreferredGroupStoreLayoutFact(
     VMIGroupStoreOp op, VMIVRegType valueType, std::string *reason) const {
-  auto fail = [&](const Twine &message) -> FailureOr<VMIGroupStoreLayoutFact> {
+  auto fail =
+      [reason](const Twine &message) -> FailureOr<VMIGroupStoreLayoutFact> {
     if (reason) {
       *reason = message.str();
     }
@@ -1926,7 +1957,8 @@ VMILayoutSupport::getPreferredGroupStoreLayoutFact(
 FailureOr<VMIGroupStoreLayoutFact>
 VMILayoutSupport::getHighPriorityGroupStoreLayoutFact(
     VMIGroupStoreOp op, VMIVRegType valueType, std::string *reason) const {
-  auto fail = [&](const Twine &message) -> FailureOr<VMIGroupStoreLayoutFact> {
+  auto fail =
+      [reason](const Twine &message) -> FailureOr<VMIGroupStoreLayoutFact> {
     if (reason) {
       *reason = message.str();
     }
@@ -1968,7 +2000,8 @@ VMILayoutSupport::getHighPriorityGroupStoreLayoutFact(
 FailureOr<VMIBitcastLayoutFact>
 VMILayoutSupport::getBitcastLayoutFact(VMIBitcastOp op,
                                        std::string *reason) const {
-  auto fail = [&](const Twine &message) -> FailureOr<VMIBitcastLayoutFact> {
+  auto fail =
+      [reason](const Twine &message) -> FailureOr<VMIBitcastLayoutFact> {
     if (reason) {
       *reason = message.str();
     }
@@ -2022,7 +2055,7 @@ FailureOr<SmallVector<VMIBitcastLayoutFact, mlir::pto::kValue4>>
 VMILayoutSupport::getBitcastLayoutFactsForLayout(
     VMIVRegType sourceType, VMIVRegType resultType, VMICastLayoutPort port,
     VMILayoutAttr layout, std::string *reason) const {
-  auto fail = [&](const Twine &message)
+  auto fail = [reason](const Twine &message)
       -> FailureOr<SmallVector<VMIBitcastLayoutFact, 4>> {
     if (reason) {
       *reason = message.str();
@@ -2086,7 +2119,8 @@ template <typename OpTy>
 static FailureOr<VMIHistogramLayoutFact>
 getHistogramLayoutFactImpl(OpTy op, ArrayRef<HistogramLayoutPattern> patterns,
                            StringRef opName, std::string *reason) {
-  auto fail = [&](const Twine &message) -> FailureOr<VMIHistogramLayoutFact> {
+  auto fail =
+      [reason](const Twine &message) -> FailureOr<VMIHistogramLayoutFact> {
     if (reason) {
       *reason = message.str();
     }
diff --git a/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp b/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp
index 3c31ab3b5d..65a7ad349e 100644
--- a/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp
+++ b/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp
@@ -260,7 +260,7 @@ lowerBinaryIgnoringMask(
 /// Lower a UNARY unified op (vneg, vabs, …) to its legacy counterpart.
 template <typename UnifiedOp>
 static LogicalResult
-lowerMaskedUnary(UnifiedOp op, OpBuilder &builder,
+lowerMaskedUnary(UnifiedOp op,
                  function_ref<Value(Location, Type, Value)> createLegacy) {
   if (hasMergePmode(op)) {
     return failure();
@@ -1450,7 +1450,7 @@ static void lowerTypedUnary(UnifiedOp op, OpBuilder &builder) {
     }
     return builder.create<IntegerOp>(loc, type, source).getResult();
   };
-  (void)lowerMaskedUnary(op, builder, createLegacy);
+  (void)lowerMaskedUnary(op, createLegacy);
 }
 
 template <typename LegacyOp, typename UnifiedOp>
@@ -1459,7 +1459,7 @@ static void lowerSimpleUnary(UnifiedOp op, OpBuilder &builder) {
                                   Value source) -> Value {
     return builder.create<LegacyOp>(loc, type, source).getResult();
   };
-  (void)lowerMaskedUnary(op, builder, createLegacy);
+  (void)lowerMaskedUnary(op, createLegacy);
 }
 
 static Value createAbsoluteValue(VMIVabsOp op, OpBuilder &builder,
@@ -1491,7 +1491,7 @@ static void lowerAbsolute(VMIVabsOp op, OpBuilder &builder) {
                                       Value source) -> Value {
     return createAbsoluteValue(op, builder, loc, type, source);
   };
-  (void)lowerMaskedUnary(op, builder, createLegacy);
+  (void)lowerMaskedUnary(op, createLegacy);
 }
 
 static void lowerLogicalNot(VMIVnotOp op, OpBuilder &builder) {
diff --git a/lib/PTO/Transforms/VMIMaskGranularityAssignment.cpp b/lib/PTO/Transforms/VMIMaskGranularityAssignment.cpp
index 6c66aa87e0..e0998c047b 100644
--- a/lib/PTO/Transforms/VMIMaskGranularityAssignment.cpp
+++ b/lib/PTO/Transforms/VMIMaskGranularityAssignment.cpp
@@ -188,7 +188,7 @@ struct MaskGranularitySolver {
   }
 
   LogicalResult collect() {
-    module.walk([&](Operation *op) {
+    module.walk([this](Operation *op) {
       for (Value result : op->getResults()) {
         addMaskValue(result);
       }
@@ -443,21 +443,25 @@ struct MaskGranularitySolver {
   }
 
   LogicalResult addExecuteRegionConstraints(scf::ExecuteRegionOp executeOp) {
-    WalkResult result = executeOp.getRegion().walk([&](scf::YieldOp yieldOp) {
-      if (yieldOp->getParentOp() != executeOp.getOperation()) {
-        return WalkResult::advance();
-      }
-      if (failed(
-              addYieldConstraints(executeOp->getResults(), yieldOp, executeOp))) {
-        return WalkResult::interrupt();
-      }
-      return WalkResult::advance();
-    });
+    WalkResult result = executeOp.getRegion().walk(
+        [this, executeOp](scf::YieldOp yieldOp) mutable {
+          const bool belongsToExecuteRegion =
+              yieldOp->getParentOp() == executeOp.getOperation();
+          if (!belongsToExecuteRegion) {
+            return WalkResult::advance();
+          }
+          if (failed(addYieldConstraints(executeOp->getResults(), yieldOp,
+                                         executeOp))) {
+            return WalkResult::interrupt();
+          }
+          return WalkResult::advance();
+        });
     return failure(result.wasInterrupted());
   }
 
   LogicalResult addIndexSwitchConstraints(scf::IndexSwitchOp indexSwitchOp) {
-    auto addBlockTerminator = [&](Block &block) -> LogicalResult {
+    auto addBlockTerminator =
+        [this, indexSwitchOp](Block &block) mutable -> LogicalResult {
       auto yieldOp = dyn_cast<scf::YieldOp>(block.getTerminator());
       if (!yieldOp) {
         return success();
@@ -478,14 +482,14 @@ struct MaskGranularitySolver {
 
   LogicalResult addWhileConstraints(scf::WhileOp whileOp) {
     return VMIControlFlowSupport::addWhileConstraints(
-        whileOp, [&](Value lhs, Value rhs, Operation *op) {
+        whileOp, [this](Value lhs, Value rhs, Operation *op) {
           return uniteEquivalentValues(lhs, rhs, op);
         });
   }
 
   LogicalResult addForConstraints(scf::ForOp forOp) {
     return VMIControlFlowSupport::addForConstraints(
-        forOp, [&](Value lhs, Value rhs, Operation *op) {
+        forOp, [this](Value lhs, Value rhs, Operation *op) {
           return uniteEquivalentValues(lhs, rhs, op);
         });
   }
@@ -532,12 +536,12 @@ struct MaskGranularitySolver {
     return success();
   }
 
-  bool hasVMIValueTypes(Operation *op) {
+  bool hasVMIValueTypes(Operation *op) const {
     return llvm::any_of(op->getOperandTypes(), containsVMIType) ||
            llvm::any_of(op->getResultTypes(), containsVMIType);
   }
 
-  bool hasVMIFunctionType(func::FuncOp func) {
+  bool hasVMIFunctionType(func::FuncOp func) const {
     FunctionType type = func.getFunctionType();
     return llvm::any_of(type.getInputs(), containsVMIType) ||
            llvm::any_of(type.getResults(), containsVMIType);
@@ -564,7 +568,8 @@ struct MaskGranularitySolver {
     }
 
     SmallVector<func::ReturnOp> returns;
-    callee.walk([&](func::ReturnOp returnOp) { returns.push_back(returnOp); });
+    callee.walk(
+        [&returns](func::ReturnOp returnOp) { returns.push_back(returnOp); });
     for (func::ReturnOp returnOp : returns) {
       for (auto [index, result] : llvm::enumerate(callOp.getResults())) {
         if (index >= returnOp.getNumOperands()) {
@@ -590,45 +595,16 @@ struct MaskGranularitySolver {
     }
   }
 
-  SmallVector<Type> getCallResultTypes(func::FuncOp func) {
-    SmallVector<Type> resultTypes;
-    bool found = false;
-    module.walk([&](func::CallOp call) {
-      if (call.getCallee() != func.getSymName()) {
-        return;
-      }
-      if (!found) {
-        resultTypes.assign(call.getResultTypes().begin(),
-                           call.getResultTypes().end());
-        found = true;
-        return;
-      }
-      if (resultTypes.size() != call.getNumResults()) {
-        return;
-      }
-      for (auto [index, type] : llvm::enumerate(call.getResultTypes())) {
-        if (index < resultTypes.size() && resultTypes[index] != type) {
-          resultTypes[index] = {};
-        }
-      }
-    });
-    return found ? resultTypes : SmallVector<Type>{};
-  }
-
   void rewriteFunctionType() {
-    module.walk([&](func::FuncOp func) {
+    module.walk([this](func::FuncOp func) {
       if (func.empty()) {
         return;
       }
 
-      SmallVector<Type> inputs;
-      inputs.reserve(func.getNumArguments());
-      for (BlockArgument arg : func.getArguments()) {
-        inputs.push_back(arg.getType());
-      }
-
+      SmallVector<Type> inputs = getFunctionInputTypes(func);
       SmallVector<Type> results;
-      SmallVector<Type> callResultTypes = getCallResultTypes(func);
+      SmallVector<Type> callResultTypes =
+          getConsistentCallResultTypes(module, func);
       auto it = firstReturnOperandsByFunc.find(func);
       if (!callResultTypes.empty()) {
         for (Type type : callResultTypes) {
@@ -693,7 +669,7 @@ struct MaskGranularitySolver {
   }
 
   Value rematerializeMaskProducer(Value value, VMIMaskType resultType,
-                                  Location loc, OpBuilder &builder) {
+                                  Location loc, OpBuilder &builder) const {
     if (auto createMask = value.getDefiningOp<VMICreateMaskOp>()) {
       return builder
           .create<VMICreateMaskOp>(loc, resultType, createMask.getActiveLanes())
diff --git a/lib/PTO/Transforms/VPTOExpandWrapperOps.cpp b/lib/PTO/Transforms/VPTOExpandWrapperOps.cpp
index 529c26b54b..5c6f0f7ef8 100644
--- a/lib/PTO/Transforms/VPTOExpandWrapperOps.cpp
+++ b/lib/PTO/Transforms/VPTOExpandWrapperOps.cpp
@@ -174,8 +174,7 @@ static Type getBufferElementType(Type type) {
   return {};
 }
 
-static Value offsetBufferPointer(Value basePtr, Type elementType,
-                                 Value elementOffset,
+static Value offsetBufferPointer(Value basePtr, Value elementOffset,
                                  PatternRewriter &rewriter, Location loc) {
   if (!basePtr) {
     return {};
@@ -603,13 +602,14 @@ static FailureOr<Value> packMadXt(Location loc, const MadXtConfig &config,
     return failure();
   }
 
-  auto constant = [&](uint64_t value) -> Value {
+  auto constant = [&rewriter, loc](uint64_t value) -> Value {
     return rewriter.create<arith::ConstantIntOp>(loc, value, mlir::pto::kValue64);
   };
-  auto shl = [&](Value value, uint64_t amount) -> Value {
+  auto shl = [&rewriter, loc, &constant](Value value,
+                                         uint64_t amount) -> Value {
     return rewriter.create<arith::ShLIOp>(loc, value, constant(amount));
   };
-  auto bitOr = [&](Value lhs, Value rhs) -> Value {
+  auto bitOr = [&rewriter, loc](Value lhs, Value rhs) -> Value {
     return rewriter.create<arith::OrIOp>(loc, lhs, rhs);
   };
 
@@ -859,6 +859,13 @@ struct LoadCbufToCbControl {
   Value dstStride;
 };
 
+struct PreparedLoadCbufOperands {
+  Value source;
+  Value destination;
+  Type elementType;
+  LoadCbufToCbControl control;
+};
+
 struct LoadCbufToMxControl {
   Value xStartPosition;
   Value yStartPosition;
@@ -951,6 +958,37 @@ deriveLoadCbufControl(const LoadCbufControlQuery &query) {
                              srcStride, dstStride};
 }
 
+template <typename LoadOp>
+static FailureOr<PreparedLoadCbufOperands>
+prepareLoadCbufOperands(LoadOp op, Value outerSize,
+                        PatternRewriter &rewriter) {
+  Location loc = op.getLoc();
+  Value source = materializeBufferPointer(op.getSource(), rewriter, loc);
+  Value destination =
+      materializeBufferPointer(op.getDestination(), rewriter, loc);
+  if (!source || !destination) {
+    return failure();
+  }
+  auto sourceType = dyn_cast<pto::PtrType>(source.getType());
+  if (!sourceType) {
+    return failure();
+  }
+
+  Type elementType = sourceType.getElementType();
+  FailureOr<LoadCbufToCbControl> control =
+      op.getMStart()
+          ? FailureOr<LoadCbufToCbControl>(LoadCbufToCbControl{
+                op.getMStart(), op.getKStart(), op.getMStep(), op.getKStep(),
+                op.getSrcStride(), op.getDstStride()})
+          : deriveLoadCbufControl(
+                {loc, outerSize, op.getK(), elementType, op.getStartRow(),
+                 op.getStartCol(), op.getTranspose(), rewriter});
+  if (failed(control)) {
+    return failure();
+  }
+  return PreparedLoadCbufOperands{source, destination, elementType, *control};
+}
+
 enum class CbufMxSide { Left, Right };
 
 struct LoadCbufMxControlQuery {
@@ -1196,8 +1234,8 @@ struct ExpandUvldPattern : public OpRewritePattern<pto::UvldOp> {
           "requires a recoverable pointer base for uvld expansion");
     }
 
-    Value loadPtr = offsetBufferPointer(basePtr, vecType.getElementType(),
-                                       op.getOffset(), rewriter, op.getLoc());
+    Value loadPtr =
+        offsetBufferPointer(basePtr, op.getOffset(), rewriter, op.getLoc());
     auto alignType = pto::AlignType::get(rewriter.getContext());
     Value align =
         rewriter.create<pto::VldasOp>(op.getLoc(), alignType, loadPtr);
@@ -1334,7 +1372,8 @@ struct ExpandDmaLoadPattern : public OpRewritePattern<pto::MteGmUbOp> {
 
     buildSoftwareLoopNest(
         rewriter, loc, loopPlan.softwareLoops, {zero, zero},
-        [&](Value srcOffset, Value dstOffset) {
+        [op, &rewriter, loc, zero, effectiveNBurst,
+         padding](Value srcOffset, Value dstOffset) mutable {
           Value source = offsetPointerByBytes(op.getSource(), srcOffset, rewriter, loc);
           Value destination =
               offsetPointerByBytes(op.getDestination(), dstOffset, rewriter, loc);
@@ -1376,7 +1415,8 @@ struct ExpandDmaStorePattern : public OpRewritePattern<pto::MteUbGmOp> {
 
     buildSoftwareLoopNest(
         rewriter, loc, loopPlan.softwareLoops, {zero, zero},
-        [&](Value srcOffset, Value dstOffset) {
+        [op, &rewriter, loc, zero, effectiveNBurst](Value srcOffset,
+                                                    Value dstOffset) mutable {
           Value source = offsetPointerByBytes(op.getSource(), srcOffset, rewriter, loc);
           Value destination =
               offsetPointerByBytes(op.getDestination(), dstOffset, rewriter, loc);
@@ -1434,7 +1474,7 @@ struct ExpandCubeLoadPattern : public OpRewritePattern<pto::MteGmL1Op> {
     DmaLoopPlan loopPlan = configureLoadToL1Loops(op, loops, one, rewriter);
     buildSoftwareLoopNest(
         rewriter, loc, loopPlan.softwareLoops, {zero, zero},
-        [&](Value srcOffset, Value dstOffset) {
+        [op, &rewriter, loc](Value srcOffset, Value dstOffset) mutable {
           Value source =
               offsetPointerByBytes(op.getSource(), srcOffset, rewriter, loc);
           Value destination = offsetPointerByBytes(op.getDestination(), dstOffset,
@@ -1510,7 +1550,7 @@ struct ExpandCubeStorePattern : public OpRewritePattern<pto::MteL1UbOp> {
                                                     loops.rend());
     buildSoftwareLoopNest(
         rewriter, loc, swLoopNestOrder, {zero, zero},
-        [&](Value srcOffset, Value dstOffset) {
+        [op, &rewriter, loc, zero](Value srcOffset, Value dstOffset) mutable {
           Value source =
               offsetPointerByBytes(op.getSource(), srcOffset, rewriter, loc);
           Value destination =
@@ -1613,43 +1653,28 @@ struct ExpandLeftLoadPattern : public OpRewritePattern<pto::MteL1L0aOp> {
   LogicalResult matchAndRewrite(pto::MteL1L0aOp op,
                                 PatternRewriter &rewriter) const override {
     Location loc = op.getLoc();
-    Value source = materializeBufferPointer(op.getSource(), rewriter, loc);
-    Value destination =
-        materializeBufferPointer(op.getDestination(), rewriter, loc);
-    auto sourceType = dyn_cast_or_null<pto::PtrType>(source.getType());
-    if (!sourceType) {
-      return rewriter.notifyMatchFailure(op, "expected typed L1 source");
-    }
-    Type elementType = sourceType.getElementType();
-    if (!destination) {
-      return rewriter.notifyMatchFailure(op, "expected pointer-like destination");
-    }
-    FailureOr<LoadCbufToCbControl> control = [&]() -> FailureOr<LoadCbufToCbControl> {
-      if (op.getMStart()) {
-        return LoadCbufToCbControl{op.getMStart(), op.getKStart(),
-                                   op.getMStep(), op.getKStep(),
-                                   op.getSrcStride(), op.getDstStride()};
-      }
-      return deriveLoadCbufControl(
-          {loc, op.getM(), op.getK(), elementType, op.getStartRow(),
-           op.getStartCol(), op.getTranspose(), rewriter});
-    }();
-    if (failed(control)) {
+    FailureOr<PreparedLoadCbufOperands> prepared =
+        prepareLoadCbufOperands(op, op.getM(), rewriter);
+    if (failed(prepared)) {
       return rewriter.notifyMatchFailure(op,
                                          "failed to derive load_cbuf_to_ca control");
     }
+    Value source = prepared->source;
+    Value destination = prepared->destination;
+    Type elementType = prepared->elementType;
+    const LoadCbufToCbControl &control = prepared->control;
     if (pto::isPTOFloat4PackedType(elementType)) {
       rewriter.create<pto::LoadCbufToCaS4Op>(
-          loc, source, destination, control->mStart,
-          control->kStart, control->mStep, control->kStep,
-          control->srcStride, control->dstStride,
+          loc, source, destination, control.mStart,
+          control.kStart, control.mStep, control.kStep,
+          control.srcStride, control.dstStride,
           rewriter.create<arith::ConstantIntOp>(loc, op.getTranspose(),
                                                 mlir::pto::kValue64));
     } else {
       auto load = rewriter.create<pto::LoadCbufToCaOp>(
-          loc, source, destination, control->mStart,
-          control->kStart, control->mStep, control->kStep,
-          control->srcStride, control->dstStride);
+          loc, source, destination, control.mStart,
+          control.kStart, control.mStep, control.kStep,
+          control.srcStride, control.dstStride);
       load->setAttr("transpose", rewriter.getBoolAttr(op.getTranspose()));
     }
     rewriter.eraseOp(op);
@@ -1663,43 +1688,28 @@ struct ExpandRightLoadPattern : public OpRewritePattern<pto::MteL1L0bOp> {
   LogicalResult matchAndRewrite(pto::MteL1L0bOp op,
                                 PatternRewriter &rewriter) const override {
     Location loc = op.getLoc();
-    Value source = materializeBufferPointer(op.getSource(), rewriter, loc);
-    Value destination =
-        materializeBufferPointer(op.getDestination(), rewriter, loc);
-    auto sourceType = dyn_cast_or_null<pto::PtrType>(source.getType());
-    if (!sourceType) {
-      return rewriter.notifyMatchFailure(op, "expected typed L1 source");
-    }
-    Type elementType = sourceType.getElementType();
-    if (!destination) {
-      return rewriter.notifyMatchFailure(op, "expected pointer-like destination");
-    }
-    FailureOr<LoadCbufToCbControl> control = [&]() -> FailureOr<LoadCbufToCbControl> {
-      if (op.getMStart()) {
-        return LoadCbufToCbControl{op.getMStart(), op.getKStart(),
-                                   op.getMStep(), op.getKStep(),
-                                   op.getSrcStride(), op.getDstStride()};
-      }
-      return deriveLoadCbufControl(
-          {loc, op.getN(), op.getK(), elementType, op.getStartRow(),
-           op.getStartCol(), op.getTranspose(), rewriter});
-    }();
-    if (failed(control)) {
+    FailureOr<PreparedLoadCbufOperands> prepared =
+        prepareLoadCbufOperands(op, op.getN(), rewriter);
+    if (failed(prepared)) {
       return rewriter.notifyMatchFailure(op,
                                          "failed to derive load_cbuf_to_cb control");
     }
+    Value source = prepared->source;
+    Value destination = prepared->destination;
+    Type elementType = prepared->elementType;
+    const LoadCbufToCbControl &control = prepared->control;
     if (pto::isPTOFloat4PackedType(elementType)) {
       rewriter.create<pto::LoadCbufToCbS4Op>(
-          loc, source, destination, control->mStart,
-          control->kStart, control->mStep, control->kStep,
-          control->srcStride, control->dstStride,
+          loc, source, destination, control.mStart,
+          control.kStart, control.mStep, control.kStep,
+          control.srcStride, control.dstStride,
           rewriter.create<arith::ConstantIntOp>(loc, op.getTranspose(),
                                                 mlir::pto::kValue64));
     } else {
       auto load = rewriter.create<pto::LoadCbufToCbOp>(
-          loc, source, destination, control->mStart,
-          control->kStart, control->mStep, control->kStep,
-          control->srcStride, control->dstStride);
+          loc, source, destination, control.mStart,
+          control.kStart, control.mStep, control.kStep,
+          control.srcStride, control.dstStride);
       load->setAttr("transpose", rewriter.getBoolAttr(op.getTranspose()));
     }
     rewriter.eraseOp(op);