fixes

slightly faster codegen

README.md

@@ -170,7 +170,7 @@ Each validation run writes artifacts in the model workspace, for example under
 The compiler currently dumps dialect snapshots such as `spatial0.mlir`,
 `spatial1_graph.mlir`, `spatial2_trivial_merged.mlir`,
 `spatial3_scheduled_no_comm.mlir`, `spatial4_scheduled.mlir`, `pim0.mlir`, `pim1_buff.mlir`,
-`pim2_coalesced.mlir`, and `pim3_folded.mlir` when an output directory is
+`pim2_folded.mlir`, and `pim3_coalesced.mlir` when an output directory is
 available.
 
 To rerun the simulator manually with tracing after validation has produced a

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 (parseOptionalRepeatCount(parser, repeatCount))
-      if (parser.parseInteger(repeatCount) || repeatCount <= 0)
+      return failure();
-        return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
-    }
     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 (parseOptionalRepeatCount(parser, repeatCount))
-        if (parser.parseInteger(repeatCount) || repeatCount <= 0)
+        return failure();
-          return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
-      }
       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 (parseOptionalRepeatCount(parser, repeatCount))
-          if (parser.parseInteger(repeatCount) || repeatCount <= 0)
+          return failure();
-            return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
-        }
         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 (parseOptionalRepeatCount(parser, repeatCount))
-          if (parser.parseInteger(repeatCount) || repeatCount <= 0)
+          return failure();
-            return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
-        }
         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 (parseOptionalRepeatCount(parser, repeatCount))
-      if (parser.parseInteger(repeatCount) || repeatCount <= 0)
+      return failure();
-        return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
-    }
     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)
+    for (unsigned number = range->first; number <= range->last; ++number)
-      operands.push_back({firstOperand.location, firstOperand.name, number});
+      operands.push_back({firstOperand.location, getInternedNumberedName(parser, range->prefix, number), 0});
     return success();
   }
 
   int64_t repeatCount = 1;
-  if (succeeded(parser.parseOptionalKeyword("x"))) {
+  if (parseOptionalRepeatCount(parser, repeatCount))
-    if (parser.parseInteger(repeatCount) || repeatCount <= 0)
+    return failure();
-      return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
-  }
   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 (parseOptionalRepeatCount(parser, repeatCount))
-      if (parser.parseInteger(repeatCount) || repeatCount <= 0)
+      return failure();
-        return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
-    }
     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 (parseOptionalRepeatCount(parser, repeatCount))
-      if (succeeded(parser.parseOptionalKeyword("x"))) {
+        return failure();
-        if (parser.parseInteger(repeatCount) || repeatCount <= 0)
-          return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
-      }
       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 (parseOptionalRepeatCount(parser, repeatCount))
-      if (parser.parseInteger(repeatCount) || repeatCount <= 0)
+      return failure();
-        return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
-    }
     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 (parseOptionalRepeatCount(parser, repeatCount))
-      if (succeeded(parser.parseOptionalKeyword("x"))) {
+        return failure();
-        if (parser.parseInteger(repeatCount) || repeatCount <= 0)
-          return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
-      }
       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 (parseOptionalRepeatCount(parser, repeatCount))
-      if (parser.parseInteger(repeatCount) || repeatCount <= 0)
+      return failure();
-        return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
-    }
     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
+    if (firstArgument.ssaName.name == lastArgument.ssaName.name
-        || firstArgument.ssaName.number > lastArgument.ssaName.number) {
+        && firstArgument.ssaName.number <= lastArgument.ssaName.number) {
-      return parser.emitError(parser.getCurrentLocation(), "invalid argument range");
+      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::ConstantOp, mlir::arith::AddIOp,
-                mlir::arith::AddIOp,
+                mlir::arith::SubIOp, mlir::arith::MulIOp,
-                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;
-
+  auto selectOp = mlir::dyn_cast<mlir::arith::SelectOp>(op);
-  if (auto selectOp = mlir::dyn_cast<mlir::arith::SelectOp>(op))
+  return selectOp && selectOp.getType().isIntOrIndex();
-    return selectOp.getType().isIntOrIndex();
-
-  return false;
 }
 
-mlir::LogicalResult
+namespace {
-walkPimCoreBlock(mlir::Block& block,
+
-                 const StaticValueKnowledge& initialKnowledge,
+enum class CoreWalkMode { ExecuteAllIterations, StructuralExtremes };
-                 llvm::function_ref<mlir::LogicalResult(mlir::Operation&, const StaticValueKnowledge&)> callback) {
+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 lower = resolveIndexValue(forOp.getLowerBound(), knowledge);
-      auto lowerBound = resolveIndexValue(forOp.getLowerBound(), knowledge);
+      auto upper = resolveIndexValue(forOp.getUpperBound(), knowledge);
-      auto upperBound = resolveIndexValue(forOp.getUpperBound(), knowledge);
       auto step = resolveIndexValue(forOp.getStep(), knowledge);
-      if (failed(lowerBound) || failed(upperBound) || failed(step) || *step <= 0) {
+      if (failed(lower) || failed(upper) || failed(step)
-        forOp.emitOpError("requires statically evaluable scf.for bounds for PIM codegen");
+          || (mode == CoreWalkMode::ExecuteAllIterations && *step <= 0)) {
-        hasFailure = true;
+        forOp.emitOpError() << "requires statically evaluable scf.for bounds for PIM " << purpose;
-        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");
         hasFailure = true;
         continue;
       }
@@ -148,46 +77,57 @@ mlir::LogicalResult walkPimCoreBlockStructurally(
         continue;
       }
 
-      llvm::SmallVector<int64_t, 2> samples;
+      mlir::Block& body = forOp.getRegion().front();
-      if (*lowerBound < *upperBound) {
+      llvm::SmallVector<mlir::Value> iterValues(forOp.getInitArgs().begin(), forOp.getInitArgs().end());
-        samples.push_back(*lowerBound);
+      auto visitIteration = [&](int64_t induction, bool carryValues) {
-        int64_t last = *lowerBound + ((*upperBound - 1 - *lowerBound) / *step) * *step;
-        if (last != *lowerBound)
-          samples.push_back(last);
-      }
-
-      for (int64_t inductionValue : samples) {
         StaticValueKnowledge loopKnowledge = knowledge;
-        loopKnowledge.indexValues[forOp.getInductionVar()] = inductionValue;
+        loopKnowledge.indexValues[forOp.getInductionVar()] = induction;
-        for (auto [iterArg, iterValue] : llvm::zip_equal(forOp.getRegionIterArgs(), forOp.getInitArgs()))
+        for (auto [index, iterArg] : llvm::enumerate(forOp.getRegionIterArgs()))
-          loopKnowledge.aliases[iterArg] = iterValue;
+          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)))
+      if (mode == CoreWalkMode::ExecuteAllIterations) {
-          hasFailure = true;
+        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))) {
+      auto condition = resolveIndexValue(ifOp.getCondition(), knowledge);
-        ifOp.emitOpError("requires statically evaluable scf.if condition for PIM verification");
+      if (failed(condition)) {
+        ifOp.emitOpError() << "requires statically evaluable scf.if condition for PIM " << purpose;
         hasFailure = true;
         continue;
       }
-
+      mlir::Region& selected = *condition != 0 ? ifOp.getThenRegion() : ifOp.getElseRegion();
-      if (!ifOp.getThenRegion().empty())
+      if (mode == CoreWalkMode::ExecuteAllIterations) {
-        if (failed(walkPimCoreBlockStructurally(ifOp.getThenRegion().front(), knowledge, callback)))
+        hasFailure |= !selected.empty() && failed(walkPimCoreBlockImpl(selected.front(), knowledge, mode, callback));
-          hasFailure = true;
+      }
-      if (!ifOp.getElseRegion().empty())
+      else {
-        if (failed(walkPimCoreBlockStructurally(ifOp.getElseRegion().front(), knowledge, callback)))
+        for (mlir::Region* region : {&ifOp.getThenRegion(), &ifOp.getElseRegion()})
-          hasFailure = true;
+          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())
+      for (mlir::Region& region : switchOp->getRegions()) {
-        if (failed(walkPimCoreBlockStructurally(region.front(), knowledge, callback)))
+        if (mode == CoreWalkMode::ExecuteAllIterations && &region != selected)
-          hasFailure = true;
+          continue;
-      auto yield = mlir::cast<mlir::scf::YieldOp>(selected->front().getTerminator());
+        hasFailure |= failed(walkPimCoreBlockImpl(region.front(), knowledge, mode, callback));
-      for (auto [result, yielded] : llvm::zip(switchOp.getResults(), yield.getOperands()))
+      }
-        knowledge.aliases[result] = resolveLoopCarriedAlias(yielded, knowledge);
+      propagateRegionResults(switchOp.getResults(), *selected, knowledge);
       continue;
     }
 
