better synchronization

better deadlock detection to also track wait/sync

src/PIM/Compiler/PimCompilerOptions.cpp

@@ -96,7 +96,7 @@ llvm::cl::opt<bool> pimEmitJson("pim-emit-json",
 
 llvm::cl::opt<bool> pimDetectCommunicationDeadlock(
   "pim-detect-communication-deadlock",
-  llvm::cl::desc("Expensively simulate the statically expanded Pim send/receive order at verification time and fail if a blocking communication deadlock is found"),
+  llvm::cl::desc("Expensively simulate statically expanded Pim SEND/RECV and exact-count SYNC/WAIT order at verification time and fail on a blocking deadlock"),
   llvm::cl::init(false),
   llvm::cl::cat(OnnxMlirOptions));
 

src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp

@@ -373,9 +373,6 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatScheduledCom
           rewriter, receiveOp->getLoc(), outputBuffer.getType(), zero,
           hostWaitLoad.getHostOffset(), outputBuffer, *hostBuffer, *sizeAttr)
           .getOutput();
-        PimSyncOp::create(
-          rewriter, receiveOp->getLoc(), hostWaitLoad.getSourceCoreId(),
-          hostWaitLoad.getAcknowledgementEventRegister());
       } else {
         received = PimReceiveOp::create(
           rewriter, receiveOp->getLoc(), outputBuffer.getType(), outputBuffer,

src/PIM/Conversion/SpatialToPim/Patterns/ChannelLowering.cpp

@@ -152,10 +152,6 @@ struct HostWaitLoadLowering : OpRewritePattern<spatial::SpatHostWaitLoadOp> {
       Value output = pim::PimMemCopyHostToDevOp::create(
         rewriter, op.getLoc(), outputBuffer.getType(), zero,
         op.getHostOffset(), outputBuffer, *hostBuffer, sizeAttr).getOutput();
-      auto sync = pim::PimSyncOp::create(
-        rewriter, op.getLoc(), op.getSourceCoreId(),
-        op.getAcknowledgementEventRegister());
-      copyRaptorDebugAttrs(op.getOperation(), sync.getOperation());
       return output;
     });
   }

src/PIM/Dialect/Pim/Passes/Transforms/Verification/VerificationPass.cpp

@@ -5,9 +5,11 @@
 #include "mlir/Pass/Pass.h"
 
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/raw_ostream.h"
 
+#include <functional>
 #include <map>
 #include <string>
 
