diff --git a/crates/ipu-codegen/src/package/benchmark.rs b/crates/ipu-codegen/src/package/benchmark.rs index a0aaf95..8c750bb 100644 --- a/crates/ipu-codegen/src/package/benchmark.rs +++ b/crates/ipu-codegen/src/package/benchmark.rs @@ -145,6 +145,74 @@ pub fn benchmark_mid_expansion( Arc::clone(&cache), &mut analysis, )?; + fn audit(graph: &crate::low::TileGraph, region: &crate::low::BlockRegion) { + use crate::low::{BlockOperation as Op, storage_root}; + use std::collections::BTreeSet; + let access = |id: crate::low::ExchangePhaseId| { + let phase = &graph.exchange_phases[id.index() as usize]; + let reads = phase.transfers.iter().map(|t| storage_root(&graph.shards, t.source.shard)).collect::>(); + let writes = phase.transfers.iter().flat_map(|t| &t.destinations).map(|v| storage_root(&graph.shards, v.shard)).collect::>(); + (reads, writes) + }; + let mut previous = None; + for (index, op) in region.operations.iter().enumerate() { + match op { + Op::Repeat(r) => { audit(graph, &r.body); previous = None; }, + Op::Exchange(next) => { + if let Some((before, at)) = previous { + let between = ®ion.operations[at+1..index]; + if between.iter().all(|op| matches!(op, Op::Copy {..})) { + let (ar, aw) = access(before); + let (br, bw) = access(*next); + let mut copies = Vec::new(); + for op in between { + let Op::Copy { tile, copy } = op else { unreachable!() }; + let c = &graph.local_copies[copy.0 as usize]; + let r = storage_root(&graph.shards, c.source); + let w = storage_root(&graph.shards, c.destination); + let (rows, bytes, ss, ds) = match c.pattern { + crate::low::CopyPattern::Contiguous => (1, c.bytes, 0, 0), + crate::low::CopyPattern::Strided { rows,row_bytes,source_stride,destination_stride } => (rows,row_bytes,source_stride,destination_stride), + }; + let overlaps = |root, offset, stride, view: &crate::low::ShardView, order| { + if storage_root(&graph.shards, view.shard) != root { return false; } + let spans = crate::low::view_byte_traversal(&graph.shards[view.shard.index() as usize], view, order).unwrap(); + spans.spans().any(|span| (0..rows).any(|row| { + let start = offset + row*stride; + start < span.offset+span.bytes && span.offset < start+bytes + })) + }; + let hazards = |phase: crate::low::ExchangePhaseId| { + let mut raw=false; let mut war=false; let mut waw=false; + for t in &graph.exchange_phases[phase.index() as usize].transfers { + let order = t.span_order(&graph.shards); + war |= overlaps(w,c.destination_offset,ds,&t.source,order); + for v in &t.destinations { + raw |= overlaps(r,c.source_offset,ss,v,order); + waw |= overlaps(w,c.destination_offset,ds,v,order); + } + } + (raw,war,waw) + }; + let ah=hazards(before); let bh=hazards(*next); + copies.push(serde_json::json!({"tile":tile,"copy":format!("{:?}",c), + "source":format!("{:?}",graph.shards[r.index() as usize].definition), + "destination":format!("{:?}",graph.shards[w.index() as usize].definition), + "precise_before":ah,"precise_after":bh,"hoist":!aw.contains(&r)&&!ar.contains(&w)&&!aw.contains(&w), + "sink":!bw.contains(&r)&&!br.contains(&w)&&!bw.contains(&w)})); + } + eprintln!("COPY_BOUNDARY {}",serde_json::json!({"before":before.index(),"after":next.index(), + "previous":format!("{:?}",graph.exchange_phases[before.index() as usize].provenance), + "next":format!("{:?}",graph.exchange_phases[next.index() as usize].provenance),"copies":copies})); + } + } + previous = Some((*next, index)); + }, + _ => {} + } + } + } + audit(&expanded, &expanded.body); let expand_ms = start.elapsed().as_secs_f64() * 1000.0; memory.insert("expanded", process_memory()); let start = Instant::now();