fix scheduling reasoning on physical cores directly

document sync experiment

.agents/invariants/PIMSIM_NN_ORACLE_INVARIANT.md

@@ -11,6 +11,10 @@ behavior defines the hardware model used for Raptor/PIMCOMP comparisons.
   simulated timing, scheduling, power, energy, or supported input programs.
 - Unsupported pimsim-nn operations must remain unsupported; do not approximate
   their timing or map them onto another operation.
+- Validation must report an explicit unsupported-op diagnostic as
+  `UNSUPPORTED`, not as a simulator failure.
+- A compiled artifact with zero active PIM cores has no meaningful
+  non-functional simulation and must be reported as `SKIP`.
 - Adapt compiler inputs to the oracle instead. For YOLO, use
   `validation/networks/pimcomp_models/yolo11n/yolo11n-pimsim-nn.onnx`, the
   dedicated pimsim-ready performance artifact with Softmax operations removed.

.agents/invariants/PIM_SYNCHRONIZATION_INVARIANT.md

@@ -0,0 +1,105 @@
+# PIM Synchronization Invariant
+
+## Scope
+
+This invariant applies to PIM `SYNC`/`WAIT` planning, event-register
+allocation, repeating core bodies, host/global-memory transfers, and pipeline
+restart synchronization.
+
+## ISA contract
+
+`SYNC event, target` and `WAIT event, expected` have static operands. Event
+registers start at zero, persist across instruction-stream restart, and are
+incremented by `SYNC`. `WAIT` succeeds only when the register equals its
+static expected count, then resets it to zero. A plan must therefore guarantee
+exactly the expected number of increments before each wait. It must not use
+generation-dependent registers, treat `WAIT` as `>=`, permit overshoot, or
+consume one event register with independent waits in the same iteration.
+
+Event registers are local to the target core. Allocation must keep READY,
+FREE/reuse, stage-zero barrier, and downstream restart events disjoint on each
+physical core.
+
+## Host/global-memory lifetime
+
+For every host-routed dependency from writer `W` to reader `R`, the repeating
+program must establish:
+
+```text
+ST_N -> READY SYNC -> READY WAIT -> LD_N
+LD_N -> reader END_N -> reader start_N+1 -> FREE SYNC
+     -> writer FREE WAIT -> ST_N+1
+```
+
+The reader emits one FREE `SYNC` per unique physical `(W, R)` relation at body
+entry, before any instruction that can block. The writer consumes each FREE
+event at most once per iteration, with a static count equal to the unique
+readers in that reuse group. The wait is placed immediately before the first
+host store that can overwrite a member of the group, after all preceding
+useful computation. Multiple transfers on one `(W, R)` relation do not create
+additional FREE signals.
+
+Relations with the same writer and first dangerous overwrite boundary belong
+to one FREE group: splitting them cannot unblock an earlier store and only
+adds waits and register pressure. After preserving the READY registers required
+by a core, planning assigns independent groups to distinct overwrite boundaries
+as remaining capacity permits. Capacity-forced merged groups are ordered by
+first overwrite and waited before their earliest member. Unused READY capacity
+must not force all readers into one early aggregate wait.
+
+Host receives lower to READY `WAIT` followed by host load only. They must not
+emit a post-load acknowledgement. Writers must not wait for host readers at
+the end of their body. The repeated reader-entry FREE signal is both the
+initial-slot bootstrap and the release of the previous generation.
+
+## Pipeline restart independence
+
+Host lifetime synchronization does not replace pipeline restart
+synchronization. At the end of every repeating stage-zero body, all unique
+stage-zero physical cores execute a dissemination barrier with
+`ceil(log2(stageZeroCoreCount))` rounds. In round `r`, rank `i` sends one
+signal to rank `(i + 2^r) mod count` and waits for exactly one signal on that
+round's destination-local event register. Each round has its own register, so
+an early signal for a later round cannot overshoot or satisfy another wait.
+
+Every downstream core sends one restart-permission signal to its release-tree
+parent at body entry, after all host FREE signals. The root signals the
+stage-zero leader; other cores signal their binary-tree parent. After the barrier, the lowest-ranked
+stage-zero core waits for the root's permission before sending one restart
+signal to it. Every downstream parent waits for exactly one restart signal
+and for one permission from each existing child before forwarding the restart
+to those children and ending its current body. This reverse permission path
+prevents a fast parent from sending generation `N+1` before a child consumes
+generation `N`, which would overshoot an exact-count restart event.
+
+Thus no downstream core can restart and advance its input generation until
+every stage-zero core has completed the protected generation, and every
+restart event receives exactly one signal between waits. Stage-zero and
+downstream core sets are disjoint; barrier, restart, and restart-permission
+registers are reserved on their respective target cores. Host READY/FREE
+allocation must use only the remaining registers and must not move, merge,
+remove, or weaken this ordering.
+
+## Ownership and verification
+
+Deferred boundary planning owns relation deduplication, event allocation,
+expected counts, and the first dangerous overwrite action. Boundary
+realization owns body-entry FREE signals and materializes planned waits at
+their exact boundaries. Channel lowering preserves those explicit operations
+and lowers host receives as READY wait plus load; it must not reconstruct
+lifetime policy after scheduling information is lost.
+
+Verification must reject non-static lane/core mappings, duplicate physical
+reader representations, register collisions, oversubscribed event capacity,
+and any plan whose generated FREE signal count differs from its static wait
+count. Static and simulator deadlock detectors must model blocking WAITs,
+their remaining matching SYNC producers, missing signals, and exact-count
+overshoot in addition to SEND/RECV cycles. Both detectors must retain the
+per-source contributions accumulated since the last successful WAIT and must not
+add a wait-for edge to a source that already supplied its share of the current
+event value merely because that source has the same static SYNC in a future
+iteration. Structural tests must inspect generated instruction streams. Happens-
+before tests must unroll at least two logical iterations and add only
+same-core restart edges plus matched communication edges. Functional
+simulation must enforce exact-count waits and include stalled or randomized
+legal schedules so correctness never depends on relative core speed.

.agents/invariants/PIPELINE_SCHEDULING_INVARIANT.md

@@ -18,17 +18,20 @@ Dynamic function inputs are stage-zero sources. Any operation that directly
 consumes one must belong to stage 0. A later stage may consume that data only
 through an explicit result forwarded by the preceding stage.
 
-Each logical core belongs to exactly one stage capacity range before physical
-placement. Those ranges cover every core but may have different sizes when the
-initial partitioner predicts a lower maximum stage interval. Physical placement
-may map a stage to arbitrary core IDs using the injected target topology.
+PEFT assigns physical cores using the injected target topology and their actual
+resident-weight capacity. Pipeline stage capacities cover every physical core
+and may have different sizes when the partitioner predicts a lower maximum
+stage interval. Before stage-local packing, the layout groups physical cores
+by bidirectional link cost; stage cores need not have consecutive IDs. Packing
+scores transfers on those physical links, and communication realization must
+preserve the resulting core identities.
 Synchronization and deferred transfers consume the explicit stage identity;
 they must not infer it from a physical core number after placement.
 
 ## Ownership
 
-Logical PEFT remains pipeline-agnostic. Stage partitioning is the first phase
-of pipeline scheduling and owns this invariant. It must construct a valid
+Physical PEFT supplies the initial placement. Stage partitioning is the first
+phase of pipeline scheduling and owns stage adjacency. It must construct a valid
 operation-level partition before physical-core packing. Operations split for
 physical capacity retain one shared stage identity. Repacking may move work
 only within its assigned stage. Deferred-transfer planning and
@@ -44,7 +47,7 @@ Before scheduled materialization, verify that:
 - every direct dynamic-function-input consumer belongs to stage 0;
 - every compute-graph edge stays within a stage or advances exactly one stage;
 - every stage-local resident-weight set fits its assigned physical core; and
-- stage capacities cover all logical cores exactly once; and
+- stage capacities cover all physical cores exactly once; and
 - physical placement is a permutation of all target cores.
 
 Pipeline scheduling tests must include an uneven physical-core layout and a

AGENTS.md

@@ -7,6 +7,7 @@ Before modifying the relevant subsystem, read:
 * `.agents/invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md`
 * `.agents/invariants/PERFORMANCE_OPTIMIZATION_INVARIANT.md`
 * `.agents/invariants/PIMSIM_NN_ORACLE_INVARIANT.md`
+* `.agents/invariants/PIM_SYNCHRONIZATION_INVARIANT.md`
 * `.agents/invariants/PIPELINE_SCHEDULING_INVARIANT.md`
 * `.agents/invariants/SPATIAL_TARGET_GENERALITY_INVARIANT.md`
 * Build commands:

backend-simulators/pim/pim-simulator/src/lib/instruction_set/isa.rs

@@ -934,11 +934,21 @@ pub fn recv(cores: &mut CPU, data: InstructionData) -> Result<InstructionStatus>
     Ok(InstructionStatus::Receiving(data))
 }
 
+#[inline(never)]
+pub fn isa_wait(functor: usize) -> bool {
+    (wait as *const () as usize) == functor
+}
+
 #[inline(never)]
 pub fn wait(cores: &mut CPU, data: InstructionData) -> Result<InstructionStatus> {
     Ok(InstructionStatus::Waiting(data))
 }
 