@@ -278,7 +280,9 @@ static bool isHostAddressableValue(Value value, const StaticValueKnowledge& know
 
 enum class CommunicationEventKind {
   Send,
-  Receive
+  Receive,
+  Sync,
+  Wait
 };
 
 struct CommunicationEvent {
@@ -286,14 +290,25 @@ struct CommunicationEvent {
   int64_t coreId = 0;
   int64_t peerCoreId = 0;
   int64_t size = 0;
+  int64_t eventRegister = 0;
+  int64_t waitValue = 0;
   uint64_t ordinal = 0;
   Operation* op = nullptr;
 };
 
 using CommunicationEventVector = SmallVector<CommunicationEvent, 0>;
+using SynchronizationEventKey = std::pair<int64_t, int64_t>;
+using SynchronizationSourceCounts =
+  DenseMap<SynchronizationEventKey, DenseMap<int64_t, int64_t>>;
 
 static StringRef getCommunicationEventKindName(CommunicationEventKind kind) {
-  return kind == CommunicationEventKind::Send ? "send" : "receive";
+  switch (kind) {
+  case CommunicationEventKind::Send: return "send";
+  case CommunicationEventKind::Receive: return "receive";
+  case CommunicationEventKind::Sync: return "sync";
+  case CommunicationEventKind::Wait: return "wait";
+  }
+  llvm_unreachable("unknown communication event kind");
 }
 
 constexpr StringLiteral kRaptorMinChannelIdAttr = "raptor.min_channel_id";
@@ -356,9 +371,15 @@ static std::string formatCommunicationEvent(const CommunicationEvent& event) {
 
   std::string text;
   llvm::raw_string_ostream os(text);
-  os << "core " << event.coreId << " " << getCommunicationEventKindName(event.kind) << " "
-     << (event.kind == CommunicationEventKind::Send ? "to" : "from") << " " << event.peerCoreId << " size "
-     << event.size << "B ordinal " << event.ordinal;
+  os << "core " << event.coreId << " " << getCommunicationEventKindName(event.kind);
+  if (event.kind == CommunicationEventKind::Send || event.kind == CommunicationEventKind::Receive)
+    os << " " << (event.kind == CommunicationEventKind::Send ? "to" : "from") << " " << event.peerCoreId
+       << " size " << event.size << "B";
+  else if (event.kind == CommunicationEventKind::Sync)
+    os << " event " << event.eventRegister << " to " << event.peerCoreId;
+  else
+    os << " event " << event.eventRegister << " value " << event.waitValue;
+  os << " ordinal " << event.ordinal;
   if (minChannelId)
     os << " min_channel " << *minChannelId;
   if (commOrder)
@@ -383,6 +404,9 @@ static std::string formatCommunicationEvent(const CommunicationEvent& event) {
 }
 
 static bool areMatchedCommunicationEvents(const CommunicationEvent& lhs, const CommunicationEvent& rhs) {
+  if ((lhs.kind != CommunicationEventKind::Send && lhs.kind != CommunicationEventKind::Receive)
+      || (rhs.kind != CommunicationEventKind::Send && rhs.kind != CommunicationEventKind::Receive))
+    return false;
   if (lhs.coreId != rhs.peerCoreId || lhs.peerCoreId != rhs.coreId || lhs.size != rhs.size)
     return false;
 
@@ -402,6 +426,27 @@ static void printCounterpartProbe(llvm::raw_ostream& os,
                                   const DenseMap<int64_t, CommunicationEventVector>& coreEvents,
                                   const DenseMap<int64_t, size_t>& programCounters,
                                   const CommunicationEvent& blockedEvent) {
+  if (blockedEvent.kind == CommunicationEventKind::Wait) {
+    os << "  SYNC probes for " << formatCommunicationEvent(blockedEvent) << "\n";
+    bool found = false;
+    for (const auto& [sourceCore, events] : coreEvents) {
+      size_t begin = programCounters.lookup(sourceCore);
+      for (size_t index = begin; index < events.size(); ++index) {
+        const CommunicationEvent& candidate = events[index];
+        if (candidate.kind != CommunicationEventKind::Sync
+            || candidate.peerCoreId != blockedEvent.coreId
+            || candidate.eventRegister != blockedEvent.eventRegister)
+          continue;
+        os << "    core " << sourceCore << " next matching SYNC at ordinal "
+           << index << " (distance +" << index - begin << ")\n";
+        found = true;
+        break;
+      }
+    }
+    if (!found)
+      os << "    no remaining matching SYNC exists\n";
+    return;
+  }
   auto peerEventsIt = coreEvents.find(blockedEvent.peerCoreId);
   if (peerEventsIt == coreEvents.end()) {
     os << "  no local stream was collected for peer core " << blockedEvent.peerCoreId << "\n";
@@ -456,7 +501,7 @@ static void printCounterpartProbe(llvm::raw_ostream& os,
 
 static CommunicationEvent makeCommunicationEvent(
   CommunicationEventKind kind, int64_t coreId, int64_t peerCoreId, int64_t size, uint64_t ordinal, Operation* op) {
-  return CommunicationEvent {kind, coreId, peerCoreId, size, ordinal, op};
+  return CommunicationEvent {kind, coreId, peerCoreId, size, 0, 0, ordinal, op};
 }
 
 static LogicalResult appendCoreCommunicationEvents(Block& block,
@@ -504,6 +549,39 @@ static LogicalResult appendCoreCommunicationEvents(Block& block,
         return success();
       }
 
+      if (auto syncOp = dyn_cast<pim::PimSyncOp>(&op)) {
+        auto targetCoreId = resolveIndexValue(syncOp.getTargetCoreId(), knowledge);
+        auto eventRegister = resolveIndexValue(syncOp.getEventRegister(), knowledge);
+        if (failed(targetCoreId) || failed(eventRegister)) {
+          diagnostics.report(&op, [](Operation* illegalOp) {
+            illegalOp->emitOpError(
+              "cannot statically resolve SYNC operands for Pim communication deadlock check");
+          });
+          return failure();
+        }
+        events.push_back(CommunicationEvent {
+          CommunicationEventKind::Sync, coreId, *targetCoreId, 0,
+          *eventRegister, 0, static_cast<uint64_t>(events.size()), &op});
+        return success();
+      }
+
+      if (auto waitOp = dyn_cast<pim::PimWaitOp>(&op)) {
+        auto eventRegister = resolveIndexValue(waitOp.getEventRegister(), knowledge);
+        auto waitValue = resolveIndexValue(waitOp.getWaitValue(), knowledge);
+        if (failed(eventRegister) || failed(waitValue)) {
+          diagnostics.report(&op, [](Operation* illegalOp) {
+            illegalOp->emitOpError(
+              "cannot statically resolve WAIT operands for Pim communication deadlock check");
+          });
+          return failure();
+        }
+        events.push_back(CommunicationEvent {
+          CommunicationEventKind::Wait, coreId, coreId, 0,
+          *eventRegister, *waitValue,
+          static_cast<uint64_t>(events.size()), &op});
+        return success();
+      }
+
       return success();
     });
 }
@@ -530,7 +608,7 @@ static void printCommunicationWindow(llvm::raw_ostream& os,
 static void printCommunicationDeadlockReport(const DenseMap<int64_t, CommunicationEventVector>& coreEvents,
                                              const DenseMap<int64_t, size_t>& programCounters,
                                              ArrayRef<int64_t> cycle) {
-  llvm::errs() << "\n=== Pim static communication deadlock report ===\n";
+  llvm::errs() << "\n=== Pim static communication/synchronization deadlock report ===\n";
   llvm::errs() << "wait cycle:";
   for (int64_t coreId : cycle)
     llvm::errs() << " " << coreId;
@@ -565,7 +643,7 @@ static void printCommunicationDeadlockReport(const DenseMap<int64_t, Communicati
       continue;
     printCommunicationWindow(llvm::errs(), coreEvents, coreId, pcIt->second);
   }
-  llvm::errs() << "=== end Pim static communication deadlock report ===\n\n";
+  llvm::errs() << "=== end Pim static communication/synchronization deadlock report ===\n\n";
 }
 
 static void emitCommunicationDeadlockCycle(ModuleOp moduleOp,
@@ -576,8 +654,8 @@ static void emitCommunicationDeadlockCycle(ModuleOp moduleOp,
 
   auto diagnostic =
     moduleOp.emitError()
-    << "Pim communication deadlock check found a blocking send/receive cycle while statically simulating the "
-       "expanded per-core communication streams; see the Pim static communication deadlock report above";
+    << "Pim communication deadlock check found a blocking SEND/RECV/WAIT cycle while statically simulating the "
+       "expanded per-core communication streams; see the static deadlock report above";
 
   for (int64_t coreId : cycle) {
     auto eventsIt = coreEvents.find(coreId);
@@ -596,46 +674,82 @@ static void emitCommunicationDeadlockCycle(ModuleOp moduleOp,
 
 static FailureOr<SmallVector<int64_t>>
 findCommunicationWaitCycle(const DenseMap<int64_t, CommunicationEventVector>& coreEvents,
-                           const DenseMap<int64_t, size_t>& programCounters) {
-  for (const auto& [startCoreId, events] : coreEvents) {
-    auto startPcIt = programCounters.find(startCoreId);
-    if (startPcIt == programCounters.end() || startPcIt->second >= events.size())
+                           const DenseMap<int64_t, size_t>& programCounters,
+                           const DenseSet<int64_t>& repeatingCores,
+                           const SynchronizationSourceCounts& sourceCounts) {
+  DenseMap<int64_t, SmallVector<int64_t>> dependencies;
+  for (const auto& [coreId, events] : coreEvents) {
+    size_t pc = programCounters.lookup(coreId);
+    if (pc >= events.size())
       continue;
-
-    DenseMap<int64_t, size_t> positionInPath;
-    SmallVector<int64_t, 8> path;
-    int64_t currentCoreId = startCoreId;
-    while (true) {
-      auto eventsIt = coreEvents.find(currentCoreId);
-      auto pcIt = programCounters.find(currentCoreId);
-      if (eventsIt == coreEvents.end() || pcIt == programCounters.end() || pcIt->second >= eventsIt->second.size())
-        break;
-
-      auto positionIt = positionInPath.find(currentCoreId);
-      if (positionIt != positionInPath.end()) {
-        SmallVector<int64_t> cycle;
-        for (size_t index = positionIt->second; index < path.size(); ++index)
-          cycle.push_back(path[index]);
-        return cycle;
-      }
-
-      positionInPath[currentCoreId] = path.size();
-      path.push_back(currentCoreId);
-      currentCoreId = eventsIt->second[pcIt->second].peerCoreId;
+    const CommunicationEvent& event = events[pc];
+    if (event.kind == CommunicationEventKind::Send
+        || event.kind == CommunicationEventKind::Receive) {
+      dependencies[coreId].push_back(event.peerCoreId);
+      continue;
+    }
+    if (event.kind != CommunicationEventKind::Wait)
+      continue;
+    int64_t blockedCoreId = coreId;
+    SynchronizationEventKey eventKey {coreId, event.eventRegister};
+    auto contributions = sourceCounts.find(eventKey);
+    for (const auto& [sourceCore, sourceEvents] : coreEvents) {
+      size_t sourcePc = programCounters.lookup(sourceCore);
+      auto matches = [&](const CommunicationEvent& candidate) {
+        return candidate.kind == CommunicationEventKind::Sync
+            && candidate.peerCoreId == blockedCoreId
+            && candidate.eventRegister == event.eventRegister;
+      };
+      int64_t signalsPerPhase = llvm::count_if(sourceEvents, matches);
+      if (repeatingCores.contains(sourceCore))
+        signalsPerPhase /= 2;
+      int64_t contributed = contributions == sourceCounts.end()
+        ? 0 : contributions->second.lookup(sourceCore);
+      if (contributed >= signalsPerPhase)
+        continue;
+      bool canSignal = llvm::any_of(
+        llvm::drop_begin(sourceEvents, sourcePc), matches);
+      if (canSignal)
+        dependencies[coreId].push_back(sourceCore);
     }
   }
 
+  DenseSet<int64_t> visited;
+  DenseMap<int64_t, size_t> positionInPath;
+  SmallVector<int64_t, 8> path;
+  std::function<std::optional<SmallVector<int64_t>>(int64_t)> visit =
+    [&](int64_t coreId) -> std::optional<SmallVector<int64_t>> {
+      auto position = positionInPath.find(coreId);
+      if (position != positionInPath.end())
+        return SmallVector<int64_t>(
+          path.begin() + position->second, path.end());
+      if (!visited.insert(coreId).second)
+        return std::nullopt;
+      positionInPath[coreId] = path.size();
+      path.push_back(coreId);
+      for (int64_t target : dependencies.lookup(coreId))
+        if (auto cycle = visit(target))
+          return cycle;
+      path.pop_back();
+      positionInPath.erase(coreId);
+      return std::nullopt;
+    };
+  for (const auto& [coreId, unused] : dependencies)
+    if (auto cycle = visit(coreId))
+      return *cycle;
   return failure();
 }
 
 static LogicalResult verifyNoStaticCommunicationDeadlock(ModuleOp moduleOp,
                                                          pim::CappedDiagnosticReporter& diagnostics) {
   DenseMap<int64_t, CommunicationEventVector> coreEvents;
+  DenseSet<int64_t> repeatingCores;
   bool hasFailure = false;
 
   for (func::FuncOp funcOp : moduleOp.getOps<func::FuncOp>()) {
     if (funcOp.isExternal())
       continue;
+    bool repeating = funcOp->hasAttr("pim.pipeline_host_buffer_bytes");
 
     for (Operation& op : funcOp.getBody().front().getOperations()) {
       if (auto coreOp = dyn_cast<pim::PimCoreOp>(&op)) {
@@ -648,6 +762,8 @@ static LogicalResult verifyNoStaticCommunicationDeadlock(ModuleOp moduleOp,
                                                  coreEvents[coreId],
                                                  diagnostics)))
           hasFailure = true;
+        if (repeating)
+          repeatingCores.insert(coreId);
         continue;
       }
 
@@ -668,8 +784,11 @@ static LogicalResult verifyNoStaticCommunicationDeadlock(ModuleOp moduleOp,
                                                      coreId,
                                                      laneKnowledge,
                                                      coreEvents[coreId],
-                                                     diagnostics)))
+                                                     diagnostics))) {
               hasFailure = true;
+            } else if (repeating) {
+              repeatingCores.insert(coreId);
+            }
         }
       }
     }
@@ -678,10 +797,20 @@ static LogicalResult verifyNoStaticCommunicationDeadlock(ModuleOp moduleOp,
   if (hasFailure)
     return failure();
 
+  for (int64_t coreId : repeatingCores) {
+    CommunicationEventVector iteration = coreEvents[coreId];
+    for (CommunicationEvent event : iteration) {
+      event.ordinal = coreEvents[coreId].size();
+      coreEvents[coreId].push_back(event);
+    }
+  }
+
   DenseMap<int64_t, size_t> programCounters;
   for (const auto& [coreId, events] : coreEvents)
     programCounters[coreId] = 0;
 
+  DenseMap<SynchronizationEventKey, int64_t> eventCounts;
+  SynchronizationSourceCounts sourceCounts;
   while (true) {
     bool madeProgress = false;
     for (const auto& [coreId, events] : coreEvents) {
@@ -690,6 +819,35 @@ static LogicalResult verifyNoStaticCommunicationDeadlock(ModuleOp moduleOp,
         continue;
 
       const CommunicationEvent& event = events[pc];
+      if (event.kind == CommunicationEventKind::Sync) {
+        SynchronizationEventKey eventKey {
+          event.peerCoreId, event.eventRegister};
+        ++eventCounts[eventKey];
+        ++sourceCounts[eventKey][coreId];
+        ++programCounters[coreId];
+        madeProgress = true;
+        continue;
+      }
+      if (event.kind == CommunicationEventKind::Wait) {
+        int64_t observed = eventCounts.lookup(
+          {coreId, event.eventRegister});
+        if (observed == event.waitValue) {
+          SynchronizationEventKey eventKey {coreId, event.eventRegister};
+          eventCounts[eventKey] = 0;
+          sourceCounts.erase(eventKey);
+          ++programCounters[coreId];
+          madeProgress = true;
+        } else if (observed > event.waitValue) {
+          auto diagnostic = event.op->emitOpError()
+            << "Pim synchronization deadlock check found exact-count WAIT overshoot on core "
+            << coreId << " event " << event.eventRegister << ": expected "
+            << event.waitValue << ", observed " << observed;
+          diagnostic.attachNote()
+            << "WAIT requires exact equality and this event cannot decrease without a successful WAIT";
+          return failure();
+        }
+        continue;
+      }
       auto peerEventsIt = coreEvents.find(event.peerCoreId);
       if (peerEventsIt == coreEvents.end())
         continue;
@@ -720,7 +878,8 @@ static LogicalResult verifyNoStaticCommunicationDeadlock(ModuleOp moduleOp,
     if (allDone)
       return success();
 
-    auto cycle = findCommunicationWaitCycle(coreEvents, programCounters);
+    auto cycle = findCommunicationWaitCycle(
+      coreEvents, programCounters, repeatingCores, sourceCounts);
     if (succeeded(cycle)) {
       emitCommunicationDeadlockCycle(moduleOp, coreEvents, programCounters, *cycle);
       return failure();
@@ -729,7 +888,7 @@ static LogicalResult verifyNoStaticCommunicationDeadlock(ModuleOp moduleOp,
     auto diagnostic =
       moduleOp.emitError()
       << "Pim communication deadlock check stalled without finding a closed wait cycle; this usually means a "
-         "send/receive peer is missing or ordered after a finished core";
+         "SEND/RECV peer or exact-count SYNC signal is missing or ordered after a finished core";
     for (const auto& [coreId, events] : coreEvents) {
       size_t pc = programCounters[coreId];
       if (pc >= events.size())

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

@@ -2,6 +2,8 @@
 #include "DeferredCommunicationScheduling.hpp"
 #include "DeferredTransferPlanning.hpp"
 
+#include "llvm/ADT/DenseSet.h"
+
 namespace onnx_mlir::spatial {
 using namespace mlir;
 namespace {
@@ -256,7 +258,7 @@ static unsigned getBarrierRoundCount(size_t coreCount) {
 
 static LogicalResult assignPipelineSynchronization(
     DeferredTransferPlan &transfers,
-    ArrayRef<BoundaryProgram> boundaries,
+    MutableArrayRef<BoundaryProgram> boundaries,
     size_t synchronizationRegisterCount) {
   bool pipelined = false;
   for (ScheduledInfo &scheduled : transfers.scheduled) {
@@ -284,14 +286,13 @@ static LogicalResult assignPipelineSynchronization(
   });
 
   DenseMap<int64_t, SmallVector<HostTransferRef>> incomingByCore;
+  DenseMap<int64_t, llvm::SmallSetVector<int64_t, 4>> readersByWriter;
   DenseMap<ExternalTransferFamily *, SmallVector<int64_t>> eventRegisters;
   DenseMap<ExternalTransferFamily *, SmallVector<int64_t>> waitValues;
-  DenseMap<ExternalTransferFamily *, SmallVector<int64_t>> acknowledgementRegisters;
   auto initialize = [&](ExternalTransferFamily &family) {
     size_t count = family.targetCores.size();
     eventRegisters.try_emplace(&family, count, 0);
     waitValues.try_emplace(&family, count, 0);
-    acknowledgementRegisters.try_emplace(&family, count, 0);
   };
   for (const BoundaryProgram &boundary : boundaries)
     for (const BoundaryInstruction &instruction : boundary.instructions) {
@@ -307,32 +308,117 @@ static LogicalResult assignPipelineSynchronization(
           int64_t source = family.sourceCores.valueAt(index);
           int64_t target = family.targetCores.valueAt(index);
           incomingByCore[target].push_back({&family, index});
-          ++transfers.hostAcknowledgementCounts[source];
+          readersByWriter[source].insert(target);
         }
       }
     }
 
+  DenseMap<int64_t, SmallVector<SmallVector<int64_t>>>
+    overwriteBoundaryReadersByWriter;
+  DenseSet<std::pair<int64_t, int64_t>> plannedRelations;
+  for (const BoundaryProgram &boundary : boundaries) {
+    DenseMap<int64_t, llvm::SmallSetVector<int64_t, 4>> readersAtBoundary;
+    for (const BoundaryInstruction &instruction : boundary.instructions) {
+      auto *send = std::get_if<EmitSendRun>(&instruction);
+      if (!send || send->slices.empty()
+          || !send->slices.front().family->hostRouted)
+        continue;
+      for (const ScheduledTransferSlice &slice : send->slices)
+        for (size_t offset = 0; offset < slice.transferCount; ++offset) {
+          size_t index = slice.familyOffset + offset;
+          int64_t writer = slice.family->sourceCores.valueAt(index);
+          int64_t reader = slice.family->targetCores.valueAt(index);
+          if (plannedRelations.insert({writer, reader}).second)
+            readersAtBoundary[writer].insert(reader);
+        }
+    }
+    for (const auto &[writer, readers] : readersAtBoundary)
+      overwriteBoundaryReadersByWriter[writer].emplace_back(
+        readers.begin(), readers.end());
+  }
+
   unsigned barrierRounds = getBarrierRoundCount(
     transfers.stageZeroCores.size());
-  bool stageZeroNeedsAcknowledgements = llvm::any_of(
-    transfers.stageZeroCores, [&](int64_t core) {
-      return transfers.hostAcknowledgementCounts.contains(core);
-    });
-  for (auto &[target, incoming] : incomingByCore) {
-    bool needsAcknowledgementRegister =
-      transfers.hostAcknowledgementCounts.contains(target);
-    bool stageZero = llvm::is_contained(transfers.stageZeroCores, target);
-    size_t reserved = stageZero
-      ? barrierRounds + (stageZeroNeedsAcknowledgements ? 1 : 0)
-      : 1 + (needsAcknowledgementRegister ? 1 : 0);
-    if (reserved >= synchronizationRegisterCount) {
-      incoming.front().family->requirement->exchange->deferred.emitOpError(
-        "pipeline synchronization leaves no event register for incoming host transfers");
+  auto reservedRegisterCount = [&](int64_t core) -> size_t {
+    if (llvm::is_contained(transfers.stageZeroCores, core)) {
+      bool releasesDownstream = !transfers.downstreamCores.empty()
+        && core == transfers.stageZeroCores.front();
+      return barrierRounds + releasesDownstream;
+    }
+    auto downstream = llvm::find(transfers.downstreamCores, core);
+    if (downstream == transfers.downstreamCores.end())
+      return 0;
+    size_t rank = downstream - transfers.downstreamCores.begin();
+    bool releasesChildren = 2 * rank + 1 < transfers.downstreamCores.size();
+    return 1 + releasesChildren;
+  };
+  for (int64_t core : transfers.stageZeroCores)
+    if (reservedRegisterCount(core) > synchronizationRegisterCount)
+      return transfers.scheduled.front().op->emitOpError(
+        "pipeline stage-zero synchronization requires more event registers than the target provides");
+  for (int64_t core : transfers.downstreamCores)
+    if (reservedRegisterCount(core) > synchronizationRegisterCount)
+      return transfers.scheduled.front().op->emitOpError(
+        "pipeline downstream restart synchronization requires more event registers than the target provides");
+  SmallVector<int64_t> synchronizationCores;
+  for (const auto &[core, incoming] : incomingByCore)
+    synchronizationCores.push_back(core);
+  for (const auto &[core, readers] : readersByWriter)
+    if (!llvm::is_contained(synchronizationCores, core))
+      synchronizationCores.push_back(core);
+  llvm::sort(synchronizationCores);
+
+  DenseMap<int64_t, SmallVector<unsigned>> freeRegisters;
+  DenseMap<int64_t, SmallVector<unsigned>> freeWaitCounts;
+  DenseMap<int64_t, DenseMap<int64_t, unsigned>> freeGroups;
+  for (int64_t core : synchronizationCores) {
+    auto incomingIt = incomingByCore.find(core);
+    ArrayRef<HostTransferRef> incoming = incomingIt == incomingByCore.end()
+      ? ArrayRef<HostTransferRef>() : ArrayRef(incomingIt->second);
+    size_t readerCount = readersByWriter.lookup(core).size();
+    bool needsFreeRegister = readerCount != 0;
+    size_t reserved = reservedRegisterCount(core);
+    if (reserved + needsFreeRegister > synchronizationRegisterCount
+        || (!incoming.empty()
+            && reserved + needsFreeRegister == synchronizationRegisterCount)) {
+      transfers.scheduled.front().op->emitOpError(
+        "pipeline synchronization leaves no event register for host communication");
       return failure();
     }
-    size_t groupCount = std::min(
-      incoming.size(), synchronizationRegisterCount - reserved);
-    // One wait consumes a complete consecutive group of producer signals.
+    size_t available = synchronizationRegisterCount - reserved;
+    size_t readyRegisterCount = std::min(
+      incoming.size(), available - needsFreeRegister);
+    auto overwriteBoundaries = overwriteBoundaryReadersByWriter.find(core);
+    size_t overwriteBoundaryCount = overwriteBoundaries
+      == overwriteBoundaryReadersByWriter.end()
+      ? 0 : overwriteBoundaries->second.size();
+    size_t freeGroupCount = std::min(
+      overwriteBoundaryCount, available - readyRegisterCount);
+    if (needsFreeRegister) {
+      size_t overwriteReaderCount = 0;
+      if (overwriteBoundaries != overwriteBoundaryReadersByWriter.end())
+        for (ArrayRef<int64_t> readers : overwriteBoundaries->second)
+          overwriteReaderCount += readers.size();
+      if (overwriteReaderCount != readerCount)
+        return transfers.scheduled.front().op->emitOpError(
+          "host reuse readers do not match writer overwrite relations");
+      freeRegisters[core].reserve(freeGroupCount);
+      freeWaitCounts[core].assign(freeGroupCount, 0);
+      for (unsigned group = 0; group < freeGroupCount; ++group)
+        freeRegisters[core].push_back(readyRegisterCount + group);
+      for (auto [ordinal, readers] :
+           llvm::enumerate(overwriteBoundaries->second)) {
+        unsigned group = ordinal * freeGroupCount / overwriteBoundaryCount;
+        for (int64_t reader : readers) {
+          freeGroups[core][reader] = group;
+          ++freeWaitCounts[core][group];
+        }
+      }
+    }
+    if (incoming.empty())
+      continue;
+
+    size_t groupCount = readyRegisterCount;
     SmallVector<size_t> groupSizes(groupCount);
     for (size_t ordinal = 0; ordinal < incoming.size(); ++ordinal)
       ++groupSizes[ordinal * groupCount / incoming.size()];
@@ -341,8 +427,6 @@ static LogicalResult assignPipelineSynchronization(
       size_t group = ordinal * groupCount / incoming.size();
       HostTransferRef transfer = incoming[ordinal];
       eventRegisters[transfer.family][transfer.index] = group;
-      acknowledgementRegisters[transfer.family][transfer.index] =
-        synchronizationRegisterCount - 1;
       if (first[group]) {
         waitValues[transfer.family][transfer.index] = groupSizes[group];
         first[group] = false;
@@ -352,27 +436,70 @@ static LogicalResult assignPipelineSynchronization(
   for (auto &[family, values] : eventRegisters) {
     family->eventRegisters = StaticIntSequence::fromValues(values);
     family->waitValues = StaticIntSequence::fromValues(waitValues[family]);
-    family->acknowledgementEventRegisters =
-      StaticIntSequence::fromValues(acknowledgementRegisters[family]);
   }
 
-  if (!transfers.stageZeroCores.empty()) {
-    size_t reserved = barrierRounds
-                    + (stageZeroNeedsAcknowledgements ? 1 : 0);
-    if (reserved > synchronizationRegisterCount)
-      return transfers.scheduled.front().op->emitOpError(
-        "pipeline stage-zero barrier requires more synchronization registers than the target provides");
-  }
-  if (!transfers.downstreamCores.empty()) {
-    bool needsAcknowledgements = llvm::any_of(
-      transfers.downstreamCores, [&](int64_t core) {
-        return transfers.hostAcknowledgementCounts.contains(core);
-      });
-    if (1 + (needsAcknowledgements ? 1 : 0)
-        > synchronizationRegisterCount)
-      return transfers.scheduled.front().op->emitOpError(
-        "pipeline stage-zero release requires more synchronization registers than the target provides");
+  SmallVector<int64_t> writers;
+  for (const auto &[writer, readers] : readersByWriter)
+    writers.push_back(writer);
+  llvm::sort(writers);
+  for (int64_t writer : writers)
+    for (int64_t reader : readersByWriter[writer]) {
+      unsigned group = freeGroups[writer].lookup(reader);
+      transfers.hostReleaseSignals[reader].push_back(
+        {writer, freeRegisters[writer][group]});
+    }
+
+  DenseSet<std::pair<int64_t, unsigned>> pendingGroups;
+  for (int64_t writer : writers)
+    for (unsigned group = 0; group < freeRegisters[writer].size(); ++group)
+      pendingGroups.insert({writer, group});
+  for (BoundaryProgram &boundary : boundaries) {
+    SmallVector<BoundaryInstruction, 0> instructions;
+    for (BoundaryInstruction &instruction : boundary.instructions) {
+      auto *send = std::get_if<EmitSendRun>(&instruction);
+      SmallVector<std::pair<unsigned, unsigned>> waits;
+      if (send && !send->slices.empty()
+          && send->slices.front().family->hostRouted) {
+        for (const ScheduledTransferSlice &slice : send->slices) {
+          ExternalTransferFamily &family = *slice.family;
+          unsigned lane = family.requirement->producer->scheduledLane;
+          if (lane >= boundary.key.first->cores.size())
+            return family.requirement->exchange->deferred.emitOpError(
+              "host reuse wait references an invalid writer lane"), failure();
+          for (size_t offset = 0; offset < slice.transferCount; ++offset) {
+            int64_t writer = family.sourceCores.valueAt(
+              slice.familyOffset + offset);
+            if (boundary.key.first->cores[lane] != writer)
+              return family.requirement->exchange->deferred.emitOpError(
+                "host reuse wait writer does not match its scheduled lane"),
+                failure();
+            int64_t reader = family.targetCores.valueAt(
+              slice.familyOffset + offset);
+            unsigned group = freeGroups[writer].lookup(reader);
+            if (!pendingGroups.erase({writer, group}))
+              continue;
+            waits.push_back({lane, group});
+          }
+        }
+      }
+      for (auto [lane, group] : waits) {
+        int64_t writer = boundary.key.first->cores[lane];
+        SmallVector<int64_t> registers(boundary.key.first->cores.size());
+        SmallVector<int64_t> counts(boundary.key.first->cores.size());
+        registers[lane] = freeRegisters[writer][group];
+        counts[lane] = freeWaitCounts[writer][group];
+        instructions.push_back(EmitHostReuseWait {
+          LaneSet::range(lane, lane + 1),
+          StaticIntSequence::fromValues(registers),
+          StaticIntSequence::fromValues(counts)});
+      }
+      instructions.push_back(std::move(instruction));
+    }
+    boundary.instructions = std::move(instructions);
   }
+  if (!pendingGroups.empty())
+    return transfers.scheduled.front().op->emitOpError(
+      "host reuse permission has no writer overwrite boundary");
   return success();
 }
 

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

@@ -34,12 +34,17 @@ struct EmitReceiveAssemblyRun {
   llvm::SmallVector<LaneSet> entryLanes;
   LaneSet lanes;
 };
+struct EmitHostReuseWait {
+  LaneSet lanes;
+  StaticIntSequence eventRegisters = StaticIntSequence::uniform(0, 1);
+  StaticIntSequence waitValues = StaticIntSequence::uniform(0, 1);
+};
 struct ProduceDeferredResult {
   DeferredExchangePlan* exchange = nullptr;
 };
 
 using BoundaryInstruction =
-  std::variant<EmitSendRun, EmitLocalCollectionRun,
+  std::variant<EmitSendRun, EmitHostReuseWait, EmitLocalCollectionRun,
                EmitLocalCollectionLoopRun, EmitReceiveAssemblyRun,
                ProduceDeferredResult>;
 struct BoundaryProgram {

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

@@ -10,6 +10,7 @@
 #include "src/Accelerators/PIM/Common/IR/StaticIntSequence.hpp"
 #include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "llvm/ADT/DenseSet.h"
 #include <array>
 namespace onnx_mlir::spatial {
 using namespace mlir;
@@ -23,7 +24,6 @@ struct LogicalTransferMetadataView {
   StaticIntSequenceChain hostOffsets;
   StaticIntSequenceChain eventRegisters;
   StaticIntSequenceChain waitValues;
-  StaticIntSequenceChain acknowledgementEventRegisters;
   StaticIntSequenceChain targetLanes;
   StaticIntSequenceChain localOffsets;
   SmallVector<StaticIntSequenceChain> projectionOffsets;
@@ -78,6 +78,29 @@ static FailureOr<Value> emitLaneCondition(const LaneSet &lanes, Value lane, unsi
   Value selected = active->emitLookup(context.constants.getIndex(0), lane, anchor, context.constants, context.rewriter, loc);
   return arith::CmpIOp::create(context.rewriter, loc, arith::CmpIPredicate::ne, selected, context.constants.getIndex(0)).getResult();
 }
+template <typename Emit>
+static LogicalResult emitForLanes(
+    const LaneSet &active, Value lane, unsigned laneCount, Operation *anchor,
+    DeferredEmissionContext &context, Location loc, Emit emit) {
+  if (active.empty())
+    return success();
+  if (!lane) {
+    if (active.contains(0))
+      emit();
+    return success();
+  }
+  auto condition = emitLaneCondition(
+    active, lane, laneCount, anchor, context, loc);
+  if (failed(condition))
+    return failure();
+  auto conditional = scf::IfOp::create(
+    context.rewriter, loc, TypeRange {}, *condition, false);
+  OpBuilder::InsertionGuard guard(context.rewriter);
+  context.rewriter.setInsertionPoint(
+    conditional.getThenRegion().front().getTerminator());
+  emit();
+  return success();
+}
 static void appendMetadata(const ScheduledTransferSlice &slice, LogicalTransferMetadataView &metadata) {
   ExternalTransferFamily &family = *slice.family;
   LaneInterval familyLanes = family.targetLanes.intervals().front();
@@ -95,8 +118,6 @@ static void appendMetadata(const ScheduledTransferSlice &slice, LogicalTransferM
     metadata.eventRegisters.append(
       family.eventRegisters, familyIndex, count);
     metadata.waitValues.append(family.waitValues, familyIndex, count);
-    metadata.acknowledgementEventRegisters.append(
-      family.acknowledgementEventRegisters, familyIndex, count);
   }
   metadata.targetLanes.append(StaticIntSequence::affine(targetLane, 1, count));
   if (family.requirement->producerLocalOffsets)
@@ -303,20 +324,15 @@ static FailureOr<Value> emitReceiveValue(ArrayRef<ScheduledTransferSlice> slices
   std::optional<StaticIntGrid> hostOffsets;
   std::optional<StaticIntGrid> eventRegisters;
   std::optional<StaticIntGrid> waitValues;
-  std::optional<StaticIntGrid> acknowledgementEventRegisters;
   if (slices.front().family->hostRouted) {
     auto offsets = buildGrid(metadata.hostOffsets);
     auto events = buildGrid(metadata.eventRegisters);
     auto waits = buildGrid(metadata.waitValues);
-    auto acknowledgements = buildGrid(
-      metadata.acknowledgementEventRegisters);
-    if (failed(offsets) || failed(events) || failed(waits)
-        || failed(acknowledgements))
+    if (failed(offsets) || failed(events) || failed(waits))
       return failure();
     hostOffsets = std::move(*offsets);
     eventRegisters = std::move(*events);
     waitValues = std::move(*waits);
-    acknowledgementEventRegisters = std::move(*acknowledgements);
   }
   Value position = lane ? lane : context.constants.getIndex(0);
   Value row = context.constants.getIndex(0);
@@ -335,8 +351,6 @@ static FailureOr<Value> emitReceiveValue(ArrayRef<ScheduledTransferSlice> slices
       eventRegisters->emitLookup(
         row, position, anchor, context.constants, context.rewriter, anchor->getLoc()),
       waitValues->emitLookup(
-        row, position, anchor, context.constants, context.rewriter, anchor->getLoc()),
-      acknowledgementEventRegisters->emitLookup(
         row, position, anchor, context.constants, context.rewriter, anchor->getLoc()));
     receive = op;
     output = op.getOutput();
@@ -406,7 +420,6 @@ static FailureOr<Value> emitReceiveAssembly(const EmitReceiveAssemblyRun &run, V
   std::optional<StaticIntGrid> hostOffsets;
   std::optional<StaticIntGrid> eventRegisters;
   std::optional<StaticIntGrid> waitValues;
-  std::optional<StaticIntGrid> acknowledgementEventRegisters;
   bool hostRouted = run.slices.front().family->hostRouted;
   auto metadataByEntry = buildRectangularReceiveMetadata(run, laneCount);
   if (succeeded(metadataByEntry)) {
@@ -424,15 +437,11 @@ static FailureOr<Value> emitReceiveAssembly(const EmitReceiveAssemblyRun &run, V
         &LogicalTransferMetadataView::eventRegisters);
       auto waits = buildRows(
         &LogicalTransferMetadataView::waitValues);
-      auto acknowledgements = buildRows(
-        &LogicalTransferMetadataView::acknowledgementEventRegisters);
-      if (failed(offsets) || failed(events) || failed(waits)
-          || failed(acknowledgements))
+      if (failed(offsets) || failed(events) || failed(waits))
         return failure();
       hostOffsets = std::move(*offsets);
       eventRegisters = std::move(*events);
       waitValues = std::move(*waits);
-      acknowledgementEventRegisters = std::move(*acknowledgements);
     }
     SmallVector<StaticIntSequence> positionRows;
     for (unsigned position : run.positions)
@@ -485,15 +494,11 @@ static FailureOr<Value> emitReceiveAssembly(const EmitReceiveAssemblyRun &run, V
         &LogicalTransferMetadataView::eventRegisters);
       auto waits = buildGrid(
         &LogicalTransferMetadataView::waitValues);
-      auto acknowledgements = buildGrid(
-        &LogicalTransferMetadataView::acknowledgementEventRegisters);
-      if (failed(offsets) || failed(events) || failed(waits)
-          || failed(acknowledgements))
+      if (failed(offsets) || failed(events) || failed(waits))
         return failure();
       hostOffsets = std::move(*offsets);
       eventRegisters = std::move(*events);
       waitValues = std::move(*waits);
-      acknowledgementEventRegisters = std::move(*acknowledgements);
     }
     SmallVector<StaticIntSequence> positionColumns;
     for (const StaticIntSequenceChain &values : positionsByLane)
@@ -527,8 +532,6 @@ static FailureOr<Value> emitReceiveAssembly(const EmitReceiveAssemblyRun &run, V
         eventRegisters->emitLookup(
           entry, runtimeLane, anchor, context.constants, context.rewriter, loc),
         waitValues->emitLookup(
-          entry, runtimeLane, anchor, context.constants, context.rewriter, loc),
-        acknowledgementEventRegisters->emitLookup(
           entry, runtimeLane, anchor, context.constants, context.rewriter, loc));
       receive = op;
       output = op.getOutput();
@@ -1083,7 +1086,8 @@ static LogicalResult emitLocalCollectionUpdate(const EmitLocalCollectionRun &upd
 }
 
 static FailureOr<SmallVector<Value>> emitInstructions(ArrayRef<BoundaryInstruction> instructions, Value lane, unsigned laneCount,
-                                                      ArrayRef<DeferredResultPlan> results, DeferredEmissionContext &context) {
+                                                      ArrayRef<DeferredResultPlan> results, ScheduledInfo &scheduled,
+                                                      DeferredEmissionContext &context) {
   SmallVector<Value> produced;
   for (size_t instructionIndex = 0;
        instructionIndex < instructions.size(); ++instructionIndex) {
@@ -1108,7 +1112,22 @@ static FailureOr<SmallVector<Value>> emitInstructions(ArrayRef<BoundaryInstructi
           }
         }
     }
-    if (auto send = std::get_if<EmitSendRun>(&instruction)) {
+    if (auto reuse = std::get_if<EmitHostReuseWait>(&instruction)) {
+      Value runtimeLane = lane ? lane : context.constants.getIndex(0);
+      Location loc = scheduled.op->getLoc();
+      if (failed(emitForLanes(
+            reuse->lanes, lane, laneCount, scheduled.op, context, loc, [&]() {
+              SpatWaitOp::create(
+                context.rewriter, loc,
+                emitStaticIntLookup(
+                  reuse->eventRegisters, runtimeLane, scheduled.op,
+                  context.constants, context.rewriter, loc),
+                emitStaticIntLookup(
+                  reuse->waitValues, runtimeLane, scheduled.op,
+                  context.constants, context.rewriter, loc));
+            })))
+        return failure();
+    } else if (auto send = std::get_if<EmitSendRun>(&instruction)) {
       if (failed(emitConditionalSendRun(*send, lane, laneCount, context)))
         return failure();
     } else if (auto update = std::get_if<EmitLocalCollectionRun>(&instruction)) {
@@ -1192,7 +1211,9 @@ static LogicalResult emitBoundary(const BoundaryProgram &boundary, ArrayRef<Defe
   if (auto batch = dyn_cast<SpatScheduledComputeBatch>(boundary.key.first->op))
     lane = *batch.getLaneArgument();
   SmallVector<DeferredExchangePlan *> exchanges = getProducedExchanges(boundary.instructions);
-  auto values = emitInstructions(boundary.instructions, lane, laneCount, results, context);
+  auto values = emitInstructions(
+    boundary.instructions, lane, laneCount, results, *boundary.key.first,
+    context);
   return failed(values) ? failure() : replaceResults(exchanges, *values, replacements);
 }
 
@@ -1203,30 +1224,129 @@ static unsigned getBarrierRoundCount(size_t coreCount) {
   return rounds;
 }
 
+static LogicalResult emitHostReleaseSynchronization(
+    DeferredTransferPlan &transfers, DeferredEmissionContext &context) {
+  struct Emission {
+    HostReleaseSignal signal;
+    LaneSet lanes;
+  };
+  DenseSet<int64_t> emittedReaders;
+  for (ScheduledInfo &scheduled : transfers.scheduled) {
+    SmallVector<Emission> emissions;
+    for (auto [laneIndex, core] : llvm::enumerate(scheduled.cores)) {
+      auto found = transfers.hostReleaseSignals.find(core);
+      if (found == transfers.hostReleaseSignals.end())
+        continue;
+      if (!emittedReaders.insert(core).second)
+        return scheduled.op->emitOpError(
+          "host release reader is represented by multiple scheduled lanes");
+      for (HostReleaseSignal signal : found->second) {
+        auto emission = llvm::find_if(emissions, [&](const Emission &item) {
+          return item.signal.writerCore == signal.writerCore
+              && item.signal.eventRegister == signal.eventRegister;
+        });
+        if (emission == emissions.end()) {
+          emissions.push_back({signal, LaneSet::range(
+            laneIndex, laneIndex + 1)});
+        } else {
+          emission->lanes = emission->lanes.unite(
+            LaneSet::range(laneIndex, laneIndex + 1));
+        }
+      }
+    }
+    if (emissions.empty())
+      continue;
+    Block *block = scheduled.blocks.front();
+    context.rewriter.setInsertionPointToStart(block);
+    Location loc = scheduled.op->getLoc();
+    Value lane;
+    if (auto batch = dyn_cast<SpatScheduledComputeBatch>(scheduled.op))
+      lane = *batch.getLaneArgument();
+    for (const Emission &emission : emissions)
+      if (failed(emitForLanes(
+            emission.lanes, lane, scheduled.cores.size(), scheduled.op,
+            context, loc, [&]() {
+              SpatSyncOp::create(
+                context.rewriter, loc,
+                context.constants.getIndex(emission.signal.writerCore),
+                context.constants.getIndex(emission.signal.eventRegister));
+            })))
+        return failure();
+  }
+  return success();
+}
+
+static LogicalResult emitRestartReleasePermissions(
+    DeferredTransferPlan &transfers, DeferredEmissionContext &context) {
+  if (transfers.downstreamCores.empty())
+    return success();
+  unsigned barrierRounds = getBarrierRoundCount(
+    transfers.stageZeroCores.size());
+  size_t restartRegister = transfers.synchronizationRegisterCount - 1;
+  size_t leaderPermissionRegister = restartRegister - barrierRounds;
+  size_t downstreamPermissionRegister = restartRegister - 1;
+  DenseMap<int64_t, unsigned> downstreamRank;
+  for (auto [rank, core] : llvm::enumerate(transfers.downstreamCores))
+    downstreamRank[core] = rank;
+
+  for (ScheduledInfo &scheduled : transfers.scheduled) {
+    LaneSet lanes;
+    SmallVector<int64_t> targets(scheduled.cores.size());
+    SmallVector<int64_t> registers(scheduled.cores.size());
+    for (auto [lane, core] : llvm::enumerate(scheduled.cores)) {
+      auto rank = downstreamRank.find(core);
+      if (rank == downstreamRank.end())
+        continue;
+      lanes = lanes.unite(LaneSet::range(lane, lane + 1));
+      if (rank->second == 0) {
+        targets[lane] = transfers.stageZeroCores.front();
+        registers[lane] = leaderPermissionRegister;
+      } else {
+        targets[lane] = transfers.downstreamCores[(rank->second - 1) / 2];
+        registers[lane] = downstreamPermissionRegister;
+      }
+    }
+    if (lanes.empty())
+      continue;
+    Block *block = scheduled.blocks.front();
+    context.rewriter.setInsertionPointToStart(block);
+    Location loc = scheduled.op->getLoc();
+    Value lane;
+    if (auto batch = dyn_cast<SpatScheduledComputeBatch>(scheduled.op))
+      lane = *batch.getLaneArgument();
+    Value runtimeLane = lane ? lane : context.constants.getIndex(0);
+    if (failed(emitForLanes(
+          lanes, lane, scheduled.cores.size(), scheduled.op, context, loc,
+          [&]() {
+            Value target = emitStaticIntLookup(
+              StaticIntSequence::fromValues(targets), runtimeLane,
+              scheduled.op, context.constants, context.rewriter, loc);
+            Value eventRegister = emitStaticIntLookup(
+              StaticIntSequence::fromValues(registers), runtimeLane,
+              scheduled.op, context.constants, context.rewriter, loc);
+            SpatSyncOp::create(
+              context.rewriter, loc, target, eventRegister);
+          })))
+      return failure();
+  }
+  return success();
+}
+
 static LogicalResult emitCompletionSynchronization(
     DeferredTransferPlan &transfers, DeferredEmissionContext &context) {
   if (transfers.synchronizationRegisterCount == 0)
     return success();
-  size_t acknowledgementRegister =
+  size_t restartRegister =
     transfers.synchronizationRegisterCount - 1;
   unsigned barrierRounds = getBarrierRoundCount(
     transfers.stageZeroCores.size());
-  bool stageZeroNeedsAcknowledgements = llvm::any_of(
-    transfers.stageZeroCores, [&](int64_t core) {
-      return transfers.hostAcknowledgementCounts.contains(core);
-    });
-  size_t firstBarrierRegister = acknowledgementRegister
-                              - (stageZeroNeedsAcknowledgements ? 1 : 0);
+  size_t firstBarrierRegister = restartRegister;
   DenseMap<int64_t, unsigned> stageZeroRank;
   for (auto [rank, core] : llvm::enumerate(transfers.stageZeroCores))
     stageZeroRank[core] = rank;
   DenseMap<int64_t, unsigned> downstreamRank;
   for (auto [rank, core] : llvm::enumerate(transfers.downstreamCores))
     downstreamRank[core] = rank;
-  auto getReleaseRegister = [&](int64_t core) {
-    return acknowledgementRegister
-         - (transfers.hostAcknowledgementCounts.contains(core) ? 1 : 0);
-  };
 
   for (ScheduledInfo &scheduled : transfers.scheduled) {
     Block *block = scheduled.blocks.front();
@@ -1236,13 +1356,10 @@ static LogicalResult emitCompletionSynchronization(
     if (auto batch = dyn_cast<SpatScheduledComputeBatch>(scheduled.op))
       lane = *batch.getLaneArgument();
 
-    SmallVector<int64_t> acknowledgementCounts, releaseRegisters;
-    SmallVector<int64_t> releaseWaitValues, leftTargets, leftRegisters;
-    SmallVector<int64_t> rightTargets, rightRegisters;
-    LaneSet barrierLanes, leaderLanes, leftLanes, rightLanes;
+    SmallVector<int64_t> leftTargets, rightTargets, childCounts;
+    LaneSet barrierLanes, leaderLanes, downstreamLanes, parentLanes,
+      leftLanes, rightLanes;
     for (auto [index, core] : llvm::enumerate(scheduled.cores)) {
-      acknowledgementCounts.push_back(
-        transfers.hostAcknowledgementCounts.lookup(core));
       if (stageZeroRank.contains(core))
         barrierLanes = barrierLanes.unite(
           LaneSet::range(index, index + 1));
@@ -1252,66 +1369,37 @@ static LogicalResult emitCompletionSynchronization(
 
       auto rank = downstreamRank.find(core);
       if (rank == downstreamRank.end()) {
-        releaseRegisters.push_back(0);
-        releaseWaitValues.push_back(0);
         leftTargets.push_back(core);
-        leftRegisters.push_back(0);
         rightTargets.push_back(core);
-        rightRegisters.push_back(0);
+        childCounts.push_back(0);
         continue;
       }
-      releaseRegisters.push_back(getReleaseRegister(core));
-      releaseWaitValues.push_back(1);
+      downstreamLanes = downstreamLanes.unite(
+        LaneSet::range(index, index + 1));
       size_t left = 2 * rank->second + 1;
       size_t right = left + 1;
+      unsigned childCount = 0;
       if (left < transfers.downstreamCores.size()) {
         int64_t child = transfers.downstreamCores[left];
         leftTargets.push_back(child);
-        leftRegisters.push_back(getReleaseRegister(child));
         leftLanes = leftLanes.unite(LaneSet::range(index, index + 1));
+        ++childCount;
       } else {
         leftTargets.push_back(core);
-        leftRegisters.push_back(0);
       }
       if (right < transfers.downstreamCores.size()) {
         int64_t child = transfers.downstreamCores[right];
         rightTargets.push_back(child);
-        rightRegisters.push_back(getReleaseRegister(child));
         rightLanes = rightLanes.unite(LaneSet::range(index, index + 1));
+        ++childCount;
       } else {
         rightTargets.push_back(core);
-        rightRegisters.push_back(0);
       }
+      childCounts.push_back(childCount);
+      if (childCount != 0)
+        parentLanes = parentLanes.unite(LaneSet::range(index, index + 1));
     }
     Value runtimeLane = lane ? lane : context.constants.getIndex(0);
-    auto emitForLanes = [&](const LaneSet &active, auto emit) -> LogicalResult {
-      if (active.empty())
-        return success();
-      if (!lane) {
-        if (active.contains(0))
-          emit();
-        return success();
-      }
-      auto condition = emitLaneCondition(
-        active, lane, scheduled.cores.size(), scheduled.op, context, loc);
-      if (failed(condition))
-        return failure();
-      auto conditional = scf::IfOp::create(
-        context.rewriter, loc, TypeRange {}, *condition, false);
-      OpBuilder::InsertionGuard guard(context.rewriter);
-      context.rewriter.setInsertionPoint(
-        conditional.getThenRegion().front().getTerminator());
-      emit();
-      return success();
-    };
-    Value acknowledgementCount = emitStaticIntLookup(
-      StaticIntSequence::fromValues(acknowledgementCounts),
-      runtimeLane, scheduled.op,
-      context.constants, context.rewriter, loc);
-    SpatWaitOp::create(
-      context.rewriter, loc,
-      context.constants.getIndex(acknowledgementRegister),
-      acknowledgementCount);
 
     // Dissemination barrier: every round doubles the covered stage-zero peers.
     auto emitBarrier = [&]() {
@@ -1341,45 +1429,64 @@ static LogicalResult emitCompletionSynchronization(
       }
     };
     if (barrierRounds > 0
-        && failed(emitForLanes(barrierLanes, emitBarrier)))
+        && failed(emitForLanes(
+          barrierLanes, lane, scheduled.cores.size(), scheduled.op, context,
+          loc, emitBarrier)))
       return failure();
 
     // Gate downstream restarts so no core advances the simulator input
     // iteration ahead of stage zero.
     if (!transfers.downstreamCores.empty()
-        && failed(emitForLanes(leaderLanes, [&]() {
+        && failed(emitForLanes(
+          leaderLanes, lane, scheduled.cores.size(), scheduled.op, context,
+          loc, [&]() {
+          SpatWaitOp::create(
+            context.rewriter, loc,
+            context.constants.getIndex(firstBarrierRegister - barrierRounds),
+            context.constants.getIndex(1));
           int64_t root = transfers.downstreamCores.front();
           SpatSyncOp::create(
             context.rewriter, loc, context.constants.getIndex(root),
-            context.constants.getIndex(getReleaseRegister(root)));
+            context.constants.getIndex(restartRegister));
         })))
       return failure();
 
-    Value releaseRegister = emitStaticIntLookup(
-      StaticIntSequence::fromValues(releaseRegisters), runtimeLane,
-      scheduled.op, context.constants, context.rewriter, loc);
-    Value releaseWaitValue = emitStaticIntLookup(
-      StaticIntSequence::fromValues(releaseWaitValues), runtimeLane,
-      scheduled.op, context.constants, context.rewriter, loc);
-    SpatWaitOp::create(
-      context.rewriter, loc, releaseRegister, releaseWaitValue);
+    if (failed(emitForLanes(
+          downstreamLanes, lane, scheduled.cores.size(), scheduled.op,
+          context, loc, [&]() {
+            SpatWaitOp::create(
+              context.rewriter, loc,
+              context.constants.getIndex(restartRegister),
+              context.constants.getIndex(1));
+          })))
+      return failure();
 
-    auto emitChild = [&](ArrayRef<int64_t> targets,
-                         ArrayRef<int64_t> registers) {
+    if (failed(emitForLanes(
+          parentLanes, lane, scheduled.cores.size(), scheduled.op, context,
+          loc, [&]() {
+            Value count = emitStaticIntLookup(
+              StaticIntSequence::fromValues(childCounts), runtimeLane,
+              scheduled.op, context.constants, context.rewriter, loc);
+            SpatWaitOp::create(
+              context.rewriter, loc,
+              context.constants.getIndex(restartRegister - 1), count);
+          })))
+      return failure();
+
+    auto emitChild = [&](ArrayRef<int64_t> targets) {
       Value target = emitStaticIntLookup(
         StaticIntSequence::fromValues(targets), runtimeLane, scheduled.op,
         context.constants, context.rewriter, loc);
-      Value eventRegister = emitStaticIntLookup(
-        StaticIntSequence::fromValues(registers), runtimeLane, scheduled.op,
-        context.constants, context.rewriter, loc);
-      SpatSyncOp::create(context.rewriter, loc, target, eventRegister);
+      SpatSyncOp::create(
+        context.rewriter, loc, target,
+        context.constants.getIndex(restartRegister));
     };
-    if (failed(emitForLanes(leftLanes, [&]() {
-          emitChild(leftTargets, leftRegisters);
-        }))
-        || failed(emitForLanes(rightLanes, [&]() {
-          emitChild(rightTargets, rightRegisters);
-        })))
+    if (failed(emitForLanes(
+          leftLanes, lane, scheduled.cores.size(), scheduled.op, context, loc,
+          [&]() { emitChild(leftTargets); }))
+        || failed(emitForLanes(
+          rightLanes, lane, scheduled.cores.size(), scheduled.op, context,
+          loc, [&]() { emitChild(rightTargets); })))
       return failure();
   }
   return success();
@@ -1399,6 +1506,10 @@ LogicalResult realizeDeferredBoundaries(ArrayRef<BoundaryProgram> boundaries, Ar
     if (failed(emitBoundary(boundary, results, context, replacements)))
       return boundary.key.first->op->emitOpError("phase 2 failed to realize a communication boundary");
   }
+  if (failed(emitRestartReleasePermissions(transfers, context)))
+    return failure();
+  if (failed(emitHostReleaseSynchronization(transfers, context)))
+    return failure();
   return emitCompletionSynchronization(transfers, context);
 }
 

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationModel.hpp

@@ -237,8 +237,6 @@ struct ExternalTransferFamily {
   StaticIntSequence hostOffsets = StaticIntSequence::uniform(0, 1);
   StaticIntSequence eventRegisters = StaticIntSequence::uniform(0, 1);
   StaticIntSequence waitValues = StaticIntSequence::uniform(1, 1);
-  StaticIntSequence acknowledgementEventRegisters =
-    StaticIntSequence::uniform(0, 1);
   bool hostRouted = false;
 };
 

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.hpp

@@ -7,6 +7,11 @@
 
 namespace onnx_mlir::spatial {
 
+struct HostReleaseSignal {
+  int64_t writerCore = -1;
+  unsigned eventRegister = 0;
+};
+
 struct DeferredTransferPlan {
   std::vector<size_t> processorStages;
   llvm::SmallVector<ScheduledInfo, 0> scheduled;
@@ -14,7 +19,8 @@ struct DeferredTransferPlan {
   llvm::DenseMap<int64_t, llvm::SmallVector<ProducedValue*>> producedByGraph;
   llvm::SmallVector<std::unique_ptr<DeferredExchangePlan>> exchanges;
   llvm::SmallVector<unsigned> stepCounts;
-  llvm::DenseMap<int64_t, unsigned> hostAcknowledgementCounts;
+  llvm::DenseMap<int64_t, llvm::SmallVector<HostReleaseSignal>>
+    hostReleaseSignals;
   llvm::SmallVector<int64_t> stageZeroCores;
   llvm::SmallVector<int64_t> downstreamCores;
   size_t synchronizationRegisterCount = 0;

src/PIM/Dialect/Spatial/Spatial.td

@@ -592,15 +592,14 @@ def SpatHostStoreSyncOp : SpatOp<"host_store_sync", []> {
 }
 
 def SpatHostWaitLoadOp : SpatOp<"host_wait_load", []> {
-  let summary = "Wait for producers, load from host memory, and acknowledge consumption";
+  let summary = "Wait for producers and load from host memory";
 
   let arguments = (ins
     Index:$sourceCoreId,
     Index:$targetCoreId,
     Index:$hostOffset,
     Index:$eventRegister,
-    Index:$waitValue,
-    Index:$acknowledgementEventRegister
+    Index:$waitValue
   );
 
   let results = (outs
@@ -610,7 +609,7 @@ def SpatHostWaitLoadOp : SpatOp<"host_wait_load", []> {
   let assemblyFormat = [{
     `from` $sourceCoreId `to` $targetCoreId
       `host_offset` $hostOffset `event` $eventRegister `count` $waitValue
-      `ack` $acknowledgementEventRegister attr-dict `:` type($output)
+      attr-dict `:` type($output)
   }];
 }