From 336f0b506e93261e9e9946fddb1576df8f7a0ba3 Mon Sep 17 00:00:00 2001 From: NiccoloN Date: Mon, 24 Aug 2026 11:58:04 +0200 Subject: [PATCH] better synchronization better deadlock detection to also track wait/sync --- .../invariants/PIMSIM_NN_ORACLE_INVARIANT.md | 4 + .../PIM_SYNCHRONIZATION_INVARIANT.md | 105 ++++++ AGENTS.md | 1 + .../src/lib/instruction_set/isa.rs | 10 + .../pim/pim-simulator/src/lib/pimcore.rs | 327 ++++++++++++++---- .../pim/pim-simulator/tests/sync.rs | 62 +++- src/PIM/Compiler/PimCompilerOptions.cpp | 2 +- .../SpatialToPim/CoreLoweringPatterns.cpp | 3 - .../SpatialToPim/Patterns/ChannelLowering.cpp | 4 - .../Verification/VerificationPass.cpp | 235 +++++++++++-- .../DeferredBoundaryPlanning.cpp | 209 ++++++++--- .../DeferredBoundaryPlanning.hpp | 7 +- .../DeferredBoundaryRealization.cpp | 321 +++++++++++------ .../DeferredCommunicationModel.hpp | 2 - .../DeferredTransferPlanning.hpp | 8 +- src/PIM/Dialect/Spatial/Spatial.td | 7 +- test/PIM/CMakeLists.txt | 10 + validation/raptor_validation/validate_one.py | 7 +- .../compare/compare_raptor_pimcomp_model.py | 81 +++-- .../pim/pimcomp/compare/run_pimcomp_models.py | 47 +-- 20 files changed, 1123 insertions(+), 329 deletions(-) create mode 100644 .agents/invariants/PIM_SYNCHRONIZATION_INVARIANT.md diff --git a/.agents/invariants/PIMSIM_NN_ORACLE_INVARIANT.md b/.agents/invariants/PIMSIM_NN_ORACLE_INVARIANT.md index 575f929..8980da5 100644 --- a/.agents/invariants/PIMSIM_NN_ORACLE_INVARIANT.md +++ b/.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. diff --git a/.agents/invariants/PIM_SYNCHRONIZATION_INVARIANT.md b/.agents/invariants/PIM_SYNCHRONIZATION_INVARIANT.md new file mode 100644 index 0000000..1f740b0 --- /dev/null +++ b/.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. diff --git a/AGENTS.md b/AGENTS.md index 7c373e2..1deb855 100644 --- a/AGENTS.md +++ b/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: diff --git a/backend-simulators/pim/pim-simulator/src/lib/instruction_set/isa.rs b/backend-simulators/pim/pim-simulator/src/lib/instruction_set/isa.rs index 5bcae69..035491b 100644 --- a/backend-simulators/pim/pim-simulator/src/lib/instruction_set/isa.rs +++ b/backend-simulators/pim/pim-simulator/src/lib/instruction_set/isa.rs @@ -886,11 +886,21 @@ pub fn recv(cores: &mut CPU, data: InstructionData) -> Result 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 { 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 { Ok(InstructionStatus::Sync(data)) diff --git a/backend-simulators/pim/pim-simulator/src/lib/pimcore.rs b/backend-simulators/pim/pim-simulator/src/lib/pimcore.rs index fb6f8d6..d5027f1 100644 --- a/backend-simulators/pim/pim-simulator/src/lib/pimcore.rs +++ b/backend-simulators/pim/pim-simulator/src/lib/pimcore.rs @@ -14,7 +14,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}, @@ -104,7 +104,13 @@ struct DeadlockInfo { states: String, } -type SyncEvents = Vec<[i32; 32]>; +#[derive(Debug, Clone, Default)] +struct SyncEvent { + count: i32, + sources: HashMap, +} + +type SyncEvents = Vec<[SyncEvent; 32]>; fn print_status(core_instructions: &[CoreInstructions]) { let mut tot_instructions = 0; @@ -182,7 +188,9 @@ impl<'a> Executable<'a> { } = self; 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 now = SystemTime::now(); @@ -222,9 +230,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 ); @@ -255,9 +265,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 ); @@ -316,18 +326,23 @@ fn store_input(cpu: &mut CPU, input: &[u8], input_regions: &[(usize, usize)]) -> Ok(()) } -fn detect_deadlock(cores_instructions: &[CoreInstructions]) -> Option { +fn detect_deadlock( + cores_instructions: &[CoreInstructions], + events: &SyncEvents, + batch_size: u32, +) -> Option { #[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; } @@ -344,94 +359,191 @@ fn detect_deadlock(cores_instructions: &[CoreInstructions]) -> Option { + 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> = 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>, + path: &mut Vec, + positions: &mut HashMap, + visited: &mut HashSet, + ) -> Option> { + 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(¤t_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::>() - .join(" -> "); - - let cycle = cycle - .iter() - .copied() - .chain(std::iter::once(waiting_for)) - .collect::>(); - 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::>() - .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::>() + .join(" -> "); + let states_msg = cycle + .iter() + .map(&format_state) + .collect::>() + .join(", "); + return Some(DeadlockInfo { + cycle: cycle_msg, + states: states_msg, + }); } } None @@ -446,7 +558,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 } @@ -454,8 +568,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 { @@ -465,3 +579,66 @@ fn handle_wait_sync( _ => false, } } + +#[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()); + } +} diff --git a/backend-simulators/pim/pim-simulator/tests/sync.rs b/backend-simulators/pim/pim-simulator/tests/sync.rs index e638e70..9edf004 100644 --- a/backend-simulators/pim/pim-simulator/tests/sync.rs +++ b/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); diff --git a/src/PIM/Compiler/PimCompilerOptions.cpp b/src/PIM/Compiler/PimCompilerOptions.cpp index 1c03f3c..5650da9 100644 --- a/src/PIM/Compiler/PimCompilerOptions.cpp +++ b/src/PIM/Compiler/PimCompilerOptions.cpp @@ -96,7 +96,7 @@ llvm::cl::opt pimEmitJson("pim-emit-json", llvm::cl::opt 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)); diff --git a/src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp b/src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp index 2fda833..199c62e 100644 --- a/src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp +++ b/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, diff --git a/src/PIM/Conversion/SpatialToPim/Patterns/ChannelLowering.cpp b/src/PIM/Conversion/SpatialToPim/Patterns/ChannelLowering.cpp index 81f4537..340eb48 100644 --- a/src/PIM/Conversion/SpatialToPim/Patterns/ChannelLowering.cpp +++ b/src/PIM/Conversion/SpatialToPim/Patterns/ChannelLowering.cpp @@ -152,10 +152,6 @@ struct HostWaitLoadLowering : OpRewritePattern { 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; }); } diff --git a/src/PIM/Dialect/Pim/Passes/Transforms/Verification/VerificationPass.cpp b/src/PIM/Dialect/Pim/Passes/Transforms/Verification/VerificationPass.cpp index 125b3b1..34bbd80 100644 --- a/src/PIM/Dialect/Pim/Passes/Transforms/Verification/VerificationPass.cpp +++ b/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 #include #include @@ -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; +using SynchronizationEventKey = std::pair; +using SynchronizationSourceCounts = + DenseMap>; 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& coreEvents, const DenseMap& 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(&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(events.size()), &op}); + return success(); + } + + if (auto waitOp = dyn_cast(&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(events.size()), &op}); + return success(); + } + return success(); }); } @@ -530,7 +608,7 @@ static void printCommunicationWindow(llvm::raw_ostream& os, static void printCommunicationDeadlockReport(const DenseMap& coreEvents, const DenseMap& programCounters, ArrayRef 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 DenseMapsecond); } - 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> findCommunicationWaitCycle(const DenseMap& coreEvents, - const DenseMap& programCounters) { - for (const auto& [startCoreId, events] : coreEvents) { - auto startPcIt = programCounters.find(startCoreId); - if (startPcIt == programCounters.end() || startPcIt->second >= events.size()) + const DenseMap& programCounters, + const DenseSet& repeatingCores, + const SynchronizationSourceCounts& sourceCounts) { + DenseMap> dependencies; + for (const auto& [coreId, events] : coreEvents) { + size_t pc = programCounters.lookup(coreId); + if (pc >= events.size()) continue; - - DenseMap positionInPath; - SmallVector 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 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 visited; + DenseMap positionInPath; + SmallVector path; + std::function>(int64_t)> visit = + [&](int64_t coreId) -> std::optional> { + auto position = positionInPath.find(coreId); + if (position != positionInPath.end()) + return SmallVector( + 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 coreEvents; + DenseSet repeatingCores; bool hasFailure = false; for (func::FuncOp funcOp : moduleOp.getOps()) { 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(&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 programCounters; for (const auto& [coreId, events] : coreEvents) programCounters[coreId] = 0; + DenseMap 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()) diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp index 72edd54..fafce40 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp +++ b/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 boundaries, + MutableArrayRef boundaries, size_t synchronizationRegisterCount) { bool pipelined = false; for (ScheduledInfo &scheduled : transfers.scheduled) { @@ -284,14 +286,13 @@ static LogicalResult assignPipelineSynchronization( }); DenseMap> incomingByCore; + DenseMap> readersByWriter; DenseMap> eventRegisters; DenseMap> waitValues; - DenseMap> 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>> + overwriteBoundaryReadersByWriter; + DenseSet> plannedRelations; + for (const BoundaryProgram &boundary : boundaries) { + DenseMap> readersAtBoundary; + for (const BoundaryInstruction &instruction : boundary.instructions) { + auto *send = std::get_if(&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 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> freeRegisters; + DenseMap> freeWaitCounts; + DenseMap> freeGroups; + for (int64_t core : synchronizationCores) { + auto incomingIt = incomingByCore.find(core); + ArrayRef incoming = incomingIt == incomingByCore.end() + ? ArrayRef() : 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 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 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 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> pendingGroups; + for (int64_t writer : writers) + for (unsigned group = 0; group < freeRegisters[writer].size(); ++group) + pendingGroups.insert({writer, group}); + for (BoundaryProgram &boundary : boundaries) { + SmallVector instructions; + for (BoundaryInstruction &instruction : boundary.instructions) { + auto *send = std::get_if(&instruction); + SmallVector> 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 registers(boundary.key.first->cores.size()); + SmallVector 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(); } diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.hpp index 2d7db41..49d48e5 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.hpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.hpp @@ -34,12 +34,17 @@ struct EmitReceiveAssemblyRun { llvm::SmallVector 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; struct BoundaryProgram { diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp index 6ca51c0..a497394 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp +++ b/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 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 projectionOffsets; @@ -78,6 +78,29 @@ static FailureOr 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 +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 emitReceiveValue(ArrayRef slices std::optional hostOffsets; std::optional eventRegisters; std::optional waitValues; - std::optional 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 emitReceiveValue(ArrayRef 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 emitReceiveAssembly(const EmitReceiveAssemblyRun &run, V std::optional hostOffsets; std::optional eventRegisters; std::optional waitValues; - std::optional acknowledgementEventRegisters; bool hostRouted = run.slices.front().family->hostRouted; auto metadataByEntry = buildRectangularReceiveMetadata(run, laneCount); if (succeeded(metadataByEntry)) { @@ -424,15 +437,11 @@ static FailureOr 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 positionRows; for (unsigned position : run.positions) @@ -485,15 +494,11 @@ static FailureOr 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 positionColumns; for (const StaticIntSequenceChain &values : positionsByLane) @@ -527,8 +532,6 @@ static FailureOr 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> emitInstructions(ArrayRef instructions, Value lane, unsigned laneCount, - ArrayRef results, DeferredEmissionContext &context) { + ArrayRef results, ScheduledInfo &scheduled, + DeferredEmissionContext &context) { SmallVector produced; for (size_t instructionIndex = 0; instructionIndex < instructions.size(); ++instructionIndex) { @@ -1108,7 +1112,22 @@ static FailureOr> emitInstructions(ArrayRef(&instruction)) { + if (auto reuse = std::get_if(&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(&instruction)) { if (failed(emitConditionalSendRun(*send, lane, laneCount, context))) return failure(); } else if (auto update = std::get_if(&instruction)) { @@ -1192,7 +1211,9 @@ static LogicalResult emitBoundary(const BoundaryProgram &boundary, ArrayRef(boundary.key.first->op)) lane = *batch.getLaneArgument(); SmallVector 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 emittedReaders; + for (ScheduledInfo &scheduled : transfers.scheduled) { + SmallVector 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(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 downstreamRank; + for (auto [rank, core] : llvm::enumerate(transfers.downstreamCores)) + downstreamRank[core] = rank; + + for (ScheduledInfo &scheduled : transfers.scheduled) { + LaneSet lanes; + SmallVector targets(scheduled.cores.size()); + SmallVector 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(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 stageZeroRank; for (auto [rank, core] : llvm::enumerate(transfers.stageZeroCores)) stageZeroRank[core] = rank; DenseMap 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(scheduled.op)) lane = *batch.getLaneArgument(); - SmallVector acknowledgementCounts, releaseRegisters; - SmallVector releaseWaitValues, leftTargets, leftRegisters; - SmallVector rightTargets, rightRegisters; - LaneSet barrierLanes, leaderLanes, leftLanes, rightLanes; + SmallVector 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 targets, - ArrayRef 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 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 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); } diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationModel.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationModel.hpp index 5f446a5..4ae6d06 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationModel.hpp +++ b/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; }; diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.hpp index 93356c3..d402f12 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.hpp +++ b/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 processorStages; llvm::SmallVector scheduled; @@ -14,7 +19,8 @@ struct DeferredTransferPlan { llvm::DenseMap> producedByGraph; llvm::SmallVector> exchanges; llvm::SmallVector stepCounts; - llvm::DenseMap hostAcknowledgementCounts; + llvm::DenseMap> + hostReleaseSignals; llvm::SmallVector stageZeroCores; llvm::SmallVector downstreamCores; size_t synchronizationRegisterCount = 0; diff --git a/src/PIM/Dialect/Spatial/Spatial.td b/src/PIM/Dialect/Spatial/Spatial.td index b0b1242..524e71b 100644 --- a/src/PIM/Dialect/Spatial/Spatial.td +++ b/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) }]; } diff --git a/test/PIM/CMakeLists.txt b/test/PIM/CMakeLists.txt index 208eff0..fe820ac 100644 --- a/test/PIM/CMakeLists.txt +++ b/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 "$" + --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) diff --git a/validation/raptor_validation/validate_one.py b/validation/raptor_validation/validate_one.py index 5558a52..dbd9c33 100644 --- a/validation/raptor_validation/validate_one.py +++ b/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, diff --git a/validation/tools/pim/pimcomp/compare/compare_raptor_pimcomp_model.py b/validation/tools/pim/pimcomp/compare/compare_raptor_pimcomp_model.py index 75f98a8..0c91d8b 100755 --- a/validation/tools/pim/pimcomp/compare/compare_raptor_pimcomp_model.py +++ b/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", diff --git a/validation/tools/pim/pimcomp/compare/run_pimcomp_models.py b/validation/tools/pim/pimcomp/compare/run_pimcomp_models.py index f9c2d40..115a62d 100755 --- a/validation/tools/pim/pimcomp/compare/run_pimcomp_models.py +++ b/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)