Fix codecheck findings

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
 

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);

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();
     }

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,23 +259,14 @@ 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()) {
     Operation *owner = use.getOwner();
     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;

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)) {

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) {

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;
     }

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) {