diff --git a/crates/ipu-codegen/src/compile.rs b/crates/ipu-codegen/src/compile.rs index 3b5b2e9c..da95e1f5 100644 --- a/crates/ipu-codegen/src/compile.rs +++ b/crates/ipu-codegen/src/compile.rs @@ -162,7 +162,7 @@ fn compile_graph( let selected = build_phase("plan_package", || { let costs = crate::estimate::MemoizedCostModel::new(&Ipu21CostModel); let fragments = crate::planner::cache::FragmentCache::default(); - let expansions = Arc::new(crate::low::expand::ExpansionCache::default()); + let expansions = Arc::new(crate::storage::GeometryCache::default()); let mut state = checkpoint::State::load(graph, config, tile_mapping)?; let mut fixed = config.clone(); let resuming = config.load_search_state.is_some(); @@ -439,7 +439,7 @@ fn remember<'a>( fn evaluate_candidate( mid: &crate::MidProgram, package: &PackageConfig, - expansions: Arc, + expansions: Arc, mut cache: crate::ExchangeScheduleCache, runtime: &[u8], ) -> PackageBuildResult { diff --git a/crates/ipu-codegen/src/compile/benchmark.rs b/crates/ipu-codegen/src/compile/benchmark.rs index 83a65188..95ac3650 100644 --- a/crates/ipu-codegen/src/compile/benchmark.rs +++ b/crates/ipu-codegen/src/compile/benchmark.rs @@ -16,7 +16,7 @@ pub struct ExpansionBenchmark { /// Whole-process Linux RSS, including caches and allocator-retained pages. pub process_memory: BTreeMap<&'static str, Option>, pub baseline: ExpansionTiming, - /// A second candidate expansion with the same search cache and fresh costing state. + /// A second candidate expansion with the same shared geometry cache. pub warm: Option, /// Matching selections are opportunities, not validated reusable graph fragments. pub selection_reuse: std::collections::BTreeMap<&'static str, SelectionReuse>, @@ -55,12 +55,8 @@ pub struct SelectionReuse { #[derive(serde::Serialize)] pub struct ExpansionTiming { pub process_memory: BTreeMap<&'static str, Option>, - /// Cache entries, hits and misses, respectively; counters are cumulative. - pub fragment_cache: (usize, u64, u64), - pub copy_geometry_cache: (usize, u64, u64), - /// Capacity-based payload estimates: copy recipes, destination facts, costing geometry. - /// Shared traversal bodies can be counted more than once; allocator overhead is excluded. - pub retained_cache_bytes: (usize, usize, usize), + /// Shared geometry counters are cumulative across cold/warm expansion and costing. + pub geometry_cache: crate::storage::GeometryCacheStats, pub mid_operations: usize, pub mid_values: usize, pub expand_ms: f64, @@ -71,8 +67,6 @@ pub struct ExpansionTiming { pub footprint_ms: f64, pub estimated_row_bytes: u64, pub maximum_transfer_chunks_per_tile: u64, - /// Distinct views, view pairs and stored receive-row descriptors. - pub geometry: (usize, usize, usize), /// Diagnostic copying costs, outside expansion/footprint timings. pub clone_shards_ms: f64, pub clone_kernel_runs_ms: f64, @@ -118,9 +112,9 @@ pub fn benchmark_mid_expansion( .ok_or_else(|| invalid("cannot cost benchmark mid program"))?; let mid_cost_ms = start.elapsed().as_secs_f64() * 1000.0; let cache = Arc::new(if cache_enabled { - crate::low::expand::ExpansionCache::default() + crate::storage::GeometryCache::default() } else { - crate::low::expand::ExpansionCache::disabled() + crate::storage::GeometryCache::disabled() }); let mut selections = std::collections::BTreeMap::< &'static str, @@ -205,24 +199,22 @@ pub fn benchmark_mid_expansion( fn measure_expansion( mid: &crate::MidProgram, config: &PipelineConfig, - cache: &Arc, + cache: &Arc, ) -> PackageBuildResult { let mut memory = BTreeMap::from([("start", process_memory())]); - let mut analysis = crate::estimate::GeometryAnalysis::default(); let start = Instant::now(); - let expanded = crate::low::expand::expand_tiles_analyzed( + let expanded = crate::low::expand::expand_tiles_cached( &mid, config.diagnostic_checkpoints, Arc::clone(&cache), - &mut analysis, )?; let expand_ms = start.elapsed().as_secs_f64() * 1000.0; memory.insert("expanded", process_memory()); let start = Instant::now(); - let low_cost = crate::estimate::program_cycles_analyzed(&expanded, None, &mut analysis)?; + let low_cost = crate::estimate::program_cycles_analyzed(&expanded, None, cache)?; let recost_ms = start.elapsed().as_secs_f64() * 1000.0; let start = Instant::now(); - let footprint = crate::estimate::program_footprint_analyzed(&expanded, &mut analysis)?; + let footprint = crate::estimate::program_footprint_analyzed(&expanded, cache)?; let footprint_ms = start.elapsed().as_secs_f64() * 1000.0; memory.insert("footprint", process_memory()); let start = Instant::now(); @@ -237,12 +229,9 @@ fn measure_expansion( let cloned = std::hint::black_box(expanded.kernel_runs.clone()); let clone_kernel_runs_ms = start.elapsed().as_secs_f64() * 1000.0; drop(cloned); - let retained = cache.retained_bytes(); let mut timing = ExpansionTiming { process_memory: BTreeMap::new(), - fragment_cache: cache.stats(), - copy_geometry_cache: cache.geometry_stats(), - retained_cache_bytes: (retained.0, retained.1, analysis.retained_bytes()), + geometry_cache: cache.stats(), mid_operations: mid.operations.len(), mid_values: mid.values.len(), expand_ms, @@ -252,7 +241,6 @@ fn measure_expansion( footprint_ms, estimated_row_bytes: footprint.estimated_row_bytes(), maximum_transfer_chunks_per_tile: footprint.maximum_transfer_chunks_per_tile, - geometry: analysis.stats(), clone_shards_ms, clone_kernel_runs_ms, shards: low.shards.len(), @@ -273,7 +261,7 @@ fn measure_expansion( .map(|t| t.destinations.len()) .sum(), }; - drop((low, expanded, analysis)); + drop((low, expanded)); memory.insert("after_plan_drop", process_memory()); timing.process_memory = memory; tracing::info!(expand_ms, tile_lists_ms, "benchmarked mid-to-low expansion"); diff --git a/crates/ipu-codegen/src/compile/screen.rs b/crates/ipu-codegen/src/compile/screen.rs index f3d28886..21b6d477 100644 --- a/crates/ipu-codegen/src/compile/screen.rs +++ b/crates/ipu-codegen/src/compile/screen.rs @@ -25,7 +25,7 @@ pub(super) fn expand_and_place( let (low, footprint) = expand_and_screen( mid, planning, - Arc::new(crate::low::expand::ExpansionCache::default()), + Arc::new(crate::storage::GeometryCache::default()), )?; let placement = place(&low)?; Ok((low, placement, footprint)) @@ -34,18 +34,16 @@ pub(super) fn expand_and_place( pub(super) fn expand_and_screen( mid: &crate::MidProgram, planning: &PipelineConfig, - cache: Arc, + cache: Arc, ) -> PackageBuildResult<(LowProgram, crate::estimate::ExchangeFootprint)> { let start = Instant::now(); - let mut analysis = crate::estimate::GeometryAnalysis::default(); - let expanded = crate::low::expand::expand_tiles_analyzed( + let expanded = crate::low::expand::expand_tiles_cached( mid, planning.diagnostic_checkpoints, - cache, - &mut analysis, + Arc::clone(&cache), )?; let expansion_ms = start.elapsed().as_millis(); - let footprint = crate::estimate::program_footprint_analyzed(&expanded, &mut analysis)?; + let footprint = crate::estimate::program_footprint_analyzed(&expanded, &cache)?; let fragments = footprint.maximum_transfer_chunks_per_tile; tracing::info!( expansion_ms, diff --git a/crates/ipu-codegen/src/estimate/geometry.rs b/crates/ipu-codegen/src/estimate/geometry.rs deleted file mode 100644 index a93ac018..00000000 --- a/crates/ipu-codegen/src/estimate/geometry.rs +++ /dev/null @@ -1,136 +0,0 @@ -//! Geometry facts shared by execution and table-storage analysis. Tile bindings, -//! transmit lanes and pointer continuation remain in the consuming phase walk. -use crate::storage::{ByteTraversal, StridedSpan, ViewGeometry}; -use crate::{CopyOrder, ExpansionResult, ShardView, TileGraph}; -use std::collections::HashMap; -use std::sync::Arc; - -const LIMIT: u32 = ipu_exchange::MAX_TRANSFER_WORDS * 4; - -pub(super) struct CopyGeometry { - pub bytes: u64, - pub fragments: u64, - pub long_fragments: u64, - pub receives: Vec, -} -impl CopyGeometry { - fn new(source: &ByteTraversal, target: &ByteTraversal) -> ExpansionResult { - let mut result = Self { - bytes: source.byte_len(), - fragments: 0, - long_fragments: 0, - receives: Vec::new(), - }; - if result.bytes != target.byte_len() { - return Err(crate::storage::StorageError::InvalidView.into()); - } - let mut append = |span: StridedSpan| { - let StridedSpan { bytes, rows, .. } = span; - if bytes == 0 { - return; - } - result.fragments += u64::from(rows) * u64::from(bytes.div_ceil(LIMIT)); - result.long_fragments += u64::from(rows) - * (u64::from(bytes / LIMIT) * u64::from(LIMIT > 256) - + u64::from(bytes % LIMIT > 256)); - result.receives.push(span); - }; - if let Some((_, target)) = source.regular_copy(target) { - append(target); - } else { - crate::for_each_copy_span(source.spans(), target.spans(), |_, offset, bytes| { - append(StridedSpan { - offset, - bytes, - rows: 1, - stride: bytes, - }); - Ok(()) - })?; - } - Ok(result) - } -} - -/// Scoped to one immutable expanded program. Integer geometry IDs make pair -/// lookup cheap and avoid retaining repeated traversal trees per recipient. -#[derive(Default)] -pub(crate) struct GeometryAnalysis { - views: HashMap<(CopyOrder, ViewGeometry), usize>, - traversals: Vec, - pairs: HashMap<(usize, usize), Arc>, -} -impl GeometryAnalysis { - pub(crate) fn retained_bytes(&self) -> usize { - self.views.capacity() * (size_of::<((CopyOrder, ViewGeometry), usize)>() + 1) - + self - .views - .keys() - .map(|(_, view)| view.heap_bytes()) - .sum::() - + self.traversals.capacity() * size_of::() - + self - .traversals - .iter() - .map(ByteTraversal::heap_bytes) - .sum::() - + self.pairs.capacity() * (size_of::<((usize, usize), Arc)>() + 1) - + self - .pairs - .values() - .map(|pair| { - size_of::() - + 2 * size_of::() - + pair.receives.capacity() * size_of::() - }) - .sum::() - } - - pub(crate) fn stats(&self) -> (usize, usize, usize) { - ( - self.traversals.len(), - self.pairs.len(), - self.pairs.values().map(|p| p.receives.len()).sum(), - ) - } - - pub(super) fn bytes(&self, view: usize) -> u64 { - self.traversals[view].byte_len() - } - pub(super) fn view( - &mut self, - program: &TileGraph, - view: &ShardView, - order: CopyOrder, - ) -> ExpansionResult { - let view = view.bind(&program.shards)?; - let key = ( - order, - ViewGeometry::new(view.shard.storage(), view.extents)?, - ); - if let Some(&id) = self.views.get(&key) { - return Ok(id); - } - let traversal = view.traversal(order)?; - let id = self.traversals.len(); - self.traversals.push(traversal); - self.views.insert(key, id); - Ok(id) - } - pub(super) fn copy( - &mut self, - program: &TileGraph, - source: usize, - target: &ShardView, - order: CopyOrder, - ) -> ExpansionResult> { - let a = source; - let b = self.view(program, target, order)?; - if let Some(pair) = self.pairs.get(&(a, b)) { - return Ok(Arc::clone(pair)); - } - let pair = Arc::new(CopyGeometry::new(&self.traversals[a], &self.traversals[b])?); - self.pairs.insert((a, b), Arc::clone(&pair)); - Ok(pair) - } -} diff --git a/crates/ipu-codegen/src/estimate/mod.rs b/crates/ipu-codegen/src/estimate/mod.rs index 219638cb..90aec792 100644 --- a/crates/ipu-codegen/src/estimate/mod.rs +++ b/crates/ipu-codegen/src/estimate/mod.rs @@ -1,8 +1,7 @@ //! Estimates share one vocabulary: tensor geometry and traffic describe work; //! memory evaluates allocations/liveness, and cycles applies IPU21 prices. -mod geometry; -pub(crate) use geometry::GeometryAnalysis; +use crate::storage::GeometryCache; mod cycles; use cycles::exchange_fragment_price; mod exchange_storage; diff --git a/crates/ipu-codegen/src/estimate/program.rs b/crates/ipu-codegen/src/estimate/program.rs index ad1d3859..566fe36f 100644 --- a/crates/ipu-codegen/src/estimate/program.rs +++ b/crates/ipu-codegen/src/estimate/program.rs @@ -140,13 +140,13 @@ pub(crate) fn program_cycles( program: &TileGraph, exchange: Option<&[u64]>, ) -> ExpansionResult { - program_cycles_analyzed(program, exchange, &mut GeometryAnalysis::default()) + program_cycles_analyzed(program, exchange, &GeometryCache::default()) } pub(crate) fn program_cycles_analyzed( program: &TileGraph, exchange: Option<&[u64]>, - geometry: &mut GeometryAnalysis, + geometry: &GeometryCache, ) -> ExpansionResult { let estimated; let phases = if let Some(costs) = exchange { @@ -210,7 +210,7 @@ pub(crate) fn program_cycles_analyzed( pub(crate) fn exchange_phase_estimate( program: &TileGraph, phase: &crate::ExchangePhase, - geometry: &mut GeometryAnalysis, + geometry: &GeometryCache, ) -> ExpansionResult<(u64, Vec)> { let mut storage = ExchangeStoragePhase::new(program.tile_count); let traffic = geometry_traffic(program, phase, Some(&mut storage), geometry)?; @@ -224,7 +224,7 @@ fn geometry_traffic( program: &TileGraph, phase: &crate::ExchangePhase, mut storage: Option<&mut ExchangeStoragePhase>, - geometry: &mut GeometryAnalysis, + geometry: &GeometryCache, ) -> ExpansionResult { #[cfg(test)] let mut expected_storage = storage.as_deref().cloned(); @@ -233,20 +233,31 @@ fn geometry_traffic( for transfer in &phase.transfers { let source = &program.shards[transfer.source.shard.index() as usize]; let order = transfer.span_order(&program.shards); - let source_geometry = geometry.view(program, &transfer.source, order)?; - let bytes = geometry.bytes(source_geometry); + let source_geometry = transfer + .source + .bind(&program.shards)? + .geometry(geometry, order)?; + let bytes = source_geometry.traversal.byte_len(); let mut outgoing_fragments = 0; let mut outgoing_long_fragments = 0; for destination in &transfer.destinations { let target = &program.shards[destination.shard.index() as usize]; - let copy = geometry.copy(program, source_geometry, destination, order)?; - outgoing_fragments = outgoing_fragments.max(copy.fragments); - outgoing_long_fragments = outgoing_long_fragments.max(copy.long_fragments); + let target_geometry = destination + .bind(&program.shards)? + .geometry(geometry, order)?; + let copy = geometry.pair(&source_geometry, &target_geometry)?; + let mut fragments = 0; + let mut long_fragments = 0; // Relative allocation/offset identities expose pointer continuation // without placement. This follows the supplied transfer order; // scheduling can change that order, pairing and control overlap. let mut resets = 0; - for row in ©.receives { + for [_, row] in ©.rows { + let limit = ipu_exchange::MAX_TRANSFER_WORDS * 4; + fragments += u64::from(row.rows) * u64::from(row.bytes.div_ceil(limit)); + long_fragments += u64::from(row.rows) + * (u64::from(row.bytes / limit) * u64::from(limit > 256) + + u64::from(row.bytes % limit > 256)); let address = (u64::from(destination.shard.index()) << 32) + u64::from(row.offset); if let Some(storage) = storage.as_deref_mut() { storage.connection_rows( @@ -269,7 +280,9 @@ fn geometry_traffic( + u64::from(row.bytes), ); } - traffic.add_receive(target.tile, bytes, copy.fragments, resets); + outgoing_fragments = outgoing_fragments.max(fragments); + outgoing_long_fragments = outgoing_long_fragments.max(long_fragments); + traffic.add_receive(target.tile, bytes, fragments, resets); } if let Some(storage) = storage.as_deref_mut() { storage.send(source.tile, outgoing_fragments, outgoing_long_fragments); @@ -307,12 +320,12 @@ fn kernel_cycles<'a>(run: &'a KernelRun) -> u64 { #[cfg(test)] pub(crate) fn program_footprint(program: &TileGraph) -> ExpansionResult { - program_footprint_analyzed(program, &mut GeometryAnalysis::default()) + program_footprint_analyzed(program, &GeometryCache::default()) } pub(crate) fn program_footprint_analyzed( program: &TileGraph, - geometry: &mut GeometryAnalysis, + geometry: &GeometryCache, ) -> ExpansionResult { // Storage belongs to a tile, not to a shared transmit lane. Count both // endpoint roles conservatively (bidi encoding may later combine them). diff --git a/crates/ipu-codegen/src/kernel/copy.rs b/crates/ipu-codegen/src/kernel/copy.rs index 4505ece3..5be7332d 100644 --- a/crates/ipu-codegen/src/kernel/copy.rs +++ b/crates/ipu-codegen/src/kernel/copy.rs @@ -4,7 +4,8 @@ use super::{KernelCall, StorageError}; use crate::low::storage::{StorageAccess, bind_storage}; -use crate::{BlockValue, BlockValueId, CopyPattern, LocalCopy}; +use crate::storage::{ByteSpan, CopyPair, StorageResult}; +use crate::{BlockValue, BlockValueId, CopyOperation, CopyPattern, LocalCopy}; use ipu_target::ipu21::WORKER_CONTEXTS; /// Coarse launch allowance when mid has not selected a local helper yet. @@ -200,6 +201,407 @@ impl CopyRun { } } +impl CopyOperation { + pub(crate) fn from_pair( + source: Buffer, + destination: Buffer, + same_buffer: bool, + pair: &CopyPair, + ) -> StorageResult> { + if pair.bytes == 0 { + return Ok(Vec::new()); + } + if let [[left, right]] = pair.rows.as_slice() + && let Some(pattern) = row_copy_pattern( + left.rows, + left.bytes, + left.offset, + right.offset, + left.stride, + right.stride, + ) + { + return Ok(vec![Self { + source, + destination, + source_offset: left.offset, + destination_offset: right.offset, + bytes: left + .bytes + .checked_mul(left.rows) + .ok_or(StorageError::Overflow)?, + pattern, + }]); + } + let mut copies = Vec::new(); + for [left, right] in &pair.rows { + for row in 0..left.rows { + copies.push(Self { + source: source.clone(), + destination: destination.clone(), + source_offset: left.offset + row * left.stride, + destination_offset: right.offset + row * right.stride, + bytes: left.bytes, + pattern: CopyPattern::Contiguous, + }); + } + } + Self::group(copies, !same_buffer) + } + + #[cfg(test)] + pub(crate) fn from_traversals( + source: Buffer, + destination: Buffer, + a: &crate::storage::ByteTraversal, + b: &crate::storage::ByteTraversal, + ) -> StorageResult> { + let same_buffer = source == destination; + Self::from_pair(source, destination, same_buffer, &CopyPair::new(a, b)?) + } + + #[cfg(test)] + pub(crate) fn from_spans( + source: Buffer, + destination: Buffer, + source_spans: impl IntoIterator, + destination_spans: impl IntoIterator, + ) -> StorageResult> { + let same_buffer = source == destination; + Self::from_pair( + source, + destination, + same_buffer, + &CopyPair::from_spans(source_spans, destination_spans)?, + ) + } + + fn group(mut copies: Vec, can_reorder: bool) -> StorageResult> { + let original = coalesce_copies(&copies); + if !can_reorder || original.len() < 2 { + return Ok(original); + } + // Local copies read one live allocation and write another. Traversing + // disjoint destination spans in physical order can turn a blocked + // transpose's many short launches into a few long strided copies. + copies.sort_by_key(|copy| (copy.destination_offset, copy.source_offset)); + if copies.windows(2).any(|pair| { + pair[0] + .destination_offset + .checked_add(pair[0].bytes) + .is_none_or(|end| end > pair[1].destination_offset) + }) { + return Ok(original); + } + let reordered = coalesce_copies(&copies); + Ok(if reordered.len() < original.len() { + reordered + } else { + original + }) + } +} + +const PARALLEL_STRIDED_COPY_MAX_BYTES: u32 = 512; + +/// Shared launch policy for symbolic rows and the irregular-span fallback. +fn row_copy_pattern( + rows: u32, + row_bytes: u32, + source: u32, + destination: u32, + source_stride: u32, + destination_stride: u32, +) -> Option { + if rows == 1 { + return Some(CopyPattern::Contiguous); + } + let geometry = [ + row_bytes, + source, + destination, + source_stride, + destination_stride, + ]; + if source_stride == 0 + || destination_stride == 0 + || !geometry.iter().all(|n| n.is_multiple_of(4)) + { + return None; + } + // Wide copies with too few rows underuse the workers in the strided kernel. + if rows < WORKER_CONTEXTS + && row_bytes.saturating_mul(rows) > PARALLEL_STRIDED_COPY_MAX_BYTES + && geometry.iter().all(|n| n.is_multiple_of(8)) + { + return None; + } + Some( + if source_stride == row_bytes && destination_stride == row_bytes { + CopyPattern::Contiguous + } else { + CopyPattern::Strided { + rows, + row_bytes, + source_stride, + destination_stride, + } + }, + ) +} + +fn coalesce_copies(copies: &[CopyOperation]) -> Vec> { + let mut coalesced = Vec::new(); + let mut index = 0; + while index < copies.len() { + let first = &copies[index]; + let Some(second) = copies.get(index + 1) else { + coalesced.push(first.clone()); + break; + }; + if first.bytes != second.bytes || first.bytes == 0 || !first.bytes.is_multiple_of(4) { + coalesced.push(first.clone()); + index += 1; + continue; + } + let source_stride = second.source_offset.saturating_sub(first.source_offset); + let destination_stride = second + .destination_offset + .saturating_sub(first.destination_offset); + if source_stride == 0 + || destination_stride == 0 + || !first.source_offset.is_multiple_of(4) + || !first.destination_offset.is_multiple_of(4) + || !source_stride.is_multiple_of(4) + || !destination_stride.is_multiple_of(4) + { + coalesced.push(first.clone()); + index += 1; + continue; + } + let mut end = index + 2; + while let Some(copy) = copies.get(end) { + let previous = &copies[end - 1]; + if copy.bytes != first.bytes + || copy.source_offset.checked_sub(previous.source_offset) != Some(source_stride) + || copy + .destination_offset + .checked_sub(previous.destination_offset) + != Some(destination_stride) + { + break; + } + end += 1; + } + let rows = u32::try_from(end - index).unwrap_or(u32::MAX); + if let Some(pattern) = row_copy_pattern( + rows, + first.bytes, + first.source_offset, + first.destination_offset, + source_stride, + destination_stride, + ) { + let mut copy = first.clone(); + copy.bytes = copy.bytes.saturating_mul(rows); + copy.pattern = pattern; + coalesced.push(copy); + } else { + coalesced.extend(copies[index..end].iter().cloned()); + } + index = end; + } + coalesced +} + +#[cfg(test)] +mod planning_tests { + use super::*; + + #[test] + fn packed_block_transpose_uses_long_rows_instead_of_tiny_launches() { + let source = [ByteSpan { + offset: 0, + bytes: 164 * 6 * 32, + }]; + let destination = (0..164) + .flat_map(|column| { + (0..6).map(move |row| ByteSpan { + offset: (row * 164 + column) * 32, + bytes: 32, + }) + }) + .collect::>(); + let copies = + CopyOperation::from_spans(0, 1, source.iter().copied(), destination.iter().copied()) + .unwrap(); + assert_eq!(copies.len(), 6); + for (row, copy) in copies.iter().enumerate() { + assert_eq!(copy.source_offset, row as u32 * 32); + assert_eq!(copy.destination_offset, row as u32 * 164 * 32); + assert_eq!( + copy.pattern, + CopyPattern::Strided { + rows: 164, + row_bytes: 32, + source_stride: 192, + destination_stride: 32, + } + ); + } + // Do not reorder a request that explicitly aliases its source. + assert_eq!( + CopyOperation::from_spans(0, 0, source.iter().copied(), destination.iter().copied()) + .unwrap() + .len(), + 164 + ); + } + + #[test] + fn copy_runs_preserve_byte_mapping_across_span_boundaries() { + let mut random = fastrand::Rng::with_seed(0x636f_7079); + for _ in 0..1000 { + let bytes = random.u32(1..=1024); + let mut spans = || { + let mut remaining = bytes; + let mut offset = random.u32(0..32); + let mut result = Vec::new(); + while remaining > 0 { + let count = random.u32(1..=remaining.min(64)); + result.push(ByteSpan { + offset, + bytes: count, + }); + offset += count + random.u32(0..16); + remaining -= count; + } + result + }; + let source = spans(); + let destination = spans(); + let flatten = |spans: &[ByteSpan]| { + spans + .iter() + .flat_map(|span| span.offset..span.offset + span.bytes) + .collect::>() + }; + let mut expected = flatten(&source) + .into_iter() + .zip(flatten(&destination)) + .collect::>(); + let copies = CopyOperation::from_spans( + 0, + 1, + source.iter().copied(), + destination.iter().copied(), + ) + .unwrap(); + let mut actual = copies + .into_iter() + .flat_map(|copy| { + let (rows, width, source_stride, destination_stride) = match copy.pattern { + CopyPattern::Contiguous => (1, copy.bytes, 0, 0), + CopyPattern::Strided { + rows, + row_bytes, + source_stride, + destination_stride, + } => (rows, row_bytes, source_stride, destination_stride), + }; + (0..rows).flat_map(move |row| { + (0..width).map(move |byte| { + ( + copy.source_offset + row * source_stride + byte, + copy.destination_offset + row * destination_stride + byte, + ) + }) + }) + }) + .collect::>(); + actual.sort_unstable(); + expected.sort_unstable(); + assert_eq!(actual, expected); + } + assert!( + CopyOperation::from_spans( + 0, + 1, + [ByteSpan { + offset: 0, + bytes: 4 + }], + [] + ) + .is_err() + ); + } + + #[test] + fn long_strided_copies_keep_all_rows_in_one_launch() { + for rows in [2, 64, 65, 164] { + let source = (0..rows) + .map(|row| ByteSpan { + offset: row * 16, + bytes: 8, + }) + .collect::>(); + let destination = (0..rows) + .map(|row| ByteSpan { + offset: row * 24, + bytes: 8, + }) + .collect::>(); + let copies = CopyOperation::from_spans( + 0, + 1, + source.iter().copied(), + destination.iter().copied(), + ) + .unwrap(); + assert_eq!(copies.len(), 1); + assert_eq!( + copies[0].pattern, + CopyPattern::Strided { + rows, + row_bytes: 8, + source_stride: 16, + destination_stride: 24, + } + ); + } + let source = [ + ByteSpan { + offset: 0, + bytes: 1024, + }, + ByteSpan { + offset: 2048, + bytes: 1024, + }, + ]; + let destination = [ + ByteSpan { + offset: 0, + bytes: 1024, + }, + ByteSpan { + offset: 4096, + bytes: 1024, + }, + ]; + let copies = + CopyOperation::from_spans(0, 1, source.iter().copied(), destination.iter().copied()) + .unwrap(); + assert_eq!(copies.len(), 2); + assert!( + copies + .iter() + .all(|copy| copy.pattern == CopyPattern::Contiguous) + ); + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/ipu-codegen/src/low/copy.rs b/crates/ipu-codegen/src/low/copy.rs index 2838deb9..3e8927a7 100644 --- a/crates/ipu-codegen/src/low/copy.rs +++ b/crates/ipu-codegen/src/low/copy.rs @@ -1,6 +1,5 @@ -//! Relative copy operations and materialization policy, before placement. - -use crate::storage::{ByteSpan, ByteTraversal, StorageError, StorageResult, for_each_copy_span}; +//! Byte-copy descriptions and explicit movement requests, before placement. +//! Local launch construction belongs to kernel::copy; this module declares the contracts. /// Requested realization of a whole-device coordinate copy. Explicit requests /// are checked by movement lowering; Automatic selects from the actual geometry. @@ -81,391 +80,3 @@ pub enum CopyPattern { destination_stride: u32, }, } - -impl CopyOperation { - pub(crate) fn from_traversals( - source: Buffer, - destination: Buffer, - source_traversal: &ByteTraversal, - destination_traversal: &ByteTraversal, - ) -> StorageResult> { - if source_traversal.byte_len() != destination_traversal.byte_len() { - return Err(StorageError::InvalidView); - } - if source_traversal.byte_len() == 0 { - return Ok(Vec::new()); - } - if let Some((left, right)) = source_traversal.regular_copy(destination_traversal) - && let Some(pattern) = row_copy_pattern( - left.rows, - left.bytes, - left.offset, - right.offset, - left.stride, - right.stride, - ) - { - return Ok(vec![Self { - source, - destination, - source_offset: left.offset, - destination_offset: right.offset, - bytes: left.bytes * left.rows, - pattern, - }]); - } - Self::from_spans( - source, - destination, - source_traversal.spans(), - destination_traversal.spans(), - ) - } - - pub(crate) fn from_spans( - source: Buffer, - destination: Buffer, - source_spans: impl IntoIterator, - destination_spans: impl IntoIterator, - ) -> StorageResult> { - let can_reorder = source != destination; - let mut copies = Vec::new(); - for_each_copy_span( - source_spans, - destination_spans, - |source_offset, destination_offset, bytes| { - copies.push(Self { - source: source.clone(), - destination: destination.clone(), - source_offset, - destination_offset, - bytes, - pattern: CopyPattern::Contiguous, - }); - Ok(()) - }, - )?; - let original = coalesce_copies(&copies); - if !can_reorder || original.len() < 2 { - return Ok(original); - } - // Local copies read one live allocation and write another. Traversing - // disjoint destination spans in physical order can turn a blocked - // transpose's many short launches into a few long strided copies. - copies.sort_by_key(|copy| (copy.destination_offset, copy.source_offset)); - if copies.windows(2).any(|pair| { - pair[0] - .destination_offset - .checked_add(pair[0].bytes) - .is_none_or(|end| end > pair[1].destination_offset) - }) { - return Ok(original); - } - let reordered = coalesce_copies(&copies); - Ok(if reordered.len() < original.len() { - reordered - } else { - original - }) - } -} - -const PARALLEL_STRIDED_COPY_MAX_BYTES: u32 = 512; - -/// Shared launch policy for symbolic rows and the irregular-span fallback. -fn row_copy_pattern( - rows: u32, - row_bytes: u32, - source: u32, - destination: u32, - source_stride: u32, - destination_stride: u32, -) -> Option { - if rows == 1 { - return Some(CopyPattern::Contiguous); - } - let geometry = [ - row_bytes, - source, - destination, - source_stride, - destination_stride, - ]; - if source_stride == 0 - || destination_stride == 0 - || !geometry.iter().all(|n| n.is_multiple_of(4)) - { - return None; - } - // Wide copies with too few rows underuse the workers in the strided kernel. - if rows < 6 - && row_bytes.saturating_mul(rows) > PARALLEL_STRIDED_COPY_MAX_BYTES - && geometry.iter().all(|n| n.is_multiple_of(8)) - { - return None; - } - Some( - if source_stride == row_bytes && destination_stride == row_bytes { - CopyPattern::Contiguous - } else { - CopyPattern::Strided { - rows, - row_bytes, - source_stride, - destination_stride, - } - }, - ) -} - -fn coalesce_copies(copies: &[CopyOperation]) -> Vec> { - let mut coalesced = Vec::new(); - let mut index = 0; - while index < copies.len() { - let first = &copies[index]; - let Some(second) = copies.get(index + 1) else { - coalesced.push(first.clone()); - break; - }; - if first.bytes != second.bytes || first.bytes == 0 || !first.bytes.is_multiple_of(4) { - coalesced.push(first.clone()); - index += 1; - continue; - } - let source_stride = second.source_offset.saturating_sub(first.source_offset); - let destination_stride = second - .destination_offset - .saturating_sub(first.destination_offset); - if source_stride == 0 - || destination_stride == 0 - || !first.source_offset.is_multiple_of(4) - || !first.destination_offset.is_multiple_of(4) - || !source_stride.is_multiple_of(4) - || !destination_stride.is_multiple_of(4) - { - coalesced.push(first.clone()); - index += 1; - continue; - } - let mut end = index + 2; - while let Some(copy) = copies.get(end) { - let previous = &copies[end - 1]; - if copy.bytes != first.bytes - || copy.source_offset.checked_sub(previous.source_offset) != Some(source_stride) - || copy - .destination_offset - .checked_sub(previous.destination_offset) - != Some(destination_stride) - { - break; - } - end += 1; - } - let rows = u32::try_from(end - index).unwrap_or(u32::MAX); - if let Some(pattern) = row_copy_pattern( - rows, - first.bytes, - first.source_offset, - first.destination_offset, - source_stride, - destination_stride, - ) { - let mut copy = first.clone(); - copy.bytes = copy.bytes.saturating_mul(rows); - copy.pattern = pattern; - coalesced.push(copy); - } else { - coalesced.extend(copies[index..end].iter().cloned()); - } - index = end; - } - coalesced -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn packed_block_transpose_uses_long_rows_instead_of_tiny_launches() { - let source = [ByteSpan { - offset: 0, - bytes: 164 * 6 * 32, - }]; - let destination = (0..164) - .flat_map(|column| { - (0..6).map(move |row| ByteSpan { - offset: (row * 164 + column) * 32, - bytes: 32, - }) - }) - .collect::>(); - let copies = - CopyOperation::from_spans(0, 1, source.iter().copied(), destination.iter().copied()) - .unwrap(); - assert_eq!(copies.len(), 6); - for (row, copy) in copies.iter().enumerate() { - assert_eq!(copy.source_offset, row as u32 * 32); - assert_eq!(copy.destination_offset, row as u32 * 164 * 32); - assert_eq!( - copy.pattern, - CopyPattern::Strided { - rows: 164, - row_bytes: 32, - source_stride: 192, - destination_stride: 32, - } - ); - } - // Do not reorder a request that explicitly aliases its source. - assert_eq!( - CopyOperation::from_spans(0, 0, source.iter().copied(), destination.iter().copied()) - .unwrap() - .len(), - 164 - ); - } - - #[test] - fn copy_runs_preserve_byte_mapping_across_span_boundaries() { - let mut random = fastrand::Rng::with_seed(0x636f_7079); - for _ in 0..1000 { - let bytes = random.u32(1..=1024); - let mut spans = || { - let mut remaining = bytes; - let mut offset = random.u32(0..32); - let mut result = Vec::new(); - while remaining > 0 { - let count = random.u32(1..=remaining.min(64)); - result.push(ByteSpan { - offset, - bytes: count, - }); - offset += count + random.u32(0..16); - remaining -= count; - } - result - }; - let source = spans(); - let destination = spans(); - let flatten = |spans: &[ByteSpan]| { - spans - .iter() - .flat_map(|span| span.offset..span.offset + span.bytes) - .collect::>() - }; - let mut expected = flatten(&source) - .into_iter() - .zip(flatten(&destination)) - .collect::>(); - let copies = CopyOperation::from_spans( - 0, - 1, - source.iter().copied(), - destination.iter().copied(), - ) - .unwrap(); - let mut actual = copies - .into_iter() - .flat_map(|copy| { - let (rows, width, source_stride, destination_stride) = match copy.pattern { - CopyPattern::Contiguous => (1, copy.bytes, 0, 0), - CopyPattern::Strided { - rows, - row_bytes, - source_stride, - destination_stride, - } => (rows, row_bytes, source_stride, destination_stride), - }; - (0..rows).flat_map(move |row| { - (0..width).map(move |byte| { - ( - copy.source_offset + row * source_stride + byte, - copy.destination_offset + row * destination_stride + byte, - ) - }) - }) - }) - .collect::>(); - actual.sort_unstable(); - expected.sort_unstable(); - assert_eq!(actual, expected); - } - assert!( - CopyOperation::from_spans( - 0, - 1, - [ByteSpan { - offset: 0, - bytes: 4 - }], - [] - ) - .is_err() - ); - } - - #[test] - fn long_strided_copies_keep_all_rows_in_one_launch() { - for rows in [2, 64, 65, 164] { - let source = (0..rows) - .map(|row| ByteSpan { - offset: row * 16, - bytes: 8, - }) - .collect::>(); - let destination = (0..rows) - .map(|row| ByteSpan { - offset: row * 24, - bytes: 8, - }) - .collect::>(); - let copies = CopyOperation::from_spans( - 0, - 1, - source.iter().copied(), - destination.iter().copied(), - ) - .unwrap(); - assert_eq!(copies.len(), 1); - assert_eq!( - copies[0].pattern, - CopyPattern::Strided { - rows, - row_bytes: 8, - source_stride: 16, - destination_stride: 24, - } - ); - } - let source = [ - ByteSpan { - offset: 0, - bytes: 1024, - }, - ByteSpan { - offset: 2048, - bytes: 1024, - }, - ]; - let destination = [ - ByteSpan { - offset: 0, - bytes: 1024, - }, - ByteSpan { - offset: 4096, - bytes: 1024, - }, - ]; - let copies = - CopyOperation::from_spans(0, 1, source.iter().copied(), destination.iter().copied()) - .unwrap(); - assert_eq!(copies.len(), 2); - assert!( - copies - .iter() - .all(|copy| copy.pattern == CopyPattern::Contiguous) - ); - } -} diff --git a/crates/ipu-codegen/src/low/expand/cache.rs b/crates/ipu-codegen/src/low/expand/cache.rs deleted file mode 100644 index 511334a2..00000000 --- a/crates/ipu-codegen/src/low/expand/cache.rs +++ /dev/null @@ -1,585 +0,0 @@ -//! Per-build caches for relative copy descriptors and destination geometry. -//! Packing selection and low graph mutation remain in movement construction. -use super::ExpansionResult; -use crate::low::{BlockValue, BlockValueId, CopyOperation, ShardView}; -use crate::storage::{CopyGeometry, CopyMapping, CopyOrder, ViewGeometry}; -use crate::tensor::{ElementOrder, Precision, ShardExtent}; -use hashbrown::HashTable; -use std::hash::{BuildHasher, Hash, Hasher}; -use std::sync::Arc; -use std::sync::Mutex; - -const MAX_ENTRIES: usize = 32768; - -// Prehashed lookup compares borrowed views; hit paths do not clone the large -// destination keys. HashTable owns probing/collision handling and reuses the -// recorded hash when growing, avoiding a second hash or a vector per bucket. -struct Memo { - entries: Mutex>, -} -struct Entries { - table: HashTable<(u64, K, Arc)>, - hits: u64, - misses: u64, -} -impl Default for Memo { - fn default() -> Self { - Self { - entries: Mutex::new(Entries { - table: HashTable::new(), - hits: 0, - misses: 0, - }), - } - } -} -impl Memo { - fn get(&self, hash: u64, matches: impl Fn(&K) -> bool) -> Option> { - let mut state = self.entries.lock().unwrap(); - let found = state - .table - .find(hash, |(_, key, _)| matches(key)) - .map(|(_, _, value)| Arc::clone(value)); - if found.is_some() { - state.hits += 1; - } else { - state.misses += 1; - } - found - } - fn insert(&self, hash: u64, key: K, value: Arc, limit: usize) { - let mut state = self.entries.lock().unwrap(); - if state.table.len() < limit - && state - .table - .find(hash, |(_, existing, _)| *existing == key) - .is_none() - { - state - .table - .insert_unique(hash, (hash, key, value), |entry| entry.0); - } - } - - fn has_capacity(&self, limit: usize) -> bool { - self.entries.lock().unwrap().table.len() < limit - } - fn stats(&self) -> (usize, u64, u64) { - let state = self.entries.lock().unwrap(); - (state.table.len(), state.hits, state.misses) - } - - fn retained_bytes(&self, heap: impl Fn(&K, &V) -> usize) -> usize { - let state = self.entries.lock().unwrap(); - state.table.capacity() * (size_of::<(u64, K, Arc)>() + 1) - + state - .table - .iter() - .map(|(_, key, value)| 2 * size_of::() + size_of::() + heap(key, value)) - .sum::() - } -} - -#[derive(PartialEq, Eq, Hash)] -struct CopyKey { - source: ViewGeometry, - destination: ViewGeometry, - order: CopyOrder, - same_buffer: bool, -} - -#[derive(PartialEq, Eq, Hash)] -struct GeometrySource { - format: (Precision, ElementOrder), - allocation: Vec, - source: Vec, - destination: Vec, -} -#[derive(PartialEq, Eq, Hash)] -struct GeometryKey { - destination: (Precision, ElementOrder), - allocation: Vec, - mappings: Vec, - order: CopyOrder, -} - -// Relative traversal depends on byte interpretation. Shape, tile ownership, -// replicas and bank class do not affect it and must not split cache entries. -fn format_key(shard: &BlockValue) -> (Precision, ElementOrder) { - ( - shard.tensor_type.format.precision, - shard.tensor_type.format.layout.order, - ) -} - -pub(crate) struct ExpansionCache { - enabled: bool, - copies: Memo>>, - geometry: Memo, -} -impl Default for ExpansionCache { - fn default() -> Self { - Self { - enabled: true, - copies: Memo::default(), - geometry: Memo::default(), - } - } -} -impl ExpansionCache { - pub(crate) fn disabled() -> Self { - Self { - enabled: false, - ..Self::default() - } - } - - pub(crate) fn stats(&self) -> (usize, u64, u64) { - self.copies.stats() - } - pub(crate) fn geometry_stats(&self) -> (usize, u64, u64) { - self.geometry.stats() - } - - /// Capacity-based retained payload, excluding allocator bookkeeping. Shared - /// traversal bodies can be counted more than once; this is an upper estimate. - pub(crate) fn retained_bytes(&self) -> (usize, usize) { - ( - self.copies.retained_bytes(|key, copies| { - key.source.heap_bytes() - + key.destination.heap_bytes() - + copies.capacity() * size_of::>() - }), - self.geometry.retained_bytes(|key, geometry| { - key.allocation.capacity() * size_of::() - + key.mappings.capacity() * size_of::() - + key - .mappings - .iter() - .map(|source| { - (source.allocation.capacity() - + source.source.capacity() - + source.destination.capacity()) - * size_of::() - }) - .sum::() - + geometry.heap_bytes() - }), - ) - } - - pub(super) fn geometry( - &self, - shards: &[BlockValue], - mappings: &[(ShardView, ShardView)], - destination: BlockValueId, - order: CopyOrder, - ) -> ExpansionResult> { - let shard = &shards[destination.index() as usize]; - for (source, target) in mappings { - source.bind(shards)?; - target.bind(shards)?; - if target.shard != destination { - return Err(super::ExpansionError::InvalidCopyPlan); - } - } - let generate = || { - let mappings = mappings - .iter() - .map(|(source, destination)| CopyMapping { - source: shards[source.shard.index() as usize].storage(), - source_extents: &source.extents, - destination_extents: &destination.extents, - }) - .collect::>(); - CopyGeometry::analyze( - shard.storage(), - &mappings, - order, - ipu_exchange::MAX_TRANSFER_WORDS * 4, - ) - }; - if !self.enabled - || (order != CopyOrder::Semantic - && mappings.len() == 1 - && mappings[0].1.extents == shard.extents) - { - return Ok(Arc::new(generate()?)); - } - let mut hash = foldhash::fast::FixedState::default().build_hasher(); - format_key(shard).hash(&mut hash); - shard.extents.hash(&mut hash); - order.hash(&mut hash); - mappings.len().hash(&mut hash); - for (source, destination) in mappings { - let input = &shards[source.shard.index() as usize]; - format_key(input).hash(&mut hash); - input.extents.hash(&mut hash); - source.extents.hash(&mut hash); - destination.extents.hash(&mut hash); - } - let hash = hash.finish(); - if let Some(geometry) = self.geometry.get(hash, |key| { - key.destination == format_key(shard) - && key.allocation == shard.extents - && key.order == order - && key.mappings.len() == mappings.len() - && key - .mappings - .iter() - .zip(mappings) - .all(|(key, (source, destination))| { - let input = &shards[source.shard.index() as usize]; - key.format == format_key(input) - && key.allocation == input.extents - && key.source == source.extents - && key.destination == destination.extents - }) - }) { - return Ok(geometry); - } - if !self.geometry.has_capacity(MAX_ENTRIES) { - return Ok(Arc::new(generate()?)); - } - let key = GeometryKey { - destination: format_key(shard), - allocation: shard.extents.clone(), - order, - mappings: mappings - .iter() - .map(|(source, destination)| { - let shard = &shards[source.shard.index() as usize]; - GeometrySource { - format: format_key(shard), - allocation: shard.extents.clone(), - source: source.extents.clone(), - destination: destination.extents.clone(), - } - }) - .collect(), - }; - let geometry = Arc::new(generate()?); - self.geometry - .insert(hash, key, Arc::clone(&geometry), MAX_ENTRIES); - Ok(geometry) - } - - pub(super) fn copy( - &self, - shards: &[BlockValue], - source: &ShardView, - destination: &ShardView, - order: CopyOrder, - ) -> ExpansionResult>>> { - let source = source.bind(shards)?; - let destination = destination.bind(shards)?; - let generate = || -> ExpansionResult>> { - let a = source.traversal(order)?; - let b = destination.traversal(order)?; - // Alias IDs can differ while their storage overlaps. Reordering - // copy spans is only safe between different backing allocations. - let copies = - CopyOperation::from_traversals(source.backing.0, destination.backing.0, &a, &b)? - .into_iter() - .map(|c| CopyOperation { - source: (), - destination: (), - source_offset: c.source_offset, - destination_offset: c.destination_offset, - bytes: c.bytes, - pattern: c.pattern, - }) - .collect(); - Ok(copies) - }; - let left = source.shard; - let right = destination.shard; - let whole_copy = source.extents == left.extents - && destination.extents == right.extents - && (order == CopyOrder::Physical - || (order == CopyOrder::Semantic - && left.tensor_type.format.layout.order == ElementOrder::RowMajor - && right.tensor_type.format.layout.order == ElementOrder::RowMajor)); - if !self.enabled || whole_copy { - return Ok(Arc::new(generate()?)); - } - let key = CopyKey { - source: ViewGeometry::new(left.storage(), source.extents)?, - destination: ViewGeometry::new(right.storage(), destination.extents)?, - order, - same_buffer: source.backing.0 == destination.backing.0, - }; - let mut hash = foldhash::fast::FixedState::default().build_hasher(); - key.hash(&mut hash); - let hash = hash.finish(); - if let Some(copies) = self.copies.get(hash, |existing| *existing == key) { - return Ok(copies); - } - let copies = generate()?; - // Generation is outside the lock, so unrelated finalists do not serialize. - let result = Arc::new(copies); - self.copies - .insert(hash, key, Arc::clone(&result), MAX_ENTRIES); - Ok(result) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::low::ShardDefinition; - use crate::tensor::{Layout, TensorType}; - - #[test] - fn memo_checks_full_keys_and_bounds_entries_even_on_hash_collisions() { - let memo = Memo::default(); - memo.insert(7, 1, Arc::new(10), 2); - memo.insert(7, 2, Arc::new(20), 2); - memo.insert(7, 3, Arc::new(30), 2); - assert_eq!(*memo.get(7, |key| *key == 1).unwrap(), 10); - assert_eq!(*memo.get(7, |key| *key == 2).unwrap(), 20); - assert!(memo.get(7, |key| *key == 3).is_none()); - assert_eq!(memo.stats(), (2, 2, 1)); - - // Growth must retain the supplied fingerprints, even though they need - // not equal a hash recomputed from the owned key. - let memo = Memo::default(); - for key in 0..1024 { - memo.insert(key % 7, key, Arc::new(key * 3), 1024); - } - for key in 0..1024 { - assert_eq!( - *memo.get(key % 7, |stored| *stored == key).unwrap(), - key * 3 - ); - } - } - - #[test] - fn copy_fragments_share_relative_geometry_across_tiles_and_origins() { - let cache = ExpansionCache::default(); - let mut first = None; - for shift in [0, 16] { - let shards = (0..2) - .map(|id| BlockValue { - id: BlockValueId(id), - tile: id as u16, - tensor_type: TensorType::new([2, 64], Precision::F16, Layout::row_sharded(1)), - extents: vec![ - ShardExtent { - axis: 0, - start: 0, - logical_end: 2, - physical_end: 2, - }, - ShardExtent { - axis: 1, - start: shift, - logical_end: shift + 8, - physical_end: shift + 8, - }, - ], - definition: ShardDefinition::Staging, - }) - .collect::>(); - let view = |id: usize| { - let mut extents = shards[id].extents.clone(); - extents[1].start += 2; - extents[1].logical_end -= 2; - extents[1].physical_end -= 2; - ShardView { - shard: BlockValueId(id as u32), - extents, - } - }; - let result = cache - .copy(&shards, &view(0), &view(1), CopyOrder::Semantic) - .unwrap(); - if let Some(first) = &first { - assert!(Arc::ptr_eq(first, &result)); - } else { - first = Some(result); - } - } - assert_eq!(cache.stats(), (1, 1, 1)); - } - - #[test] - fn destination_geometry_ignores_ownership_but_tracks_byte_interpretation() { - let mut shards = (0..2) - .map(|id| BlockValue { - id: BlockValueId(id), - tile: id as u16, - tensor_type: TensorType::new([2, 64], Precision::F16, Layout::row_sharded(1)), - extents: vec![ - ShardExtent { - axis: 0, - start: 0, - logical_end: 2, - physical_end: 2, - }, - ShardExtent { - axis: 1, - start: 0, - logical_end: 64, - physical_end: 64, - }, - ], - definition: ShardDefinition::Staging, - }) - .collect::>(); - let views = shards - .iter() - .map(|shard| ShardView { - shard: shard.id, - extents: shard.extents.clone(), - }) - .collect::>(); - let mappings = [(views[0].clone(), views[1].clone())]; - let cache = ExpansionCache::default(); - let before = cache - .geometry(&shards, &mappings, BlockValueId(1), CopyOrder::Semantic) - .unwrap(); - for shard in &mut shards { - shard.tile += 2; - shard.tensor_type.shape.0[0] *= 2; - shard.tensor_type.format.layout = Layout::row_sharded(2); - shard.tensor_type.format.layout.memory_class = crate::MemoryClass::Ipu21Interleaved; - } - let remapped = cache - .geometry(&shards, &mappings, BlockValueId(1), CopyOrder::Semantic) - .unwrap(); - assert!(Arc::ptr_eq(&before, &remapped)); - for shard in &mut shards { - shard.tensor_type.format.precision = Precision::F32; - } - let wider = cache - .geometry(&shards, &mappings, BlockValueId(1), CopyOrder::Semantic) - .unwrap(); - assert!(!Arc::ptr_eq(&before, &wider)); - assert_eq!(wider.bytes, before.bytes * 2); - assert_eq!(cache.geometry_stats(), (2, 1, 2)); - } - #[test] - fn copy_preparation_preserves_order_when_distinct_values_share_storage() { - fn mapping(copies: &[CopyOperation<()>]) -> Vec<(u32, u32)> { - copies - .iter() - .flat_map(|copy| { - let (rows, bytes, source_stride, destination_stride) = match copy.pattern { - crate::CopyPattern::Contiguous => (1, copy.bytes, 0, 0), - crate::CopyPattern::Strided { - rows, - row_bytes, - source_stride, - destination_stride, - } => (rows, row_bytes, source_stride, destination_stride), - }; - (0..rows).flat_map(move |row| { - (0..bytes).map(move |byte| { - ( - copy.source_offset + row * source_stride + byte, - copy.destination_offset + row * destination_stride + byte, - ) - }) - }) - }) - .collect() - } - let orders = [ - ElementOrder::RowMajor, - ElementOrder::Amp(crate::AmpOrder::Left), - ElementOrder::Amp(crate::AmpOrder::Output), - ElementOrder::Amp(crate::AmpOrder::TransposedLeft), - ElementOrder::Amp(crate::AmpOrder::TransposedOutput), - ElementOrder::BlockMajor(crate::BlockMajorOrder::Matrix { - row_block: 16, - column_block: 32, - }), - ]; - let mut reordered = 0; - for precision in [Precision::F16, Precision::F32] { - for source_order in orders { - for destination_order in orders { - let mut shards = [source_order, destination_order] - .into_iter() - .enumerate() - .map(|(index, order)| { - let mut layout = Layout::row_sharded(1); - layout.order = order; - BlockValue { - id: BlockValueId(index as u32), - tile: 0, - tensor_type: TensorType::new([48, 96], precision, layout), - extents: [48, 96] - .into_iter() - .enumerate() - .map(|(axis, end)| ShardExtent { - axis: axis as u16, - start: 0, - logical_end: end, - physical_end: end, - }) - .collect(), - definition: ShardDefinition::Staging, - } - }) - .collect::>(); - let view = |id| ShardView { - shard: BlockValueId(id), - extents: vec![ - ShardExtent { - axis: 0, - start: 0, - logical_end: 32, - physical_end: 32, - }, - ShardExtent { - axis: 1, - start: 0, - logical_end: 64, - physical_end: 64, - }, - ], - }; - let (source, destination) = (view(0), view(1)); - let a = source - .bind(&shards) - .unwrap() - .traversal(CopyOrder::Semantic) - .unwrap(); - let b = destination - .bind(&shards) - .unwrap() - .traversal(CopyOrder::Semantic) - .unwrap(); - let mut expected = Vec::new(); - crate::storage::for_each_copy_span(a.spans(), b.spans(), |a, b, bytes| { - expected.extend((0..bytes).map(|byte| (a + byte, b + byte))); - Ok(()) - }) - .unwrap(); - let cache = ExpansionCache::default(); - let separate = cache - .copy(&shards, &source, &destination, CopyOrder::Semantic) - .unwrap(); - reordered += usize::from(mapping(&separate) != expected); - for offset in [0, 64] { - shards[1].definition = ShardDefinition::ShiftedAlias { - source: BlockValueId(0), - offset, - }; - let shared = cache - .copy(&shards, &source, &destination, CopyOrder::Semantic) - .unwrap(); - assert_eq!(mapping(&shared), expected); - } - } - } - } - assert!( - reordered > 0, - "fixture must expose unsafe reuse of a reordered disjoint copy" - ); - } -} diff --git a/crates/ipu-codegen/src/low/expand/copies.rs b/crates/ipu-codegen/src/low/expand/copies.rs index 072ab60e..2a3cffeb 100644 --- a/crates/ipu-codegen/src/low/expand/copies.rs +++ b/crates/ipu-codegen/src/low/expand/copies.rs @@ -1,9 +1,8 @@ -//! Bind relative copy descriptors to the selected low storage identities. - +//! Bind copy geometry to low storage identities, then select local launches. use super::*; pub(super) fn append_span_copies( - cache: &ExpansionCache, + cache: &crate::storage::GeometryCache, shards: &[BlockValue], source: &ShardView, destination: &ShardView, @@ -11,23 +10,18 @@ pub(super) fn append_span_copies( copies: &mut Vec<(u16, LocalCopy)>, order: CopyOrder, ) -> ExpansionResult<()> { + let a = source.bind(shards)?; + let b = destination.bind(shards)?; + let pair = cache.pair(&a.geometry(cache, order)?, &b.geometry(cache, order)?)?; copies.extend( - cache - .copy(shards, source, destination, order)? - .iter() - .map(|c| { - ( - tile, - LocalCopy { - source: source.shard, - destination: destination.shard, - source_offset: c.source_offset, - destination_offset: c.destination_offset, - bytes: c.bytes, - pattern: c.pattern, - }, - ) - }), + LocalCopy::from_pair( + source.shard, + destination.shard, + a.backing.0 == b.backing.0, + &pair, + )? + .into_iter() + .map(|copy| (tile, copy)), ); Ok(()) } diff --git a/crates/ipu-codegen/src/low/expand/mod.rs b/crates/ipu-codegen/src/low/expand/mod.rs index 20c7435c..9b8ff13d 100644 --- a/crates/ipu-codegen/src/low/expand/mod.rs +++ b/crates/ipu-codegen/src/low/expand/mod.rs @@ -1,8 +1,7 @@ //! Expand selected whole-device primitives into tile-local calls and movement. -mod cache; mod emit; -pub(crate) use cache::ExpansionCache; +use crate::storage::GeometryCache; mod compute; mod buffers; @@ -69,35 +68,20 @@ pub(crate) fn expand_tiles( graph: &MidProgram, checkpoints: bool, ) -> ExpansionResult> { - expand_tiles_cached(graph, checkpoints, Arc::new(ExpansionCache::default())) + expand_tiles_cached(graph, checkpoints, Arc::new(GeometryCache::default())) } -#[cfg(test)] pub(crate) fn expand_tiles_cached( graph: &MidProgram, checkpoints: bool, - cache: Arc, -) -> ExpansionResult> { - expand_tiles_analyzed( - graph, - checkpoints, - cache, - &mut crate::estimate::GeometryAnalysis::default(), - ) -} - -pub(crate) fn expand_tiles_analyzed( - graph: &MidProgram, - checkpoints: bool, - cache: Arc, - analysis: &mut crate::estimate::GeometryAnalysis, + cache: Arc, ) -> ExpansionResult> { if graph.tile_count == 0 { return Err(ExpansionError::EmptyTileGroup); } let start = Instant::now(); let mut state = TileGraphBuilder::new(graph)?; - state.cache = cache; + state.cache = Arc::clone(&cache); let body = state.build_region(&graph.operations, checkpoints)?; for value in graph .inputs @@ -141,7 +125,7 @@ pub(crate) fn expand_tiles_analyzed( crate::low::passes::simplify(&mut program)?; let simplify_time = start.elapsed(); let start = Instant::now(); - relay::select(&mut program, analysis)?; + relay::select(&mut program, &cache)?; let relay_time = start.elapsed(); crate::low::initialization::omit_unread_fp8_input_padding(&mut program); crate::low::initialization::reuse_finite_padding(&mut program); @@ -157,7 +141,7 @@ pub(crate) fn expand_tiles_analyzed( } struct TileGraphBuilder { - cache: Arc, + cache: Arc, tile_count: u16, storage_groups: Vec, shards: Vec, @@ -175,7 +159,7 @@ impl TileGraphBuilder { fn new(graph: &MidProgram) -> ExpansionResult { let tile_count = graph.tile_count; let mut state = Self { - cache: Arc::new(ExpansionCache::default()), + cache: Arc::new(GeometryCache::default()), tile_count, storage_groups: graph .values diff --git a/crates/ipu-codegen/src/low/expand/movement.rs b/crates/ipu-codegen/src/low/expand/movement.rs index 947a748c..9f6b1467 100644 --- a/crates/ipu-codegen/src/low/expand/movement.rs +++ b/crates/ipu-codegen/src/low/expand/movement.rs @@ -605,10 +605,28 @@ impl TileGraphBuilder { |(mappings, order)| (mappings, order, order), ); let transfers = batch.transfers.entry(exchange_order).or_default(); - let geometry = - self.cache - .geometry(&self.shards, &mappings, destination_shard, copy_order)?; let destination = &self.shards[destination_shard.index() as usize]; + let geometric_mappings = mappings + .iter() + .map(|(source, target)| { + let source = source.bind(&self.shards)?; + target.bind(&self.shards)?; + if target.shard != destination_shard { + return Err(ExpansionError::InvalidCopyPlan); + } + Ok(crate::storage::CopyMapping { + source: source.shard.storage(), + source_extents: source.extents, + destination_extents: &target.extents, + }) + }) + .collect::>>()?; + let geometry = self.cache.destination( + destination.storage(), + &geometric_mappings, + copy_order, + ipu_exchange::MAX_TRANSFER_WORDS * 4, + )?; let preparation = select_destination_packing( &destination.tensor_type, &destination.extents, @@ -857,7 +875,7 @@ struct CopyStaging { fn select_destination_packing( destination: &TensorType, extents: &[ShardExtent], - geometry: &crate::storage::CopyGeometry, + geometry: &crate::storage::DestinationGeometry, policy: PackingPolicy, ) -> ExpansionResult> { let transform = geometry.semantic diff --git a/crates/ipu-codegen/src/low/expand/relay.rs b/crates/ipu-codegen/src/low/expand/relay.rs index ecc9f50e..4867ff8f 100644 --- a/crates/ipu-codegen/src/low/expand/relay.rs +++ b/crates/ipu-codegen/src/low/expand/relay.rs @@ -5,7 +5,7 @@ use crate::low::storage::storage_root; pub(super) fn select( program: &mut TileGraph, - analysis: &mut crate::estimate::GeometryAnalysis, + analysis: &crate::storage::GeometryCache, ) -> ExpansionResult<()> { for index in 0..program.exchange_phases.len() { if !program.exchange_phases[index] @@ -406,11 +406,7 @@ mod tests { fn relays_preserve_native_bytes_and_enter_normal_costing() { let mut program = fixture(); let expected = execute(&program); - select( - &mut program, - &mut crate::estimate::GeometryAnalysis::default(), - ) - .unwrap(); + select(&mut program, &crate::storage::GeometryCache::default()).unwrap(); assert_eq!(program.shards.len(), 52); assert_eq!(program.exchange_phases.len(), 1); assert_eq!(execute(&program), expected); @@ -438,11 +434,7 @@ mod tests { } } let before = program.clone(); - select( - &mut program, - &mut crate::estimate::GeometryAnalysis::default(), - ) - .unwrap(); + select(&mut program, &crate::storage::GeometryCache::default()).unwrap(); assert_eq!(program, before); } } diff --git a/crates/ipu-codegen/src/low/expand/tests.rs b/crates/ipu-codegen/src/low/expand/tests.rs index 34b73b03..ae61698d 100644 --- a/crates/ipu-codegen/src/low/expand/tests.rs +++ b/crates/ipu-codegen/src/low/expand/tests.rs @@ -7,8 +7,8 @@ fn lower_to_tiles( ) -> super::ExpansionResult { let mut graph = graph.clone(); graph.compose_copies(); - let cache = Arc::new(ExpansionCache::default()); - let expected = super::expand_tiles_cached(&graph, true, Arc::new(ExpansionCache::disabled()))?; + let cache = Arc::new(GeometryCache::default()); + let expected = super::expand_tiles_cached(&graph, true, Arc::new(GeometryCache::disabled()))?; for _ in 0..2 { let cached = super::expand_tiles_cached(&graph, true, Arc::clone(&cache))?; assert_eq!(cached, expected, "cache changed the complete low graph"); @@ -1507,7 +1507,7 @@ fn randomized_broadcast_adds_schedule_remote_singleton_views() { #[test] fn randomized_blocked_gemms_expand_to_tile_kernel_phases() { - let shared_cache = Arc::new(ExpansionCache::default()); + let shared_cache = Arc::new(GeometryCache::default()); let mut random = fastrand::Rng::with_seed(0x6765_6d6d); for case in 0..CASES { let tiles = 1_u16 << random.u32(0..=3); @@ -1527,7 +1527,7 @@ fn randomized_blocked_gemms_expand_to_tile_kernel_phases() { .with_input(right, format(tiles)); let mid = lower(&graph, &config, &Ipu21CostModel).unwrap(); let uncached = - super::expand_tiles_cached(&mid, true, Arc::new(ExpansionCache::disabled())).unwrap(); + super::expand_tiles_cached(&mid, true, Arc::new(GeometryCache::disabled())).unwrap(); for _ in 0..2 { let cached = super::expand_tiles_cached(&mid, true, Arc::clone(&shared_cache)).unwrap(); assert_eq!(cached, uncached, "cache changed graph in case {case}"); diff --git a/crates/ipu-codegen/src/low/storage.rs b/crates/ipu-codegen/src/low/storage.rs index 4df64980..5841afa4 100644 --- a/crates/ipu-codegen/src/low/storage.rs +++ b/crates/ipu-codegen/src/low/storage.rs @@ -73,6 +73,14 @@ impl ShardView { } impl BoundView<'_> { + pub(crate) fn geometry( + &self, + cache: &crate::storage::GeometryCache, + order: CopyOrder, + ) -> StorageResult> { + cache.view(self.shard.storage(), self.extents, order) + } + pub(crate) fn traversal(&self, order: CopyOrder) -> StorageResult { let storage = self.shard.storage(); match order { diff --git a/crates/ipu-codegen/src/storage.rs b/crates/ipu-codegen/src/storage.rs index 283a92e4..36c5822f 100644 --- a/crates/ipu-codegen/src/storage.rs +++ b/crates/ipu-codegen/src/storage.rs @@ -2,7 +2,12 @@ mod movement; pub use movement::CopyOrder; -pub(crate) use movement::{CopyGeometry, CopyMapping, for_each_copy_span, uncovered_bytes}; +pub(crate) use movement::{for_each_copy_span, uncovered_bytes}; +mod geometry; +pub use geometry::GeometryCacheStats; +pub(crate) use geometry::{ + CopyMapping, CopyPair, DestinationGeometry, GeometryCache, GeometryView, +}; mod traversal; pub(crate) use traversal::{ByteTraversal, StridedSpan, byte_traversal, panel_byte_traversal}; diff --git a/crates/ipu-codegen/src/storage/movement.rs b/crates/ipu-codegen/src/storage/movement.rs index 391f00b0..057b6a27 100644 --- a/crates/ipu-codegen/src/storage/movement.rs +++ b/crates/ipu-codegen/src/storage/movement.rs @@ -6,7 +6,6 @@ use super::{ storage_bytes, }; use crate::tensor::ShardExtent; -use std::sync::OnceLock; #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)] pub enum CopyOrder { @@ -19,94 +18,6 @@ pub enum CopyOrder { Panels, } -#[derive(Clone, Copy)] -pub(crate) struct CopyMapping<'a> { - pub source: TensorStorage<'a>, - pub source_extents: &'a [ShardExtent], - pub destination_extents: &'a [ShardExtent], -} - -/// Facts about populating one allocation. Alignment and fragment facts concern -/// semantic traversal; physical/panel requests retain their selected traversal. -pub(crate) struct CopyGeometry { - pub bytes: u32, - coverage: ByteTraversal, - uncovered: OnceLock>>, - pub fragments: Option, - pub semantic: bool, - pub destination_word_aligned: bool, - pub same_element_order: bool, - pub padding: bool, -} - -impl CopyGeometry { - pub(crate) fn heap_bytes(&self) -> usize { - self.coverage.heap_bytes() - + self - .uncovered - .get() - .and_then(|result| result.as_ref().ok()) - .map_or(0, |spans| spans.capacity() * size_of::()) - } - - pub(crate) fn analyze( - destination: TensorStorage<'_>, - mappings: &[CopyMapping<'_>], - order: CopyOrder, - max_fragment_bytes: u32, - ) -> StorageResult { - if max_fragment_bytes == 0 { - return Err(StorageError::InvalidView); - } - let mut fragments = (order == CopyOrder::Semantic).then_some(0u64); - let mut destination_unaligned = false; - if order == CopyOrder::Semantic { - for mapping in mappings { - let source = byte_traversal(mapping.source, mapping.source_extents, false)?; - let target = byte_traversal(destination, mapping.destination_extents, false)?; - destination_unaligned |= !target.word_aligned(); - fragments = match fragments { - Some(count) if source.word_aligned() && !destination_unaligned => source - .copy_fragments(&target, max_fragment_bytes) - .ok() - .map(|n| count.saturating_add(n)), - _ => None, - }; - } - } - Ok(Self { - bytes: storage_bytes(destination)?, - coverage: copy_coverage( - destination, - mappings.iter().map(|m| m.destination_extents), - order, - )?, - uncovered: OnceLock::new(), - fragments, - semantic: order == CopyOrder::Semantic, - destination_word_aligned: !destination_unaligned, - same_element_order: mappings.iter().all(|mapping| { - mapping.source.format.layout.order == destination.format.layout.order - }), - padding: destination - .extents - .iter() - .any(|extent| extent.physical_end > extent.logical_end), - }) - } - - /// Keep coverage symbolic until a realization actually needs clearing. - /// A packing kernel writes its padding itself, so scanning its span union - /// would be wasted work. The cached fact is independent of that choice. - pub(crate) fn uncovered(&self) -> StorageResult<&[ByteSpan]> { - self.uncovered - .get_or_init(|| uncovered_ranges(self.bytes, &self.coverage)) - .as_ref() - .map(Vec::as_slice) - .map_err(Clone::clone) - } -} - /// Exact byte coverage, not summed volume: overlapping mappings cannot hide holes. /// Any alignment required by a clearing implementation belongs to its caller. pub(crate) fn uncovered_bytes<'a>( @@ -140,7 +51,10 @@ fn copy_coverage<'a>( Ok(ByteTraversal::physical_union(covered)) } -fn uncovered_ranges(bytes: u32, covered: &ByteTraversal) -> StorageResult> { +pub(super) fn uncovered_ranges( + bytes: u32, + covered: &ByteTraversal, +) -> StorageResult> { let mut cursor = 0u32; let mut holes: Vec = Vec::new(); for span in covered.spans().chain(std::iter::once(ByteSpan {