+#[inline(never)]
+pub fn isa_sync(functor: usize) -> bool {
+    (sync as *const () as usize) == functor
+}
+
 #[inline(never)]
 pub fn sync(cores: &mut CPU, data: InstructionData) -> Result<InstructionStatus> {
     Ok(InstructionStatus::Sync(data))

backend-simulators/pim/pim-simulator/src/lib/pimcore.rs

@@ -16,7 +16,7 @@ use crate::{
     cpu::CPU,
     instruction_set::{
         Instruction, InstructionStatus, Instructions,
-        isa::{NAMES, functor_to_name, isa_recv, isa_send},
+        isa::{NAMES, functor_to_name, isa_recv, isa_send, isa_sync, isa_wait},
     },
     memory_manager::type_traits::TryToUsize,
     send_recv::{SendRecv, handle_send_recv},
@@ -411,7 +411,13 @@ struct DeadlockInfo {
     states: String,
 }
 
-type SyncEvents = Vec<[i32; 32]>;
+#[derive(Debug, Clone, Default)]
+struct SyncEvent {
+    count: i32,
+    sources: HashMap<i32, i32>,
+}
+
+type SyncEvents = Vec<[SyncEvent; 32]>;
 
 fn print_status(core_instructions: &[CoreInstructions]) {
     let mut tot_instructions = 0;
@@ -543,7 +549,9 @@ impl<'a> Executable<'a> {
         let mut barrier_iteration = None;
         let mut cpu_progressed = 0;
         let max_core = cpu.num_core();
-        let mut sync_events: SyncEvents = vec![[0; 32]; max_core];
+        let mut sync_events: SyncEvents = (0..max_core)
+            .map(|_| std::array::from_fn(|_| SyncEvent::default()))
+            .collect();
         let mut cpu_index = 0;
         let mut cycle = 0;
         let mut scheduler_no_progress = 0usize;
@@ -684,9 +692,11 @@ impl<'a> Executable<'a> {
                     }
                     if (now.elapsed().unwrap() > Duration::from_secs(5)) {
                         print_status(cores_instructions);
-                        if let Some(deadlock) = detect_deadlock(cores_instructions) {
+                        if let Some(deadlock) =
+                            detect_deadlock(cores_instructions, &sync_events, batch_size)
+                        {
                             bail!(
-                                "Deadlock cycle detected: {} [{}]",
+                                "Communication deadlock detected: {} [{}]",
                                 deadlock.cycle,
                                 deadlock.states
                             );
@@ -771,9 +781,9 @@ impl<'a> Executable<'a> {
         }
         print_status(cores_instructions);
 
-        if let Some(deadlock) = detect_deadlock(cores_instructions) {
+        if let Some(deadlock) = detect_deadlock(cores_instructions, &sync_events, batch_size) {
             bail!(
-                "Deadlock cycle detected: {} [{}]",
+                "Communication deadlock detected: {} [{}]",
                 deadlock.cycle,
                 deadlock.states
             );
@@ -839,18 +849,23 @@ fn store_input(
     Ok(())
 }
 
-fn detect_deadlock(cores_instructions: &[CoreInstructions]) -> Option<DeadlockInfo> {
+fn detect_deadlock(
+    cores_instructions: &[CoreInstructions],
+    events: &SyncEvents,
+    batch_size: u32,
+) -> Option<DeadlockInfo> {
     #[derive(Debug, PartialEq, Eq)]
     enum CoreState {
         SendingTo(i32, i32),
         ReceivingFrom(i32, i32),
+        WaitingEvent(i32, i32, i32),
         Working,
         Halted,
     }
 
     let mut states = HashMap::new();
 
-    for core_inst in cores_instructions.iter() {
+    for (core, core_inst) in cores_instructions.iter().enumerate() {
         if core_inst.program_counter >= core_inst.instructions.len() {
             continue;
         }
@@ -867,94 +882,191 @@ fn detect_deadlock(cores_instructions: &[CoreInstructions]) -> Option<DeadlockIn
             );
         } else if isa_send(functor_address) {
             states.insert(this_core, CoreState::SendingTo(target_core, data.imm_len()));
+        } else if isa_wait(functor_address) {
+            let register = data.offset_select();
+            states.insert(
+                this_core,
+                CoreState::WaitingEvent(
+                    register,
+                    data.offset_value(),
+                    events[core][register as usize].count,
+                ),
+            );
         } else {
             states.insert(this_core, CoreState::Working);
         }
     }
 
-    let mut wait_for = HashMap::new();
+    let event_sources = |target: i32, register: i32| {
+        let mut sources = Vec::new();
+        let mut signal_count = 0usize;
+        let event = &events[target as usize][register as usize];
+        for core_inst in cores_instructions {
+            let matches = |instruction: &&Instruction| {
+                isa_sync(instruction.functor as usize)
+                    && instruction.data.get_core_immcore().1 == target
+                    && instruction.data.offset_select() == register
+            };
+            let remaining = core_inst.instructions[core_inst.program_counter..]
+                .iter()
+                .filter(matches)
+                .count();
+            let per_iteration = core_inst.instructions.iter().filter(matches).count();
+            if per_iteration == 0 {
+                continue;
+            }
+            let future_iterations = batch_size.saturating_sub(core_inst.current_iteration + 1);
+            let source = core_inst.instructions.iter().find(|instruction| {
+                isa_sync(instruction.functor as usize)
+                    && instruction.data.get_core_immcore().1 == target
+                    && instruction.data.offset_select() == register
+            });
+            let source = source.unwrap().data.get_core_immcore().0;
+            let contributed = event.sources.get(&source).copied().unwrap_or(0) as usize;
+            let needed = per_iteration.saturating_sub(contributed);
+            let count = (remaining + per_iteration * future_iterations as usize).min(needed);
+            if count != 0 {
+                sources.push(source);
+                signal_count += count;
+            }
+        }
+        sources.sort_unstable();
+        sources.dedup();
+        (sources, signal_count)
+    };
+
+    let format_state = |core: &i32| {
+        let position = cores_instructions.get(*core as usize);
+        let location = position.map_or_else(
+            || format!("core {}", core - 1),
+            |instructions| {
+                format!(
+                    "core {} iteration {} pc {}",
+                    core - 1,
+                    instructions.current_iteration,
+                    instructions.program_counter
+                )
+            },
+        );
+        match states.get(core).unwrap_or(&CoreState::Halted) {
+            CoreState::SendingTo(target, size) => {
+                format!("{location} send {}B -> {}", size, target - 1)
+            }
+            CoreState::ReceivingFrom(source, size) => {
+                format!("{location} recv {}B <- {}", size, source - 1)
+            }
+            CoreState::WaitingEvent(register, expected, observed) => {
+                format!("{location} wait event {register} == {expected} (observed {observed})")
+            }
+            CoreState::Working => format!("{location} working"),
+            CoreState::Halted => format!("{location} halted"),
+        }
+    };
+
+    let mut wait_for: HashMap<i32, Vec<i32>> = HashMap::new();
 
     for (&core_id, state) in states.iter() {
         match state {
             CoreState::SendingTo(target_core, size) => {
                 let target_state = states.get(target_core).unwrap_or(&CoreState::Halted);
                 if target_state != &CoreState::ReceivingFrom(core_id, *size) {
-                    wait_for.insert(core_id, *target_core);
+                    wait_for.insert(core_id, vec![*target_core]);
                 }
             }
             CoreState::ReceivingFrom(target_core, size) => {
                 let target_state = states.get(target_core).unwrap_or(&CoreState::Halted);
                 if target_state != &CoreState::SendingTo(core_id, *size) {
-                    wait_for.insert(core_id, *target_core);
+                    wait_for.insert(core_id, vec![*target_core]);
                 }
             }
+            CoreState::WaitingEvent(register, expected, observed) => {
+                if observed > expected {
+                    return Some(DeadlockInfo {
+                        cycle: format!(
+                            "core {} WAIT event {} overshot exact value {} with {}",
+                            core_id - 1,
+                            register,
+                            expected,
+                            observed
+                        ),
+                        states: format_state(&core_id),
+                    });
+                }
+                if observed == expected {
+                    continue;
+                }
+                let (sources, remaining_signals) = event_sources(core_id, *register);
+                if *observed as usize + remaining_signals < *expected as usize {
+                    return Some(DeadlockInfo {
+                        cycle: format!(
+                            "core {} WAIT event {} needs {} but only {} signal(s) can arrive",
+                            core_id - 1,
+                            register,
+                            expected,
+                            *observed as usize + remaining_signals
+                        ),
+                        states: format_state(&core_id),
+                    });
+                }
+                wait_for.insert(core_id, sources);
+            }
             CoreState::Working | CoreState::Halted => {}
         }
     }
 
-    let mut visited = HashSet::new();
-
-    for &start_core in wait_for.keys() {
-        if visited.contains(&start_core) {
-            continue;
+    fn find_cycle(
+        core: i32,
+        wait_for: &HashMap<i32, Vec<i32>>,
+        path: &mut Vec<i32>,
+        positions: &mut HashMap<i32, usize>,
+        visited: &mut HashSet<i32>,
+    ) -> Option<Vec<i32>> {
+        if let Some(position) = positions.get(&core) {
+            return Some(path[*position..].to_vec());
         }
+        if !visited.insert(core) {
+            return None;
+        }
+        positions.insert(core, path.len());
+        path.push(core);
+        if let Some(targets) = wait_for.get(&core) {
+            for target in targets {
+                if let Some(cycle) = find_cycle(*target, wait_for, path, positions, visited) {
+                    return Some(cycle);
+                }
+            }
+        }
+        path.pop();
+        positions.remove(&core);
+        None
+    }
 
+    let mut visited = HashSet::new();
+    for start_core in wait_for.keys() {
         let mut path = Vec::new();
-        let mut current_core = start_core;
-        let mut in_path = HashSet::new();
-
-        while let Some(&waiting_for) = wait_for.get(&current_core) {
-            path.push(current_core);
-            in_path.insert(current_core);
-            visited.insert(current_core);
-
-            // Found a closed loop!
-            if in_path.contains(&waiting_for) {
-                let cycle_start = path.iter().position(|&c| c == waiting_for).unwrap();
-                let cycle = &path[cycle_start..];
-                let format_core = |core: &i32| (core - 1).to_string();
-
-                let cycle_str = cycle
-                    .iter()
-                    .map(format_core)
-                    .collect::<Vec<_>>()
-                    .join(" -> ");
-
-                let cycle = cycle
-                    .iter()
-                    .copied()
-                    .chain(std::iter::once(waiting_for))
-                    .collect::<Vec<_>>();
-                let cycle_msg = format!("{} -> {}", cycle_str, waiting_for - 1);
-                let states_msg = cycle
-                    .iter()
-                    .filter_map(|core| {
-                        states.get(core).map(|state| match state {
-                            CoreState::SendingTo(target, size) => {
-                                format!("core {} send {}B -> {}", core - 1, size, target - 1)
-                            }
-                            CoreState::ReceivingFrom(source, size) => {
-                                format!("core {} recv {}B <- {}", core - 1, size, source - 1)
-                            }
-                            CoreState::Working => format!("core {} working", core - 1),
-                            CoreState::Halted => format!("core {} halted", core - 1),
-                        })
-                    })
-                    .collect::<Vec<_>>()
-                    .join(", ");
-
-                return Some(DeadlockInfo {
-                    cycle: cycle_msg,
-                    states: states_msg,
-                });
-            }
-
-            // Hit a known branch that didn't result in a cycle
-            if visited.contains(&waiting_for) {
-                break;
-            }
-
-            current_core = waiting_for;
+        let mut positions = HashMap::new();
+        if let Some(cycle) = find_cycle(
+            *start_core,
+            &wait_for,
+            &mut path,
+            &mut positions,
+            &mut visited,
+        ) {
+            let cycle_msg = cycle
+                .iter()
+                .chain(std::iter::once(&cycle[0]))
+                .map(|core| (core - 1).to_string())
+                .collect::<Vec<_>>()
+                .join(" -> ");
+            let states_msg = cycle
+                .iter()
+                .map(&format_state)
+                .collect::<Vec<_>>()
+                .join(", ");
+            return Some(DeadlockInfo {
+                cycle: cycle_msg,
+                states: states_msg,
+            });
         }
     }
     None
@@ -969,7 +1081,9 @@ fn handle_wait_sync(
         InstructionStatus::Sync(data) => {
             let (source, target) = data.get_core_immcore();
             let register = data.offset_select() as usize;
-            events[target as usize][register] += 1;
+            let event = &mut events[target as usize][register];
+            event.count += 1;
+            *event.sources.entry(source).or_default() += 1;
             core_instructions[source as usize].program_counter += 1;
             true
         }
@@ -977,8 +1091,8 @@ fn handle_wait_sync(
             let core = data.core_indx() as usize;
             let register = data.offset_select() as usize;
             let value = data.offset_value();
-            if events[core][register] >= value {
-                events[core][register] -= value;
+            if events[core][register].count == value {
+                events[core][register] = SyncEvent::default();
                 core_instructions[core].program_counter += 1;
                 true
             } else {
@@ -1128,3 +1242,66 @@ mod scheduler_tests {
         );
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::instruction_set::{
+        InstructionsBuilder,
+        instruction_data::InstructionDataBuilder,
+        isa::{sync, wait},
+    };
+
+    fn wait_then_sync(core: i32, target: i32) -> CoreInstructions {
+        let mut instructions = InstructionsBuilder::new();
+        let mut data = InstructionDataBuilder::new();
+        data.set_core_indx(core).fix_core_indx();
+        instructions.make_inst(wait, data.set_offset_select_value(0, 2).build());
+        instructions.make_inst(
+            sync,
+            data.set_imm_core(target)
+                .set_offset_select_value(1, 0)
+                .build(),
+        );
+        CoreInstructions::from(instructions.build())
+    }
+
+    fn sync_then_wait(core: i32, target: i32) -> CoreInstructions {
+        let mut instructions = InstructionsBuilder::new();
+        let mut data = InstructionDataBuilder::new();
+        data.set_core_indx(core).fix_core_indx();
+        instructions.make_inst(
+            sync,
+            data.set_imm_core(target)
+                .set_offset_select_value(0, 0)
+                .build(),
+        );
+        instructions.make_inst(wait, data.set_offset_select_value(1, 1).build());
+        CoreInstructions::from(instructions.build())
+    }
+
+    #[test]
+    fn contributed_sync_source_is_not_a_wait_dependency() {
+        let mut writer = wait_then_sync(1, 2);
+        writer.current_iteration = 1;
+        let mut contributed_reader = sync_then_wait(2, 1);
+        contributed_reader.current_iteration = 1;
+        contributed_reader.program_counter = 1;
+        let pending_reader = sync_then_wait(3, 1);
+        let cores = vec![
+            CoreInstructions::empty(),
+            writer,
+            contributed_reader,
+            pending_reader,
+        ];
+        let mut events: SyncEvents = (0..cores.len())
+            .map(|_| std::array::from_fn(|_| SyncEvent::default()))
+            .collect();
+        events[1][0].count = 1;
+        events[1][0].sources.insert(2, 1);
+
+        assert!(detect_deadlock(&cores, &events, 3).is_none());
+        events[1][0].sources.clear();
+        assert!(detect_deadlock(&cores, &events, 3).is_some());
+    }
+}

backend-simulators/pim/pim-simulator/tests/sync.rs

@@ -297,7 +297,7 @@ fn multiple_send_recv_test() {
 }
 
 #[test]
-fn sync_wait_tokens_test() {
+fn sync_wait_exact_count_resets_test() {
     let cpu = common::empty_cpu(2);
     let mut cores = CoreInstructionsBuilder::new(2);
     let mut instructions = InstructionsBuilder::new();
@@ -313,14 +313,68 @@ fn sync_wait_tokens_test() {
     cores.set_core(1, instructions.build());
 
     data.set_core_indx(2).fix_core_indx();
-    for _ in 0..2 {
-        instructions.make_inst(wait, data.set_offset_select_value(0, 1).build());
-    }
+    instructions.make_inst(wait, data.set_offset_select_value(0, 2).build());
     cores.set_core(2, instructions.build());
 
     Executable::new(cpu, cores.build()).execute().unwrap();
 }
 
+#[test]
+fn sync_wait_rejects_overshoot() {
+    let cpu = common::empty_cpu(3);
+    let mut cores = CoreInstructionsBuilder::new(3);
+    let mut instructions = InstructionsBuilder::new();
+    let mut data = InstructionDataBuilder::new();
+
+    data.set_core_indx(1).fix_core_indx();
+    instructions.make_inst(
+        sync,
+        data.set_imm_core(3).set_offset_select_value(0, 0).build(),
+    );
+    cores.set_core(1, instructions.build());
+
+    data.set_core_indx(2).fix_core_indx();
+    instructions.make_inst(
+        sync,
+        data.set_imm_core(3).set_offset_select_value(0, 0).build(),
+    );
+    cores.set_core(2, instructions.build());
+
+    data.set_core_indx(3).fix_core_indx();
+    instructions.make_inst(wait, data.set_offset_select_value(0, 1).build());
+    cores.set_core(3, instructions.build());
+
+    let error = Executable::new(cpu, cores.build()).execute().unwrap_err();
+    assert!(error.to_string().contains("overshot exact value"));
+}
+
+#[test]
+fn sync_wait_deadlock_cycle_is_reported() {
+    let cpu = common::empty_cpu(2);
+    let mut cores = CoreInstructionsBuilder::new(2);
+    let mut instructions = InstructionsBuilder::new();
+    let mut data = InstructionDataBuilder::new();
+
+    data.set_core_indx(1).fix_core_indx();
+    instructions.make_inst(wait, data.set_offset_select_value(0, 1).build());
+    instructions.make_inst(
+        sync,
+        data.set_imm_core(2).set_offset_select_value(0, 0).build(),
+    );
+    cores.set_core(1, instructions.build());
+
+    data.set_core_indx(2).fix_core_indx();
+    instructions.make_inst(wait, data.set_offset_select_value(0, 1).build());
+    instructions.make_inst(
+        sync,
+        data.set_imm_core(1).set_offset_select_value(0, 0).build(),
+    );
+    cores.set_core(2, instructions.build());
+
+    let error = Executable::new(cpu, cores.build()).execute().unwrap_err();
+    assert!(error.to_string().contains("wait event"));
+}
+
 #[test]
 fn blocked_transfers_do_not_starve_sync_producer() {
     let cpu = common::empty_cpu(4);

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/DeferredCommunicationRealization.cpp

@@ -6,7 +6,6 @@
 #include "DeferredCommunicationRealization.hpp"
 #include "DeferredCommunicationScheduling.hpp"
 #include "DeferredTransferPlanning.hpp"
-#include "Scheduling/PeftScheduler.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 
@@ -34,76 +33,6 @@ static LogicalResult verifyNoEscapingRegionValues(Operation* owner, StringRef ph
                               << escapingUser->getName() << " at " << escapingUser->getLoc();
 }
 
-static LogicalResult placeLogicalProcessorsOnPhysicalCores(
-    DeferredTransferPlan& plan, const SchedulingTarget& target,
-    size_t pipelineStages) {
-  std::vector<Cost> logicalTrafficFlits(target.processorCount * target.processorCount, 0);
-  for (const std::unique_ptr<DeferredExchangePlan>& exchange : plan.exchanges)
-    for (const ExternalTransferFamily& transfer : exchange->external) {
-      auto fragmentType = dyn_cast<ShapedType>(transfer.requirement->publicationFragmentType);
-      if (!fragmentType || !fragmentType.hasStaticShape())
-        return exchange->deferred.emitOpError("physical core placement requires a static transfer fragment");
-      auto fragmentBytes = pim::getCheckedShapedTypeSizeInBytes(
-        fragmentType, exchange->deferred, "physical core placement transfer fragment");
-      if (failed(fragmentBytes))
-        return failure();
-      Cost flits = static_cast<Cost>(*fragmentBytes) / target.transferWidthBytes
-                 + (*fragmentBytes % target.transferWidthBytes != 0);
-      for (size_t index = 0; index < transfer.sourceCores.size(); ++index) {
-        size_t sourceLogicalProcessor = static_cast<size_t>(transfer.sourceCores.valueAt(index));
-        size_t targetLogicalProcessor = static_cast<size_t>(transfer.targetCores.valueAt(index));
-        Cost& traffic = logicalTrafficFlits[sourceLogicalProcessor * target.processorCount + targetLogicalProcessor];
-        traffic = checkedAdd(traffic, flits);
-      }
-    }
-
-  std::vector<size_t> placementGroups;
-  if (pipelineStages > 1) {
-    if (plan.processorStages.size() != target.processorCount)
-      return failure();
-    placementGroups = plan.processorStages;
-  }
-  std::vector<size_t> physicalCoreForLogicalProcessor =
-    mapLogicalProcessorsToPhysicalCores(
-      logicalTrafficFlits, target, placementGroups);
-  auto getPhysicalCore = [&](int64_t logicalProcessor) {
-    assert(logicalProcessor >= 0 && static_cast<size_t>(logicalProcessor) < physicalCoreForLogicalProcessor.size()
-           && "logical processor is outside the scheduling target");
-    return static_cast<int64_t>(physicalCoreForLogicalProcessor[logicalProcessor]);
-  };
-  auto remap = [&](StaticIntSequence& logicalProcessors) {
-    SmallVector<int64_t> physicalCores;
-    physicalCores.reserve(logicalProcessors.size());
-    for (size_t index = 0; index < logicalProcessors.size(); ++index)
-      physicalCores.push_back(getPhysicalCore(logicalProcessors.valueAt(index)));
-    logicalProcessors = StaticIntSequence::fromValues(physicalCores);
-  };
-
-  for (ScheduledInfo& scheduled : plan.scheduled) {
-    for (int64_t& logicalProcessor : scheduled.cores)
-      logicalProcessor = getPhysicalCore(logicalProcessor);
-    if (isa<SpatScheduledCompute>(scheduled.op)) {
-      scheduled.op->setAttr(
-        kCoreIdAttrName, IntegerAttr::get(IntegerType::get(scheduled.op->getContext(), 32), scheduled.cores.front()));
-    }
-    else {
-      SmallVector<int32_t> physicalCores;
-      physicalCores.reserve(scheduled.cores.size());
-      for (int64_t physicalCore : scheduled.cores)
-        physicalCores.push_back(static_cast<int32_t>(physicalCore));
-      scheduled.op->setAttr(kCoreIdsAttrName, DenseI32ArrayAttr::get(scheduled.op->getContext(), physicalCores));
-    }
-  }
-  for (const std::unique_ptr<ProducedValue>& produced : plan.producedStorage)
-    produced->core = getPhysicalCore(produced->core);
-  for (const std::unique_ptr<DeferredExchangePlan>& exchange : plan.exchanges)
-    for (ExternalTransferFamily& transfer : exchange->external) {
-      remap(transfer.sourceCores);
-      remap(transfer.targetCores);
-    }
-  return success();
-}
-
 static LogicalResult replaceFinalGraphPublications(func::FuncOp funcOp, DeferredTransferPlan& plan) {
   for (Operation& op : funcOp.getOps()) {
     if (!isa<SpatGraphCompute, SpatGraphComputeBatch>(op))
@@ -227,9 +156,6 @@ LogicalResult realizeDeferredCommunication(func::FuncOp funcOp,
     funcOp, materialization, pipelineStages, target.processorCount);
   if (failed(transfers))
     return funcOp.emitOpError("phase 2 failed to build symbolic transfer families");
-  if (failed(placeLogicalProcessorsOnPhysicalCores(
-        *transfers, target, pipelineStages)))
-    return failure();
   if (transfers->pipelineHostBufferBytes != 0) {
     auto bytes = pim::checkedCast<int64_t>(
       transfers->pipelineHostBufferBytes, funcOp,

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

@@ -285,7 +285,7 @@ FailureOr<ScheduledCommunicationPlan> scheduleDeferredCommunication(func::FuncOp
     enqueue(stream);
   auto advance = [&] {
     bool changed = false;
-    while (!advanceable.empty()) {
+    while (!advanceable.empty() && ready.empty()) {
       unsigned stream = advanceable.front();
       advanceable.pop();
       StreamProgress& progress = streams[stream];
@@ -305,6 +305,8 @@ FailureOr<ScheduledCommunicationPlan> scheduleDeferredCommunication(func::FuncOp
   ScheduledCommunicationPlan result;
   unsigned finishedGroups = 0;
   while (finishedGroups != groups.size()) {
+    while (!ready.empty() && groups[ready.top()].scheduled)
+      ready.pop();
     bool progressed = advance();
     std::optional<unsigned> chosen;
     unsigned bestExtension = 0;

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/Passes/Transforms/MergeComputeNodes/ScheduledComputeReport.cpp

@@ -111,8 +111,16 @@ void dumpScheduledComputeReport(ModuleOp moduleOp, func::FuncOp funcOp, const Me
      << "  materialized homogeneous runs: " << materializedRuns << "\n"
      << "  largest run: " << largestRun << "\n"
      << "  instances compacted: " << instancesCompacted << "\n"
-     << "  compatible runs rejected: 0\n\n"
-     << "Materialized scheduled ops\n";
+     << "  compatible runs rejected: 0\n";
+  llvm::MapVector<size_t, SmallVector<size_t>> stageCores;
+  for (auto [core, stage] : llvm::enumerate(schedule.processorStages))
+    stageCores[stage].push_back(core);
+  for (const auto &[stage, cores] : stageCores) {
+    os << "  pipeline stage " << stage << " cores=";
+    printIndexedList(os, ArrayRef<size_t>(cores));
+    os << "\n";
+  }
+  os << "\nMaterialized scheduled ops\n";
 
   for (const ScheduledMaterializationRecord &record : records) {
     bool batch = isa<SpatScheduledComputeBatch>(record.scheduledOp);

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

@@ -10,8 +10,6 @@
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.hpp"
 #include "src/Accelerators/PIM/Passes/PIMPasses.h"
 
-#include <limits>
-
 using namespace mlir;
 
 namespace onnx_mlir {
@@ -34,38 +32,6 @@ static FailureOr<func::FuncOp> requireEntry(ModuleOp moduleOp) {
   return *entry;
 }
 
-static SchedulingTarget getPipelineSchedulingTarget(
-    const SchedulingTarget& physicalTarget, size_t pipelineStages) {
-  if (pipelineStages == 1)
-    return physicalTarget;
-
-  PipelineCoreLayout layout(physicalTarget.processorCount, pipelineStages);
-  SchedulingTarget schedulingTarget = physicalTarget;
-  schedulingTarget.processorCount = layout.getLogicalProcessorCount();
-  schedulingTarget.residentWeightCapacity = checkedMultiply(
-    physicalTarget.residentWeightCapacity, pipelineStages);
-  schedulingTarget.interProcessorLatencyNs.assign(
-    schedulingTarget.processorCount * schedulingTarget.processorCount, 0);
-  Cost latencySum = 0;
-  size_t pairCount = 0;
-  for (size_t source = 0; source < schedulingTarget.processorCount; ++source)
-    for (size_t destination = 0;
-         destination < schedulingTarget.processorCount; ++destination) {
-      Cost latency = physicalTarget.getInterProcessorLatencyNs(
-        source, destination);
-      schedulingTarget.interProcessorLatencyNs[
-        source * schedulingTarget.processorCount + destination] = latency;
-      if (source != destination) {
-        latencySum = checkedAdd(latencySum, latency);
-        ++pairCount;
-      }
-    }
-  schedulingTarget.averageInterProcessorLatencyNs = pairCount == 0
-    ? 0
-    : (latencySum + pairCount - 1) / pairCount;
-  return schedulingTarget;
-}
-
 struct ScheduleAndRealizeSpatialPass final
     : PassWrapper<ScheduleAndRealizeSpatialPass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ScheduleAndRealizeSpatialPass)
@@ -92,9 +58,7 @@ struct ScheduleAndRealizeSpatialPass final
     PipelineCoreLayout pipelineLayout(target.processorCount, pipelineStages);
     if (!pipelineLayout.isValid()
         || (pipelineStages > 1
-            && target.synchronizationRegisterCount == 0)
-        || target.residentWeightCapacity
-             > std::numeric_limits<size_t>::max() / pipelineStages) {
+            && target.synchronizationRegisterCount == 0)) {
       moduleOp.emitError("ScheduleAndRealizeSpatial requires valid pipeline stages and resource counts");
       signalPassFailure();
       return;
@@ -106,21 +70,15 @@ struct ScheduleAndRealizeSpatialPass final
     }
     func::FuncOp entryFunc = *entry;
 
-    SchedulingTarget schedulingTarget = getPipelineSchedulingTarget(
-      target, pipelineStages);
     ComputeGraph scheduledGraph;
     MergeScheduleResult schedule;
     for (;;) {
-      MergeSchedulingAnalysis analysis(
-        entryFunc, schedulingTarget,
-        pipelineStages > 1 ? target.processorCount : 0);
-      scheduledGraph = analysis.getGraph();
-      schedule = std::move(analysis.getResult());
+      scheduledGraph = buildComputeGraph(entryFunc, target);
       std::string pipelineError;
       if (pipelineStages > 1) {
         FailureOr<PipelineWorkloadPreparation> preparation =
           preparePipelineWorkload(
-            scheduledGraph, schedule, pipelineStages, target, pipelineError);
+            scheduledGraph, pipelineStages, target, pipelineError);
         if (failed(preparation)) {
           moduleOp.emitError() << pipelineError;
           signalPassFailure();
@@ -129,6 +87,9 @@ struct ScheduleAndRealizeSpatialPass final
         if (*preparation == PipelineWorkloadPreparation::Changed)
           continue;
       }
+      MergeSchedulingAnalysis analysis(
+        scheduledGraph, target, entryFunc.getContext());
+      schedule = std::move(analysis.getResult());
       if (succeeded(applyPipelineScheduling(
             scheduledGraph, schedule, pipelineStages, target, pipelineError)))
         break;

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp

@@ -88,22 +88,21 @@ void verifySchedule(const ComputeGraph& graph,
 
 } // namespace
 
-MergeSchedulingAnalysis::MergeSchedulingAnalysis(mlir::Operation* op,
+MergeSchedulingAnalysis::MergeSchedulingAnalysis(const ComputeGraph& computeGraph,
                                                  const SchedulingTarget& schedulingTarget,
-                                                 size_t partitionCount)
-: entryOp(op), target(schedulingTarget), computePartitionCount(partitionCount) {
+                                                 mlir::MLIRContext* context)
+: context(context), target(schedulingTarget), graph(computeGraph) {
   result = run();
 }
 
 MergeScheduleResult MergeSchedulingAnalysis::run() {
-  graph = buildComputeGraph(entryOp, target, computePartitionCount);
   if (!verifyAcyclic(graph))
     llvm::report_fatal_error("merge scheduling: compute graph is cyclic");
 
   MergeScheduleResult schedule = runPeftScheduler(
     graph, PeftScheduleOptions {
       target,
-      entryOp->getContext()});
+      context});
   verifySchedule(graph, schedule, target);
   return schedule;
 }

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.hpp

@@ -11,17 +11,15 @@ namespace spatial {
 
 class MergeSchedulingAnalysis {
 public:
-  MergeSchedulingAnalysis(mlir::Operation* op,
+  MergeSchedulingAnalysis(const ComputeGraph& graph,
                           const SchedulingTarget& target,
-                          size_t computePartitionCount = 0);
+                          mlir::MLIRContext* context);
   MergeScheduleResult& getResult() { return result; }
-  const ComputeGraph& getGraph() const { return graph; }
 
 private:
-  mlir::Operation* entryOp = nullptr;
+  mlir::MLIRContext* context = nullptr;
   const SchedulingTarget& target;
-  size_t computePartitionCount = 0;
-  ComputeGraph graph;
+  const ComputeGraph& graph;
   MergeScheduleResult result;
 
   MergeScheduleResult run();

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp

@@ -5,7 +5,6 @@
 #include "llvm/Support/FormatVariadic.h"
 
 #include <limits>
-#include <numeric>
 #include <optional>
 #include <queue>
 #include <tuple>
@@ -42,11 +41,6 @@ struct TopologyModel {
   }
 };
 
-Time getAverageTransferTime(const TransferCost& transferCost, const SchedulingTarget& target) {
-  return checkedAdd(transferCost.fixed,
-                    checkedMultiply(transferCost.networkFlits, target.averageInterProcessorLatencyNs));
-}
-
 std::vector<std::vector<size_t>> buildReverseLevels(const ComputeGraph& graph) {
   std::vector<size_t> remainingSuccessors(graph.nodes.size(), 0);
   std::queue<size_t> readySinks;
@@ -83,24 +77,6 @@ std::vector<std::vector<size_t>> buildReverseLevels(const ComputeGraph& graph) {
   return reverseLevels;
 }
 
-void verifyOctTableSize(size_t nodeCount, size_t processorCount) {
-  constexpr size_t kMaxOctTableBytes = 1ull << 35;
-  if (nodeCount == 0 || processorCount == 0)
-    return;
-  if (processorCount > std::numeric_limits<size_t>::max() / sizeof(Time))
-    llvm::report_fatal_error("PEFT scheduler: OCT table size overflow");
-  size_t rowBytes = processorCount * sizeof(Time);
-  if (nodeCount > std::numeric_limits<size_t>::max() / rowBytes)
-    llvm::report_fatal_error("PEFT scheduler: OCT table size overflow");
-  size_t totalBytes = nodeCount * rowBytes;
-  if (totalBytes > kMaxOctTableBytes) {
-    std::string message = llvm::formatv("PEFT scheduler: OCT table would require {0} MiB, exceeding the 1024 MiB guard",
-                                        totalBytes / (1024 * 1024))
-                            .str();
-    llvm::report_fatal_error(llvm::StringRef(message));
-  }
-}
-
 bool hasHighResidentWeightPressure(const ComputeGraph& graph, size_t processorCount, size_t residentWeightCapacity) {
   if (residentWeightCapacity > std::numeric_limits<size_t>::max() / processorCount)
     return false;
@@ -243,79 +219,6 @@ FailureOr<LanePublicationSignatures> buildLanePublicationSignatures(SpatComputeB
 
 } // namespace
 
-std::vector<size_t> mapLogicalProcessorsToPhysicalCores(ArrayRef<Cost> logicalTrafficFlits,
-                                                        const SchedulingTarget& target,
-                                                        ArrayRef<size_t> placementGroups) {
-  const size_t processorCount = target.processorCount;
-  assert(logicalTrafficFlits.size() == processorCount * processorCount
-         && "logical traffic matrix must cover every processor pair");
-  assert((placementGroups.empty() || placementGroups.size() == processorCount)
-         && "physical placement groups must cover every processor");
-  std::vector<size_t> physicalCoreForLogicalProcessor(processorCount);
-  std::iota(physicalCoreForLogicalProcessor.begin(), physicalCoreForLogicalProcessor.end(), 0);
-
-  auto transferCost = [&](size_t sourceLogicalProcessor,
-                          size_t targetLogicalProcessor,
-                          size_t sourcePhysicalCore,
-                          size_t targetPhysicalCore) {
-    Cost traffic = logicalTrafficFlits[sourceLogicalProcessor * processorCount + targetLogicalProcessor];
-    return checkedMultiply(traffic, target.getInterProcessorLatencyNs(sourcePhysicalCore, targetPhysicalCore));
-  };
-
-  for (size_t logicalProcessor = 0; logicalProcessor < processorCount; ++logicalProcessor) {
-    size_t bestPeerLogicalProcessor = logicalProcessor;
-    Cost bestSaving = 0;
-    for (size_t peerLogicalProcessor = 0; peerLogicalProcessor < processorCount; ++peerLogicalProcessor) {
-      if (peerLogicalProcessor == logicalProcessor)
-        continue;
-      if (!placementGroups.empty()
-          && placementGroups[peerLogicalProcessor]
-               != placementGroups[logicalProcessor])
-        continue;
-      size_t physicalCore = physicalCoreForLogicalProcessor[logicalProcessor];
-      size_t peerPhysicalCore = physicalCoreForLogicalProcessor[peerLogicalProcessor];
-      Cost currentCost = 0;
-      Cost swappedCost = 0;
-      for (size_t otherLogicalProcessor = 0; otherLogicalProcessor < processorCount; ++otherLogicalProcessor) {
-        if (otherLogicalProcessor == logicalProcessor || otherLogicalProcessor == peerLogicalProcessor)
-          continue;
-        size_t otherPhysicalCore = physicalCoreForLogicalProcessor[otherLogicalProcessor];
-        currentCost = checkedAdd(
-          currentCost, transferCost(logicalProcessor, otherLogicalProcessor, physicalCore, otherPhysicalCore));
-        currentCost = checkedAdd(
-          currentCost, transferCost(otherLogicalProcessor, logicalProcessor, otherPhysicalCore, physicalCore));
-        currentCost = checkedAdd(
-          currentCost, transferCost(peerLogicalProcessor, otherLogicalProcessor, peerPhysicalCore, otherPhysicalCore));
-        currentCost = checkedAdd(
-          currentCost, transferCost(otherLogicalProcessor, peerLogicalProcessor, otherPhysicalCore, peerPhysicalCore));
-        swappedCost = checkedAdd(
-          swappedCost, transferCost(logicalProcessor, otherLogicalProcessor, peerPhysicalCore, otherPhysicalCore));
-        swappedCost = checkedAdd(
-          swappedCost, transferCost(otherLogicalProcessor, logicalProcessor, otherPhysicalCore, peerPhysicalCore));
-        swappedCost = checkedAdd(
-          swappedCost, transferCost(peerLogicalProcessor, otherLogicalProcessor, physicalCore, otherPhysicalCore));
-        swappedCost = checkedAdd(
-          swappedCost, transferCost(otherLogicalProcessor, peerLogicalProcessor, otherPhysicalCore, physicalCore));
-      }
-      currentCost =
-        checkedAdd(currentCost, transferCost(logicalProcessor, peerLogicalProcessor, physicalCore, peerPhysicalCore));
-      currentCost =
-        checkedAdd(currentCost, transferCost(peerLogicalProcessor, logicalProcessor, peerPhysicalCore, physicalCore));
-      swappedCost =
-        checkedAdd(swappedCost, transferCost(logicalProcessor, peerLogicalProcessor, peerPhysicalCore, physicalCore));
-      swappedCost =
-        checkedAdd(swappedCost, transferCost(peerLogicalProcessor, logicalProcessor, physicalCore, peerPhysicalCore));
-      if (currentCost > swappedCost && currentCost - swappedCost > bestSaving) {
-        bestSaving = currentCost - swappedCost;
-        bestPeerLogicalProcessor = peerLogicalProcessor;
-      }
-    }
-    std::swap(physicalCoreForLogicalProcessor[logicalProcessor],
-              physicalCoreForLogicalProcessor[bestPeerLogicalProcessor]);
-  }
-  return physicalCoreForLogicalProcessor;
-}
-
 MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftScheduleOptions& options) {
   const size_t nodeCount = graph.nodes.size();
   const size_t processorCount = options.target.processorCount;
@@ -327,36 +230,17 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
   std::vector<ResidentWeightSet> capacityReservations = planResidentWeightReservations(
     graph, processorCount, options.target.residentWeightCapacity, topology, preferWeightReuse);
 
-  verifyOctTableSize(nodeCount, processorCount);
   std::vector<std::vector<size_t>> reverseLevels = buildReverseLevels(graph);
 
-  // MOCK: Replace this with your actual heterogeneous cost lookup.
-  // If graph.nodes[task] is modified to hold a vector of costs per processor, access it here.
-  auto getComputeCost = [&](size_t task, size_t processor) -> Time { return graph.nodes[task].cost; };
-  std::vector<Time> oct(nodeCount * processorCount, 0);
-  std::vector<Time> minOctPlusComp(nodeCount, 0);
-
-  // 1. O(P(E+V)) Heterogeneous OCT Calculation
+  // Compute costs are identical on every core. The optimistic successor can
+  // stay on the same core for zero transfer cost, so every physical OCT column
+  // is identical even with nonuniform links. Store that exact lower bound once.
+  std::vector<Time> oct(nodeCount, 0);
   for (const std::vector<size_t>& levelNodes : reverseLevels) {
     auto computeNodeOct = [&](size_t levelIndex) {
       size_t task = levelNodes[levelIndex];
-      std::vector<Time> maxVals(processorCount, 0);
-
-      for (const auto& [succ, comm] : graph.successors[task]) {
-        Time valDifferentCpu = addOrMax(minOctPlusComp[succ], getAverageTransferTime(comm, options.target));
-        for (size_t processor = 0; processor < processorCount; ++processor) {
-          Time valSameCpu = addOrMax(oct[succ * processorCount + processor], getComputeCost(succ, processor));
-          Time bestSucc = std::min(valSameCpu, valDifferentCpu);
-          maxVals[processor] = std::max(maxVals[processor], bestSucc);
-        }
-      }
-
-      Time minForPreds = std::numeric_limits<Time>::max();
-      for (size_t processor = 0; processor < processorCount; ++processor) {
-        oct[task * processorCount + processor] = maxVals[processor];
-        minForPreds = std::min(minForPreds, addOrMax(maxVals[processor], getComputeCost(task, processor)));
-      }
-      minOctPlusComp[task] = minForPreds == std::numeric_limits<Time>::max() ? 0 : minForPreds;
+      for (const auto& [succ, comm] : graph.successors[task])
+        oct[task] = std::max(oct[task], addOrMax(oct[succ], graph.nodes[succ].cost));
     };
 
     if (options.context != nullptr)
@@ -373,9 +257,7 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
   };
   std::vector<RankEntry> ranks(nodeCount);
   auto computeRank = [&](size_t node) {
-    long double rank = 0.0L;
-    for (size_t processor = 0; processor < processorCount; ++processor)
-      rank += static_cast<long double>(oct[node * processorCount + processor]);
+    long double rank = static_cast<long double>(oct[node]);
     ranks[node] = {rank, node, graph.nodes[node].originalOrder};
   };
 
@@ -405,7 +287,6 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
     }
 
     std::vector<char> scheduled(nodeCount, false);
-    std::vector<ResidentWeightSet> reservations = capacityReservations;
     std::vector<ResidentWeightSet> processorResidentWeights(processorCount);
     std::vector<ScheduledTask> schedules(nodeCount);
     std::vector<std::vector<size_t>> tasksByProcessor(processorCount);
@@ -435,9 +316,9 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
         unsigned int overlapCount =
           countResidentWeightOverlap(processorResidentWeights[processor], graph.nodes[task].residentWeights);
         size_t residentWeightUnion =
-          getResidentWeightUnionSize(reservations[processor], graph.nodes[task].residentWeights);
+          getResidentWeightUnionSize(capacityReservations[processor], graph.nodes[task].residentWeights);
         smallestResidentWeightUnion = std::min(smallestResidentWeightUnion, residentWeightUnion);
-        if (!graph.nodes[task].residentWeights.empty() && residentWeightUnion > options.target.residentWeightCapacity) {
+        if (!graph.nodes[task].residentWeights.empty() && residentWeightUnion > capacityReservations[processor].size()) {
           residentWeightRejected = true;
           continue;
         }
@@ -448,7 +329,7 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
           dataReady = std::max(dataReady, addOrMax(predSchedule.endTime, commPenalty));
         }
 
-        Time computeCost = getComputeCost(task, processor);
+        Time computeCost = graph.nodes[task].cost;
         Time est = dataReady;
         Time currentEnd = 0;
         bool foundGap = false;
@@ -467,7 +348,7 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
           est = std::max(currentEnd, dataReady);
 
         Time eft = addOrMax(est, computeCost);
-        Time oeft = addOrMax(eft, oct[task * processorCount + processor]);
+        Time oeft = addOrMax(eft, oct[task]);
         size_t centerDistance = topology.getCenterDistance(processor);
         size_t taskCount = tasksByProcessor[processor].size();
         bool betterResidentWeightChoice =
@@ -519,7 +400,7 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
           const ComputeInstance& instance = graph.nodes[task].instance;
           std::string message =
             llvm::formatv("PEFT scheduler: no valid processor for task {0} (lanes {1}..{2}, {3} distinct weights); "
-                          "smallest processor union is {4}, exceeding resident-weight capacity {5}",
+                          "no physical reservation contains the task weights (smallest union {4}, core capacity {5})",
                           graph.nodes[task].originalOrder,
                           instance.laneStart,
                           instance.laneStart + instance.laneCount,
@@ -539,7 +420,6 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
       schedules[task] = {bestProcessor, bestEst, bestEft};
       scheduled[task] = true;
       ++scheduledCount;
-      insertResidentWeights(reservations[bestProcessor], graph.nodes[task].residentWeights);
       insertResidentWeights(processorResidentWeights[bestProcessor], graph.nodes[task].residentWeights);
 
       auto& timeline = timelineByProcessor[bestProcessor];

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.hpp

@@ -21,16 +21,11 @@ inline Time getPeftTransferTime(const TransferCost& transferCost,
   if (sourceProcessor == targetProcessor)
     return 0;
   return checkedAdd(transferCost.fixed,
-                    checkedMultiply(transferCost.networkFlits, target.averageInterProcessorLatencyNs));
+                    checkedMultiply(transferCost.networkFlits, target.getInterProcessorLatencyNs(sourceProcessor, targetProcessor)));
 }
 
-// PEFT assigns logical processors. Physical core IDs are chosen only after
-// materialization exposes the exact transfer traffic.
+// PEFT assigns physical cores using the injected target topology.
 MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftScheduleOptions& options);
 
-std::vector<size_t> mapLogicalProcessorsToPhysicalCores(llvm::ArrayRef<Cost> logicalTrafficFlits,
-                                                        const SchedulingTarget& target,
-                                                        llvm::ArrayRef<size_t> placementGroups = {});
-
 } // namespace spatial
 } // namespace onnx_mlir

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.cpp

@@ -49,21 +49,8 @@ struct PipelineTaskModel {
   std::vector<TaskList> predecessors;
 };
 
-static bool fits(const ComputeGraph &graph,
-                 const TaskCosts &taskCosts,
-                 const TaskList &tasks,
-                 Cost maximumCost,
-                 size_t residentWeightCapacity,
-                 size_t maximumPacks);
-static Cost findMaximumPackCost(const ComputeGraph &graph,
-                                const TaskCosts &taskCosts,
-                                const TaskList &tasks,
-                                size_t residentWeightCapacity,
-                                size_t pipelineStages);
-
 static BlueprintAssemblyInfo getBlueprintAssemblyInfo(
     const ComputeGraph &graph, const ComputeGraphNode &node,
-    size_t consumerCpu, const MergeScheduleResult &schedule,
     const SchedulingTarget &target) {
   BlueprintAssemblyInfo info;
   if (!node.instance.op)
@@ -112,12 +99,7 @@ static BlueprintAssemblyInfo getBlueprintAssemblyInfo(
           elements, static_cast<Cost>(fragmentSizes[fragment * rank + dimension]));
       Cost bytes = (checkedMultiply(elements, target.computeBitwidth) + 7) / 8;
       TransferCost transferCost = getTransferCostFromBytes(bytes, target);
-      size_t producerCpu = schedule.computeToCpuMap.lookup(producerInstance);
-      Cost transferCycles = (producerCpu == consumerCpu
-                               ? transferCost.fixed
-                               : getPeftTransferTime(
-                                   transferCost, producerCpu, consumerCpu, target))
-                          / target.processorPeriodNs;
+      Cost transferCycles = transferCost.fixed / target.processorPeriodNs;
       info.schedulingCost = checkedAdd(
         info.schedulingCost, checkedAdd(elements, transferCycles));
       info.work = checkedAdd(info.work, elements);
@@ -127,8 +109,7 @@ static BlueprintAssemblyInfo getBlueprintAssemblyInfo(
 }
 
 static PipelineTaskModel getPipelineTaskModel(
-    const ComputeGraph &graph, const MergeScheduleResult &schedule,
-    const SchedulingTarget &target) {
+    const ComputeGraph &graph, const SchedulingTarget &target) {
   PipelineTaskModel model;
   model.schedulingCosts.reserve(graph.nodes.size());
   model.assemblyCosts.reserve(graph.nodes.size());
@@ -143,9 +124,8 @@ static PipelineTaskModel getPipelineTaskModel(
     }
 
   for (size_t task = 0; task < graph.nodes.size(); ++task) {
-    size_t cpu = schedule.computeToCpuMap.lookup(graph.nodes[task].instance);
     BlueprintAssemblyInfo blueprint = getBlueprintAssemblyInfo(
-      graph, graph.nodes[task], cpu, schedule, target);
+      graph, graph.nodes[task], target);
     model.schedulingCosts[task] = checkedAdd(
       model.schedulingCosts[task], blueprint.schedulingCost);
     model.assemblyCosts[task] = checkedAdd(
@@ -156,25 +136,11 @@ static PipelineTaskModel getPipelineTaskModel(
     for (const auto &[predecessor, transferCost] : graph.predecessors[task]) {
       if (!llvm::is_contained(model.predecessors[task], predecessor))
         model.predecessors[task].push_back(predecessor);
-      size_t predecessorCpu =
-        schedule.computeToCpuMap.lookup(graph.nodes[predecessor].instance);
-      if (predecessorCpu == cpu)
-        continue;
-      model.schedulingCosts[task] = checkedAdd(
-        model.schedulingCosts[task], getPeftTransferTime(
-                                       transferCost, predecessorCpu, cpu, target)
-                                     / target.processorPeriodNs);
     }
   }
   return model;
 }
 
-static TaskCosts getPipelineTaskCosts(const ComputeGraph &graph,
-                                      const MergeScheduleResult &schedule,
-                                      const SchedulingTarget &target) {
-  return getPipelineTaskModel(graph, schedule, target).schedulingCosts;
-}
-
 static void copyPipelineComputeAttrs(SpatGraphComputeBatch source,
                                      SpatGraphComputeBatch destination) {
   for (NamedAttribute attr : source->getAttrs()) {
@@ -609,123 +575,23 @@ static LogicalResult splitBatchCompute(SpatGraphComputeBatch batch,
 }
 
 static FailureOr<PipelineWorkloadPreparation> preparePipelineWorkloadImpl(
-    const ComputeGraph &graph, const MergeScheduleResult &schedule,
-    size_t pipelineStages, const SchedulingTarget &physicalTarget,
-    std::string &error) {
-  size_t groupSize = schedule.processorCount;
-  std::vector<TaskList> tasksByCpu(groupSize);
-  for (size_t task = 0; task < graph.nodes.size(); ++task) {
-    auto cpu = schedule.computeToCpuMap.find(graph.nodes[task].instance);
-    if (cpu == schedule.computeToCpuMap.end() || cpu->second >= groupSize) {
-      error = "pipeline split received an incomplete PEFT schedule for task "
-            + std::to_string(task) + " (cpu "
-            + (cpu == schedule.computeToCpuMap.end()
-                 ? std::string("missing")
-                 : std::to_string(cpu->second))
-            + ", schedule processors " + std::to_string(groupSize) + ")";
-      return failure();
-    }
-    tasksByCpu[cpu->second].push_back(task);
-  }
-  for (TaskList &tasks : tasksByCpu)
-    llvm::sort(tasks, [&](size_t lhs, size_t rhs) {
-      return schedule.computeToCpuSlotMap.lookup(graph.nodes[lhs].instance)
-           < schedule.computeToCpuSlotMap.lookup(graph.nodes[rhs].instance);
-    });
-  TaskCosts taskCosts = getPipelineTaskCosts(graph, schedule, physicalTarget);
-
-  for (const TaskList &tasks : tasksByCpu) {
-    if (tasks.empty())
+    const ComputeGraph &graph, size_t pipelineStages,
+    const SchedulingTarget &physicalTarget, std::string &error) {
+  for (const ComputeGraphNode &node : graph.nodes) {
+    if (node.residentWeights.size() <= physicalTarget.residentWeightCapacity)
       continue;
-    Cost maximumCost = findMaximumPackCost(
-      graph, taskCosts, tasks, physicalTarget.residentWeightCapacity, pipelineStages);
-    if (fits(graph, taskCosts, tasks, maximumCost,
-             physicalTarget.residentWeightCapacity, pipelineStages))
-      continue;
-
-    SmallVector<size_t, 8> candidates(tasks.begin(), tasks.end());
-    llvm::stable_sort(candidates, [&](size_t lhs, size_t rhs) {
-      if (graph.nodes[lhs].residentWeights.size()
-          != graph.nodes[rhs].residentWeights.size())
-        return graph.nodes[lhs].residentWeights.size()
-             > graph.nodes[rhs].residentWeights.size();
-      return taskCosts[lhs] > taskCosts[rhs];
-    });
-    std::string candidateError;
-    for (size_t task : candidates) {
-      auto batch = dyn_cast<SpatGraphComputeBatch>(
-        graph.nodes[task].instance.op);
-      if (!batch || batch->hasAttr("pipeline.split"))
-        continue;
-      std::string currentError;
-      if (succeeded(splitBatchCompute(
-            batch, pipelineStages, physicalTarget, currentError)))
-        return PipelineWorkloadPreparation::Changed;
-      if (!currentError.empty())
-        candidateError = currentError;
-    }
-    error = candidateError.empty()
-      ? "pipeline scheduling could not find a splittable graph compute"
-      : candidateError;
+    auto batch = dyn_cast_or_null<SpatGraphComputeBatch>(node.instance.op);
+    if (batch && succeeded(splitBatchCompute(
+                   batch, pipelineStages, physicalTarget, error)))
+      return PipelineWorkloadPreparation::Changed;
+    if (error.empty())
+      error = "pipeline scheduling cannot split a compute instance to fit one "
+              "physical core's crossbars";
     return failure();
   }
-
   return PipelineWorkloadPreparation::Ready;
 }
 
-bool fits(const ComputeGraph& graph,
-          const TaskCosts& taskCosts,
-          const TaskList& tasks,
-          Cost maximumCost,
-          size_t residentWeightCapacity,
-          size_t maximumPacks) {
-  size_t packs = 1;
-  Cost cost = 0;
-  ResidentWeightSet weights;
-  bool packEmpty = true;
-  for (size_t task : tasks) {
-    const ComputeGraphNode& node = graph.nodes[task];
-    Cost taskCost = taskCosts[task];
-    if (node.residentWeights.size() > residentWeightCapacity)
-      return false;
-    bool startsNewPack = !packEmpty
-                      && (cost > maximumCost - taskCost
-                          || getResidentWeightUnionSize(weights, node.residentWeights) > residentWeightCapacity);
-    if (startsNewPack) {
-      if (++packs > maximumPacks)
-        return false;
-      cost = 0;
-      weights.clear();
-      packEmpty = true;
-    }
-    cost = checkedAdd(cost, taskCost);
-    insertResidentWeights(weights, node.residentWeights);
-    packEmpty = false;
-  }
-  return true;
-}
-
-Cost findMaximumPackCost(const ComputeGraph& graph,
-                         const TaskCosts& taskCosts,
-                         const TaskList& tasks,
-                         size_t residentWeightCapacity,
-                         size_t pipelineStages) {
-  Cost low = 0;
-  Cost high = 0;
-  for (size_t task : tasks) {
-    low = std::max(low, taskCosts[task]);
-    high = checkedAdd(high, taskCosts[task]);
-  }
-  while (low < high) {
-    Cost middle = low + (high - low) / 2;
-    if (fits(graph, taskCosts, tasks, middle, residentWeightCapacity, pipelineStages))
-      high = middle;
-    else
-      low = middle + 1;
-  }
-  return low;
-}
-
 static Cost findMaximumIndexedPackCost(
     const TaskCosts &taskCosts,
     const std::vector<TaskList> &taskWeightIds, size_t weightCount,
@@ -1139,11 +1005,12 @@ static bool packPipelineStage(
     std::vector<TaskList> &tasksByCpu, const PipelineCoreLayout &layout,
     ArrayRef<size_t> topologicalPosition, size_t stage,
     size_t residentWeightCapacity, const SchedulingTarget &target,
-    std::vector<size_t> &taskCpus) {
+    std::vector<size_t> &taskCpus, std::vector<Time> &taskEndTimes,
+    bool prioritizeWeightReuse = false) {
   PipelineStageRange range = layout.getStageRange(stage);
   TaskList tasks;
   for (size_t cpu = range.begin; cpu < range.begin + range.size; ++cpu)
-    llvm::append_range(tasks, tasksByCpu[cpu]);
+    llvm::append_range(tasks, tasksByCpu[layout.getPhysicalCore(cpu)]);
   llvm::sort(tasks, [&](size_t lhs, size_t rhs) {
     return topologicalPosition[lhs] < topologicalPosition[rhs];
   });
@@ -1151,8 +1018,10 @@ static bool packPipelineStage(
   std::vector<ResidentWeightSet> weights(range.size);
   TaskCosts loads(range.size);
   TaskCosts assemblyLoads(range.size);
+  std::vector<Time> coreReady(range.size);
   for (size_t task : tasks) {
     std::optional<size_t> bestCore;
+    Time bestEndTime = 0;
     using PackScore = std::tuple<Cost, Time, Cost, Cost, size_t>;
     std::optional<PackScore> bestScore;
     for (size_t core = 0; core < range.size; ++core) {
@@ -1165,20 +1034,29 @@ static bool packPipelineStage(
         assemblyLoads[core], assemblyCosts[task]);
       Cost schedulingLoad = checkedAdd(
         loads[core], schedulingCosts[task]);
-      Time transferTime = 0;
-      size_t candidateCpu = range.begin + core;
+      Time startTime = coreReady[core];
+      size_t candidateCpu = layout.getPhysicalCore(range.begin + core);
+      // Only same-stage edges use direct sends/receives. Pipeline boundaries
+      // use host buffers and belong to different samples in steady state.
       for (const auto &[predecessor, transferCost] :
            graph.predecessors[task])
-        if (taskCpus[predecessor] < target.processorCount)
-          transferTime = checkedAdd(
-            transferTime, getPeftTransferTime(
-                            transferCost, taskCpus[predecessor],
-                            candidateCpu, target));
-      PackScore score {
-        assemblyLoad, transferTime, schedulingLoad, addedWeights, core};
+        if (taskCpus[predecessor] < target.processorCount
+            && layout.getStageForCore(taskCpus[predecessor]) == stage) {
+          Time transfer = getPeftTransferTime(
+            transferCost, taskCpus[predecessor], candidateCpu, target);
+          startTime = std::max(
+            startTime, addOrMax(taskEndTimes[predecessor], transfer));
+        }
+      Time endTime = addOrMax(startTime, schedulingCosts[task]);
+      PackScore score = prioritizeWeightReuse
+        ? PackScore {
+            addedWeights, endTime, assemblyLoad, schedulingLoad, core}
+        : PackScore {
+            endTime, assemblyLoad, schedulingLoad, addedWeights, core};
       if (!bestScore || score < *bestScore) {
         bestCore = core;
         bestScore = score;
+        bestEndTime = endTime;
       }
     }
     if (!bestCore)
@@ -1189,10 +1067,12 @@ static bool packPipelineStage(
     loads[*bestCore] = checkedAdd(loads[*bestCore], schedulingCosts[task]);
     assemblyLoads[*bestCore] = checkedAdd(
       assemblyLoads[*bestCore], assemblyCosts[task]);
-    taskCpus[task] = range.begin + *bestCore;
+    taskCpus[task] = layout.getPhysicalCore(range.begin + *bestCore);
+    taskEndTimes[task] = bestEndTime;
+    coreReady[*bestCore] = taskEndTimes[task];
   }
   for (size_t core = 0; core < range.size; ++core)
-    tasksByCpu[range.begin + core] = std::move(packed[core]);
+    tasksByCpu[layout.getPhysicalCore(range.begin + core)] = std::move(packed[core]);
   return true;
 }
 
@@ -1234,11 +1114,16 @@ static LogicalResult packPipelineStages(
 
   const TaskCosts &balanceCosts = getPipelineBalanceCosts(graph, model);
   std::vector<size_t> taskCpus(graph.nodes.size(), target.processorCount);
+  std::vector<Time> taskEndTimes(graph.nodes.size());
   for (size_t stage = 0; stage < pipelineStages; ++stage)
     if (!packPipelineStage(
           graph, model.schedulingCosts, balanceCosts, tasksByCpu,
           layout, topologicalPosition, stage, residentWeightCapacity,
-          target, taskCpus)) {
+          target, taskCpus, taskEndTimes)
+        && !packPipelineStage(
+          graph, model.schedulingCosts, balanceCosts, tasksByCpu,
+          layout, topologicalPosition, stage, residentWeightCapacity,
+          target, taskCpus, taskEndTimes, /*prioritizeWeightReuse=*/true)) {
       failedStage = stage;
       error = "pipeline scheduling cannot pack dependency-monotone stage "
             + std::to_string(stage)
@@ -1316,7 +1201,7 @@ mlir::LogicalResult assignPipelineCores(const ComputeGraph& graph,
   PipelineCoreLayout balancedLayout(
     physicalTarget.processorCount, pipelineStages);
   if (!balancedLayout.isValid()
-      || groupSize != balancedLayout.getLogicalProcessorCount()) {
+      || groupSize != physicalTarget.processorCount) {
     error = "pipeline scheduling received an incompatible physical core layout";
     return mlir::failure();
   }
@@ -1331,7 +1216,7 @@ mlir::LogicalResult assignPipelineCores(const ComputeGraph& graph,
     }
   }
   PipelineTaskModel taskModel = getPipelineTaskModel(
-    graph, schedule, physicalTarget);
+    graph, physicalTarget);
   FailureOr<PipelineStageAssignment> assignment = assignPipelineStages(
     graph, taskModel, balancedLayout,
     physicalTarget.residentWeightCapacity, error);
@@ -1342,13 +1227,13 @@ mlir::LogicalResult assignPipelineCores(const ComputeGraph& graph,
   std::string packingError;
   bool packed = false;
   for (size_t attempt = 0; attempt < physicalTarget.processorCount; ++attempt) {
-    PipelineCoreLayout candidateLayout(assignment->stageSizes);
+    PipelineCoreLayout candidateLayout(assignment->stageSizes, physicalTarget);
     for (TaskList &tasks : tasksByPhysicalCpu)
       tasks.clear();
     for (size_t task = 0; task < graph.nodes.size(); ++task) {
       PipelineStageRange range =
         candidateLayout.getStageRange(assignment->taskStages[task]);
-      tasksByPhysicalCpu[range.begin].push_back(task);
+      tasksByPhysicalCpu[candidateLayout.getPhysicalCore(range.begin)].push_back(task);
     }
     size_t failedStage = 0;
     if (succeeded(packPipelineStages(
@@ -1379,7 +1264,7 @@ mlir::LogicalResult assignPipelineCores(const ComputeGraph& graph,
     error = packingError;
     return failure();
   }
-  PipelineCoreLayout pipelineLayout(assignment->stageSizes);
+  PipelineCoreLayout pipelineLayout(assignment->stageSizes, physicalTarget);
   if (failed(verifyPipelineStageAssignment(
         graph, taskModel, tasksByPhysicalCpu, pipelineLayout, error)))
     return failure();
@@ -1389,8 +1274,8 @@ mlir::LogicalResult assignPipelineCores(const ComputeGraph& graph,
   schedule.processorStages.resize(physicalTarget.processorCount);
   for (size_t stage = 0; stage < pipelineLayout.getStageCount(); ++stage) {
     PipelineStageRange range = pipelineLayout.getStageRange(stage);
-    std::fill_n(
-      schedule.processorStages.begin() + range.begin, range.size, stage);
+    for (size_t index = range.begin; index < range.begin + range.size; ++index)
+      schedule.processorStages[pipelineLayout.getPhysicalCore(index)] = stage;
   }
   schedule.computeToCpuSlotMap.clear();
   schedule.computeToAestMap.clear();
@@ -1464,6 +1349,48 @@ mlir::LogicalResult assignPipelineCores(const ComputeGraph& graph,
 
 } // namespace
 
+PipelineCoreLayout::PipelineCoreLayout(ArrayRef<size_t> sizes,
+                                       const SchedulingTarget &target)
+    : PipelineCoreLayout(sizes) {
+  assert(processorCount == target.processorCount);
+  std::vector<bool> assigned(processorCount);
+  std::vector<Cost> remainingDistance(processorCount);
+  auto distanceBetween = [&](size_t source, size_t destination) {
+    return checkedAdd(target.getInterProcessorLatencyNs(source, destination),
+                      target.getInterProcessorLatencyNs(destination, source));
+  };
+  for (size_t core = 0; core < processorCount; ++core)
+    for (size_t other = 0; other < processorCount; ++other)
+      remainingDistance[core] = checkedAdd(
+        remainingDistance[core], distanceBetween(core, other));
+  physicalCores.reserve(processorCount);
+  physicalCoreStages.resize(processorCount);
+  // ponytail: greedy O(P^2) clustering; use a graph partitioner only if measured
+  // intra-stage communication warrants a more expensive placement search.
+  for (auto [stage, size] : llvm::enumerate(stageSizes)) {
+    std::vector<Cost> distance(processorCount);
+    for (size_t index = 0; index < size; ++index) {
+      size_t best = processorCount;
+      for (size_t core = 0; core < processorCount; ++core)
+        if (!assigned[core]
+            && (best == processorCount
+                || (index == 0 ? remainingDistance[core] > remainingDistance[best]
+                               : distance[core] < distance[best])))
+          best = core;
+      assert(best != processorCount);
+      assigned[best] = true;
+      physicalCores.push_back(best);
+      physicalCoreStages[best] = stage;
+      for (size_t core = 0; core < processorCount; ++core)
+        if (!assigned[core]) {
+          Cost link = distanceBetween(best, core);
+          distance[core] = checkedAdd(distance[core], link);
+          remainingDistance[core] -= link;
+        }
+    }
+  }
+}
+
 mlir::LogicalResult applyPipelineScheduling(const ComputeGraph& graph,
                                             MergeScheduleResult& schedule,
                                             size_t pipelineStages,
@@ -1475,7 +1402,7 @@ mlir::LogicalResult applyPipelineScheduling(const ComputeGraph& graph,
     physicalTarget.processorCount, pipelineStages);
   if (!pipelineLayout.isValid() || schedule.processorCount == 0
       || schedule.processorCount
-           != pipelineLayout.getLogicalProcessorCount()) {
+           != physicalTarget.processorCount) {
     error = "pipeline scheduling requires a valid balanced physical core layout";
     return mlir::failure();
   }
@@ -1483,11 +1410,10 @@ mlir::LogicalResult applyPipelineScheduling(const ComputeGraph& graph,
 }
 
 mlir::FailureOr<PipelineWorkloadPreparation> preparePipelineWorkload(
-    const ComputeGraph &graph, const MergeScheduleResult &schedule,
-    size_t pipelineStages, const SchedulingTarget &physicalTarget,
-    std::string &error) {
+    const ComputeGraph &graph, size_t pipelineStages,
+    const SchedulingTarget &physicalTarget, std::string &error) {
   return preparePipelineWorkloadImpl(
-    graph, schedule, pipelineStages, physicalTarget, error);
+    graph, pipelineStages, physicalTarget, error);
 }
 
 } // namespace onnx_mlir::spatial

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.hpp

@@ -40,23 +40,24 @@ public:
           stageSizes.begin(), stageSizes.end(), size_t {0})),
         stageSizes(stageSizes.begin(), stageSizes.end()) {}
 
+  PipelineCoreLayout(llvm::ArrayRef<size_t> stageSizes,
+                     const SchedulingTarget &target);
+
   bool isValid() const {
     return !stageSizes.empty()
         && llvm::none_of(stageSizes, [](size_t size) { return size == 0; });
   }
 
-  size_t getLogicalProcessorCount() const {
-    return isValid()
-      ? *std::min_element(stageSizes.begin(), stageSizes.end())
-      : 0;
-  }
-
   size_t getStageCount() const { return stageSizes.size(); }
 
   size_t getProcessorCount() const { return processorCount; }
 
   llvm::ArrayRef<size_t> getStageSizes() const { return stageSizes; }
 
+  size_t getPhysicalCore(size_t position) const {
+    return physicalCores.empty() ? position : physicalCores[position];
+  }
+
   PipelineStageRange getStageRange(size_t stage) const {
     return {std::accumulate(
               stageSizes.begin(), stageSizes.begin() + stage, size_t {0}),
@@ -66,6 +67,8 @@ public:
   std::optional<size_t> getStageForCore(size_t core) const {
     if (!isValid() || core >= processorCount)
       return std::nullopt;
+    if (!physicalCoreStages.empty())
+      return physicalCoreStages[core];
     size_t end = 0;
     for (auto [stage, size] : llvm::enumerate(stageSizes)) {
       end += size;
@@ -78,6 +81,8 @@ public:
 private:
   size_t processorCount;
   std::vector<size_t> stageSizes;
+  std::vector<size_t> physicalCores;
+  std::vector<size_t> physicalCoreStages;
 };
 
 mlir::LogicalResult applyPipelineScheduling(const ComputeGraph& graph,
@@ -92,8 +97,8 @@ enum class PipelineWorkloadPreparation {
 };
 
 mlir::FailureOr<PipelineWorkloadPreparation> preparePipelineWorkload(
-    const ComputeGraph& graph, const MergeScheduleResult& schedule,
-    size_t pipelineStages, const SchedulingTarget& physicalTarget,
+    const ComputeGraph& graph, size_t pipelineStages,
+    const SchedulingTarget& physicalTarget,
     std::string& error);
 
 } // namespace onnx_mlir::spatial

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)
   }];
 }
 

test/PIM/CMakeLists.txt

@@ -36,3 +36,13 @@ add_pim_unittest(SpatialSchedulingTargetTest
   LINK_LIBS PRIVATE
   OMPimCompilerUtils
 )
+
+add_test(
+  NAME PimHostReuseSynchronizationTest
+  COMMAND "${PYTHON_EXECUTABLE}"
+          "${CMAKE_SOURCE_DIR}/validation/tools/pim/pimcomp/compare/test_PIMCOMP_adversarial_memory_sync.py"
+          --compiler "$<TARGET_FILE:onnx-mlir>"
+          --simple-model "${CMAKE_SOURCE_DIR}/validation/operations/relu/after_conv/relu_after_conv.onnx"
+          --grouped-model "${CMAKE_SOURCE_DIR}/validation/operations/conv/relu_conv_store/conv_relu_conv_store.onnx"
+)
+set_tests_properties(PimHostReuseSynchronizationTest PROPERTIES LABELS pim-unittest)

test/PIM/SpatialSchedulingTargetTest.cpp

@@ -3,15 +3,156 @@
 #include <string>
 #include <vector>
 
+#include "mlir/IR/BuiltinOps.h"
+
+#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.hpp"
 
 using namespace onnx_mlir::spatial;
 
+static void testPhysicalPeft() {
+  TransferCost transfer {.fixed = 50, .networkFlits = 4};
+  SchedulingTarget line;
+  line.processorCount = 3;
+  line.interProcessorLatencyNs = {
+    0,
+    1,
+    10,
+    1,
+    0,
+    1,
+    10,
+    1,
+    0,
+  };
+  line.averageInterProcessorLatencyNs = 4;
+  assert(getPeftTransferTime(transfer, 0, 1, line) == 54);
+  assert(getPeftTransferTime(transfer, 0, 2, line) == 90);
+  SchedulingTarget reversed = line;
+  reversed.interProcessorLatencyNs = {0, 10, 1, 10, 0, 1, 1, 1, 0};
+  assert(getPeftTransferTime(transfer, 0, 1, reversed) == 90);
+  assert(getPeftTransferTime(transfer, 0, 2, reversed) == 54);
+
+  mlir::MLIRContext context;
+  mlir::OwningOpRef<mlir::ModuleOp> owner =
+    mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
+  ComputeGraph chain;
+  chain.successors.resize(3);
+  chain.predecessors.resize(3);
+  for (uint32_t task = 0; task < 3; ++task) {
+    ComputeInstance instance {owner->getOperation(), task, 1};
+    ResidentWeight weight;
+    weight.opaqueLane = task;
+    chain.nodes.push_back({instance, 10, {weight}, task});
+    if (task != 0) {
+      chain.successors[task - 1].push_back({task, transfer});
+      chain.predecessors[task].push_back({task - 1, transfer});
+    }
+  }
+  line.residentWeightCapacity = reversed.residentWeightCapacity = 3;
+  for (const SchedulingTarget &target : {line, reversed}) {
+    auto scheduled = runPeftScheduler(chain, {target, &context});
+    size_t firstCore = scheduled.computeToCpuMap.lookup(chain.nodes[0].instance);
+    assert(firstCore == (target.interProcessorLatencyNs == line.interProcessorLatencyNs ? 1 : 2));
+    for (size_t task = 1; task < chain.nodes.size(); ++task) {
+      auto previous = chain.nodes[task - 1].instance;
+      auto current = chain.nodes[task].instance;
+      size_t source = scheduled.computeToCpuMap.lookup(previous);
+      size_t destination = scheduled.computeToCpuMap.lookup(current);
+      assert(source != destination);
+      assert(scheduled.computeToAestMap.lookup(current)
+             == scheduled.computeToAestMap.lookup(previous) + 10
+                  + getPeftTransferTime(transfer, source, destination, target));
+    }
+  }
+}
+
+static void testPipelineInputArrival() {
+  ComputeGraph graph;
+  graph.successors.resize(4);
+  graph.predecessors.resize(4);
+  MergeScheduleResult schedule;
+  schedule.processorCount = 4;
+  for (uint32_t task = 0; task < 4; ++task) {
+    ComputeInstance instance {nullptr, task, 1};
+    ResidentWeight weight;
+    weight.opaqueLane = task;
+    graph.nodes.push_back({instance, 1, {weight}, task});
+    schedule.dominanceOrderCompute.push_back(instance);
+    schedule.computeToCpuMap[instance] = task;
+    schedule.computeToCpuSlotMap[instance] = 0;
+    if (task != 0) {
+      TransferCost transfer {.fixed = 0, .networkFlits = 100};
+      graph.predecessors[task].push_back({task - 1, transfer});
+      graph.successors[task - 1].push_back({task, transfer});
+    }
+  }
+  SchedulingTarget target;
+  target.processorCount = 4;
+  target.residentWeightCapacity = 4;
+  target.interProcessorLatencyNs.assign(16, 1);
+  for (size_t core = 0; core < 4; ++core)
+    target.interProcessorLatencyNs[core * 4 + core] = 0;
+  std::string error;
+  assert(mlir::succeeded(applyPipelineScheduling(graph, schedule, 2, target, error)));
+  for (size_t first : {0, 2}) {
+    size_t source = schedule.computeToCpuMap.lookup(graph.nodes[first].instance);
+    size_t destination = schedule.computeToCpuMap.lookup(graph.nodes[first + 1].instance);
+    assert(schedule.processorStages[source] == first / 2);
+    // An idle remote core would finish later than waiting for the local core.
+    assert(source == destination);
+  }
+}
+
+static void testReadyCommunication() {
+  mlir::MLIRContext context;
+  mlir::OwningOpRef<mlir::ModuleOp> owner =
+    mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
+  auto type = mlir::RankedTensorType::get({1}, mlir::Float32Type::get(&context));
+  mlir::Block payloads;
+  auto payload = payloads.addArgument(type, owner->getLoc());
+  DeferredTransferPlan plan;
+  plan.stepCounts = {3, 1, 1};
+  plan.scheduled.resize(3);
+  for (auto [index, info] : llvm::enumerate(plan.scheduled)) {
+    info.op = owner->getOperation();
+    info.streamIds.push_back(index);
+  }
+  ProducedValue producer;
+  producer.scheduled = &plan.scheduled[0];
+  producer.payload = payload;
+  // Two consumers exercise removal of already-scheduled ready-queue entries.
+  for (unsigned target = 1; target <= 2; ++target) {
+    auto exchange = std::make_unique<DeferredExchangePlan>();
+    exchange->target = &plan.scheduled[target];
+    exchange->exchangeId = target;
+    RequirementFamily requirement;
+    requirement.producer = &producer;
+    requirement.publicationFragmentType = type;
+    exchange->requirements.push_back(requirement);
+    ExternalTransferFamily transfer;
+    transfer.requirement = &exchange->requirements.front();
+    transfer.sourceScheduled = producer.scheduled;
+    transfer.targetScheduled = exchange->target;
+    transfer.targetLanes = LaneSet::all(1);
+    transfer.targetStreams = onnx_mlir::StaticIntSequence::uniform(target, 1);
+    exchange->external.push_back(transfer);
+    plan.exchanges.push_back(std::move(exchange));
+  }
+  auto scheduled = scheduleDeferredCommunication({}, plan);
+  assert(mlir::succeeded(scheduled));
+  assert(scheduled->slices.size() == 2);
+  for (const auto &slice : scheduled->slices)
+    assert(slice.sourceInsertionStep == 1);
+}
+
 int main() {
+  testPipelineInputArrival();
+  testReadyCommunication();
   PipelineCoreLayout unevenLayout(138, 4);
   assert(unevenLayout.isValid());
-  assert(unevenLayout.getLogicalProcessorCount() == 34);
   assert(unevenLayout.getStageRange(0).begin == 0);
   assert(unevenLayout.getStageRange(0).size == 35);
   assert(unevenLayout.getStageRange(1).begin == 35);
@@ -56,43 +197,7 @@ int main() {
   assert(fast.getInterProcessorLatencyNs(0, 1) == 3);
   assert(slow.getInterProcessorLatencyNs(0, 1) == 10);
 
-  SchedulingTarget line;
-  line.processorCount = 3;
-  line.interProcessorLatencyNs = {
-    0,
-    1,
-    10,
-    1,
-    0,
-    1,
-    10,
-    1,
-    0,
-  };
-  std::vector<Cost> logicalTrafficFlits(9, 0);
-  logicalTrafficFlits[2] = 100;
-  assert(mapLogicalProcessorsToPhysicalCores(logicalTrafficFlits, line) == std::vector<size_t>({1, 0, 2}));
-
-  SchedulingTarget alreadyPlaced = line;
-  alreadyPlaced.interProcessorLatencyNs = {
-    0,
-    10,
-    1,
-    10,
-    0,
-    1,
-    1,
-    1,
-    0,
-  };
-  assert(mapLogicalProcessorsToPhysicalCores(logicalTrafficFlits, alreadyPlaced) == std::vector<size_t>({0, 1, 2}));
-
-  std::vector<size_t> placementGroups {0, 1, 1};
-  std::vector<size_t> groupedPlacement = mapLogicalProcessorsToPhysicalCores(
-    logicalTrafficFlits, line, placementGroups);
-  for (size_t processor = 0; processor < groupedPlacement.size(); ++processor)
-    assert(placementGroups[processor]
-           == placementGroups[groupedPlacement[processor]]);
+  testPhysicalPeft();
 
   ComputeGraph graph;
   graph.successors.resize(6);
@@ -108,16 +213,16 @@ int main() {
     graph.instanceToIndex[instance] = task;
   }
 
-  MergeScheduleResult logicalSchedule;
-  logicalSchedule.processorCount = 2;
-  logicalSchedule.dominanceOrderCompute.reserve(graph.nodes.size());
+  MergeScheduleResult physicalSchedule;
+  physicalSchedule.processorCount = 4;
+  physicalSchedule.dominanceOrderCompute.reserve(graph.nodes.size());
   for (size_t task = 0; task < graph.nodes.size(); ++task) {
     const ComputeInstance& instance = graph.nodes[task].instance;
-    logicalSchedule.dominanceOrderCompute.push_back(instance);
+    physicalSchedule.dominanceOrderCompute.push_back(instance);
     size_t cpu = task < 4 ? 0 : 1;
-    logicalSchedule.computeToCpuMap[instance] = cpu;
-    logicalSchedule.computeToCpuSlotMap[instance] = task < 4 ? task : task - 4;
-    logicalSchedule.computeToAestMap[instance] = task;
+    physicalSchedule.computeToCpuMap[instance] = cpu;
+    physicalSchedule.computeToCpuSlotMap[instance] = task < 4 ? task : task - 4;
+    physicalSchedule.computeToAestMap[instance] = task;
   }
 
   SchedulingTarget physical = fast;
@@ -129,16 +234,46 @@ int main() {
     3, 3, 0, 3,
     3, 3, 3, 0,
   };
+  SchedulingTarget clustered = physical;
+  clustered.interProcessorLatencyNs = {
+    0, 20, 1, 20,
+    20, 0, 20, 1,
+    1, 20, 0, 20,
+    20, 1, 20, 0,
+  };
+  PipelineCoreLayout clusteredLayout(std::vector<size_t> {2, 2}, clustered);
+  assert(clusteredLayout.getPhysicalCore(0) == 0);
+  assert(clusteredLayout.getPhysicalCore(1) == 2);
+  assert(clusteredLayout.getStageForCore(0) == 0);
+  assert(clusteredLayout.getStageForCore(2) == 0);
+  assert(clusteredLayout.getStageForCore(1) == 1);
+  assert(clusteredLayout.getStageForCore(3) == 1);
+  PipelineCoreLayout uniformLayout(std::vector<size_t> {2, 2}, physical);
+  assert(uniformLayout.getPhysicalCore(1) == 1);
   std::string pipelineError;
+  ComputeGraph preparationGraph = graph;
+  ResidentWeight extraWeight;
+  extraWeight.opaqueLane = graph.nodes.size();
+  preparationGraph.nodes[2].residentWeights.push_back(extraWeight);
+  auto preparation = preparePipelineWorkload(
+    preparationGraph, 2, physical, pipelineError);
+  assert(mlir::succeeded(preparation));
+  assert(*preparation == PipelineWorkloadPreparation::Ready);
+  extraWeight.opaqueLane++;
+  preparationGraph.nodes[2].residentWeights.push_back(extraWeight);
+  assert(mlir::failed(preparePipelineWorkload(
+    preparationGraph, 2, physical, pipelineError)));
+  assert(pipelineError.find("physical core's crossbars") != std::string::npos);
+  pipelineError.clear();
   ComputeGraph emptyGraph;
   MergeScheduleResult emptySchedule;
-  emptySchedule.processorCount = 2;
+  emptySchedule.processorCount = 4;
   assert(mlir::succeeded(applyPipelineScheduling(
     emptyGraph, emptySchedule, 2, physical, pipelineError)));
   assert(emptySchedule.processorCount == physical.processorCount);
   assert(emptySchedule.processorStages == std::vector<size_t>({0, 0, 1, 1}));
 
-  MergeScheduleResult pipelineSchedule = logicalSchedule;
+  MergeScheduleResult pipelineSchedule = physicalSchedule;
   assert(mlir::succeeded(applyPipelineScheduling(
     graph, pipelineSchedule, 2, physical, pipelineError)));
   assert(pipelineSchedule.processorCount == 4);
@@ -167,7 +302,8 @@ int main() {
   fourStagePhysical.interProcessorLatencyNs.assign(64, 3);
   for (size_t core = 0; core < 8; ++core)
     fourStagePhysical.interProcessorLatencyNs[core * 8 + core] = 0;
-  MergeScheduleResult fourStageSchedule = logicalSchedule;
+  MergeScheduleResult fourStageSchedule = physicalSchedule;
+  fourStageSchedule.processorCount = 8;
   assert(mlir::succeeded(applyPipelineScheduling(
     graph, fourStageSchedule, 4, fourStagePhysical, pipelineError)));
   for (size_t task = 0; task < graph.nodes.size(); ++task)
@@ -191,7 +327,7 @@ int main() {
     {4, TransferCost {.fixed = 0, .networkFlits = 1}});
   const Cost communicationCosts[] = {6, 4, 6, 4, 1};
   MergeScheduleResult communicationSchedule;
-  communicationSchedule.processorCount = 2;
+  communicationSchedule.processorCount = 4;
   for (uint32_t task = 0; task < 5; ++task) {
     ComputeInstance instance {nullptr, task, 1};
     ResidentWeight weight;
@@ -214,6 +350,9 @@ int main() {
 
   SchedulingTarget slowPipeline = fastPipeline;
   slowPipeline.averageInterProcessorLatencyNs = 10;
+  slowPipeline.interProcessorLatencyNs.assign(16, 10);
+  for (size_t core = 0; core < 4; ++core)
+    slowPipeline.interProcessorLatencyNs[core * 4 + core] = 0;
   MergeScheduleResult slowCommunicationSchedule = communicationSchedule;
   assert(mlir::succeeded(applyPipelineScheduling(
     communicationGraph, slowCommunicationSchedule, 2, slowPipeline, pipelineError)));

validation/raptor_validation/validate_one.py

@@ -75,7 +75,7 @@ PIMSIM_FAILED = "ERROR"
 PIMSIM_UNSUPPORTED = "UNSUPPORTED"
 PIMSIM_SKIPPED = "SKIP"
 PIMSIM_NOT_RUN = "-"
-PIMSIM_UNSUPPORTED_VSOFTMAX = "Pimsim does not support opcode vsoftmax"
+PIMSIM_UNSUPPORTED_VSOFTMAX = "does not support opcode vsoftmax"
 
 
 class PimSimUnsupportedError(RuntimeError):
@@ -563,7 +563,10 @@ def validate_execution(
         "Run non-functional simulation", name,
     )
     config_path = execution["pimsim_config"]
-    if state["compiled"] and pimsim_nn_build_dir is not None and config_path is not None:
+    if state["compiled"] and state["resource_metrics"].get("used_core_count") == 0:
+        state["pimsim_status"] = PIMSIM_SKIPPED
+        print_info(reporter, "Pimsim non-functional simulation skipped: no active cores")
+    elif state["compiled"] and pimsim_nn_build_dir is not None and config_path is not None:
         try:
             state["metrics"] = run_pimsim_nn(
                 pimsim_nn_build_dir, pim_dir, config_path, name,

validation/tools/pim/pimcomp/compare/compare_raptor_pimcomp_model.py

@@ -16,6 +16,7 @@ import sys
 import time
 import types
 from collections import Counter
+from concurrent.futures import ThreadPoolExecutor
 from dataclasses import asdict, dataclass
 from pathlib import Path
 from tempfile import TemporaryDirectory
@@ -242,7 +243,7 @@ def reference_inputs_exist(
 
 
 def reference_batch_dirs(root: Path, batch_size: int) -> list[Path]:
-    return [root / f"batch_{index:06d}" for index in range(batch_size)]
+    return [root / f"batch_{index:06d}/outputs" for index in range(batch_size)]
 
 
 def reference_batch_outputs_exist(
@@ -442,6 +443,7 @@ def generate_reference_batch_outputs(
     runner_build_dir: Path,
     model_path: Path,
     input_batch: list[list[np.ndarray]],
+    outputs_desc: list[tuple[int, str, int, list[int]]],
     steps: list[StepRecord],
     args: argparse.Namespace,
     out_dir: Path,
@@ -450,20 +452,26 @@ def generate_reference_batch_outputs(
 ) -> list[Path]:
     if print_header:
         print_step("Run reference")
-    references = []
-    for index, sample in enumerate(input_batch):
-        references.append(
-            generate_reference_outputs(
-                runner_path,
-                runner_build_dir,
-                model_path,
-                sample,
-                steps,
-                args,
-                out_dir / f"batch_{index:06d}",
-                print_header=False,
-            )
+    references = reference_batch_dirs(out_dir, len(input_batch))
+    missing = [
+        index for index, reference in enumerate(references)
+        if not reference_outputs_exist(outputs_desc, reference)
+    ]
+
+    def generate(index: int) -> None:
+        generate_reference_outputs(
+            runner_path,
+            runner_build_dir,
+            model_path,
+            input_batch[index],
+            steps,
+            args,
+            out_dir / f"batch_{index:06d}",
+            print_header=False,
         )
+
+    with ThreadPoolExecutor(max_workers=args.jobs) as executor:
+        list(executor.map(generate, missing))
     return references
 
 
@@ -490,6 +498,7 @@ def prepare_reference_batch_outputs(
             runner_build_dir,
             model_path,
             input_batch,
+            outputs_desc,
             steps,
             args,
             out_dir,
@@ -512,16 +521,29 @@ def prepare_common_artifacts(
     outputs_ready = reference_outputs_exist(outputs_desc, outputs_dir)
     if inputs_ready:
         arrays_in_order = load_saved_inputs(inputs_desc, inputs_dir)
+    input_batch = generate_input_batch(
+        inputs_desc, arrays_in_order, args.batch_size, args.seed)
+    batch_dir = common_dir / f"reference_seed_{args.seed}"
+    first_batch_dir = batch_dir / "batch_000000"
+    if inputs_ready and outputs_ready and not reference_outputs_exist(
+        outputs_desc, first_batch_dir / "outputs"
+    ):
+        shutil.copytree(inputs_dir, first_batch_dir / "inputs", dirs_exist_ok=True)
+        shutil.copytree(outputs_dir, first_batch_dir / "outputs", dirs_exist_ok=True)
+    references = prepare_reference_batch_outputs(
+        runner_path,
+        runner_path.parent,
+        model_path,
+        input_batch,
+        outputs_desc,
+        steps,
+        args,
+        batch_dir,
+    )
     if not (inputs_ready and outputs_ready):
-        generate_reference_outputs(
-            runner_path,
-            runner_path.parent,
-            model_path,
-            arrays_in_order,
-            steps,
-            args,
-            common_dir,
-        )
+        shutil.copytree(
+            references[0].parent / "inputs", inputs_dir, dirs_exist_ok=True)
+        shutil.copytree(references[0], outputs_dir, dirs_exist_ok=True)
 def compile_raptor_target(
     model_path: Path,
     out_dir: Path,
@@ -1409,7 +1431,14 @@ def main():
     parser.add_argument("--mesh-cols", type=int)
     parser.add_argument("--pimsim-time-ms", type=int, default=1000)
     parser.add_argument("--pimsim-mode", choices=["latency", "throughput"], default="latency")
-    parser.add_argument("--batch-size", type=int, default=128)
+    parser.add_argument("--batch-size", type=int, default=64)
+    parser.add_argument(
+        "-j",
+        "--jobs",
+        type=int,
+        default=4,
+        help="Maximum parallel native reference runner processes (default: 4).",
+    )
     parser.add_argument("--pimcomp-pipeline", choices=["element", "batch"])
     parser.add_argument("--pimcomp-model-name", help="Use a Pimcomp built-in model name such as vgg16.")
     parser.add_argument(
@@ -1466,6 +1495,8 @@ def main():
         parser.error("--pimsim-time-ms must be positive")
     if args.batch_size <= 0:
         parser.error("--batch-size must be positive")
+    if args.jobs <= 0:
+        parser.error("--jobs must be positive")
     if args.pimsim_mode == "throughput" and args.batch_size < 2:
         parser.error("throughput mode requires batch size greater than 1")
     if args.timeout_seconds < 0:
@@ -1674,7 +1705,7 @@ def main():
         write_input_batch_csv(out_dir / "inputs.csv", input_batch)
         raptor_input_bins = write_input_batch_binaries(input_batch, out_dir / "simulation/raptor_inputs")
         if args.pimsim_mode == "throughput":
-            batch_reference_dir = common_dir / f"reference_batch_{args.batch_size}_seed_{args.seed}"
+            batch_reference_dir = common_dir / f"reference_seed_{args.seed}"
             throughput_references = try_stage(
                 failures,
                 "Run reference",

validation/tools/pim/pimcomp/compare/run_pimcomp_models.py

@@ -261,23 +261,9 @@ def performance_values(performance: dict) -> dict[str, float | None]:
     }
 
 
-def comparison_passed(report: dict, compiler: str | None = None) -> bool:
-    other_compiler = "PIMCOMP" if compiler == "raptor" else "RAPTOR"
-    if any(
-        compiler is None or other_compiler not in failure.get("stage", "").upper()
-        for failure in report.get("failures", [])
-    ):
-        return False
-    compilers = (compiler,) if compiler is not None else ("raptor", "pimcomp")
-    for name in compilers:
-        result = report.get(f"{name}_validation") or {}
-        if result.get("status") != "done" or not result.get("passed"):
-            return False
-    for name in compilers:
-        performance = report.get(f"{name}_performance") or {}
-        if performance.get("error") or performance.get("skipped"):
-            return False
-    return True
+def comparison_passed(report: dict) -> bool:
+    result = report.get("raptor_validation") or {}
+    return result.get("status") == "done" and bool(result.get("passed"))
 
 
 def functional_validation_status(result: dict | None) -> str:
@@ -404,13 +390,18 @@ def comparison_command(
         *[f"--raptor-extra-arg={arg}" for arg in raptor_extra_args],
         "--timeout-seconds",
         str(timeout),
-        "--fail-on-error",
         *([] if fast else ["--no-fast"]),
         *reuse_args,
     ]
 
 
-def prepare_common_command(model: Path, common_dir: Path, timeout: float) -> list[str]:
+def prepare_common_command(
+    model: Path,
+    common_dir: Path,
+    timeout: float,
+    batch_size: int,
+    jobs: int,
+) -> list[str]:
     return [
         sys.executable,
         str(COMPARE),
@@ -421,6 +412,10 @@ def prepare_common_command(model: Path, common_dir: Path, timeout: float) -> lis
         "--common-dir",
         str(common_dir),
         "--prepare-common",
+        "--batch-size",
+        str(batch_size),
+        "--jobs",
+        str(jobs),
         "--timeout-seconds",
         str(timeout),
     ]
@@ -587,8 +582,8 @@ def main() -> int:
     parser.add_argument(
         "--batch-size",
         type=int,
-        default=128,
-        help="functional throughput batch size (default: 128).",
+        default=64,
+        help="functional throughput and shared reference batch size (default: 64).",
     )
     parser.add_argument(
         "--timeout-seconds",
@@ -664,6 +659,11 @@ def main() -> int:
         }
         comparisons_by_arch[arch] = comparisons
         configs_by_arch[arch] = configs
+    reference_batch_size = max(
+        1 if mode == "latency" else args.batch_size
+        for comparisons in comparisons_by_arch.values()
+        for mode, _, _ in comparisons
+    )
 
     missing = [
         str(path)
@@ -706,6 +706,7 @@ def main() -> int:
     print(f"Modes: {', '.join(args.mode)}")
     print(f"Throughput Pimsim time: {args.pimsim_time_ms} ms")
     print(f"Max parallel jobs: {args.jobs}")
+    print(f"Shared reference batch: {reference_batch_size}")
     print(
         f"Comparison jobs: "
         f"{sum(len(args.models) * len(comparisons) for comparisons in comparisons_by_arch.values())}"
@@ -722,6 +723,8 @@ def main() -> int:
                         FUNCTIONAL_MODELS[name],
                         common_dir(out_dir, name, common_root),
                         args.timeout_seconds,
+                        reference_batch_size,
+                        args.jobs,
                     ),
                     dry_run=args.dry_run,
                 )
@@ -849,10 +852,8 @@ def main() -> int:
                 report_path = result_dir(
                     out_dir, name, arch, mode, pipeline, args.ablation_variant
                 ) / "pimcomp/comparison_report.json"
-                compiler = "raptor" if args.raptor_only else args.only
                 if not report_path.exists() or not comparison_passed(
                     json.loads(report_path.read_text(encoding="utf-8")),
-                    compiler,
                 ):
                     if label not in failed:
                         failed.append(label)

validation/tools/pim/pimcomp/correctness/README.md

@@ -1,4 +1,47 @@
-# Pimcomp batch correctness reproduction
+# Pimcomp correctness experiments
+
+## Synchronization ordering
+
+The synchronization experiment checks whether repeated PIMCOMP and Raptor
+programs preserve cross-core global-memory generations under legal execution
+schedules. It builds a small two-convolution model for Arch-A, Arch-B, and
+Arch-C, finds cross-core `ST`/`LD` dependencies, and compares greedy,
+bounded-stall, randomized, and adversarial runs using the Rust simulator's
+provenance trace. Byte-identical batch inputs isolate intermediate-memory
+ordering from the separate host-input lifetime issue.
+
+The experiment records static dependency evidence, dynamic provenance,
+functional output comparisons, architecture contract classifications, and
+per-artifact PIMCOMP/Raptor conclusions. Diagnostic scheduling changes only
+the functional simulator's execution order; it does not modify PIMCOMP or the
+`pimsim-nn` performance oracle.
+
+Prerequisites are the repository virtual environment, the release Raptor
+compiler, built PIMCOMP frontend/backend, and the existing `pimsim-nn` build.
+Run from the repository root with a new or empty output directory:
+
+```bash
+.venv/bin/python \
+  validation/tools/pim/pimcomp/correctness/synchronization/run_experiment.py \
+  --out-dir /tmp/pimcomp-adversarial-sync \
+  --batch-size 4 --seed 0 --self-check
+```
+
+The directory contains:
+
+- [`run_experiment.py`](synchronization/run_experiment.py): orchestration,
+  classification, reporting, and self-checks.
+- [`global_memory.py`](synchronization/global_memory.py): artifact compilation,
+  dependency analysis, simulator execution, and provenance helpers.
+- [`architecture_contract.py`](synchronization/architecture_contract.py) and
+  [`architecture_evidence.json`](synchronization/architecture_evidence.json):
+  conservative architecture-contract evidence and labels.
+
+The main outputs are `adversarial_memory_sync_report.json` and
+`adversarial_memory_sync_report.md`, with per-architecture evidence below the
+same output root.
+
+## Pimcomp batch prefill reproduction
 
 Pimcomp's batch scheduler currently emits an incomplete standalone program for
 models containing post operations. The generated `VerificationInfo.json` uses a
@@ -24,14 +67,14 @@ schedule.
 Run the default reproduction from the repository root:
 
 ```bash
-.venv/bin/python validation/tools/pim/pimcomp/correctness/run_prefill_experiment.py
+.venv/bin/python validation/tools/pim/pimcomp/correctness/prefill/run_experiment.py
 ```
 
 The launcher accepts an alternate comparison directory, model, work directory,
 and shared reference-artifact directory:
 
 ```bash
-.venv/bin/python validation/tools/pim/pimcomp/correctness/run_prefill_experiment.py \
+.venv/bin/python validation/tools/pim/pimcomp/correctness/prefill/run_experiment.py \
   validation/networks/pimcomp_models/googlenet/artifacts/arch-a/throughput/pipeline2 \
   validation/networks/pimcomp_models/googlenet/googlenet-12-pimsim-nn.onnx \
   validation/networks/pimcomp_models/googlenet/artifacts/arch-a/throughput/pipeline2/correctness/prefill \
@@ -42,7 +85,7 @@ Without the optional work-directory argument, the experiment uses the same
 `correctness/prefill/` directory below the comparison artifacts.
 
 It runs the exported artifact once with its original memory image and once
-with [`prefill_batch_memory.py`](prefill_batch_memory.py), then compares both
+with [`batch_memory.py`](prefill/batch_memory.py), then compares both
 outputs with the recorded native reference. The expected GoogLeNet result is a
 baseline maximum difference near `6.70705` and a prefilled maximum difference
 near `4.05e-6`.

validation/tools/pim/pimcomp/correctness/prefill/run_experiment.py

@@ -11,11 +11,11 @@ from pathlib import Path
 
 import numpy as np
 
-from prefill_batch_memory import prefill_batch_memory
+from batch_memory import prefill_batch_memory
 
 
 SCRIPT_DIR = Path(__file__).resolve().parent
-REPO_ROOT = SCRIPT_DIR.parents[4]
+REPO_ROOT = SCRIPT_DIR.parents[5]
 DEFAULT_COMPARISON_DIR = REPO_ROOT / "validation/networks/pimcomp_models/googlenet/artifacts/arch-a/throughput/pipeline2"
 DEFAULT_MODEL = REPO_ROOT / "validation/networks/pimcomp_models/googlenet/googlenet-12-pimsim-nn.onnx"
 DEFAULT_COMMON_DIR = REPO_ROOT / "validation/networks/pimcomp_models/googlenet/artifacts/common"

validation/tools/pim/pimcomp/correctness/synchronization/architecture_contract.py

@@ -15,12 +15,13 @@ from pathlib import Path
 from typing import Any
 
 
-REPO = Path(__file__).resolve().parents[5]
+REPO = Path(__file__).resolve().parents[6]
 VALIDATION = REPO / "validation"
 PIMCOMP_ROOT = REPO / "third_party/PIMCOMP-NN"
 PIMSIM_ROOT = REPO / "backend-simulators/pim/pimsim-nn"
 RUST_ROOT = REPO / "backend-simulators/pim/pim-simulator"
-EVIDENCE_PATH = Path(__file__).with_name("pimcomp_architecture_sync_evidence.json")
+COMPARE_DIR = REPO / "validation/tools/pim/pimcomp/compare"
+EVIDENCE_PATH = Path(__file__).with_name("architecture_evidence.json")
 
 
 def _default_evidence() -> dict[str, Any]:
@@ -161,9 +162,8 @@ def make_identical_inputs(model: Path, batch_size: int, out: Path) -> list[Path]
 
     import sys
 
-    compare_dir = Path(__file__).resolve().parent
-    if str(compare_dir) not in sys.path:
-        sys.path.insert(0, str(compare_dir))
+    if str(COMPARE_DIR) not in sys.path:
+        sys.path.insert(0, str(COMPARE_DIR))
     import compare_raptor_pimcomp_model as compare  # noqa: PLC0415
 
     inputs, _ = compare.onnx_io(model)

validation/tools/pim/pimcomp/correctness/synchronization/global_memory.py

@@ -16,17 +16,17 @@ from typing import Any
 import numpy as np
 
 
-REPO = Path(__file__).resolve().parents[5]
+REPO = Path(__file__).resolve().parents[6]
 VALIDATION = REPO / "validation"
 CONFIG_ROOT = VALIDATION / "pimsim_configs/pimcomp"
 PIMCOMP_ROOT = REPO / "third_party/PIMCOMP-NN"
 PIMSIM_NN_ROOT = REPO / "backend-simulators/pim/pimsim-nn"
 RUST_ROOT = REPO / "backend-simulators/pim/pim-simulator"
 RUST_BINARY = RUST_ROOT / "target/release/pim-simulator"
-COMPARE_SCRIPT = Path(__file__).with_name("compare_raptor_pimcomp_model.py")
+COMPARE_SCRIPT = REPO / "validation/tools/pim/pimcomp/compare/compare_raptor_pimcomp_model.py"
 PYTHON = REPO / ".venv/bin/python"
 
-sys.path.insert(0, str(Path(__file__).resolve().parent))
+sys.path.insert(0, str(COMPARE_SCRIPT.parent))
 import compare_raptor_pimcomp_model as compare  # noqa: E402
 
 
@@ -182,7 +182,7 @@ def compile_artifact(args: Any, model: Path, architecture_out: Path, throughput_
         "--pimsim-time-ms", "1000",
         "--batch-size", str(args.batch_size),
         "--pimcomp-pipeline", "batch",
-        "--pimcomp-replication", "balance",
+        "--pimcomp-replication", "GA",
         "--raptor-extra-arg=--pipeline=4",
         "--seed", str(args.seed),
         "--timeout-seconds", str(args.timeout),

validation/tools/pim/pimcomp/correctness/synchronization/run_experiment.py

@@ -3,7 +3,7 @@
 
 One documented invocation is:
 
-    .venv/bin/python validation/tools/pim/pimcomp/correctness/test_PIMCOMP_adversarial_memory_sync.py \
+    .venv/bin/python validation/tools/pim/pimcomp/correctness/synchronization/run_experiment.py \
         --out-dir /tmp/pimcomp-adversarial-sync --batch-size 4 --seed 0 --self-check
 
 The experiment uses identical external input bytes for throughput runs so the
@@ -24,16 +24,16 @@ from pathlib import Path
 from typing import Any
 
 
-REPO = Path(__file__).resolve().parents[5]
+REPO = Path(__file__).resolve().parents[6]
 SCRIPT = Path(__file__).resolve()
-HELPER_DIR = REPO / "validation/tools/pim/pimcomp/compare"
-GLOBAL_SCRIPT = HELPER_DIR / "test_PIMCOMP_global_memory_sync.py"
-AUDIT_SCRIPT = HELPER_DIR / "test_PIMCOMP_architecture_sync_contract.py"
+EXPERIMENT_DIR = SCRIPT.parent
+GLOBAL_SCRIPT = EXPERIMENT_DIR / "global_memory.py"
+AUDIT_SCRIPT = EXPERIMENT_DIR / "architecture_contract.py"
 PYTHON = REPO / ".venv/bin/python"
 
-sys.path.insert(0, str(HELPER_DIR))
-import test_PIMCOMP_architecture_sync_contract as audit  # noqa: E402
-import test_PIMCOMP_global_memory_sync as global_sync  # noqa: E402
+sys.path.insert(0, str(EXPERIMENT_DIR))
+import architecture_contract as audit  # noqa: E402
+import global_memory as global_sync  # noqa: E402
 
 
 INVALID = "INVALID_SYNCHRONIZATION_REPRODUCER"