better synchronization

better deadlock detection to also track wait/sync
This commit is contained in:
NiccoloN
2026-08-24 11:58:04 +02:00
parent d634484df2
commit 336f0b506e
20 changed files with 1123 additions and 329 deletions
@@ -11,6 +11,10 @@ behavior defines the hardware model used for Raptor/PIMCOMP comparisons.
simulated timing, scheduling, power, energy, or supported input programs. simulated timing, scheduling, power, energy, or supported input programs.
- Unsupported pimsim-nn operations must remain unsupported; do not approximate - Unsupported pimsim-nn operations must remain unsupported; do not approximate
their timing or map them onto another operation. 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 - Adapt compiler inputs to the oracle instead. For YOLO, use
`validation/networks/pimcomp_models/yolo11n/yolo11n-pimsim-nn.onnx`, the `validation/networks/pimcomp_models/yolo11n/yolo11n-pimsim-nn.onnx`, the
dedicated pimsim-ready performance artifact with Softmax operations removed. dedicated pimsim-ready performance artifact with Softmax operations removed.
@@ -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.
+1
View File
@@ -7,6 +7,7 @@ Before modifying the relevant subsystem, read:
* `.agents/invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md` * `.agents/invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md`
* `.agents/invariants/PERFORMANCE_OPTIMIZATION_INVARIANT.md` * `.agents/invariants/PERFORMANCE_OPTIMIZATION_INVARIANT.md`
* `.agents/invariants/PIMSIM_NN_ORACLE_INVARIANT.md` * `.agents/invariants/PIMSIM_NN_ORACLE_INVARIANT.md`
* `.agents/invariants/PIM_SYNCHRONIZATION_INVARIANT.md`
* `.agents/invariants/PIPELINE_SCHEDULING_INVARIANT.md` * `.agents/invariants/PIPELINE_SCHEDULING_INVARIANT.md`
* `.agents/invariants/SPATIAL_TARGET_GENERALITY_INVARIANT.md` * `.agents/invariants/SPATIAL_TARGET_GENERALITY_INVARIANT.md`
* Build commands: * Build commands:
@@ -886,11 +886,21 @@ pub fn recv(cores: &mut CPU, data: InstructionData) -> Result<InstructionStatus>
Ok(InstructionStatus::Receiving(data)) Ok(InstructionStatus::Receiving(data))
} }
#[inline(never)]
pub fn isa_wait(functor: usize) -> bool {
(wait as *const () as usize) == functor
}
#[inline(never)] #[inline(never)]
pub fn wait(cores: &mut CPU, data: InstructionData) -> Result<InstructionStatus> { pub fn wait(cores: &mut CPU, data: InstructionData) -> Result<InstructionStatus> {
Ok(InstructionStatus::Waiting(data)) Ok(InstructionStatus::Waiting(data))
} }
#[inline(never)]
pub fn isa_sync(functor: usize) -> bool {
(sync as *const () as usize) == functor
}
#[inline(never)] #[inline(never)]
pub fn sync(cores: &mut CPU, data: InstructionData) -> Result<InstructionStatus> { pub fn sync(cores: &mut CPU, data: InstructionData) -> Result<InstructionStatus> {
Ok(InstructionStatus::Sync(data)) Ok(InstructionStatus::Sync(data))
@@ -14,7 +14,7 @@ use crate::{
cpu::CPU, cpu::CPU,
instruction_set::{ instruction_set::{
Instruction, InstructionStatus, Instructions, 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, memory_manager::type_traits::TryToUsize,
send_recv::{SendRecv, handle_send_recv}, send_recv::{SendRecv, handle_send_recv},
@@ -104,7 +104,13 @@ struct DeadlockInfo {
states: String, 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]) { fn print_status(core_instructions: &[CoreInstructions]) {
let mut tot_instructions = 0; let mut tot_instructions = 0;
@@ -182,7 +188,9 @@ impl<'a> Executable<'a> {
} = self; } = self;
let mut cpu_progressed = 0; let mut cpu_progressed = 0;
let max_core = cpu.num_core(); 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 cpu_index = 0;
let mut now = SystemTime::now(); let mut now = SystemTime::now();
@@ -222,9 +230,11 @@ impl<'a> Executable<'a> {
} }
if (now.elapsed().unwrap() > Duration::from_secs(5)) { if (now.elapsed().unwrap() > Duration::from_secs(5)) {
print_status(cores_instructions); 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!( bail!(
"Deadlock cycle detected: {} [{}]", "Communication deadlock detected: {} [{}]",
deadlock.cycle, deadlock.cycle,
deadlock.states deadlock.states
); );
@@ -255,9 +265,9 @@ impl<'a> Executable<'a> {
} }
print_status(cores_instructions); 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!( bail!(
"Deadlock cycle detected: {} [{}]", "Communication deadlock detected: {} [{}]",
deadlock.cycle, deadlock.cycle,
deadlock.states deadlock.states
); );
@@ -316,18 +326,23 @@ fn store_input(cpu: &mut CPU, input: &[u8], input_regions: &[(usize, usize)]) ->
Ok(()) 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)] #[derive(Debug, PartialEq, Eq)]
enum CoreState { enum CoreState {
SendingTo(i32, i32), SendingTo(i32, i32),
ReceivingFrom(i32, i32), ReceivingFrom(i32, i32),
WaitingEvent(i32, i32, i32),
Working, Working,
Halted, Halted,
} }
let mut states = HashMap::new(); 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() { if core_inst.program_counter >= core_inst.instructions.len() {
continue; continue;
} }
@@ -344,94 +359,191 @@ fn detect_deadlock(cores_instructions: &[CoreInstructions]) -> Option<DeadlockIn
); );
} else if isa_send(functor_address) { } else if isa_send(functor_address) {
states.insert(this_core, CoreState::SendingTo(target_core, data.imm_len())); 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 { } else {
states.insert(this_core, CoreState::Working); 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() { for (&core_id, state) in states.iter() {
match state { match state {
CoreState::SendingTo(target_core, size) => { CoreState::SendingTo(target_core, size) => {
let target_state = states.get(target_core).unwrap_or(&CoreState::Halted); let target_state = states.get(target_core).unwrap_or(&CoreState::Halted);
if target_state != &CoreState::ReceivingFrom(core_id, *size) { 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) => { CoreState::ReceivingFrom(target_core, size) => {
let target_state = states.get(target_core).unwrap_or(&CoreState::Halted); let target_state = states.get(target_core).unwrap_or(&CoreState::Halted);
if target_state != &CoreState::SendingTo(core_id, *size) { 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 => {} CoreState::Working | CoreState::Halted => {}
} }
} }
let mut visited = HashSet::new(); fn find_cycle(
core: i32,
for &start_core in wait_for.keys() { wait_for: &HashMap<i32, Vec<i32>>,
if visited.contains(&start_core) { path: &mut Vec<i32>,
continue; 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 path = Vec::new();
let mut current_core = start_core; let mut positions = HashMap::new();
let mut in_path = HashSet::new(); if let Some(cycle) = find_cycle(
*start_core,
while let Some(&waiting_for) = wait_for.get(&current_core) { &wait_for,
path.push(current_core); &mut path,
in_path.insert(current_core); &mut positions,
visited.insert(current_core); &mut visited,
) {
// Found a closed loop! let cycle_msg = cycle
if in_path.contains(&waiting_for) { .iter()
let cycle_start = path.iter().position(|&c| c == waiting_for).unwrap(); .chain(std::iter::once(&cycle[0]))
let cycle = &path[cycle_start..]; .map(|core| (core - 1).to_string())
let format_core = |core: &i32| (core - 1).to_string(); .collect::<Vec<_>>()
.join(" -> ");
let cycle_str = cycle let states_msg = cycle
.iter() .iter()
.map(format_core) .map(&format_state)
.collect::<Vec<_>>() .collect::<Vec<_>>()
.join(" -> "); .join(", ");
return Some(DeadlockInfo {
let cycle = cycle cycle: cycle_msg,
.iter() states: states_msg,
.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;
} }
} }
None None
@@ -446,7 +558,9 @@ fn handle_wait_sync(
InstructionStatus::Sync(data) => { InstructionStatus::Sync(data) => {
let (source, target) = data.get_core_immcore(); let (source, target) = data.get_core_immcore();
let register = data.offset_select() as usize; 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; core_instructions[source as usize].program_counter += 1;
true true
} }
@@ -454,8 +568,8 @@ fn handle_wait_sync(
let core = data.core_indx() as usize; let core = data.core_indx() as usize;
let register = data.offset_select() as usize; let register = data.offset_select() as usize;
let value = data.offset_value(); let value = data.offset_value();
if events[core][register] >= value { if events[core][register].count == value {
events[core][register] -= value; events[core][register] = SyncEvent::default();
core_instructions[core].program_counter += 1; core_instructions[core].program_counter += 1;
true true
} else { } else {
@@ -465,3 +579,66 @@ fn handle_wait_sync(
_ => false, _ => 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());
}
}
@@ -297,7 +297,7 @@ fn multiple_send_recv_test() {
} }
#[test] #[test]
fn sync_wait_tokens_test() { fn sync_wait_exact_count_resets_test() {
let cpu = common::empty_cpu(2); let cpu = common::empty_cpu(2);
let mut cores = CoreInstructionsBuilder::new(2); let mut cores = CoreInstructionsBuilder::new(2);
let mut instructions = InstructionsBuilder::new(); let mut instructions = InstructionsBuilder::new();
@@ -313,14 +313,68 @@ fn sync_wait_tokens_test() {
cores.set_core(1, instructions.build()); cores.set_core(1, instructions.build());
data.set_core_indx(2).fix_core_indx(); data.set_core_indx(2).fix_core_indx();
for _ in 0..2 { instructions.make_inst(wait, data.set_offset_select_value(0, 2).build());
instructions.make_inst(wait, data.set_offset_select_value(0, 1).build());
}
cores.set_core(2, instructions.build()); cores.set_core(2, instructions.build());
Executable::new(cpu, cores.build()).execute().unwrap(); 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] #[test]
fn blocked_transfers_do_not_starve_sync_producer() { fn blocked_transfers_do_not_starve_sync_producer() {
let cpu = common::empty_cpu(4); let cpu = common::empty_cpu(4);
+1 -1
View File
@@ -96,7 +96,7 @@ llvm::cl::opt<bool> pimEmitJson("pim-emit-json",
llvm::cl::opt<bool> pimDetectCommunicationDeadlock( llvm::cl::opt<bool> pimDetectCommunicationDeadlock(
"pim-detect-communication-deadlock", "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::init(false),
llvm::cl::cat(OnnxMlirOptions)); llvm::cl::cat(OnnxMlirOptions));
@@ -373,9 +373,6 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatScheduledCom
rewriter, receiveOp->getLoc(), outputBuffer.getType(), zero, rewriter, receiveOp->getLoc(), outputBuffer.getType(), zero,
hostWaitLoad.getHostOffset(), outputBuffer, *hostBuffer, *sizeAttr) hostWaitLoad.getHostOffset(), outputBuffer, *hostBuffer, *sizeAttr)
.getOutput(); .getOutput();
PimSyncOp::create(
rewriter, receiveOp->getLoc(), hostWaitLoad.getSourceCoreId(),
hostWaitLoad.getAcknowledgementEventRegister());
} else { } else {
received = PimReceiveOp::create( received = PimReceiveOp::create(
rewriter, receiveOp->getLoc(), outputBuffer.getType(), outputBuffer, rewriter, receiveOp->getLoc(), outputBuffer.getType(), outputBuffer,
@@ -152,10 +152,6 @@ struct HostWaitLoadLowering : OpRewritePattern<spatial::SpatHostWaitLoadOp> {
Value output = pim::PimMemCopyHostToDevOp::create( Value output = pim::PimMemCopyHostToDevOp::create(
rewriter, op.getLoc(), outputBuffer.getType(), zero, rewriter, op.getLoc(), outputBuffer.getType(), zero,
op.getHostOffset(), outputBuffer, *hostBuffer, sizeAttr).getOutput(); 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; return output;
}); });
} }
@@ -5,9 +5,11 @@
#include "mlir/Pass/Pass.h" #include "mlir/Pass/Pass.h"
#include "llvm/ADT/STLExtras.h" #include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/Support/FormatVariadic.h" #include "llvm/Support/FormatVariadic.h"
#include "llvm/Support/raw_ostream.h" #include "llvm/Support/raw_ostream.h"
#include <functional>
#include <map> #include <map>
#include <string> #include <string>
@@ -278,7 +280,9 @@ static bool isHostAddressableValue(Value value, const StaticValueKnowledge& know
enum class CommunicationEventKind { enum class CommunicationEventKind {
Send, Send,
Receive Receive,
Sync,
Wait
}; };
struct CommunicationEvent { struct CommunicationEvent {
@@ -286,14 +290,25 @@ struct CommunicationEvent {
int64_t coreId = 0; int64_t coreId = 0;
int64_t peerCoreId = 0; int64_t peerCoreId = 0;
int64_t size = 0; int64_t size = 0;
int64_t eventRegister = 0;
int64_t waitValue = 0;
uint64_t ordinal = 0; uint64_t ordinal = 0;
Operation* op = nullptr; Operation* op = nullptr;
}; };
using CommunicationEventVector = SmallVector<CommunicationEvent, 0>; 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) { 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"; constexpr StringLiteral kRaptorMinChannelIdAttr = "raptor.min_channel_id";
@@ -356,9 +371,15 @@ static std::string formatCommunicationEvent(const CommunicationEvent& event) {
std::string text; std::string text;
llvm::raw_string_ostream os(text); llvm::raw_string_ostream os(text);
os << "core " << event.coreId << " " << getCommunicationEventKindName(event.kind) << " " os << "core " << event.coreId << " " << getCommunicationEventKindName(event.kind);
<< (event.kind == CommunicationEventKind::Send ? "to" : "from") << " " << event.peerCoreId << " size " if (event.kind == CommunicationEventKind::Send || event.kind == CommunicationEventKind::Receive)
<< event.size << "B ordinal " << event.ordinal; 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) if (minChannelId)
os << " min_channel " << *minChannelId; os << " min_channel " << *minChannelId;
if (commOrder) if (commOrder)
@@ -383,6 +404,9 @@ static std::string formatCommunicationEvent(const CommunicationEvent& event) {
} }
static bool areMatchedCommunicationEvents(const CommunicationEvent& lhs, const CommunicationEvent& rhs) { 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) if (lhs.coreId != rhs.peerCoreId || lhs.peerCoreId != rhs.coreId || lhs.size != rhs.size)
return false; return false;
@@ -402,6 +426,27 @@ static void printCounterpartProbe(llvm::raw_ostream& os,
const DenseMap<int64_t, CommunicationEventVector>& coreEvents, const DenseMap<int64_t, CommunicationEventVector>& coreEvents,
const DenseMap<int64_t, size_t>& programCounters, const DenseMap<int64_t, size_t>& programCounters,
const CommunicationEvent& blockedEvent) { 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); auto peerEventsIt = coreEvents.find(blockedEvent.peerCoreId);
if (peerEventsIt == coreEvents.end()) { if (peerEventsIt == coreEvents.end()) {
os << " no local stream was collected for peer core " << blockedEvent.peerCoreId << "\n"; 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( static CommunicationEvent makeCommunicationEvent(
CommunicationEventKind kind, int64_t coreId, int64_t peerCoreId, int64_t size, uint64_t ordinal, Operation* op) { 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, static LogicalResult appendCoreCommunicationEvents(Block& block,
@@ -504,6 +549,39 @@ static LogicalResult appendCoreCommunicationEvents(Block& block,
return success(); 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(); return success();
}); });
} }
@@ -530,7 +608,7 @@ static void printCommunicationWindow(llvm::raw_ostream& os,
static void printCommunicationDeadlockReport(const DenseMap<int64_t, CommunicationEventVector>& coreEvents, static void printCommunicationDeadlockReport(const DenseMap<int64_t, CommunicationEventVector>& coreEvents,
const DenseMap<int64_t, size_t>& programCounters, const DenseMap<int64_t, size_t>& programCounters,
ArrayRef<int64_t> cycle) { 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:"; llvm::errs() << "wait cycle:";
for (int64_t coreId : cycle) for (int64_t coreId : cycle)
llvm::errs() << " " << coreId; llvm::errs() << " " << coreId;
@@ -565,7 +643,7 @@ static void printCommunicationDeadlockReport(const DenseMap<int64_t, Communicati
continue; continue;
printCommunicationWindow(llvm::errs(), coreEvents, coreId, pcIt->second); 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, static void emitCommunicationDeadlockCycle(ModuleOp moduleOp,
@@ -576,8 +654,8 @@ static void emitCommunicationDeadlockCycle(ModuleOp moduleOp,
auto diagnostic = auto diagnostic =
moduleOp.emitError() moduleOp.emitError()
<< "Pim communication deadlock check found a blocking send/receive cycle while statically simulating the " << "Pim communication deadlock check found a blocking SEND/RECV/WAIT cycle while statically simulating the "
"expanded per-core communication streams; see the Pim static communication deadlock report above"; "expanded per-core communication streams; see the static deadlock report above";
for (int64_t coreId : cycle) { for (int64_t coreId : cycle) {
auto eventsIt = coreEvents.find(coreId); auto eventsIt = coreEvents.find(coreId);
@@ -596,46 +674,82 @@ static void emitCommunicationDeadlockCycle(ModuleOp moduleOp,
static FailureOr<SmallVector<int64_t>> static FailureOr<SmallVector<int64_t>>
findCommunicationWaitCycle(const DenseMap<int64_t, CommunicationEventVector>& coreEvents, findCommunicationWaitCycle(const DenseMap<int64_t, CommunicationEventVector>& coreEvents,
const DenseMap<int64_t, size_t>& programCounters) { const DenseMap<int64_t, size_t>& programCounters,
for (const auto& [startCoreId, events] : coreEvents) { const DenseSet<int64_t>& repeatingCores,
auto startPcIt = programCounters.find(startCoreId); const SynchronizationSourceCounts& sourceCounts) {
if (startPcIt == programCounters.end() || startPcIt->second >= events.size()) DenseMap<int64_t, SmallVector<int64_t>> dependencies;
for (const auto& [coreId, events] : coreEvents) {
size_t pc = programCounters.lookup(coreId);
if (pc >= events.size())
continue; continue;
const CommunicationEvent& event = events[pc];
DenseMap<int64_t, size_t> positionInPath; if (event.kind == CommunicationEventKind::Send
SmallVector<int64_t, 8> path; || event.kind == CommunicationEventKind::Receive) {
int64_t currentCoreId = startCoreId; dependencies[coreId].push_back(event.peerCoreId);
while (true) { continue;
auto eventsIt = coreEvents.find(currentCoreId); }
auto pcIt = programCounters.find(currentCoreId); if (event.kind != CommunicationEventKind::Wait)
if (eventsIt == coreEvents.end() || pcIt == programCounters.end() || pcIt->second >= eventsIt->second.size()) continue;
break; int64_t blockedCoreId = coreId;
SynchronizationEventKey eventKey {coreId, event.eventRegister};
auto positionIt = positionInPath.find(currentCoreId); auto contributions = sourceCounts.find(eventKey);
if (positionIt != positionInPath.end()) { for (const auto& [sourceCore, sourceEvents] : coreEvents) {
SmallVector<int64_t> cycle; size_t sourcePc = programCounters.lookup(sourceCore);
for (size_t index = positionIt->second; index < path.size(); ++index) auto matches = [&](const CommunicationEvent& candidate) {
cycle.push_back(path[index]); return candidate.kind == CommunicationEventKind::Sync
return cycle; && candidate.peerCoreId == blockedCoreId
} && candidate.eventRegister == event.eventRegister;
};
positionInPath[currentCoreId] = path.size(); int64_t signalsPerPhase = llvm::count_if(sourceEvents, matches);
path.push_back(currentCoreId); if (repeatingCores.contains(sourceCore))
currentCoreId = eventsIt->second[pcIt->second].peerCoreId; 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(); return failure();
} }
static LogicalResult verifyNoStaticCommunicationDeadlock(ModuleOp moduleOp, static LogicalResult verifyNoStaticCommunicationDeadlock(ModuleOp moduleOp,
pim::CappedDiagnosticReporter& diagnostics) { pim::CappedDiagnosticReporter& diagnostics) {
DenseMap<int64_t, CommunicationEventVector> coreEvents; DenseMap<int64_t, CommunicationEventVector> coreEvents;
DenseSet<int64_t> repeatingCores;
bool hasFailure = false; bool hasFailure = false;
for (func::FuncOp funcOp : moduleOp.getOps<func::FuncOp>()) { for (func::FuncOp funcOp : moduleOp.getOps<func::FuncOp>()) {
if (funcOp.isExternal()) if (funcOp.isExternal())
continue; continue;
bool repeating = funcOp->hasAttr("pim.pipeline_host_buffer_bytes");
for (Operation& op : funcOp.getBody().front().getOperations()) { for (Operation& op : funcOp.getBody().front().getOperations()) {
if (auto coreOp = dyn_cast<pim::PimCoreOp>(&op)) { if (auto coreOp = dyn_cast<pim::PimCoreOp>(&op)) {
@@ -648,6 +762,8 @@ static LogicalResult verifyNoStaticCommunicationDeadlock(ModuleOp moduleOp,
coreEvents[coreId], coreEvents[coreId],
diagnostics))) diagnostics)))
hasFailure = true; hasFailure = true;
if (repeating)
repeatingCores.insert(coreId);
continue; continue;
} }
@@ -668,8 +784,11 @@ static LogicalResult verifyNoStaticCommunicationDeadlock(ModuleOp moduleOp,
coreId, coreId,
laneKnowledge, laneKnowledge,
coreEvents[coreId], coreEvents[coreId],
diagnostics))) diagnostics))) {
hasFailure = true; hasFailure = true;
} else if (repeating) {
repeatingCores.insert(coreId);
}
} }
} }
} }
@@ -678,10 +797,20 @@ static LogicalResult verifyNoStaticCommunicationDeadlock(ModuleOp moduleOp,
if (hasFailure) if (hasFailure)
return failure(); 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; DenseMap<int64_t, size_t> programCounters;
for (const auto& [coreId, events] : coreEvents) for (const auto& [coreId, events] : coreEvents)
programCounters[coreId] = 0; programCounters[coreId] = 0;
DenseMap<SynchronizationEventKey, int64_t> eventCounts;
SynchronizationSourceCounts sourceCounts;
while (true) { while (true) {
bool madeProgress = false; bool madeProgress = false;
for (const auto& [coreId, events] : coreEvents) { for (const auto& [coreId, events] : coreEvents) {
@@ -690,6 +819,35 @@ static LogicalResult verifyNoStaticCommunicationDeadlock(ModuleOp moduleOp,
continue; continue;
const CommunicationEvent& event = events[pc]; 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); auto peerEventsIt = coreEvents.find(event.peerCoreId);
if (peerEventsIt == coreEvents.end()) if (peerEventsIt == coreEvents.end())
continue; continue;
@@ -720,7 +878,8 @@ static LogicalResult verifyNoStaticCommunicationDeadlock(ModuleOp moduleOp,
if (allDone) if (allDone)
return success(); return success();
auto cycle = findCommunicationWaitCycle(coreEvents, programCounters); auto cycle = findCommunicationWaitCycle(
coreEvents, programCounters, repeatingCores, sourceCounts);
if (succeeded(cycle)) { if (succeeded(cycle)) {
emitCommunicationDeadlockCycle(moduleOp, coreEvents, programCounters, *cycle); emitCommunicationDeadlockCycle(moduleOp, coreEvents, programCounters, *cycle);
return failure(); return failure();
@@ -729,7 +888,7 @@ static LogicalResult verifyNoStaticCommunicationDeadlock(ModuleOp moduleOp,
auto diagnostic = auto diagnostic =
moduleOp.emitError() moduleOp.emitError()
<< "Pim communication deadlock check stalled without finding a closed wait cycle; this usually means a " << "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) { for (const auto& [coreId, events] : coreEvents) {
size_t pc = programCounters[coreId]; size_t pc = programCounters[coreId];
if (pc >= events.size()) if (pc >= events.size())
@@ -2,6 +2,8 @@
#include "DeferredCommunicationScheduling.hpp" #include "DeferredCommunicationScheduling.hpp"
#include "DeferredTransferPlanning.hpp" #include "DeferredTransferPlanning.hpp"
#include "llvm/ADT/DenseSet.h"
namespace onnx_mlir::spatial { namespace onnx_mlir::spatial {
using namespace mlir; using namespace mlir;
namespace { namespace {
@@ -256,7 +258,7 @@ static unsigned getBarrierRoundCount(size_t coreCount) {
static LogicalResult assignPipelineSynchronization( static LogicalResult assignPipelineSynchronization(
DeferredTransferPlan &transfers, DeferredTransferPlan &transfers,
ArrayRef<BoundaryProgram> boundaries, MutableArrayRef<BoundaryProgram> boundaries,
size_t synchronizationRegisterCount) { size_t synchronizationRegisterCount) {
bool pipelined = false; bool pipelined = false;
for (ScheduledInfo &scheduled : transfers.scheduled) { for (ScheduledInfo &scheduled : transfers.scheduled) {
@@ -284,14 +286,13 @@ static LogicalResult assignPipelineSynchronization(
}); });
DenseMap<int64_t, SmallVector<HostTransferRef>> incomingByCore; 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>> eventRegisters;
DenseMap<ExternalTransferFamily *, SmallVector<int64_t>> waitValues; DenseMap<ExternalTransferFamily *, SmallVector<int64_t>> waitValues;
DenseMap<ExternalTransferFamily *, SmallVector<int64_t>> acknowledgementRegisters;
auto initialize = [&](ExternalTransferFamily &family) { auto initialize = [&](ExternalTransferFamily &family) {
size_t count = family.targetCores.size(); size_t count = family.targetCores.size();
eventRegisters.try_emplace(&family, count, 0); eventRegisters.try_emplace(&family, count, 0);
waitValues.try_emplace(&family, count, 0); waitValues.try_emplace(&family, count, 0);
acknowledgementRegisters.try_emplace(&family, count, 0);
}; };
for (const BoundaryProgram &boundary : boundaries) for (const BoundaryProgram &boundary : boundaries)
for (const BoundaryInstruction &instruction : boundary.instructions) { for (const BoundaryInstruction &instruction : boundary.instructions) {
@@ -307,32 +308,117 @@ static LogicalResult assignPipelineSynchronization(
int64_t source = family.sourceCores.valueAt(index); int64_t source = family.sourceCores.valueAt(index);
int64_t target = family.targetCores.valueAt(index); int64_t target = family.targetCores.valueAt(index);
incomingByCore[target].push_back({&family, 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( unsigned barrierRounds = getBarrierRoundCount(
transfers.stageZeroCores.size()); transfers.stageZeroCores.size());
bool stageZeroNeedsAcknowledgements = llvm::any_of( auto reservedRegisterCount = [&](int64_t core) -> size_t {
transfers.stageZeroCores, [&](int64_t core) { if (llvm::is_contained(transfers.stageZeroCores, core)) {
return transfers.hostAcknowledgementCounts.contains(core); bool releasesDownstream = !transfers.downstreamCores.empty()
}); && core == transfers.stageZeroCores.front();
for (auto &[target, incoming] : incomingByCore) { return barrierRounds + releasesDownstream;
bool needsAcknowledgementRegister = }
transfers.hostAcknowledgementCounts.contains(target); auto downstream = llvm::find(transfers.downstreamCores, core);
bool stageZero = llvm::is_contained(transfers.stageZeroCores, target); if (downstream == transfers.downstreamCores.end())
size_t reserved = stageZero return 0;
? barrierRounds + (stageZeroNeedsAcknowledgements ? 1 : 0) size_t rank = downstream - transfers.downstreamCores.begin();
: 1 + (needsAcknowledgementRegister ? 1 : 0); bool releasesChildren = 2 * rank + 1 < transfers.downstreamCores.size();
if (reserved >= synchronizationRegisterCount) { return 1 + releasesChildren;
incoming.front().family->requirement->exchange->deferred.emitOpError( };
"pipeline synchronization leaves no event register for incoming host transfers"); 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(); return failure();
} }
size_t groupCount = std::min( size_t available = synchronizationRegisterCount - reserved;
incoming.size(), synchronizationRegisterCount - reserved); size_t readyRegisterCount = std::min(
// One wait consumes a complete consecutive group of producer signals. 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); SmallVector<size_t> groupSizes(groupCount);
for (size_t ordinal = 0; ordinal < incoming.size(); ++ordinal) for (size_t ordinal = 0; ordinal < incoming.size(); ++ordinal)
++groupSizes[ordinal * groupCount / incoming.size()]; ++groupSizes[ordinal * groupCount / incoming.size()];
@@ -341,8 +427,6 @@ static LogicalResult assignPipelineSynchronization(
size_t group = ordinal * groupCount / incoming.size(); size_t group = ordinal * groupCount / incoming.size();
HostTransferRef transfer = incoming[ordinal]; HostTransferRef transfer = incoming[ordinal];
eventRegisters[transfer.family][transfer.index] = group; eventRegisters[transfer.family][transfer.index] = group;
acknowledgementRegisters[transfer.family][transfer.index] =
synchronizationRegisterCount - 1;
if (first[group]) { if (first[group]) {
waitValues[transfer.family][transfer.index] = groupSizes[group]; waitValues[transfer.family][transfer.index] = groupSizes[group];
first[group] = false; first[group] = false;
@@ -352,27 +436,70 @@ static LogicalResult assignPipelineSynchronization(
for (auto &[family, values] : eventRegisters) { for (auto &[family, values] : eventRegisters) {
family->eventRegisters = StaticIntSequence::fromValues(values); family->eventRegisters = StaticIntSequence::fromValues(values);
family->waitValues = StaticIntSequence::fromValues(waitValues[family]); family->waitValues = StaticIntSequence::fromValues(waitValues[family]);
family->acknowledgementEventRegisters =
StaticIntSequence::fromValues(acknowledgementRegisters[family]);
} }
if (!transfers.stageZeroCores.empty()) { SmallVector<int64_t> writers;
size_t reserved = barrierRounds for (const auto &[writer, readers] : readersByWriter)
+ (stageZeroNeedsAcknowledgements ? 1 : 0); writers.push_back(writer);
if (reserved > synchronizationRegisterCount) llvm::sort(writers);
return transfers.scheduled.front().op->emitOpError( for (int64_t writer : writers)
"pipeline stage-zero barrier requires more synchronization registers than the target provides"); for (int64_t reader : readersByWriter[writer]) {
} unsigned group = freeGroups[writer].lookup(reader);
if (!transfers.downstreamCores.empty()) { transfers.hostReleaseSignals[reader].push_back(
bool needsAcknowledgements = llvm::any_of( {writer, freeRegisters[writer][group]});
transfers.downstreamCores, [&](int64_t core) { }
return transfers.hostAcknowledgementCounts.contains(core);
}); DenseSet<std::pair<int64_t, unsigned>> pendingGroups;
if (1 + (needsAcknowledgements ? 1 : 0) for (int64_t writer : writers)
> synchronizationRegisterCount) for (unsigned group = 0; group < freeRegisters[writer].size(); ++group)
return transfers.scheduled.front().op->emitOpError( pendingGroups.insert({writer, group});
"pipeline stage-zero release requires more synchronization registers than the target provides"); 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(); return success();
} }
@@ -34,12 +34,17 @@ struct EmitReceiveAssemblyRun {
llvm::SmallVector<LaneSet> entryLanes; llvm::SmallVector<LaneSet> entryLanes;
LaneSet lanes; LaneSet lanes;
}; };
struct EmitHostReuseWait {
LaneSet lanes;
StaticIntSequence eventRegisters = StaticIntSequence::uniform(0, 1);
StaticIntSequence waitValues = StaticIntSequence::uniform(0, 1);
};
struct ProduceDeferredResult { struct ProduceDeferredResult {
DeferredExchangePlan* exchange = nullptr; DeferredExchangePlan* exchange = nullptr;
}; };
using BoundaryInstruction = using BoundaryInstruction =
std::variant<EmitSendRun, EmitLocalCollectionRun, std::variant<EmitSendRun, EmitHostReuseWait, EmitLocalCollectionRun,
EmitLocalCollectionLoopRun, EmitReceiveAssemblyRun, EmitLocalCollectionLoopRun, EmitReceiveAssemblyRun,
ProduceDeferredResult>; ProduceDeferredResult>;
struct BoundaryProgram { struct BoundaryProgram {
@@ -10,6 +10,7 @@
#include "src/Accelerators/PIM/Common/IR/StaticIntSequence.hpp" #include "src/Accelerators/PIM/Common/IR/StaticIntSequence.hpp"
#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp" #include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
#include "src/Accelerators/PIM/Common/PimCommon.hpp" #include "src/Accelerators/PIM/Common/PimCommon.hpp"
#include "llvm/ADT/DenseSet.h"
#include <array> #include <array>
namespace onnx_mlir::spatial { namespace onnx_mlir::spatial {
using namespace mlir; using namespace mlir;
@@ -23,7 +24,6 @@ struct LogicalTransferMetadataView {
StaticIntSequenceChain hostOffsets; StaticIntSequenceChain hostOffsets;
StaticIntSequenceChain eventRegisters; StaticIntSequenceChain eventRegisters;
StaticIntSequenceChain waitValues; StaticIntSequenceChain waitValues;
StaticIntSequenceChain acknowledgementEventRegisters;
StaticIntSequenceChain targetLanes; StaticIntSequenceChain targetLanes;
StaticIntSequenceChain localOffsets; StaticIntSequenceChain localOffsets;
SmallVector<StaticIntSequenceChain> projectionOffsets; 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); 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(); 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) { static void appendMetadata(const ScheduledTransferSlice &slice, LogicalTransferMetadataView &metadata) {
ExternalTransferFamily &family = *slice.family; ExternalTransferFamily &family = *slice.family;
LaneInterval familyLanes = family.targetLanes.intervals().front(); LaneInterval familyLanes = family.targetLanes.intervals().front();
@@ -95,8 +118,6 @@ static void appendMetadata(const ScheduledTransferSlice &slice, LogicalTransferM
metadata.eventRegisters.append( metadata.eventRegisters.append(
family.eventRegisters, familyIndex, count); family.eventRegisters, familyIndex, count);
metadata.waitValues.append(family.waitValues, familyIndex, count); metadata.waitValues.append(family.waitValues, familyIndex, count);
metadata.acknowledgementEventRegisters.append(
family.acknowledgementEventRegisters, familyIndex, count);
} }
metadata.targetLanes.append(StaticIntSequence::affine(targetLane, 1, count)); metadata.targetLanes.append(StaticIntSequence::affine(targetLane, 1, count));
if (family.requirement->producerLocalOffsets) if (family.requirement->producerLocalOffsets)
@@ -303,20 +324,15 @@ static FailureOr<Value> emitReceiveValue(ArrayRef<ScheduledTransferSlice> slices
std::optional<StaticIntGrid> hostOffsets; std::optional<StaticIntGrid> hostOffsets;
std::optional<StaticIntGrid> eventRegisters; std::optional<StaticIntGrid> eventRegisters;
std::optional<StaticIntGrid> waitValues; std::optional<StaticIntGrid> waitValues;
std::optional<StaticIntGrid> acknowledgementEventRegisters;
if (slices.front().family->hostRouted) { if (slices.front().family->hostRouted) {
auto offsets = buildGrid(metadata.hostOffsets); auto offsets = buildGrid(metadata.hostOffsets);
auto events = buildGrid(metadata.eventRegisters); auto events = buildGrid(metadata.eventRegisters);
auto waits = buildGrid(metadata.waitValues); auto waits = buildGrid(metadata.waitValues);
auto acknowledgements = buildGrid( if (failed(offsets) || failed(events) || failed(waits))
metadata.acknowledgementEventRegisters);
if (failed(offsets) || failed(events) || failed(waits)
|| failed(acknowledgements))
return failure(); return failure();
hostOffsets = std::move(*offsets); hostOffsets = std::move(*offsets);
eventRegisters = std::move(*events); eventRegisters = std::move(*events);
waitValues = std::move(*waits); waitValues = std::move(*waits);
acknowledgementEventRegisters = std::move(*acknowledgements);
} }
Value position = lane ? lane : context.constants.getIndex(0); Value position = lane ? lane : context.constants.getIndex(0);
Value row = context.constants.getIndex(0); Value row = context.constants.getIndex(0);
@@ -335,8 +351,6 @@ static FailureOr<Value> emitReceiveValue(ArrayRef<ScheduledTransferSlice> slices
eventRegisters->emitLookup( eventRegisters->emitLookup(
row, position, anchor, context.constants, context.rewriter, anchor->getLoc()), row, position, anchor, context.constants, context.rewriter, anchor->getLoc()),
waitValues->emitLookup( waitValues->emitLookup(
row, position, anchor, context.constants, context.rewriter, anchor->getLoc()),
acknowledgementEventRegisters->emitLookup(
row, position, anchor, context.constants, context.rewriter, anchor->getLoc())); row, position, anchor, context.constants, context.rewriter, anchor->getLoc()));
receive = op; receive = op;
output = op.getOutput(); output = op.getOutput();
@@ -406,7 +420,6 @@ static FailureOr<Value> emitReceiveAssembly(const EmitReceiveAssemblyRun &run, V
std::optional<StaticIntGrid> hostOffsets; std::optional<StaticIntGrid> hostOffsets;
std::optional<StaticIntGrid> eventRegisters; std::optional<StaticIntGrid> eventRegisters;
std::optional<StaticIntGrid> waitValues; std::optional<StaticIntGrid> waitValues;
std::optional<StaticIntGrid> acknowledgementEventRegisters;
bool hostRouted = run.slices.front().family->hostRouted; bool hostRouted = run.slices.front().family->hostRouted;
auto metadataByEntry = buildRectangularReceiveMetadata(run, laneCount); auto metadataByEntry = buildRectangularReceiveMetadata(run, laneCount);
if (succeeded(metadataByEntry)) { if (succeeded(metadataByEntry)) {
@@ -424,15 +437,11 @@ static FailureOr<Value> emitReceiveAssembly(const EmitReceiveAssemblyRun &run, V
&LogicalTransferMetadataView::eventRegisters); &LogicalTransferMetadataView::eventRegisters);
auto waits = buildRows( auto waits = buildRows(
&LogicalTransferMetadataView::waitValues); &LogicalTransferMetadataView::waitValues);
auto acknowledgements = buildRows( if (failed(offsets) || failed(events) || failed(waits))
&LogicalTransferMetadataView::acknowledgementEventRegisters);
if (failed(offsets) || failed(events) || failed(waits)
|| failed(acknowledgements))
return failure(); return failure();
hostOffsets = std::move(*offsets); hostOffsets = std::move(*offsets);
eventRegisters = std::move(*events); eventRegisters = std::move(*events);
waitValues = std::move(*waits); waitValues = std::move(*waits);
acknowledgementEventRegisters = std::move(*acknowledgements);
} }
SmallVector<StaticIntSequence> positionRows; SmallVector<StaticIntSequence> positionRows;
for (unsigned position : run.positions) for (unsigned position : run.positions)
@@ -485,15 +494,11 @@ static FailureOr<Value> emitReceiveAssembly(const EmitReceiveAssemblyRun &run, V
&LogicalTransferMetadataView::eventRegisters); &LogicalTransferMetadataView::eventRegisters);
auto waits = buildGrid( auto waits = buildGrid(
&LogicalTransferMetadataView::waitValues); &LogicalTransferMetadataView::waitValues);
auto acknowledgements = buildGrid( if (failed(offsets) || failed(events) || failed(waits))
&LogicalTransferMetadataView::acknowledgementEventRegisters);
if (failed(offsets) || failed(events) || failed(waits)
|| failed(acknowledgements))
return failure(); return failure();
hostOffsets = std::move(*offsets); hostOffsets = std::move(*offsets);
eventRegisters = std::move(*events); eventRegisters = std::move(*events);
waitValues = std::move(*waits); waitValues = std::move(*waits);
acknowledgementEventRegisters = std::move(*acknowledgements);
} }
SmallVector<StaticIntSequence> positionColumns; SmallVector<StaticIntSequence> positionColumns;
for (const StaticIntSequenceChain &values : positionsByLane) for (const StaticIntSequenceChain &values : positionsByLane)
@@ -527,8 +532,6 @@ static FailureOr<Value> emitReceiveAssembly(const EmitReceiveAssemblyRun &run, V
eventRegisters->emitLookup( eventRegisters->emitLookup(
entry, runtimeLane, anchor, context.constants, context.rewriter, loc), entry, runtimeLane, anchor, context.constants, context.rewriter, loc),
waitValues->emitLookup( waitValues->emitLookup(
entry, runtimeLane, anchor, context.constants, context.rewriter, loc),
acknowledgementEventRegisters->emitLookup(
entry, runtimeLane, anchor, context.constants, context.rewriter, loc)); entry, runtimeLane, anchor, context.constants, context.rewriter, loc));
receive = op; receive = op;
output = op.getOutput(); 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, 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; SmallVector<Value> produced;
for (size_t instructionIndex = 0; for (size_t instructionIndex = 0;
instructionIndex < instructions.size(); ++instructionIndex) { 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))) if (failed(emitConditionalSendRun(*send, lane, laneCount, context)))
return failure(); return failure();
} else if (auto update = std::get_if<EmitLocalCollectionRun>(&instruction)) { } 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)) if (auto batch = dyn_cast<SpatScheduledComputeBatch>(boundary.key.first->op))
lane = *batch.getLaneArgument(); lane = *batch.getLaneArgument();
SmallVector<DeferredExchangePlan *> exchanges = getProducedExchanges(boundary.instructions); 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); return failed(values) ? failure() : replaceResults(exchanges, *values, replacements);
} }
@@ -1203,30 +1224,129 @@ static unsigned getBarrierRoundCount(size_t coreCount) {
return rounds; 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( static LogicalResult emitCompletionSynchronization(
DeferredTransferPlan &transfers, DeferredEmissionContext &context) { DeferredTransferPlan &transfers, DeferredEmissionContext &context) {
if (transfers.synchronizationRegisterCount == 0) if (transfers.synchronizationRegisterCount == 0)
return success(); return success();
size_t acknowledgementRegister = size_t restartRegister =
transfers.synchronizationRegisterCount - 1; transfers.synchronizationRegisterCount - 1;
unsigned barrierRounds = getBarrierRoundCount( unsigned barrierRounds = getBarrierRoundCount(
transfers.stageZeroCores.size()); transfers.stageZeroCores.size());
bool stageZeroNeedsAcknowledgements = llvm::any_of( size_t firstBarrierRegister = restartRegister;
transfers.stageZeroCores, [&](int64_t core) {
return transfers.hostAcknowledgementCounts.contains(core);
});
size_t firstBarrierRegister = acknowledgementRegister
- (stageZeroNeedsAcknowledgements ? 1 : 0);
DenseMap<int64_t, unsigned> stageZeroRank; DenseMap<int64_t, unsigned> stageZeroRank;
for (auto [rank, core] : llvm::enumerate(transfers.stageZeroCores)) for (auto [rank, core] : llvm::enumerate(transfers.stageZeroCores))
stageZeroRank[core] = rank; stageZeroRank[core] = rank;
DenseMap<int64_t, unsigned> downstreamRank; DenseMap<int64_t, unsigned> downstreamRank;
for (auto [rank, core] : llvm::enumerate(transfers.downstreamCores)) for (auto [rank, core] : llvm::enumerate(transfers.downstreamCores))
downstreamRank[core] = rank; downstreamRank[core] = rank;
auto getReleaseRegister = [&](int64_t core) {
return acknowledgementRegister
- (transfers.hostAcknowledgementCounts.contains(core) ? 1 : 0);
};
for (ScheduledInfo &scheduled : transfers.scheduled) { for (ScheduledInfo &scheduled : transfers.scheduled) {
Block *block = scheduled.blocks.front(); Block *block = scheduled.blocks.front();
@@ -1236,13 +1356,10 @@ static LogicalResult emitCompletionSynchronization(
if (auto batch = dyn_cast<SpatScheduledComputeBatch>(scheduled.op)) if (auto batch = dyn_cast<SpatScheduledComputeBatch>(scheduled.op))
lane = *batch.getLaneArgument(); lane = *batch.getLaneArgument();
SmallVector<int64_t> acknowledgementCounts, releaseRegisters; SmallVector<int64_t> leftTargets, rightTargets, childCounts;
SmallVector<int64_t> releaseWaitValues, leftTargets, leftRegisters; LaneSet barrierLanes, leaderLanes, downstreamLanes, parentLanes,
SmallVector<int64_t> rightTargets, rightRegisters; leftLanes, rightLanes;
LaneSet barrierLanes, leaderLanes, leftLanes, rightLanes;
for (auto [index, core] : llvm::enumerate(scheduled.cores)) { for (auto [index, core] : llvm::enumerate(scheduled.cores)) {
acknowledgementCounts.push_back(
transfers.hostAcknowledgementCounts.lookup(core));
if (stageZeroRank.contains(core)) if (stageZeroRank.contains(core))
barrierLanes = barrierLanes.unite( barrierLanes = barrierLanes.unite(
LaneSet::range(index, index + 1)); LaneSet::range(index, index + 1));
@@ -1252,66 +1369,37 @@ static LogicalResult emitCompletionSynchronization(
auto rank = downstreamRank.find(core); auto rank = downstreamRank.find(core);
if (rank == downstreamRank.end()) { if (rank == downstreamRank.end()) {
releaseRegisters.push_back(0);
releaseWaitValues.push_back(0);
leftTargets.push_back(core); leftTargets.push_back(core);
leftRegisters.push_back(0);
rightTargets.push_back(core); rightTargets.push_back(core);
rightRegisters.push_back(0); childCounts.push_back(0);
continue; continue;
} }
releaseRegisters.push_back(getReleaseRegister(core)); downstreamLanes = downstreamLanes.unite(
releaseWaitValues.push_back(1); LaneSet::range(index, index + 1));
size_t left = 2 * rank->second + 1; size_t left = 2 * rank->second + 1;
size_t right = left + 1; size_t right = left + 1;
unsigned childCount = 0;
if (left < transfers.downstreamCores.size()) { if (left < transfers.downstreamCores.size()) {
int64_t child = transfers.downstreamCores[left]; int64_t child = transfers.downstreamCores[left];
leftTargets.push_back(child); leftTargets.push_back(child);
leftRegisters.push_back(getReleaseRegister(child));
leftLanes = leftLanes.unite(LaneSet::range(index, index + 1)); leftLanes = leftLanes.unite(LaneSet::range(index, index + 1));
++childCount;
} else { } else {
leftTargets.push_back(core); leftTargets.push_back(core);
leftRegisters.push_back(0);
} }
if (right < transfers.downstreamCores.size()) { if (right < transfers.downstreamCores.size()) {
int64_t child = transfers.downstreamCores[right]; int64_t child = transfers.downstreamCores[right];
rightTargets.push_back(child); rightTargets.push_back(child);
rightRegisters.push_back(getReleaseRegister(child));
rightLanes = rightLanes.unite(LaneSet::range(index, index + 1)); rightLanes = rightLanes.unite(LaneSet::range(index, index + 1));
++childCount;
} else { } else {
rightTargets.push_back(core); 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); 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. // Dissemination barrier: every round doubles the covered stage-zero peers.
auto emitBarrier = [&]() { auto emitBarrier = [&]() {
@@ -1341,45 +1429,64 @@ static LogicalResult emitCompletionSynchronization(
} }
}; };
if (barrierRounds > 0 if (barrierRounds > 0
&& failed(emitForLanes(barrierLanes, emitBarrier))) && failed(emitForLanes(
barrierLanes, lane, scheduled.cores.size(), scheduled.op, context,
loc, emitBarrier)))
return failure(); return failure();
// Gate downstream restarts so no core advances the simulator input // Gate downstream restarts so no core advances the simulator input
// iteration ahead of stage zero. // iteration ahead of stage zero.
if (!transfers.downstreamCores.empty() 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(); int64_t root = transfers.downstreamCores.front();
SpatSyncOp::create( SpatSyncOp::create(
context.rewriter, loc, context.constants.getIndex(root), context.rewriter, loc, context.constants.getIndex(root),
context.constants.getIndex(getReleaseRegister(root))); context.constants.getIndex(restartRegister));
}))) })))
return failure(); return failure();
Value releaseRegister = emitStaticIntLookup( if (failed(emitForLanes(
StaticIntSequence::fromValues(releaseRegisters), runtimeLane, downstreamLanes, lane, scheduled.cores.size(), scheduled.op,
scheduled.op, context.constants, context.rewriter, loc); context, loc, [&]() {
Value releaseWaitValue = emitStaticIntLookup( SpatWaitOp::create(
StaticIntSequence::fromValues(releaseWaitValues), runtimeLane, context.rewriter, loc,
scheduled.op, context.constants, context.rewriter, loc); context.constants.getIndex(restartRegister),
SpatWaitOp::create( context.constants.getIndex(1));
context.rewriter, loc, releaseRegister, releaseWaitValue); })))
return failure();
auto emitChild = [&](ArrayRef<int64_t> targets, if (failed(emitForLanes(
ArrayRef<int64_t> registers) { 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( Value target = emitStaticIntLookup(
StaticIntSequence::fromValues(targets), runtimeLane, scheduled.op, StaticIntSequence::fromValues(targets), runtimeLane, scheduled.op,
context.constants, context.rewriter, loc); context.constants, context.rewriter, loc);
Value eventRegister = emitStaticIntLookup( SpatSyncOp::create(
StaticIntSequence::fromValues(registers), runtimeLane, scheduled.op, context.rewriter, loc, target,
context.constants, context.rewriter, loc); context.constants.getIndex(restartRegister));
SpatSyncOp::create(context.rewriter, loc, target, eventRegister);
}; };
if (failed(emitForLanes(leftLanes, [&]() { if (failed(emitForLanes(
emitChild(leftTargets, leftRegisters); leftLanes, lane, scheduled.cores.size(), scheduled.op, context, loc,
})) [&]() { emitChild(leftTargets); }))
|| failed(emitForLanes(rightLanes, [&]() { || failed(emitForLanes(
emitChild(rightTargets, rightRegisters); rightLanes, lane, scheduled.cores.size(), scheduled.op, context,
}))) loc, [&]() { emitChild(rightTargets); })))
return failure(); return failure();
} }
return success(); return success();
@@ -1399,6 +1506,10 @@ LogicalResult realizeDeferredBoundaries(ArrayRef<BoundaryProgram> boundaries, Ar
if (failed(emitBoundary(boundary, results, context, replacements))) if (failed(emitBoundary(boundary, results, context, replacements)))
return boundary.key.first->op->emitOpError("phase 2 failed to realize a communication boundary"); 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); return emitCompletionSynchronization(transfers, context);
} }
@@ -237,8 +237,6 @@ struct ExternalTransferFamily {
StaticIntSequence hostOffsets = StaticIntSequence::uniform(0, 1); StaticIntSequence hostOffsets = StaticIntSequence::uniform(0, 1);
StaticIntSequence eventRegisters = StaticIntSequence::uniform(0, 1); StaticIntSequence eventRegisters = StaticIntSequence::uniform(0, 1);
StaticIntSequence waitValues = StaticIntSequence::uniform(1, 1); StaticIntSequence waitValues = StaticIntSequence::uniform(1, 1);
StaticIntSequence acknowledgementEventRegisters =
StaticIntSequence::uniform(0, 1);
bool hostRouted = false; bool hostRouted = false;
}; };
@@ -7,6 +7,11 @@
namespace onnx_mlir::spatial { namespace onnx_mlir::spatial {
struct HostReleaseSignal {
int64_t writerCore = -1;
unsigned eventRegister = 0;
};
struct DeferredTransferPlan { struct DeferredTransferPlan {
std::vector<size_t> processorStages; std::vector<size_t> processorStages;
llvm::SmallVector<ScheduledInfo, 0> scheduled; llvm::SmallVector<ScheduledInfo, 0> scheduled;
@@ -14,7 +19,8 @@ struct DeferredTransferPlan {
llvm::DenseMap<int64_t, llvm::SmallVector<ProducedValue*>> producedByGraph; llvm::DenseMap<int64_t, llvm::SmallVector<ProducedValue*>> producedByGraph;
llvm::SmallVector<std::unique_ptr<DeferredExchangePlan>> exchanges; llvm::SmallVector<std::unique_ptr<DeferredExchangePlan>> exchanges;
llvm::SmallVector<unsigned> stepCounts; 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> stageZeroCores;
llvm::SmallVector<int64_t> downstreamCores; llvm::SmallVector<int64_t> downstreamCores;
size_t synchronizationRegisterCount = 0; size_t synchronizationRegisterCount = 0;
+3 -4
View File
@@ -592,15 +592,14 @@ def SpatHostStoreSyncOp : SpatOp<"host_store_sync", []> {
} }
def SpatHostWaitLoadOp : SpatOp<"host_wait_load", []> { 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 let arguments = (ins
Index:$sourceCoreId, Index:$sourceCoreId,
Index:$targetCoreId, Index:$targetCoreId,
Index:$hostOffset, Index:$hostOffset,
Index:$eventRegister, Index:$eventRegister,
Index:$waitValue, Index:$waitValue
Index:$acknowledgementEventRegister
); );
let results = (outs let results = (outs
@@ -610,7 +609,7 @@ def SpatHostWaitLoadOp : SpatOp<"host_wait_load", []> {
let assemblyFormat = [{ let assemblyFormat = [{
`from` $sourceCoreId `to` $targetCoreId `from` $sourceCoreId `to` $targetCoreId
`host_offset` $hostOffset `event` $eventRegister `count` $waitValue `host_offset` $hostOffset `event` $eventRegister `count` $waitValue
`ack` $acknowledgementEventRegister attr-dict `:` type($output) attr-dict `:` type($output)
}]; }];
} }
+10
View File
@@ -36,3 +36,13 @@ add_pim_unittest(SpatialSchedulingTargetTest
LINK_LIBS PRIVATE LINK_LIBS PRIVATE
OMPimCompilerUtils 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)
+5 -2
View File
@@ -75,7 +75,7 @@ PIMSIM_FAILED = "ERROR"
PIMSIM_UNSUPPORTED = "UNSUPPORTED" PIMSIM_UNSUPPORTED = "UNSUPPORTED"
PIMSIM_SKIPPED = "SKIP" PIMSIM_SKIPPED = "SKIP"
PIMSIM_NOT_RUN = "-" PIMSIM_NOT_RUN = "-"
PIMSIM_UNSUPPORTED_VSOFTMAX = "Pimsim does not support opcode vsoftmax" PIMSIM_UNSUPPORTED_VSOFTMAX = "does not support opcode vsoftmax"
class PimSimUnsupportedError(RuntimeError): class PimSimUnsupportedError(RuntimeError):
@@ -563,7 +563,10 @@ def validate_execution(
"Run non-functional simulation", name, "Run non-functional simulation", name,
) )
config_path = execution["pimsim_config"] 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: try:
state["metrics"] = run_pimsim_nn( state["metrics"] = run_pimsim_nn(
pimsim_nn_build_dir, pim_dir, config_path, name, pimsim_nn_build_dir, pim_dir, config_path, name,
@@ -16,6 +16,7 @@ import sys
import time import time
import types import types
from collections import Counter from collections import Counter
from concurrent.futures import ThreadPoolExecutor
from dataclasses import asdict, dataclass from dataclasses import asdict, dataclass
from pathlib import Path from pathlib import Path
from tempfile import TemporaryDirectory from tempfile import TemporaryDirectory
@@ -242,7 +243,7 @@ def reference_inputs_exist(
def reference_batch_dirs(root: Path, batch_size: int) -> list[Path]: 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( def reference_batch_outputs_exist(
@@ -442,6 +443,7 @@ def generate_reference_batch_outputs(
runner_build_dir: Path, runner_build_dir: Path,
model_path: Path, model_path: Path,
input_batch: list[list[np.ndarray]], input_batch: list[list[np.ndarray]],
outputs_desc: list[tuple[int, str, int, list[int]]],
steps: list[StepRecord], steps: list[StepRecord],
args: argparse.Namespace, args: argparse.Namespace,
out_dir: Path, out_dir: Path,
@@ -450,20 +452,26 @@ def generate_reference_batch_outputs(
) -> list[Path]: ) -> list[Path]:
if print_header: if print_header:
print_step("Run reference") print_step("Run reference")
references = [] references = reference_batch_dirs(out_dir, len(input_batch))
for index, sample in enumerate(input_batch): missing = [
references.append( index for index, reference in enumerate(references)
generate_reference_outputs( if not reference_outputs_exist(outputs_desc, reference)
runner_path, ]
runner_build_dir,
model_path, def generate(index: int) -> None:
sample, generate_reference_outputs(
steps, runner_path,
args, runner_build_dir,
out_dir / f"batch_{index:06d}", model_path,
print_header=False, 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 return references
@@ -490,6 +498,7 @@ def prepare_reference_batch_outputs(
runner_build_dir, runner_build_dir,
model_path, model_path,
input_batch, input_batch,
outputs_desc,
steps, steps,
args, args,
out_dir, out_dir,
@@ -512,16 +521,29 @@ def prepare_common_artifacts(
outputs_ready = reference_outputs_exist(outputs_desc, outputs_dir) outputs_ready = reference_outputs_exist(outputs_desc, outputs_dir)
if inputs_ready: if inputs_ready:
arrays_in_order = load_saved_inputs(inputs_desc, inputs_dir) 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): if not (inputs_ready and outputs_ready):
generate_reference_outputs( shutil.copytree(
runner_path, references[0].parent / "inputs", inputs_dir, dirs_exist_ok=True)
runner_path.parent, shutil.copytree(references[0], outputs_dir, dirs_exist_ok=True)
model_path,
arrays_in_order,
steps,
args,
common_dir,
)
def compile_raptor_target( def compile_raptor_target(
model_path: Path, model_path: Path,
out_dir: Path, out_dir: Path,
@@ -1409,7 +1431,14 @@ def main():
parser.add_argument("--mesh-cols", type=int) parser.add_argument("--mesh-cols", type=int)
parser.add_argument("--pimsim-time-ms", type=int, default=1000) parser.add_argument("--pimsim-time-ms", type=int, default=1000)
parser.add_argument("--pimsim-mode", choices=["latency", "throughput"], default="latency") 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-pipeline", choices=["element", "batch"])
parser.add_argument("--pimcomp-model-name", help="Use a Pimcomp built-in model name such as vgg16.") parser.add_argument("--pimcomp-model-name", help="Use a Pimcomp built-in model name such as vgg16.")
parser.add_argument( parser.add_argument(
@@ -1466,6 +1495,8 @@ def main():
parser.error("--pimsim-time-ms must be positive") parser.error("--pimsim-time-ms must be positive")
if args.batch_size <= 0: if args.batch_size <= 0:
parser.error("--batch-size must be positive") 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: if args.pimsim_mode == "throughput" and args.batch_size < 2:
parser.error("throughput mode requires batch size greater than 1") parser.error("throughput mode requires batch size greater than 1")
if args.timeout_seconds < 0: if args.timeout_seconds < 0:
@@ -1674,7 +1705,7 @@ def main():
write_input_batch_csv(out_dir / "inputs.csv", input_batch) write_input_batch_csv(out_dir / "inputs.csv", input_batch)
raptor_input_bins = write_input_batch_binaries(input_batch, out_dir / "simulation/raptor_inputs") raptor_input_bins = write_input_batch_binaries(input_batch, out_dir / "simulation/raptor_inputs")
if args.pimsim_mode == "throughput": 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( throughput_references = try_stage(
failures, failures,
"Run reference", "Run reference",
@@ -261,23 +261,9 @@ def performance_values(performance: dict) -> dict[str, float | None]:
} }
def comparison_passed(report: dict, compiler: str | None = None) -> bool: def comparison_passed(report: dict) -> bool:
other_compiler = "PIMCOMP" if compiler == "raptor" else "RAPTOR" result = report.get("raptor_validation") or {}
if any( return result.get("status") == "done" and bool(result.get("passed"))
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 functional_validation_status(result: dict | None) -> str: 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], *[f"--raptor-extra-arg={arg}" for arg in raptor_extra_args],
"--timeout-seconds", "--timeout-seconds",
str(timeout), str(timeout),
"--fail-on-error",
*([] if fast else ["--no-fast"]), *([] if fast else ["--no-fast"]),
*reuse_args, *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 [ return [
sys.executable, sys.executable,
str(COMPARE), str(COMPARE),
@@ -421,6 +412,10 @@ def prepare_common_command(model: Path, common_dir: Path, timeout: float) -> lis
"--common-dir", "--common-dir",
str(common_dir), str(common_dir),
"--prepare-common", "--prepare-common",
"--batch-size",
str(batch_size),
"--jobs",
str(jobs),
"--timeout-seconds", "--timeout-seconds",
str(timeout), str(timeout),
] ]
@@ -587,8 +582,8 @@ def main() -> int:
parser.add_argument( parser.add_argument(
"--batch-size", "--batch-size",
type=int, type=int,
default=128, default=64,
help="functional throughput batch size (default: 128).", help="functional throughput and shared reference batch size (default: 64).",
) )
parser.add_argument( parser.add_argument(
"--timeout-seconds", "--timeout-seconds",
@@ -664,6 +659,11 @@ def main() -> int:
} }
comparisons_by_arch[arch] = comparisons comparisons_by_arch[arch] = comparisons
configs_by_arch[arch] = configs 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 = [ missing = [
str(path) str(path)
@@ -706,6 +706,7 @@ def main() -> int:
print(f"Modes: {', '.join(args.mode)}") print(f"Modes: {', '.join(args.mode)}")
print(f"Throughput Pimsim time: {args.pimsim_time_ms} ms") print(f"Throughput Pimsim time: {args.pimsim_time_ms} ms")
print(f"Max parallel jobs: {args.jobs}") print(f"Max parallel jobs: {args.jobs}")
print(f"Shared reference batch: {reference_batch_size}")
print( print(
f"Comparison jobs: " f"Comparison jobs: "
f"{sum(len(args.models) * len(comparisons) for comparisons in comparisons_by_arch.values())}" f"{sum(len(args.models) * len(comparisons) for comparisons in comparisons_by_arch.values())}"
@@ -722,6 +723,8 @@ def main() -> int:
FUNCTIONAL_MODELS[name], FUNCTIONAL_MODELS[name],
common_dir(out_dir, name, common_root), common_dir(out_dir, name, common_root),
args.timeout_seconds, args.timeout_seconds,
reference_batch_size,
args.jobs,
), ),
dry_run=args.dry_run, dry_run=args.dry_run,
) )
@@ -849,10 +852,8 @@ def main() -> int:
report_path = result_dir( report_path = result_dir(
out_dir, name, arch, mode, pipeline, args.ablation_variant out_dir, name, arch, mode, pipeline, args.ablation_variant
) / "pimcomp/comparison_report.json" ) / "pimcomp/comparison_report.json"
compiler = "raptor" if args.raptor_only else args.only
if not report_path.exists() or not comparison_passed( if not report_path.exists() or not comparison_passed(
json.loads(report_path.read_text(encoding="utf-8")), json.loads(report_path.read_text(encoding="utf-8")),
compiler,
): ):
if label not in failed: if label not in failed:
failed.append(label) failed.append(label)