fixes

slightly faster codegen

src/PIM/Common/IR/CompactAsmUtils.hpp

@@ -1,6 +1,7 @@
 #ifndef ONNX_MLIR_PIM_COMPACT_ASM_UTILS_HPP
 #define ONNX_MLIR_PIM_COMPACT_ASM_UTILS_HPP
 
+#include "mlir/IR/Builders.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/Value.h"
 #include "mlir/Support/LLVM.h"
@@ -9,8 +10,11 @@
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/Twine.h"
 #include "llvm/Support/LogicalResult.h"
 
+#include <optional>
+
 namespace onnx_mlir {
 namespace compact_asm {
 
@@ -21,6 +25,59 @@ enum class ListDelimiter {
   Paren
 };
 
+struct NumericSuffixRange {
+  StringRef prefix;
+  unsigned first = 0;
+  unsigned last = 0;
+};
+
+inline std::optional<NumericSuffixRange> getNumericSuffixRange(StringRef firstName, StringRef lastName) {
+  auto split = [](StringRef name) -> std::optional<std::pair<StringRef, unsigned>> {
+    size_t suffixStart = name.size();
+    while (suffixStart > 0 && name[suffixStart - 1] >= '0' && name[suffixStart - 1] <= '9')
+      --suffixStart;
+    if (suffixStart == name.size())
+      return std::nullopt;
+
+    unsigned number = 0;
+    if (name.drop_front(suffixStart).getAsInteger(10, number))
+      return std::nullopt;
+    return std::pair(name.take_front(suffixStart), number);
+  };
+
+  auto first = split(firstName);
+  auto last = split(lastName);
+  if (!first || !last || first->first != last->first || first->second > last->second)
+    return std::nullopt;
+  return NumericSuffixRange {first->first, first->second, last->second};
+}
+
+inline StringRef getInternedNumberedName(OpAsmParser& parser, StringRef prefix, unsigned number) {
+  return parser.getBuilder().getStringAttr((Twine(prefix) + Twine(number)).str()).getValue();
+}
+
+inline ParseResult parseOptionalRepeatCount(OpAsmParser& parser, int64_t& repeatCount) {
+  repeatCount = 1;
+  if (succeeded(parser.parseOptionalKeyword("x"))) {
+    if (parser.parseInteger(repeatCount) || repeatCount <= 0)
+      return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
+    return success();
+  }
+
+  const char* tokenStart = parser.getCurrentLocation().getPointer();
+  if (!tokenStart || tokenStart[0] != 'x' || tokenStart[1] < '0' || tokenStart[1] > '9')
+    return success();
+
+  StringRef fusedRepeat;
+  if (failed(parser.parseOptionalKeyword(&fusedRepeat)))
+    return failure();
+  if (!fusedRepeat.consume_front("x") || fusedRepeat.empty()
+      || fusedRepeat.getAsInteger(10, repeatCount) || repeatCount <= 0) {
+    return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
+  }
+  return success();
+}
+
 inline ParseResult parseOpenDelimiter(OpAsmParser& parser, ListDelimiter delimiter) {
   if (delimiter == ListDelimiter::Square)
     return parser.parseLSquare();
@@ -59,10 +116,8 @@ inline ParseResult parseCompressedRepeatedList(OpAsmParser& parser,
       return failure();
 
     int64_t repeatCount = 1;
-    if (succeeded(parser.parseOptionalKeyword("x"))) {
-      if (parser.parseInteger(repeatCount) || repeatCount <= 0)
-        return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
-    }
+    if (parseOptionalRepeatCount(parser, repeatCount))
+      return failure();
     for (int64_t index = 0; index < repeatCount; ++index)
       entries.push_back(entry);
 
@@ -86,10 +141,8 @@ parseCompressedIntegerEntries(OpAsmParser& parser, ListDelimiter delimiter, Smal
         return failure();
 
       int64_t repeatCount = 1;
-      if (succeeded(parser.parseOptionalKeyword("x"))) {
-        if (parser.parseInteger(repeatCount) || repeatCount <= 0)
-          return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
-      }
+      if (parseOptionalRepeatCount(parser, repeatCount))
+        return failure();
       for (int64_t repeat = 0; repeat < repeatCount; ++repeat)
         llvm::append_range(values, subgroup);
     }
@@ -109,10 +162,8 @@ parseCompressedIntegerEntries(OpAsmParser& parser, ListDelimiter delimiter, Smal
             return parser.emitError(parser.getCurrentLocation(), "step after 'by' must be positive");
         }
         int64_t repeatCount = 1;
-        if (succeeded(parser.parseOptionalKeyword("x"))) {
-          if (parser.parseInteger(repeatCount) || repeatCount <= 0)
-            return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
-        }
+        if (parseOptionalRepeatCount(parser, repeatCount))
+          return failure();
         if ((last - first) % step != 0) {
           return parser.emitError(parser.getCurrentLocation(),
                                   "range end must be reachable from start using the given step");
@@ -124,10 +175,8 @@ parseCompressedIntegerEntries(OpAsmParser& parser, ListDelimiter delimiter, Smal
       }
       else {
         int64_t repeatCount = 1;
-        if (succeeded(parser.parseOptionalKeyword("x"))) {
-          if (parser.parseInteger(repeatCount) || repeatCount <= 0)
-            return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
-        }
+        if (parseOptionalRepeatCount(parser, repeatCount))
+          return failure();
         for (int64_t index = 0; index < repeatCount; ++index)
           values.push_back(static_cast<IntT>(first));
       }
@@ -406,10 +455,8 @@ inline ParseResult parseCompressedTypeSequence(OpAsmParser& parser, SmallVectorI
 
   auto appendType = [&](Type type) -> ParseResult {
     int64_t repeatCount = 1;
-    if (succeeded(parser.parseOptionalKeyword("x"))) {
-      if (parser.parseInteger(repeatCount) || repeatCount <= 0)
-        return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
-    }
+    if (parseOptionalRepeatCount(parser, repeatCount))
+      return failure();
     for (int64_t index = 0; index < repeatCount; ++index)
       types.push_back(type);
     return success();
@@ -432,18 +479,23 @@ inline ParseResult parseCompressedOperandEntryWithFirst(OpAsmParser& parser,
     OpAsmParser::UnresolvedOperand lastOperand;
     if (parser.parseOperand(lastOperand))
       return failure();
-    if (firstOperand.name != lastOperand.name || firstOperand.number > lastOperand.number)
+    if (firstOperand.name == lastOperand.name && firstOperand.number <= lastOperand.number) {
+      for (unsigned number = firstOperand.number; number <= lastOperand.number; ++number)
+        operands.push_back({firstOperand.location, firstOperand.name, number});
+      return success();
+    }
+
+    auto range = getNumericSuffixRange(firstOperand.name, lastOperand.name);
+    if (!range || firstOperand.number != 0 || lastOperand.number != 0)
       return parser.emitError(parser.getCurrentLocation(), "invalid operand range");
-    for (unsigned number = firstOperand.number; number <= lastOperand.number; ++number)
-      operands.push_back({firstOperand.location, firstOperand.name, number});
+    for (unsigned number = range->first; number <= range->last; ++number)
+      operands.push_back({firstOperand.location, getInternedNumberedName(parser, range->prefix, number), 0});
     return success();
   }
 
   int64_t repeatCount = 1;
-  if (succeeded(parser.parseOptionalKeyword("x"))) {
-    if (parser.parseInteger(repeatCount) || repeatCount <= 0)
-      return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
-  }
+  if (parseOptionalRepeatCount(parser, repeatCount))
+    return failure();
   for (int64_t index = 0; index < repeatCount; ++index)
     operands.push_back(firstOperand);
   return success();
@@ -560,10 +612,8 @@ inline ParseResult parseCompressedOrTupleOperandList(OpAsmParser& parser,
       return failure();
 
     int64_t repeatCount = 1;
-    if (succeeded(parser.parseOptionalKeyword("x"))) {
-      if (parser.parseInteger(repeatCount) || repeatCount <= 0)
-        return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
-    }
+    if (parseOptionalRepeatCount(parser, repeatCount))
+      return failure();
     for (int64_t repeat = 0; repeat < repeatCount; ++repeat)
       llvm::append_range(operands, tupleOperands);
 
@@ -574,11 +624,8 @@ inline ParseResult parseCompressedOrTupleOperandList(OpAsmParser& parser,
       if (parseCompressedOperandSequence(parser, tupleOperands) || parser.parseRParen())
         return failure();
 
-      repeatCount = 1;
-      if (succeeded(parser.parseOptionalKeyword("x"))) {
-        if (parser.parseInteger(repeatCount) || repeatCount <= 0)
-          return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
-      }
+      if (parseOptionalRepeatCount(parser, repeatCount))
+        return failure();
       for (int64_t repeat = 0; repeat < repeatCount; ++repeat)
         llvm::append_range(operands, tupleOperands);
     }
@@ -608,10 +655,8 @@ parseCompressedOrTupleTypeList(OpAsmParser& parser, ListDelimiter delimiter, Sma
       return failure();
 
     int64_t repeatCount = 1;
-    if (succeeded(parser.parseOptionalKeyword("x"))) {
-      if (parser.parseInteger(repeatCount) || repeatCount <= 0)
-        return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
-    }
+    if (parseOptionalRepeatCount(parser, repeatCount))
+      return failure();
     for (int64_t repeat = 0; repeat < repeatCount; ++repeat)
       llvm::append_range(types, tupleTypes);
 
@@ -622,11 +667,8 @@ parseCompressedOrTupleTypeList(OpAsmParser& parser, ListDelimiter delimiter, Sma
       if (parseCompressedTypeSequence(parser, tupleTypes, /*allowEmpty=*/false) || parser.parseRParen())
         return failure();
 
-      repeatCount = 1;
-      if (succeeded(parser.parseOptionalKeyword("x"))) {
-        if (parser.parseInteger(repeatCount) || repeatCount <= 0)
-          return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
-      }
+      if (parseOptionalRepeatCount(parser, repeatCount))
+        return failure();
       for (int64_t repeat = 0; repeat < repeatCount; ++repeat)
         llvm::append_range(types, tupleTypes);
     }
@@ -639,10 +681,8 @@ parseCompressedOrTupleTypeList(OpAsmParser& parser, ListDelimiter delimiter, Sma
       return failure();
 
     int64_t repeatCount = 1;
-    if (succeeded(parser.parseOptionalKeyword("x"))) {
-      if (parser.parseInteger(repeatCount) || repeatCount <= 0)
-        return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
-    }
+    if (parseOptionalRepeatCount(parser, repeatCount))
+      return failure();
     for (int64_t repeat = 0; repeat < repeatCount; ++repeat)
       types.push_back(type);
 
@@ -678,13 +718,22 @@ inline ParseResult parseCompressedArgumentEntryWithFirst(OpAsmParser& parser,
     OpAsmParser::Argument lastArgument;
     if (parser.parseArgument(lastArgument))
       return failure();
-    if (firstArgument.ssaName.name != lastArgument.ssaName.name
-        || firstArgument.ssaName.number > lastArgument.ssaName.number) {
-      return parser.emitError(parser.getCurrentLocation(), "invalid argument range");
+    if (firstArgument.ssaName.name == lastArgument.ssaName.name
+        && firstArgument.ssaName.number <= lastArgument.ssaName.number) {
+      for (unsigned number = firstArgument.ssaName.number; number <= lastArgument.ssaName.number; ++number) {
+        OpAsmParser::Argument argument;
+        argument.ssaName = {firstArgument.ssaName.location, firstArgument.ssaName.name, number};
+        arguments.push_back(argument);
+      }
+      return success();
     }
-    for (unsigned number = firstArgument.ssaName.number; number <= lastArgument.ssaName.number; ++number) {
+
+    auto range = getNumericSuffixRange(firstArgument.ssaName.name, lastArgument.ssaName.name);
+    if (!range || firstArgument.ssaName.number != 0 || lastArgument.ssaName.number != 0)
+      return parser.emitError(parser.getCurrentLocation(), "invalid argument range");
+    for (unsigned number = range->first; number <= range->last; ++number) {
       OpAsmParser::Argument argument;
-      argument.ssaName = {firstArgument.ssaName.location, firstArgument.ssaName.name, number};
+      argument.ssaName = {firstArgument.ssaName.location, getInternedNumberedName(parser, range->prefix, number), 0};
       arguments.push_back(argument);
     }
     return success();

src/PIM/Common/IR/CoreBlockUtils.cpp

@@ -3,7 +3,7 @@
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 
-#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/STLExtras.h"
 
 #include "src/Accelerators/PIM/Common/IR/CoreBlockUtils.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
@@ -12,10 +12,8 @@ namespace onnx_mlir {
 
 bool isCoreStaticAddressOp(mlir::Operation* op) {
   if (mlir::isa<mlir::affine::AffineApplyOp,
-                mlir::arith::ConstantOp,
-                mlir::arith::AddIOp,
-                mlir::arith::SubIOp,
-                mlir::arith::MulIOp,
+                mlir::arith::ConstantOp, mlir::arith::AddIOp,
+                mlir::arith::SubIOp, mlir::arith::MulIOp,
                 mlir::arith::DivUIOp,
                 mlir::arith::DivSIOp,
                 mlir::arith::MinUIOp,
@@ -29,19 +27,33 @@ bool isCoreStaticAddressOp(mlir::Operation* op) {
                 mlir::memref::CollapseShapeOp,
                 mlir::memref::ExpandShapeOp>(op))
     return true;
-
-  if (auto selectOp = mlir::dyn_cast<mlir::arith::SelectOp>(op))
-    return selectOp.getType().isIntOrIndex();
-
-  return false;
+  auto selectOp = mlir::dyn_cast<mlir::arith::SelectOp>(op);
+  return selectOp && selectOp.getType().isIntOrIndex();
 }
 
-mlir::LogicalResult
-walkPimCoreBlock(mlir::Block& block,
-                 const StaticValueKnowledge& initialKnowledge,
-                 llvm::function_ref<mlir::LogicalResult(mlir::Operation&, const StaticValueKnowledge&)> callback) {
+namespace {
+
+enum class CoreWalkMode { ExecuteAllIterations, StructuralExtremes };
+using CoreWalkCallback =
+  llvm::function_ref<mlir::LogicalResult(mlir::Operation&, const StaticValueKnowledge&)>;
+
+static void propagateRegionResults(mlir::ValueRange results,
+                                   mlir::Region& region,
+                                   StaticValueKnowledge& knowledge) {
+  if (region.empty())
+    return;
+  auto yield = mlir::cast<mlir::scf::YieldOp>(region.front().getTerminator());
+  for (auto [result, yielded] : llvm::zip(results, yield.getOperands()))
+    knowledge.aliases[result] = resolveLoopCarriedAlias(yielded, knowledge);
+}
+
+static mlir::LogicalResult walkPimCoreBlockImpl(mlir::Block& block,
+                                                const StaticValueKnowledge& initialKnowledge,
+                                                CoreWalkMode mode,
+                                                CoreWalkCallback callback) {
   bool hasFailure = false;
   StaticValueKnowledge knowledge = initialKnowledge;
+  llvm::StringRef purpose = mode == CoreWalkMode::ExecuteAllIterations ? "codegen" : "verification";
   for (mlir::Operation& op : block) {
     if (mlir::isa<pim::PimHaltOp, mlir::scf::YieldOp>(op) || isCoreStaticAddressOp(&op))
       continue;
@@ -50,95 +62,12 @@ walkPimCoreBlock(mlir::Block& block,
       continue;
 
     if (auto forOp = mlir::dyn_cast<mlir::scf::ForOp>(op)) {
-      mlir::Block& loopBody = forOp.getRegion().front();
-      auto lowerBound = resolveIndexValue(forOp.getLowerBound(), knowledge);
-      auto upperBound = resolveIndexValue(forOp.getUpperBound(), knowledge);
+      auto lower = resolveIndexValue(forOp.getLowerBound(), knowledge);
+      auto upper = resolveIndexValue(forOp.getUpperBound(), knowledge);
       auto step = resolveIndexValue(forOp.getStep(), knowledge);
-      if (failed(lowerBound) || failed(upperBound) || failed(step) || *step <= 0) {
-        forOp.emitOpError("requires statically evaluable scf.for bounds for PIM codegen");
-        hasFailure = true;
-        continue;
-      }
-
-      llvm::SmallVector<mlir::Value> iterValues(forOp.getInitArgs().begin(), forOp.getInitArgs().end());
-      for (int64_t inductionValue = *lowerBound; inductionValue < *upperBound; inductionValue += *step) {
-        StaticValueKnowledge loopKnowledge = knowledge;
-        loopKnowledge.indexValues[forOp.getInductionVar()] = inductionValue;
-        for (auto [iterArg, iterValue] : llvm::zip_equal(forOp.getRegionIterArgs(), iterValues))
-          loopKnowledge.aliases[iterArg] = iterValue;
-
-        if (failed(walkPimCoreBlock(loopBody, loopKnowledge, callback)))
-          hasFailure = true;
-
-        auto yieldOp = mlir::cast<mlir::scf::YieldOp>(loopBody.getTerminator());
-        for (auto [index, yieldedValue] : llvm::enumerate(yieldOp.getOperands()))
-          iterValues[index] = resolveLoopCarriedAlias(yieldedValue, loopKnowledge);
-      }
-      continue;
-    }
-
-    if (auto ifOp = mlir::dyn_cast<mlir::scf::IfOp>(op)) {
-      auto condition = resolveIndexValue(ifOp.getCondition(), knowledge);
-      if (failed(condition)) {
-        ifOp.emitOpError("requires statically evaluable scf.if condition for PIM codegen");
-        hasFailure = true;
-        continue;
-      }
-
-      mlir::Region& selectedRegion = *condition != 0 ? ifOp.getThenRegion() : ifOp.getElseRegion();
-      if (!selectedRegion.empty())
-        if (failed(walkPimCoreBlock(selectedRegion.front(), knowledge, callback)))
-          hasFailure = true;
-      continue;
-    }
-
-    if (auto switchOp = mlir::dyn_cast<mlir::scf::IndexSwitchOp>(op)) {
-      auto selector = resolveIndexValue(switchOp.getArg(), knowledge);
-      if (failed(selector)) {
-        switchOp.emitOpError("requires a statically evaluable scf.index_switch selector for PIM codegen");
-        hasFailure = true;
-        continue;
-      }
-      mlir::Region* selected = &switchOp.getDefaultRegion();
-      for (auto [caseValue, caseRegion] : llvm::zip(switchOp.getCases(), switchOp.getCaseRegions()))
-        if (caseValue == *selector) {
-          selected = &caseRegion;
-          break;
-        }
-      if (failed(walkPimCoreBlock(selected->front(), knowledge, callback)))
-        hasFailure = true;
-      auto yield = mlir::cast<mlir::scf::YieldOp>(selected->front().getTerminator());
-      for (auto [result, yielded] : llvm::zip(switchOp.getResults(), yield.getOperands()))
-        knowledge.aliases[result] = resolveLoopCarriedAlias(yielded, knowledge);
-      continue;
-    }
-
-    if (failed(callback(op, knowledge)))
-      hasFailure = true;
-  }
-  return mlir::success(!hasFailure);
-}
-
-mlir::LogicalResult walkPimCoreBlockStructurally(
-  mlir::Block& block,
-  const StaticValueKnowledge& initialKnowledge,
-  llvm::function_ref<mlir::LogicalResult(mlir::Operation&, const StaticValueKnowledge&)> callback) {
-  bool hasFailure = false;
-  StaticValueKnowledge knowledge = initialKnowledge;
-  for (mlir::Operation& op : block) {
-    if (mlir::isa<pim::PimHaltOp, mlir::scf::YieldOp>(op) || isCoreStaticAddressOp(&op))
-      continue;
-    if (auto loadOp = mlir::dyn_cast<mlir::memref::LoadOp>(op);
-        loadOp && succeeded(resolveIndexValue(loadOp.getResult(), knowledge)))
-      continue;
-
-    if (auto forOp = mlir::dyn_cast<mlir::scf::ForOp>(op)) {
-      mlir::Block& loopBody = forOp.getRegion().front();
-      auto lowerBound = resolveIndexValue(forOp.getLowerBound(), knowledge);
-      auto upperBound = resolveIndexValue(forOp.getUpperBound(), knowledge);
-      auto step = resolveIndexValue(forOp.getStep(), knowledge);
-      if (failed(lowerBound) || failed(upperBound) || failed(step)) {
-        forOp.emitOpError("requires statically evaluable scf.for bounds for PIM verification");
+      if (failed(lower) || failed(upper) || failed(step)
+          || (mode == CoreWalkMode::ExecuteAllIterations && *step <= 0)) {
+        forOp.emitOpError() << "requires statically evaluable scf.for bounds for PIM " << purpose;
         hasFailure = true;
         continue;
       }
@@ -148,46 +77,57 @@ mlir::LogicalResult walkPimCoreBlockStructurally(
         continue;
       }
 
-      llvm::SmallVector<int64_t, 2> samples;
-      if (*lowerBound < *upperBound) {
-        samples.push_back(*lowerBound);
-        int64_t last = *lowerBound + ((*upperBound - 1 - *lowerBound) / *step) * *step;
-        if (last != *lowerBound)
-          samples.push_back(last);
-      }
-
-      for (int64_t inductionValue : samples) {
+      mlir::Block& body = forOp.getRegion().front();
+      llvm::SmallVector<mlir::Value> iterValues(forOp.getInitArgs().begin(), forOp.getInitArgs().end());
+      auto visitIteration = [&](int64_t induction, bool carryValues) {
         StaticValueKnowledge loopKnowledge = knowledge;
-        loopKnowledge.indexValues[forOp.getInductionVar()] = inductionValue;
-        for (auto [iterArg, iterValue] : llvm::zip_equal(forOp.getRegionIterArgs(), forOp.getInitArgs()))
-          loopKnowledge.aliases[iterArg] = iterValue;
+        loopKnowledge.indexValues[forOp.getInductionVar()] = induction;
+        for (auto [index, iterArg] : llvm::enumerate(forOp.getRegionIterArgs()))
+          loopKnowledge.aliases[iterArg] = carryValues ? iterValues[index] : forOp.getInitArgs()[index];
+        hasFailure |= failed(walkPimCoreBlockImpl(body, loopKnowledge, mode, callback));
+        auto yield = mlir::cast<mlir::scf::YieldOp>(body.getTerminator());
+        for (auto [index, yielded] : llvm::enumerate(yield.getOperands()))
+          iterValues[index] = resolveLoopCarriedAlias(yielded, loopKnowledge);
+      };
 
-        if (failed(walkPimCoreBlockStructurally(loopBody, loopKnowledge, callback)))
-          hasFailure = true;
+      if (mode == CoreWalkMode::ExecuteAllIterations) {
+        for (int64_t induction = *lower; induction < *upper; induction += *step)
+          visitIteration(induction, true);
       }
+      else if (*lower < *upper) {
+        visitIteration(*lower, false);
+        int64_t last = *lower + ((*upper - 1 - *lower) / *step) * *step;
+        if (last != *lower)
+          visitIteration(last, false);
+      }
+      for (auto [result, value] : llvm::zip(forOp.getResults(), iterValues))
+        knowledge.aliases[result] = value;
       continue;
     }
 
     if (auto ifOp = mlir::dyn_cast<mlir::scf::IfOp>(op)) {
-      if (failed(resolveIndexValue(ifOp.getCondition(), knowledge))) {
-        ifOp.emitOpError("requires statically evaluable scf.if condition for PIM verification");
+      auto condition = resolveIndexValue(ifOp.getCondition(), knowledge);
+      if (failed(condition)) {
+        ifOp.emitOpError() << "requires statically evaluable scf.if condition for PIM " << purpose;
         hasFailure = true;
         continue;
       }
-
-      if (!ifOp.getThenRegion().empty())
-        if (failed(walkPimCoreBlockStructurally(ifOp.getThenRegion().front(), knowledge, callback)))
-          hasFailure = true;
-      if (!ifOp.getElseRegion().empty())
-        if (failed(walkPimCoreBlockStructurally(ifOp.getElseRegion().front(), knowledge, callback)))
-          hasFailure = true;
+      mlir::Region& selected = *condition != 0 ? ifOp.getThenRegion() : ifOp.getElseRegion();
+      if (mode == CoreWalkMode::ExecuteAllIterations) {
+        hasFailure |= !selected.empty() && failed(walkPimCoreBlockImpl(selected.front(), knowledge, mode, callback));
+      }
+      else {
+        for (mlir::Region* region : {&ifOp.getThenRegion(), &ifOp.getElseRegion()})
+          hasFailure |= !region->empty() && failed(walkPimCoreBlockImpl(region->front(), knowledge, mode, callback));
+      }
+      propagateRegionResults(ifOp.getResults(), selected, knowledge);
       continue;
     }
 
     if (auto switchOp = mlir::dyn_cast<mlir::scf::IndexSwitchOp>(op)) {
       auto selector = resolveIndexValue(switchOp.getArg(), knowledge);
       if (failed(selector)) {
-        switchOp.emitOpError("requires a statically evaluable scf.index_switch selector for PIM verification");
+        switchOp.emitOpError() << "requires a statically evaluable scf.index_switch selector for PIM " << purpose;
         hasFailure = true;
         continue;
       }
@@ -197,19 +137,34 @@ mlir::LogicalResult walkPimCoreBlockStructurally(
           selected = &caseRegion;
           break;
         }
-      for (mlir::Region& region : switchOp->getRegions())
-        if (failed(walkPimCoreBlockStructurally(region.front(), knowledge, callback)))
-          hasFailure = true;
-      auto yield = mlir::cast<mlir::scf::YieldOp>(selected->front().getTerminator());
-      for (auto [result, yielded] : llvm::zip(switchOp.getResults(), yield.getOperands()))
-        knowledge.aliases[result] = resolveLoopCarriedAlias(yielded, knowledge);
+      for (mlir::Region& region : switchOp->getRegions()) {
+        if (mode == CoreWalkMode::ExecuteAllIterations && &region != selected)
+          continue;
+        hasFailure |= failed(walkPimCoreBlockImpl(region.front(), knowledge, mode, callback));
+      }
+      propagateRegionResults(switchOp.getResults(), *selected, knowledge);
       continue;
     }
 
-    if (failed(callback(op, knowledge)))
-      hasFailure = true;
+    hasFailure |= failed(callback(op, knowledge));
   }
   return mlir::success(!hasFailure);
 }
 
+} // namespace
+
+mlir::LogicalResult
+walkPimCoreBlock(mlir::Block& block,
+                 const StaticValueKnowledge& knowledge,
+                 llvm::function_ref<mlir::LogicalResult(mlir::Operation&, const StaticValueKnowledge&)> callback) {
+  return walkPimCoreBlockImpl(block, knowledge, CoreWalkMode::ExecuteAllIterations, callback);
+}
+
+mlir::LogicalResult walkPimCoreBlockStructurally(
+  mlir::Block& block,
+  const StaticValueKnowledge& knowledge,
+  llvm::function_ref<mlir::LogicalResult(mlir::Operation&, const StaticValueKnowledge&)> callback) {
+  return walkPimCoreBlockImpl(block, knowledge, CoreWalkMode::StructuralExtremes, callback);
+}
+
 } // namespace onnx_mlir

src/PIM/Common/IR/WeightUtils.cpp

@@ -146,17 +146,17 @@ void walkPimMvmVmmWeightUses(mlir::Operation* root, llvm::function_ref<void(mlir
   });
 }
 
-std::optional<unsigned> resolveWeightIndex(mlir::Operation* weightOwner, mlir::Value weight) {
+std::optional<unsigned> resolveWeightIndex(mlir::Operation* coreLikeOp, mlir::Value weight) {
   weight = stripMemRefAddressingOps(weight);
 
-  if (auto coreOp = mlir::dyn_cast_or_null<pim::PimCoreOp>(weightOwner)) {
+  if (auto coreOp = mlir::dyn_cast_or_null<pim::PimCoreOp>(coreLikeOp)) {
     for (unsigned weightIndex = 0; weightIndex < coreOp.getWeights().size(); ++weightIndex)
       if (coreOp.getWeightArgument(weightIndex) == weight)
         return weightIndex;
     return std::nullopt;
   }
 
-  if (auto coreBatchOp = mlir::dyn_cast_or_null<pim::PimCoreBatchOp>(weightOwner)) {
+  if (auto coreBatchOp = mlir::dyn_cast_or_null<pim::PimCoreBatchOp>(coreLikeOp)) {
     for (unsigned weightIndex = 0; weightIndex < coreBatchOp.getWeights().size(); ++weightIndex)
       if (coreBatchOp.getWeightArgument(weightIndex) == weight)
         return weightIndex;
@@ -167,7 +167,7 @@ std::optional<unsigned> resolveWeightIndex(mlir::Operation* weightOwner, mlir::V
 }
 
 llvm::FailureOr<ResolvedWeightView>
-resolveWeightView(mlir::Operation* weightOwner, mlir::Value weight, const StaticValueKnowledge& knowledge) {
+resolveWeightView(mlir::Operation* coreLikeOp, mlir::Value weight, const StaticValueKnowledge& knowledge) {
   llvm::SmallVector<mlir::Operation*> viewOps;
   mlir::Value current = weight;
 
@@ -179,7 +179,7 @@ resolveWeightView(mlir::Operation* weightOwner, mlir::Value weight, const Static
 
     if (auto defOp = current.getDefiningOp()) {
       if (auto getGlobalOp = mlir::dyn_cast<mlir::memref::GetGlobalOp>(defOp)) {
-        auto moduleOp = weightOwner ? weightOwner->getParentOfType<mlir::ModuleOp>() : mlir::ModuleOp {};
+        auto moduleOp = coreLikeOp ? coreLikeOp->getParentOfType<mlir::ModuleOp>() : mlir::ModuleOp {};
         auto globalOp = lookupGlobalForGetGlobal(moduleOp, getGlobalOp);
         if (!globalOp || !globalOp.getInitialValue())
           return mlir::failure();
@@ -297,15 +297,15 @@ resolveWeightView(mlir::Operation* weightOwner, mlir::Value weight, const Static
       continue;
     }
 
-    auto weightIndex = resolveWeightIndex(weightOwner, current);
+    auto weightIndex = resolveWeightIndex(coreLikeOp, current);
     if (!weightIndex)
       return mlir::failure();
 
-    if (auto coreOp = mlir::dyn_cast_or_null<pim::PimCoreOp>(weightOwner)) {
+    if (auto coreOp = mlir::dyn_cast_or_null<pim::PimCoreOp>(coreLikeOp)) {
       current = coreOp.getWeights()[*weightIndex];
       continue;
     }
-    if (auto coreBatchOp = mlir::dyn_cast_or_null<pim::PimCoreBatchOp>(weightOwner)) {
+    if (auto coreBatchOp = mlir::dyn_cast_or_null<pim::PimCoreBatchOp>(coreLikeOp)) {
       current = coreBatchOp.getWeights()[*weightIndex];
       continue;
     }

src/PIM/Common/IR/WeightUtils.hpp

@@ -45,9 +45,9 @@ bool hasOnlySpatialMvmVmmWeightUses(mlir::Value value);
 /// passes can identify globals that must remain weight-backed.
 void walkPimMvmVmmWeightUses(mlir::Operation* root, llvm::function_ref<void(mlir::OpOperand&)> callback);
 
-std::optional<unsigned> resolveWeightIndex(mlir::Operation* weightOwner, mlir::Value weight);
+std::optional<unsigned> resolveWeightIndex(mlir::Operation* coreLikeOp, mlir::Value weight);
 llvm::FailureOr<ResolvedWeightView>
-resolveWeightView(mlir::Operation* weightOwner, mlir::Value weight, const StaticValueKnowledge& knowledge = {});
+resolveWeightView(mlir::Operation* coreLikeOp, mlir::Value weight, const StaticValueKnowledge& knowledge = {});
 
 template <typename CoreLikeOpTy>
 llvm::SmallVector<unsigned, 8> getUsedWeightIndices(CoreLikeOpTy coreLikeOp) {