-    if (failed(callback(op, knowledge)))
+    hasFailure |= failed(callback(op, knowledge));
-      hasFailure = true;
   }
   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) {

src/PIM/Compiler/CMakeLists.txt

@@ -17,6 +17,7 @@ add_pim_library(OMPimCompilerUtils
   PimCompilerUtils.cpp
   PimArtifactWriter.cpp
   PimCodeGen.cpp
+  PimCoreProgram.cpp
   PimMemoryLiveness.cpp
   PimWeightEmitter.cpp
 

src/PIM/Compiler/PimArtifactWriter.cpp

@@ -7,7 +7,7 @@
 #include <algorithm>
 #include <cassert>
 #include <cstring>
-#include <vector>
+#include <limits>
 
 #include "src/Accelerators/PIM/Common/IR/WeightUtils.hpp"
 #include "src/Accelerators/PIM/Compiler/PimArtifactWriter.hpp"
@@ -20,6 +20,65 @@ using namespace mlir;
 
 namespace onnx_mlir {
 
+PimInstructionWriter::PimInstructionWriter(raw_pwrite_stream& stream)
+: stream(stream), buffer(kBufferSize) {
+  pim_binary::writeHeader(stream);
+}
+
+PimInstructionWriter::PimInstructionWriter(raw_fd_ostream& stream)
+: PimInstructionWriter(static_cast<raw_pwrite_stream&>(stream)) {
+  fileStream = &stream;
+}
+
+bool PimInstructionWriter::consumeStreamError() {
+  if (!fileStream || !fileStream->has_error())
+    return false;
+  fileStream->clear_error();
+  hasFailure = true;
+  return true;
+}
+
+LogicalResult PimInstructionWriter::flushBuffer() {
+  if (hasFailure)
+    return failure();
+  if (bufferedBytes != 0) {
+    stream.write(buffer.data(), bufferedBytes);
+    bufferedBytes = 0;
+  }
+  return consumeStreamError() ? failure() : success();
+}
+
+LogicalResult PimInstructionWriter::append(const pim_binary::InstructionRecord& record) {
+  if (hasFailure || finalized || count == std::numeric_limits<uint32_t>::max()) {
+    hasFailure = true;
+    return failure();
+  }
+
+  if (bufferedBytes + pim_binary::kRecordSize > buffer.size() && failed(flushBuffer()))
+    return failure();
+
+  pim_binary::EncodedInstruction encoded = pim_binary::encodeInstructionRecord(record);
+  std::memcpy(buffer.data() + bufferedBytes, encoded.data(), encoded.size());
+  bufferedBytes += encoded.size();
+  ++count;
+  return success();
+}
+
+LogicalResult PimInstructionWriter::finalize() {
+  if (finalized)
+    return failure();
+  finalized = true;
+  if (hasFailure || failed(flushBuffer()))
+    return failure();
+
+  stream.flush();
+  if (consumeStreamError())
+    return failure();
+  pim_binary::patchInstructionCount(stream, count);
+  stream.flush();
+  return consumeStreamError() ? failure() : success();
+}
+
 OnnxMlirCompilerErrorCodes
 writeMemoryBinary(ModuleOp moduleOp, func::FuncOp funcOp, PimAcceleratorMemory& memory, StringRef outputDirPath) {
   auto memoryFilePath = (outputDirPath + "/memory.bin").str();

src/PIM/Compiler/PimArtifactWriter.hpp

@@ -2,16 +2,45 @@
 
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/BuiltinOps.h"
+#include "mlir/Support/LogicalResult.h"
 
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/JSON.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <vector>
 
 #include "onnx-mlir/Compiler/OMCompilerTypes.h"
+#include "src/Accelerators/PIM/Compiler/PimBinaryFormat.hpp"
 
 namespace onnx_mlir {
 
 class PimAcceleratorMemory;
 
+class PimInstructionWriter {
+public:
+  explicit PimInstructionWriter(llvm::raw_pwrite_stream& stream);
+  explicit PimInstructionWriter(llvm::raw_fd_ostream& stream);
+
+  mlir::LogicalResult append(const pim_binary::InstructionRecord& record);
+  mlir::LogicalResult finalize();
+
+private:
+  static constexpr size_t kBufferSize = 256 * 1024;
+  mlir::LogicalResult flushBuffer();
+  bool consumeStreamError();
+
+  llvm::raw_pwrite_stream& stream;
+  llvm::raw_fd_ostream* fileStream = nullptr;
+  std::vector<char> buffer;
+  size_t bufferedBytes = 0;
+  uint32_t count = 0;
+  bool finalized = false;
+  bool hasFailure = false;
+};
+
 OnnxMlirCompilerErrorCodes writeMemoryBinary(mlir::ModuleOp moduleOp,
                                              mlir::func::FuncOp funcOp,
                                              PimAcceleratorMemory& memory,

src/PIM/Compiler/PimBinaryFormat.hpp

@@ -1,5 +1,6 @@
 #pragma once
 
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/JSON.h"
@@ -53,6 +54,14 @@ enum class Opcode : uint32_t {
   sync = 32,
 };
 
+inline constexpr size_t kOpcodeCount = static_cast<size_t>(Opcode::sync) + 1;
+inline constexpr std::array<llvm::StringLiteral, kOpcodeCount> kOpcodeNames = {
+  "nop",   "sldi", "sld",   "sadd",  "ssub", "smul", "saddi", "smuli", "setbw", "mvmul", "vvadd",
+  "vvsub", "vvmul", "vvdmul", "vvmax", "vvsll", "vvsra", "vavg", "vrelu", "vtanh", "vsigm", "vsoftmax",
+  "vmv",   "vrsu", "vrsl",  "ld",    "st",   "lldi", "lmv",   "send",  "recv",  "wait",  "sync",
+};
+static_assert(kOpcodeNames.size() == kOpcodeCount);
+
 struct InstructionRecord {
   Opcode opcode = Opcode::nop;
   uint8_t rd = 0;
@@ -64,14 +73,29 @@ struct InstructionRecord {
   uint8_t flags = 0;
 };
 
+using EncodedInstruction = std::array<char, kRecordSize>;
+static_assert(kRecordSize == 20);
+static_assert(sizeof(EncodedInstruction) == kRecordSize);
+
+inline EncodedInstruction encodeInstructionRecord(const InstructionRecord& record) {
+  EncodedInstruction encoded = {};
+  encoded[0] = static_cast<char>(static_cast<uint8_t>(record.opcode));
+  encoded[1] = static_cast<char>(record.rd);
+  encoded[2] = static_cast<char>(record.r1);
+  encoded[3] = static_cast<char>(record.flags);
+  llvm::support::endian::write32le(encoded.data() + 4, static_cast<uint32_t>(record.r2OrImm));
+  llvm::support::endian::write32le(encoded.data() + 8, static_cast<uint32_t>(record.generic1));
+  llvm::support::endian::write32le(encoded.data() + 12, static_cast<uint32_t>(record.generic2));
+  llvm::support::endian::write32le(encoded.data() + 16, static_cast<uint32_t>(record.generic3));
+  return encoded;
+}
+
 inline void writeUint32LE(llvm::raw_ostream& os, uint32_t value) {
   std::array<char, sizeof(uint32_t)> bytes;
   llvm::support::endian::write32le(bytes.data(), value);
   os.write(bytes.data(), bytes.size());
 }
 
-inline void writeInt32LE(llvm::raw_ostream& os, int32_t value) { writeUint32LE(os, static_cast<uint32_t>(value)); }
-
 inline void writeHeader(llvm::raw_ostream& os) {
   os.write(kMagic, sizeof(kMagic));
   writeUint32LE(os, kVersion);
@@ -84,17 +108,6 @@ inline void patchInstructionCount(llvm::raw_pwrite_stream& os, uint32_t instruct
   os.pwrite(bytes.data(), bytes.size(), kCountOffset);
 }
 
-inline void writeInstructionRecord(llvm::raw_ostream& os, const InstructionRecord& record) {
-  os << static_cast<char>(static_cast<uint8_t>(record.opcode));
-  os << static_cast<char>(record.rd);
-  os << static_cast<char>(record.r1);
-  os << static_cast<char>(record.flags);
-  writeInt32LE(os, record.r2OrImm);
-  writeInt32LE(os, record.generic1);
-  writeInt32LE(os, record.generic2);
-  writeInt32LE(os, record.generic3);
-}
-
 inline int32_t toI32(int64_t value) { return onnx_mlir::pim::checkedI32OrCrash(value, "binary field"); }
 
 inline uint8_t toU8(int64_t value) {
@@ -107,113 +120,64 @@ inline int32_t getOptionalInt(const llvm::json::Object& object, llvm::StringRef
   return defaultValue;
 }
 
+struct InstructionJsonFormat {
+  bool rd;
+  bool r1;
+  bool offset;
+  llvm::StringLiteral r2;
+  llvm::StringLiteral generic1;
+  llvm::StringLiteral generic2;
+  llvm::StringLiteral generic3;
+};
+
+inline constexpr std::array<InstructionJsonFormat, kOpcodeCount> kInstructionJsonFormats = {{
+  {false, false, false, "",     "",     "",      ""    }, // nop
+  {true,  false, false, "imm",  "",     "",      ""    }, // sldi
+  {true,  true,  true,  "",     "",     "",      ""    }, // sld
+  {true,  true,  false, "rs2",  "",     "",      ""    }, // sadd
+  {true,  true,  false, "rs2",  "",     "",      ""    }, // ssub
+  {true,  true,  false, "rs2",  "",     "",      ""    }, // smul
+  {true,  true,  false, "imm",  "",     "",      ""    }, // saddi
+  {true,  true,  false, "imm",  "",     "",      ""    }, // smuli
+  {false, false, false, "",     "ibiw", "obiw",  ""    }, // setbw
+  {true,  true,  false, "mbiw", "relu", "group", ""    }, // mvmul
+  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vvadd
+  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vvsub
+  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vvmul
+  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vvdmul
+  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vvmax
+  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vvsll
+  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vvsra
+  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vavg
+  {true,  true,  true,  "",     "",     "",      "len" }, // vrelu
+  {true,  true,  true,  "",     "",     "",      "len" }, // vtanh
+  {true,  true,  true,  "",     "",     "",      "len" }, // vsigm
+  {true,  true,  true,  "",     "",     "",      "len" }, // vsoftmax
+  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vmv
+  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vrsu
+  {true,  true,  true,  "rs2",  "",     "",      "len" }, // vrsl
+  {true,  true,  true,  "",     "",     "",      "size"}, // ld
+  {true,  true,  true,  "",     "",     "",      "size"}, // st
+  {true,  false, true,  "imm",  "",     "",      "len" }, // lldi
+  {true,  true,  true,  "",     "",     "",      "len" }, // lmv
+  {true,  false, true,  "core", "",     "",      "size"}, // send
+  {true,  false, true,  "core", "",     "",      "size"}, // recv
+  {false, false, false, "",     "",     "",      ""    }, // wait
+  {false, false, false, "",     "",     "",      ""    }, // sync
+}};
+static_assert(kInstructionJsonFormats.size() == kOpcodeCount);
+
 inline Opcode opcodeFromString(llvm::StringRef opName) {
-  if (opName == "nop")
+  for (auto [index, name] : llvm::enumerate(kOpcodeNames))
-    return Opcode::nop;
+    if (opName == name)
-  if (opName == "sldi")
+      return static_cast<Opcode>(index);
-    return Opcode::sldi;
-  if (opName == "sld")
-    return Opcode::sld;
-  if (opName == "sadd")
-    return Opcode::sadd;
-  if (opName == "ssub")
-    return Opcode::ssub;
-  if (opName == "smul")
-    return Opcode::smul;
-  if (opName == "saddi")
-    return Opcode::saddi;
-  if (opName == "smuli")
-    return Opcode::smuli;
-  if (opName == "setbw")
-    return Opcode::setbw;
-  if (opName == "mvmul")
-    return Opcode::mvmul;
-  if (opName == "vvadd")
-    return Opcode::vvadd;
-  if (opName == "vvsub")
-    return Opcode::vvsub;
-  if (opName == "vvmul")
-    return Opcode::vvmul;
-  if (opName == "vvdmul")
-    return Opcode::vvdmul;
-  if (opName == "vvmax")
-    return Opcode::vvmax;
-  if (opName == "vvsll")
-    return Opcode::vvsll;
-  if (opName == "vvsra")
-    return Opcode::vvsra;
-  if (opName == "vavg")
-    return Opcode::vavg;
-  if (opName == "vrelu")
-    return Opcode::vrelu;
-  if (opName == "vtanh")
-    return Opcode::vtanh;
-  if (opName == "vsigm")
-    return Opcode::vsigm;
-  if (opName == "vsoftmax")
-    return Opcode::vsoftmax;
-  if (opName == "vmv")
-    return Opcode::vmv;
-  if (opName == "vrsu")
-    return Opcode::vrsu;
-  if (opName == "vrsl")
-    return Opcode::vrsl;
-  if (opName == "ld")
-    return Opcode::ld;
-  if (opName == "st")
-    return Opcode::st;
-  if (opName == "lldi")
-    return Opcode::lldi;
-  if (opName == "lmv")
-    return Opcode::lmv;
-  if (opName == "send")
-    return Opcode::send;
-  if (opName == "recv")
-    return Opcode::recv;
-  if (opName == "wait")
-    return Opcode::wait;
-  if (opName == "sync")
-    return Opcode::sync;
   llvm_unreachable("Unsupported PIM binary opcode");
 }
 
 inline llvm::StringRef opcodeToString(Opcode opcode) {
-  switch (opcode) {
+  size_t index = static_cast<size_t>(opcode);
-  case Opcode::nop:      return "nop";
+  assert(index < kOpcodeNames.size() && "Unsupported PIM binary opcode");
-  case Opcode::sldi:     return "sldi";
+  return kOpcodeNames[index];
-  case Opcode::sld:      return "sld";
-  case Opcode::sadd:     return "sadd";
-  case Opcode::ssub:     return "ssub";
-  case Opcode::smul:     return "smul";
-  case Opcode::saddi:    return "saddi";
-  case Opcode::smuli:    return "smuli";
-  case Opcode::setbw:    return "setbw";
-  case Opcode::mvmul:    return "mvmul";
-  case Opcode::vvadd:    return "vvadd";
-  case Opcode::vvsub:    return "vvsub";
-  case Opcode::vvmul:    return "vvmul";
-  case Opcode::vvdmul:   return "vvdmul";
-  case Opcode::vvmax:    return "vvmax";
-  case Opcode::vvsll:    return "vvsll";
-  case Opcode::vvsra:    return "vvsra";
-  case Opcode::vavg:     return "vavg";
-  case Opcode::vrelu:    return "vrelu";
-  case Opcode::vtanh:    return "vtanh";
-  case Opcode::vsigm:    return "vsigm";
-  case Opcode::vsoftmax: return "vsoftmax";
-  case Opcode::vmv:      return "vmv";
-  case Opcode::vrsu:     return "vrsu";
-  case Opcode::vrsl:     return "vrsl";
-  case Opcode::ld:       return "ld";
-  case Opcode::st:       return "st";
-  case Opcode::lldi:     return "lldi";
-  case Opcode::lmv:      return "lmv";
-  case Opcode::send:     return "send";
-  case Opcode::recv:     return "recv";
-  case Opcode::wait:     return "wait";
-  case Opcode::sync:     return "sync";
-  }
-  llvm_unreachable("Unsupported PIM binary opcode");
 }
 
 inline InstructionRecord makeInstructionRecord(const llvm::json::Object& instruction) {
@@ -221,43 +185,25 @@ inline InstructionRecord makeInstructionRecord(const llvm::json::Object& instruc
   std::optional<llvm::StringRef> opName = instruction.getString("op");
   assert(opName && "Missing op field in PIM instruction");
   record.opcode = opcodeFromString(*opName);
-  record.rd = toU8(getOptionalInt(instruction, "rd"));
+  const auto& format = kInstructionJsonFormats[static_cast<size_t>(record.opcode)];
-  record.r1 = toU8(getOptionalInt(instruction, "rs1"));
+  if (format.rd)
-
+    record.rd = toU8(getOptionalInt(instruction, "rd"));
-  switch (record.opcode) {
+  if (format.r1)
-  case Opcode::sldi:
+    record.r1 = toU8(getOptionalInt(instruction, "rs1"));
-  case Opcode::saddi:
+  if (!format.r2.empty())
-  case Opcode::smuli:
+    record.r2OrImm = getOptionalInt(instruction, format.r2);
-  case Opcode::lldi:  record.r2OrImm = getOptionalInt(instruction, "imm"); break;
+  if (!format.generic1.empty())
-  case Opcode::mvmul:
+    record.generic1 = getOptionalInt(instruction, format.generic1);
-    record.r2OrImm = getOptionalInt(instruction, "mbiw");
+  if (!format.generic2.empty())
-    record.generic1 = getOptionalInt(instruction, "relu");
+    record.generic2 = getOptionalInt(instruction, format.generic2);
-    record.generic2 = getOptionalInt(instruction, "group");
+  if (format.offset) {
-    break;
-  case Opcode::setbw:
-    record.generic1 = getOptionalInt(instruction, "ibiw");
-    record.generic2 = getOptionalInt(instruction, "obiw");
-    break;
-  case Opcode::send:
-  case Opcode::recv:
-    record.r2OrImm = getOptionalInt(instruction, "core");
-    record.generic3 = getOptionalInt(instruction, "size");
-    break;
-  default: record.r2OrImm = getOptionalInt(instruction, "rs2"); break;
-  }
-
-  if (record.opcode != Opcode::mvmul && record.opcode != Opcode::setbw) {
     if (auto* offsetValue = instruction.getObject("offset")) {
       record.generic1 = getOptionalInt(*offsetValue, "offset_select");
       record.generic2 = getOptionalInt(*offsetValue, "offset_value");
     }
   }
-
+  if (!format.generic3.empty())
-  if (instruction.get("len"))
+    record.generic3 = getOptionalInt(instruction, format.generic3);
-    record.generic3 = getOptionalInt(instruction, "len");
-  else if (instruction.get("size") && record.opcode != Opcode::send && record.opcode != Opcode::recv)
-    record.generic3 = getOptionalInt(instruction, "size");
-
   return record;
 }
 
@@ -271,98 +217,21 @@ inline llvm::json::Object makeInstructionJson(const InstructionRecord& record) {
     offset["offset_value"] = offsetValue;
     instruction["offset"] = std::move(offset);
   };
-
+  const auto& format = kInstructionJsonFormats[static_cast<size_t>(record.opcode)];
-  switch (record.opcode) {
+  if (format.rd)
-  case Opcode::sldi:
-    instruction["rd"] = static_cast<int64_t>(record.rd);
-    instruction["imm"] = record.r2OrImm;
-    break;
-  case Opcode::sld:
     instruction["rd"] = static_cast<int64_t>(record.rd);
+  if (format.r1)
     instruction["rs1"] = static_cast<int64_t>(record.r1);
+  if (!format.r2.empty())
+    instruction[format.r2] = record.r2OrImm;
+  if (!format.generic1.empty())
+    instruction[format.generic1] = record.generic1;
+  if (!format.generic2.empty())
+    instruction[format.generic2] = record.generic2;
+  if (format.offset)
     addOffset(record.generic1, record.generic2);
-    break;
+  if (!format.generic3.empty())
-  case Opcode::sadd:
+    instruction[format.generic3] = record.generic3;
-  case Opcode::ssub:
-  case Opcode::smul:
-    instruction["rd"] = static_cast<int64_t>(record.rd);
-    instruction["rs1"] = static_cast<int64_t>(record.r1);
-    instruction["rs2"] = record.r2OrImm;
-    break;
-  case Opcode::saddi:
-  case Opcode::smuli:
-    instruction["rd"] = static_cast<int64_t>(record.rd);
-    instruction["rs1"] = static_cast<int64_t>(record.r1);
-    instruction["imm"] = record.r2OrImm;
-    break;
-  case Opcode::setbw:
-    instruction["ibiw"] = record.generic1;
-    instruction["obiw"] = record.generic2;
-    break;
-  case Opcode::mvmul:
-    instruction["rd"] = static_cast<int64_t>(record.rd);
-    instruction["rs1"] = static_cast<int64_t>(record.r1);
-    instruction["mbiw"] = record.r2OrImm;
-    instruction["relu"] = record.generic1;
-    instruction["group"] = record.generic2;
-    break;
-  case Opcode::vvadd:
-  case Opcode::vvsub:
-  case Opcode::vvmul:
-  case Opcode::vvdmul:
-  case Opcode::vvmax:
-  case Opcode::vvsll:
-  case Opcode::vvsra:
-  case Opcode::vavg:
-  case Opcode::vmv:
-  case Opcode::vrsu:
-  case Opcode::vrsl:
-    instruction["rd"] = static_cast<int64_t>(record.rd);
-    instruction["rs1"] = static_cast<int64_t>(record.r1);
-    instruction["rs2"] = record.r2OrImm;
-    addOffset(record.generic1, record.generic2);
-    instruction["len"] = record.generic3;
-    break;
-  case Opcode::vrelu:
-  case Opcode::vtanh:
-  case Opcode::vsigm:
-  case Opcode::vsoftmax:
-    instruction["rd"] = static_cast<int64_t>(record.rd);
-    instruction["rs1"] = static_cast<int64_t>(record.r1);
-    addOffset(record.generic1, record.generic2);
-    instruction["len"] = record.generic3;
-    break;
-  case Opcode::ld:
-  case Opcode::st:
-    instruction["rd"] = static_cast<int64_t>(record.rd);
-    instruction["rs1"] = static_cast<int64_t>(record.r1);
-    addOffset(record.generic1, record.generic2);
-    instruction["size"] = record.generic3;
-    break;
-  case Opcode::lldi:
-    instruction["rd"] = static_cast<int64_t>(record.rd);
-    instruction["imm"] = record.r2OrImm;
-    addOffset(record.generic1, record.generic2);
-    instruction["len"] = record.generic3;
-    break;
-  case Opcode::lmv:
-    instruction["rd"] = static_cast<int64_t>(record.rd);
-    instruction["rs1"] = static_cast<int64_t>(record.r1);
-    addOffset(record.generic1, record.generic2);
-    instruction["len"] = record.generic3;
-    break;
-  case Opcode::send:
-  case Opcode::recv:
-    instruction["rd"] = static_cast<int64_t>(record.rd);
-    instruction["core"] = record.r2OrImm;
-    addOffset(record.generic1, record.generic2);
-    instruction["size"] = record.generic3;
-    break;
-  case Opcode::wait:
-  case Opcode::sync:
-  case Opcode::nop:  break;
-  }
-
   return instruction;
 }
 

src/PIM/Compiler/PimCodeGen.hpp

@@ -24,6 +24,10 @@
 
 namespace onnx_mlir {
 
+struct CompiledCoreProgram;
+struct CompiledTransposePlan;
+class PimInstructionWriter;
+
 struct MemEntry {
   size_t address;
   size_t size;
@@ -35,9 +39,7 @@ struct PhysicalSlotInfo {
   size_t size = 0;
 };
 
-struct MemoryPlanArtifacts {
+using MemoryPlanArtifacts = std::string;
-  std::string textReport;
-};
 
 struct MemoryValueKey {
   mlir::Value value;
@@ -98,6 +100,7 @@ class PimMemory {
   size_t nextPhysicalSlotId = 0;
 
   MemEntry* gatherMemEntry(mlir::Value value, std::optional<unsigned> lane = std::nullopt);
+  size_t allocateAddress(size_t size, const MemoryValueKey& key);
   void allocateGatheredMemory();
   void allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memEntry, MemoryReportKind reportKind);
   PhysicalSlotInfo allocatePhysicalSlot(size_t slotSize, const MemoryValueKey& key);
@@ -110,7 +113,6 @@ public:
   void allocateCore(mlir::Operation* op, std::optional<unsigned> lane = std::nullopt);
   MemoryReportRow getReportRow() const;
   const MemoryPlanArtifacts& getLivenessArtifacts() const { return livenessArtifacts; }
-  void remove(mlir::Value val);
 
   size_t getFirstAvailableAddress() const { return firstAvailableAddress; }
   MemEntry getMemEntry(const MemoryValueKey& key) const;
@@ -153,12 +155,11 @@ public:
                          uint64_t totalAllocaBytes);
   void setTotalWeightBytes(uint64_t bytes) { totalWeightBytes = bytes; }
   void flushReport();
-  void clean(mlir::Operation* op);
 };
 
 struct CoreEmissionJob {
   mlir::Operation* coreLikeOp = nullptr;
-  size_t originalCoreId = 0;
+  const CompiledCoreProgram* program = nullptr;
   size_t emittedCoreId = 0;
   llvm::SmallVector<unsigned, 4> lanes;
   std::optional<uint64_t> batchReportId;
@@ -166,11 +167,10 @@ struct CoreEmissionJob {
 
 class PimCodeGen {
   PimAcceleratorMemory& memory;
-  llvm::raw_fd_ostream& coreBinaryStream;
+  PimInstructionWriter& instructionWriter;
   llvm::raw_fd_ostream* coreJsonStream;
   const llvm::DenseMap<size_t, size_t>& emittedCoreIds;
   std::optional<unsigned> batchLane;
-  mutable uint32_t emittedInstructionCount = 0;
   mutable std::array<std::optional<int32_t>, 256> scalarRegisterValues = {};
 
   size_t addressOf(mlir::Value value, const StaticValueKnowledge& knowledge) const {
@@ -181,30 +181,44 @@ class PimCodeGen {
   void emitInstruction(const pim_binary::InstructionRecord& instruction) const;
   void updateScalarRegisterCache(const pim_binary::InstructionRecord& instruction) const;
 
+  void genSetRegisterImmediate(uint8_t registerNumber, int32_t immediate) const;
   void genSetRegisterImmediateUnsigned(size_t registerNumber, size_t immediate) const;
   void setupRd(size_t rdAddress, size_t rdOffset) const;
   void setupRdRs1(size_t rdAddress, size_t rdOffset, size_t rs1Address, size_t rs1Offset) const;
   void setupRdRs1Rs2(
     size_t rdAddress, size_t rdOffset, size_t rs1Address, size_t rs1Offset, size_t rs2Address, size_t rs2Offset) const;
 
-  void emitMemCopyOp(mlir::StringRef opName,
+  void emitMemCopyOp(pim_binary::Opcode opcode,
                      size_t rdAddr,
                      size_t rdOffset,
                      size_t rs1Addr,
                      size_t rs1Offset,
                      size_t size,
                      mlir::StringRef sizeFieldName = "size") const;
-  void emitCommunicationOp(mlir::StringRef opName, size_t bufferAddr, size_t coreId, size_t size) const;
+  void emitCommunicationOp(pim_binary::Opcode opcode, size_t bufferAddr, size_t coreId, size_t size) const;
   void emitMvmOp(size_t groupId, size_t rdAddr, size_t rdOffset, size_t rs1Addr, size_t rs1Offset) const;
 
 public:
+  void emitBinaryVectorOp(pim_binary::Opcode opcode,
+                          mlir::Value output,
+                          mlir::Value lhs,
+                          mlir::Value rhs,
+                          size_t byteSize,
+                          const StaticValueKnowledge& knowledge) const;
+  void emitUnaryVectorOp(pim_binary::Opcode opcode,
+                         mlir::Value output,
+                         mlir::Value input,
+                         size_t byteSize,
+                         const StaticValueKnowledge& knowledge,
+                         int32_t r2OrImm = 0,
+                         int32_t generic1 = 0) const;
+
   PimCodeGen(PimAcceleratorMemory& memory,
-             llvm::raw_fd_ostream& coreBinary,
+             PimInstructionWriter& instructionWriter,
              llvm::raw_fd_ostream* coreJson,
              const llvm::DenseMap<size_t, size_t>& emittedCoreIds)
-  : memory(memory), coreBinaryStream(coreBinary), coreJsonStream(coreJson), emittedCoreIds(emittedCoreIds) {}
+  : memory(memory), instructionWriter(instructionWriter), coreJsonStream(coreJson), emittedCoreIds(emittedCoreIds) {}
 
-  uint32_t getEmittedInstructionCount() const { return emittedInstructionCount; }
   void setBatchLane(std::optional<unsigned> lane) { batchLane = lane; }
   llvm::FailureOr<int64_t> indexOf(mlir::Value value, const StaticValueKnowledge& knowledge) const {
     return memory.getIndexValue(value, knowledge);
@@ -221,18 +235,7 @@ public:
   template <typename MVMTy>
   void codeGenMVMLikeOp(size_t mvmId, MVMTy mvmLikeOp, bool transposeMatrix, const StaticValueKnowledge& knowledge);
 
-  void codeGenVVAddOp(pim::PimVVAddOp vvaddOp, const StaticValueKnowledge& knowledge) const;
+  void codeGenTransposeOp(const CompiledTransposePlan& plan, const StaticValueKnowledge& knowledge) const;
-  void codeGenVVSubOp(pim::PimVVSubOp vvsubOp, const StaticValueKnowledge& knowledge) const;
-  void codeGenVVMulOp(pim::PimVVMulOp vvmulOp, const StaticValueKnowledge& knowledge) const;
-  void codeGenVVMaxOp(pim::PimVVMaxOp vvmaxOp, const StaticValueKnowledge& knowledge) const;
-  void codeGenVVDMulOp(pim::PimVVDMulOp vvdmulOp, const StaticValueKnowledge& knowledge) const;
-  void codeGenVAvgOp(pim::PimVAvgOp vavgOp, const StaticValueKnowledge& knowledge) const;
-  void codeGenVReluOp(pim::PimVReluOp vreluOp, const StaticValueKnowledge& knowledge) const;
-  void codeGenVTanhOp(pim::PimVTanhOp vtanhOp, const StaticValueKnowledge& knowledge) const;
-  void codeGenVSigmOp(pim::PimVSigmOp vsigmOp, const StaticValueKnowledge& knowledge) const;
-  void codeGenVSoftmaxOp(pim::PimVSoftmaxOp vsoftmaxOp, const StaticValueKnowledge& knowledge) const;
-  void codeGetGlobalOp(mlir::memref::GetGlobalOp getGlobalOp, const StaticValueKnowledge& knowledge) const;
-  void codeGenTransposeOp(pim::PimTransposeOp transposeOp, const StaticValueKnowledge& knowledge) const;
 };
 
 OnnxMlirCompilerErrorCodes compileToPimCode(mlir::ModuleOp& moduleOpRef, std::string& outputDirName);

src/PIM/Compiler/PimCompilerUtils.cpp

@@ -21,7 +21,7 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
   verifyExplicitPimCoreCount();
 
   if (pimOnlyCodegen) {
-    // Skip all the lowering passes and directly generate code for PIM.
+    pm.addPass(createEmitPimCodePass());
     return;
   }
 

src/PIM/Compiler/PimCoreProgram.cpp

@@ -0,0 +1,209 @@
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/IR/BuiltinTypes.h"
+
+#include "llvm/ADT/STLExtras.h"
+
+#include "src/Accelerators/PIM/Common/IR/CoreBlockUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
+#include "src/Accelerators/PIM/Compiler/PimCoreProgram.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
+
+using namespace llvm;
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace {
+
+static FailureOr<CompiledCoreOpKind> classifyCompiledCoreOpKind(Operation& op) {
+  if (isa<pim::PimMemCopyHostToDevOp>(op)) return CompiledCoreOpKind::Load;
+  if (isa<pim::PimMemCopyDevToHostOp>(op)) return CompiledCoreOpKind::Store;
+  if (isa<pim::PimMemCopyOp>(op)) return CompiledCoreOpKind::Lmv;
+  if (isa<pim::PimReceiveOp>(op)) return CompiledCoreOpKind::Receive;
+  if (isa<pim::PimSendOp>(op)) return CompiledCoreOpKind::Send;
+  if (isa<pim::PimConcatOp>(op)) return CompiledCoreOpKind::Concat;
+  if (isa<pim::PimVMMOp>(op)) return CompiledCoreOpKind::Vmm;
+  if (isa<pim::PimTransposeOp>(op)) return CompiledCoreOpKind::Transpose;
+  if (isa<pim::PimVVAddOp>(op)) return CompiledCoreOpKind::VVAdd;
+  if (isa<pim::PimVVSubOp>(op)) return CompiledCoreOpKind::VVSub;
+  if (isa<pim::PimVVMulOp>(op)) return CompiledCoreOpKind::VVMul;
+  if (isa<pim::PimVVMaxOp>(op)) return CompiledCoreOpKind::VVMax;
+  if (isa<pim::PimVVDMulOp>(op)) return CompiledCoreOpKind::VVDMul;
+  if (isa<pim::PimVAvgOp>(op)) return CompiledCoreOpKind::VAvg;
+  if (isa<pim::PimVReluOp>(op)) return CompiledCoreOpKind::VRelu;
+  if (isa<pim::PimVTanhOp>(op)) return CompiledCoreOpKind::VTanh;
+  if (isa<pim::PimVSigmOp>(op)) return CompiledCoreOpKind::VSigm;
+  if (isa<pim::PimVSoftmaxOp>(op)) return CompiledCoreOpKind::VSoftmax;
+  return failure();
+}
+
+static bool isStoragePreservingTranspose(ArrayRef<size_t> sourceShape, ArrayRef<int64_t> permutation) {
+  SmallVector<unsigned> sourceNonUnitDims;
+  SmallVector<unsigned> destinationSourceNonUnitDims;
+  for (auto [dim, size] : llvm::enumerate(sourceShape))
+    if (size != 1)
+      sourceNonUnitDims.push_back(dim);
+  for (int64_t sourceDim : permutation)
+    if (sourceShape[sourceDim] != 1)
+      destinationSourceNonUnitDims.push_back(static_cast<unsigned>(sourceDim));
+  return sourceNonUnitDims == destinationSourceNonUnitDims;
+}
+
+static FailureOr<CompiledTransposePlan> compileTransposePlan(pim::PimTransposeOp transposeOp) {
+  auto sourceType = cast<ShapedType>(transposeOp.getInput().getType());
+  ArrayRef<int64_t> sourceShape = sourceType.getShape();
+  size_t rank = sourceShape.size();
+  CompiledTransposePlan plan;
+  plan.source = transposeOp.getInput();
+  plan.destination = transposeOp.getOutputBuffer();
+  plan.elementBytes = getElementTypeSizeInBytes(sourceType.getElementType());
+  auto totalElements = pim::checkedSize(sourceType.getNumElements(), transposeOp, "transpose elements");
+  if (failed(totalElements)) return failure();
+  plan.totalElements = *totalElements;
+  auto totalBytes = pim::checkedMul(plan.totalElements, plan.elementBytes, transposeOp, "transpose byte size");
+  if (failed(totalBytes)) return failure();
+  plan.totalBytes = *totalBytes;
+
+  SmallVector<int64_t> permutation = map_to_vector(transposeOp.getPermutation().getAsRange<IntegerAttr>(),
+                                                   [](IntegerAttr attr) { return attr.getInt(); });
+  if (permutation.size() != rank) {
+    transposeOp.emitOpError("requires permutation rank to match source rank for PIM codegen");
+    return failure();
+  }
+
+  SmallVector<size_t> destinationShape(rank);
+  plan.destinationStrides.assign(rank, 1);
+  plan.destinationDimensionForSource.assign(rank, 0);
+  plan.destinationRewinds.assign(rank, 0);
+  SmallVector<bool> seenSourceDimensions(rank, false);
+  for (size_t dim = 0; dim < rank; ++dim) {
+    auto size = pim::checkedSize(sourceShape[dim], transposeOp, "transpose source dimension");
+    if (failed(size)) return failure();
+    plan.sourceShape.push_back(*size);
+  }
+  for (auto [destinationDim, sourceDim] : llvm::enumerate(permutation)) {
+    if (sourceDim < 0 || static_cast<size_t>(sourceDim) >= rank || seenSourceDimensions[sourceDim]) {
+      transposeOp.emitOpError("requires a valid permutation containing each source dimension exactly once");
+      return failure();
+    }
+    seenSourceDimensions[sourceDim] = true;
+    destinationShape[destinationDim] = plan.sourceShape[sourceDim];
+    plan.destinationDimensionForSource[sourceDim] = destinationDim;
+  }
+  for (size_t dim = rank; dim > 1; --dim) {
+    auto stride = pim::checkedMul(
+      plan.destinationStrides[dim - 1], destinationShape[dim - 1], transposeOp, "transpose destination stride");
+    if (failed(stride)) return failure();
+    plan.destinationStrides[dim - 2] = *stride;
+  }
+  for (size_t sourceDim = 0; sourceDim < rank; ++sourceDim) {
+    auto rewind = pim::checkedMul(plan.sourceShape[sourceDim],
+                                  plan.destinationStrides[plan.destinationDimensionForSource[sourceDim]],
+                                  transposeOp,
+                                  "transpose destination rewind");
+    if (failed(rewind)) return failure();
+    plan.destinationRewinds[sourceDim] = *rewind;
+  }
+  plan.storagePreserving = isStoragePreservingTranspose(plan.sourceShape, permutation);
+  return plan;
+}
+
+static LogicalResult compileCoreEmissionPlan(Block& block, SmallVectorImpl<CompiledCoreNode>& plan) {
+  for (Operation& op : block) {
+    if (isa<pim::PimHaltOp, scf::YieldOp, memref::GetGlobalOp>(op) || isCoreStaticAddressOp(&op))
+      continue;
+    if (auto loadOp = dyn_cast<memref::LoadOp>(op); loadOp && succeeded(compileIndexExpr(loadOp.getResult())))
+      continue;
+
+    if (auto forOp = dyn_cast<scf::ForOp>(op)) {
+      auto lower = compileIndexExpr(forOp.getLowerBound());
+      auto upper = compileIndexExpr(forOp.getUpperBound());
+      auto step = compileIndexExpr(forOp.getStep());
+      if (failed(lower) || failed(upper) || failed(step)) {
+        forOp.emitOpError("requires statically evaluable scf.for bounds for PIM codegen");
+        return failure();
+      }
+      CompiledCoreNode node;
+      node.kind = CompiledCoreNode::Kind::Loop;
+      node.op = forOp;
+      node.lowerBound = *lower;
+      node.upperBound = *upper;
+      node.step = *step;
+      node.loopBody = std::make_unique<SmallVector<CompiledCoreNode, 8>>();
+      if (failed(compileCoreEmissionPlan(forOp.getRegion().front(), *node.loopBody))) return failure();
+      plan.push_back(std::move(node));
+      continue;
+    }
+    if (auto ifOp = dyn_cast<scf::IfOp>(op)) {
+      auto condition = compileIndexExpr(ifOp.getCondition());
+      if (failed(condition)) {
+        ifOp.emitOpError("requires statically evaluable scf.if condition for PIM codegen");
+        return failure();
+      }
+      CompiledCoreNode node;
+      node.kind = CompiledCoreNode::Kind::If;
+      node.op = ifOp;
+      node.condition = *condition;
+      node.thenBody = std::make_unique<SmallVector<CompiledCoreNode, 8>>();
+      node.elseBody = std::make_unique<SmallVector<CompiledCoreNode, 8>>();
+      if (failed(compileCoreEmissionPlan(ifOp.getThenRegion().front(), *node.thenBody))) return failure();
+      if (!ifOp.getElseRegion().empty()
+          && failed(compileCoreEmissionPlan(ifOp.getElseRegion().front(), *node.elseBody)))
+        return failure();
+      plan.push_back(std::move(node));
+      continue;
+    }
+    if (auto switchOp = dyn_cast<scf::IndexSwitchOp>(op)) {
+      auto selector = compileIndexExpr(switchOp.getArg());
+      if (failed(selector)) {
+        switchOp.emitOpError("requires a statically evaluable scf.index_switch selector for PIM codegen");
+        return failure();
+      }
+      CompiledCoreNode node;
+      node.kind = CompiledCoreNode::Kind::IndexSwitch;
+      node.op = switchOp;
+      node.condition = *selector;
+      llvm::append_range(node.caseValues, switchOp.getCases());
+      for (Region& region : switchOp.getCaseRegions()) {
+        auto body = std::make_unique<SmallVector<CompiledCoreNode, 8>>();
+        if (failed(compileCoreEmissionPlan(region.front(), *body))) return failure();
+        node.caseBodies.push_back(std::move(body));
+      }
+      node.defaultBody = std::make_unique<SmallVector<CompiledCoreNode, 8>>();
+      if (failed(compileCoreEmissionPlan(switchOp.getDefaultRegion().front(), *node.defaultBody))) return failure();
+      plan.push_back(std::move(node));
+      continue;
+    }
+
+    auto opKind = classifyCompiledCoreOpKind(op);
+    if (failed(opKind)) {
+      InFlightDiagnostic diagnostic = op.emitError() << "unsupported codegen for op '" << op.getName() << "'";
+      if (auto coreOp = op.getParentOfType<pim::PimCoreOp>())
+        diagnostic << " inside pim.core " << coreOp.getCoreId();
+      else if (auto batchOp = op.getParentOfType<pim::PimCoreBatchOp>())
+        diagnostic << " inside pim.core_batch with laneCount " << batchOp.getLaneCount();
+      return failure();
+    }
+    CompiledCoreNode node;
+    node.op = &op;
+    node.opKind = *opKind;
+    if (*opKind == CompiledCoreOpKind::Transpose) {
+      auto transposePlan = compileTransposePlan(cast<pim::PimTransposeOp>(op));
+      if (failed(transposePlan)) return failure();
+      node.transposePlan = *transposePlan;
+    }
+    plan.push_back(std::move(node));
+  }
+  return success();
+}
+
+} // namespace
+
+LogicalResult compileCoreProgram(Operation* coreLikeOp, CompiledCoreProgram& program) {
+  Block& block = isa<pim::PimCoreOp>(coreLikeOp) ? cast<pim::PimCoreOp>(coreLikeOp).getBody().front()
+                                                 : cast<pim::PimCoreBatchOp>(coreLikeOp).getBody().front();
+  return compileCoreEmissionPlan(block, program.nodes);
+}
+
+} // namespace onnx_mlir

src/PIM/Compiler/PimCoreProgram.hpp

@@ -0,0 +1,74 @@
+#pragma once
+
+#include "mlir/IR/Operation.h"
+#include "mlir/Support/LogicalResult.h"
+
+#include "llvm/ADT/SmallVector.h"
+
+#include <memory>
+#include <optional>
+
+#include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp"
+
+namespace onnx_mlir {
+
+struct CompiledTransposePlan {
+  mlir::Value source;
+  mlir::Value destination;
+  size_t elementBytes = 0;
+  size_t totalElements = 0;
+  size_t totalBytes = 0;
+  llvm::SmallVector<size_t> sourceShape;
+  llvm::SmallVector<size_t> destinationStrides;
+  llvm::SmallVector<unsigned> destinationDimensionForSource;
+  llvm::SmallVector<size_t> destinationRewinds;
+  bool storagePreserving = false;
+};
+
+enum class CompiledCoreOpKind : uint8_t {
+  Load,
+  Store,
+  Lmv,
+  Receive,
+  Send,
+  Concat,
+  Vmm,
+  Transpose,
+  VVAdd,
+  VVSub,
+  VVMul,
+  VVMax,
+  VVDMul,
+  VAvg,
+  VRelu,
+  VTanh,
+  VSigm,
+  VSoftmax
+};
+
+struct CompiledCoreNode {
+  enum class Kind : uint8_t { Op, Loop, If, IndexSwitch };
+
+  Kind kind = Kind::Op;
+  mlir::Operation* op = nullptr;
+  CompiledCoreOpKind opKind = CompiledCoreOpKind::Load;
+  CompiledIndexExpr lowerBound;
+  CompiledIndexExpr upperBound;
+  CompiledIndexExpr step;
+  CompiledIndexExpr condition;
+  std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> loopBody;
+  std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> thenBody;
+  std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> elseBody;
+  llvm::SmallVector<int64_t> caseValues;
+  llvm::SmallVector<std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>>> caseBodies;
+  std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> defaultBody;
+  std::optional<CompiledTransposePlan> transposePlan;
+};
+
+struct CompiledCoreProgram {
+  llvm::SmallVector<CompiledCoreNode, 32> nodes;
+};
+
+mlir::LogicalResult compileCoreProgram(mlir::Operation* coreLikeOp, CompiledCoreProgram& program);
+
+} // namespace onnx_mlir

src/PIM/Compiler/PimMemoryLiveness.cpp

@@ -42,8 +42,6 @@ static MemoryValueKey getMemoryValueKey(mlir::Value value, std::optional<unsigne
 struct MemoryTouchInterval {
   uint64_t start = 0;
   uint64_t end = 0;
-  Operation* startOp = nullptr;
-  Operation* endOp = nullptr;
   Operation* firstTouchOp = nullptr;
   Operation* lastTouchOp = nullptr;
   uint64_t firstTouchPosition = 0;
@@ -218,22 +216,18 @@ static void appendAliasDescription(llvm::SmallVectorImpl<std::string>& aliases,
 struct OrderedTouchRange {
   uint64_t start = 0;
   uint64_t end = 0;
-  Operation* startOp = nullptr;
-  Operation* endOp = nullptr;
   bool escapedLoop = false;
 };
 
 static OrderedTouchRange
 getEffectiveTouchRange(mlir::Value definingValue, Operation* user, const OperationOrdering& ordering) {
-  OrderedTouchRange range {ordering.position.lookup(user), ordering.position.lookup(user), user, user, false};
+  OrderedTouchRange range {ordering.position.lookup(user), ordering.position.lookup(user), false};
   for (Operation* current = user; current; current = current->getParentOp()) {
     auto forOp = dyn_cast<scf::ForOp>(current);
     if (!forOp || isWithin(definingValue, &forOp.getRegion()))
       continue;
     range.start = std::min(range.start, ordering.position.lookup(forOp));
     range.end = std::max(range.end, ordering.subtreeEnd.lookup(forOp));
-    range.startOp = forOp;
-    range.endOp = forOp;
     range.escapedLoop = true;
   }
   return range;
@@ -247,8 +241,6 @@ computeMemoryTouchInterval(memref::AllocOp allocOp,
   MemoryTouchInterval interval;
   interval.start = ordering.position.lookup(allocOp);
   interval.end = interval.start;
-  interval.startOp = allocOp;
-  interval.endOp = allocOp;
   auto recordAlias = [&](mlir::Value value) {
     if (includeAliasDescriptions)
       appendAliasDescription(interval.aliasesFollowed, value);
@@ -350,18 +342,14 @@ computeMemoryTouchInterval(memref::AllocOp allocOp,
         if (!interval.hasRuntimeUse) {
           interval.start = range.start;
           interval.end = range.end;
-          interval.startOp = range.startOp;
-          interval.endOp = range.endOp;
           interval.hasRuntimeUse = true;
         }
         else {
           if (range.start < interval.start) {
             interval.start = range.start;
-            interval.startOp = range.startOp;
           }
           if (range.end > interval.end) {
             interval.end = range.end;
-            interval.endOp = range.endOp;
           }
         }
         continue;
@@ -380,8 +368,6 @@ computeMemoryTouchInterval(memref::AllocOp allocOp,
     interval.endUsedFallback = true;
     interval.start = ordering.position.lookup(allocOp);
     interval.end = fallbackEnd;
-    interval.startOp = allocOp;
-    interval.endOp = allocOp->getParentOp();
     interval.firstTouchPosition = interval.start;
     interval.lastTouchPosition = interval.end;
     addFallbackReason(interval.fallbackReason, "no runtime memory touch");
@@ -390,7 +376,6 @@ computeMemoryTouchInterval(memref::AllocOp allocOp,
 
   if (interval.endUsedFallback) {
     interval.end = std::max(interval.end, fallbackEnd);
-    interval.endOp = allocOp->getParentOp();
   }
 
   return interval;
@@ -447,8 +432,6 @@ SmallVector<LocalAllocInterval, 0> onnx_mlir::buildLocalAllocIntervals(Operation
     interval.start = touchInterval.start;
     interval.end = touchInterval.end;
     interval.size = *checkedSize;
-    interval.startOp = touchInterval.startOp;
-    interval.endOp = touchInterval.endOp;
     interval.firstTouchOp = touchInterval.firstTouchOp;
     interval.lastTouchOp = touchInterval.lastTouchOp;
     interval.firstTouchPosition = touchInterval.firstTouchPosition;
@@ -574,7 +557,7 @@ MemoryPlanArtifacts onnx_mlir::buildMemoryPlanArtifacts(Operation* coreLikeOp,
   double savedPercent =
     totalLogicalBytes == 0 ? 0.0 : 100.0 * static_cast<double>(savedBytes) / static_cast<double>(totalLogicalBytes);
 
-  raw_string_ostream os(artifacts.textReport);
+  raw_string_ostream os(artifacts);
   os << "=== PIM Memory Liveness Report ===\n";
   os << "Op: " << coreLikeOp->getName() << "\n";
   if (lane)

src/PIM/Compiler/PimMemoryLiveness.hpp

@@ -21,8 +21,6 @@ struct LocalAllocInterval {
   uint64_t start = 0;
   uint64_t end = 0;
   size_t size = 0;
-  mlir::Operation* startOp = nullptr;
-  mlir::Operation* endOp = nullptr;
   mlir::Operation* firstTouchOp = nullptr;
   mlir::Operation* lastTouchOp = nullptr;
   uint64_t firstTouchPosition = 0;

src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp

@@ -2180,39 +2180,55 @@ static Value createIm2colRows(const ConvLoweringState& state,
   auto elemType = preparedInput.type.getElementType();
   auto packedRowType = RankedTensorType::get(
     {plan.effectiveMaxParallelPixels * plan.patchSize}, elemType, plan.gemmInputRowsType.getEncoding());
-  auto zeroAttr = DenseElementsAttr::get(packedRowType, rewriter.getZeroAttr(elemType));
+  auto patchType = RankedTensorType::get({1, state.numChannelsIn, state.wHeight, state.wWidth}, elemType);
-  Value zeroRow = getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), zeroAttr, packedRowType);
+  auto patchRowType = RankedTensorType::get({plan.patchSize}, elemType);
+  bool hasPartialLane = plan.chunkNumPatches % plan.effectiveMaxParallelPixels != 0;
+  SmallVector<Value> im2colInputs {preparedInput.value};
   auto im2colComputeOp = createSpatComputeBatch(
     rewriter,
     loc,
     TypeRange {plan.gemmInputRowsType},
     plan.packedNumRows,
     {},
-    ValueRange {preparedInput.value, zeroRow},
+    im2colInputs,
     [&](detail::SpatComputeBatchBodyArgs args) {
       Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
       Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
       Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
       Value cPack = getOrCreateIndexConstant(rewriter, anchorOp, plan.effectiveMaxParallelPixels);
-      Value cNumPatches = getOrCreateIndexConstant(rewriter, anchorOp, plan.chunkNumPatches);
       Value laneStart = affineMulConst(rewriter, loc, args.lane, plan.effectiveMaxParallelPixels, anchorOp);
-      Value remaining = arith::SubIOp::create(rewriter, loc, cNumPatches, laneStart);
+      Value lanePatches = cPack;
-      Value isPartial = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::ult, remaining, cPack);
+      if (hasPartialLane) {
-      Value lanePatches = arith::SelectOp::create(rewriter, loc, isPartial, remaining, cPack);
+        Value cNumPatches = getOrCreateIndexConstant(rewriter, anchorOp, plan.chunkNumPatches);
-      auto patchType = RankedTensorType::get({1, state.numChannelsIn, state.wHeight, state.wWidth}, elemType);
+        Value remaining = arith::SubIOp::create(rewriter, loc, cNumPatches, laneStart);
-      auto patchRowType = RankedTensorType::get({plan.patchSize}, elemType);
+        Value isPartial = arith::CmpIOp::create(
-
+          rewriter, loc, arith::CmpIPredicate::ult, remaining, cPack);
+        lanePatches = arith::SelectOp::create(rewriter, loc, isPartial, remaining, cPack);
+      }
+      Value rowInit = tensor::EmptyOp::create(rewriter, loc, packedRowType.getShape(), elemType);
+      if (hasPartialLane) {
+        auto zeroAttr = cast<TypedAttr>(rewriter.getZeroAttr(elemType));
+        rowInit = linalg::MapOp::create(
+          rewriter, loc, ValueRange {}, rowInit,
+          [&](OpBuilder& builder, Location nestedLoc, ValueRange) {
+            Value zero = arith::ConstantOp::create(builder, nestedLoc, zeroAttr);
+            linalg::YieldOp::create(builder, nestedLoc, zero);
+          }).getResult().front();
+      }
       auto rowLoop = buildNormalizedScfFor(
         rewriter,
         loc,
         c0,
         lanePatches,
         c1,
-        ValueRange {args.inputs[1]},
+        ValueRange {rowInit},
         [&](OpBuilder&, Location nestedLoc, Value copyIndex, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
           Value patchIndex = arith::AddIOp::create(rewriter, nestedLoc, laneStart, copyIndex);
-          Value batchIndex =
+          Value batchIndex = state.batchSize == 1
-            affineAddFloorDivConst(rewriter, nestedLoc, patchIndex, plan.chunkStart, plan.numPatchesPerBatch, anchorOp);
+            ? c0
+            : affineAddFloorDivConst(
+                rewriter, nestedLoc, patchIndex, plan.chunkStart,
+                plan.numPatchesPerBatch, anchorOp);
           Value batchPatchIndex =
             affineAddModConst(rewriter, nestedLoc, patchIndex, plan.chunkStart, plan.numPatchesPerBatch, anchorOp);
           Value outHeightIndex = affineFloorDivConst(rewriter, nestedLoc, batchPatchIndex, state.outWidth, anchorOp);
@@ -2249,7 +2265,8 @@ static Value createIm2colRows(const ConvLoweringState& state,
         });
       if (failed(rowLoop))
         return failure();
-      publishGraphBatchPhysicalFragment(rewriter, loc, rowLoop->results.front(), args.outputs.front(), args.lane);
+      Value row = rowLoop->results.front();
+      publishGraphBatchPhysicalFragment(rewriter, loc, row, args.outputs.front(), args.lane);
       return success();
     });
 

src/PIM/Conversion/SpatialToPim/Common.cpp

@@ -1,6 +1,7 @@
 #include "mlir/IR/ValueRange.h"
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 
 #include "llvm/ADT/STLExtras.h"
 
@@ -59,6 +60,23 @@ bool hasLaterUserInBlock(mlir::Value value, Operation* operation) {
   return false;
 }
 
+static bool isTensorView(mlir::Value value) {
+  return isa_and_nonnull<tensor::CastOp,
+                         tensor::CollapseShapeOp,
+                         tensor::ExpandShapeOp,
+                         tensor::ExtractSliceOp,
+                         tensor::ReshapeOp>(value.getDefiningOp());
+}
+
+static bool isLoopCarriedOutput(mlir::Value operand, Operation* operation) {
+  auto argument = dyn_cast<BlockArgument>(operand);
+  if (!argument || argument.getArgNumber() == 0 || operation->getBlock() != argument.getOwner())
+    return false;
+  auto loop = dyn_cast_or_null<scf::ForOp>(argument.getOwner()->getParentOp());
+  return loop && cast<scf::YieldOp>(loop.getBody()->getTerminator())
+                   .getOperand(argument.getArgNumber() - 1) == operation->getResult(0);
+}
+
 mlir::Value getBestOutputTensorFromOperandsOrAllocate(RewriterBase& rewriter, Operation* operation) {
   assert("Only support operations with a single result" && operation->getNumResults() == 1);
   mlir::Value result = operation->getResult(0);
@@ -67,7 +85,11 @@ mlir::Value getBestOutputTensorFromOperandsOrAllocate(RewriterBase& rewriter, Op
 
   SmallVector<mlir::Value> operands = getOpOperandsSortedByUses(operation);
   auto validOperands = make_filter_range(operands, [operation, resultType](mlir::Value operand) {
-    return operand.getType() == resultType && !hasLaterUserInBlock(operand, operation);
+    return operand.getType() == resultType
+      && (!isa<BlockArgument>(operand) || isLoopCarriedOutput(operand, operation))
+      && !operand.getDefiningOp<arith::ConstantOp>()
+      && !isTensorView(operand)
+      && !hasLaterUserInBlock(operand, operation);
   });
   auto bestOperand = validOperands.begin();
 

src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp

@@ -2,6 +2,7 @@
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SCF/Utils/Utils.h"
@@ -112,7 +113,9 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
                          memref::MemRefDialect,
                          scf::SCFDialect,
                          BuiltinDialect>();
-  target.addLegalOp<spatial::SpatConcatOp,
+  target.addLegalOp<linalg::MapOp,
+                    linalg::YieldOp,
+                    spatial::SpatConcatOp,
                     spatial::SpatChannelReceiveOp,
                     spatial::SpatChannelSendOp,
                     spatial::SpatExtractRowsOp>();
@@ -186,7 +189,9 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
                                  memref::MemRefDialect,
                                  scf::SCFDialect,
                                  BuiltinDialect>();
-  coreBodyTarget.addLegalOp<spatial::SpatConcatOp,
+  coreBodyTarget.addLegalOp<linalg::MapOp,
+                            linalg::YieldOp,
+                            spatial::SpatConcatOp,
                             spatial::SpatChannelReceiveOp,
                             spatial::SpatChannelSendOp,
                             spatial::SpatExtractRowsOp>();
@@ -234,7 +239,7 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
                                       memref::MemRefDialect,
                                       scf::SCFDialect,
                                       BuiltinDialect>();
-  communicationTarget.addLegalOp<ModuleOp>();
+  communicationTarget.addLegalOp<ModuleOp, linalg::MapOp, linalg::YieldOp>();
   communicationTarget.addIllegalOp<spatial::SpatConcatOp,
                                    spatial::SpatChannelReceiveOp,
                                    spatial::SpatChannelSendOp,

src/PIM/Dialect/Pim/Transforms/Bufferization/OpBufferizationInterfaces.cpp

@@ -35,8 +35,54 @@ static StaticValueKnowledge getEnclosingBufferizationKnowledge(Operation* op) {
   return knowledge;
 }
 
+template <typename ConcreteModel, typename ConcreteOp>
+struct CopyOpInterface
+: DstBufferizableOpInterfaceExternalModel<ConcreteModel, ConcreteOp> {
+  bool bufferizesToMemoryRead(Operation* op, OpOperand& opOperand, const AnalysisState& state) const {
+    return cast<DestinationStyleOpInterface>(op).isDpsInput(&opOperand);
+  }
+
+  bool bufferizesToMemoryWrite(Operation* op, OpOperand& opOperand, const AnalysisState& state) const {
+    return cast<DestinationStyleOpInterface>(op).isDpsInit(&opOperand);
+  }
+
+  AliasingValueList getAliasingValues(Operation* op,
+                                      OpOperand& opOperand,
+                                      const AnalysisState& state) const {
+    auto dstOp = cast<DestinationStyleOpInterface>(op);
+    if (dstOp.isDpsInit(&opOperand))
+      return {{dstOp.getTiedOpResult(&opOperand), BufferRelation::Equivalent}};
+    return {};
+  }
+
+  AliasingOpOperandList getAliasingOpOperands(Operation* op,
+                                              Value value,
+                                              const AnalysisState& state) const {
+    auto result = dyn_cast<OpResult>(value);
+    if (!result || result.getDefiningOp() != op)
+      return {};
+    return {{cast<DestinationStyleOpInterface>(op).getTiedOpOperand(result), BufferRelation::Equivalent}};
+  }
+
+  bool mustBufferizeInPlace(Operation* op, OpOperand& opOperand, const AnalysisState& state) const {
+    return isa<UnrankedTensorType>(opOperand.get().getType());
+  }
+
+  bool isWritable(Operation* op, Value value, const AnalysisState& state) const { return isa<OpResult>(value); }
+
+  bool isNotConflicting(Operation* op,
+                        OpOperand* read,
+                        OpOperand* write,
+                        const AnalysisState& state) const {
+    if (read->getOwner() != op || write->getOwner() != op)
+      return false;
+    auto dstOp = cast<DestinationStyleOpInterface>(op);
+    return dstOp.isDpsInput(read) && dstOp.isDpsInit(write);
+  }
+};
+
 struct MemCopyHostToDevOpInterface
-: DstBufferizableOpInterfaceExternalModel<MemCopyHostToDevOpInterface, PimMemCopyHostToDevOp> {
+: CopyOpInterface<MemCopyHostToDevOpInterface, PimMemCopyHostToDevOp> {
   LogicalResult bufferize(Operation* op,
                           RewriterBase& rewriter,
                           const BufferizationOptions& options,
@@ -68,7 +114,7 @@ struct MemCopyHostToDevOpInterface
 };
 
 struct MemCopyDevToHostOpInterface
-: DstBufferizableOpInterfaceExternalModel<MemCopyDevToHostOpInterface, PimMemCopyDevToHostOp> {
+: CopyOpInterface<MemCopyDevToHostOpInterface, PimMemCopyDevToHostOp> {
   LogicalResult bufferize(Operation* op,
                           RewriterBase& rewriter,
                           const BufferizationOptions& options,
@@ -99,11 +145,7 @@ struct MemCopyDevToHostOpInterface
   }
 };
 
-struct MemCopyOpInterface : DstBufferizableOpInterfaceExternalModel<MemCopyOpInterface, PimMemCopyOp> {
+struct MemCopyOpInterface : CopyOpInterface<MemCopyOpInterface, PimMemCopyOp> {
-  bool bufferizesToMemoryRead(Operation* op, OpOperand& opOperand, const AnalysisState& state) const {
-    return !cast<DestinationStyleOpInterface>(op).isDpsInit(&opOperand);
-  }
-
   LogicalResult bufferize(Operation* op,
                           RewriterBase& rewriter,
                           const BufferizationOptions& options,

src/PIM/Dialect/Pim/Transforms/Bufferization/PimBufferizationPass.cpp

@@ -1,5 +1,4 @@
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
-#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h"
 #include "mlir/Dialect/Bufferization/Transforms/Transforms.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
@@ -7,16 +6,15 @@
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "mlir/IR/Dominance.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
-#include "mlir/Rewrite/PatternApplicator.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 #include "llvm/ADT/SmallPtrSet.h"
-#include "llvm/ADT/SmallSet.h"
 #include "llvm/Support/Casting.h"
 
 #include "Common/PimCommon.hpp"
+#include "Common/Support/Diagnostics.hpp"
 #include "Compiler/PimCodeGen.hpp"
 #include "Dialect/Pim/PimOps.hpp"
 #include "Dialect/Pim/Transforms/Bufferization/Common.hpp"
@@ -57,7 +55,10 @@ static StaticValueKnowledge seedCoreBatchKnowledge(pim::PimCoreBatchOp coreBatch
 }
 
 static LogicalResult
-lowerMemRefCopyToPimCopy(memref::CopyOp copyOp, PatternRewriter& rewriter, const StaticValueKnowledge& knowledge) {
+lowerMemRefCopyToPimCopy(memref::CopyOp copyOp,
+                         Value zeroOffset,
+                         PatternRewriter& rewriter,
+                         const StaticValueKnowledge& knowledge) {
   if (!copyOp->getParentOfType<pim::PimCoreOp>() && !copyOp->getParentOfType<pim::PimCoreBatchOp>())
     return failure();
 
@@ -68,7 +69,6 @@ lowerMemRefCopyToPimCopy(memref::CopyOp copyOp, PatternRewriter& rewriter, const
   if (sourceType.getElementType() != targetType.getElementType())
     return failure();
 
-  Value zeroOffset = getOrCreateIndexConstant(rewriter, copyOp, 0);
   auto sizeAttr = getMemRefSizeInBytesAttr(rewriter, copyOp.getOperation(), copyOp.getSource());
   if (failed(sizeAttr))
     return failure();
@@ -121,58 +121,35 @@ lowerMemRefCopyToPimCopy(memref::CopyOp copyOp, PatternRewriter& rewriter, const
   return success();
 }
 
-static LogicalResult verifyLoweredPimCopy(pim::PimMemCopyHostToDevOp copyOp,
+enum class ExpectedPimCopyDirection { HostToDevice, DeviceToHost, DeviceToDevice };
-                                          const StaticValueKnowledge& knowledge,
-                                          bool emitDiagnostic) {
-  bool sourceIsHost = isHostBackedPimAddress(copyOp.getHostSource(), knowledge);
-  bool targetIsHost = isHostBackedPimAddress(copyOp.getDeviceTarget(), knowledge);
-  bool sourceIsDevice = isDeviceLocalPimAddress(copyOp.getHostSource(), knowledge);
-  bool targetIsDevice = isDeviceLocalPimAddress(copyOp.getDeviceTarget(), knowledge);
-  if (!sourceIsHost || !targetIsDevice || targetIsHost || sourceIsDevice) {
-    if (emitDiagnostic)
-      copyOp.emitOpError() << "pim.memcp_hd requires a host-backed source and a device-local target: source="
-                           << copyOp.getHostSource() << " host=" << sourceIsHost << " device=" << sourceIsDevice
-                           << ", target=" << copyOp.getDeviceTarget() << " host=" << targetIsHost
-                           << " device=" << targetIsDevice;
-    return failure();
-  }
-  return success();
-}
 
-static LogicalResult verifyLoweredPimCopy(pim::PimMemCopyDevToHostOp copyOp,
+static LogicalResult verifyPimCopyEndpoints(Operation* copy,
-                                          const StaticValueKnowledge& knowledge,
+                                            Value source,
-                                          bool emitDiagnostic) {
+                                            Value target,
-  bool sourceIsHost = isHostBackedPimAddress(copyOp.getDeviceSource(), knowledge);
+                                            ExpectedPimCopyDirection direction,
-  bool targetIsHost = isHostBackedPimAddress(copyOp.getHostTarget(), knowledge);
+                                            const StaticValueKnowledge& knowledge,
-  bool sourceIsDevice = isDeviceLocalPimAddress(copyOp.getDeviceSource(), knowledge);
+                                            bool emitDiagnostic) {
-  bool targetIsDevice = isDeviceLocalPimAddress(copyOp.getHostTarget(), knowledge);
+  struct ExpectedEndpoints {
-  if (!targetIsHost || !sourceIsDevice || sourceIsHost || targetIsDevice) {
+    bool sourceHost, sourceDevice, targetHost, targetDevice;
-    if (emitDiagnostic)
+    StringLiteral description;
-      copyOp.emitOpError() << "pim.memcp_dh requires a device-local source and a host-backed target: source="
+  };
-                           << copyOp.getDeviceSource() << " host=" << sourceIsHost << " device=" << sourceIsDevice
+  static constexpr ExpectedEndpoints expected[] = {
-                           << ", target=" << copyOp.getHostTarget() << " host=" << targetIsHost
+    {true, false, false, true, "a host-backed source and a device-local target"},
-                           << " device=" << targetIsDevice;
+    {false, true, true, false, "a device-local source and a host-backed target"},
-    return failure();
+    {false, true, false, true, "device-local source and target operands"},
-  }
+  };
-  return success();
+  const auto& endpoints = expected[static_cast<unsigned>(direction)];
-}
+  bool sourceHost = isHostBackedPimAddress(source, knowledge);
-
+  bool sourceDevice = isDeviceLocalPimAddress(source, knowledge);
-static LogicalResult verifyLoweredPimCopy(pim::PimMemCopyOp copyOp,
+  bool targetHost = isHostBackedPimAddress(target, knowledge);
-                                          const StaticValueKnowledge& knowledge,
+  bool targetDevice = isDeviceLocalPimAddress(target, knowledge);
-                                          bool emitDiagnostic) {
+  bool valid = sourceHost == endpoints.sourceHost && sourceDevice == endpoints.sourceDevice
-  bool sourceIsHost = isHostBackedPimAddress(copyOp.getSource(), knowledge);
+    && targetHost == endpoints.targetHost && targetDevice == endpoints.targetDevice;
-  bool targetIsHost = isHostBackedPimAddress(copyOp.getTarget(), knowledge);
+  if (!valid && emitDiagnostic)
-  bool sourceIsDevice = isDeviceLocalPimAddress(copyOp.getSource(), knowledge);
+    copy->emitOpError() << "requires " << endpoints.description << ": source=" << source << " host=" << sourceHost
-  bool targetIsDevice = isDeviceLocalPimAddress(copyOp.getTarget(), knowledge);
+                        << " device=" << sourceDevice << ", target=" << target << " host=" << targetHost
-  if (!sourceIsDevice || !targetIsDevice || sourceIsHost || targetIsHost) {
+                        << " device=" << targetDevice;
-    if (emitDiagnostic)
+  return success(valid);
-      copyOp.emitOpError() << "pim.memcp requires device-local source and target operands: source="
-                           << copyOp.getSource() << " host=" << sourceIsHost << " device=" << sourceIsDevice
-                           << ", target=" << copyOp.getTarget() << " host=" << targetIsHost
-                           << " device=" << targetIsDevice;
-    return failure();
-  }
-  return success();
 }
 
 struct PimBufferizationPass : PassWrapper<PimBufferizationPass, OperationPass<ModuleOp>> {
@@ -191,14 +168,6 @@ private:
   LogicalResult verifyPimCopyAddressSpaces(ModuleOp moduleOp) const;
 };
 
-static LogicalResult applyPatternsOnce(Operation* op, PatternApplicator& applicator, PatternRewriter& rewriter) {
-  if (!op || !op->getBlock())
-    return failure();
-
-  rewriter.setInsertionPoint(op);
-  return applicator.matchAndRewrite(op, rewriter);
-}
-
 static void materializeWritableConstantDestinations(func::FuncOp funcOp) {
   SmallVector<OpOperand*> constantBackedRoots;
   llvm::SmallPtrSet<OpOperand*, 8> seenRoots;
@@ -237,69 +206,44 @@ static void materializeWritableConstantDestinations(func::FuncOp funcOp) {
   }
 }
 
-static LogicalResult verifyConflictFreePimCoreWrites(
+static LogicalResult verifyPimCoresNeedNoTensorCopies(
-  func::FuncOp funcOp, const bufferization::OneShotBufferizationOptions& options) {
+  ModuleOp module, const bufferization::OneShotBufferizationOptions& baseOptions) {
-  bufferization::AnalysisState analysisState(options);
+  static constexpr StringLiteral kExistingAlloc = "raptor.existing_core_alloc";
-  DominanceInfo dominance(funcOp);
+  OwningOpRef<ModuleOp> clone = module.clone();
-  size_t violationCount = 0;
+  clone->walk([&](bufferization::AllocTensorOp alloc) {
+    if (alloc->getParentOfType<pim::PimCoreOp>()
+        || alloc->getParentOfType<pim::PimCoreBatchOp>())
+      alloc->setAttr(kExistingAlloc, UnitAttr::get(module.getContext()));
+  });
 
-  auto verifyCore = [&](Operation* coreOp) {
+  auto options = baseOptions;
-    coreOp->walk([&](Operation* writeOp) {
+  options.bufferizeFunctionBoundaries = false;
-      for (OpOperand& write : writeOp->getOpOperands()) {
+  options.opFilter.allowOperation([](Operation* op) {
-        if (!isa<TensorType>(write.get().getType()) || !analysisState.bufferizesToMemoryWrite(write))
+    return isa<pim::PimCoreOp, pim::PimCoreBatchOp>(op)
-          continue;
+        || op->getParentOfType<pim::PimCoreOp>()
+        || op->getParentOfType<pim::PimCoreBatchOp>();
+  });
 
-        SmallVector<Value, 8> worklist {write.get()};
+  bufferization::BufferizationState state;
-        llvm::SmallDenseSet<Value, 8> visited;
+  if (failed(bufferization::insertTensorCopies(*clone, options, state))) {
-        bool hasConflict = false;
+    module.emitError("official one-shot analysis failed while verifying PIM core copy freedom");
-        Value conflictingAlias;
+    return failure();
-        Operation* conflictingUse = nullptr;
+  }
-        while (!worklist.empty() && !hasConflict) {
-          Value alias = worklist.pop_back_val();
-          if (!visited.insert(alias).second)
-            continue;
 
-          for (OpOperand& use : alias.getUses()) {
+  CappedDiagnosticReporter diagnostics;
-            bool usePrecedesWrite = false;
+  clone->walk([&](bufferization::AllocTensorOp alloc) {
-            for (Operation* ancestor = use.getOwner(); ancestor; ancestor = ancestor->getParentOp())
+    if (alloc->hasAttr(kExistingAlloc)
-              if (ancestor == writeOp || dominance.properlyDominates(ancestor, writeOp)) {
+        || (!alloc->getParentOfType<pim::PimCoreOp>()
-                usePrecedesWrite = true;
+            && !alloc->getParentOfType<pim::PimCoreBatchOp>()))
-                break;
+      return;
-              }
+    Operation* requiredBy = alloc->getUsers().empty()
-            if (usePrecedesWrite || analysisState.insideMutuallyExclusiveRegions(use.getOwner(), writeOp))
+      ? alloc.getOperation() : *alloc->getUsers().begin();
-              continue;
+    diagnostics.report(requiredBy, [](Operation* op) {
-            hasConflict = true;
+      op->emitOpError("official one-shot bufferization requires a tensor copy inside a PIM core");
-            conflictingAlias = alias;
-            conflictingUse = use.getOwner();
-            break;
-          }
-
-          if (alias.getDefiningOp<tensor::ExtractSliceOp>()
-              && llvm::all_of(alias.getUses(), [&](OpOperand& use) {
-                   return use.getOwner() == writeOp;
-                 }))
-            continue;
-          if (isa<OpResult>(alias))
-            for (bufferization::AliasingOpOperand tied : analysisState.getAliasingOpOperands(alias).getAliases())
-              worklist.push_back(tied.opOperand->get());
-        }
-
-        if (!hasConflict)
-          continue;
-        if (violationCount++ == 0)
-          writeOp->emitOpError() << "PIM core tensor write may modify an alias used later: operand #"
-                                 << write.getOperandNumber() << " (" << write.get() << "), alias="
-                                 << conflictingAlias << ", later use=" << *conflictingUse;
-      }
     });
-  };
+  });
-
+  diagnostics.emitSuppressedSummary(module, "required PIM core tensor copies");
-  funcOp.walk([&](pim::PimCoreOp coreOp) { verifyCore(coreOp); });
+  return success(!diagnostics.hasFailure());
-  funcOp.walk([&](pim::PimCoreBatchOp coreOp) { verifyCore(coreOp); });
-  if (violationCount != 0)
-    funcOp.emitError() << "found " << violationCount
-                       << " non-linear PIM tensor write(s); the first is reported above";
-  return success(violationCount == 0);
 }
 
 } // namespace
@@ -314,7 +258,7 @@ void PimBufferizationPass::runOnOperation() {
   options.setFunctionBoundaryTypeConversion(bufferization::LayoutMapOption::IdentityLayoutMap);
 
   materializeWritableConstantDestinations(funcOp);
-  if (failed(verifyConflictFreePimCoreWrites(funcOp, options))) {
+  if (failed(verifyPimCoresNeedNoTensorCopies(moduleOp, options))) {
     signalPassFailure();
     return;
   }
@@ -325,14 +269,8 @@ void PimBufferizationPass::runOnOperation() {
         || op->getParentOfType<pim::PimCoreBatchOp>();
   });
   bufferization::BufferizationState state;
-  if (failed(bufferization::insertTensorCopies(
+  if (failed(bufferization::insertTensorCopies(moduleOp, hostOptions, state))
-        moduleOp, hostOptions, state))) {
+      || failed(bufferization::bufferizeModuleOp(moduleOp, options, state))) {
-    moduleOp.emitError("Failed to bufferize PIM and Spatial ops");
-    signalPassFailure();
-    return;
-  }
-  if (failed(bufferization::bufferizeModuleOp(
-        moduleOp, options, state))) {
     moduleOp.emitError("Failed to bufferize PIM and Spatial ops");
     signalPassFailure();
     return;
@@ -369,15 +307,16 @@ void PimBufferizationPass::runOnOperation() {
           if (auto copyOp = dyn_cast<memref::CopyOp>(&op))
             addCopyOp(copyOp, opKnowledge);
           return success();
-        });
+      });
     }
   });
 
   bool hasFailed = false;
+  Value zeroOffset = getOrCreateIndexConstant(rewriter, funcOp, 0);
   for (const MemRefCopyWorkItem& workItem : copyWorklist) {
     memref::CopyOp copyOp = workItem.copyOp;
     rewriter.setInsertionPoint(copyOp);
-    if (failed(lowerMemRefCopyToPimCopy(copyOp, rewriter, workItem.knowledge)))
+    if (failed(lowerMemRefCopyToPimCopy(copyOp, zeroOffset, rewriter, workItem.knowledge)))
       hasFailed = true;
   }
   if (hasFailed) {
@@ -387,32 +326,10 @@ void PimBufferizationPass::runOnOperation() {
 
   RewritePatternSet contiguityPatterns(ctx);
   populatePimContiguityNormalizationPatterns(contiguityPatterns);
-  FrozenRewritePatternSet frozenContiguityPatterns(std::move(contiguityPatterns));
-  PatternApplicator contiguityApplicator(frozenContiguityPatterns);
-  contiguityApplicator.applyDefaultCostModel();
 
-  SmallVector<Operation*> contiguityWorklist;
+  GreedyRewriteConfig contiguityConfig;
-  moduleOp.walk([&](Operation* op) {
+  contiguityConfig.enableFolding(false);
-    if (isa<pim::PimMemCopyOp, pim::PimMemCopyHostToDevOp, pim::PimMemCopyDevToHostOp>(op))
+  if (failed(applyPatternsGreedily(moduleOp, std::move(contiguityPatterns), contiguityConfig))) {
-      contiguityWorklist.push_back(op);
-  });
-
-  hasFailed = false;
-  for (Operation* op : contiguityWorklist) {
-    if (auto copyOp = dyn_cast<pim::PimMemCopyOp>(op); copyOp && pim::isNormalizedCopyOp(copyOp))
-      continue;
-    if (auto copyOp = dyn_cast<pim::PimMemCopyHostToDevOp>(op); copyOp && pim::isNormalizedCopyOp(copyOp))
-      continue;
-    if (auto copyOp = dyn_cast<pim::PimMemCopyDevToHostOp>(op); copyOp && pim::isNormalizedCopyOp(copyOp))
-      continue;
-
-    if (failed(applyPatternsOnce(op, contiguityApplicator, rewriter))) {
-      op->emitOpError("failed to normalize PIM copy contiguity");
-      hasFailed = true;
-    }
-  }
-
-  if (hasFailed) {
     moduleOp.emitError("failed to normalize PIM copy contiguity during bufferization");
     signalPassFailure();
     return;
@@ -548,13 +465,19 @@ LogicalResult PimBufferizationPass::verifyPimCopyAddressSpaces(ModuleOp moduleOp
     (void) walkPimCoreBlockStructurally(
       coreLikeOp.getBody().front(), initialKnowledge, [&](Operation& op, const StaticValueKnowledge& knowledge) {
         if (auto copyOp = dyn_cast<pim::PimMemCopyOp>(&op);
-            copyOp && failed(verifyLoweredPimCopy(copyOp, knowledge, failureCount == 0)))
+            copyOp && failed(verifyPimCopyEndpoints(copyOp, copyOp.getSource(), copyOp.getTarget(),
+                                                    ExpectedPimCopyDirection::DeviceToDevice,
+                                                    knowledge, failureCount == 0)))
           ++failureCount;
         if (auto copyOp = dyn_cast<pim::PimMemCopyHostToDevOp>(&op);
-            copyOp && failed(verifyLoweredPimCopy(copyOp, knowledge, failureCount == 0)))
+            copyOp && failed(verifyPimCopyEndpoints(copyOp, copyOp.getHostSource(), copyOp.getDeviceTarget(),
+                                                    ExpectedPimCopyDirection::HostToDevice,
+                                                    knowledge, failureCount == 0)))
           ++failureCount;
         if (auto copyOp = dyn_cast<pim::PimMemCopyDevToHostOp>(&op);
-            copyOp && failed(verifyLoweredPimCopy(copyOp, knowledge, failureCount == 0)))
+            copyOp && failed(verifyPimCopyEndpoints(copyOp, copyOp.getDeviceSource(), copyOp.getHostTarget(),
+                                                    ExpectedPimCopyDirection::DeviceToHost,
+                                                    knowledge, failureCount == 0)))
           ++failureCount;
         return success();
       });

src/PIM/Dialect/Pim/Transforms/HostConstantFolding/HostConstantFoldingPass.cpp

@@ -41,7 +41,7 @@ struct HostConstantFoldingPass : PassWrapper<HostConstantFoldingPass, OperationP
       return;
     }
 
-    dumpModule(moduleOp, "pim3_folded");
+    dumpModule(moduleOp, "pim2_folded");
   }
 
   std::shared_ptr<const FrozenRewritePatternSet> patterns;

src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescingPass.cpp

@@ -203,7 +203,7 @@ struct PimMemoryCoalescingPass : PassWrapper<PimMemoryCoalescingPass, OperationP
     }
 
     emitReport(reportEntries);
-    dumpModule(getOperation(), "pim2_coalesced");
+    dumpModule(getOperation(), "pim3_coalesced");
   }
 };
 

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp

@@ -142,6 +142,71 @@ static std::optional<EmitLocalCollectionRun> buildLocalConcat(
        ? std::optional<EmitLocalCollectionRun>(std::move(run)) : std::nullopt;
 }
 
+static bool canLoopLocalCollection(
+    const EmitLocalCollectionRun &update, unsigned targetLaneCount) {
+  if (!update.collection || update.concatenatePayloads
+      || update.families.size() != 1
+      || !(update.lanes == LaneSet::all(targetLaneCount)))
+    return false;
+  RequirementFamily &requirement = *update.families.front()->requirement;
+  return requirement.targetLanes == update.lanes
+      && !requirement.producerProjection
+      && (!requirement.producerLocalOffsets
+          || requirement.producerLocalOffsets->size() == targetLaneCount);
+}
+
+static bool haveSameLocalCollectionContract(
+    const EmitLocalCollectionRun &lhs,
+    const EmitLocalCollectionRun &rhs) {
+  if (lhs.collection != rhs.collection)
+    return false;
+  RequirementFamily &left = *lhs.families.front()->requirement;
+  RequirementFamily &right = *rhs.families.front()->requirement;
+  if (left.producer->payload != right.producer->payload
+      || left.publicationFragmentType != right.publicationFragmentType)
+    return false;
+  if (lhs.collection->key.kind != FragmentCollectionKind::InsertAssembly)
+    return true;
+  const auto &entries =
+    lhs.collection->key.exchange->program.insertAssembly->entries;
+  const auto &leftEntry = entries[lhs.collectionPosition];
+  const auto &rightEntry = entries[rhs.collectionPosition];
+  return leftEntry.sourceTransform == rightEntry.sourceTransform
+      && leftEntry.sourceType == rightEntry.sourceType;
+}
+
+static void appendLocalUpdates(
+    BoundaryProgram &boundary,
+    SmallVectorImpl<EmitLocalCollectionRun> &updates,
+    unsigned targetLaneCount) {
+  for (size_t index = 0; index < updates.size();) {
+    EmitLocalCollectionRun &first = updates[index];
+    if (!canLoopLocalCollection(first, targetLaneCount)) {
+      boundary.instructions.push_back(std::move(first));
+      ++index;
+      continue;
+    }
+    size_t end = index + 1;
+    while (end < updates.size()
+           && canLoopLocalCollection(updates[end], targetLaneCount)
+           && haveSameLocalCollectionContract(first, updates[end]))
+      ++end;
+    if (end - index == 1) {
+      boundary.instructions.push_back(std::move(first));
+      ++index;
+      continue;
+    }
+    EmitLocalCollectionLoopRun run;
+    run.collection = first.collection;
+    run.lanes = first.lanes;
+    for (; index < end; ++index) {
+      run.positions.push_back(updates[index].collectionPosition);
+      run.families.push_back(updates[index].families.front());
+    }
+    boundary.instructions.push_back(std::move(run));
+  }
+}
+
 static void appendReceive(BoundaryProgram &boundary,
                           const ScheduledTransferSlice &slice,
                           CollectionTarget target) {
@@ -288,8 +353,7 @@ FailureOr<DeferredBoundaryPlan> buildDeferredBoundaryPlan(
       if (failed(addCoverage(*local.requirement, local.targetLanes, coverage)))
         return failure();
     }
-    for (EmitLocalCollectionRun &update : localUpdates)
+    appendLocalUpdates(boundary, localUpdates, exchange->targetLaneCount);
-      boundary.instructions.push_back(std::move(update));
     for (RequirementFamily &requirement : exchange->requirements)
       if (!(coverage.lookup(&requirement) == requirement.targetLanes))
         return exchange->deferred.emitOpError(

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.hpp

@@ -20,6 +20,12 @@ struct EmitLocalCollectionRun {
   LaneSet lanes;
   bool concatenatePayloads = false;
 };
+struct EmitLocalCollectionLoopRun {
+  const FragmentCollectionPlan* collection = nullptr;
+  llvm::SmallVector<unsigned> positions;
+  llvm::SmallVector<LocalAvailabilityFamily*> families;
+  LaneSet lanes;
+};
 struct EmitReceiveAssemblyRun {
   const FragmentCollectionPlan* collection = nullptr;
   llvm::SmallVector<ScheduledTransferSlice> slices;
@@ -33,7 +39,8 @@ struct ProduceDeferredResult {
 };
 
 using BoundaryInstruction =
-  std::variant<EmitSendRun, EmitLocalCollectionRun, EmitReceiveAssemblyRun,
+  std::variant<EmitSendRun, EmitLocalCollectionRun,
+               EmitLocalCollectionLoopRun, EmitReceiveAssemblyRun,
                ProduceDeferredResult>;
 struct BoundaryProgram {
   BoundaryKey key;

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp

@@ -507,6 +507,137 @@ static FailureOr<Value> transformAssemblySource(Value fragment, const DeferredIn
   llvm_unreachable("unknown deferred assembly source transform");
 }
 
+static FailureOr<Value> materializeLoopedLocalAssemblySource(
+    RequirementFamily &requirement,
+    const DeferredInsertAssemblyEntryTemplate &entry,
+    Value localOffset, DeferredExchangePlan &exchange,
+    DeferredEmissionContext &context) {
+  Value payload = requirement.producer->payload;
+  auto payloadType = dyn_cast<RankedTensorType>(payload.getType());
+  RankedTensorType sourceType = entry.sourceType;
+  if (entry.sourceTransform
+        == DeferredAssemblySourceTransform::RemoveLeadingUnitDimension
+      && payloadType && sourceType
+      && payloadType.getRank() > sourceType.getRank()
+      && payloadType.getElementType() == sourceType.getElementType()
+      && payloadType.getShape().take_back(sourceType.getRank())
+           == sourceType.getShape()) {
+    MixedSliceGeometry geometry;
+    int64_t rankDifference = payloadType.getRank() - sourceType.getRank();
+    geometry.offsets.assign(payloadType.getRank(),
+                            context.rewriter.getIndexAttr(0));
+    if (payload.getType() != requirement.publicationFragmentType)
+      geometry.offsets.front() = localOffset;
+    geometry.sizes.assign(rankDifference, context.rewriter.getIndexAttr(1));
+    for (int64_t dimension : sourceType.getShape())
+      geometry.sizes.push_back(context.rewriter.getIndexAttr(dimension));
+    geometry.strides.assign(payloadType.getRank(),
+                            context.rewriter.getIndexAttr(1));
+    return extractMixedSliceOrIdentity(
+      context.rewriter, exchange.deferred.getLoc(), payload, sourceType,
+      geometry);
+  }
+  auto fragment = materializeSendPayload(
+    requirement, localOffset, nullptr, context,
+    exchange.deferred.getLoc());
+  if (failed(fragment))
+    return failure();
+  return transformAssemblySource(*fragment, entry, exchange, context);
+}
+
+static LogicalResult emitLoopedLocalCollectionUpdate(
+    const EmitLocalCollectionLoopRun &run, Value lane, unsigned laneCount,
+    const DeferredResultPlan &resultPlan, DeferredEmissionContext &context) {
+  if (!run.collection || run.positions.size() < 2
+      || run.positions.size() != run.families.size())
+    return failure();
+  const FragmentCollectionPlan &collection = *run.collection;
+  DeferredExchangePlan &exchange = *collection.key.exchange;
+  const DeferredInsertAssemblyEntryTemplate *entry = nullptr;
+  if (collection.key.kind == FragmentCollectionKind::InsertAssembly)
+    entry = &exchange.program.insertAssembly
+               ->entries[run.positions.front()];
+  SmallVector<StaticIntSequence> offsetRows;
+  SmallVector<int64_t> positions;
+  offsetRows.reserve(run.families.size());
+  positions.reserve(run.positions.size());
+  for (auto [position, family] : llvm::zip_equal(run.positions,
+                                                 run.families)) {
+    RequirementFamily &requirement = *family->requirement;
+    offsetRows.push_back(requirement.producerLocalOffsets
+      ? *requirement.producerLocalOffsets
+      : StaticIntSequence::uniform(0, laneCount));
+    positions.push_back(position);
+  }
+  auto localOffsets = StaticIntGrid::fromRows(offsetRows);
+  if (failed(localOffsets))
+    return failure();
+  Operation *anchor = exchange.deferred;
+  Location loc = anchor->getLoc();
+  Value runtimeLane = lane ? lane : context.constants.getIndex(0);
+  Value current = context.fragmentCollections.lookup(collection.key);
+  if (!current)
+    current = createCollectionInitial(collection, context);
+  auto loop = buildNormalizedScfFor(
+    context.rewriter, loc, context.constants.getIndex(0),
+    context.constants.getIndex(run.positions.size()),
+    context.constants.getIndex(1), ValueRange {current},
+    [&](OpBuilder &, Location, Value action, ValueRange iterArgs,
+        SmallVectorImpl<Value> &yielded) -> LogicalResult {
+      Value position = lookup(positions, action, anchor, context, loc);
+      Value localOffset = localOffsets->emitLookup(
+        action, runtimeLane, anchor, context.constants, context.rewriter, loc);
+      RequirementFamily &requirement = *run.families.front()->requirement;
+      FailureOr<Value> source = entry
+        ? materializeLoopedLocalAssemblySource(
+            requirement, *entry, localOffset, exchange, context)
+        : materializeSendPayload(
+            requirement, localOffset, nullptr, context, loc);
+      if (failed(source))
+        return failure();
+      Value next;
+      if (entry) {
+        if (source->getType() != entry->sourceType)
+          return failure();
+        next = insertMixedSlice(
+          context.rewriter, loc, *source, iterArgs.front(),
+          lookupGeometry(resultPlan.assemblyGeometry, position, runtimeLane,
+                         anchor, context, loc));
+      } else {
+        bool grouped = collection.key.kind
+                    == FragmentCollectionKind::GroupedLeaf;
+        Value specialization = context.constants.getIndex(0);
+        Value leafPosition = position;
+        if (grouped) {
+          Value divisor = context.constants.getIndex(collection.positionCount);
+          specialization = arith::DivUIOp::create(
+            context.rewriter, loc, position, divisor);
+          leafPosition = arith::RemUIOp::create(
+            context.rewriter, loc, position, divisor);
+        }
+        unsigned leafIndex = collection.key.leafIndex;
+        const DeferredProjectionLeafTemplate &leaf =
+          exchange.program.leaves[leafIndex];
+        auto inserted = insertProjectionFragment(
+          *source, specialization, leafPosition,
+          grouped ? specialization : context.constants.getIndex(0),
+          runtimeLane, iterArgs.front(), leaf,
+          resultPlan.innerGeometry[leafIndex], exchange, grouped, context);
+        if (failed(inserted))
+          return failure();
+        next = *inserted;
+      }
+      if (!next)
+        return failure();
+      yielded.push_back(next);
+      return success();
+    });
+  if (failed(loop))
+    return failure();
+  context.fragmentCollections[collection.key] = loop->results.front();
+  return success();
+}
+
 static LogicalResult emitInsertAssemblyUpdate(const EmitReceiveAssemblyRun &run, Value lane, unsigned laneCount,
                                               const DeferredResultPlan &resultPlan, DeferredEmissionContext &context) {
   const FragmentCollectionPlan &collection = *run.collection;
@@ -683,6 +814,18 @@ static FailureOr<SmallVector<Value>> emitInstructions(ArrayRef<BoundaryInstructi
       if (failed(emitted))
         return exchange->deferred.emitOpError(
           "failed to update fragment collection from local availability"), failure();
+    } else if (auto update =
+                 std::get_if<EmitLocalCollectionLoopRun>(&instruction)) {
+      DeferredExchangePlan *exchange = update->collection->key.exchange;
+      const DeferredResultPlan *resultPlan = findResultPlan(results, exchange);
+      LogicalResult emitted = resultPlan
+        ? emitLoopedLocalCollectionUpdate(
+            *update, lane, laneCount, *resultPlan, context)
+        : failure();
+      if (failed(emitted))
+        return exchange->deferred.emitOpError(
+          "failed to update fragment collection from local assembly run"),
+          failure();
     } else if (auto assembly = std::get_if<EmitReceiveAssemblyRun>(&instruction)) {
       DeferredExchangePlan *exchange = assembly->collection->key.exchange;
       const DeferredResultPlan *resultPlan = findResultPlan(results, exchange);

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp

@@ -168,12 +168,16 @@ static LogicalResult materializeResultfulBatchRun(
 
       IRMapping mapper;
       mapper.map(*batch.getLaneArgument(), originalLane);
-      Value localLane = runLaneCount == 1
+      Value localLane;
-        ? getOrCreateIndexConstant(rewriter, batch.getOperation(), 0)
+      if (runLaneCount == 1)
-        : arith::SubIOp::create(
+        localLane = getOrCreateIndexConstant(rewriter, batch.getOperation(), 0);
-            builder, bodyLoc, originalLane,
+      else if (first.laneStart == 0)
-            getOrCreateIndexConstant(
+        localLane = originalLane;
-              rewriter, batch.getOperation(), first.laneStart));
+      else
+        localLane = arith::SubIOp::create(
+          builder, bodyLoc, originalLane,
+          getOrCreateIndexConstant(
+            rewriter, batch.getOperation(), first.laneStart));
       for (auto [index, weight] : llvm::enumerate(batch.getWeights()))
         mapper.map(*batch.getWeightArgument(index), getBlockOperand(block, scheduledWeights, weight));
       SmallVector<DeferredInputPlan> inputPlans;