fix: add pre-loop event init and tail event drain for back-edge sync (#428)

The generated C++ Cube kernel hangs with --enable-insert-sync because:

1. Back-edge wait_flag ops hoisted to loop heads fire on iteration 0
   against uninitialized event registers. This patch adds a pre-loop
   set_flag for each back-edge wait to initialize the event before the
   first iteration.

2. The tail sync (ptoas_auto_sync_tail with kBarrierAll) only emits
   pipe_barrier(PIPE_ALL) which does not drain pending event flag
   registers. This patch adds explicit wait_flag calls at the function
   tail to drain all back-edge event dependencies.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: zhangstevenunity

include/PTO/Transforms/InsertSync/SyncCodegen.h

@@ -57,6 +57,19 @@ class SyncCodegen {
 
   // Insert the compiler tail-clean barrier right before function return.
   void AppendAutoSyncTailBarrierIfNeeded(IRRewriter &rewriter);
+
+  // [Fix #428] Collect back-edge waits at loop heads that need pre-loop
+  // set_flag initialization to avoid hangs on iteration 0.
+  void CollectBackEdgeLoopHeadWaits();
+
+  // [Fix #428] Emit set_flag ops before a for-loop to initialize event
+  // registers for back-edge wait_flag ops at the loop head.
+  void EmitPreLoopEventInit(IRRewriter &rewriter, Operation *op);
+
+  // [Fix #428] Emit explicit wait_flag ops before function return to drain
+  // all pending back-edge event dependencies that pipe_barrier(PIPE_ALL)
+  // alone does not cover.
+  void EmitTailEventDrain(IRRewriter &rewriter, func::ReturnOp ret);
 
   void CreateSetWaitOpForSingleBuffer(IRRewriter &rewriter, Operation *op,
                                       SyncOperation *sync, bool beforeInsert);
@@ -92,6 +105,10 @@ class SyncCodegen {
 
   // Deferred tail-clean barrier requested by sync analysis.
   bool pendingAutoSyncTailBarrier_ = false;
+
+  // [Fix #428] Map from scf.for Operation* to the back-edge wait ops at
+  // its loop head that need pre-loop set_flag initialization.
+  DenseMap<const Operation *, SmallVector<SyncOperation *, 4>> preLoopInitWaits_;
 };
 
 } // namespace pto

lib/PTO/Transforms/InsertSync/SyncCodegen.cpp

@@ -66,14 +66,25 @@ static void MergeSyncList(SyncOps &dstList, const SyncOps &srcList) {
 void SyncCodegen::Run() {
   MLIRContext *ctx = func_->getContext();
   IRRewriter rewriter(ctx);
-  
+
   UpdateOpInsertSync(rewriter);
- 
-  // [Optional Debug] 这里的 Debug 打印可以保留或注释掉
-  // ...
- 
+
+  // [Fix #428] Collect back-edge wait ops at loop heads that need pre-loop
+  // set_flag initialization. A back-edge wait_flag at LOOP_BEGIN fires on
+  // iteration 0, but the matching set_flag only fires at the end of iteration
+  // 0. Without a pre-loop set_flag, the event register is uninitialized and
+  // the hardware hangs.
+  CollectBackEdgeLoopHeadWaits();
+
   func_->walk<WalkOrder::PreOrder>([&](Operation *op) {
     if (op2InsertSync.count(op)) {
+      // [Fix #428] Before emitting the normal pipeBefore waits for a loop op,
+      // emit pre-loop set_flag initialization for any back-edge wait_flag that
+      // would otherwise fire against an uninitialized event register.
+      if (isa<scf::ForOp>(op)) {
+        EmitPreLoopEventInit(rewriter, op);
+      }
+
       // 处理 PRE Sync
       for (auto &syncBefore : op2InsertSync[op].pipeBefore) {
         SyncInsert(rewriter, op, syncBefore, true);
@@ -299,6 +310,15 @@ void SyncCodegen::AppendAutoSyncTailBarrierIfNeeded(IRRewriter &rewriter) {
 
   auto pipeAllAttr = getPipeAttr(rewriter, PipelineType::PIPE_ALL);
   for (auto ret : returns) {
+    // [Fix #428] Before the tail barrier, emit explicit wait_flag ops to
+    // drain all pending back-edge event dependencies. pipe_barrier(PIPE_ALL)
+    // waits for in-flight pipe operations but does NOT drain event flag
+    // registers. Without explicit wait_flag calls, stale event state can
+    // leak to the next kernel invocation.
+    rewriter.setInsertionPoint(ret);
+    EmitTailEventDrain(rewriter, ret);
+
+    // Re-set insertion point before ret (after any drain ops we just emitted)
     rewriter.setInsertionPoint(ret);
     auto barrier = rewriter.create<pto::BarrierOp>(ret.getLoc(), pipeAllAttr);
     barrier->setAttr("pto.auto_sync_tail_barrier", rewriter.getUnitAttr());
@@ -332,6 +352,84 @@ void SyncCodegen::CreateSetWaitOpForSingleBuffer(IRRewriter &rewriter,
     rewriter.create<pto::SetFlagOp>(op->getLoc(), srcPipe, dstPipe, eventId);
   }
 }
+
+// ==============================================================================
+// [Fix #428] Tail event drain for back-edge sync events
+// ==============================================================================
+
+void SyncCodegen::EmitTailEventDrain(IRRewriter &rewriter,
+                                     func::ReturnOp ret) {
+  // Collect all unique (srcPipe, dstPipe, eventId) triples from back-edge
+  // syncs across all loops. These events may still be pending when the kernel
+  // reaches the return statement and must be explicitly drained.
+  //
+  // We use a set of tuples to deduplicate — the same event may appear in
+  // multiple loops or be shared via widen.
+  struct EventKey {
+    PipelineType src;
+    PipelineType dst;
+    int eventId;
+    bool operator<(const EventKey &o) const {
+      if (src != o.src) return static_cast<unsigned>(src) < static_cast<unsigned>(o.src);
+      if (dst != o.dst) return static_cast<unsigned>(dst) < static_cast<unsigned>(o.dst);
+      return eventId < o.eventId;
+    }
+    bool operator==(const EventKey &o) const {
+      return src == o.src && dst == o.dst && eventId == o.eventId;
+    }
+  };
+
+  SmallVector<EventKey> drainEvents;
+  auto addUnique = [&](PipelineType src, PipelineType dst, int eid) {
+    EventKey key{src, dst, eid};
+    for (auto &existing : drainEvents) {
+      if (existing == key)
+        return;
+    }
+    drainEvents.push_back(key);
+  };
+
+  // Scan all LOOP_END elements for back-edge set/wait pairs with allocated
+  // event IDs.
+  for (auto &pair : preLoopInitWaits_) {
+    for (auto *waitSync : pair.second) {
+      if (waitSync->uselessSync || waitSync->eventIds.empty())
+        continue;
+      addUnique(waitSync->GetActualSrcPipe(), waitSync->GetActualDstPipe(),
+                waitSync->eventIds[0]);
+    }
+  }
+
+  // Also scan the last element's pipeAfter for any set_flag ops that might
+  // leave events pending (these are the "syncEnd" phantom pairs from
+  // UpdateBackwardMatchSync that sink to the function tail).
+  if (!syncIR_.empty()) {
+    for (auto *sync : syncIR_.back()->pipeAfter) {
+      if (sync->uselessSync || sync->eventIds.empty())
+        continue;
+      if (sync->isSyncWaitType()) {
+        addUnique(sync->GetActualSrcPipe(), sync->GetActualDstPipe(),
+                  sync->eventIds[0]);
+      }
+    }
+  }
+
+  if (drainEvents.empty())
+    return;
+
+  // Sort for deterministic output.
+  llvm::sort(drainEvents);
+
+  LLVM_DEBUG(llvm::dbgs() << "[Fix #428] Emitting " << drainEvents.size()
+                          << " tail event drain wait_flag ops\n");
+
+  for (auto &ev : drainEvents) {
+    auto srcPipe = getPipeAttr(rewriter, ev.src);
+    auto dstPipe = getPipeAttr(rewriter, ev.dst);
+    auto eventId = getEventAttr(rewriter, ev.eventId);
+    rewriter.create<pto::WaitFlagOp>(ret.getLoc(), srcPipe, dstPipe, eventId);
+  }
+}
 
 void SyncCodegen::CreateSetWaitOpForMultiBuffer(IRRewriter &rewriter,
                                                 Operation *op,
@@ -399,3 +497,64 @@ Value SyncCodegen::GetBufferSelected(IRRewriter &rewriter, Operation *op,
   SyncIndex2SelectBuffer[sync->GetSyncIndex()] = selected;
   return selected;
 }
+
+// ==============================================================================
+// [Fix #428] Pre-loop event initialization for back-edge sync
+// ==============================================================================
+
+void SyncCodegen::CollectBackEdgeLoopHeadWaits() {
+  for (auto &nowElement : syncIR_) {
+    auto *loopElement = dyn_cast<LoopInstanceElement>(nowElement.get());
+    if (!loopElement || loopElement->getLoopKind() != KindOfLoop::LOOP_END)
+      continue;
+
+    // Look at the LOOP_BEGIN node's pipeBefore — these are waits that
+    // MoveSyncState hoisted to the loop head.
+    auto *loopBegin =
+        dyn_cast<LoopInstanceElement>(syncIR_[loopElement->beginId].get());
+    if (!loopBegin)
+      continue;
+
+    for (auto *sync : loopBegin->pipeBefore) {
+      if (sync->uselessSync)
+        continue;
+      if (!sync->isSyncWaitType())
+        continue;
+      if (sync->eventIds.empty())
+        continue;
+      // This is a wait at loop head with an allocated event ID.
+      // It needs a pre-loop set_flag to initialize the event register.
+      // Record {Operation* forOp -> SyncOperation* wait} for later emission.
+      if (loopElement->elementOp) {
+        preLoopInitWaits_[loopElement->elementOp].push_back(sync);
+      }
+    }
+  }
+}
+
+void SyncCodegen::EmitPreLoopEventInit(IRRewriter &rewriter, Operation *op) {
+  auto it = preLoopInitWaits_.find(op);
+  if (it == preLoopInitWaits_.end())
+    return;
+
+  // For each back-edge wait at the loop head, emit a set_flag before the
+  // for loop to initialize the event register. This ensures that on iteration
+  // 0, the wait_flag finds a valid (already-set) event instead of hanging.
+  rewriter.setInsertionPoint(op);
+  for (auto *waitSync : it->second) {
+    if (waitSync->uselessSync || waitSync->eventIds.empty())
+      continue;
+
+    auto srcPipe = getPipeAttr(rewriter, waitSync->GetActualSrcPipe());
+    auto dstPipe = getPipeAttr(rewriter, waitSync->GetActualDstPipe());
+    auto eventId = getEventAttr(rewriter, waitSync->eventIds[0]);
+
+    LLVM_DEBUG(llvm::dbgs()
+               << "[Fix #428] Emitting pre-loop set_flag("
+               << static_cast<unsigned>(waitSync->GetActualSrcPipe()) << ", "
+               << static_cast<unsigned>(waitSync->GetActualDstPipe()) << ", "
+               << waitSync->eventIds[0] << ") before loop\n");
+
+    rewriter.create<pto::SetFlagOp>(op->getLoc(), srcPipe, dstPipe, eventId);
+  }
+}