fixes

slightly faster codegen

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> {
-  bool bufferizesToMemoryRead(Operation* op, OpOperand& opOperand, const AnalysisState& state) const {
-    return !cast<DestinationStyleOpInterface>(op).isDpsInit(&opOperand);
-  }
-
+struct MemCopyOpInterface : CopyOpInterface<MemCopyOpInterface, PimMemCopyOp> {
   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,
-                                          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();
-}
+enum class ExpectedPimCopyDirection { HostToDevice, DeviceToHost, DeviceToDevice };
 
-static LogicalResult verifyLoweredPimCopy(pim::PimMemCopyDevToHostOp copyOp,
-                                          const StaticValueKnowledge& knowledge,
-                                          bool emitDiagnostic) {
-  bool sourceIsHost = isHostBackedPimAddress(copyOp.getDeviceSource(), knowledge);
-  bool targetIsHost = isHostBackedPimAddress(copyOp.getHostTarget(), knowledge);
-  bool sourceIsDevice = isDeviceLocalPimAddress(copyOp.getDeviceSource(), knowledge);
-  bool targetIsDevice = isDeviceLocalPimAddress(copyOp.getHostTarget(), knowledge);
-  if (!targetIsHost || !sourceIsDevice || sourceIsHost || targetIsDevice) {
-    if (emitDiagnostic)
-      copyOp.emitOpError() << "pim.memcp_dh requires a device-local source and a host-backed target: source="
-                           << copyOp.getDeviceSource() << " host=" << sourceIsHost << " device=" << sourceIsDevice
-                           << ", target=" << copyOp.getHostTarget() << " host=" << targetIsHost
-                           << " device=" << targetIsDevice;
-    return failure();
-  }
-  return success();
-}
-
-static LogicalResult verifyLoweredPimCopy(pim::PimMemCopyOp copyOp,
-                                          const StaticValueKnowledge& knowledge,
-                                          bool emitDiagnostic) {
-  bool sourceIsHost = isHostBackedPimAddress(copyOp.getSource(), knowledge);
-  bool targetIsHost = isHostBackedPimAddress(copyOp.getTarget(), knowledge);
-  bool sourceIsDevice = isDeviceLocalPimAddress(copyOp.getSource(), knowledge);
-  bool targetIsDevice = isDeviceLocalPimAddress(copyOp.getTarget(), knowledge);
-  if (!sourceIsDevice || !targetIsDevice || sourceIsHost || targetIsHost) {
-    if (emitDiagnostic)
-      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();
+static LogicalResult verifyPimCopyEndpoints(Operation* copy,
+                                            Value source,
+                                            Value target,
+                                            ExpectedPimCopyDirection direction,
+                                            const StaticValueKnowledge& knowledge,
+                                            bool emitDiagnostic) {
+  struct ExpectedEndpoints {
+    bool sourceHost, sourceDevice, targetHost, targetDevice;
+    StringLiteral description;
+  };
+  static constexpr ExpectedEndpoints expected[] = {
+    {true, false, false, true, "a host-backed source and a device-local target"},
+    {false, true, true, false, "a device-local source and a host-backed target"},
+    {false, true, false, true, "device-local source and target operands"},
+  };
+  const auto& endpoints = expected[static_cast<unsigned>(direction)];
+  bool sourceHost = isHostBackedPimAddress(source, knowledge);
+  bool sourceDevice = isDeviceLocalPimAddress(source, knowledge);
+  bool targetHost = isHostBackedPimAddress(target, knowledge);
+  bool targetDevice = isDeviceLocalPimAddress(target, knowledge);
+  bool valid = sourceHost == endpoints.sourceHost && sourceDevice == endpoints.sourceDevice
+    && targetHost == endpoints.targetHost && targetDevice == endpoints.targetDevice;
+  if (!valid && emitDiagnostic)
+    copy->emitOpError() << "requires " << endpoints.description << ": source=" << source << " host=" << sourceHost
+                        << " device=" << sourceDevice << ", target=" << target << " host=" << targetHost
+                        << " device=" << targetDevice;
+  return success(valid);
 }
 
 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(
-  func::FuncOp funcOp, const bufferization::OneShotBufferizationOptions& options) {
-  bufferization::AnalysisState analysisState(options);
-  DominanceInfo dominance(funcOp);
-  size_t violationCount = 0;
+static LogicalResult verifyPimCoresNeedNoTensorCopies(
+  ModuleOp module, const bufferization::OneShotBufferizationOptions& baseOptions) {
+  static constexpr StringLiteral kExistingAlloc = "raptor.existing_core_alloc";
+  OwningOpRef<ModuleOp> clone = module.clone();
+  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) {
-    coreOp->walk([&](Operation* writeOp) {
-      for (OpOperand& write : writeOp->getOpOperands()) {
-        if (!isa<TensorType>(write.get().getType()) || !analysisState.bufferizesToMemoryWrite(write))
-          continue;
+  auto options = baseOptions;
+  options.bufferizeFunctionBoundaries = false;
+  options.opFilter.allowOperation([](Operation* op) {
+    return isa<pim::PimCoreOp, pim::PimCoreBatchOp>(op)
+        || op->getParentOfType<pim::PimCoreOp>()
+        || op->getParentOfType<pim::PimCoreBatchOp>();
+  });
 
-        SmallVector<Value, 8> worklist {write.get()};
-        llvm::SmallDenseSet<Value, 8> visited;
-        bool hasConflict = false;
-        Value conflictingAlias;
-        Operation* conflictingUse = nullptr;
-        while (!worklist.empty() && !hasConflict) {
-          Value alias = worklist.pop_back_val();
-          if (!visited.insert(alias).second)
-            continue;
+  bufferization::BufferizationState state;
+  if (failed(bufferization::insertTensorCopies(*clone, options, state))) {
+    module.emitError("official one-shot analysis failed while verifying PIM core copy freedom");
+    return failure();
+  }
 
-          for (OpOperand& use : alias.getUses()) {
-            bool usePrecedesWrite = false;
-            for (Operation* ancestor = use.getOwner(); ancestor; ancestor = ancestor->getParentOp())
-              if (ancestor == writeOp || dominance.properlyDominates(ancestor, writeOp)) {
-                usePrecedesWrite = true;
-                break;
-              }
-            if (usePrecedesWrite || analysisState.insideMutuallyExclusiveRegions(use.getOwner(), writeOp))
-              continue;
-            hasConflict = true;
-            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;
-      }
+  CappedDiagnosticReporter diagnostics;
+  clone->walk([&](bufferization::AllocTensorOp alloc) {
+    if (alloc->hasAttr(kExistingAlloc)
+        || (!alloc->getParentOfType<pim::PimCoreOp>()
+            && !alloc->getParentOfType<pim::PimCoreBatchOp>()))
+      return;
+    Operation* requiredBy = alloc->getUsers().empty()
+      ? alloc.getOperation() : *alloc->getUsers().begin();
+    diagnostics.report(requiredBy, [](Operation* op) {
+      op->emitOpError("official one-shot bufferization requires a tensor copy inside a PIM core");
     });
-  };
-
-  funcOp.walk([&](pim::PimCoreOp coreOp) { verifyCore(coreOp); });
-  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);
+  });
+  diagnostics.emitSuppressedSummary(module, "required PIM core tensor copies");
+  return success(!diagnostics.hasFailure());
 }
 
 } // 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(
-        moduleOp, hostOptions, state))) {
-    moduleOp.emitError("Failed to bufferize PIM and Spatial ops");
-    signalPassFailure();
-    return;
-  }
-  if (failed(bufferization::bufferizeModuleOp(
-        moduleOp, options, state))) {
+  if (failed(bufferization::insertTensorCopies(moduleOp, hostOptions, state))
+      || 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;
-  moduleOp.walk([&](Operation* op) {
-    if (isa<pim::PimMemCopyOp, pim::PimMemCopyHostToDevOp, pim::PimMemCopyDevToHostOp>(op))
-      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) {
+  GreedyRewriteConfig contiguityConfig;
+  contiguityConfig.enableFolding(false);
+  if (failed(applyPatternsGreedily(moduleOp, std::move(contiguityPatterns), contiguityConfig))) {
     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)
-      boundary.instructions.push_back(std::move(update));
+    appendLocalUpdates(boundary, localUpdates, exchange->targetLaneCount);
     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
-        ? getOrCreateIndexConstant(rewriter, batch.getOperation(), 0)
-        : arith::SubIOp::create(
-            builder, bodyLoc, originalLane,
-            getOrCreateIndexConstant(
-              rewriter, batch.getOperation(), first.laneStart));
+      Value localLane;
+      if (runLaneCount == 1)
+        localLane = getOrCreateIndexConstant(rewriter, batch.getOperation(), 0);
+      else if (first.laneStart == 0)
+        localLane = originalLane;
+      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;