diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..4034936 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,6 @@ +- It's more important that this done right than that this can be done quickly, and you can stop and ask me things if necessary. +- Add as few additional functions or data structures as possible to achieve your goals. +- No layer violations. If a change would require moving things between compiler layers or mixing behaviour between layers, ask before making it. +- Unit tests must be randomized and property-based, and not just test exact equivalence to some value. +- If you think there's a "cleaner final boundary" for a feature or change, surface it. +- If making a specified refactor, check all files using the old version to check that the refactor is done completely. diff --git a/Cargo.lock b/Cargo.lock index ad3c406..ee333dd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -329,6 +329,7 @@ dependencies = [ "ipu-package", "ipu-profile", "ipu-runtime", + "ipu-target", "serde_json", "tracing", "tracing-subscriber", @@ -342,8 +343,8 @@ dependencies = [ "foldhash", "ipu-driver", "ipu-elf", - "ipu-exchange", "ipu-package", + "ipu-target", "rayon", "serde", "thiserror", @@ -356,6 +357,7 @@ version = "0.1.0" dependencies = [ "fastrand", "ipu-package", + "ipu-target", "libc", "object", "thiserror", @@ -375,16 +377,6 @@ dependencies = [ "tracing", ] -[[package]] -name = "ipu-exchange" -version = "0.1.0" -dependencies = [ - "fastrand", - "serde", - "thiserror", - "tracing", -] - [[package]] name = "ipu-package" version = "0.1.0" @@ -392,6 +384,7 @@ dependencies = [ "capnp", "capnpc", "fastrand", + "ipu-target", "thiserror", "tracing", ] @@ -411,10 +404,19 @@ version = "0.1.0" dependencies = [ "ipu-driver", "ipu-package", - "thiserror", "tracing-subscriber", ] +[[package]] +name = "ipu-target" +version = "0.1.0" +dependencies = [ + "fastrand", + "serde", + "thiserror", + "tracing", +] + [[package]] name = "ipu-tests" version = "0.1.0" @@ -427,9 +429,9 @@ dependencies = [ "ipu-codegen", "ipu-driver", "ipu-elf", - "ipu-exchange", "ipu-package", "ipu-runtime", + "ipu-target", "rand_distr", "rand_xoshiro", "serde_json", diff --git a/Cargo.toml b/Cargo.toml index 4511359..f3a8abf 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [workspace] members = [ - "crates/ipu-exchange", + "crates/ipu-target", "crates/ipu-codegen", "crates/ipu-package", "crates/ipu-elf", diff --git a/README.md b/README.md index be10d7e..f2ea2f0 100644 --- a/README.md +++ b/README.md @@ -1,14 +1,15 @@ # ipu-stack -`ipu-stack` is a small collection of runtime, packaging, code-generation, and -early graph-lowering components for Graphcore IPU21 devices. It does not yet -contain the allocator or complete graph-to-device compiler. +`ipu-stack` is a collection of graph lowering, code generation, packaging, and +runtime components for Graphcore IPU21 devices. ## Components -- `ipu-exchange` generates and encodes device and host exchange programs. -- `ipu-codegen` emits straight-line supervisor code from caller-resolved tile - programs. +- `ipu-target` owns IPU21 topology and memory geometry, exchange generation and + parsing, instruction encoding, finalized tile programs, and supervisor-code + emission. +- `ipu-codegen` plans and lowers graphs and builds loadable packages using that + target API. - `ipu-elf` compiles Graphcore tile sources and links Colossus ELF objects. - `ipu-package` reads, writes, and validates `.ipuexe` application packages and cycle profiles. @@ -30,8 +31,8 @@ graph is shaped structured SSA. Its separate mid-level lowering selects precision and layout with a toy cost model and inserts explicit casts and rearrangements. Mid-to-low lowering then produces logical per-tile shard work, kernel runs, synchronized exchanges, and structured repeats. Package -construction produces completion-only tile programs until SRAM placement, -exchange encoding, and kernel-symbol selection are implemented. The config +construction resolves SRAM placement, exchange encoding, and kernel symbols +before `ipu-target` emits each finalized tile program. The config uses one shared `PipelineConfig` for target, tile count, input formats, operator catalog, scheduling, and profiling. `PackageConfig` adds the toolchain, static runtime source, and build directory. * Unmerged path TODO diff --git a/crates/ipu-cli/Cargo.toml b/crates/ipu-cli/Cargo.toml index 32617f1..0c4e673 100644 --- a/crates/ipu-cli/Cargo.toml +++ b/crates/ipu-cli/Cargo.toml @@ -18,6 +18,7 @@ ipu-elf = { path = "../ipu-elf" } ipu-package = { path = "../ipu-package" } ipu-profile = { path = "../ipu-profile" } ipu-runtime = { path = "../ipu-runtime" } +ipu-target = { path = "../ipu-target" } serde_json.workspace = true tracing.workspace = true tracing-subscriber.workspace = true * Unmerged path crates/ipu-cli/src/main.rs diff --git a/crates/ipu-cli/src/profile_report.html b/crates/ipu-cli/src/profile_report.html index 9e70b39..80eb080 100644 --- a/crates/ipu-cli/src/profile_report.html +++ b/crates/ipu-cli/src/profile_report.html @@ -158,7 +158,7 @@ canvas { tx rx tx/rx - partner bus + partner bus exchange idle sync @@ -225,7 +225,7 @@ const colors = { send: "#17854c", receive: "#287db2", sendReceive: "#6b5596", - partnerBusy: "#9a547f", + partnerBusy: "#ffa756", exchangeIdle: "#8ebeb4", compute: "#cf3f4a", synchronization: "#c58b24", diff --git a/crates/ipu-codegen/Cargo.toml b/crates/ipu-codegen/Cargo.toml index dfa72f0..5f86789 100644 --- a/crates/ipu-codegen/Cargo.toml +++ b/crates/ipu-codegen/Cargo.toml @@ -6,7 +6,7 @@ license.workspace = true [dependencies] foldhash.workspace = true -ipu-exchange = { path = "../ipu-exchange" } +ipu-target = { path = "../ipu-target" } ipu-driver = { path = "../ipu-driver" } ipu-elf = { path = "../ipu-elf" } ipu-package = { path = "../ipu-package" } diff --git a/crates/ipu-codegen/src/config.rs b/crates/ipu-codegen/src/config.rs new file mode 100644 index 0000000..9a0ca01 --- /dev/null +++ b/crates/ipu-codegen/src/config.rs @@ -0,0 +1,341 @@ +//! Compiler target and planner search configuration. + +use crate::graph::{TensorShape, ValueId}; +use crate::layout::{MemoryClass, TensorFormat}; +use crate::operator::{GemmPlanConstraint, Precision}; +use ipu_target::hardware::HardwareTarget; +use std::collections::BTreeMap; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct PipelineConfig { + pub target: HardwareTarget, + pub tile_count: u16, + pub inputs: BTreeMap, + /// Graph-boundary tensors whose layout may be selected by their first + /// consumer. Precision remains fixed, while packaging exposes the chosen + /// physical layout directly through the host binding. + pub automatic_inputs: BTreeMap, + pub search_domain: PlannerSearchDomain, + pub planning_beam_width: usize, + /// Complete beam finalists ranked with the physical exchange scheduler. + pub exchange_schedule_finalists: usize, + /// SRAM retained for exchange tables, profiling, host commands, and code. + pub standard_memory_reservation_bytes: u64, + /// Maximum SRAM per tile available to planned values and the reservation. + pub tile_memory_budget_bytes: u64, + pub profiling: ProfilingConfig, + pub diagnostic_checkpoints: bool, + pub conversion_streaming: ConversionStreamingPolicy, + pub(crate) resolved_active_tile_counts: Vec, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum ConversionStreamingPolicy { + Never, + #[default] + WhenRequired, + Always, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum AttentionStrategy { + #[default] + Automatic, + Flash, + Materialized, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum OperatorClass { + Gemm, + Gelu, + Add, + Attention, +} + +/// Search axes shared by semantic operator plan generators. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct PlannerSearchDomain { + /// `None` derives useful tile counts from graph shapes; `Some` restricts + /// planning to the explicitly supplied counts. + pub(crate) active_tiles: Option>, + pub(crate) operator_precisions: BTreeMap>, + pub(crate) weight_memory_classes: Vec, + pub(crate) attention_strategy: AttentionStrategy, + pub(crate) gemm_plan_constraints: Vec, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum ProfilingConfig { + #[default] + Disabled, + Overall, + Full, +} + +impl std::fmt::Display for AttentionStrategy { + fn fmt(&self, output: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + output.write_str(match self { + Self::Automatic => "auto", + Self::Flash => "flash", + Self::Materialized => "materialized", + }) + } +} + +impl std::str::FromStr for AttentionStrategy { + type Err = &'static str; + + fn from_str(value: &str) -> Result { + match value { + "auto" | "automatic" => Ok(Self::Automatic), + "flash" => Ok(Self::Flash), + "materialized" => Ok(Self::Materialized), + _ => Err("expected auto, flash, or materialized"), + } + } +} + +impl Default for PlannerSearchDomain { + fn default() -> Self { + Self { + active_tiles: None, + operator_precisions: BTreeMap::from([ + (OperatorClass::Gemm, vec![Precision::F16, Precision::F32]), + (OperatorClass::Gelu, vec![Precision::F16, Precision::F32]), + (OperatorClass::Add, vec![Precision::F16, Precision::F32]), + (OperatorClass::Attention, vec![Precision::F16]), + ]), + weight_memory_classes: vec![MemoryClass::Standard, MemoryClass::Interleaved], + attention_strategy: AttentionStrategy::Automatic, + gemm_plan_constraints: Vec::new(), + } + } +} + +impl PlannerSearchDomain { + pub(crate) fn precisions(&self, operator: OperatorClass) -> &[Precision] { + self.operator_precisions + .get(&operator) + .map_or(&[], Vec::as_slice) + } + + pub(crate) fn permits_precision(&self, operator: OperatorClass, precision: Precision) -> bool { + self.precisions(operator).contains(&precision) + } + + pub(crate) fn permits_weight_memory(&self, memory_class: MemoryClass) -> bool { + self.weight_memory_classes.contains(&memory_class) + } + + pub(crate) fn active_tile_counts<'a>( + &self, + capacity: u16, + shapes: impl IntoIterator, + ) -> Vec { + match &self.active_tiles { + None => { + let mut counts = candidate_active_tile_counts(capacity); + for count in shape_aware_active_tile_counts(capacity, shapes) { + if !counts.contains(&count) { + counts.push(count); + } + } + counts + } + Some(counts) => counts + .iter() + .copied() + .filter(|&count| count <= capacity) + .collect(), + } + } + + pub fn with_operator_precisions( + mut self, + operator: OperatorClass, + precisions: impl IntoIterator, + ) -> Self { + let mut unique = Vec::new(); + for precision in precisions { + if !unique.contains(&precision) { + unique.push(precision); + } + } + self.operator_precisions.insert(operator, unique); + self + } + + pub fn with_weight_memory_classes( + mut self, + classes: impl IntoIterator, + ) -> Self { + let mut unique = Vec::new(); + for class in classes { + if !unique.contains(&class) { + unique.push(class); + } + } + self.weight_memory_classes = unique; + self + } + + pub fn with_active_tile_counts(mut self, counts: impl IntoIterator) -> Self { + let mut active_tiles = Vec::new(); + for count in counts { + if count != 0 && !active_tiles.contains(&count) { + active_tiles.push(count); + } + } + self.active_tiles = Some(active_tiles); + self + } + + pub fn with_attention_strategy(mut self, strategy: AttentionStrategy) -> Self { + self.attention_strategy = strategy; + self + } + + pub fn with_gemm_plan_constraint(mut self, constraint: GemmPlanConstraint) -> Self { + self.gemm_plan_constraints + .retain(|existing| existing.source_operation != constraint.source_operation); + self.gemm_plan_constraints.push(constraint); + self + } +} + +impl ProfilingConfig { + pub const fn records_overall_time(self) -> bool { + !matches!(self, Self::Disabled) + } + + pub const fn records_steps(self) -> bool { + matches!(self, Self::Full) + } +} + +impl std::fmt::Display for ProfilingConfig { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str(match self { + Self::Disabled => "none", + Self::Overall => "overall", + Self::Full => "full", + }) + } +} + +impl std::str::FromStr for ProfilingConfig { + type Err = &'static str; + + fn from_str(value: &str) -> Result { + match value { + "none" | "disabled" => Ok(Self::Disabled), + "overall" => Ok(Self::Overall), + "full" => Ok(Self::Full), + _ => Err("profiling mode must be one of: none, overall, full"), + } + } +} + +impl PipelineConfig { + pub fn new(tile_count: u16) -> Self { + let target = HardwareTarget::Ipu21; + let memory = target.memory_constraints(); + Self { + target, + tile_count, + inputs: BTreeMap::new(), + automatic_inputs: BTreeMap::new(), + search_domain: PlannerSearchDomain::default(), + planning_beam_width: 64, + exchange_schedule_finalists: 1, + standard_memory_reservation_bytes: memory.default_standard_reservation_bytes, + tile_memory_budget_bytes: memory.total_bytes, + profiling: ProfilingConfig::default(), + diagnostic_checkpoints: false, + conversion_streaming: ConversionStreamingPolicy::WhenRequired, + resolved_active_tile_counts: Vec::new(), + } + } + + pub fn with_input(mut self, value: ValueId, format: TensorFormat) -> Self { + self.inputs.insert(value, format); + self.automatic_inputs.remove(&value); + self + } + + pub fn with_automatic_input(mut self, value: ValueId, precision: Precision) -> Self { + self.inputs.remove(&value); + self.automatic_inputs.insert(value, precision); + self + } + + pub fn with_planning_beam_width(mut self, width: usize) -> Self { + self.planning_beam_width = width.max(1); + self + } + + pub fn with_exchange_schedule_finalists(mut self, finalists: usize) -> Self { + self.exchange_schedule_finalists = finalists.max(1); + self + } + + pub fn with_search_domain(mut self, search_domain: PlannerSearchDomain) -> Self { + self.search_domain = search_domain; + self + } + + pub fn with_standard_memory_reservation(mut self, bytes: u64) -> Self { + self.standard_memory_reservation_bytes = bytes; + self + } + + pub fn with_tile_memory_budget(mut self, bytes: u64) -> Self { + self.tile_memory_budget_bytes = bytes; + self + } +} + +pub(crate) fn candidate_active_tile_counts(capacity: u16) -> Vec { + if capacity == 0 { + return vec![0]; + } + let mut counts = vec![capacity]; + // Power-of-two subsets provide progressively smaller fallback grids. + let mut power = 1u16; + while let Some(next) = power.checked_mul(2) { + if next > capacity { + break; + } + power = next; + } + loop { + if !counts.contains(&power) { + counts.push(power); + } + if power == 1 { + break; + } + power /= 2; + } + counts +} + +pub(crate) fn shape_aware_active_tile_counts<'a>( + capacity: u16, + shapes: impl IntoIterator, +) -> Vec { + let minimum = capacity.div_ceil(2); + let mut counts = shapes + .into_iter() + .flat_map(|shape| shape.0.iter().copied()) + .filter_map(|extent| { + let extent = u16::try_from(extent).ok()?; + (extent > 1 && extent <= capacity).then(|| capacity / extent * extent) + }) + .filter(|&count| count >= minimum && count < capacity) + .collect::>(); + counts.sort_unstable_by(|left, right| right.cmp(left)); + counts.dedup(); + counts +} diff --git a/crates/ipu-codegen/src/conversion.rs b/crates/ipu-codegen/src/conversion.rs new file mode 100644 index 0000000..57d1add --- /dev/null +++ b/crates/ipu-codegen/src/conversion.rs @@ -0,0 +1,1244 @@ +//! Address-independent layout-conversion routes and copy geometry. + +use crate::graph::TensorShape; +use crate::layout::{ + AMP_COLUMN_MICRO, Layout, LayoutError, NativeKernelOrder, StorageOrder, TensorRegion, + TensorType, +}; +use crate::mid::{MidGraph, MidOperation, MidOperationKind, MidValue}; +use crate::operator::Precision; +use crate::storage::{StorageError, amp_micro_dimension, physical_byte_offset}; +use std::collections::BTreeMap; + +/// Address-independent strategy for materializing a format conversion. +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum ConversionStrategy { + LocalKernel, + DirectRetile, + DirectLogical, + StageLogicalThenTransform, +} + +impl ConversionStrategy { + pub const fn uses_intersections(self) -> bool { + !matches!(self, Self::LocalKernel) + } +} + +pub fn layout_conversion_strategy( + precision: Precision, + from: &Layout, + to: &Layout, +) -> ConversionStrategy { + if from.order == to.order { + ConversionStrategy::DirectRetile + } else if precision == Precision::F32 + || matches!(from.order, StorageOrder::Linear) + && matches!( + to.order, + StorageOrder::Native(NativeKernelOrder::Left | NativeKernelOrder::Output) + ) + || matches!(to.order, StorageOrder::Linear) + || matches!( + (from.order, to.order), + (StorageOrder::Blocked(source), StorageOrder::Blocked(destination)) + if source.axes == destination.axes + && source.permutation == destination.permutation + ) + { + ConversionStrategy::DirectLogical + } else { + ConversionStrategy::StageLogicalThenTransform + } +} + +/// One arbitrary-rank affine copy nest within a conversion route. +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct CopyGeometry { + pub source_offset: u32, + pub destination_offset: u32, + pub contiguous_bytes: u32, + /// Inner to outer affine dimensions. + pub dimensions: Vec, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct CopyDimension { + pub count: u32, + pub source_stride: u32, + pub destination_stride: u32, +} + +impl CopyGeometry { + pub fn copy_count(&self) -> u64 { + self.dimensions.iter().fold(1, |copies, dimension| { + copies.saturating_mul(u64::from(dimension.count)) + }) + } + + pub fn bytes(&self) -> u64 { + u64::from(self.contiguous_bytes).saturating_mul(self.copy_count()) + } + + pub(crate) fn offsets(&self) -> Option> { + let mut offsets = vec![(self.source_offset, self.destination_offset)]; + for dimension in &self.dimensions { + let inner = offsets.clone(); + offsets.clear(); + offsets.reserve(inner.len().checked_mul(dimension.count as usize)?); + for index in 0..dimension.count { + let source_delta = index.checked_mul(dimension.source_stride)?; + let destination_delta = index.checked_mul(dimension.destination_stride)?; + for &(source, destination) in &inner { + offsets.push(( + source.checked_add(source_delta)?, + destination.checked_add(destination_delta)?, + )); + } + } + } + Some(offsets) + } +} + +/// A semantic intersection between one resolved source shard and one +/// resolved destination shard. Regions remain address-independent; low +/// lowering binds them to the corresponding shard IDs. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ConversionMapping { + pub source_shard: u32, + pub source_tile: u16, + pub source_storage: TensorRegion, + pub destination_shard: u32, + pub destination_tile: u16, + pub destination_storage: TensorRegion, + pub source_region: TensorRegion, + pub destination_region: TensorRegion, + /// Local copies from the source layout into word-aligned transfer staging. + pub source_copies: Vec, + pub copies: Vec, + /// Local copies from word-aligned transfer staging into the planned + /// destination storage. + pub destination_copies: Vec, +} + +/// A logical view whose physical materialization may be deferred until a +/// consumer requests bounded slices. +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)] +pub enum DeferredTransform { + /// Split the final input axis into `parts` equal-width slices and fold the + /// new part axis into the leading output axis. + SplitLastAxisIntoLeading { parts: u32 }, +} + +impl DeferredTransform { + pub fn map_slices( + self, + source_shape: &TensorShape, + output_shape: &TensorShape, + output: &TensorRegion, + ) -> Option> { + let streams = output.extents.first()?; + (streams.start..streams.logical_end) + .map(|stream| { + let mut slice = output.clone(); + slice.extents[0].start = stream; + slice.extents[0].logical_end = stream.checked_add(1)?; + slice.extents[0].physical_end = stream.checked_add(1)?; + self.map_slice(source_shape, output_shape, &slice) + }) + .collect() + } + + pub(crate) fn map_slice( + self, + source_shape: &TensorShape, + output_shape: &TensorShape, + output: &TensorRegion, + ) -> Option { + let Self::SplitLastAxisIntoLeading { parts } = self; + let [source_batch, source_rows, source_columns] = source_shape.0.as_slice() else { + return None; + }; + let [output_streams, output_rows, output_columns] = output_shape.0.as_slice() else { + return None; + }; + let [stream, rows, columns] = output.extents.as_slice() else { + return None; + }; + if parts == 0 + || stream.axis != 0 + || rows.axis != 1 + || columns.axis != 2 + || stream.logical_end != stream.start.checked_add(1)? + || *output_streams != source_batch.checked_mul(parts)? + || output_rows != source_rows + || source_columns != &output_columns.checked_mul(parts)? + || stream.logical_end > *output_streams + || rows.logical_end > *output_rows + || columns.logical_end > *output_columns + { + return None; + } + let batch = stream.start / parts; + let column_base = (stream.start % parts).checked_mul(*output_columns)?; + Some(DeferredSliceMapping { + source: TensorRegion::new([ + crate::ShardExtent { + axis: 0, + start: batch, + logical_end: batch.checked_add(1)?, + physical_end: batch.checked_add(1)?, + }, + *rows, + crate::ShardExtent { + axis: 2, + start: column_base.checked_add(columns.start)?, + logical_end: column_base.checked_add(columns.logical_end)?, + physical_end: column_base.checked_add(columns.logical_end)?, + }, + ]), + destination: output.clone(), + source_axes: vec![None, Some(1), Some(2)], + }) + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct DeferredSliceMapping { + pub source: TensorRegion, + pub destination: TensorRegion, + /// Source axis corresponding to each destination axis. `None` selects a + /// fixed source coordinate. + pub source_axes: Vec>, +} + +impl DeferredSliceMapping { + pub fn project_source(&self, source: &TensorRegion) -> Option { + if source.len() != self.source.len() || self.source_axes.len() != self.destination.len() { + return None; + } + self.source_axes + .iter() + .zip(self.destination.iter()) + .enumerate() + .map(|(axis, (source_axis, destination))| { + let Some(source_axis) = source_axis else { + return Some(*destination); + }; + let selected = source.get(*source_axis)?; + let base = self.source.get(*source_axis)?.start; + Some(crate::ShardExtent { + axis: u16::try_from(axis).ok()?, + start: destination + .start + .checked_add(selected.start.checked_sub(base)?)?, + logical_end: destination + .start + .checked_add(selected.logical_end.checked_sub(base)?)?, + physical_end: destination + .start + .checked_add(selected.logical_end.checked_sub(base)?)?, + }) + }) + .collect() + } +} + +#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)] +pub enum ConversionGeometryError { + #[error(transparent)] + Layout(#[from] LayoutError), + #[error(transparent)] + Storage(#[from] StorageError), + #[error("layout conversion has no regular block geometry")] + Unsupported, + #[error("layout conversion geometry overflowed")] + Overflow, +} + +/// Resolves ownership intersections and their regular local/remote copy +/// nests from tensor shape and layouts. No SRAM addresses or low-level shard +/// identities participate in this decision. +pub(crate) fn plan_conversion( + shape: &TensorShape, + precision: Precision, + from: &Layout, + to: &Layout, + strategy: ConversionStrategy, +) -> Result, ConversionGeometryError> { + if strategy == ConversionStrategy::LocalKernel { + return Ok(Vec::new()); + } + let sources = from.resolve(shape)?.shard_extents(); + let destinations = to.resolve(shape)?.shard_extents(); + let mut mappings = Vec::new(); + for (destination_index, destination) in destinations.into_iter().enumerate() { + let mut intersections = BTreeMap::::new(); + for (source_index, source) in sources.iter().enumerate() { + let physical_retile = + from.order == to.order && strategy == ConversionStrategy::DirectRetile; + let source_region = if physical_retile { + source.extents.physical() + } else { + source.extents.logical() + }; + let destination_region = if physical_retile { + destination.extents.physical() + } else { + destination.extents.logical() + }; + let Some(region) = source_region.intersection(&destination_region) else { + continue; + }; + let selected = intersections + .entry(region) + .or_insert((source_index, source)); + if source.tile == destination.tile { + *selected = (source_index, source); + } + } + for (region, (source_index, source)) in intersections { + let destination_storage = if strategy == ConversionStrategy::StageLogicalThenTransform { + destination.extents.logical() + } else { + destination.extents.clone() + }; + mappings.push(plan_mapping( + precision, + from.order, + to.order, + strategy, + ConversionMapping { + source_shard: u32::try_from(source_index) + .map_err(|_| ConversionGeometryError::Overflow)?, + source_tile: source.tile, + source_storage: source.extents.clone(), + destination_shard: u32::try_from(destination_index) + .map_err(|_| ConversionGeometryError::Overflow)?, + destination_tile: destination.tile, + destination_storage, + source_region: region.clone(), + destination_region: region, + source_copies: Vec::new(), + copies: Vec::new(), + destination_copies: Vec::new(), + }, + )?); + } + } + if mappings.is_empty() { + return Err(ConversionGeometryError::Unsupported); + } + Ok(mappings) +} + +pub(crate) fn plan_view_conversion( + source: &TensorType, + destination: &TensorType, + transform: DeferredTransform, +) -> Result<(ConversionStrategy, Vec), ConversionGeometryError> { + let strategy = layout_conversion_strategy( + source.format.precision, + &source.format.layout, + &destination.format.layout, + ); + let sources = source.format.layout.resolve(&source.shape)?.shard_extents(); + let destinations = destination + .format + .layout + .resolve(&destination.shape)? + .shard_extents(); + let mut mappings = Vec::new(); + for (destination_index, output) in destinations.into_iter().enumerate() { + let destination_storage = if strategy == ConversionStrategy::StageLogicalThenTransform { + output.extents.logical() + } else { + output.extents.clone() + }; + for view in transform + .map_slices(&source.shape, &destination.shape, &output.extents.logical()) + .ok_or(ConversionGeometryError::Unsupported)? + { + let mut intersections = BTreeMap::::new(); + for (source_index, input) in sources.iter().enumerate() { + let Some(region) = input.extents.logical().intersection(&view.source) else { + continue; + }; + let selected = intersections.entry(region).or_insert((source_index, input)); + if input.tile == output.tile { + *selected = (source_index, input); + } + } + for (source_region, (source_index, input)) in intersections { + let destination_region = view + .project_source(&source_region) + .ok_or(ConversionGeometryError::Unsupported)?; + mappings.push(plan_mapping( + source.format.precision, + source.format.layout.order, + destination.format.layout.order, + strategy, + ConversionMapping { + source_shard: u32::try_from(source_index) + .map_err(|_| ConversionGeometryError::Overflow)?, + source_tile: input.tile, + source_storage: input.extents.clone(), + destination_shard: u32::try_from(destination_index) + .map_err(|_| ConversionGeometryError::Overflow)?, + destination_tile: output.tile, + destination_storage: destination_storage.clone(), + source_region, + destination_region, + source_copies: Vec::new(), + copies: Vec::new(), + destination_copies: Vec::new(), + }, + )?); + } + } + } + if mappings.is_empty() { + return Err(ConversionGeometryError::Unsupported); + } + Ok((strategy, mappings)) +} + +pub(crate) fn finalize_conversion_plans( + graph: &mut MidGraph, +) -> Result<(), ConversionGeometryError> { + finalize_operations(&mut graph.operations, &graph.values) +} + +fn finalize_operations( + operations: &mut [MidOperation], + values: &[MidValue], +) -> Result<(), ConversionGeometryError> { + for operation in operations { + match &mut operation.kind { + MidOperationKind::Convert(transform, strategy, _, mappings) => { + let source = &values + .get(operation.inputs[0].index() as usize) + .ok_or(ConversionGeometryError::Unsupported)? + .tensor_type; + let destination = &values + .get(operation.results[0].index() as usize) + .ok_or(ConversionGeometryError::Unsupported)? + .tensor_type; + *mappings = if let Some(transform) = transform { + let (resolved, planned) = + plan_view_conversion(source, destination, *transform)?; + if resolved != *strategy { + return Err(ConversionGeometryError::Unsupported); + } + planned + } else { + plan_conversion( + &destination.shape, + destination.format.precision, + &source.format.layout, + &destination.format.layout, + *strategy, + )? + }; + } + MidOperationKind::Repeat(repeat) => { + finalize_operations(&mut repeat.body.operations, values)?; + } + MidOperationKind::Operator(_) | MidOperationKind::CastPrecision => {} + } + } + Ok(()) +} + +fn plan_mapping( + precision: Precision, + from: StorageOrder, + to: StorageOrder, + strategy: ConversionStrategy, + mut mapping: ConversionMapping, +) -> Result { + let (source_order, destination_order, logical_order) = match strategy { + ConversionStrategy::DirectRetile if from == to => (from, to, false), + ConversionStrategy::DirectLogical => (from, to, true), + ConversionStrategy::StageLogicalThenTransform => { + (from, StorageOrder::Linear, from != StorageOrder::Linear) + } + ConversionStrategy::DirectRetile | ConversionStrategy::LocalKernel => { + return Err(ConversionGeometryError::Unsupported); + } + }; + let direct = copy_geometries_between( + precision, + ( + source_order, + &mapping.source_storage, + &mapping.source_region, + ), + ( + destination_order, + &mapping.destination_storage, + &mapping.destination_region, + ), + mapping.source_tile == mapping.destination_tile, + logical_order, + ); + let (source_copies, copies, destination_copies) = match direct { + Ok(copies) => (Vec::new(), copies, Vec::new()), + Err(ConversionGeometryError::Unsupported) => { + let logical_storage = mapping.destination_region.logical(); + let source_copies = copy_geometries_between( + precision, + ( + source_order, + &mapping.source_storage, + &mapping.source_region, + ), + ( + StorageOrder::Linear, + &logical_storage, + &mapping.destination_region, + ), + true, + source_order != StorageOrder::Linear, + )?; + let destination_copies = copy_geometries( + precision, + StorageOrder::Linear, + destination_order, + &logical_storage, + &mapping.destination_storage, + &mapping.destination_region, + true, + destination_order != StorageOrder::Linear, + )?; + let bytes = mapping + .source_region + .logical_elements() + .checked_mul(precision.bytes()) + .and_then(|bytes| u32::try_from(bytes).ok()) + .ok_or(ConversionGeometryError::Overflow)?; + ( + source_copies, + vec![CopyGeometry { + source_offset: 0, + destination_offset: 0, + contiguous_bytes: bytes.div_ceil(4) * 4, + dimensions: Vec::new(), + }], + destination_copies, + ) + } + Err(error) => return Err(error), + }; + mapping.source_copies = source_copies; + mapping.copies = copies; + mapping.destination_copies = destination_copies; + Ok(mapping) +} + +fn copy_geometries( + precision: Precision, + source_order: StorageOrder, + destination_order: StorageOrder, + source_storage: &TensorRegion, + destination_storage: &TensorRegion, + region: &TensorRegion, + local: bool, + logical_order: bool, +) -> Result, ConversionGeometryError> { + if source_storage.len() != destination_storage.len() || source_storage.len() != region.len() { + return Err(ConversionGeometryError::Unsupported); + } + let dimensions = if logical_order { + vec![1; region.len()] + } else { + if source_order != destination_order { + return Err(ConversionGeometryError::Unsupported); + } + physical_block_dimensions(source_order, precision, region)? + }; + let geometries = match affine_geometries( + source_order, + destination_order, + precision, + source_storage, + destination_storage, + region, + &dimensions, + logical_order, + ) { + Ok(geometries) => geometries, + Err(ConversionGeometryError::Unsupported) if local && logical_order => { + let scalar = vec![1; region.len()]; + affine_geometries( + source_order, + destination_order, + precision, + source_storage, + destination_storage, + region, + &scalar, + true, + )? + } + Err(error) => return Err(error), + }; + let alignment = if local { precision.bytes() as u32 } else { 4 }; + for geometry in &geometries { + if !geometry.source_offset.is_multiple_of(alignment) + || !geometry.destination_offset.is_multiple_of(alignment) + || geometry.contiguous_bytes == 0 + || !geometry.contiguous_bytes.is_multiple_of(alignment) + { + return Err(ConversionGeometryError::Unsupported); + } + } + Ok(geometries) +} + +fn copy_geometries_between( + precision: Precision, + source: (StorageOrder, &TensorRegion, &TensorRegion), + destination: (StorageOrder, &TensorRegion, &TensorRegion), + local: bool, + logical_order: bool, +) -> Result, ConversionGeometryError> { + let (source_order, source_storage, source_region) = source; + let (destination_order, destination_storage, destination_region) = destination; + if source_storage.len() != destination_storage.len() + || source_storage.len() != source_region.len() + || source_region.len() != destination_region.len() + { + return Err(ConversionGeometryError::Unsupported); + } + let mut normalized_source = Vec::with_capacity(source_region.len()); + let mut normalized_destination = Vec::with_capacity(source_region.len()); + let mut normalized_region = Vec::with_capacity(source_region.len()); + for (axis, (((source_storage, destination_storage), source), destination)) in source_storage + .iter() + .zip(destination_storage.iter()) + .zip(source_region.iter()) + .zip(destination_region.iter()) + .enumerate() + { + let source_width = source.logical_end - source.start; + if source_width != destination.logical_end - destination.start + || source.start < source_storage.start + || destination.start < destination_storage.start + { + return Err(ConversionGeometryError::Unsupported); + } + let source_delta = source.start - source_storage.start; + let destination_delta = destination.start - destination_storage.start; + let common_start = source_delta.max(destination_delta); + let source_start = common_start - source_delta; + let destination_start = common_start - destination_delta; + let axis = u16::try_from(axis).map_err(|_| ConversionGeometryError::Overflow)?; + let extent = |start: u32, + storage: &crate::ShardExtent| + -> Result { + Ok(crate::ShardExtent { + axis, + start, + logical_end: start + .checked_add(storage.logical_end - storage.start) + .ok_or(ConversionGeometryError::Overflow)?, + physical_end: start + .checked_add(storage.physical_end - storage.start) + .ok_or(ConversionGeometryError::Overflow)?, + }) + }; + normalized_source.push(extent(source_start, source_storage)?); + normalized_destination.push(extent(destination_start, destination_storage)?); + normalized_region.push(crate::ShardExtent { + axis, + start: common_start, + logical_end: common_start + .checked_add(source_width) + .ok_or(ConversionGeometryError::Overflow)?, + physical_end: common_start + .checked_add(source_width) + .ok_or(ConversionGeometryError::Overflow)?, + }); + } + copy_geometries( + precision, + source_order, + destination_order, + &normalized_source.into(), + &normalized_destination.into(), + &normalized_region.into(), + local, + logical_order, + ) +} + +fn affine_geometries( + source_order: StorageOrder, + destination_order: StorageOrder, + precision: Precision, + source_storage: &TensorRegion, + destination_storage: &TensorRegion, + region: &TensorRegion, + dimensions: &[u32], + logical_order: bool, +) -> Result, ConversionGeometryError> { + let mut pending = vec![region.clone()]; + let mut geometries = Vec::new(); + while let Some(region) = pending.pop() { + if logical_order { + let cells = [ + ( + affine_cell_dimensions(source_order, precision, region.len())?, + source_storage, + ), + ( + affine_cell_dimensions(destination_order, precision, region.len())?, + destination_storage, + ), + ]; + let mut split = None; + 'orders: for (dimensions, storage) in cells { + for (axis, (&dimension, extent)) in + dimensions.iter().zip(storage.iter()).enumerate() + { + if dimension == 0 { + continue; + } + let local = region[axis].start - extent.start; + let boundary = extent + .start + .checked_add( + local + .checked_div(dimension) + .and_then(|cell| cell.checked_add(1)) + .and_then(|cell| cell.checked_mul(dimension)) + .ok_or(ConversionGeometryError::Overflow)?, + ) + .ok_or(ConversionGeometryError::Overflow)?; + if boundary < region[axis].logical_end { + split = Some((axis, boundary)); + break 'orders; + } + } + } + if let Some((axis, boundary)) = split { + let mut first = region.clone(); + let mut second = region; + first[axis].logical_end = boundary; + first[axis].physical_end = boundary; + second[axis].start = boundary; + pending.push(second); + pending.push(first); + continue; + } + } + match affine_geometry( + source_order, + destination_order, + precision, + source_storage, + destination_storage, + ®ion, + dimensions, + ) { + Ok(geometry) => geometries.push(geometry), + Err(ConversionGeometryError::Unsupported) => { + let Some((axis, count)) = region + .iter() + .enumerate() + .filter_map(|(axis, extent)| { + let count = (extent.logical_end - extent.start) / dimensions[axis]; + (count > 1).then_some((axis, count)) + }) + .max_by_key(|&(_, count)| count) + else { + return Err(ConversionGeometryError::Unsupported); + }; + let split = region[axis].start + count.div_ceil(2) * dimensions[axis]; + let mut first = region.clone(); + let mut second = region; + first[axis].logical_end = split; + first[axis].physical_end = split; + second[axis].start = split; + pending.push(second); + pending.push(first); + } + Err(error) => return Err(error), + } + } + compact_geometries(geometries) +} + +fn compact_geometries( + mut geometries: Vec, +) -> Result, ConversionGeometryError> { + geometries.sort_by_key(|geometry| geometry.source_offset); + loop { + let mut compacted = Vec::new(); + let mut merged = false; + let mut index = 0; + while index < geometries.len() { + let mut geometry = geometries[index].clone(); + let mut end = index + 1; + let Some(second) = geometries.get(end).filter(|second| { + second.contiguous_bytes == geometry.contiguous_bytes + && second.dimensions == geometry.dimensions + }) else { + compacted.push(geometry); + index = end; + continue; + }; + let Some(source_stride) = second.source_offset.checked_sub(geometry.source_offset) + else { + compacted.push(geometry); + index = end; + continue; + }; + let Some(destination_stride) = second + .destination_offset + .checked_sub(geometry.destination_offset) + else { + compacted.push(geometry); + index = end; + continue; + }; + if source_stride == 0 || destination_stride == 0 { + compacted.push(geometry); + index = end; + continue; + } + end += 1; + while let Some(next) = geometries.get(end) { + let previous = &geometries[end - 1]; + if next.contiguous_bytes != geometry.contiguous_bytes + || next.dimensions != geometry.dimensions + || next.source_offset.checked_sub(previous.source_offset) != Some(source_stride) + || next + .destination_offset + .checked_sub(previous.destination_offset) + != Some(destination_stride) + { + break; + } + end += 1; + } + geometry.dimensions.push(CopyDimension { + count: u32::try_from(end - index).map_err(|_| ConversionGeometryError::Overflow)?, + source_stride, + destination_stride, + }); + compacted.push(geometry); + merged = true; + index = end; + } + geometries = compacted; + if !merged { + return Ok(geometries); + } + } +} + +fn affine_cell_dimensions( + order: StorageOrder, + precision: Precision, + rank: usize, +) -> Result, ConversionGeometryError> { + let mut dimensions = vec![0; rank]; + if rank < 2 || order == StorageOrder::Linear { + return Ok(dimensions); + } + let row = rank - 2; + let column = rank - 1; + match order { + StorageOrder::Linear => {} + StorageOrder::Native(NativeKernelOrder::Left) => { + dimensions[column] = amp_micro_dimension(precision); + } + StorageOrder::Native(NativeKernelOrder::TransposedLeft) => { + dimensions[row] = amp_micro_dimension(precision); + } + StorageOrder::Native(NativeKernelOrder::TransposedRight) => { + dimensions[row] = amp_micro_dimension(precision); + dimensions[column] = AMP_COLUMN_MICRO; + } + StorageOrder::Native(NativeKernelOrder::Output) => dimensions[column] = 2, + StorageOrder::Native(NativeKernelOrder::TransposedOutput) => dimensions[row] = 2, + StorageOrder::Blocked(order) => { + let [row, column] = order.physical_axes(rank)?; + dimensions[row] = amp_micro_dimension(precision); + dimensions[column] = u32::from(order.block_shape[1]); + } + } + Ok(dimensions) +} + +fn affine_geometry( + source_order: StorageOrder, + destination_order: StorageOrder, + precision: Precision, + source_storage: &TensorRegion, + destination_storage: &TensorRegion, + region: &TensorRegion, + dimensions: &[u32], +) -> Result { + let rank = region.len(); + if rank == 0 + || source_storage.len() != rank + || destination_storage.len() != rank + || dimensions.len() != rank + { + return Err(ConversionGeometryError::Unsupported); + } + for (axis, (((source, destination), region), &block)) in source_storage + .iter() + .zip(destination_storage.iter()) + .zip(region.iter()) + .zip(dimensions) + .enumerate() + { + let source_must_align = !(source_order == StorageOrder::Linear && axis + 1 == rank); + let destination_must_align = + !(destination_order == StorageOrder::Linear && axis + 1 == rank); + if source.axis != region.axis + || destination.axis != region.axis + || region.start < source.start + || region.start < destination.start + || region.logical_end > source.physical_end + || region.logical_end > destination.physical_end + || block == 0 + || source_must_align && !(region.start - source.start).is_multiple_of(block) + || destination_must_align && !(region.start - destination.start).is_multiple_of(block) + || !(region.logical_end - region.start).is_multiple_of(block) + { + return Err(ConversionGeometryError::Unsupported); + } + } + let block_bytes = dimensions + .iter() + .try_fold(1u32, |elements, dimension| elements.checked_mul(*dimension)) + .and_then(|elements| elements.checked_mul(precision.bytes() as u32)) + .ok_or(ConversionGeometryError::Overflow)?; + let starts = region.iter().map(|extent| extent.start).collect::>(); + let source_offset = physical_byte_offset(source_order, precision, source_storage, &starts)?; + let destination_offset = + physical_byte_offset(destination_order, precision, destination_storage, &starts)?; + let mut copy_dimensions = Vec::new(); + let mut axis_strides = vec![None; rank]; + for axis in 0..rank { + let count = (region[axis].logical_end - region[axis].start) / dimensions[axis]; + if count <= 1 { + continue; + } + let mut next = starts.clone(); + next[axis] += dimensions[axis]; + let source_stride = physical_byte_offset(source_order, precision, source_storage, &next)? + .checked_sub(source_offset) + .ok_or(ConversionGeometryError::Unsupported)?; + let destination_stride = + physical_byte_offset(destination_order, precision, destination_storage, &next)? + .checked_sub(destination_offset) + .ok_or(ConversionGeometryError::Unsupported)?; + let mut last = starts.clone(); + last[axis] += (count - 1) * dimensions[axis]; + let source_end = source_offset + .checked_add((count - 1) * source_stride) + .ok_or(ConversionGeometryError::Overflow)?; + let destination_end = destination_offset + .checked_add((count - 1) * destination_stride) + .ok_or(ConversionGeometryError::Overflow)?; + if physical_byte_offset(source_order, precision, source_storage, &last)? != source_end + || physical_byte_offset(destination_order, precision, destination_storage, &last)? + != destination_end + { + return Err(ConversionGeometryError::Unsupported); + } + axis_strides[axis] = Some((source_stride, destination_stride)); + copy_dimensions.push(CopyDimension { + count, + source_stride, + destination_stride, + }); + } + let mut last = starts; + let mut expected_source = source_offset; + let mut expected_destination = destination_offset; + for (axis, &dimension) in dimensions.iter().enumerate() { + let count = (region[axis].logical_end - region[axis].start) / dimension; + last[axis] += (count - 1) * dimension; + if let Some((source_stride, destination_stride)) = axis_strides[axis] { + expected_source = expected_source + .checked_add((count - 1) * source_stride) + .ok_or(ConversionGeometryError::Overflow)?; + expected_destination = expected_destination + .checked_add((count - 1) * destination_stride) + .ok_or(ConversionGeometryError::Overflow)?; + } + } + if physical_byte_offset(source_order, precision, source_storage, &last)? != expected_source + || physical_byte_offset(destination_order, precision, destination_storage, &last)? + != expected_destination + { + return Err(ConversionGeometryError::Unsupported); + } + copy_dimensions + .sort_by_key(|dimension| dimension.source_stride.max(dimension.destination_stride)); + Ok(CopyGeometry { + source_offset, + destination_offset, + contiguous_bytes: block_bytes, + dimensions: copy_dimensions, + }) +} + +fn physical_block_dimensions( + order: StorageOrder, + precision: Precision, + region: &TensorRegion, +) -> Result, ConversionGeometryError> { + let rank = region.len(); + let mut dimensions = vec![1u32; rank]; + if rank == 1 { + dimensions[0] = region[0].logical_end - region[0].start; + } else { + let row = rank - 2; + let column = rank - 1; + match order { + StorageOrder::Linear => { + dimensions[column] = region[column].logical_end - region[column].start; + } + StorageOrder::Native(NativeKernelOrder::Left) => { + dimensions[column] = amp_micro_dimension(precision); + } + StorageOrder::Native(NativeKernelOrder::Output) => { + dimensions[column] = AMP_COLUMN_MICRO; + } + StorageOrder::Native(NativeKernelOrder::TransposedLeft) => { + dimensions[row] = amp_micro_dimension(precision); + } + StorageOrder::Native(NativeKernelOrder::TransposedOutput) => { + dimensions[row] = AMP_COLUMN_MICRO; + } + StorageOrder::Native(NativeKernelOrder::TransposedRight) => { + dimensions[row] = amp_micro_dimension(precision); + dimensions[column] = AMP_COLUMN_MICRO; + } + StorageOrder::Blocked(order) => { + let [row, column] = order.physical_axes(rank)?; + dimensions[row] = amp_micro_dimension(precision); + dimensions[column] = u32::from(order.block_shape[1]); + } + } + } + Ok(dimensions) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::layout::{ + AxisTiling, BlockedOrder, MemoryClass, Padding, ShardExtent, TensorAxis, TensorTiling, + }; + use std::collections::BTreeMap; + + #[test] + fn randomized_affine_geometries_preserve_logical_elements() { + let mut random = fastrand::Rng::with_seed(0x636f_7079_6765_6f6d); + for case in 0..128 { + let rank = random.usize(2..=4); + let mut shape = (0..rank).map(|_| random.u32(1..=3)).collect::>(); + shape[rank - 2] = 16 * random.u32(1..=3); + shape[rank - 1] = 16 * random.u32(1..=3); + let first_axis = random.usize(..rank); + let second_axis = (first_axis + random.usize(1..rank)) % rank; + shape[first_axis] = 16 * random.u32(1..=3); + shape[second_axis] = 16 * random.u32(1..=3); + let arbitrary_blocked = StorageOrder::Blocked(BlockedOrder { + axes: [ + TensorAxis::FromStart(first_axis as u16), + TensorAxis::FromStart(second_axis as u16), + ], + block_shape: [16, 16], + permutation: if random.bool() { [0, 1] } else { [1, 0] }, + }); + let orders = [ + StorageOrder::Linear, + StorageOrder::Blocked(BlockedOrder::matrix(16, 16)), + StorageOrder::Blocked(BlockedOrder::transposed_matrix(16, 16)), + arbitrary_blocked, + StorageOrder::Native(NativeKernelOrder::Left), + StorageOrder::Native(NativeKernelOrder::TransposedLeft), + StorageOrder::Native(NativeKernelOrder::TransposedRight), + StorageOrder::Native(NativeKernelOrder::Output), + StorageOrder::Native(NativeKernelOrder::TransposedOutput), + ]; + let source_starts = shape + .iter() + .map(|_| 16 * random.u32(0..=1)) + .collect::>(); + let destination_starts = shape + .iter() + .map(|_| 16 * random.u32(0..=1)) + .collect::>(); + let mut storage = |starts: &[u32]| { + shape + .iter() + .zip(starts) + .enumerate() + .map(|(axis, (&width, &start))| ShardExtent { + axis: axis as u16, + start: 0, + logical_end: start + width + 16 * random.u32(0..=1), + physical_end: start + width + 16, + }) + .collect::>() + .into() + }; + let region = |starts: &[u32]| { + shape + .iter() + .zip(starts) + .enumerate() + .map(|(axis, (&width, &start))| ShardExtent { + axis: axis as u16, + start, + logical_end: start + width, + physical_end: start + width, + }) + .collect::>() + .into() + }; + let source_storage = storage(&source_starts); + let destination_storage = storage(&destination_starts); + let source_region = region(&source_starts); + let destination_region = region(&destination_starts); + let precision = if random.bool() { + Precision::F16 + } else { + Precision::F32 + }; + let source_order = orders[random.usize(..orders.len())]; + let destination_order = orders[random.usize(..orders.len())]; + let geometries = copy_geometries_between( + precision, + (source_order, &source_storage, &source_region), + (destination_order, &destination_storage, &destination_region), + true, + true, + ) + .unwrap_or_else(|error| { + panic!("case {case}: {source_order:?} -> {destination_order:?}: {error}") + }); + let element_bytes = precision.bytes() as u32; + let mut copied = BTreeMap::new(); + for geometry in geometries { + for (source, destination) in geometry.offsets().unwrap() { + for byte in (0..geometry.contiguous_bytes).step_by(element_bytes as usize) { + assert_eq!( + copied.insert(source + byte, destination + byte), + None, + "case {case}: duplicate source element" + ); + } + } + } + let mut coordinates = vec![0; rank]; + loop { + let source_coordinates = coordinates + .iter() + .zip(&source_starts) + .map(|(coordinate, start)| coordinate + start) + .collect::>(); + let destination_coordinates = coordinates + .iter() + .zip(&destination_starts) + .map(|(coordinate, start)| coordinate + start) + .collect::>(); + let source = physical_byte_offset( + source_order, + precision, + &source_storage, + &source_coordinates, + ) + .unwrap(); + let destination = physical_byte_offset( + destination_order, + precision, + &destination_storage, + &destination_coordinates, + ) + .unwrap(); + assert_eq!( + copied.remove(&source), + Some(destination), + "case {case}: {shape:?} {precision:?} {source_order:?} -> {destination_order:?} at {coordinates:?}" + ); + let mut axis = rank; + loop { + if axis == 0 { + assert!(copied.is_empty(), "case {case}: extra copied elements"); + break; + } + axis -= 1; + coordinates[axis] += 1; + if coordinates[axis] < shape[axis] { + break; + } + coordinates[axis] = 0; + } + if axis == 0 && coordinates.iter().all(|&coordinate| coordinate == 0) { + break; + } + } + } + } + + #[test] + fn randomized_native_results_can_be_linearized_directly() { + let mut random = fastrand::Rng::with_seed(0x7374_6167_6564_6c69); + for case in 0..64 { + let row_partitions = random.u16(1..=8); + let column_partitions = random.u16(1..=8); + let rows = 16 * u32::from(row_partitions) * random.u32(1..=4) - random.u32(0..16); + let columns = 2 * u32::from(column_partitions) * random.u32(1..=8) - random.u32(0..2); + let tile_count = row_partitions * column_partitions; + let source = Layout { + order: StorageOrder::Native(NativeKernelOrder::TransposedLeft), + tiling: TensorTiling { + tile_count, + replicas: 1, + axes: vec![ + AxisTiling::new(TensorAxis::FromEnd(2), row_partitions, 16, Padding::Zero) + .with_tile_stride(1), + AxisTiling::new( + TensorAxis::FromEnd(1), + column_partitions, + 2, + Padding::Zero, + ) + .with_tile_stride(row_partitions), + ], + }, + memory_class: MemoryClass::Interleaved, + }; + let destination = Layout::logical_linear(tile_count, 4); + let shape = if random.bool() { + TensorShape::new([rows, columns]) + } else { + TensorShape::new([1, rows, columns]) + }; + if source.resolve(&shape).is_err() || destination.resolve(&shape).is_err() { + continue; + } + let strategy = layout_conversion_strategy(Precision::F16, &source, &destination); + assert_eq!( + strategy, + ConversionStrategy::DirectLogical, + "random case {case}" + ); + plan_conversion(&shape, Precision::F16, &source, &destination, strategy) + .unwrap_or_else(|error| panic!("random case {case}: {error}")); + } + } +} diff --git a/crates/ipu-codegen/src/cost/exchange.rs b/crates/ipu-codegen/src/cost/exchange.rs new file mode 100644 index 0000000..96064a8 --- /dev/null +++ b/crates/ipu-codegen/src/cost/exchange.rs @@ -0,0 +1,149 @@ +//! Exchange traffic accounting and conversion to target costs. + +use crate::metrics::{CostEstimate, ExchangeFootprint}; +use ipu_target::hardware::HardwareTarget; + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub(super) struct ExchangeEndpointLoad { + pub bytes: u64, + pub fragments: u64, +} + +impl ExchangeEndpointLoad { + fn add(&mut self, bytes: u64, fragments: u64) { + self.bytes = self.bytes.saturating_add(bytes); + self.fragments = self.fragments.saturating_add(fragments); + } +} + +/// Resource-indexed work for transfers sharing an exchange phase. Sends from +/// an adjacent tile pair occupy one bus; receives are independent per tile. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub(super) struct ExchangeEndpointTraffic { + pub outgoing_buses: Vec, + pub incoming_tiles: Vec, +} + +impl ExchangeEndpointTraffic { + pub(super) fn from_maxima( + outgoing_bytes: u64, + incoming_bytes: u64, + outgoing_fragments: u64, + incoming_fragments: u64, + ) -> Self { + let mut traffic = Self::default(); + traffic.add_outgoing(0, outgoing_bytes, outgoing_fragments); + traffic.add_incoming(0, incoming_bytes, incoming_fragments); + traffic + } + + pub(super) fn add_outgoing(&mut self, bus: u16, bytes: u64, fragments: u64) { + add_endpoint_load(&mut self.outgoing_buses, bus, bytes, fragments); + } + + pub(super) fn add_incoming(&mut self, tile: u16, bytes: u64, fragments: u64) { + add_endpoint_load(&mut self.incoming_tiles, tile, bytes, fragments); + } + + pub(super) fn merge(&mut self, other: &Self) { + for (bus, load) in other.outgoing_buses.iter().copied().enumerate() { + self.add_outgoing(bus as u16, load.bytes, load.fragments); + } + for (tile, load) in other.incoming_tiles.iter().copied().enumerate() { + self.add_incoming(tile as u16, load.bytes, load.fragments); + } + } + + pub(super) fn maximum_outgoing_bytes(&self) -> u64 { + endpoint_maxima(&self.outgoing_buses).0 + } + + pub(super) fn maximum_incoming_bytes(&self) -> u64 { + endpoint_maxima(&self.incoming_tiles).0 + } + + pub(super) fn maximum_payload_bytes(&self) -> u64 { + self.maximum_outgoing_bytes() + .max(self.maximum_incoming_bytes()) + } + + pub(super) fn maximum_outgoing_fragments(&self) -> u64 { + endpoint_maxima(&self.outgoing_buses).1 + } + + pub(super) fn maximum_incoming_fragments(&self) -> u64 { + endpoint_maxima(&self.incoming_tiles).1 + } + + pub(super) fn maximum_fragments(&self) -> u64 { + self.maximum_outgoing_fragments() + .max(self.maximum_incoming_fragments()) + } + + pub(super) fn is_empty(&self) -> bool { + self.maximum_payload_bytes() == 0 + } +} + +fn endpoint_maxima(loads: &[ExchangeEndpointLoad]) -> (u64, u64) { + loads.iter().fold((0, 0), |(bytes, fragments), load| { + (bytes.max(load.bytes), fragments.max(load.fragments)) + }) +} + +fn add_endpoint_load( + loads: &mut Vec, + endpoint: u16, + bytes: u64, + fragments: u64, +) { + if bytes == 0 && fragments == 0 { + return; + } + loads.resize( + loads.len().max(usize::from(endpoint).saturating_add(1)), + ExchangeEndpointLoad::default(), + ); + loads[usize::from(endpoint)].add(bytes, fragments); +} + +pub(super) fn cost( + traffic: &ExchangeEndpointTraffic, + phases: u64, + target: HardwareTarget, +) -> CostEstimate { + if traffic.is_empty() || phases == 0 { + return CostEstimate::default(); + } + let costs = target.costs(); + let payload_cycles = traffic + .maximum_payload_bytes() + .div_ceil(costs.exchange_bytes_per_cycle); + let cutover_cycles = traffic + .maximum_fragments() + .saturating_mul(costs.logical_fragment_cycles); + let cycles = payload_cycles + .max(cutover_cycles) + .saturating_add(phases.saturating_mul(costs.exchange_phase_cycles)); + CostEstimate { + cycles, + exchange_cycles: cycles, + exchange_footprint: footprint(traffic, phases, target), + } +} + +fn footprint( + traffic: &ExchangeEndpointTraffic, + phases: u64, + target: HardwareTarget, +) -> ExchangeFootprint { + let transfer_bytes = u64::from(target.exchange().maximum_transfer_words) * 4; + ExchangeFootprint { + phases, + maximum_transfer_chunks_per_tile: traffic + .maximum_payload_bytes() + .div_ceil(transfer_bytes) + .max(traffic.maximum_fragments()) + .max(phases), + } +} * Unmerged path crates/ipu-codegen/src/cost/kernel.rs diff --git a/crates/ipu-codegen/src/cost/memoization.rs b/crates/ipu-codegen/src/cost/memoization.rs new file mode 100644 index 0000000..bbef2a7 --- /dev/null +++ b/crates/ipu-codegen/src/cost/memoization.rs @@ -0,0 +1,162 @@ +//! Thread-safe memoization for repeated conversion estimates. + +use super::kernel::CostModel; +use crate::OperatorSchedule; +use crate::conversion::{ConversionStrategy, DeferredTransform, layout_conversion_strategy}; +use crate::graph::TensorShape; +use crate::layout::{Layout, TensorType}; +use crate::metrics::{CostEstimate, ExchangeFootprint}; +use crate::operator::Precision; +use foldhash::fast::FixedState; +use ipu_target::hardware::HardwareTarget; +use std::collections::HashMap; +use std::sync::Mutex; + +pub(crate) struct MemoizedCostModel<'a, C> { + inner: &'a C, + spatial_capacity: u16, + rearrangements: Mutex, +} + +type RearrangementKey = (TensorType, TensorType, ConversionStrategy); +type RearrangementCache = HashMap; + +impl<'a, C> MemoizedCostModel<'a, C> { + pub(crate) fn new(inner: &'a C, spatial_capacity: u16) -> Self { + Self { + inner, + spatial_capacity, + rearrangements: Mutex::new(HashMap::default()), + } + } + + fn rearrangement( + &self, + source: TensorType, + destination: TensorType, + strategy: ConversionStrategy, + ) -> CostEstimate + where + C: CostModel, + { + let active_tiles = source + .format + .layout + .tiling + .tile_count + .max(destination.format.layout.tiling.tile_count); + let key = (source.clone(), destination.clone(), strategy); + let mut cache = self.rearrangements.lock().unwrap(); + if let Some(&cost) = cache.get(&key) { + return cost; + } + let mut cost = self + .inner + .rearrangement_cost(&source, &destination, strategy); + cost.cycles = cost + .cycles + .saturating_mul(u64::from(self.spatial_capacity)) + .div_ceil(u64::from(active_tiles)); + cache.insert(key, cost); + cost + } +} + +impl CostModel for MemoizedCostModel<'_, C> { + fn target(&self) -> HardwareTarget { + self.inner.target() + } + + fn operator_cycles( + &self, + schedule: &OperatorSchedule, + inputs: &[TensorType], + output: &TensorType, + ) -> u64 { + self.inner.operator_cycles(schedule, inputs, output) + } + + fn cast_cycles(&self, input: &TensorType, to: Precision) -> u64 { + self.inner.cast_cycles(input, to) + } + + fn layout_conversion_cost( + &self, + shape: &TensorShape, + precision: Precision, + from: &Layout, + to: &Layout, + ) -> CostEstimate { + let requested = layout_conversion_strategy(precision, from, to); + let source = TensorType::new(shape.0.clone(), precision, from.clone()); + let destination = TensorType::new(shape.0.clone(), precision, to.clone()); + self.rearrangement(source, destination, requested) + } + + fn operator_exchange_cycles( + &self, + schedule: &OperatorSchedule, + inputs: &[TensorType], + output: &TensorType, + ) -> u64 { + self.inner + .operator_exchange_cycles(schedule, inputs, output) + } + + fn operator_exchange_footprint( + &self, + schedule: &OperatorSchedule, + inputs: &[TensorType], + output: &TensorType, + ) -> ExchangeFootprint { + self.inner + .operator_exchange_footprint(schedule, inputs, output) + } + + fn deferred_input_cycles( + &self, + transform: DeferredTransform, + source: &TensorType, + logical_output: &TensorType, + consumer_input: &TensorType, + consumer_dispatch: &OperatorSchedule, + producer_cycles: u64, + ) -> u64 { + self.inner.deferred_input_cycles( + transform, + source, + logical_output, + consumer_input, + consumer_dispatch, + producer_cycles, + ) + } + + fn deferred_input_exchange_cycles( + &self, + transform: DeferredTransform, + source: &TensorType, + logical_output: &TensorType, + consumer_input: &TensorType, + consumer_dispatch: &OperatorSchedule, + producer_cycles: u64, + ) -> u64 { + self.inner.deferred_input_exchange_cycles( + transform, + source, + logical_output, + consumer_input, + consumer_dispatch, + producer_cycles, + ) + } + + fn rearrangement_cost( + &self, + source: &TensorType, + destination: &TensorType, + strategy: ConversionStrategy, + ) -> CostEstimate { + self.rearrangement(source.clone(), destination.clone(), strategy) + } +} diff --git a/crates/ipu-codegen/src/cost/mod.rs b/crates/ipu-codegen/src/cost/mod.rs new file mode 100644 index 0000000..a7669b9 --- /dev/null +++ b/crates/ipu-codegen/src/cost/mod.rs @@ -0,0 +1,18 @@ +//! Resource estimation and target-specific planning costs. + +mod exchange; +mod kernel; +mod memoization; +mod parallel_reduction; +mod resources; + +pub(crate) use kernel::row_major_pack_cycles; +pub use kernel::{CostModel, Ipu21CostModel}; +pub(crate) use memoization::MemoizedCostModel; +pub(crate) use parallel_reduction::parallel_reduction_preselection_metrics; +pub(crate) use resources::{ + conversion_memory_estimate, operator_memory_estimate, region_peak_memory, + region_peak_memory_with_multiplicity, +}; +#[cfg(test)] +pub(crate) use resources::{maximum_shard_bytes, physical_elements}; diff --git a/crates/ipu-codegen/src/cost/parallel_reduction.rs b/crates/ipu-codegen/src/cost/parallel_reduction.rs new file mode 100644 index 0000000..de1a1fe --- /dev/null +++ b/crates/ipu-codegen/src/cost/parallel_reduction.rs @@ -0,0 +1,84 @@ +//! Cheap parallel-reduction grid filtering before concrete layout expansion. + +use crate::layout::{AMP_COLUMN_MICRO, TensorType}; +use crate::metrics::{CostEstimate, MemoryPeaks, RegionMetrics}; +use crate::operator::{GemmBlockShape, GemmGrid, GemmOrientation, Precision}; +use ipu_target::hardware::HardwareTarget; + +pub(crate) fn parallel_reduction_preselection_metrics( + target: HardwareTarget, + block: GemmBlockShape, + grid: GemmGrid, + orientation: GemmOrientation, + inputs: &[TensorType], + output_precision: Precision, +) -> Option { + let [left, right] = inputs else { return None }; + if left.shape.0.len() < 2 || right.shape.0.len() < 2 { + return None; + } + let outer_rows = left.shape.0[..left.shape.0.len() - 2] + .iter() + .fold(1u64, |product, &extent| { + product.saturating_mul(u64::from(extent)) + }); + let [physical_left, physical_right] = orientation.physical_order([left, right]); + let logical_rows = physical_left.shape.0[orientation + .row_axis() + .resolve(physical_left.shape.0.len()) + .ok()?]; + let local_rows = logical_rows.div_ceil(u32::from(grid.rows)); + let local_columns = block.output_columns.div_ceil(AMP_COLUMN_MICRO); + let local_inner = block.inner.div_ceil(AMP_COLUMN_MICRO); + let costs = target.costs(); + let compute = u64::from(local_columns) + .saturating_mul(u64::from(local_inner)) + .saturating_mul( + outer_rows + .saturating_mul(u64::from(local_rows)) + .saturating_mul(4) + .saturating_add(costs.amp_grid_search_setup_cycles), + ); + let communication = u64::from(local_columns) + .saturating_mul(u64::from(local_inner)) + .saturating_add(u64::from(local_rows).saturating_mul(u64::from(local_inner))) + .saturating_add( + u64::from(local_rows) + .saturating_mul(u64::from(local_columns)) + .saturating_mul(u64::from(grid.inner.saturating_sub(1))), + ); + let left_precision = physical_left.format.precision; + let right_precision = physical_right.format.precision; + let left_bytes = outer_rows + .saturating_mul(u64::from(local_rows)) + .saturating_mul(u64::from(local_inner)) + .saturating_mul(u64::from(AMP_COLUMN_MICRO)) + .saturating_mul(left_precision.bytes()); + let right_bytes = u64::from(local_columns) + .saturating_mul(u64::from(AMP_COLUMN_MICRO)) + .saturating_mul(u64::from(local_inner)) + .saturating_mul(u64::from(AMP_COLUMN_MICRO)) + .saturating_mul(right_precision.bytes()); + let partial_bytes = outer_rows + .saturating_mul(u64::from(local_rows)) + .saturating_mul(u64::from(local_columns)) + .saturating_mul(u64::from(AMP_COLUMN_MICRO)) + .saturating_mul(output_precision.bytes()); + let compute_interleaved = left_bytes + .saturating_add(right_bytes) + .saturating_add(partial_bytes); + let reduction_standard = partial_bytes.saturating_mul(3); + Some(RegionMetrics { + cost: CostEstimate { + cycles: compute.saturating_add(communication), + exchange_cycles: communication, + ..CostEstimate::default() + }, + memory: MemoryPeaks { + standard: reduction_standard, + interleaved: compute_interleaved, + total: compute_interleaved.max(reduction_standard.saturating_add(partial_bytes)), + ..MemoryPeaks::default() + }, + }) +} diff --git a/crates/ipu-codegen/src/cost/resources.rs b/crates/ipu-codegen/src/cost/resources.rs new file mode 100644 index 0000000..b0a4806 --- /dev/null +++ b/crates/ipu-codegen/src/cost/resources.rs @@ -0,0 +1,1312 @@ +//! Structural memory, capacity, and communication estimates. + +use super::exchange::ExchangeEndpointTraffic; +use crate::graph::TensorShape; +use crate::layout::{ + AMP_COLUMN_MICRO, AMP_INNER_BLOCK, Layout, MemoryClass, StorageOrder, TensorRegion, TensorType, +}; +use crate::metrics::{MemoryEstimate, MemoryPeaks, MemoryUsage}; +use crate::mid::{MidOperation, MidOperationKind, MidValue, MidValueId}; +use crate::operator::{ + AllocationRequirements, MemoryElementRequirement, MemoryOperand, OperandMaterialization, + Precision, +}; +use crate::{OperatorSchedule, ScheduleStep}; +use ipu_target::hardware::HardwareTarget; +use std::collections::{BTreeMap, BTreeSet, HashMap}; +use std::sync::Arc; + +fn layout_extents(shape: &TensorShape, layout: &Layout) -> Option> { + Some( + layout + .resolve(shape) + .ok()? + .shard_extents() + .into_iter() + .map(|shard| (shard.tile, shard.extents.logical())) + .collect(), + ) +} + +pub(crate) fn physical_elements(shape: &TensorShape, layout: &Layout) -> u64 { + layout.resolve(shape).map_or_else( + |_| { + shape + .elements() + .saturating_mul(u64::from(layout.tiling.replicas)) + }, + |resolved| resolved.total_elements(), + ) +} + +pub(crate) fn maximum_shard_bytes(tensor: &TensorType) -> u64 { + let Ok(resolved) = tensor.format.layout.resolve(&tensor.shape) else { + return u64::MAX; + }; + resolved + .maximum_tile_elements() + .saturating_mul(tensor.format.precision.bytes()) +} + +/// Mean physical storage assigned to one active spatial tile. Replicated +/// layouts include each replica in both the total storage and tile count. +pub(crate) fn average_shard_bytes(tensor: &TensorType) -> u64 { + tensor + .format + .layout + .resolve(&tensor.shape) + .map_or(u64::MAX, |resolved| { + resolved + .total_elements() + .saturating_mul(tensor.format.precision.bytes()) + .div_ceil(u64::from(resolved.tile_count()).max(1)) + }) +} + +pub(crate) fn maximum_axis_shard_extent(tensor: &TensorType, axis: usize) -> u64 { + tensor + .format + .layout + .resolve(&tensor.shape) + .ok() + .and_then(|resolved| resolved.maximum_axis_extent(axis)) + .map_or(u64::MAX, u64::from) +} + +pub(crate) fn gemm_partial_tensor(schedule: &OperatorSchedule, output: &TensorType) -> TensorType { + let Some(ScheduleStep::Gemm(plan)) = schedule.steps.first() else { + return output.clone(); + }; + plan.partial_tensor(output) + .unwrap_or_else(|| output.clone()) +} + +pub(crate) fn tensor_memory(tensor: &TensorType) -> MemoryUsage { + let mut usage = MemoryUsage::default(); + usage.add_class( + tensor.format.layout.memory_class, + maximum_shard_bytes(tensor), + ); + usage +} + +/// Interleaved working-set lower bound used to discard impossible parallel +/// GEMM grids before their concrete layouts and staging policies are expanded. +/// Exact candidates are checked again by [`operator_memory_estimate`]. +fn allocation_memory(tensor: &TensorType, requirement: AllocationRequirements) -> MemoryUsage { + let mut bytes = + maximum_shard_bytes(tensor).saturating_add(u64::from(requirement.access_tail_bytes)); + if requirement.memory_element == MemoryElementRequirement::Distinct { + let element = match tensor.format.layout.memory_class { + MemoryClass::Standard => ipu_target::memory::TILE_MEMORY_ELEMENT_SIZE, + MemoryClass::Interleaved => ipu_target::memory::IPU21_INTERLEAVED_ELEMENT_SIZE, + }; + bytes = bytes.div_ceil(u64::from(element)) * u64::from(element); + } + let mut usage = MemoryUsage::default(); + usage.add_class(tensor.format.layout.memory_class, bytes); + usage +} + +fn allocation_requirements( + operations: &[MidOperation], +) -> BTreeMap { + let mut requirements = BTreeMap::::new(); + for operation in operations { + if let Some(plan) = operation.operator_plan() { + for (&id, operand) in operation.inputs.iter().zip(&plan.requirements.inputs) { + let requirement = requirements.entry(id).or_default(); + requirement.merge(operand.allocation); + } + if let Some(&id) = operation.results.first() { + let requirement = requirements.entry(id).or_default(); + requirement.merge(plan.requirements.output.allocation); + } + for operands in &plan.requirements.memory_space.distinct_element_groups { + for operand in operands { + let id = match operand { + MemoryOperand::Output => operation.results.first().copied(), + MemoryOperand::Input(index) => { + operation.inputs.get(usize::from(*index)).copied() + } + }; + if let Some(id) = id { + requirements + .entry(id) + .or_default() + .require_distinct_element(); + } + } + } + } + if operation.conversion().is_some() { + for id in operation + .inputs + .first() + .into_iter() + .chain(operation.results.first()) + { + let requirement = requirements.entry(*id).or_default(); + requirement.alignment = requirement.alignment.max(8); + } + } + } + requirements +} + +fn value_allocation( + id: MidValueId, + values: &[MidValue], + requirements: &BTreeMap, +) -> MemoryUsage { + allocation_memory( + &values[id.index() as usize].tensor_type, + requirements.get(&id).copied().unwrap_or_default(), + ) +} + +fn maximum_standard_allocation( + ids: &BTreeSet, + values: &[MidValue], + requirements: &BTreeMap, +) -> u64 { + ids.iter() + .map(|&id| value_allocation(id, values, requirements).standard) + .max() + .unwrap_or(0) +} + +pub(crate) fn operator_memory_estimate( + schedule: &OperatorSchedule, + inputs: &[TensorType], + output: &TensorType, +) -> MemoryEstimate { + let requirements = &schedule.requirements; + let live = inputs.iter().zip(&requirements.inputs).fold( + tensor_memory(output), + |usage, (input, requirement)| { + if requirement.materialization == OperandMaterialization::DispatchSlices { + usage + } else { + usage.saturating_add(tensor_memory(input)) + } + }, + ); + let mut temporary = MemoryUsage::default(); + let mut maximum_standard_temporary_allocation = 0u64; + let gemm = match schedule.steps.first() { + Some(ScheduleStep::Gemm(gemm)) => Some(gemm), + _ => None, + }; + let reduction_staging = schedule.steps.get(1).and_then(|step| match step { + ScheduleStep::Reduce { staging, .. } => Some(*staging), + _ => None, + }); + let attention = match schedule.steps.as_slice() { + [ScheduleStep::Attention(attention)] => Some(attention), + _ => None, + }; + if let (Some(plan), Some(first), Some(second)) = (gemm, inputs.first(), inputs.get(1)) + && plan.geometry.compute.inner > 1 + { + let compute = plan.geometry.compute; + let orientation = plan.geometry.orientation; + let output_column_block = plan.geometry.block.output_columns; + let [left, right] = orientation.physical_order([first, second]); + let left_requirement = requirements.inputs.get(orientation.physical_left_input()); + let right_rank = right.shape.0.len(); + let right_inner_axis = orientation + .row_axis() + .resolve(right_rank) + .expect("GEMM operand rank validated"); + let right_column_axis = orientation + .column_axis() + .resolve(right_rank) + .expect("GEMM operand rank validated"); + let inner_blocks = right.shape.0[right_inner_axis].div_ceil(AMP_INNER_BLOCK); + let column_blocks = right.shape.0[right_column_axis].div_ceil(output_column_block); + let right_staging = u64::from(inner_blocks.div_ceil(u32::from(compute.inner))) + .saturating_mul(u64::from(AMP_INNER_BLOCK)) + .saturating_mul(u64::from( + column_blocks.div_ceil(u32::from(compute.columns)), + )) + .saturating_mul(u64::from(output_column_block)) + .saturating_mul(right.format.precision.bytes()); + let mut convolution = MemoryUsage::default(); + convolution.add_class(MemoryClass::Interleaved, right_staging); + if left_requirement.is_some_and(|requirement| { + requirement.materialization == OperandMaterialization::DispatchSlices + }) { + let requirement = left_requirement.expect("checked requirement"); + let mut left_staging = maximum_shard_bytes(left) + .saturating_add(u64::from(requirement.allocation.access_tail_bytes)); + let left_must_be_distinct = requirements + .memory_space + .distinct_element_groups + .iter() + .any(|operands| { + operands.contains(&MemoryOperand::Input( + orientation.physical_left_input() as u16 + )) + }); + if left_must_be_distinct { + left_staging = left_staging + .div_ceil(u64::from(ipu_target::memory::TILE_MEMORY_ELEMENT_SIZE)) + .saturating_mul(u64::from(ipu_target::memory::TILE_MEMORY_ELEMENT_SIZE)); + } + convolution.add_class(left.format.layout.memory_class, left_staging); + if left.format.layout.memory_class == MemoryClass::Standard { + maximum_standard_temporary_allocation = + maximum_standard_temporary_allocation.max(left_staging); + } + } + // Compute retains one local partial alongside operand staging. The + // later reduction ping-pongs an accumulator and result while its + // staging policy bounds the simultaneously resident remote partials. + let partial_bytes = maximum_shard_bytes(&gemm_partial_tensor(schedule, output)); + let reduction_partial_bytes = if plan.geometry.result + != (crate::GemmResultGrid { + rows: compute.rows, + columns: compute.columns, + }) { + maximum_shard_bytes(output) + } else { + partial_bytes + }; + convolution.interleaved = convolution.interleaved.saturating_add(partial_bytes); + let staged_remote_partials = match reduction_staging { + Some(crate::ReductionStaging::Complete) => compute.inner.saturating_sub(1), + Some(crate::ReductionStaging::Streamed) => 1, + None => 0, + }; + let reduction = MemoryUsage { + standard: reduction_partial_bytes + .saturating_mul(u64::from(staged_remote_partials).saturating_add(2)), + interleaved: partial_bytes, + }; + temporary = MemoryUsage { + standard: convolution.standard.max(reduction.standard), + interleaved: convolution.interleaved.max(reduction.interleaved), + }; + } + if let (Some(plan), Some(left), Some(requirement)) = + (gemm, inputs.first(), requirements.inputs.first()) + && requirement.materialization == OperandMaterialization::DispatchSlices + && plan.geometry.compute.inner == 1 + { + let inner = left.shape.0.last().copied().map_or(1, u64::from).max(1); + let bytes = maximum_shard_bytes(left) + .div_ceil(inner) + .saturating_mul(u64::from(plan.geometry.block.inner)) + .saturating_add(u64::from(requirement.allocation.access_tail_bytes)); + temporary.add_class(left.format.layout.memory_class, bytes); + } + if let (Some(plan), Some(right), Some(requirement)) = + (gemm, inputs.get(1), requirements.inputs.get(1)) + && right.format.precision == Precision::F16 + && plan.geometry.compute.inner == 1 + && gemm_uses_panel_buffer(schedule, right, output) + { + // Each local output-column panel has one final kernel buffer reused + // across K phases. Remote bytes can be exchanged directly into it. + let output_columns = + maximum_axis_shard_extent(output, output.shape.0.len().saturating_sub(1)); + let panels = output_columns.div_ceil(u64::from(plan.geometry.block.output_columns)); + let bytes = panels + .saturating_mul(u64::from(plan.geometry.block.inner)) + .saturating_mul(u64::from(plan.geometry.block.output_columns)) + .saturating_mul(right.format.precision.bytes()); + temporary.add_class(requirement.local_staging.memory_class(), bytes); + } + if let Some(crate::AttentionMap { + blocking: + crate::AttentionBlocking::Flash { + query_rows: query_block_rows, + key_rows: key_block_rows, + }, + query_dimension: padded_query_dimension, + value_dimension: padded_value_dimension, + .. + }) = attention + { + let element_bytes = inputs.first().map_or(Precision::F16.bytes(), |input| { + input.format.precision.bytes() + }); + let key_rows = inputs + .get(1) + .and_then(|key| key.shape.0.get(key.shape.0.len().saturating_sub(2))) + .copied() + .map_or(1, u64::from); + let blocks = key_rows.div_ceil(u64::from(*key_block_rows).max(1)); + let panels_per_block = u64::from( + padded_query_dimension + .div_ceil(AMP_COLUMN_MICRO) + .saturating_add(padded_value_dimension.div_ceil(AMP_COLUMN_MICRO)), + ); + let query_rows = output + .shape + .0 + .get(output.shape.0.len().saturating_sub(2)) + .copied() + .map_or(1, u64::from); + let query_partitions = query_rows + .div_ceil(u64::from(*query_block_rows).max(1)) + .max(1); + let prepared_panels_per_owner = blocks + .saturating_mul(panels_per_block) + .div_ceil(query_partitions); + let panel_bytes = u64::from(*key_block_rows) + .saturating_mul(u64::from(AMP_COLUMN_MICRO)) + .saturating_mul(element_bytes); + // Every attention tile retains the current K and V panels. Prepared + // panels are spread over query-tile owners; each needs both its + // row-major gather buffer and its packed source until consumption. + temporary.standard = temporary.standard.saturating_add( + u64::from(*key_block_rows) + .saturating_mul(u64::from( + padded_query_dimension.saturating_add(*padded_value_dimension), + )) + .saturating_mul(element_bytes) + .saturating_add( + prepared_panels_per_owner + .saturating_mul(panel_bytes) + .saturating_mul(2), + ) + .saturating_add( + u64::from(*query_block_rows) + .saturating_mul(u64::from(key_block_rows.saturating_add(16))) + .saturating_mul(element_bytes), + ), + ); + temporary.interleaved = temporary.interleaved.saturating_add( + u64::from(*query_block_rows) + .saturating_mul(u64::from((*padded_value_dimension).max(*key_block_rows))) + .saturating_mul(Precision::F32.bytes()), + ); + } + if let Some(crate::AttentionMap { + blocking: + crate::AttentionBlocking::Materialized { + query_rows: query_block_rows, + padded_key_rows, + }, + query_dimension: padded_query_dimension, + value_dimension: padded_value_dimension, + .. + }) = attention + { + let element_bytes = inputs.first().map_or(Precision::F16.bytes(), |input| { + input.format.precision.bytes() + }); + let key_rows = inputs + .get(1) + .and_then(|key| key.shape.0.get(key.shape.0.len().saturating_sub(2))) + .copied() + .map_or(1, u64::from); + let blocks = key_rows.div_ceil(u64::from(AMP_INNER_BLOCK)); + let panels_per_block = u64::from( + padded_query_dimension + .div_ceil(AMP_COLUMN_MICRO) + .saturating_add(padded_value_dimension.div_ceil(AMP_COLUMN_MICRO)), + ); + let query_rows = output + .shape + .0 + .get(output.shape.0.len().saturating_sub(2)) + .copied() + .map_or(1, u64::from); + let query_partitions = query_rows + .div_ceil(u64::from(*query_block_rows).max(1)) + .max(1); + let prepared_panels_per_owner = blocks + .saturating_mul(panels_per_block) + .div_ceil(query_partitions); + let panel_bytes = u64::from(AMP_INNER_BLOCK) + .saturating_mul(u64::from(AMP_COLUMN_MICRO)) + .saturating_mul(element_bytes); + let operand_staging = u64::from(*padded_key_rows) + .saturating_mul(u64::from( + (*padded_query_dimension).max(*padded_value_dimension), + )) + .saturating_mul(element_bytes); + maximum_standard_temporary_allocation = + maximum_standard_temporary_allocation.max(operand_staging); + let probability_state = u64::from(*query_block_rows) + .saturating_mul(u64::from(padded_key_rows + AMP_COLUMN_MICRO)) + .saturating_mul(element_bytes); + temporary.standard = temporary.standard.saturating_add( + operand_staging + .saturating_add( + prepared_panels_per_owner + .saturating_mul(panel_bytes) + .saturating_mul(2), + ) + .saturating_add( + (probability_state > operand_staging) + .then_some(probability_state) + .unwrap_or(0), + ), + ); + temporary.interleaved = temporary.interleaved.saturating_add( + u64::from(*query_block_rows) + .saturating_mul(u64::from((*padded_value_dimension).max(*padded_key_rows))) + .saturating_mul(element_bytes), + ); + } + MemoryEstimate { + live, + temporary, + peak: live.saturating_add(temporary), + maximum_standard_temporary_allocation, + } +} + +pub(crate) fn gemm_uses_panel_buffer( + schedule: &OperatorSchedule, + right: &TensorType, + output: &TensorType, +) -> bool { + let Some(ScheduleStep::Gemm(plan)) = schedule.steps.first() else { + return false; + }; + let inner_block = plan.geometry.block.inner; + let orientation = plan.geometry.orientation; + let rank = right.shape.0.len(); + let output_rank = output.shape.0.len(); + if rank < 2 || output_rank < 2 { + return true; + } + let streamed = right + .format + .layout + .tiling + .axes + .iter() + .any(|axis| axis.axis == orientation.row_axis() && axis.partitions > 1); + if streamed { + return true; + } + if right.format.layout.memory_class == MemoryClass::Interleaved { + return false; + } + let k = right.shape.0[orientation.row_axis().resolve(rank).unwrap()]; + let columns = maximum_axis_shard_extent( + output, + orientation.column_axis().resolve(output_rank).unwrap(), + ); + k > inner_block && columns > 16 +} + +pub(crate) fn gemm_requires_panel_repacking( + schedule: &OperatorSchedule, + right: &TensorType, + output: &TensorType, +) -> bool { + gemm_uses_panel_buffer(schedule, right, output) + && !matches!(right.format.layout.order, StorageOrder::Blocked(_)) +} + +pub(crate) fn gemm_exchange_phase_count( + schedule: &OperatorSchedule, + inputs: &[TensorType], + _output: &TensorType, +) -> u64 { + let Some(ScheduleStep::Gemm(plan)) = schedule.steps.first() else { + return 0; + }; + let inner_block = plan.geometry.block.inner; + let orientation = plan.geometry.orientation; + let Some(left) = inputs.get(orientation.physical_left_input()) else { + return 0; + }; + let Some(&inner) = orientation + .column_axis() + .resolve(left.shape.0.len()) + .ok() + .and_then(|axis| left.shape.0.get(axis)) + else { + return 0; + }; + u64::from(inner).div_ceil(u64::from(inner_block)) +} + +pub(crate) fn conversion_memory_estimate( + input: &TensorType, + output: &TensorType, + strategy: crate::ConversionStrategy, +) -> MemoryEstimate { + let live = tensor_memory(input).saturating_add(tensor_memory(output)); + let mut staging_by_tile = HashMap::::new(); + let mut maximum_standard_temporary_allocation = 0; + if strategy == crate::ConversionStrategy::StageLogicalThenTransform { + for shard in output + .format + .layout + .resolve(&output.shape) + .into_iter() + .flat_map(|layout| layout.shard_extents()) + { + let bytes = shard + .extents + .logical_elements() + .saturating_mul(output.format.precision.bytes()); + let tile_bytes = staging_by_tile.entry(shard.tile).or_default(); + *tile_bytes = tile_bytes.saturating_add(bytes); + maximum_standard_temporary_allocation = + maximum_standard_temporary_allocation.max(bytes); + } + } + let geometry_staging = if strategy.uses_intersections() + && input.format.layout.tiling != output.format.layout.tiling + && input + .format + .layout + .tiling + .tile_count + .max(output.format.layout.tiling.tile_count) + > 1 + && input.format.layout.order != output.format.layout.order + { + let source = maximum_shard_bytes(input); + let destination = maximum_shard_bytes(output); + maximum_standard_temporary_allocation = maximum_standard_temporary_allocation + .max(source) + .max(destination); + source.saturating_add(destination) + } else { + 0 + }; + let temporary = MemoryUsage { + standard: staging_by_tile + .into_values() + .max() + .unwrap_or(0) + .saturating_add(geometry_staging), + interleaved: 0, + }; + MemoryEstimate { + live, + temporary, + peak: live.saturating_add(temporary), + maximum_standard_temporary_allocation, + } +} + +pub(crate) fn region_peak_memory( + initial: &[MidValueId], + operations: &[MidOperation], + outputs: &[MidValueId], + values: &[MidValue], + target: HardwareTarget, +) -> MemoryPeaks { + region_peak_memory_with_multiplicity( + initial, + operations, + outputs, + values, + &BTreeMap::new(), + target, + ) +} + +pub(crate) fn region_peak_memory_with_multiplicity( + initial: &[MidValueId], + operations: &[MidOperation], + outputs: &[MidValueId], + values: &[MidValue], + allocation_multiplicity: &BTreeMap, + target: HardwareTarget, +) -> MemoryPeaks { + let constraints = target.memory_constraints(); + let requirements = allocation_requirements(operations); + let streamed_aliases = operations + .iter() + .filter_map(|operation| { + if operation.conversion()?.1 != OperandMaterialization::DispatchSlices { + return None; + } + Some((*operation.results.first()?, *operation.inputs.first()?)) + }) + .collect::>(); + let mut uses = BTreeMap::::new(); + for input in operations.iter().flat_map(operation_value_inputs) { + *uses.entry(*input).or_default() += 1; + } + for output in outputs { + *uses.entry(*output).or_default() += 1; + } + let mut live_values = BTreeSet::new(); + for id in initial { + live_values.insert(*id); + } + let mut peaks = MemoryPeaks::default(); + let observe = |peaks: &mut MemoryPeaks, ids: &BTreeSet, temporary: MemoryUsage| { + let roots = ids + .iter() + .map(|id| allocation_root(*id, &streamed_aliases)) + .collect::>(); + let live = roots.iter().fold(MemoryUsage::default(), |usage, id| { + let allocation = value_allocation(*id, values, &requirements); + let copies = u64::from(allocation_multiplicity.get(id).copied().unwrap_or(1)); + usage.saturating_add(MemoryUsage { + standard: allocation.standard.saturating_mul(copies), + interleaved: allocation.interleaved.saturating_mul(copies), + }) + }); + peaks.observe( + live.saturating_add(temporary), + maximum_standard_allocation(&roots, values, &requirements), + constraints, + ); + }; + observe(&mut peaks, &live_values, MemoryUsage::default()); + for operation in operations { + let mut during_values = live_values.clone(); + for result in &operation.results { + during_values.insert(*result); + } + let roots = during_values + .iter() + .map(|id| allocation_root(*id, &streamed_aliases)) + .collect::>(); + let live = roots.iter().fold(MemoryUsage::default(), |usage, id| { + let allocation = value_allocation(*id, values, &requirements); + let copies = u64::from(allocation_multiplicity.get(id).copied().unwrap_or(1)); + usage.saturating_add(MemoryUsage { + standard: allocation.standard.saturating_mul(copies), + interleaved: allocation.interleaved.saturating_mul(copies), + }) + }); + peaks.observe( + live.saturating_add(operation.metrics.memory.temporary), + maximum_standard_allocation(&roots, values, &requirements).max( + operation + .metrics + .memory + .maximum_standard_temporary_allocation, + ), + constraints, + ); + for input in operation_value_inputs(operation) { + if let Some(remaining) = uses.get_mut(input) { + *remaining = remaining.saturating_sub(1); + if *remaining == 0 { + live_values.remove(input); + } + } + } + for result in &operation.results { + if uses.get(result).copied().unwrap_or(0) != 0 { + live_values.insert(*result); + } + } + } + observe(&mut peaks, &live_values, MemoryUsage::default()); + let exchange_rows = operations + .iter() + .map(|operation| operation.metrics.cost.exchange_row_bytes(target)) + .fold(0u64, u64::saturating_add); + peaks.exchange_rows = exchange_rows; + peaks.standard = peaks.standard.saturating_add(exchange_rows); + peaks.total = peaks.total.saturating_add(exchange_rows); + peaks +} + +fn operation_value_inputs(operation: &MidOperation) -> Vec<&MidValueId> { + let mut inputs = operation.inputs.iter().collect::>(); + if let MidOperationKind::Repeat(repeat) = &operation.kind { + inputs.extend(repeat.iterated_inputs.iter().flatten()); + } + inputs +} + +fn allocation_root(mut id: MidValueId, aliases: &BTreeMap) -> MidValueId { + while let Some(source) = aliases.get(&id) { + id = *source; + } + id +} + +pub(crate) fn gemm_exchange_endpoint_traffic( + schedule: &OperatorSchedule, + inputs: &[TensorType], + compute_output: &TensorType, + target: HardwareTarget, +) -> Option { + let Some(ScheduleStep::Gemm(plan)) = schedule.steps.first() else { + return Some(ExchangeEndpointTraffic::default()); + }; + let [first, second] = inputs else { + return None; + }; + let [left, right] = plan.geometry.orientation.physical_order([first, second]); + if plan.geometry.compute.inner > 1 { + // The parallel schedule grid contains a K axis which is deliberately + // absent from `compute_output`: every K group produces a partial with + // the same logical output extent. Consequently, matching operand and + // partial-output tile numbers cannot determine locality. The physical + // left operand is replicated across output-column groups and the + // physical right operand across output-row groups; only a shortfall in + // those explicit replica counts creates operator-internal traffic. + return Some(parallel_gemm_operand_traffic( + left, + plan.geometry.compute.columns, + right, + plan.geometry.compute.rows, + )); + } + let orientation = plan.geometry.orientation; + let left_rank = left.shape.0.len(); + let right_rank = right.shape.0.len(); + let output_rank = compute_output.shape.0.len(); + if left_rank < 2 || right_rank < 2 || output_rank < 2 { + return None; + } + let left_row_axis = orientation.row_axis().resolve(left_rank).ok()?; + let left_inner_axis = orientation.column_axis().resolve(left_rank).ok()?; + let right_inner_axis = orientation.row_axis().resolve(right_rank).ok()?; + let right_column_axis = orientation.column_axis().resolve(right_rank).ok()?; + let output_row_axis = orientation.row_axis().resolve(output_rank).ok()?; + let output_column_axis = orientation.column_axis().resolve(output_rank).ok()?; + let output_plans = tile_axis_plans(compute_output)?; + let left_plan = GemmOperandTrafficPlan::new( + left, + compute_output, + left_row_axis, + left_inner_axis, + output_row_axis, + )?; + let right_plan = GemmOperandTrafficPlan::new( + right, + compute_output, + right_column_axis, + right_inner_axis, + output_column_axis, + )?; + let transfer_bytes = u64::from(target.exchange().maximum_transfer_words) * 4; + let mut traffic = ExchangeEndpointTraffic::default(); + let mut left_is_remote = false; + let mut right_is_remote = false; + for tile in 0..compute_output.format.layout.tiling.tile_count { + let left_remote = left_plan.remote_bytes(tile, &output_plans, None)?; + let right_remote = right_plan.remote_bytes(tile, &output_plans, None)?; + left_is_remote |= left_remote != 0; + right_is_remote |= right_remote != 0; + let incoming = left_remote.saturating_add(right_remote); + traffic.add_incoming(tile, incoming, incoming.div_ceil(transfer_bytes)); + } + add_operand_outgoing_bus_work(&mut traffic, left, left_is_remote, transfer_bytes); + add_operand_outgoing_bus_work(&mut traffic, right, right_is_remote, transfer_bytes); + Some(traffic) +} + +fn parallel_gemm_operand_traffic( + left: &TensorType, + left_required_replicas: u16, + right: &TensorType, + right_required_replicas: u16, +) -> ExchangeEndpointTraffic { + let mut traffic = ExchangeEndpointTraffic::default(); + for (operand, required_replicas) in [ + (left, left_required_replicas), + (right, right_required_replicas), + ] { + if operand.format.layout.tiling.replicas >= required_replicas { + continue; + } + let stored_replicas = operand.format.layout.tiling.replicas.max(1); + let base_tiles = operand + .format + .layout + .tiling + .tile_count + .checked_div(stored_replicas) + .unwrap_or(0); + let Some(tile_count) = base_tiles.checked_mul(required_replicas) else { + return ExchangeEndpointTraffic::from_maxima( + u64::MAX / 16, + u64::MAX / 16, + u64::MAX / 16, + u64::MAX / 16, + ); + }; + let mut consumer_layout = operand.format.layout.clone(); + consumer_layout.tiling.tile_count = tile_count; + consumer_layout.tiling.replicas = required_replicas; + let Some(replication) = replica_shortfall_traffic(operand, &consumer_layout) else { + return ExchangeEndpointTraffic::from_maxima( + u64::MAX / 16, + u64::MAX / 16, + u64::MAX / 16, + u64::MAX / 16, + ); + }; + traffic.merge(&replication); + } + traffic +} + +/// Replicating an otherwise unchanged layout only matches identical shard +/// extents. Account for those roles directly instead of running the general +/// all-pairs layout-intersection algorithm for every GEMM candidate. +fn replica_shortfall_traffic( + operand: &TensorType, + consumer_layout: &Layout, +) -> Option { + let sources = layout_extents(&operand.shape, &operand.format.layout)?; + let destinations = layout_extents(&operand.shape, consumer_layout)?; + let mut source_groups = HashMap::>::new(); + for (tile, extents) in sources { + source_groups.entry(extents).or_default().push(tile); + } + let mut destination_groups = HashMap::>::new(); + for (tile, extents) in destinations { + destination_groups.entry(extents).or_default().push(tile); + } + + let mut traffic = ExchangeEndpointTraffic::default(); + for (extents, destination_tiles) in destination_groups { + let source_tiles = source_groups.get(&extents)?; + let remote_destinations = destination_tiles + .into_iter() + .filter(|tile| source_tiles.binary_search(tile).is_err()) + .collect::>(); + if remote_destinations.is_empty() { + continue; + } + let bytes = extents + .logical_elements() + .saturating_mul(operand.format.precision.bytes()); + traffic.add_outgoing(source_tiles[0] / 2, bytes, 1); + for tile in remote_destinations { + traffic.add_incoming(tile, bytes, 1); + } + } + Some(traffic) +} + +struct GemmOperandTrafficPlan<'a> { + operand: &'a TensorType, + operand_plans: Vec, + operand_spatial_axis: usize, + operand_inner_axis: usize, + output_spatial_axis: usize, + rank_offset: usize, +} + +impl<'a> GemmOperandTrafficPlan<'a> { + fn new( + operand: &'a TensorType, + output: &TensorType, + operand_spatial_axis: usize, + operand_inner_axis: usize, + output_spatial_axis: usize, + ) -> Option { + let rank_offset = output.shape.0.len().checked_sub(operand.shape.0.len())?; + let operand_plans = tile_axis_plans(operand)?; + Some(Self { + operand, + operand_plans, + operand_spatial_axis, + operand_inner_axis, + output_spatial_axis, + rank_offset, + }) + } + + fn remote_bytes( + &self, + tile: u16, + output_plans: &[TileAxisPlan], + inner_plan: Option<&TileAxisPlan>, + ) -> Option { + let mut required_elements = 1u64; + let mut local_elements = 1u64; + for axis in 0..self.operand.shape.0.len() { + let extent = self.operand.shape.0[axis]; + let required = if axis == self.operand_spatial_axis { + clipped_range(output_plans[self.output_spatial_axis].range(tile), extent) + } else if axis == self.operand_inner_axis { + inner_plan.map_or(0..extent, |plan| clipped_range(plan.range(tile), extent)) + } else if extent == 1 { + 0..1 + } else { + let output_axis = axis.checked_add(self.rank_offset)?; + clipped_range(output_plans[output_axis].range(tile), extent) + }; + required_elements = required_elements + .saturating_mul(u64::from(required.end.saturating_sub(required.start))); + let local_length = if tile < self.operand.format.layout.tiling.tile_count { + let local = clipped_range(self.operand_plans[axis].range(tile), extent); + u64::from( + required + .end + .min(local.end) + .saturating_sub(required.start.max(local.start)), + ) + } else { + 0 + }; + local_elements = local_elements.saturating_mul(local_length); + } + Some( + required_elements + .saturating_sub(local_elements) + .saturating_mul(self.operand.format.precision.bytes()), + ) + } +} + +fn clipped_range(range: std::ops::Range, extent: u32) -> std::ops::Range { + range.start.min(extent)..range.end.min(extent) +} + +fn add_operand_outgoing_bus_work( + traffic: &mut ExchangeEndpointTraffic, + operand: &TensorType, + remote: bool, + transfer_bytes: u64, +) { + if !remote { + return; + } + for tile in 0..operand.format.layout.tiling.tile_count { + let bytes = maximum_shard_bytes(operand); + traffic.add_outgoing(tile / 2, bytes, bytes.div_ceil(transfer_bytes)); + } +} + +#[derive(Clone)] +struct TileAxisPlan { + layout: Arc, + axis: usize, +} + +impl TileAxisPlan { + fn range(&self, tile: u16) -> std::ops::Range { + self.layout + .tile_axis_range(tile, self.axis) + .expect("axis plan was resolved for this tile") + } +} + +fn tile_axis_plans(tensor: &TensorType) -> Option> { + let layout = Arc::new(tensor.format.layout.resolve(&tensor.shape).ok()?); + (0..tensor.shape.0.len()) + .map(|axis| { + layout.maximum_axis_extent(axis)?; + Some(TileAxisPlan { + layout: Arc::clone(&layout), + axis, + }) + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{OperandRequirement, OperatorRequirements}; + + fn output_stationary_schedule() -> OperatorSchedule { + OperatorSchedule { + operator: crate::MidOperator::Gemm { + options: crate::GemmOptions::default(), + multiply: Precision::F16, + accumulate: crate::AccumulationPrecision::F16, + }, + steps: vec![ScheduleStep::Gemm(crate::GemmMap { + inputs: [ + crate::ScheduleValue::Input(0), + crate::ScheduleValue::Input(1), + ], + output: crate::ScheduleValue::Output, + kernel: crate::GemmKernelFamily { + multiply: Precision::F16, + accumulate: crate::AccumulationPrecision::F16, + weights: crate::GemmWeightLoad::Standard, + }, + geometry: crate::GemmGeometry { + block: crate::GemmBlockShape { + inner: AMP_INNER_BLOCK, + output_columns: crate::layout::AMP_OUTPUT_COLUMN_BLOCK, + }, + orientation: crate::GemmOrientation::Normal, + compute: crate::GemmGrid { + rows: 1, + columns: 1, + inner: 1, + }, + result: crate::GemmResultGrid { + rows: 1, + columns: 1, + }, + order: crate::GridOrder::ColumnsFast, + }, + })], + requirements: OperatorRequirements { + inputs: Vec::new(), + output: OperandRequirement::new( + crate::TensorFormat { + precision: Precision::F16, + layout: Layout::row_major(crate::TensorTiling::replicated(1)), + }, + 8, + ), + output_aliasing: crate::OutputAliasing::Fresh, + memory_space: crate::MemorySpaceRequirements::default(), + }, + } + } + + fn parallel_reduction_schedule( + row_partitions: u16, + column_partitions: u16, + inner_partitions: u16, + ) -> OperatorSchedule { + let mut schedule = output_stationary_schedule(); + let Some(ScheduleStep::Gemm(plan)) = schedule.steps.first_mut() else { + unreachable!(); + }; + plan.geometry.result = crate::GemmResultGrid { + rows: row_partitions, + columns: column_partitions, + }; + plan.geometry.compute = crate::GemmGrid { + rows: row_partitions, + columns: column_partitions, + inner: inner_partitions, + }; + plan.output = crate::ScheduleValue::Temporary(0); + schedule.steps.push(ScheduleStep::Reduce { + input: crate::ScheduleValue::Temporary(0), + output: crate::ScheduleValue::Output, + staging: crate::ReductionStaging::Streamed, + }); + schedule + } + + #[test] + fn randomized_average_shard_storage_covers_spatial_work() { + let mut random = fastrand::Rng::with_seed(0x7370_6174_6961_6c77); + for case in 0..32 { + let row_partitions = 1_u16 << random.u32(0..=4); + let column_partitions = 1_u16 << random.u32(0..=4); + let tiles = row_partitions * column_partitions; + let rows = u32::from(row_partitions) * random.u32(1..=8); + let columns = u32::from(column_partitions) * 64 * random.u32(1..=4); + let tensor = TensorType::new( + [rows, columns], + Precision::F16, + Layout::amp_output_grid( + crate::GemmOrientation::Normal, + 64, + tiles, + row_partitions, + column_partitions, + crate::operator::GridOrder::ColumnsFast, + ), + ); + let total = physical_elements(&tensor.shape, &tensor.format.layout) + .saturating_mul(tensor.format.precision.bytes()); + let average = average_shard_bytes(&tensor); + assert!( + average.saturating_mul(u64::from(tiles)) >= total, + "case {case}" + ); + assert!(average <= maximum_shard_bytes(&tensor), "case {case}"); + } + } + + #[test] + fn randomized_gemm_endpoint_traffic_tracks_both_exchange_directions() { + let mut random = fastrand::Rng::with_seed(0x6269_6469_7265_6374); + for case in 0..32 { + let row_partitions = 1_u16 << random.u32(1..=3); + let column_partitions = 1_u16 << random.u32(1..=3); + let tiles = row_partitions * column_partitions; + let rows = u32::from(row_partitions) * random.u32(1..=8); + let inner = AMP_INNER_BLOCK * random.u32(1..=4); + let columns = u32::from(column_partitions) + * crate::layout::AMP_OUTPUT_COLUMN_BLOCK + * random.u32(1..=3); + let output = TensorType::new( + [1, rows, columns], + Precision::F16, + Layout::amp_output_grid( + crate::GemmOrientation::Normal, + crate::layout::AMP_OUTPUT_COLUMN_BLOCK, + tiles, + row_partitions, + column_partitions, + crate::operator::GridOrder::ColumnsFast, + ), + ); + let local_left = TensorType::new( + [1, rows, inner], + Precision::F16, + Layout::amp_left_grid( + AMP_INNER_BLOCK as u16, + tiles, + row_partitions, + column_partitions, + crate::operator::GridOrder::ColumnsFast, + ), + ); + let local_right = TensorType::new( + [1, inner, columns], + Precision::F16, + Layout::block_major_matrix_grid( + AMP_INNER_BLOCK as u16, + crate::layout::AMP_OUTPUT_COLUMN_BLOCK, + tiles, + row_partitions, + column_partitions, + crate::operator::GridOrder::ColumnsFast, + ), + ); + let schedule = output_stationary_schedule(); + let local = gemm_exchange_endpoint_traffic( + &schedule, + &[local_left, local_right], + &output, + HardwareTarget::Ipu21, + ) + .unwrap(); + assert!(local.is_empty(), "case {case}"); + + let sharded_left = TensorType::new( + [1, rows, inner], + Precision::F16, + Layout::amp_left(AMP_INNER_BLOCK as u16, row_partitions), + ); + let sharded_right = TensorType::new( + [1, inner, columns], + Precision::F16, + Layout::block_major_matrix_storage( + crate::GemmOrientation::Normal, + AMP_INNER_BLOCK as u16, + crate::layout::AMP_OUTPUT_COLUMN_BLOCK, + column_partitions, + 1, + 1, + MemoryClass::Standard, + ), + ); + let remote = gemm_exchange_endpoint_traffic( + &schedule, + &[sharded_left, sharded_right], + &output, + HardwareTarget::Ipu21, + ) + .unwrap(); + assert!(remote.maximum_outgoing_bytes() != 0, "case {case}"); + assert!(remote.maximum_incoming_bytes() != 0, "case {case}"); + assert_eq!( + remote.maximum_payload_bytes(), + remote + .maximum_outgoing_bytes() + .max(remote.maximum_incoming_bytes()), + "case {case}" + ); + } + } + + #[test] + fn randomized_parallel_gemm_traffic_tracks_replica_shortfalls() { + let mut random = fastrand::Rng::with_seed(0x7265_706c_6963_6173); + for case in 0..32 { + let row_partitions = random.u16(2..=5); + let column_partitions = random.u16(2..=5); + let inner_partitions = random.u16(2..=5); + let tiles = row_partitions * column_partitions * inner_partitions; + let inner_block = 16 * random.u16(1..=4); + let column_block = 16 * random.u32(1..=4); + let rows = u32::from(row_partitions) * random.u32(1..=8); + let inner = u32::from(inner_partitions) * u32::from(inner_block); + let columns = u32::from(column_partitions) * column_block; + let left = TensorType::new( + [1, rows, inner], + Precision::F16, + Layout::amp_left_parallel_grid( + crate::GemmOrientation::Normal, + inner_block, + tiles, + row_partitions, + column_partitions, + inner_partitions, + ), + ); + let resident_right = TensorType::new( + [1, inner, columns], + Precision::F16, + Layout::block_major_matrix_storage( + crate::GemmOrientation::Normal, + inner_block, + column_block, + column_partitions, + inner_partitions, + row_partitions, + MemoryClass::Standard, + ), + ); + let compute_output = TensorType::new( + [1, rows, columns], + Precision::F16, + Layout::amp_left_result_grid( + crate::GemmOrientation::Normal, + column_block, + row_partitions * column_partitions, + row_partitions, + column_partitions, + crate::operator::GridOrder::ColumnsFast, + ), + ); + let schedule = + parallel_reduction_schedule(row_partitions, column_partitions, inner_partitions); + let resident = gemm_exchange_endpoint_traffic( + &schedule, + &[left.clone(), resident_right], + &compute_output, + HardwareTarget::Ipu21, + ) + .unwrap(); + assert!(resident.is_empty(), "case {case}"); + + let sharded_right = TensorType::new( + [1, inner, columns], + Precision::F16, + Layout::block_major_matrix_storage( + crate::GemmOrientation::Normal, + inner_block, + column_block, + column_partitions, + inner_partitions, + 1, + MemoryClass::Standard, + ), + ); + let expected_incoming = maximum_shard_bytes(&sharded_right); + let expected_outgoing = expected_incoming.saturating_mul(u64::from( + sharded_right.format.layout.tiling.tile_count.min(2), + )); + let streamed = gemm_exchange_endpoint_traffic( + &schedule, + &[left, sharded_right], + &compute_output, + HardwareTarget::Ipu21, + ) + .unwrap(); + assert_eq!( + streamed.maximum_incoming_bytes(), + expected_incoming, + "case {case}" + ); + assert_eq!( + streamed.maximum_outgoing_bytes(), + expected_outgoing, + "case {case}" + ); + } + } +} * Unmerged path crates/ipu-codegen/src/exchange.rs diff --git a/crates/ipu-codegen/src/host.rs b/crates/ipu-codegen/src/host.rs index 9989928..d84c1cd 100644 --- a/crates/ipu-codegen/src/host.rs +++ b/crates/ipu-codegen/src/host.rs @@ -1,16 +1,15 @@ -use crate::{HostPhase, HostProgram}; use ipu_package::{ - Binding, HostCall, HostExchange, HostPage, HostSlice, RegionSlice, SEGMENT_EXECUTE, - SEGMENT_READ, Segment, + AddressRegion, Binding, HostCall, HostExchange, HostPage, HostSlice, RegionSlice, + SEGMENT_EXECUTE, SEGMENT_READ, Segment, }; +use ipu_target::program::{HostPhase, HostProgram}; +use ipu_target::{exchange::ExchangeConstants, hardware::HardwareTarget}; use std::collections::{BTreeMap, BTreeSet, HashMap, VecDeque}; use super::package::{PackageBuildResult, invalid}; -const HOST_DATA_START: u32 = ipu_exchange::HOST_PAGE_BYTES; -const HOST_PACKET_ADDRESS: u32 = ipu_exchange::EXCHANGE_WINDOW_BASE; -const HOST_CLOSE_ADDRESS: u32 = ipu_exchange::EXCHANGE_WINDOW_BASE + 0x160; -const HOST_STAGING_ADDRESS: u32 = ipu_exchange::EXCHANGE_WINDOW_BASE + 0x180; +const HOST_CLOSE_OFFSET: u32 = 0x160; +const HOST_STAGING_OFFSET: u32 = 0x180; #[derive(Clone, Copy)] enum Direction { @@ -54,13 +53,15 @@ pub(crate) struct HostPackagePlan { } pub(crate) fn plan( + target: HardwareTarget, weights: &[Binding], inputs: &[Binding], outputs: &[Binding], execution_tiles: u16, base: u32, - data_ranges: &[Vec<(u32, u32)>], + data_ranges: &[Vec], ) -> PackageBuildResult { + let exchange = target.exchange(); if data_ranges.len() != usize::from(execution_tiles) { return Err(invalid("host plan has no data ranges for every tile")); } @@ -77,9 +78,9 @@ pub(crate) fn plan( let mut weight_cursor = 0; let mut input_cursor = 0; let mut output_cursor = 0; - let pending_weights = collect(weights, Direction::ToTile, &mut weight_cursor)?; - let pending_inputs = collect(inputs, Direction::ToTile, &mut input_cursor)?; - let pending_outputs = collect(outputs, Direction::ToHost, &mut output_cursor)?; + let pending_weights = collect(weights, Direction::ToTile, &mut weight_cursor, exchange)?; + let pending_inputs = collect(inputs, Direction::ToTile, &mut input_cursor, exchange)?; + let pending_outputs = collect(outputs, Direction::ToHost, &mut output_cursor, exchange)?; let participating = pending_weights .iter() .chain(&pending_inputs) @@ -91,17 +92,20 @@ pub(crate) fn plan( .enumerate() .map(|(slot, tile)| Ok((tile, u32::try_from(slot)?))) .collect::>>()?; - let (mut weight_phases, weight_slices, weight_ends) = batch(pending_weights, &slots)?; - let (mut input_phases, input_slices, input_ends) = batch(pending_inputs, &slots)?; - let (output_phases, output_slices, output_ends) = batch(pending_outputs, &slots)?; + let (mut weight_phases, weight_slices, weight_ends) = + batch(pending_weights, &slots, exchange)?; + let (mut input_phases, input_slices, input_ends) = + batch(pending_inputs, &slots, exchange)?; + let (output_phases, output_slices, output_ends) = + batch(pending_outputs, &slots, exchange)?; for transfer in weight_phases .iter_mut() .chain(&mut input_phases) .flat_map(|phase| &mut phase.transfers) { transfer.copy_destination = Some(transfer.tile_address); - transfer.tile_address = HOST_STAGING_ADDRESS; - ipu_exchange::plan_host_to_tile( + transfer.tile_address = exchange.window_base + HOST_STAGING_OFFSET; + ipu_target::exchange::plan_host_to_tile( transfer.physical_tile, transfer.tile_address, transfer.host_offset, @@ -109,7 +113,7 @@ pub(crate) fn plan( )?; } for transfer in output_phases.iter().flat_map(|phase| &phase.transfers) { - ipu_exchange::plan_tile_to_host( + ipu_target::exchange::plan_tile_to_host( transfer.physical_tile, transfer.tile_address, transfer.host_offset, @@ -132,6 +136,7 @@ pub(crate) fn plan( &phases, base, &data_ranges[usize::from(physical_tile)], + exchange, )?; maximum_end = maximum_end.max(planned.end); let weight_end = weight_phases.len(); @@ -172,7 +177,7 @@ pub(crate) fn plan( input_batch_ends: input_ends, output_batch_ends: output_ends, }); - let data_bytes = u64::from(ipu_exchange::HOST_PAGE_BYTES) + let data_bytes = u64::from(exchange.host_page_bytes) .checked_mul(u64::try_from(slots.len().max(1))?) .ok_or_else(|| invalid("host page arena overflow"))?; Ok(HostPackagePlan { @@ -185,7 +190,7 @@ pub(crate) fn plan( pages: vec![ HostPage { index: 0, - size: u64::from(ipu_exchange::HOST_PAGE_BYTES), + size: u64::from(exchange.host_page_bytes), }, HostPage { index: 1, @@ -196,7 +201,7 @@ pub(crate) fn plan( calls, }, end: maximum_end, - staging_address: HOST_STAGING_ADDRESS, + staging_address: exchange.window_base + HOST_STAGING_OFFSET, }) } @@ -204,12 +209,19 @@ fn collect( bindings: &[Binding], direction: Direction, cursor: &mut u64, + exchange: &ExchangeConstants, ) -> PackageBuildResult> { let mut result = Vec::new(); for binding in bindings { let base = *cursor; for slice in &binding.slices { - append_slice(&mut result, direction, slice, base)?; + append_slice( + &mut result, + direction, + slice, + base, + exchange.host_page_bytes, + )?; } *cursor = cursor .checked_add(binding_size(binding)?) @@ -223,6 +235,7 @@ fn append_slice( direction: Direction, slice: &RegionSlice, file_base: u64, + host_page_bytes: u32, ) -> PackageBuildResult<()> { let mut tile_address = slice.tile_address; let mut file_offset = file_base @@ -230,7 +243,7 @@ fn append_slice( .ok_or_else(|| invalid("host file offset overflow"))?; let mut remaining = u32::try_from(slice.size)?; while remaining != 0 { - let bytes = remaining.min(ipu_exchange::HOST_PAGE_BYTES); + let bytes = remaining.min(host_page_bytes); result.push(PendingTransfer { transfer: Transfer { direction, @@ -266,6 +279,7 @@ fn binding_size(binding: &Binding) -> PackageBuildResult { fn batch( pending: Vec, slots: &BTreeMap, + exchange: &ExchangeConstants, ) -> PackageBuildResult<(Vec, Vec, Vec)> { let mut queues = BTreeMap::>::new(); for transfer in pending { @@ -284,9 +298,10 @@ fn batch( continue; }; let page_offset = slots[&tile] - .checked_mul(ipu_exchange::HOST_PAGE_BYTES) + .checked_mul(exchange.host_page_bytes) .ok_or_else(|| invalid("host page offset overflow"))?; - pending.transfer.host_offset = HOST_DATA_START + pending.transfer.host_offset = exchange + .host_page_bytes .checked_add(page_offset) .ok_or_else(|| invalid("host exchange offset overflow"))?; slices.push(HostSlice { @@ -313,7 +328,8 @@ fn plan_tile( physical_tile: u16, phases: &[Phase], base: u32, - data_ranges: &[(u32, u32)], + data_ranges: &[AddressRegion], + exchange: &ExchangeConstants, ) -> PackageBuildResult { let follower = align_up(base, 8)?; let mut cursor = follower + 12; @@ -334,7 +350,7 @@ fn plan_tile( }); continue; } - let (instructions, packet_words) = phase_instructions(physical_tile, phase)?; + let (instructions, packet_words) = phase_instructions(physical_tile, phase, exchange)?; cursor = align_up(cursor, 8)?; let address = cursor; let data = words(&instructions); @@ -352,13 +368,13 @@ fn plan_tile( let packet = PacketCopy { source: packet_source, destination: if xreq_targets(physical_tile, phase)?.is_empty() { - HOST_PACKET_ADDRESS + 8 + exchange.window_base + 8 } else { - HOST_PACKET_ADDRESS + exchange.window_base }, words: u32::try_from(packet_words.len())?, }; - let descriptors = descriptor_words(physical_tile, phase, packet)?; + let descriptors = descriptor_words(physical_tile, phase, packet, exchange)?; let descriptor_data = words(&descriptors); let table = data_arena.allocate(u32::try_from(descriptor_data.len())?, 4)?; segments.push(segment(table, descriptor_data, SEGMENT_READ)); @@ -376,11 +392,11 @@ fn plan_tile( } struct DataArena { - ranges: Vec<(u32, u32)>, + ranges: Vec, } impl DataArena { - fn new(ranges: &[(u32, u32)]) -> Self { + fn new(ranges: &[AddressRegion]) -> Self { Self { ranges: ranges.to_vec(), } @@ -391,10 +407,10 @@ impl DataArena { .ranges .iter() .enumerate() - .filter_map(|(index, &(base, limit))| { - let start = align_up(base, alignment).ok()?; + .filter_map(|(index, range)| { + let start = align_up(range.start, alignment).ok()?; let end = start.checked_add(bytes)?; - (end <= limit).then_some((limit - end, index, start, end)) + (end <= range.end).then_some((range.end - end, index, start, end)) }) .min_by_key(|candidate| (candidate.0, candidate.2)) .ok_or_else(|| { @@ -403,14 +419,14 @@ impl DataArena { )) })?; let (_, index, start, end) = candidate; - let (base, limit) = self.ranges.remove(index); - if base < start { - self.ranges.push((base, start)); + let range = self.ranges.remove(index); + if range.start < start { + self.ranges.push(AddressRegion::new(range.start, start)); } - if end < limit { - self.ranges.push((end, limit)); + if end < range.end { + self.ranges.push(AddressRegion::new(end, range.end)); } - self.ranges.sort_unstable(); + self.ranges.sort_unstable_by_key(|range| range.start); Ok(start) } } @@ -418,14 +434,18 @@ impl DataArena { fn phase_instructions( physical_tile: u16, phase: &Phase, + exchange: &ExchangeConstants, ) -> PackageBuildResult<(Vec, Vec)> { let target = target(physical_tile, phase) - .map(|transfer| target_program(transfer, HOST_PACKET_ADDRESS + 8)) + .map(|transfer| target_program(transfer, exchange.window_base + 8, exchange)) .transpose()?; let targets = xreq_targets(physical_tile, phase)?; let xreq = (!targets.is_empty()) .then(|| { - ipu_exchange::assemble_host_xreq_program_for_targets(&targets, HOST_PACKET_ADDRESS) + ipu_target::exchange::assemble_host_xreq_program_for_targets( + &targets, + exchange.window_base, + ) }) .transpose()?; Ok(match (target, xreq) { @@ -433,20 +453,20 @@ fn phase_instructions( let mut packets = xreq.packet_words; packets.extend_from_slice(&target.packet_words); ( - ipu_exchange::wrap_combined_host_operation( + ipu_target::exchange::wrap_combined_host_operation( physical_tile, &target.instructions, - HOST_PACKET_ADDRESS, + exchange.window_base, )?, packets, ) } (None, Some(xreq)) => ( - ipu_exchange::wrap_host_xreq_operation(physical_tile, &xreq.instructions)?, + ipu_target::exchange::wrap_host_xreq_operation(physical_tile, &xreq.instructions)?, xreq.packet_words, ), (Some(target), None) => ( - ipu_exchange::wrap_host_target_operation(physical_tile, &target.instructions)?, + ipu_target::exchange::wrap_host_target_operation(physical_tile, &target.instructions)?, target.packet_words, ), (None, None) => return Err(invalid("active host phase has no work")), @@ -456,22 +476,23 @@ fn phase_instructions( fn target_program( transfer: Transfer, packet_address: u32, -) -> PackageBuildResult { + exchange: &ExchangeConstants, +) -> PackageBuildResult { Ok(match transfer.direction { - Direction::ToTile => ipu_exchange::assemble_host_to_tile_target_program( + Direction::ToTile => ipu_target::exchange::assemble_host_to_tile_target_program( transfer.physical_tile, transfer.tile_address, transfer.host_offset, transfer.bytes, packet_address, )?, - Direction::ToHost => ipu_exchange::assemble_tile_to_host_target_program( + Direction::ToHost => ipu_target::exchange::assemble_tile_to_host_target_program( transfer.physical_tile, transfer.tile_address, transfer.host_offset, transfer.bytes, packet_address, - HOST_CLOSE_ADDRESS, + exchange.window_base + HOST_CLOSE_OFFSET, )?, }) } @@ -480,6 +501,7 @@ fn descriptor_words( physical_tile: u16, phase: &Phase, packet: PacketCopy, + exchange: &ExchangeConstants, ) -> PackageBuildResult> { let target = target(physical_tile, phase); let copy_words = target @@ -489,8 +511,8 @@ fn descriptor_words( return Err(invalid("host descriptor is not encodable")); } let packet_destination = match packet.destination { - HOST_PACKET_ADDRESS => 0, - address if address == HOST_PACKET_ADDRESS + 8 => 1 << 23, + address if address == exchange.window_base => 0, + address if address == exchange.window_base + 8 => 1 << 23, _ => return Err(invalid("host packet destination is not encodable")), }; Ok(vec![ @@ -515,7 +537,7 @@ fn xreq_targets(physical_tile: u16, phase: &Phase) -> PackageBuildResult { Some(Ok(transfer.physical_tile)) } @@ -533,9 +555,9 @@ fn active(physical_tile: u16, phase: &Phase) -> bool { fn inactive_instructions() -> Vec { vec![ - ipu_exchange::sans(1), - ipu_exchange::SYNC_ANS_INSTRUCTION, - ipu_exchange::RETURN_M10_INSTRUCTION, + ipu_target::instruction::sans(1), + ipu_target::instruction::SYNC_ANS_INSTRUCTION, + ipu_target::instruction::RETURN_M10_INSTRUCTION, ] } * Unmerged path crates/ipu-codegen/src/kernel/build.rs diff --git a/crates/ipu-codegen/src/kernel/materialize.rs b/crates/ipu-codegen/src/kernel/materialize.rs new file mode 100644 index 0000000..00de461 --- /dev/null +++ b/crates/ipu-codegen/src/kernel/materialize.rs @@ -0,0 +1,78 @@ +use super::{KernelBuildPlan, KernelMaterializationError}; +use crate::{KernelRun, LowShard, LowShardId, view_byte_spans}; +use ipu_target::program::{ComputeStep, StepProfile, TileAddress}; +use std::collections::BTreeMap; + +/// Resolves one scheduled call after placement has assigned each shard base. +/// Layout conversion supplies the byte offset; the build plan supplies the +/// linked specialization and ABI scalar values. +pub fn materialize_kernel_run( + run: &KernelRun, + shards: &[LowShard], + shard_addresses: &BTreeMap, + plan: &KernelBuildPlan, + overrides: &BTreeMap, +) -> Result { + let call = plan.call(run)?; + let resolve = |view: &crate::ShardView| { + let shard = shards.get(view.shard.index() as usize).ok_or( + KernelMaterializationError::UnplacedShard(view.shard.index()), + )?; + let spans = view_byte_spans(shard, view)?; + let [span] = spans.as_slice() else { + return Err(KernelMaterializationError::FragmentedView { + shard: view.shard.index(), + spans: spans.len(), + }); + }; + let base = overrides.get(&view.shard).copied().unwrap_or_else(|| { + TileAddress::Absolute( + shard_addresses + .get(&view.shard) + .copied() + .unwrap_or_default(), + ) + }); + if !overrides.contains_key(&view.shard) && !shard_addresses.contains_key(&view.shard) { + return Err(KernelMaterializationError::UnplacedShard( + view.shard.index(), + )); + } + add_address_offset(base, span.offset) + }; + let output_address = resolve(&run.output)?; + let input_addresses = run + .inputs + .iter() + .map(|operand| resolve(&operand.views[0])) + .collect::, _>>()?; + Ok(ComputeStep { + symbol: call.symbol, + output_address, + input_addresses, + arguments: call.arguments, + profile: StepProfile::default(), + }) +} + +fn add_address_offset( + address: TileAddress, + offset: u32, +) -> Result { + Ok(match address { + TileAddress::Absolute(address) => TileAddress::Absolute( + address + .checked_add(offset) + .ok_or(KernelMaterializationError::AddressOverflow)?, + ), + TileAddress::RepeatPointer { + index, + offset: existing, + } => TileAddress::RepeatPointer { + index, + offset: existing + .checked_add(offset) + .ok_or(KernelMaterializationError::AddressOverflow)?, + }, + }) +} * Unmerged path crates/ipu-codegen/src/kernel/mod.rs diff --git a/crates/ipu-codegen/src/kernel/spec.rs b/crates/ipu-codegen/src/kernel/spec.rs new file mode 100644 index 0000000..13fdca5 --- /dev/null +++ b/crates/ipu-codegen/src/kernel/spec.rs @@ -0,0 +1,76 @@ +use crate::{AccumulationPrecision, GemmKernelMode, GemmWeightLoad, Layout, Precision}; + +/// A concrete tile-local callable produced during low lowering. +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)] +pub enum TileKernelSpec { + FillZero, + Gemm { + multiply: Precision, + accumulate: AccumulationPrecision, + mode: GemmKernelMode, + weights: GemmWeightLoad, + inner_block: u32, + output_columns: u32, + rows: u32, + }, + Gelu, + ReductionSum { + partials: u16, + }, + Add, + AttentionSoftmax { + query_rows: u32, + head_dimension: u32, + key_columns: u32, + padded_key_columns: u32, + }, + AttentionMerge { + query_rows: u32, + value_dimension: u32, + padded_value_dimension: u32, + key_block_columns: u32, + initial: bool, + final_block: bool, + }, + Cast { + from: Precision, + to: Precision, + }, + Rearrange { + from: Layout, + to: Layout, + matrices: u32, + logical_rows: u32, + physical_rows: u32, + logical_columns: u32, + physical_columns: u32, + }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum KernelSymbols { + Exact(&'static str), + Planned, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum KernelAvailability { + Implemented, + Required, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ScalarArgument { + pub register: u8, + pub name: &'static str, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct KernelAbi { + pub symbols: KernelSymbols, + pub availability: KernelAvailability, + pub output_register: u8, + pub input_registers: Vec, + pub scalar_arguments: Vec, + pub return_register: u8, +} diff --git a/crates/ipu-codegen/src/layout.rs b/crates/ipu-codegen/src/layout.rs new file mode 100644 index 0000000..a72e8ef --- /dev/null +++ b/crates/ipu-codegen/src/layout.rs @@ -0,0 +1,1637 @@ +use crate::graph::TensorShape; +use crate::operator::{GemmOrientation, GridOrder, Precision}; +use std::ops::Range; + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum TensorAxis { + FromStart(u16), + FromEnd(u16), + /// Grain-aligned ownership intervals in canonical logical element order. + Linear, +} + +impl TensorAxis { + pub fn resolve(self, rank: usize) -> Result { + match self { + Self::FromStart(axis) if usize::from(axis) < rank => Ok(usize::from(axis)), + Self::FromEnd(axis) if axis != 0 && usize::from(axis) <= rank => { + Ok(rank - usize::from(axis)) + } + _ => Err(LayoutError::AxisOutOfRange { axis: self, rank }), + } + } +} + +/// Kernel-native AMP register or accumulator-drain order. +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum NativeKernelOrder { + Left, + /// A semantic `[K, N]` matrix packed as the left operand `[N, K]`. + TransposedLeft, + /// Semantic `[key, channel]` storage packed as the right operand of + /// `query * key.transpose()`. + TransposedRight, + Output, + /// A semantic `[M, N]` output packed as the physical output `[N, M]`. + TransposedOutput, +} + +/// Ordinary elements grouped into rectangular blocks over two semantic axes. +/// `permutation` selects which named axis is the physical row and column. +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct BlockedOrder { + pub axes: [TensorAxis; 2], + pub block_shape: [u16; 2], + pub permutation: [u8; 2], +} + +impl BlockedOrder { + const MATRIX_AXES: [TensorAxis; 2] = [TensorAxis::FromEnd(2), TensorAxis::FromEnd(1)]; + + pub const fn matrix(row_block: u16, column_block: u16) -> Self { + Self { + axes: Self::MATRIX_AXES, + block_shape: [row_block, column_block], + permutation: [0, 1], + } + } + + pub const fn transposed_matrix(row_block: u16, column_block: u16) -> Self { + Self { + axes: Self::MATRIX_AXES, + block_shape: [row_block, column_block], + permutation: [1, 0], + } + } + + pub fn physical_axes(self, rank: usize) -> Result<[usize; 2], LayoutError> { + let axes = [self.axes[0].resolve(rank)?, self.axes[1].resolve(rank)?]; + let [row, column] = self.permutation; + if row > 1 || column > 1 || row == column || axes[0] == axes[1] { + return Err(LayoutError::InvalidStorageOrder); + } + Ok([axes[usize::from(row)], axes[usize::from(column)]]) + } + + pub fn is_matrix(self) -> bool { + self.axes == Self::MATRIX_AXES && self.permutation == [0, 1] + } + + pub fn is_transposed_matrix(self) -> bool { + self.axes == Self::MATRIX_AXES && self.permutation == [1, 0] + } +} + +pub const AMP_INNER_BLOCK: u32 = 64; +pub(crate) const AMP_NARROW_OUTPUT_COLUMN_BLOCK: u32 = 32; +pub const AMP_OUTPUT_COLUMN_BLOCK: u32 = 64; +pub(crate) const AMP_WIDE_OUTPUT_COLUMN_BLOCK: u32 = 128; +pub const AMP_COLUMN_MICRO: u32 = 16; + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum StorageOrder { + Linear, + Blocked(BlockedOrder), + Native(NativeKernelOrder), +} + +/// Physical traversal within one 16-by-16 F16 matrix micro-panel. Layouts +/// with the same order can exchange whole panels while changing their outer +/// ownership and panel sequence, without an intermediate rearrangement. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum F16MicroPanelOrder { + RowsThenColumns, + ColumnsThenRows, +} + +impl StorageOrder { + /// This packing is consumed as contiguous K-major panels, while a generic + /// intersection rearrangement produces rectangular tensor-coordinate + /// views. It must therefore be selected for an automatic input or produced + /// by a specialized operator/local staging path. + pub(crate) fn requires_direct_population(&self) -> bool { + matches!( + self, + Self::Blocked(order) if order.permutation == [1, 0] + ) || matches!(self, Self::Native(NativeKernelOrder::TransposedRight)) + } + + /// Whether a row-major logical staging shard can be transformed locally + /// into this order by the generated conversion kernels. + pub(crate) fn supports_row_major_population(self) -> bool { + match self { + Self::Linear + | Self::Native(NativeKernelOrder::Left | NativeKernelOrder::TransposedRight) => true, + Self::Blocked(order) => order.is_matrix(), + Self::Native(_) => false, + } + } + + pub(crate) const fn f16_micro_panel_order(self) -> Option { + match self { + Self::Native(NativeKernelOrder::Left | NativeKernelOrder::TransposedRight) => { + Some(F16MicroPanelOrder::RowsThenColumns) + } + Self::Native(NativeKernelOrder::TransposedLeft) | Self::Blocked(_) => { + Some(F16MicroPanelOrder::ColumnsThenRows) + } + Self::Linear + | Self::Native(NativeKernelOrder::Output | NativeKernelOrder::TransposedOutput) => None, + } + } + + /// Smallest column span which remains a self-contained physical fragment + /// when canonical linear ownership divides a matrix into row segments. + pub(crate) fn retained_linear_column_grain(self, precision: Precision) -> Option { + match self { + Self::Linear => Some(1), + Self::Native(NativeKernelOrder::Left) => Some(match precision { + Precision::F8F143 { .. } => 32, + Precision::F16 => 16, + Precision::F32 => 8, + }), + Self::Native(NativeKernelOrder::Output) => Some(AMP_COLUMN_MICRO), + Self::Blocked(_) | Self::Native(_) => None, + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum MemoryClass { + Standard, + Interleaved, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum Padding { + Reject, + Zero, +} + +/// Blocking and distribution of one logical tensor axis. +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct AxisTiling { + pub axis: TensorAxis, + /// Number of contiguous partitions distributed across the tile group. + pub partitions: u16, + /// Number of equal semantic groups which are padded independently before + /// partitioning. Partitions must subdivide groups evenly. This keeps, for + /// example, attention-head boundaries intact without coupling the number + /// of column shards to the number of heads. + pub padding_groups: u16, + /// Required physical block multiple. One imposes no blocking constraint. + pub block_size: u32, + /// Physical extent multiple, independently of the grain distributed + /// between partitions. This permits fine-grained ownership of an axis + /// whose producer operates on wider padded blocks. + pub padding_multiple: u32, + /// Physical extent multiple applied independently to every partition. + /// Unlike `padding_multiple`, this does not change semantic partition + /// boundaries: padding belongs to the allocation owned by that partition. + pub shard_padding_multiple: u32, + pub padding: Padding, + /// Optional physical-tile stride for this partition coordinate. When + /// absent, axes are packed after the replica coordinate and preceding + /// axes. Explicit strides allow operands of one operator to share a 2-D + /// tile grid while replicating along different grid dimensions. + pub tile_stride: Option, +} + +impl AxisTiling { + pub const fn new(axis: TensorAxis, partitions: u16, block_size: u32, padding: Padding) -> Self { + Self { + axis, + partitions, + padding_groups: 1, + block_size, + padding_multiple: block_size, + shard_padding_multiple: 1, + padding, + tile_stride: None, + } + } + + pub const fn with_tile_stride(mut self, tile_stride: u16) -> Self { + self.tile_stride = Some(tile_stride); + self + } + + pub const fn with_padding_multiple(mut self, padding_multiple: u32) -> Self { + self.padding_multiple = padding_multiple; + self + } + + pub const fn with_shard_padding_multiple(mut self, shard_padding_multiple: u32) -> Self { + self.shard_padding_multiple = shard_padding_multiple; + self + } + + pub const fn with_padding_groups(mut self, padding_groups: u16) -> Self { + self.padding_groups = padding_groups; + self + } +} + +/// Logical tile group and the tensor axes distributed or blocked within it. +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct TensorTiling { + pub tile_count: u16, + pub replicas: u16, + pub axes: Vec, +} + +impl TensorTiling { + pub fn linear(tile_count: u16, grain: u32) -> Self { + Self { + tile_count, + replicas: 1, + axes: vec![AxisTiling::new( + TensorAxis::Linear, + tile_count, + grain, + Padding::Reject, + )], + } + } + + pub fn linear_grain(&self) -> Option { + match self.axes.as_slice() { + [axis] if axis.axis == TensorAxis::Linear => Some(axis.block_size), + _ => None, + } + } + + pub fn replicated(tile_count: u16) -> Self { + Self { + tile_count, + replicas: tile_count, + axes: Vec::new(), + } + } + + pub fn sharded(axis: TensorAxis, tile_count: u16) -> Self { + Self { + tile_count, + replicas: 1, + axes: vec![AxisTiling::new(axis, tile_count, 1, Padding::Reject)], + } + } + + pub(crate) fn axis_strides(&self) -> Result, LayoutError> { + let mut packed_stride = u32::from(self.replicas); + self.axes + .iter() + .map(|axis| { + let stride = axis.tile_stride.map_or(packed_stride, u32::from); + packed_stride = packed_stride + .checked_mul(u32::from(axis.partitions)) + .ok_or(LayoutError::TileCountOverflow)?; + if stride == 0 { + return Err(LayoutError::EmptyAxisTiling); + } + Ok(stride) + }) + .collect() + } +} + +/// Layout decisions which constrain operators and exchange generation without +/// assigning physical tile identities or SRAM addresses. +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct Layout { + pub order: StorageOrder, + pub tiling: TensorTiling, + pub memory_class: MemoryClass, +} + +impl Layout { + pub fn row_major(tiling: TensorTiling) -> Self { + Self { + order: StorageOrder::Linear, + tiling, + memory_class: MemoryClass::Standard, + } + } + + pub fn logical_linear(tile_count: u16, grain: u32) -> Self { + Self::row_major(TensorTiling::linear(tile_count, grain)) + } + + pub fn with_retained_order_linear_ownership(&self, tile_count: u16, grain: u32) -> Self { + Self { + order: self.order, + tiling: TensorTiling::linear(tile_count, grain), + memory_class: self.memory_class, + } + } + + pub fn row_sharded(tile_count: u16) -> Self { + Self::row_major(TensorTiling::sharded(TensorAxis::FromEnd(2), tile_count)) + } + + pub fn head_sharded(tile_count: u16) -> Self { + Self::row_major(TensorTiling::sharded(TensorAxis::FromEnd(3), tile_count)) + } + + fn attention_tiling(heads: u16, query_partitions: u16) -> TensorTiling { + TensorTiling { + tile_count: heads.saturating_mul(query_partitions), + replicas: 1, + axes: vec![ + AxisTiling::new(TensorAxis::FromEnd(2), query_partitions, 1, Padding::Reject) + .with_tile_stride(heads), + AxisTiling::new(TensorAxis::FromEnd(3), heads, 1, Padding::Reject) + .with_tile_stride(1), + ], + } + } + + pub fn attention_query(heads: u16, query_partitions: u16) -> Self { + let mut tiling = Self::attention_tiling(heads, query_partitions); + tiling.axes.push(AxisTiling::new( + TensorAxis::FromEnd(1), + 1, + AMP_COLUMN_MICRO, + Padding::Zero, + )); + Self { + order: StorageOrder::Native(NativeKernelOrder::Left), + tiling, + memory_class: MemoryClass::Standard, + } + } + + pub fn attention_key(heads: u16, key_partitions: u16) -> Self { + let axes = vec![ + AxisTiling::new(TensorAxis::FromEnd(3), heads, 1, Padding::Reject).with_tile_stride(1), + AxisTiling::new( + TensorAxis::FromEnd(2), + key_partitions, + AMP_INNER_BLOCK, + Padding::Zero, + ) + .with_tile_stride(heads), + AxisTiling::new(TensorAxis::FromEnd(1), 1, AMP_COLUMN_MICRO, Padding::Zero), + ]; + Self { + order: StorageOrder::Native(NativeKernelOrder::TransposedRight), + tiling: TensorTiling { + tile_count: heads.saturating_mul(key_partitions), + replicas: 1, + axes, + }, + memory_class: MemoryClass::Standard, + } + } + + pub fn attention_block_major_key_value(heads: u16, key_partitions: u16) -> Self { + let mut layout = Self::attention_key(heads, key_partitions); + layout.order = StorageOrder::Blocked(BlockedOrder::matrix( + AMP_INNER_BLOCK as u16, + AMP_COLUMN_MICRO as u16, + )); + layout + } + + pub fn attention_output(heads: u16, query_partitions: u16) -> Self { + let mut tiling = Self::attention_tiling(heads, query_partitions); + tiling.axes.push(AxisTiling::new( + TensorAxis::FromEnd(1), + 1, + AMP_COLUMN_MICRO, + Padding::Zero, + )); + Self::row_major(tiling) + } + + pub fn amp_left(inner: u16, tile_count: u16) -> Self { + Self { + order: StorageOrder::Native(NativeKernelOrder::Left), + tiling: TensorTiling { + tile_count, + replicas: 1, + axes: vec![ + AxisTiling::new(TensorAxis::FromEnd(2), tile_count, 1, Padding::Reject), + AxisTiling::new(TensorAxis::FromEnd(1), 1, u32::from(inner), Padding::Zero), + ], + }, + memory_class: MemoryClass::Standard, + } + } + + pub fn block_major_matrix(row_block: u16, tile_count: u16) -> Self { + Self::block_major_matrix_storage( + GemmOrientation::Normal, + row_block, + AMP_OUTPUT_COLUMN_BLOCK, + tile_count, + 1, + 1, + MemoryClass::Standard, + ) + } + + pub fn amp_output(tile_count: u16) -> Self { + Self { + order: StorageOrder::Native(NativeKernelOrder::Output), + tiling: TensorTiling { + tile_count, + replicas: 1, + axes: vec![ + AxisTiling::new(TensorAxis::FromEnd(2), tile_count, 1, Padding::Reject), + AxisTiling::new(TensorAxis::FromEnd(1), 1, 64, Padding::Zero), + ], + }, + memory_class: MemoryClass::Interleaved, + } + } + + /// F16 AMP result stored in the same within-panel order as a following + /// left operand. The GEMM coefficient routing makes the native accumulator + /// drain land in this order without a post-compute permutation. + pub fn amp_left_result(tile_count: u16) -> Self { + let mut layout = Self::amp_output(tile_count); + layout.order = StorageOrder::Native(NativeKernelOrder::Left); + layout + } + + /// AMP left operand on a row-by-column tile grid. The row shard is + /// replicated across column groups so it is local to every output shard. + pub fn amp_left_grid( + inner: u16, + tile_count: u16, + row_partitions: u16, + column_partitions: u16, + grid_order: GridOrder, + ) -> Self { + if column_partitions == 1 && row_partitions == tile_count { + return Self::amp_left(inner, tile_count); + } + Self { + order: StorageOrder::Native(NativeKernelOrder::Left), + tiling: TensorTiling { + tile_count, + replicas: column_partitions, + axes: vec![ + AxisTiling::new(TensorAxis::FromEnd(2), row_partitions, 1, Padding::Reject) + .with_tile_stride(match grid_order { + GridOrder::ColumnsFast => column_partitions, + GridOrder::RowsFast => 1, + }), + AxisTiling::new(TensorAxis::FromEnd(1), 1, u32::from(inner), Padding::Zero), + ], + }, + memory_class: MemoryClass::Standard, + } + } + + /// AMP left operand for a row-by-column-by-K compute grid. K and rows + /// are true shards; the column coordinate is a replica because the same + /// activation range is consumed by each output-column group. + pub fn amp_left_parallel_grid( + orientation: GemmOrientation, + inner: u16, + tile_count: u16, + row_partitions: u16, + column_partitions: u16, + inner_partitions: u16, + ) -> Self { + Self { + order: StorageOrder::Native(match orientation { + GemmOrientation::Normal => NativeKernelOrder::Left, + GemmOrientation::Swapped => NativeKernelOrder::TransposedLeft, + }), + tiling: TensorTiling { + tile_count, + replicas: column_partitions, + axes: vec![ + AxisTiling::new( + orientation.column_axis(), + inner_partitions, + u32::from(inner), + Padding::Zero, + ), + AxisTiling::new(orientation.row_axis(), row_partitions, 1, Padding::Reject), + ], + }, + memory_class: MemoryClass::Standard, + } + } + + /// Block-major matrix storage on a row-by-column tile grid. Each column + /// shard is replicated across row groups so it is local to every consumer. + pub fn block_major_matrix_grid( + inner: u16, + output_column_block: u32, + tile_count: u16, + row_partitions: u16, + column_partitions: u16, + grid_order: GridOrder, + ) -> Self { + Self { + order: StorageOrder::Blocked(BlockedOrder::matrix(inner, AMP_COLUMN_MICRO as u16)), + tiling: TensorTiling { + tile_count, + replicas: row_partitions, + axes: vec![ + AxisTiling::new(TensorAxis::FromEnd(2), 1, u32::from(inner), Padding::Zero), + AxisTiling::new( + TensorAxis::FromEnd(1), + column_partitions, + output_column_block, + Padding::Zero, + ) + .with_tile_stride(match grid_order { + GridOrder::ColumnsFast => 1, + GridOrder::RowsFast => row_partitions, + }), + ], + }, + memory_class: MemoryClass::Standard, + } + } + + /// Matrix storage with complete row-by-column blocks contiguous in the + /// selected memory class. Column and row sharding select the owner set; + /// `copies` controls persistent replication independently of consumers. + pub fn block_major_matrix_storage( + orientation: GemmOrientation, + inner_block: u16, + output_column_block: u32, + column_partitions: u16, + inner_partitions: u16, + copies: u16, + memory_class: MemoryClass, + ) -> Self { + let tile_count = column_partitions + .checked_mul(inner_partitions) + .and_then(|tiles| tiles.checked_mul(copies)) + .unwrap_or(0); + Self { + order: StorageOrder::Blocked(match orientation { + GemmOrientation::Normal => { + BlockedOrder::matrix(inner_block, AMP_COLUMN_MICRO as u16) + } + GemmOrientation::Swapped => { + BlockedOrder::transposed_matrix(inner_block, AMP_COLUMN_MICRO as u16) + } + }), + tiling: TensorTiling { + tile_count, + replicas: copies, + axes: vec![ + AxisTiling::new( + orientation.column_axis(), + column_partitions, + output_column_block, + Padding::Zero, + ) + .with_tile_stride(1), + AxisTiling::new( + orientation.row_axis(), + inner_partitions, + u32::from(inner_block), + Padding::Zero, + ), + ], + }, + memory_class, + } + } + + /// AMP output distributed over both matrix axes on one tile grid. + pub fn amp_output_grid( + orientation: GemmOrientation, + output_column_block: u32, + tile_count: u16, + row_partitions: u16, + column_partitions: u16, + grid_order: GridOrder, + ) -> Self { + if orientation == GemmOrientation::Normal + && output_column_block == AMP_OUTPUT_COLUMN_BLOCK + && column_partitions == 1 + && row_partitions == tile_count + { + return Self::amp_output(tile_count); + } + Self { + order: StorageOrder::Native(match orientation { + GemmOrientation::Normal => NativeKernelOrder::Output, + GemmOrientation::Swapped => NativeKernelOrder::TransposedOutput, + }), + tiling: TensorTiling { + tile_count, + replicas: 1, + axes: vec![ + AxisTiling::new( + orientation.column_axis(), + column_partitions, + output_column_block, + Padding::Zero, + ) + .with_tile_stride(match grid_order { + GridOrder::ColumnsFast => 1, + GridOrder::RowsFast => row_partitions, + }), + AxisTiling::new(orientation.row_axis(), row_partitions, 1, Padding::Reject) + .with_tile_stride(match grid_order { + GridOrder::ColumnsFast => column_partitions, + GridOrder::RowsFast => 1, + }), + ], + }, + memory_class: MemoryClass::Interleaved, + } + } + + pub fn amp_left_result_grid( + orientation: GemmOrientation, + output_column_block: u32, + tile_count: u16, + row_partitions: u16, + column_partitions: u16, + grid_order: GridOrder, + ) -> Self { + let mut layout = Self::amp_output_grid( + orientation, + output_column_block, + tile_count, + row_partitions, + column_partitions, + grid_order, + ); + layout.order = StorageOrder::Native(match orientation { + GemmOrientation::Normal => NativeKernelOrder::Left, + GemmOrientation::Swapped => NativeKernelOrder::TransposedLeft, + }); + layout + } + + /// AMP output storage sharded by rows and replicated across column groups. + pub fn amp_output_replicated_grid( + tile_count: u16, + row_partitions: u16, + column_replicas: u16, + ) -> Self { + if column_replicas == 1 && row_partitions == tile_count { + return Self::amp_output(tile_count); + } + Self { + order: StorageOrder::Native(NativeKernelOrder::Output), + tiling: TensorTiling { + tile_count, + replicas: column_replicas, + axes: vec![ + AxisTiling::new(TensorAxis::FromEnd(2), row_partitions, 1, Padding::Reject), + AxisTiling::new(TensorAxis::FromEnd(1), 1, 64, Padding::Zero), + ], + }, + memory_class: MemoryClass::Interleaved, + } + } +} + +#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)] +pub enum LayoutError { + #[error("layout has an empty tile group")] + EmptyTileGroup, + #[error("axis tiling must have nonzero partitions and block size")] + EmptyAxisTiling, + #[error("tile count calculation overflowed")] + TileCountOverflow, + #[error("axis {axis:?} is outside rank {rank}")] + AxisOutOfRange { axis: TensorAxis, rank: usize }, + #[error("tensor rank {0} cannot be represented by shard axis identifiers")] + RankTooLarge(usize), + #[error("storage axes or permutation are invalid")] + InvalidStorageOrder, + #[error("axis {0} is tiled more than once")] + DuplicateAxis(usize), + #[error("axis {axis} extent {extent} is not divisible by block size {block_size}")] + IndivisibleAxis { + axis: usize, + extent: u32, + block_size: u32, + }, + #[error("shard extent {extent} is not divisible by block size {block_size}")] + IndivisibleShard { extent: u32, block_size: u32 }, + #[error( + "axis extent {extent} cannot be divided into {groups} padding groups and {partitions} partitions" + )] + InvalidPaddingGroups { + groups: u16, + partitions: u16, + extent: u32, + }, + #[error("padded extent for axis {0} overflowed")] + ExtentOverflow(usize), + #[error("layout declares {declared} tiles but its tiling implies {implied}")] + TileCountMismatch { declared: u16, implied: u32 }, + #[error("tile strides do not form the declared partition and replica mapping")] + InvalidTileMapping, +} + +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct TensorFormat { + pub precision: Precision, + pub layout: Layout, +} + +impl TensorFormat { + pub(crate) fn supports_f16_micro_panel_exchange(&self, destination: &Self) -> bool { + self.precision == Precision::F16 + && destination.precision == Precision::F16 + && self.layout.order.f16_micro_panel_order().is_some() + && self.layout.order.f16_micro_panel_order() + == destination.layout.order.f16_micro_panel_order() + } +} + +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct TensorType { + pub shape: TensorShape, + pub format: TensorFormat, +} + +impl TensorType { + pub fn new(shape: impl IntoIterator, precision: Precision, layout: Layout) -> Self { + Self { + shape: TensorShape::new(shape), + format: TensorFormat { precision, layout }, + } + } +} + +/// Half-open bounds along one tensor axis. `physical_end` includes any zero +/// padding while `logical_end` never exceeds the semantic tensor shape. +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct ShardExtent { + pub axis: u16, + pub start: u32, + pub logical_end: u32, + pub physical_end: u32, +} + +#[derive(Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct TensorRegion { + pub extents: Vec, +} + +impl TensorRegion { + pub fn new(extents: impl IntoIterator) -> Self { + Self { + extents: extents.into_iter().collect(), + } + } + + pub fn logical_bounds(bounds: impl IntoIterator) -> Option { + bounds + .into_iter() + .enumerate() + .map(|(axis, (start, end))| { + Some(ShardExtent { + axis: u16::try_from(axis).ok()?, + start, + logical_end: end, + physical_end: end, + }) + }) + .collect::>>() + .map(Self::new) + } + + pub fn logical(&self) -> Self { + Self::new(self.extents.iter().map(|extent| ShardExtent { + physical_end: extent.logical_end, + ..*extent + })) + } + + pub fn physical(&self) -> Self { + Self::new(self.extents.iter().map(|extent| ShardExtent { + logical_end: extent.physical_end, + ..*extent + })) + } + + pub fn intersection(&self, other: &Self) -> Option { + if self.extents.len() != other.extents.len() { + return None; + } + self.extents + .iter() + .zip(&other.extents) + .map(|(left, right)| { + if left.axis != right.axis { + return None; + } + let start = left.start.max(right.start); + let end = left.logical_end.min(right.logical_end); + (start < end).then_some(ShardExtent { + axis: left.axis, + start, + logical_end: end, + physical_end: end, + }) + }) + .collect::>>() + .map(Self::new) + } + + pub fn logical_elements(&self) -> u64 { + self.extents.iter().fold(1, |elements, extent| { + elements.saturating_mul(u64::from(extent.logical_end.saturating_sub(extent.start))) + }) + } +} + +impl std::ops::Deref for TensorRegion { + type Target = [ShardExtent]; + + fn deref(&self) -> &Self::Target { + &self.extents + } +} + +impl std::ops::DerefMut for TensorRegion { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.extents + } +} + +impl<'a> IntoIterator for &'a TensorRegion { + type Item = &'a ShardExtent; + type IntoIter = std::slice::Iter<'a, ShardExtent>; + + fn into_iter(self) -> Self::IntoIter { + self.extents.iter() + } +} + +impl<'a> IntoIterator for &'a mut TensorRegion { + type Item = &'a mut ShardExtent; + type IntoIter = std::slice::IterMut<'a, ShardExtent>; + + fn into_iter(self) -> Self::IntoIter { + self.extents.iter_mut() + } +} + +impl From> for TensorRegion { + fn from(extents: Vec) -> Self { + Self { extents } + } +} + +impl FromIterator for TensorRegion { + fn from_iter>(iter: T) -> Self { + Self::new(iter) + } +} + +/// One rectangular region owned by a tile. Linear ownership can give a tile +/// several regions when its interval crosses row boundaries. +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) struct ResolvedShard { + pub tile: u16, + pub extents: TensorRegion, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +struct ResolvedAxis { + index: usize, + tiling: AxisTiling, + tile_stride: u32, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum ResolvedOwnership { + Linear { grain: u32 }, + Axes, +} + +/// Shape-dependent ownership implied by a [`Layout`]. +/// +/// This is the canonical source for padded dimensions, shard ranges, and +/// physical allocation sizes. Element order is intentionally left unresolved: +/// storage code maps these tensor regions into row-major, block-major, or AMP +/// byte spans. +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) struct ResolvedLayout { + shape: TensorShape, + padded_shape: TensorShape, + tile_count: u16, + replicas: u16, + axes: Vec, + ownership: ResolvedOwnership, +} + +impl Layout { + pub(crate) fn resolve(&self, shape: &TensorShape) -> Result { + if self.tiling.tile_count == 0 || self.tiling.replicas == 0 { + return Err(LayoutError::EmptyTileGroup); + } + if shape.0.len() > usize::from(u16::MAX) { + return Err(LayoutError::RankTooLarge(shape.0.len())); + } + if let StorageOrder::Blocked(order) = self.order { + order.physical_axes(shape.0.len())?; + if order.block_shape.contains(&0) { + return Err(LayoutError::InvalidStorageOrder); + } + } + if let Some(grain) = self.tiling.linear_grain() { + let elements = shape.elements(); + if shape.0.is_empty() + || grain == 0 + || elements / u64::from(grain) < u64::from(self.tiling.tile_count) + || !elements.is_multiple_of(u64::from(grain)) + { + return Err(LayoutError::EmptyAxisTiling); + } + return Ok(ResolvedLayout { + shape: shape.clone(), + padded_shape: shape.clone(), + tile_count: self.tiling.tile_count, + replicas: self.tiling.replicas, + axes: Vec::new(), + ownership: ResolvedOwnership::Linear { grain }, + }); + } + + let mut used_tiles = u32::from(self.tiling.replicas); + let mut dimensions = shape.0.clone(); + let mut used_axes = Vec::with_capacity(self.tiling.axes.len()); + for tiling in &self.tiling.axes { + if tiling.partitions == 0 + || tiling.padding_groups == 0 + || tiling.block_size == 0 + || tiling.padding_multiple == 0 + || tiling.shard_padding_multiple == 0 + { + return Err(LayoutError::EmptyAxisTiling); + } + used_tiles = used_tiles + .checked_mul(u32::from(tiling.partitions)) + .ok_or(LayoutError::TileCountOverflow)?; + let axis = tiling.axis.resolve(dimensions.len())?; + if used_axes.contains(&axis) { + return Err(LayoutError::DuplicateAxis(axis)); + } + used_axes.push(axis); + let extent = dimensions[axis]; + if !u32::from(tiling.partitions).is_multiple_of(u32::from(tiling.padding_groups)) + || !extent.is_multiple_of(u32::from(tiling.padding_groups)) + { + return Err(LayoutError::InvalidPaddingGroups { + groups: tiling.padding_groups, + partitions: tiling.partitions, + extent, + }); + } + let group_extent = extent / u32::from(tiling.padding_groups); + let remainder = group_extent % tiling.padding_multiple; + if remainder != 0 && tiling.padding == Padding::Reject { + return Err(LayoutError::IndivisibleAxis { + axis, + extent: group_extent, + block_size: tiling.padding_multiple, + }); + } + if remainder != 0 { + let padded_group_extent = group_extent + .checked_add(tiling.padding_multiple - remainder) + .ok_or(LayoutError::ExtentOverflow(axis))?; + dimensions[axis] = padded_group_extent + .checked_mul(u32::from(tiling.padding_groups)) + .ok_or(LayoutError::ExtentOverflow(axis))?; + } + let padded_group_extent = dimensions[axis] / u32::from(tiling.padding_groups); + if !padded_group_extent.is_multiple_of(tiling.block_size) { + return Err(LayoutError::IndivisibleAxis { + axis, + extent: padded_group_extent, + block_size: tiling.block_size, + }); + } + } + if used_tiles != u32::from(self.tiling.tile_count) { + return Err(LayoutError::TileCountMismatch { + declared: self.tiling.tile_count, + implied: used_tiles, + }); + } + + let strides = self.tiling.axis_strides()?; + validate_tile_mapping( + &self.tiling.axes, + self.tiling.replicas, + &strides, + self.tiling.tile_count, + )?; + let axes = self + .tiling + .axes + .iter() + .copied() + .zip(strides) + .map(|(tiling, tile_stride)| { + Ok(ResolvedAxis { + index: tiling.axis.resolve(dimensions.len())?, + tiling, + tile_stride, + }) + }) + .collect::, LayoutError>>()?; + let resolved = ResolvedLayout { + shape: shape.clone(), + padded_shape: TensorShape(dimensions), + tile_count: self.tiling.tile_count, + replicas: self.tiling.replicas, + axes, + ownership: ResolvedOwnership::Axes, + }; + // Every shard has one of two adjacent block counts. Validate one long, + // one short, and the final coordinate rather than scanning all tiles. + for axis in &resolved.axes { + let (partitions_per_group, _, _, long_shards) = + axis.partition_geometry(&resolved.padded_shape); + axis.bounds(&resolved.shape, &resolved.padded_shape, 0)?; + if long_shards < partitions_per_group { + axis.bounds(&resolved.shape, &resolved.padded_shape, long_shards)?; + } + axis.bounds( + &resolved.shape, + &resolved.padded_shape, + u32::from(axis.tiling.partitions) - 1, + )?; + } + Ok(resolved) + } +} + +impl ResolvedLayout { + pub(crate) fn padded_shape(&self) -> &TensorShape { + &self.padded_shape + } + + pub(crate) const fn tile_count(&self) -> u16 { + self.tile_count + } + + pub(crate) fn total_elements(&self) -> u64 { + match self.ownership { + ResolvedOwnership::Linear { .. } => self.shape.elements(), + ResolvedOwnership::Axes => self + .padded_shape + .0 + .iter() + .enumerate() + .map(|(index, &extent)| { + self.axis(index).map_or(u64::from(extent), |axis| { + axis.total_physical_extent(&self.padded_shape) + }) + }) + .product::() + .saturating_mul(u64::from(self.replicas)), + } + } + + pub(crate) fn maximum_tile_elements(&self) -> u64 { + match self.ownership { + ResolvedOwnership::Linear { grain } => { + let grains = self.shape.elements() / u64::from(grain); + grains + .div_ceil(u64::from(self.tile_count)) + .saturating_mul(u64::from(grain)) + } + ResolvedOwnership::Axes => self + .padded_shape + .0 + .iter() + .enumerate() + .map(|(index, &extent)| { + self.axis(index).map_or(u64::from(extent), |axis| { + u64::from(axis.maximum_physical_extent(&self.padded_shape)) + }) + }) + .product(), + } + } + + pub(crate) fn tile_elements(&self, tile: u16) -> Option { + if tile >= self.tile_count { + return None; + } + match self.ownership { + ResolvedOwnership::Linear { grain } => { + let grains = self.shape.elements() / u64::from(grain); + let short = grains / u64::from(self.tile_count); + let long = grains % u64::from(self.tile_count); + Some((short + u64::from(u64::from(tile) < long)).saturating_mul(u64::from(grain))) + } + ResolvedOwnership::Axes => self.padded_shape.0.iter().enumerate().try_fold( + 1_u64, + |elements, (index, &extent)| { + let width = self.axis(index).map_or(u64::from(extent), |axis| { + let coordinate = (u32::from(tile) / axis.tile_stride) + % u32::from(axis.tiling.partitions); + let (start, _, end) = axis + .bounds(&self.shape, &self.padded_shape, coordinate) + .expect("resolved axis bounds remain valid"); + u64::from(end - start) + }); + elements.checked_mul(width) + }, + ), + } + } + + pub(crate) fn has_empty_shards(&self) -> bool { + match self.ownership { + ResolvedOwnership::Linear { .. } => false, + ResolvedOwnership::Axes => self.axes.iter().any(|axis| { + let (partitions_per_group, _, _, _) = axis.partition_geometry(&self.padded_shape); + let (start, logical_end, _) = axis + .bounds(&self.shape, &self.padded_shape, partitions_per_group - 1) + .expect("resolved axis bounds remain valid"); + start == logical_end + }), + } + } + + pub(crate) fn maximum_axis_extent(&self, axis: usize) -> Option { + let &padded_extent = self.padded_shape.0.get(axis)?; + if matches!(self.ownership, ResolvedOwnership::Linear { .. }) { + return Some(padded_extent); + } + Some(self.axis(axis).map_or(padded_extent, |resolved| { + resolved.maximum_physical_extent(&self.padded_shape) + })) + } + + /// Physical range of one axis assigned to `tile`. Linear ownership may + /// cross higher-dimensional rows, so it conservatively returns the whole + /// axis just as the previous traffic estimator did. + pub(crate) fn tile_axis_range(&self, tile: u16, axis: usize) -> Option> { + if tile >= self.tile_count { + return None; + } + let &padded_extent = self.padded_shape.0.get(axis)?; + if matches!(self.ownership, ResolvedOwnership::Linear { .. }) { + return Some(0..padded_extent); + } + Some(self.axis(axis).map_or(0..padded_extent, |resolved| { + let coordinate = + (u32::from(tile) / resolved.tile_stride) % u32::from(resolved.tiling.partitions); + let (start, _, physical_end) = resolved + .bounds(&self.shape, &self.padded_shape, coordinate) + .expect("resolved axis bounds remain valid"); + start..physical_end + })) + } + + pub(crate) fn axis_bounds(&self, axis: usize, coordinate: u32) -> Option<(u32, u32, u32)> { + self.axis(axis)? + .bounds(&self.shape, &self.padded_shape, coordinate) + .ok() + } + + pub(crate) fn shard_extents(&self) -> Vec { + let shards = match self.ownership { + ResolvedOwnership::Linear { grain } => self.linear_shard_extents(grain), + ResolvedOwnership::Axes => (0..self.tile_count) + .map(|tile| { + let extents = self + .padded_shape + .0 + .iter() + .enumerate() + .map(|(index, &padded_extent)| { + let (start, logical_end, physical_end) = self.axis(index).map_or( + (0, self.shape.0[index], padded_extent), + |axis| { + let coordinate = (u32::from(tile) / axis.tile_stride) + % u32::from(axis.tiling.partitions); + axis.bounds(&self.shape, &self.padded_shape, coordinate) + .expect("resolved axis bounds remain valid") + }, + ); + ShardExtent { + axis: index as u16, + start, + logical_end, + physical_end, + } + }) + .collect(); + ResolvedShard { tile, extents } + }) + .collect(), + }; + debug_assert!({ + let mut materialized = vec![0_u64; usize::from(self.tile_count)]; + for shard in &shards { + let elements = shard.extents.iter().fold(1_u64, |elements, extent| { + elements.saturating_mul(u64::from(extent.physical_end - extent.start)) + }); + materialized[usize::from(shard.tile)] = + materialized[usize::from(shard.tile)].saturating_add(elements); + } + materialized + .into_iter() + .enumerate() + .all(|(tile, elements)| self.tile_elements(tile as u16) == Some(elements)) + }); + shards + } + + fn axis(&self, index: usize) -> Option<&ResolvedAxis> { + self.axes.iter().find(|axis| axis.index == index) + } + + fn linear_shard_extents(&self, grain: u32) -> Vec { + let elements = self.shape.elements(); + let grains = elements / u64::from(grain); + let tiles = u64::from(self.tile_count); + let rank = self.shape.0.len(); + let width = u64::from(self.shape.0[rank - 1]); + let mut shards = Vec::new(); + for tile in 0..self.tile_count { + let start_grain = + u64::from(tile) * (grains / tiles) + u64::from(tile).min(grains % tiles); + let tile_grains = grains / tiles + u64::from(u64::from(tile) < grains % tiles); + let start = start_grain * u64::from(grain); + let end = start + tile_grains * u64::from(grain); + let first_row = start / width; + let last_row = end.div_ceil(width); + for row in first_row..last_row { + let mut coordinates = vec![0u32; rank.saturating_sub(1)]; + let mut linear_row = row; + for axis in (0..rank.saturating_sub(1)).rev() { + let extent = u64::from(self.shape.0[axis]); + coordinates[axis] = (linear_row % extent) as u32; + linear_row /= extent; + } + let mut extents = coordinates + .into_iter() + .enumerate() + .map(|(axis, coordinate)| ShardExtent { + axis: axis as u16, + start: coordinate, + logical_end: coordinate + 1, + physical_end: coordinate + 1, + }) + .collect::>(); + let column_start = if row == first_row { start % width } else { 0 }; + let column_end = if row + 1 == last_row && !end.is_multiple_of(width) { + end % width + } else { + width + }; + extents.push(ShardExtent { + axis: (rank - 1) as u16, + start: column_start as u32, + logical_end: column_end as u32, + physical_end: column_end as u32, + }); + shards.push(ResolvedShard { + tile, + extents: extents.into(), + }); + } + } + shards + } +} + +impl ResolvedAxis { + fn partition_geometry(self, padded_shape: &TensorShape) -> (u32, u32, u32, u32) { + let groups = u32::from(self.tiling.padding_groups); + let partitions_per_group = u32::from(self.tiling.partitions) / groups; + let padded_group_extent = padded_shape.0[self.index] / groups; + let blocks = padded_group_extent / self.tiling.block_size; + let short_blocks = blocks / partitions_per_group; + let long_shards = blocks % partitions_per_group; + (partitions_per_group, blocks, short_blocks, long_shards) + } + + fn padded_block_width(self, blocks: u32) -> u32 { + let allocated = blocks * self.tiling.block_size; + allocated + .div_ceil(self.tiling.shard_padding_multiple) + .saturating_mul(self.tiling.shard_padding_multiple) + } + + fn maximum_physical_extent(self, padded_shape: &TensorShape) -> u32 { + let (_, _, short_blocks, long_shards) = self.partition_geometry(padded_shape); + self.padded_block_width(short_blocks + u32::from(long_shards != 0)) + } + + fn total_physical_extent(self, padded_shape: &TensorShape) -> u64 { + let (partitions_per_group, _, short_blocks, long_shards) = + self.partition_geometry(padded_shape); + let short_shards = partitions_per_group - long_shards; + let group_width = u64::from(self.padded_block_width(short_blocks)) + .saturating_mul(u64::from(short_shards)) + .saturating_add( + u64::from(self.padded_block_width(short_blocks + 1)) + .saturating_mul(u64::from(long_shards)), + ); + group_width.saturating_mul(u64::from(self.tiling.padding_groups)) + } + + fn bounds( + self, + shape: &TensorShape, + padded_shape: &TensorShape, + coordinate: u32, + ) -> Result<(u32, u32, u32), LayoutError> { + resolve_axis_bounds( + self.tiling, + padded_shape.0[self.index], + shape.0[self.index], + coordinate, + self.index, + ) + } +} + +fn resolve_axis_bounds( + tiling: AxisTiling, + padded_extent: u32, + logical_extent: u32, + coordinate: u32, + axis: usize, +) -> Result<(u32, u32, u32), LayoutError> { + let partitions = u32::from(tiling.partitions); + let groups = u32::from(tiling.padding_groups); + if groups == 0 + || coordinate >= partitions + || !partitions.is_multiple_of(groups) + || !padded_extent.is_multiple_of(groups) + || !logical_extent.is_multiple_of(groups) + { + return Err(LayoutError::InvalidPaddingGroups { + groups: tiling.padding_groups, + partitions: tiling.partitions, + extent: logical_extent, + }); + } + let partitions_per_group = partitions / groups; + let group = coordinate / partitions_per_group; + let coordinate_in_group = coordinate % partitions_per_group; + let padded_group_extent = padded_extent / groups; + let logical_group_extent = logical_extent / groups; + if !padded_group_extent.is_multiple_of(tiling.block_size) { + return Err(LayoutError::IndivisibleAxis { + axis, + extent: padded_group_extent, + block_size: tiling.block_size, + }); + } + let blocks = padded_group_extent / tiling.block_size; + let short_size = blocks / partitions_per_group; + let long_shards = blocks % partitions_per_group; + let start_blocks = coordinate_in_group * short_size + coordinate_in_group.min(long_shards); + let shard_blocks = short_size + u32::from(coordinate_in_group < long_shards); + let start_in_group = start_blocks + .checked_mul(tiling.block_size) + .ok_or(LayoutError::ExtentOverflow(axis))?; + let allocated = shard_blocks + .checked_mul(tiling.block_size) + .ok_or(LayoutError::ExtentOverflow(axis))?; + let remainder = allocated % tiling.shard_padding_multiple; + if remainder != 0 && tiling.padding == Padding::Reject { + return Err(LayoutError::IndivisibleShard { + extent: allocated, + block_size: tiling.shard_padding_multiple, + }); + } + let physical_width = if remainder == 0 { + allocated + } else { + allocated + .checked_add(tiling.shard_padding_multiple - remainder) + .ok_or(LayoutError::ExtentOverflow(axis))? + }; + let group_logical_base = group + .checked_mul(logical_group_extent) + .ok_or(LayoutError::ExtentOverflow(axis))?; + let start = group_logical_base + .checked_add(start_in_group) + .ok_or(LayoutError::ExtentOverflow(axis))?; + let logical_end = group_logical_base + .checked_add( + start_in_group + .checked_add(allocated) + .ok_or(LayoutError::ExtentOverflow(axis))? + .min(logical_group_extent) + .max(start_in_group), + ) + .ok_or(LayoutError::ExtentOverflow(axis))?; + let physical_end = start + .checked_add(physical_width) + .ok_or(LayoutError::ExtentOverflow(axis))?; + Ok((start, logical_end, physical_end)) +} + +fn validate_tile_mapping( + axes: &[AxisTiling], + replicas: u16, + strides: &[u32], + tile_count: u16, +) -> Result<(), LayoutError> { + if has_regular_tile_mapping(axes, replicas, strides) { + return Ok(()); + } + let coordinate_count = axes + .iter() + .try_fold(1usize, |count, axis| { + count.checked_mul(usize::from(axis.partitions)) + }) + .ok_or(LayoutError::TileCountOverflow)?; + let mut coordinate_copies = vec![0u16; coordinate_count]; + for tile in 0..tile_count { + let coordinate = axes + .iter() + .zip(strides) + .try_fold(0usize, |coordinate, (axis, stride)| { + coordinate + .checked_mul(usize::from(axis.partitions)) + .and_then(|coordinate| { + coordinate.checked_add( + ((u32::from(tile) / stride) % u32::from(axis.partitions)) as usize, + ) + }) + }) + .ok_or(LayoutError::TileCountOverflow)?; + coordinate_copies[coordinate] = coordinate_copies[coordinate] + .checked_add(1) + .ok_or(LayoutError::TileCountOverflow)?; + } + if coordinate_copies.iter().any(|copies| *copies != replicas) { + return Err(LayoutError::InvalidTileMapping); + } + Ok(()) +} + +fn has_regular_tile_mapping(axes: &[AxisTiling], replicas: u16, strides: &[u32]) -> bool { + let mut digits = axes + .iter() + .zip(strides) + .filter(|(axis, _)| axis.partitions > 1) + .map(|(axis, &stride)| (stride, u32::from(axis.partitions))) + .collect::>(); + digits.sort_unstable(); + let Some(&(base, _)) = digits.first() else { + return true; + }; + if base == 0 || !u32::from(replicas).is_multiple_of(base) { + return false; + } + let mut expected_stride = base; + for (stride, partitions) in digits { + if stride != expected_stride { + return false; + } + let Some(next) = expected_stride.checked_mul(partitions) else { + return false; + }; + expected_stride = next; + } + true +} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::BTreeSet; + + const RANDOM_CASES: usize = 128; + + #[test] + fn randomized_axis_layout_resolution_matches_materialized_ownership() { + let mut random = fastrand::Rng::with_seed(0x7265_736f_6c76_6564); + for case in 0..RANDOM_CASES { + let rank = random.usize(2..=5); + let first_axis = random.usize(0..rank); + let second_axis = (first_axis + random.usize(1..rank)) % rank; + let first_partitions = random.u16(1..=8); + let second_partitions = random.u16(1..=8); + let replicas = random.u16(1..=3); + let first_block = random.u32(1..=16); + let second_block = random.u32(1..=16); + let mut shape = (0..rank).map(|_| random.u32(1..=97)).collect::>(); + shape[first_axis] = random.u32(1..=97); + shape[second_axis] = random.u32(1..=97); + let layout = Layout { + order: crate::layout::StorageOrder::Linear, + tiling: TensorTiling { + tile_count: first_partitions * second_partitions * replicas, + replicas, + axes: vec![ + AxisTiling::new( + TensorAxis::FromStart(first_axis as u16), + first_partitions, + first_block, + Padding::Zero, + ) + .with_shard_padding_multiple(random.u32(1..=16)), + AxisTiling::new( + TensorAxis::FromStart(second_axis as u16), + second_partitions, + second_block, + Padding::Zero, + ) + .with_shard_padding_multiple(random.u32(1..=16)), + ], + }, + memory_class: MemoryClass::Standard, + }; + let resolved = layout.resolve(&TensorShape(shape.clone())).unwrap(); + let shards = resolved.shard_extents(); + assert_eq!( + shards.len(), + usize::from(resolved.tile_count()), + "case {case}" + ); + + let tile_elements = shards + .iter() + .map(|shard| { + shard.extents.iter().fold(1_u64, |elements, extent| { + elements * u64::from(extent.physical_end - extent.start) + }) + }) + .collect::>(); + assert_eq!( + tile_elements.iter().sum::(), + resolved.total_elements(), + "case {case}" + ); + assert_eq!( + tile_elements.iter().copied().max().unwrap(), + resolved.maximum_tile_elements(), + "case {case}" + ); + for (tile, &elements) in tile_elements.iter().enumerate() { + assert_eq!( + resolved.tile_elements(tile as u16), + Some(elements), + "case {case}" + ); + } + + for axis in [first_axis, second_axis] { + let ranges = shards + .iter() + .map(|shard| { + let extent = shard.extents[axis]; + (extent.start, extent.logical_end) + }) + .filter(|(start, end)| start < end) + .collect::>(); + let mut cursor = 0; + for (start, end) in ranges { + assert_eq!(start, cursor, "case {case}, axis {axis}"); + cursor = end; + } + assert_eq!(cursor, shape[axis], "case {case}, axis {axis}"); + } + } + } + + #[test] + fn randomized_linear_resolution_covers_each_logical_element_once() { + let mut random = fastrand::Rng::with_seed(0x6c69_6e65_6172_697a); + for case in 0..RANDOM_CASES { + let rank = random.usize(1..=5); + let shape = TensorShape((0..rank).map(|_| random.u32(1..=8)).collect::>()); + let elements = shape.elements(); + let divisors = (1..=elements.min(32)) + .filter(|grain| elements.is_multiple_of(*grain)) + .collect::>(); + let grain = divisors[random.usize(..divisors.len())] as u32; + let grains = elements / u64::from(grain); + let tile_count = random.u16(1..=u16::try_from(grains.min(32)).unwrap()); + let layout = Layout::row_major(TensorTiling::linear(tile_count, grain)); + let resolved = layout.resolve(&shape).unwrap(); + let shards = resolved.shard_extents(); + let materialized = shards.iter().flat_map(|shard| &shard.extents).count(); + assert!(materialized >= usize::from(tile_count), "case {case}"); + assert_eq!( + shards + .iter() + .map(|shard| shard.extents.iter().fold(1_u64, |count, extent| { + count * u64::from(extent.logical_end - extent.start) + })) + .sum::(), + elements, + "case {case}" + ); + assert_eq!(resolved.total_elements(), elements, "case {case}"); + } + } +} * Unmerged path crates/ipu-codegen/src/lib.rs diff --git a/crates/ipu-codegen/src/low/attention.rs b/crates/ipu-codegen/src/low/attention.rs new file mode 100644 index 0000000..6960ba9 --- /dev/null +++ b/crates/ipu-codegen/src/low/attention.rs @@ -0,0 +1,1130 @@ +use super::*; + +struct AttentionLowering { + key_shards: Vec, + value_shards: Vec, + key_rows: u32, + tasks: Vec, + prepared: Vec, + exchange_provenance: WorkProvenance, + kernel_provenance: WorkProvenance, +} + +impl LoweringState { + fn build_attention_tasks( + &mut self, + query: MidValueId, + result: MidValueId, + shape: AttentionBufferShape, + ) -> LowLoweringResult> { + let value_row_block = u16::try_from(shape.physical_staging_rows) + .map_err(|_| LowLoweringError::InvalidOperatorPlan)?; + let outputs = self.value_shards(result)?.to_vec(); + let mut tasks = Vec::with_capacity(outputs.len()); + for output in outputs { + let tile = self.shards[output.index() as usize].tile; + let rank = self.shards[output.index() as usize].extents.len(); + if rank != 3 { + return Err(LowLoweringError::InvalidOperatorPlan); + } + let rows = self.shards[output.index() as usize].extents[rank - 2].physical_end + - self.shards[output.index() as usize].extents[rank - 2].start; + let value_dimension = *self.shards[output.index() as usize] + .tensor_type + .shape + .0 + .last() + .ok_or(LowLoweringError::InvalidOperatorPlan)?; + if rows == 0 || rows > shape.query_block_rows { + return Err(LowLoweringError::InvalidOperatorPlan); + } + let canonical_query = self.local_shard(query, tile)?; + let query_dimension = *self.shards[canonical_query.index() as usize] + .tensor_type + .shape + .0 + .last() + .ok_or(LowLoweringError::InvalidOperatorPlan)?; + let deferred_query = self.deferred_view(query).is_some(); + let query_shard = if deferred_query { + self.push_matrix_buffer( + tile, + rows, + rows, + query_dimension, + shape.padded_query_dimension, + StorageOrder::Native(NativeKernelOrder::Left), + )? + } else { + canonical_query + }; + let query_region = TensorRegion::logical_bounds([ + ( + self.shards[output.index() as usize].extents[rank - 3].start, + self.shards[output.index() as usize].extents[rank - 3].start + 1, + ), + ( + self.shards[output.index() as usize].extents[rank - 2].start, + self.shards[output.index() as usize].extents[rank - 2].start + rows, + ), + (0, query_dimension), + ]) + .ok_or(LowLoweringError::IdOverflow)?; + let direct_query = deferred_query + && (self.deferred_supports_physical_exchange(query, query_shard) + || self.mappings_benefit_from_word_exchange( + &self.deferred_region_mappings(query, &query_region, query_shard)?, + query_shard, + )?); + let query_receive = (deferred_query && !direct_query) + .then(|| { + self.push_matrix_buffer( + tile, + rows, + rows, + query_dimension, + query_dimension, + StorageOrder::Linear, + ) + }) + .transpose()?; + let scratch = self.push_attention_scratch( + tile, + rows, + shape.scratch_columns, + Precision::F16, + StorageOrder::Native(NativeKernelOrder::Left), + MemoryClass::Interleaved, + )?; + let key_staging = self.push_matrix_buffer( + tile, + shape.logical_staging_rows, + shape.physical_staging_rows, + query_dimension, + shape.padded_query_dimension, + StorageOrder::Native(NativeKernelOrder::TransposedRight), + )?; + self.shards[key_staging.index() as usize].definition = ShardDefinition::ExchangeStaging; + let weights = self.push_attention_scratch( + tile, + rows, + shape.state_columns, + Precision::F16, + StorageOrder::Native(NativeKernelOrder::Left), + MemoryClass::Standard, + )?; + if shape.reuse_key_staging_for_state + && crate::shard_storage_bytes(&self.shards[weights.index() as usize])? + <= crate::shard_storage_bytes(&self.shards[key_staging.index() as usize])? + { + // Materialized QK consumes the packed K matrix before softmax + // starts. Reinterpret that now-dead standard-memory allocation + // as probabilities plus row state so the PV kernel retains its + // proven standard-load path without a second large buffer. + self.shards[weights.index() as usize].definition = + ShardDefinition::Alias(key_staging); + } + let value_staging = self.push_matrix_buffer( + tile, + shape.logical_staging_rows, + shape.physical_staging_rows, + value_dimension, + shape.padded_value_dimension, + StorageOrder::Blocked(BlockedOrder::matrix( + value_row_block, + AMP_COLUMN_MICRO as u16, + )), + )?; + self.shards[value_staging.index() as usize].definition = + ShardDefinition::ExchangeStaging; + tasks.push(AttentionTask { + tile, + head: self.shards[output.index() as usize].extents[rank - 3].start, + query_row_start: self.shards[output.index() as usize].extents[rank - 2].start, + query_rows: rows, + query_dimension, + value_dimension, + query: query_shard, + query_receive, + output, + scratch, + weights, + key_staging, + value_staging, + }); + } + Ok(tasks) + } + + fn materialize_attention_queries( + &mut self, + query: MidValueId, + tasks: &[AttentionTask], + provenance: WorkProvenance, + tiles: &mut [TileWorkList], + ) -> LowLoweringResult<()> { + if self.deferred_view(query).is_none() { + return Ok(()); + } + let mut transfers = BTreeMap::>::new(); + let mut local_copies = Vec::new(); + let physical = tasks.iter().all(|task| task.query_receive.is_none()); + let physical = physical + && tasks + .iter() + .all(|task| self.deferred_supports_physical_exchange(query, task.query)); + let direct = tasks.iter().all(|task| task.query_receive.is_none()); + if direct && !physical { + for task in tasks { + if self.shard_has_padding(task.query) { + self.append_fill_zero(tiles, task.query, provenance.clone())?; + } + } + } + for task in tasks { + let region = TensorRegion::logical_bounds([ + (task.head, task.head + 1), + (task.query_row_start, task.query_row_start + task.query_rows), + (0, task.query_dimension), + ]) + .ok_or(LowLoweringError::IdOverflow)?; + self.materialize_deferred_region( + query, + ®ion, + task.query_receive.unwrap_or(task.query), + if physical { + ExchangeOrder::Physical + } else { + ExchangeOrder::Semantic + }, + &mut transfers, + &mut local_copies, + )?; + } + for (tile, copy) in local_copies { + self.append_local_copy(tiles, tile, copy)?; + } + let order = if physical { + ExchangeOrder::Physical + } else { + ExchangeOrder::Semantic + }; + self.append_phase( + transfers, + provenance, + |source| (source, order.clone()), + tiles, + )?; + if !physical { + for task in tasks { + self.append_attention_rearrange( + tiles, + task.tile, + task.query_receive + .ok_or(LowLoweringError::InvalidOperatorPlan)?, + task.query, + WorkProvenance { + operation: provenance.operation, + value: provenance.value, + reason: WorkReason::OperatorKernel, + }, + )?; + } + } + Ok(()) + } + + fn prepare_attention_blocks( + &mut self, + key: MidValueId, + value: MidValueId, + tasks: &[AttentionTask], + key_rows: u32, + shape: AttentionBufferShape, + provenance: WorkProvenance, + tiles: &mut [TileWorkList], + ) -> LowLoweringResult> { + let block_rows = shape.panel_rows; + let padded_query_dimension = shape.padded_query_dimension; + let padded_value_dimension = shape.padded_value_dimension; + let blocks = key_rows.div_ceil(block_rows); + let key_destinations = tasks.iter().fold( + BTreeMap::>::new(), + |mut destinations, task| { + destinations + .entry(task.head) + .or_default() + .push(task.key_staging); + destinations + }, + ); + let value_destinations = tasks.iter().fold( + BTreeMap::>::new(), + |mut destinations, task| { + destinations + .entry(task.head) + .or_default() + .push(task.value_staging); + destinations + }, + ); + let key_panel_count = padded_query_dimension.div_ceil(AMP_COLUMN_MICRO); + let value_panel_count = padded_value_dimension.div_ceil(AMP_COLUMN_MICRO); + let mut semantic_gathers = BTreeMap::>::new(); + let mut physical_gathers = BTreeMap::>::new(); + let mut prepared = Vec::new(); + for block in 0..blocks { + let row_start = block * block_rows; + let valid_rows = key_rows.saturating_sub(row_start).min(block_rows); + let owner_offset = block.saturating_mul(key_panel_count + value_panel_count); + let key_panels = self.prepare_distributed_attention_panels( + key, + &key_destinations, + row_start, + valid_rows, + tasks[0].query_dimension, + padded_query_dimension, + StorageOrder::Native(NativeKernelOrder::TransposedRight), + owner_offset, + &mut semantic_gathers, + &mut physical_gathers, + provenance.clone(), + tiles, + )?; + let row_block = + u16::try_from(block_rows).map_err(|_| LowLoweringError::InvalidOperatorPlan)?; + let value_panels = self.prepare_distributed_attention_panels( + value, + &value_destinations, + row_start, + valid_rows, + tasks[0].value_dimension, + padded_value_dimension, + StorageOrder::Blocked(BlockedOrder::matrix(row_block, AMP_COLUMN_MICRO as u16)), + owner_offset + key_panel_count, + &mut semantic_gathers, + &mut physical_gathers, + provenance.clone(), + tiles, + )?; + prepared.push(PreparedAttentionBlock { + row_start, + valid_rows, + key_panels, + value_panels, + }); + } + let gathers = + semantic_gathers + .into_iter() + .map(|(source, destinations)| ((source, ExchangeOrder::Semantic), destinations)) + .chain(physical_gathers.into_iter().map(|(source, destinations)| { + ((source, ExchangeOrder::Physical), destinations) + })); + self.append_phase(gathers, provenance, |ordered| ordered, tiles)?; + for block in &prepared { + for panel in block.key_panels.iter().chain(&block.value_panels) { + if let Some(row_major) = panel.row_major { + self.append_attention_rearrange( + tiles, + panel.tile, + row_major, + panel.packed, + WorkProvenance { + operation: provenance.operation, + value: provenance.value, + reason: WorkReason::OperatorKernel, + }, + )?; + } + } + } + Ok(prepared) + } + + fn prepare_attention( + &mut self, + operation: &MidOperation, + plan: &crate::AttentionMap, + tiles: &mut [TileWorkList], + ) -> LowLoweringResult { + let [query, key, value] = operation.inputs.as_slice() else { + return Err(LowLoweringError::InvalidOperatorPlan); + }; + let [result] = operation.results.as_slice() else { + return Err(LowLoweringError::ResultArity); + }; + let key_shards = self.value_shards(*key)?.to_vec(); + let value_shards = self.value_shards(*value)?.to_vec(); + let deferred = self.deferred_view(*key).is_some(); + if key_shards.len() != value_shards.len() + || deferred != self.deferred_view(*value).is_some() + { + return Err(LowLoweringError::InvalidOperatorPlan); + } + let key_rows = self.shards[key_shards[0].index() as usize] + .tensor_type + .shape + .0[1]; + if key_rows == 0 { + return Err(LowLoweringError::InvalidOperatorPlan); + } + if let crate::AttentionBlocking::Materialized { + padded_key_rows, .. + } = plan.blocking + && key_rows > padded_key_rows + { + return Err(LowLoweringError::InvalidOperatorPlan); + } + let shape = AttentionBufferShape::from_plan(plan, key_rows); + let tasks = self.build_attention_tasks(*query, *result, shape)?; + if tasks.is_empty() { + return Err(LowLoweringError::InvalidOperatorPlan); + } + let exchange_provenance = WorkProvenance { + operation: operation.source, + value: Some(*key), + reason: WorkReason::OperatorInputs, + }; + let kernel_provenance = WorkProvenance { + operation: operation.source, + value: Some(*result), + reason: WorkReason::OperatorKernel, + }; + self.materialize_attention_queries(*query, &tasks, exchange_provenance, tiles)?; + let prepared = if deferred { + self.prepare_attention_blocks( + *key, + *value, + &tasks, + key_rows, + shape, + exchange_provenance, + tiles, + )? + } else { + Vec::new() + }; + Ok(AttentionLowering { + key_shards, + value_shards, + key_rows, + tasks, + prepared, + exchange_provenance, + kernel_provenance, + }) + } + + pub(super) fn lower_blocked_attention( + &mut self, + operation: &MidOperation, + plan: &crate::AttentionMap, + requirements: &OperatorRequirements, + tiles: &mut [TileWorkList], + ) -> LowLoweringResult<()> { + let crate::AttentionBlocking::Flash { + query_rows: query_block_rows, + key_rows: key_block_rows, + } = plan.blocking + else { + return Err(LowLoweringError::InvalidOperatorPlan); + }; + let [query_key_block, probability_value_block] = plan.gemm_blocks(); + let padded_value_dimension = plan.value_dimension; + if key_block_rows != AMP_INNER_BLOCK || query_block_rows == 0 { + return Err(LowLoweringError::InvalidOperatorPlan); + } + let AttentionLowering { + key_shards, + value_shards, + key_rows, + tasks, + prepared: prepared_blocks, + exchange_provenance, + kernel_provenance, + } = self.prepare_attention(operation, plan, tiles)?; + let deferred_key_value = !prepared_blocks.is_empty(); + let blocks = usize::try_from(key_rows.div_ceil(key_block_rows)) + .map_err(|_| LowLoweringError::IdOverflow)?; + if blocks == 0 { + return Err(LowLoweringError::InvalidOperatorPlan); + } + for block in 0..blocks { + let block_start = + u32::try_from(block).map_err(|_| LowLoweringError::IdOverflow)? * key_block_rows; + let mut transfers = BTreeMap::>::new(); + let mut task_sources = Vec::with_capacity(tasks.len()); + if deferred_key_value { + let prepared = prepared_blocks + .get(block) + .ok_or(LowLoweringError::InvalidOperatorPlan)?; + self.append_prepared_panel_broadcasts( + &prepared.key_panels, + 0, + &mut transfers, + tiles, + )?; + self.append_prepared_panel_broadcasts( + &prepared.value_panels, + 0, + &mut transfers, + tiles, + )?; + self.append_phase(transfers, exchange_provenance, semantic_exchange, tiles)?; + task_sources.extend( + tasks + .iter() + .map(|task| (task.key_staging, task.value_staging, prepared.valid_rows)), + ); + } else { + for task in &tasks { + let source_matches = |candidate: &&LowShardId| { + let shard = &self.shards[candidate.index() as usize]; + shard.extents[0].start == task.head && shard.extents[1].start == block_start + }; + let key_source = *key_shards + .iter() + .find(source_matches) + .ok_or(LowLoweringError::InvalidOperatorPlan)?; + let value_source = *value_shards + .iter() + .find(source_matches) + .ok_or(LowLoweringError::InvalidOperatorPlan)?; + let valid_key_rows = self.shards[key_source.index() as usize].extents[1] + .logical_end + .saturating_sub(block_start); + let mut operands = Vec::with_capacity(2); + for (source, destination) in [ + (key_source, task.key_staging), + (value_source, task.value_staging), + ] { + if self.shards[source.index() as usize].tile == task.tile { + operands.push(source); + } else { + transfers + .entry(self.full_view(source)) + .or_default() + .push(self.full_view(destination)); + operands.push(destination); + } + } + task_sources.push((operands[0], operands[1], valid_key_rows)); + } + self.append_phase(transfers, exchange_provenance, semantic_exchange, tiles)?; + } + for (task, (key_operand, value_operand, valid_key_rows)) in + tasks.iter().zip(task_sources) + { + let score_view = self.narrow_view(task.scratch, &[(1, 0, key_block_rows)])?; + self.append_kernel( + tiles, + task.tile, + KernelRun::new( + kernel_provenance, + gemm_kernel_spec( + plan.kernel, + GemmKernelMode::Initialize, + query_key_block, + task.query_rows, + ), + vec![ + KernelOperand { + views: vec![self.full_view(task.query)], + }, + KernelOperand { + views: vec![self.full_view(key_operand)], + }, + ], + score_view.clone(), + KernelRequirements::Operator(requirements.clone()), + ), + )?; + self.append_kernel( + tiles, + task.tile, + KernelRun::new( + kernel_provenance, + TileKernelSpec::AttentionSoftmax { + query_rows: task.query_rows, + head_dimension: task.query_dimension, + key_columns: valid_key_rows, + padded_key_columns: key_block_rows, + }, + vec![KernelOperand { + views: vec![score_view], + }], + self.full_view(task.weights), + KernelRequirements::Operator(requirements.clone()), + ), + )?; + let probability_view = self.narrow_view(task.weights, &[(1, 0, key_block_rows)])?; + let block_value_view = + self.narrow_view(task.scratch, &[(1, 0, padded_value_dimension)])?; + self.append_kernel( + tiles, + task.tile, + KernelRun::new( + kernel_provenance, + gemm_kernel_spec( + plan.kernel, + GemmKernelMode::Initialize, + probability_value_block, + task.query_rows, + ), + vec![ + KernelOperand { + views: vec![probability_view], + }, + KernelOperand { + views: vec![self.full_view(value_operand)], + }, + ], + block_value_view.clone(), + KernelRequirements::Operator(requirements.clone()), + ), + )?; + self.append_kernel( + tiles, + task.tile, + KernelRun::new( + kernel_provenance, + TileKernelSpec::AttentionMerge { + query_rows: task.query_rows, + value_dimension: task.value_dimension, + padded_value_dimension, + key_block_columns: key_block_rows, + initial: block == 0, + final_block: block + 1 == blocks, + }, + vec![ + KernelOperand { + views: vec![block_value_view], + }, + KernelOperand { + views: vec![self.full_view(task.weights)], + }, + ], + self.full_view(task.output), + KernelRequirements::Operator(requirements.clone()), + ), + )?; + } + } + Ok(()) + } + + fn append_materialized_attention_input( + &mut self, + operand: AttentionOperand, + sources: &[LowShardId], + prepared: &[PreparedAttentionBlock], + tasks: &[AttentionTask], + provenance: WorkProvenance, + tiles: &mut [TileWorkList], + ) -> LowLoweringResult<()> { + let mut transfers = BTreeMap::>::new(); + if prepared.is_empty() { + for task in tasks { + let destination = match operand { + AttentionOperand::Key => task.key_staging, + AttentionOperand::Value => task.value_staging, + }; + let matching_sources = sources + .iter() + .copied() + .filter(|source| { + self.shards[source.index() as usize].extents[0].start == task.head + }) + .collect::>(); + for source in matching_sources { + let source_view = self.full_view(source); + let row_extent = self.shards[source.index() as usize].extents[1]; + let destination_view = self.narrow_view( + destination, + &[(0, row_extent.start, row_extent.physical_end)], + )?; + let source_tile = self.shards[source.index() as usize].tile; + if source_tile == task.tile { + let mut copies = Vec::new(); + append_span_copies( + &self.shards, + &source_view, + &destination_view, + task.tile, + &mut copies, + )?; + for (tile, copy) in copies { + self.append_local_copy(tiles, tile, copy)?; + } + } else { + transfers + .entry(source_view) + .or_default() + .push(destination_view); + } + } + } + } else { + for block in prepared { + let panels = match operand { + AttentionOperand::Key => &block.key_panels, + AttentionOperand::Value => &block.value_panels, + }; + self.append_prepared_panel_broadcasts( + panels, + block.row_start, + &mut transfers, + tiles, + )?; + } + } + self.append_phase(transfers, provenance, physical_exchange, tiles) + } + + pub(super) fn lower_materialized_attention( + &mut self, + operation: &MidOperation, + plan: &crate::AttentionMap, + requirements: &OperatorRequirements, + tiles: &mut [TileWorkList], + ) -> LowLoweringResult<()> { + let crate::AttentionBlocking::Materialized { + query_rows: query_block_rows, + padded_key_rows, + } = plan.blocking + else { + return Err(LowLoweringError::InvalidOperatorPlan); + }; + let [query_key_block, probability_value_block] = plan.gemm_blocks(); + let padded_value_dimension = plan.value_dimension; + if query_block_rows == 0 + || padded_key_rows == 0 + || !padded_key_rows.is_multiple_of(AMP_INNER_BLOCK) + { + return Err(LowLoweringError::InvalidOperatorPlan); + } + let AttentionLowering { + key_shards, + value_shards, + key_rows, + tasks, + prepared, + exchange_provenance, + kernel_provenance, + .. + } = self.prepare_attention(operation, plan, tiles)?; + self.append_materialized_attention_input( + AttentionOperand::Key, + &key_shards, + &prepared, + &tasks, + exchange_provenance, + tiles, + )?; + for task in &tasks { + let scores = self.narrow_view(task.scratch, &[(1, 0, padded_key_rows)])?; + self.append_kernel( + tiles, + task.tile, + KernelRun::new( + kernel_provenance, + gemm_kernel_spec( + plan.kernel, + GemmKernelMode::Initialize, + query_key_block, + task.query_rows, + ), + vec![ + KernelOperand { + views: vec![self.full_view(task.query)], + }, + KernelOperand { + views: vec![self.full_view(task.key_staging)], + }, + ], + scores.clone(), + KernelRequirements::Operator(requirements.clone()), + ), + )?; + self.append_kernel( + tiles, + task.tile, + KernelRun::new( + kernel_provenance, + TileKernelSpec::AttentionSoftmax { + query_rows: task.query_rows, + head_dimension: task.query_dimension, + key_columns: key_rows, + padded_key_columns: padded_key_rows, + }, + vec![KernelOperand { + views: vec![scores], + }], + self.full_view(task.weights), + KernelRequirements::Operator(requirements.clone()), + ), + )?; + } + self.append_materialized_attention_input( + AttentionOperand::Value, + &value_shards, + &prepared, + &tasks, + exchange_provenance, + tiles, + )?; + for task in &tasks { + let probabilities = self.narrow_view(task.weights, &[(1, 0, padded_key_rows)])?; + let block_value = self.narrow_view(task.scratch, &[(1, 0, padded_value_dimension)])?; + self.append_kernel( + tiles, + task.tile, + KernelRun::new( + kernel_provenance, + gemm_kernel_spec( + plan.kernel, + GemmKernelMode::Initialize, + probability_value_block, + task.query_rows, + ), + vec![ + KernelOperand { + views: vec![probabilities], + }, + KernelOperand { + views: vec![self.full_view(task.value_staging)], + }, + ], + block_value.clone(), + KernelRequirements::Operator(requirements.clone()), + ), + )?; + self.append_kernel( + tiles, + task.tile, + KernelRun::new( + kernel_provenance, + TileKernelSpec::AttentionMerge { + query_rows: task.query_rows, + value_dimension: task.value_dimension, + padded_value_dimension, + key_block_columns: padded_key_rows, + initial: true, + final_block: true, + }, + vec![ + KernelOperand { + views: vec![block_value], + }, + KernelOperand { + views: vec![self.full_view(task.weights)], + }, + ], + self.full_view(task.output), + KernelRequirements::Operator(requirements.clone()), + ), + )?; + } + Ok(()) + } + + fn prepare_distributed_attention_panels( + &mut self, + value: MidValueId, + destinations: &BTreeMap>, + block_start: u32, + valid_rows: u32, + logical_columns: u32, + physical_columns: u32, + order: StorageOrder, + owner_offset: u32, + semantic_gathers: &mut BTreeMap>, + physical_gathers: &mut BTreeMap>, + provenance: WorkProvenance, + tiles: &mut [TileWorkList], + ) -> LowLoweringResult> { + let panels = physical_columns.div_ceil(AMP_COLUMN_MICRO); + if panels == 0 || valid_rows == 0 { + return Err(LowLoweringError::InvalidOperatorPlan); + } + let mut packed_panels = Vec::new(); + for (&stream, stream_destinations) in destinations { + for panel in 0..panels { + let column_start = panel * AMP_COLUMN_MICRO; + let panel_columns = logical_columns + .saturating_sub(column_start) + .min(AMP_COLUMN_MICRO); + if panel_columns == 0 { + continue; + } + let owner = usize::try_from(owner_offset.saturating_add(panel)) + .map_err(|_| LowLoweringError::IdOverflow)? + % stream_destinations.len(); + let tile = stream_destinations[owner]; + let tile = self.shards[tile.index() as usize].tile; + let packed = self.push_matrix_buffer( + tile, + valid_rows, + AMP_INNER_BLOCK, + panel_columns, + AMP_COLUMN_MICRO, + order, + )?; + let region = TensorRegion::logical_bounds([ + (stream, stream + 1), + (block_start, block_start + valid_rows), + (column_start, column_start + panel_columns), + ]) + .ok_or(LowLoweringError::IdOverflow)?; + let physical = self.deferred_supports_physical_exchange(value, packed); + let word_exchange = !physical + && self.mappings_benefit_from_word_exchange( + &self.deferred_region_mappings(value, ®ion, packed)?, + packed, + )?; + if word_exchange && self.shard_has_padding(packed) { + self.append_fill_zero(tiles, packed, provenance.clone())?; + } + let row_major = if physical || word_exchange { + None + } else { + Some(self.push_matrix_buffer( + tile, + valid_rows, + valid_rows, + panel_columns, + panel_columns, + StorageOrder::Linear, + )?) + }; + let gather_destination = row_major.unwrap_or(packed); + let mut local_copies = Vec::new(); + self.materialize_deferred_region( + value, + ®ion, + gather_destination, + if physical { + ExchangeOrder::Physical + } else { + ExchangeOrder::Semantic + }, + if physical { + physical_gathers + } else { + semantic_gathers + }, + &mut local_copies, + )?; + for (tile, copy) in local_copies { + self.append_local_copy(tiles, tile, copy)?; + } + packed_panels.push(PreparedDistributedPanel { + panel, + row_major, + packed, + tile, + destinations: stream_destinations.clone(), + }); + } + } + Ok(packed_panels) + } + + fn append_prepared_panel_broadcasts( + &mut self, + panels: &[PreparedDistributedPanel], + destination_row_start: u32, + broadcasts: &mut BTreeMap>, + tiles: &mut [TileWorkList], + ) -> LowLoweringResult<()> { + for panel in panels { + let source = self.full_view(panel.packed); + let source_rows = source.extents[0].physical_end - source.extents[0].start; + let column_start = panel.panel * AMP_COLUMN_MICRO; + for &destination in &panel.destinations { + let destination_tile = self.shards[destination.index() as usize].tile; + let destination_view = self.narrow_view( + destination, + &[ + ( + 0, + destination_row_start, + destination_row_start + source_rows, + ), + (1, column_start, column_start + AMP_COLUMN_MICRO), + ], + )?; + if panel.tile == destination_tile { + let mut copies = Vec::new(); + append_span_copies( + &self.shards, + &source, + &destination_view, + panel.tile, + &mut copies, + )?; + for (tile, copy) in copies { + self.append_local_copy(tiles, tile, copy)?; + } + } else { + broadcasts + .entry(source.clone()) + .or_insert_with(Vec::new) + .push(destination_view); + } + } + } + Ok(()) + } + + fn append_attention_rearrange( + &mut self, + tiles: &mut [TileWorkList], + tile: u16, + source: LowShardId, + destination: LowShardId, + provenance: WorkProvenance, + ) -> LowLoweringResult<()> { + let input = self.shards[source.index() as usize] + .tensor_type + .format + .clone(); + let output = self.shards[destination.index() as usize] + .tensor_type + .format + .clone(); + self.append_kernel( + tiles, + tile, + KernelRun::new( + provenance, + rearrange_kernel_spec( + input.layout.clone(), + output.layout.clone(), + &self.full_view(source), + &self.full_view(destination), + )?, + vec![KernelOperand { + views: vec![self.full_view(source)], + }], + self.full_view(destination), + KernelRequirements::Conversion { + input: OperandRequirement::new(input, 2), + output: OperandRequirement::new(output, 2), + memory_space: MemorySpaceRequirements::default(), + }, + ), + ) + } + + fn push_attention_scratch( + &mut self, + tile: u16, + rows: u32, + columns: u32, + precision: Precision, + order: StorageOrder, + memory_class: MemoryClass, + ) -> LowLoweringResult { + self.push_shard(LowShard { + id: LowShardId(0), + tile, + tensor_type: TensorType::new( + [rows, columns], + precision, + Layout { + order: order.clone(), + tiling: TensorTiling::replicated(1), + memory_class, + }, + ), + extents: vec![ + ShardExtent { + axis: 0, + start: 0, + logical_end: rows, + physical_end: rows, + }, + ShardExtent { + axis: 1, + start: 0, + logical_end: columns, + physical_end: columns, + }, + ] + .into(), + definition: ShardDefinition::Staging, + }) + } + + pub(super) fn push_packed_buffer( + &mut self, + tile: u16, + elements: u32, + precision: Precision, + definition: ShardDefinition, + ) -> LowLoweringResult { + self.push_shard(LowShard { + id: LowShardId(0), + tile, + tensor_type: TensorType::new( + [elements], + precision, + Layout { + order: StorageOrder::Linear, + tiling: TensorTiling::replicated(1), + memory_class: MemoryClass::Standard, + }, + ), + extents: vec![ShardExtent { + axis: 0, + start: 0, + logical_end: elements, + physical_end: elements, + }] + .into(), + definition, + }) + } + + pub(super) fn push_matrix_buffer( + &mut self, + tile: u16, + logical_rows: u32, + physical_rows: u32, + logical_columns: u32, + physical_columns: u32, + order: StorageOrder, + ) -> LowLoweringResult { + self.push_shard(LowShard { + id: LowShardId(0), + tile, + tensor_type: TensorType::new( + [logical_rows, logical_columns], + Precision::F16, + Layout { + order, + tiling: TensorTiling::replicated(1), + memory_class: MemoryClass::Standard, + }, + ), + extents: vec![ + ShardExtent { + axis: 0, + start: 0, + logical_end: logical_rows, + physical_end: physical_rows, + }, + ShardExtent { + axis: 1, + start: 0, + logical_end: logical_columns, + physical_end: physical_columns, + }, + ] + .into(), + definition: ShardDefinition::Staging, + }) + } +} diff --git a/crates/ipu-codegen/src/low/gemm.rs b/crates/ipu-codegen/src/low/gemm.rs new file mode 100644 index 0000000..942469d --- /dev/null +++ b/crates/ipu-codegen/src/low/gemm.rs @@ -0,0 +1,1000 @@ +use super::*; + +struct GemmLowering { + left_value: MidValueId, + right_value: MidValueId, + output_value: MidValueId, + left_shards: Vec, + right_shards: Vec, + output_shards: Vec, + left_rank: usize, + right_rank: usize, + output_rank: usize, + block: crate::GemmBlockShape, + kernel: crate::GemmKernelFamily, +} + +impl GemmLowering { + fn bind( + state: &LoweringState, + operation: &MidOperation, + plan: &crate::GemmMap, + ) -> LowLoweringResult { + let [left_value, right_value] = operation.inputs.as_slice() else { + return Err(LowLoweringError::InvalidOperatorPlan); + }; + let [output_value] = operation.results.as_slice() else { + return Err(LowLoweringError::ResultArity); + }; + let inner_block = plan.geometry.block.inner; + let output_column_block = plan.geometry.block.output_columns; + if inner_block == 0 || output_column_block == 0 { + return Err(LowLoweringError::InvalidOperatorPlan); + } + let left_shards = state.value_shards(*left_value)?.to_vec(); + let right_shards = state.value_shards(*right_value)?.to_vec(); + let output_shards = state.value_shards(*output_value)?.to_vec(); + let left_rank = state.shards[left_shards[0].index() as usize].extents.len(); + let right_rank = state.shards[right_shards[0].index() as usize].extents.len(); + let output_rank = state.shards[output_shards[0].index() as usize] + .extents + .len(); + if left_rank < 2 || right_rank < 2 || output_rank < 2 { + return Err(LowLoweringError::InvalidOperatorPlan); + } + Ok(GemmLowering { + left_value: *left_value, + right_value: *right_value, + output_value: *output_value, + left_shards, + right_shards, + output_shards, + left_rank, + right_rank, + output_rank, + block: plan.geometry.block, + kernel: plan.kernel, + }) + } +} + +impl LoweringState { + pub(super) fn lower_blocked_gemm( + &mut self, + operation: &MidOperation, + plan: &crate::GemmMap, + reduction_staging: Option, + requirements: &OperatorRequirements, + tiles: &mut [TileWorkList], + ) -> LowLoweringResult<()> { + if plan.inputs != [ScheduleValue::Input(0), ScheduleValue::Input(1)] + || (plan.geometry.compute.inner > 1 && plan.output != ScheduleValue::Temporary(0)) + || (plan.geometry.compute.inner == 1 && plan.output != ScheduleValue::Output) + { + return Err(LowLoweringError::InvalidOperatorPlan); + } + let gemm = GemmLowering::bind(self, operation, plan)?; + if plan.geometry.compute.inner > 1 { + return self.lower_parallel_reduction_gemm( + operation, + plan, + reduction_staging.ok_or(LowLoweringError::InvalidOperatorPlan)?, + requirements, + gemm, + tiles, + ); + } + if reduction_staging.is_some() { + return Err(LowLoweringError::InvalidOperatorPlan); + } + if plan.geometry.orientation != crate::GemmOrientation::Normal { + return Err(LowLoweringError::InvalidOperatorPlan); + } + let left_type = &self.shards[gemm.left_shards[0].index() as usize].tensor_type; + let output_type = &self.shards[gemm.output_shards[0].index() as usize].tensor_type; + let inner_extent = left_type + .format + .layout + .resolve(&left_type.shape)? + .padded_shape() + .0[gemm.left_rank - 1]; + let column_extent = output_type + .format + .layout + .resolve(&output_type.shape)? + .padded_shape() + .0[gemm.output_rank - 1]; + if !inner_extent.is_multiple_of(gemm.block.inner) + || !column_extent.is_multiple_of(gemm.block.output_columns) + { + return Err(LowLoweringError::InvalidOperatorPlan); + } + self.lower_output_stationary_gemm( + operation, + gemm, + [inner_extent, column_extent], + requirements, + tiles, + ) + } + + fn lower_parallel_reduction_gemm( + &mut self, + operation: &MidOperation, + plan: &crate::GemmMap, + reduction_staging: ReductionStaging, + requirements: &OperatorRequirements, + gemm: GemmLowering, + tiles: &mut [TileWorkList], + ) -> LowLoweringResult<()> { + let (inner_block, output_column_block) = (gemm.block.inner, gemm.block.output_columns); + let orientation = plan.geometry.orientation; + let kernel_family = gemm.kernel; + let row_partitions = plan.geometry.compute.rows; + let column_partitions = plan.geometry.compute.columns; + let inner_partitions = plan.geometry.compute.inner; + let result_row_partitions = plan + .geometry + .result + .rows + .checked_div(row_partitions) + .unwrap_or(0); + let result_column_partitions = plan + .geometry + .result + .columns + .checked_div(column_partitions) + .unwrap_or(0); + let semantic_left_value = gemm.left_value; + let semantic_right_value = gemm.right_value; + let output_value = gemm.output_value; + let semantic_left_shards = gemm.left_shards; + let semantic_right_shards = gemm.right_shards; + let output_shards = gemm.output_shards; + let semantic_left_rank = gemm.left_rank; + let semantic_right_rank = gemm.right_rank; + let output_rank = gemm.output_rank; + if row_partitions == 0 + || column_partitions == 0 + || inner_partitions < 2 + || result_row_partitions == 0 + || result_column_partitions == 0 + || result_row_partitions.saturating_mul(result_column_partitions) > inner_partitions + { + return Err(LowLoweringError::InvalidOperatorPlan); + } + let [left_value, right_value] = + orientation.physical_order([&semantic_left_value, &semantic_right_value]); + let [left_shards, right_shards] = + orientation.physical_order([semantic_left_shards, semantic_right_shards]); + let [left_rank, right_rank] = + orientation.physical_order([semantic_left_rank, semantic_right_rank]); + let [left_requirement, right_requirement] = + orientation.physical_order([&requirements.inputs[0], &requirements.inputs[1]]); + let mut kernel_requirements = requirements.clone(); + if orientation == crate::GemmOrientation::Swapped { + kernel_requirements.inputs.swap(0, 1); + } + let output_value = &output_value; + let left_row_axis = orientation.row_axis().resolve(left_rank)?; + let left_inner_axis = orientation.column_axis().resolve(left_rank)?; + let right_inner_axis = orientation.row_axis().resolve(right_rank)?; + let right_column_axis = orientation.column_axis().resolve(right_rank)?; + let output_row_axis = orientation.row_axis().resolve(output_rank)?; + let output_column_axis = orientation.column_axis().resolve(output_rank)?; + let output_type = self.shards[output_shards[0].index() as usize] + .tensor_type + .clone(); + let partial_type = plan + .partial_tensor(&output_type) + .ok_or(LowLoweringError::InvalidOperatorPlan)?; + let partial_layout = partial_type.format.layout.resolve(&partial_type.shape)?; + let columns = (0..u32::from(column_partitions)) + .map(|partition| { + partial_layout + .axis_bounds(output_column_axis, partition) + .ok_or(LowLoweringError::InvalidOperatorPlan) + }) + .collect::>>()?; + + let mut replica_groups = BTreeMap::>::new(); + for left in left_shards.iter().copied() { + let key = self.shards[left.index() as usize].extents.physical(); + replica_groups.entry(key).or_default().push(left); + } + let mut replica_columns = BTreeMap::::new(); + for replicas in replica_groups.values_mut() { + replicas.sort_unstable_by_key(|shard| self.shards[shard.index() as usize].tile); + if replicas.len() != usize::from(column_partitions) { + return Err(LowLoweringError::InvalidOperatorPlan); + } + for (column, shard) in replicas.iter().copied().enumerate() { + replica_columns.insert( + shard, + u16::try_from(column).map_err(|_| LowLoweringError::IdOverflow)?, + ); + } + } + { + let mut transfers = BTreeMap::>::new(); + let mut local_copies = Vec::<(u16, LocalCopy)>::new(); + let mut gemm_runs = Vec::<(u16, KernelRun)>::new(); + let mut partials = BTreeMap::>::new(); + let mut resident_lefts = BTreeMap::::new(); + let mut weight_staging = BTreeMap::<(u16, LowShardId), LowShardId>::new(); + for (output_column, &(column_start, logical_column_end, column_end)) in + columns.iter().enumerate() + { + let output_column = + u32::try_from(output_column).map_err(|_| LowLoweringError::IdOverflow)?; + let local_output_columns = column_end - column_start; + if local_output_columns == 0 + || local_output_columns > output_column_block + || !local_output_columns.is_multiple_of(crate::layout::AMP_COLUMN_MICRO) + { + return Err(LowLoweringError::InvalidOperatorPlan); + } + for left in left_shards.iter().copied() { + let left_shard = self.shards[left.index() as usize].clone(); + let resident_left = if let Some(view) = resident_lefts.get(&left) { + view.clone() + } else { + let restrictions = left_shard + .extents + .iter() + .enumerate() + .map(|(axis, extent)| (axis, extent.start, extent.physical_end)) + .collect::>(); + let view = self.schedule_input_view( + *left_value, + left_shard.tile, + &restrictions, + &mut transfers, + &mut local_copies, + )?; + if left_requirement.materialization + != crate::OperandMaterialization::DispatchSlices + && view.shard != left + { + return Err(LowLoweringError::InvalidOperatorPlan); + } + resident_lefts.insert(left, view.clone()); + view + }; + let inner = left_shard.extents[left_inner_axis]; + if !(inner.physical_end - inner.start).is_multiple_of(inner_block) { + return Err(LowLoweringError::InvalidOperatorPlan); + } + if replica_columns.get(&left).copied().map(u32::from) != Some(output_column) { + continue; + } + let partial_layout = partial_type.format.layout.resolve(&partial_type.shape)?; + let padded_output = partial_layout.padded_shape(); + let mut extents = partial_type + .shape + .0 + .iter() + .zip(&padded_output.0) + .enumerate() + .map(|(axis, (&logical_end, &physical_end))| ShardExtent { + axis: u16::try_from(axis).unwrap_or(u16::MAX), + start: 0, + logical_end, + physical_end, + }) + .collect::>(); + if orientation == crate::GemmOrientation::Normal { + for axis in 0..output_rank.saturating_sub(2) { + extents[axis] = left_shard.extents[axis]; + extents[axis].axis = + u16::try_from(axis).map_err(|_| LowLoweringError::IdOverflow)?; + } + } + extents[output_row_axis] = left_shard.extents[left_row_axis]; + extents[output_row_axis].axis = + u16::try_from(output_row_axis).map_err(|_| LowLoweringError::IdOverflow)?; + extents[output_column_axis] = ShardExtent { + axis: u16::try_from(output_column_axis) + .map_err(|_| LowLoweringError::IdOverflow)?, + start: column_start, + logical_end: logical_column_end, + physical_end: column_end, + }; + let partial_key = TensorRegion::new(extents.iter().copied()).physical(); + let direct_output = output_shards.iter().copied().find(|output| { + let shard = &self.shards[output.index() as usize]; + shard.tile == left_shard.tile + && shard.tensor_type.format.layout.order + == partial_type.format.layout.order + && shard.tensor_type.format.layout.memory_class + == partial_type.format.layout.memory_class + && shard.extents.iter().zip(&extents).all(|(owner, partial)| { + owner.start <= partial.start + && owner.logical_end >= partial.logical_end + && owner.physical_end >= partial.physical_end + }) + }); + let partial = if let Some(output) = direct_output { + ShardView { + shard: output, + extents: extents.clone().into(), + } + } else { + let partial = self.push_shard(LowShard { + id: LowShardId(0), + tile: left_shard.tile, + tensor_type: partial_type.clone(), + extents: extents.into(), + definition: ShardDefinition::Staging, + })?; + self.full_view(partial) + }; + partials + .entry(partial_key) + .or_default() + .push((left_shard.tile, partial.clone())); + + let source_panel_block = right_requirement.format.layout.order.clone(); + let source_panel_block = match source_panel_block { + StorageOrder::Blocked(order) => u32::from(order.block_shape[0]), + _ => AMP_INNER_BLOCK, + }; + let first_panel_end = inner + .start + .saturating_add(source_panel_block) + .min(inner.physical_end); + let first_source = self + .matrix_shards_for_block( + &right_shards, + right_column_axis, + right_inner_axis, + column_start, + column_end, + inner.start, + first_panel_end, + ) + .next() + .ok_or(LowLoweringError::InvalidOperatorPlan)?; + let mut weight_type = self.shards[first_source.index() as usize] + .tensor_type + .clone(); + weight_type.format.layout.memory_class = + right_requirement.local_staging.memory_class(); + let mut weight_extents = + self.shards[first_source.index() as usize].extents.clone(); + weight_extents[right_inner_axis].start = inner.start; + weight_extents[right_inner_axis].logical_end = inner.logical_end; + weight_extents[right_inner_axis].physical_end = inner.physical_end; + let source_inner = + self.shards[first_source.index() as usize].extents[right_inner_axis]; + let source_covers_compute_inner = source_inner.start <= inner.start + && source_inner.physical_end >= inner.physical_end; + let stage_local_sources = right_requirement.local_staging.stages_local(); + let weights = if self.shards[first_source.index() as usize].tile + == left_shard.tile + && source_covers_compute_inner + && !stage_local_sources + { + None + } else { + let key = (left_shard.tile, first_source); + if let Some(staging) = weight_staging.get(&key).copied() { + Some(staging) + } else { + let staging = self.push_shard(LowShard { + id: LowShardId(0), + tile: left_shard.tile, + tensor_type: weight_type, + extents: weight_extents, + definition: ShardDefinition::ExchangeStaging, + })?; + weight_staging.insert(key, staging); + Some(staging) + } + }; + + for (block_index, inner_start) in (inner.start..inner.physical_end) + .step_by(inner_block as usize) + .enumerate() + { + let inner_end = inner_start + inner_block; + let mut sources = Vec::new(); + for panel_start in + (inner_start..inner_end).step_by(source_panel_block as usize) + { + let panel_end = panel_start + source_panel_block; + let source = self + .matrix_shards_for_block( + &right_shards, + right_column_axis, + right_inner_axis, + column_start, + column_end, + panel_start, + panel_end, + ) + .next() + .ok_or(LowLoweringError::InvalidOperatorPlan)?; + let target_view = self.narrow_view( + source, + &[ + (right_inner_axis, panel_start, panel_end), + (right_column_axis, column_start, column_end), + ], + )?; + let source_is_local = + self.shards[source.index() as usize].tile == left_shard.tile; + let consume_direct = source_is_local && !stage_local_sources; + if !consume_direct { + let destination_view = self.narrow_view( + weights.ok_or(LowLoweringError::InvalidOperatorPlan)?, + &[ + (right_inner_axis, panel_start, panel_end), + (right_column_axis, column_start, column_end), + ], + )?; + if source_is_local { + append_logical_span_copies( + &self.shards, + &target_view, + &destination_view, + left_shard.tile, + &mut local_copies, + )?; + } else { + transfers + .entry(target_view.clone()) + .or_default() + .push(destination_view); + } + } + sources.push((target_view, consume_direct)); + } + + let split = sources.len() > 1 && sources.iter().any(|(_, local)| *local); + let invocations = if split { + sources + .into_iter() + .enumerate() + .map(|(panel, (source, local))| { + let start = inner_start + + u32::try_from(panel) + .map_err(|_| LowLoweringError::IdOverflow)? + * source_panel_block; + Ok((start, start + source_panel_block, local.then_some(source))) + }) + .collect::>>()? + } else { + let direct = (sources.len() == 1 && sources[0].1) + .then(|| sources.pop().expect("one direct source").0); + vec![(inner_start, inner_end, direct)] + }; + for (panel, (panel_start, panel_end, direct)) in + invocations.into_iter().enumerate() + { + let left_view = self.narrow_view( + resident_left.shard, + &[(left_inner_axis, panel_start, panel_end)], + )?; + let staged = weights.ok_or(LowLoweringError::InvalidOperatorPlan); + let selected = direct + .as_ref() + .map_or_else(|| staged, |view| Ok(view.shard))?; + let mode = if block_index == 0 && panel == 0 { + GemmKernelMode::Initialize + } else { + GemmKernelMode::Accumulate + }; + let mut family = kernel_family; + family.weights = if self.shards[selected.index() as usize] + .tensor_type + .format + .layout + .memory_class + == crate::MemoryClass::Standard + { + crate::GemmWeightLoad::Standard + } else { + crate::GemmWeightLoad::Interleaved + }; + let kernel = gemm_kernel_spec( + family, + mode, + crate::GemmBlockShape { + inner: panel_end - panel_start, + output_columns: local_output_columns, + }, + gemm_kernel_rows( + &partial, + kernel_requirements.output.format.layout.order, + )?, + ); + let weight_view = direct.map_or_else( + || { + self.narrow_view( + selected, + &[ + (right_inner_axis, panel_start, panel_end), + (right_column_axis, column_start, column_end), + ], + ) + }, + Ok, + )?; + gemm_runs.push(( + left_shard.tile, + KernelRun::new( + WorkProvenance { + operation: operation.source, + value: Some(*output_value), + reason: WorkReason::OperatorKernel, + }, + kernel, + vec![ + KernelOperand { + views: vec![left_view], + }, + KernelOperand { + views: vec![weight_view], + }, + ], + partial.clone(), + KernelRequirements::Operator(kernel_requirements.clone()), + ), + )); + } + } + } + } + self.append_phase( + transfers, + WorkProvenance { + operation: operation.source, + value: Some(*right_value), + reason: WorkReason::OperatorInput { + input: orientation.physical_right_input() as u16, + }, + }, + semantic_exchange, + tiles, + )?; + for (tile, copy) in local_copies { + self.append_local_copy(tiles, tile, copy)?; + } + for (tile, run) in gemm_runs { + self.append_kernel(tiles, tile, run)?; + } + + tracing::debug!( + partial_groups = partials.len(), + output_shards = output_shards.len(), + result_row_partitions, + result_column_partitions, + "prepared parallel GEMM partials" + ); + + let remote_partials_per_stage = match reduction_staging { + crate::ReductionStaging::Complete => inner_partitions.saturating_sub(1), + crate::ReductionStaging::Streamed => 1, + }; + let reduction_stages = inner_partitions + .saturating_sub(1) + .div_ceil(remote_partials_per_stage.max(1)); + let mut reduction_transfers = (0..reduction_stages) + .map(|_| BTreeMap::>::new()) + .collect::>(); + let mut seed_copies = Vec::<(u16, LocalCopy)>::new(); + let mut reduction_runs = (0..reduction_stages) + .map(|_| Vec::<(u16, KernelRun)>::new()) + .collect::>(); + let mut result_copies = Vec::<(u16, LocalCopy)>::new(); + let mut reduction_roots = 0usize; + for contributors in partials.into_values() { + let Some((_, complete)) = contributors.first() else { + return Err(LowLoweringError::InvalidOperatorPlan); + }; + let expected = complete + .extents + .iter() + .try_fold(1u64, |elements, extent| { + elements.checked_mul(u64::from(extent.physical_end - extent.start)) + }) + .ok_or(LowLoweringError::IdOverflow)?; + let mut covered = 0u64; + for output in output_shards.iter().copied() { + let owner = self.shards[output.index() as usize].clone(); + let intersection = + intersect_extents_with_shared_padding(&owner.extents, &complete.extents); + let Some(intersection) = intersection else { + continue; + }; + let elements = intersection + .iter() + .try_fold(1u32, |elements, extent| { + elements.checked_mul(extent.physical_end - extent.start) + }) + .ok_or(LowLoweringError::IdOverflow)?; + if elements == 0 || !elements.is_multiple_of(8) { + return Err(LowLoweringError::InvalidOperatorPlan); + } + covered = covered + .checked_add(u64::from(elements)) + .ok_or(LowLoweringError::IdOverflow)?; + + let initial = self.push_packed_buffer( + owner.tile, + elements, + Precision::F16, + ShardDefinition::Staging, + )?; + let remote_elements = elements + .checked_mul(u32::from(remote_partials_per_stage)) + .ok_or(LowLoweringError::IdOverflow)?; + let remote = self.push_packed_buffer( + owner.tile, + remote_elements, + Precision::F16, + ShardDefinition::ExchangeStaging, + )?; + let result = self.push_packed_buffer( + owner.tile, + elements, + Precision::F16, + ShardDefinition::Staging, + )?; + let seed = contributors + .iter() + .position(|(tile, _)| *tile == owner.tile) + .unwrap_or(0); + let source_view = |partial: &ShardView| ShardView { + shard: partial.shard, + extents: intersection.clone().into(), + }; + let seed_source = source_view(&contributors[seed].1); + if contributors[seed].0 == owner.tile { + append_span_copies( + &self.shards, + &seed_source, + &self.full_view(initial), + owner.tile, + &mut seed_copies, + )?; + } else { + reduction_transfers[0] + .entry(seed_source) + .or_default() + .push(self.full_view(initial)); + } + + let remote_contributors = contributors + .iter() + .enumerate() + .filter(|(index, _)| *index != seed) + .map(|(_, (_, partial))| partial) + .collect::>(); + for (stage, chunk) in remote_contributors + .chunks(usize::from(remote_partials_per_stage)) + .enumerate() + { + for (slot, partial) in chunk.iter().enumerate() { + let start = u32::try_from(slot) + .map_err(|_| LowLoweringError::IdOverflow)? + .checked_mul(elements) + .ok_or(LowLoweringError::IdOverflow)?; + let end = start + .checked_add(elements) + .ok_or(LowLoweringError::IdOverflow)?; + reduction_transfers[stage] + .entry(source_view(partial)) + .or_default() + .push(ShardView { + shard: remote, + extents: vec![ShardExtent { + axis: 0, + start, + logical_end: end, + physical_end: end, + }] + .into(), + }); + } + let (accumulator, stage_result) = if stage.is_multiple_of(2) { + (initial, result) + } else { + (result, initial) + }; + reduction_runs[stage].push(( + owner.tile, + KernelRun::new( + WorkProvenance { + operation: operation.source, + value: Some(*output_value), + reason: WorkReason::OperatorKernel, + }, + TileKernelSpec::ReductionSum { + partials: u16::try_from(chunk.len() + 1) + .map_err(|_| LowLoweringError::IdOverflow)?, + }, + vec![ + KernelOperand { + views: vec![self.full_view(accumulator)], + }, + KernelOperand { + views: vec![self.full_view(remote)], + }, + ], + self.full_view(stage_result), + KernelRequirements::Operator(requirements.clone()), + ), + )); + } + let final_result = if usize::from(reduction_stages).is_multiple_of(2) { + initial + } else { + result + }; + append_span_copies( + &self.shards, + &self.full_view(final_result), + &ShardView { + shard: output, + extents: intersection.into(), + }, + owner.tile, + &mut result_copies, + )?; + reduction_roots += 1; + } + if covered != expected { + return Err(LowLoweringError::InvalidOperatorPlan); + } + } + for (stage, (transfers, runs)) in reduction_transfers + .into_iter() + .zip(reduction_runs) + .enumerate() + { + self.append_phase( + transfers, + WorkProvenance { + operation: operation.source, + value: Some(*output_value), + reason: WorkReason::OperatorInputs, + }, + physical_exchange, + tiles, + )?; + if stage == 0 { + for (tile, copy) in seed_copies.drain(..) { + self.append_local_copy(tiles, tile, copy)?; + } + } + for (tile, run) in runs { + self.append_kernel(tiles, tile, run)?; + } + } + for (tile, copy) in result_copies { + self.append_local_copy(tiles, tile, copy)?; + } + tracing::debug!(reduction_roots, "materialized packed parallel reduction"); + } + Ok(()) + } + + fn lower_output_stationary_gemm( + &mut self, + operation: &MidOperation, + gemm: GemmLowering, + [inner_extent, column_extent]: [u32; 2], + requirements: &OperatorRequirements, + tiles: &mut [TileWorkList], + ) -> LowLoweringResult<()> { + let staging_memory_class = requirements.inputs[1].local_staging.memory_class(); + let mut remote_staging = BTreeMap::<(u16, u32), LowShardId>::new(); + let mut local_staging = BTreeMap::<(u16, u32), LowShardId>::new(); + + for inner_start in (0..inner_extent).step_by(gemm.block.inner as usize) { + let inner_end = inner_start + gemm.block.inner; + let mut transfers = BTreeMap::>::new(); + let mut local_copies = Vec::<(u16, LocalCopy)>::new(); + let mut runs = Vec::with_capacity(gemm.output_shards.len()); + let mut left_views = BTreeMap::::new(); + for column_start in (0..column_extent).step_by(gemm.block.output_columns as usize) { + let column_end = column_start + gemm.block.output_columns; + let right_candidates = self + .right_shards_for_block( + &gemm.right_shards, + column_start, + column_end, + inner_start, + inner_end, + ) + .collect::>(); + if right_candidates.is_empty() { + return Err(LowLoweringError::InvalidOperatorPlan); + } + let column_outputs = gemm + .output_shards + .iter() + .copied() + .filter(|output| { + let extents = &self.shards[output.index() as usize].extents; + let columns = extents[extents.len() - 1]; + columns.start <= column_start && columns.physical_end >= column_end + }) + .collect::>(); + if column_outputs.is_empty() { + return Err(LowLoweringError::InvalidOperatorPlan); + } + for output in column_outputs { + let tile = self.shards[output.index() as usize].tile; + let left_view = if let Some(view) = left_views.get(&tile) { + view.clone() + } else { + let view = self.schedule_input_view( + gemm.left_value, + tile, + &[(gemm.left_rank - 1, inner_start, inner_end)], + &mut transfers, + &mut local_copies, + )?; + left_views.insert(tile, view.clone()); + view + }; + let right = self + .prefer_local_shard(&right_candidates, tile) + .ok_or(LowLoweringError::InvalidOperatorPlan)?; + let right_rank = self.shards[right.index() as usize].extents.len(); + let right_view = self.narrow_view( + right, + &[ + (right_rank - 2, inner_start, inner_end), + (right_rank - 1, column_start, column_end), + ], + )?; + let local = self.shards[right.index() as usize].tile == tile; + let local_spans = local + .then(|| view_byte_spans(&self.shards[right.index() as usize], &right_view)) + .transpose()?; + let resident_view = + if local_spans.as_ref().is_some_and(|spans| spans.len() == 1) + && !requirements.inputs[1].local_staging.stages_local() + { + right_view.clone() + } else { + let resident = if local { + let existing = local_staging.get(&(tile, column_start)).copied(); + let mut tensor_type = + self.shards[right.index() as usize].tensor_type.clone(); + tensor_type.format.layout.memory_class = staging_memory_class; + let resident = self.push_shard(LowShard { + id: LowShardId(0), + tile, + tensor_type, + extents: right_view.extents.clone(), + definition: existing + .map(ShardDefinition::Alias) + .unwrap_or(ShardDefinition::LocalCopy(right)), + })?; + local_staging + .entry((tile, column_start)) + .or_insert(resident); + resident + } else { + let selected_staging = &mut remote_staging; + if let Some(resident) = + selected_staging.get(&(tile, column_start)).copied() + { + resident + } else { + let mut tensor_type = + self.shards[right.index() as usize].tensor_type.clone(); + tensor_type.format.layout.memory_class = staging_memory_class; + let resident = self.push_shard(LowShard { + id: LowShardId(0), + tile, + tensor_type, + extents: right_view.extents.clone(), + definition: ShardDefinition::ExchangeStaging, + })?; + selected_staging.insert((tile, column_start), resident); + resident + } + }; + if let Some(spans) = local_spans { + let mut destination_offset = 0u32; + for span in spans { + local_copies.push(( + tile, + LocalCopy { + source: right, + source_offset: span.offset, + destination: resident, + destination_offset, + bytes: span.bytes, + pattern: LocalCopyPattern::Contiguous, + }, + )); + destination_offset = destination_offset + .checked_add(span.bytes) + .ok_or(LowLoweringError::IdOverflow)?; + } + } else { + transfers + .entry(right_view.clone()) + .or_default() + .push(self.full_view(resident)); + } + self.full_view(resident) + }; + let output_view = self + .narrow_view(output, &[(gemm.output_rank - 1, column_start, column_end)])?; + let mode = if inner_start == 0 { + GemmKernelMode::Initialize + } else { + GemmKernelMode::Accumulate + }; + let mut family = gemm.kernel; + if self.shards[resident_view.shard.index() as usize] + .tensor_type + .format + .layout + .memory_class + == crate::MemoryClass::Interleaved + { + family.weights = crate::GemmWeightLoad::Interleaved; + } + let kernel = gemm_kernel_spec( + family, + mode, + gemm.block, + gemm_kernel_rows(&output_view, requirements.output.format.layout.order)?, + ); + runs.push(( + tile, + KernelRun::new( + WorkProvenance { + operation: operation.source, + value: Some(gemm.output_value), + reason: WorkReason::OperatorKernel, + }, + kernel, + vec![ + KernelOperand { + views: vec![left_view], + }, + KernelOperand { + views: vec![resident_view], + }, + ], + output_view, + KernelRequirements::Operator(requirements.clone()), + ), + )); + } + } + self.append_phase( + transfers, + WorkProvenance { + operation: operation.source, + value: (!self.deferred_values.contains_key(&gemm.left_value)) + .then_some(gemm.right_value), + reason: if self.deferred_values.contains_key(&gemm.left_value) { + WorkReason::OperatorInputs + } else { + WorkReason::OperatorInput { input: 1 } + }, + }, + semantic_exchange, + tiles, + )?; + for (tile, copy) in local_copies { + self.append_local_copy(tiles, tile, copy)?; + } + for (tile, run) in runs { + self.append_kernel(tiles, tile, run)?; + } + } + Ok(()) + } +} * Unmerged path crates/ipu-codegen/src/low/mod.rs diff --git a/crates/ipu-codegen/src/low/pointwise.rs b/crates/ipu-codegen/src/low/pointwise.rs new file mode 100644 index 0000000..51c15fc --- /dev/null +++ b/crates/ipu-codegen/src/low/pointwise.rs @@ -0,0 +1,219 @@ +use super::*; + +impl LoweringState { + pub(super) fn lower_schedule( + &mut self, + operation: &MidOperation, + schedule: &OperatorSchedule, + requirements: &OperatorRequirements, + tiles: &mut [TileWorkList], + ) -> LowLoweringResult<()> { + let [ScheduleStep::KernelMap(map)] = schedule.steps.as_slice() else { + return Err(LowLoweringError::InvalidOperatorPlan); + }; + if map.output != ScheduleValue::Output || map.inputs.len() != operation.inputs.len() { + return Err(LowLoweringError::InvalidOperatorPlan); + } + let [result] = operation.results.as_slice() else { + return Err(LowLoweringError::ResultArity); + }; + let outputs = self.value_shards(*result)?.to_vec(); + let mut wave_transfers = Vec::>>::new(); + let mut wave_runs = Vec::>::new(); + for output in outputs { + if self.shards[output.index() as usize] + .extents + .iter() + .any(|extent| extent.start == extent.physical_end) + { + continue; + } + let tile = self.shards[output.index() as usize].tile; + let sources = map + .inputs + .iter() + .map(|(value, access)| { + let ScheduleValue::Input(index) = value else { + return Err(LowLoweringError::InvalidOperatorPlan); + }; + let input = operation + .inputs + .get(usize::from(*index)) + .ok_or(LowLoweringError::InvalidOperatorPlan)?; + Ok(match access { + ScheduleAccess::LogicalOverlap => self + .value_shards(*input)? + .iter() + .find_map(|source| self.broadcast_view(*source, output)) + .ok_or(LowLoweringError::InvalidOperatorPlan)?, + ScheduleAccess::TileLocal => { + let output_extents = &self.shards[output.index() as usize].extents; + let source = self + .value_shards(*input)? + .iter() + .copied() + .find(|source| { + let source = &self.shards[source.index() as usize]; + source.tile == tile && source.extents == *output_extents + }) + .ok_or(LowLoweringError::InvalidOperatorPlan)?; + self.full_view(source) + } + }) + }) + .collect::>>()?; + let chunks = vec![self.shards[output.index() as usize].extents.clone()]; + for (wave, output_extents) in chunks.into_iter().enumerate() { + if wave_transfers.len() <= wave { + wave_transfers.push(BTreeMap::new()); + wave_runs.push(Vec::new()); + } + let inputs = sources + .iter() + .enumerate() + .map(|(index, source)| { + let access = map.inputs[index].1; + let source_view = match access { + ScheduleAccess::LogicalOverlap => self + .broadcast_view_for_extents(source.shard, output, &output_extents) + .ok_or(LowLoweringError::InvalidOperatorPlan)?, + ScheduleAccess::TileLocal => source.clone(), + }; + let view = if self.shards[source_view.shard.index() as usize].tile == tile { + source_view + } else { + let copy = self.push_shard(LowShard { + id: LowShardId(0), + tile, + tensor_type: self.shards[source_view.shard.index() as usize] + .tensor_type + .clone(), + extents: source_view.extents.clone(), + definition: ShardDefinition::ExchangeStaging, + })?; + wave_transfers[wave] + .entry(source_view) + .or_default() + .push(self.full_view(copy)); + self.full_view(copy) + }; + Ok(KernelOperand { views: vec![view] }) + }) + .collect::>()?; + wave_runs[wave].push(( + tile, + KernelRun::new( + WorkProvenance { + operation: operation.source, + value: operation.results.first().copied(), + reason: WorkReason::OperatorKernel, + }, + map.kernel.clone(), + inputs, + ShardView { + shard: output, + extents: output_extents, + }, + KernelRequirements::Operator(requirements.clone()), + ), + )); + } + } + for (transfers, runs) in wave_transfers.into_iter().zip(wave_runs) { + self.append_phase( + transfers, + WorkProvenance { + operation: operation.source, + value: None, + reason: WorkReason::OperatorInputs, + }, + semantic_exchange, + tiles, + )?; + for (tile, run) in runs { + self.append_kernel(tiles, tile, run)?; + } + } + Ok(()) + } + + fn broadcast_view(&self, source: LowShardId, output: LowShardId) -> Option { + self.broadcast_view_for_extents( + source, + output, + &self.shards[output.index() as usize].extents, + ) + } + + fn broadcast_view_for_extents( + &self, + source: LowShardId, + output: LowShardId, + output_extents: &[ShardExtent], + ) -> Option { + let source_shard = &self.shards[source.index() as usize]; + let output_shard = &self.shards[output.index() as usize]; + let source_rank = source_shard.extents.len(); + let output_rank = output_shard.extents.len(); + if source_rank > output_rank { + return None; + } + let offset = output_rank - source_rank; + let mut extents = source_shard.extents.clone(); + for (axis, extent) in extents.iter_mut().enumerate() { + let dimension = source_shard.tensor_type.shape.0[axis]; + if dimension == 1 { + if extent.start != 0 || extent.logical_end == 0 { + return None; + } + extent.start = 0; + extent.logical_end = 1; + extent.physical_end = 1; + } else { + let required = output_extents[offset + axis]; + if extent.start > required.start || extent.logical_end < required.logical_end { + return None; + } + extent.start = required.start; + extent.logical_end = required.logical_end; + extent.physical_end = required.logical_end; + } + } + Some(ShardView { + shard: source, + extents, + }) + } + + pub(super) fn schedule_input_view( + &mut self, + value: MidValueId, + tile: u16, + ranges: &[(usize, u32, u32)], + transfers: &mut BTreeMap>, + local_copies: &mut Vec<(u16, LocalCopy)>, + ) -> LowLoweringResult { + let target = self.local_shard(value, tile)?; + let target_view = self.narrow_view(target, ranges)?; + if !self.deferred_values.contains_key(&value) { + return Ok(target_view); + } + + let staging = self.push_shard(LowShard { + id: LowShardId(0), + tile, + tensor_type: self.shards[target.index() as usize].tensor_type.clone(), + extents: target_view.extents.clone(), + definition: ShardDefinition::ExchangeStaging, + })?; + self.materialize_deferred_region( + value, + &target_view.extents.logical(), + staging, + ExchangeOrder::Semantic, + transfers, + local_copies, + )?; + Ok(self.full_view(staging)) + } +} diff --git a/crates/ipu-codegen/src/low/tests.rs b/crates/ipu-codegen/src/low/tests.rs new file mode 100644 index 0000000..dddf68f --- /dev/null +++ b/crates/ipu-codegen/src/low/tests.rs @@ -0,0 +1,1490 @@ +use super::*; +use crate::{ + AxisTiling, ComputeGraph, GridOrder, Ipu21CostModel, Layout, MemoryClass, Padding, + PipelineConfig, PlannerSearchDomain, Precision, StorageOrder, TensorAxis, TensorFormat, + TensorTiling, TileKernelSpec, lower, +}; +use std::collections::BTreeSet; + +const CASES: usize = 32; + +fn format(tiles: u16) -> TensorFormat { + TensorFormat { + precision: Precision::F16, + layout: Layout::row_sharded(tiles), + } +} + +#[test] +fn randomized_linear_shards_cover_flat_storage_once_in_balanced_grains() { + let mut random = fastrand::Rng::with_seed(0x666c_6174_5f73_6864); + for case in 0..CASES { + let rank = random.usize(2..=4); + let grain = 1_u32 << random.u32(1..=5); + let mut shape = (0..rank - 1).map(|_| random.u32(1..=5)).collect::>(); + shape.push(grain * random.u32(1..=8)); + let elements = shape + .iter() + .map(|&extent| u64::from(extent)) + .product::(); + let grains = elements / u64::from(grain); + let tiles = random.u16(1..=u16::try_from(grains.min(64)).unwrap()); + let tensor = TensorType::new( + shape.clone(), + Precision::F16, + Layout::logical_linear(tiles, grain), + ); + let shards = shard_extents(&tensor).unwrap(); + let mut coverage = vec![0_u8; usize::try_from(elements).unwrap()]; + let mut tile_elements = vec![0_u64; usize::from(tiles)]; + for (tile, extents) in shards { + assert_eq!(extents.len(), rank, "case {case}"); + assert!( + extents[..rank - 1] + .iter() + .all(|extent| extent.logical_end == extent.start + 1), + "case {case}" + ); + let mut row = 0_u64; + for (axis, extent) in extents[..rank - 1].iter().enumerate() { + row = row * u64::from(shape[axis]) + u64::from(extent.start); + } + let columns = &extents[rank - 1]; + let width = u64::from(shape[rank - 1]); + for column in columns.start..columns.logical_end { + let index = usize::try_from(row * width + u64::from(column)).unwrap(); + coverage[index] += 1; + tile_elements[usize::from(tile)] += 1; + } + } + assert!(coverage.into_iter().all(|count| count == 1), "case {case}"); + assert!( + tile_elements + .iter() + .all(|count| count % u64::from(grain) == 0), + "case {case}" + ); + assert!( + tile_elements.iter().max().unwrap() - tile_elements.iter().min().unwrap() + <= u64::from(grain), + "case {case}" + ); + } +} + +#[test] +fn randomized_parallel_reduction_gemms_lower_to_packed_reductions() { + let mut random = fastrand::Rng::with_seed(0x7472_6565_5f6b_7370); + for case in 0..CASES { + let output_columns = [64, 128][random.usize(0..2)]; + let inner_partitions = random.u16(2..=4); + let column_partitions = random.u16(1..=3); + let row_partitions = random.u16(inner_partitions..=8); + let tiles = inner_partitions * column_partitions * row_partitions; + let row_distributed_result = inner_partitions == 2 || random.bool(); + let rows_per_partition = if row_distributed_result { + u32::from(inner_partitions) * 2 + } else { + 2 + }; + let rows = u32::from(row_partitions) * rows_per_partition; + let inner = u32::from(inner_partitions) + * 64 + * random.u32(1..=u32::from(row_partitions / inner_partitions)); + let columns = u32::from(column_partitions) * output_columns; + let mut graph = ComputeGraph::new(); + let left = graph.host_input("left", [1, rows, inner]).unwrap(); + let right = graph.parameter("right", [1, inner, columns]).unwrap(); + let product = graph.gemm(left, right).unwrap(); + graph.set_outputs([product]).unwrap(); + let left_format = TensorFormat { + precision: Precision::F16, + layout: Layout::amp_left_parallel_grid( + crate::GemmOrientation::Normal, + 64, + tiles, + row_partitions, + column_partitions, + inner_partitions, + ), + }; + let right_format = TensorFormat { + precision: Precision::F16, + layout: Layout::block_major_matrix_storage( + crate::GemmOrientation::Normal, + 64, + output_columns, + column_partitions, + inner_partitions, + 1, + MemoryClass::Interleaved, + ), + }; + let (result_row_partitions, result_column_partitions) = if random.bool() { + (1, 1) + } else if row_distributed_result { + (inner_partitions, 1) + } else { + (1, inner_partitions) + }; + let reduction_staging = if random.bool() { + crate::ReductionStaging::Complete + } else { + crate::ReductionStaging::Streamed + }; + let config = PipelineConfig::new(tiles) + .with_search_domain( + PlannerSearchDomain::default() + .with_active_tile_counts([tiles]) + .with_operator_precisions(crate::OperatorClass::Gemm, [Precision::F16]) + .with_gemm_plan_constraint(crate::GemmPlanConstraint { + source_operation: 0, + geometry: crate::GemmGeometry { + block: crate::GemmBlockShape { + inner: inner + .div_ceil(u32::from(inner_partitions)) + .div_ceil(crate::layout::AMP_COLUMN_MICRO) + * crate::layout::AMP_COLUMN_MICRO, + output_columns, + }, + orientation: crate::GemmOrientation::Normal, + compute: crate::GemmGrid { + rows: row_partitions, + columns: column_partitions, + inner: inner_partitions, + }, + result: crate::GemmResultGrid { + rows: row_partitions.saturating_mul(result_row_partitions), + columns: column_partitions.saturating_mul(result_column_partitions), + }, + order: crate::GridOrder::ColumnsFast, + }, + reduction_staging: Some(reduction_staging), + weight_memory_class: MemoryClass::Interleaved, + local_weight_staging: crate::LocalOperandStaging::Direct( + MemoryClass::Interleaved, + ), + }), + ) + .with_input(left, left_format) + .with_input(right, right_format); + let mid = lower(&graph, &config, &Ipu21CostModel) + .unwrap_or_else(|error| { + panic!( + "case {case}: {error}; rows={rows} inner={inner} columns={columns} grid={row_partitions}x{column_partitions}x{inner_partitions} result={result_row_partitions}x{result_column_partitions} block={output_columns} staging={reduction_staging:?}" + ) + }); + let low = lower_to_tiles(&mid, &config) + .unwrap_or_else(|error| { + panic!( + "case {case}: {error}; rows={rows} inner={inner} columns={columns} grid={row_partitions}x{column_partitions}x{inner_partitions}" + ) + }); + let reduction_runs = low + .kernel_runs + .iter() + .filter(|run| matches!(run.kernel, TileKernelSpec::ReductionSum { .. })) + .collect::>(); + assert!(!reduction_runs.is_empty(), "case {case}"); + assert!( + reduction_runs.iter().all(|run| { + matches!( + run.kernel, + TileKernelSpec::ReductionSum { partials } + if partials == match reduction_staging { + crate::ReductionStaging::Complete => inner_partitions, + crate::ReductionStaging::Streamed => 2, + } + ) && run.inputs.len() == 2 + }), + "case {case}" + ); + assert!( + low.exchange_phases.len() <= usize::from(inner_partitions).saturating_add(2), + "case {case}" + ); + if (result_row_partitions, result_column_partitions) != (1, 1) { + let output_shards = low.outputs[0] + .shards + .iter() + .copied() + .collect::>(); + let packed_results = reduction_runs + .iter() + .map(|run| run.output.shard) + .collect::>(); + let copied_outputs = low + .local_copies + .iter() + .filter(|copy| packed_results.contains(©.source)) + .map(|copy| copy.destination) + .collect::>(); + assert!( + output_shards.is_subset(&copied_outputs), + "case {case}: every distributed result shard must receive a packed result" + ); + } + } +} + +#[test] +fn randomized_parameter_owner_groups_pack_independently_of_compute_tiles() { + let mut random = fastrand::Rng::with_seed(0x7061_7261_6d73); + for case in 0..CASES { + let owner_tiles = 1_u16 << random.u32(1..=3); + let compute_tiles = owner_tiles * 2; + let inner = u32::from(owner_tiles) * 64; + let rows = u32::from(compute_tiles) * random.u32(1..=4); + let mut graph = ComputeGraph::new(); + let left = graph.host_input("left", [rows, inner]).unwrap(); + let right0 = graph.parameter("right.0", [inner, 64]).unwrap(); + let right1 = graph.parameter("right.1", [inner, 64]).unwrap(); + let output0 = graph.gemm(left, right0).unwrap(); + let output1 = graph.gemm(left, right1).unwrap(); + graph.set_outputs([output0, output1]).unwrap(); + + let left_format = TensorFormat { + precision: Precision::F16, + layout: Layout::amp_left(64, compute_tiles), + }; + let right_format = TensorFormat { + precision: Precision::F16, + layout: Layout::block_major_matrix_storage( + crate::GemmOrientation::Normal, + 64, + 64, + 1, + owner_tiles, + 1, + MemoryClass::Standard, + ), + }; + let config = PipelineConfig::new(compute_tiles) + .with_search_domain( + PlannerSearchDomain::default() + .with_active_tile_counts([compute_tiles]) + .with_operator_precisions(crate::OperatorClass::Gemm, [Precision::F16]) + .with_weight_memory_classes([MemoryClass::Standard]), + ) + .with_input(left, left_format.clone()) + .with_input(right0, right_format.clone()) + .with_input(right1, right_format.clone()); + + let mid = lower(&graph, &config, &Ipu21CostModel).unwrap(); + let low = lower_to_tiles(&mid, &config).unwrap(); + let parameter_tiles = |name: &str| { + low.inputs + .iter() + .find(|input| input.name == name) + .unwrap() + .shards + .iter() + .map(|shard| low.shards[shard.index() as usize].tile) + .collect::>() + }; + let first = parameter_tiles("right.0"); + let second = parameter_tiles("right.1"); + assert_eq!(first.len(), usize::from(owner_tiles), "case {case}"); + assert_eq!(second.len(), usize::from(owner_tiles), "case {case}"); + assert!(first.is_disjoint(&second), "case {case}"); + } +} + +#[test] +fn randomized_pointwise_schedule_skips_empty_output_shards() { + let mut random = fastrand::Rng::with_seed(0x656d_7074); + for case in 0..CASES { + let tiles = random.u16(2..=32); + let rows = random.u32(1..u32::from(tiles)); + let columns = random.u32(1..=32) * 2; + let tensor_format = format(tiles); + let mut graph = ComputeGraph::new(); + let input = graph.host_input("input", [rows, columns]).unwrap(); + let output = graph.gelu(input).unwrap(); + graph.set_outputs([output]).unwrap(); + let config = PipelineConfig::new(tiles) + .with_search_domain(PlannerSearchDomain::default().with_active_tile_counts([tiles])) + .with_input(input, tensor_format); + + let mid = lower(&graph, &config, &Ipu21CostModel).unwrap(); + let low = lower_to_tiles(&mid, &config).unwrap(); + let runs = low + .tiles + .iter() + .flat_map(|tile| low.work(tile)) + .filter_map(|work| match work { + TileWorkRef::Kernel(run) => Some(run), + _ => None, + }) + .collect::>(); + assert_eq!(runs.len(), rows as usize, "random case {case}"); + assert!(runs.iter().all(|run| { + run.output + .extents + .iter() + .all(|extent| extent.start < extent.physical_end) + })); + } +} + +#[test] +fn randomized_schedule_streaming_defers_one_use_rearrangements() { + let mut random = fastrand::Rng::with_seed(0x7374_7265_616d); + for case in 0..8 { + let batch = random.u32(1..=4); + let tokens = 16; + let mut graph = ComputeGraph::new(); + let input = graph.host_input("input", [batch, tokens, 64]).unwrap(); + let up = graph.parameter("up", [1, 64, 256]).unwrap(); + let down = graph.parameter("down", [1, 256, 64]).unwrap(); + let hidden = graph.gemm(input, up).unwrap(); + let hidden = graph.gelu(hidden).unwrap(); + let output = graph.gemm(hidden, down).unwrap(); + graph.set_outputs([output]).unwrap(); + let mut config = PipelineConfig::new(16) + .with_search_domain(PlannerSearchDomain::default().with_active_tile_counts([16])) + .with_automatic_input(input, Precision::F16) + .with_automatic_input(up, Precision::F16) + .with_automatic_input(down, Precision::F16); + config.conversion_streaming = crate::ConversionStreamingPolicy::Always; + + let mid = lower(&graph, &config, &Ipu21CostModel).unwrap(); + let deferred = mid + .operations + .iter() + .filter_map(|operation| { + operation.conversion().and_then(|(_, materialization)| { + (materialization == crate::OperandMaterialization::DispatchSlices) + .then(|| operation.results[0]) + }) + }) + .collect::>(); + assert!(!deferred.is_empty(), "case {case}"); + let consumers = mid + .operations + .iter() + .filter(|operation| { + operation + .inputs + .iter() + .any(|input| deferred.contains(input)) + }) + .filter_map(|operation| operation.source) + .collect::>(); + + let low = lower_to_tiles(&mid, &config).unwrap(); + for run in low + .tiles + .iter() + .flat_map(|tile| low.work(tile)) + .filter_map(|work| match work { + TileWorkRef::Kernel(run) if matches!(run.kernel, TileKernelSpec::Gemm { .. }) => { + Some(run) + } + _ => None, + }) + { + let output = &low.shards[run.output.shard.index() as usize]; + let flattens_outer_rows = matches!( + output.tensor_type.format.layout.order, + StorageOrder::Native(NativeKernelOrder::Left | NativeKernelOrder::Output) + ); + assert!( + flattens_outer_rows + || run.output.extents[..run.output.extents.len() - 2] + .iter() + .all(|extent| extent.physical_end - extent.start == 1), + "case {case}" + ); + } + assert!( + low.exchange_phases + .iter() + .all(|phase| phase.provenance.reason != WorkReason::LayoutRearrangement), + "case {case}" + ); + assert!( + low.exchange_phases + .iter() + .any(|phase| phase.provenance.reason == WorkReason::OperatorInputs), + "case {case}" + ); + assert!( + low.shards + .iter() + .filter(|shard| shard.definition == ShardDefinition::Unmaterialized) + .count() + >= 16 * deferred.len(), + "case {case}" + ); + for run in &low.kernel_runs { + if run + .provenance + .operation + .is_some_and(|operation| consumers.contains(&operation)) + { + let input = &run.inputs[0].views[0]; + assert_ne!( + low.shards[input.shard.index() as usize].definition, + ShardDefinition::Unmaterialized, + "case {case}" + ); + let inner = input.extents.last().unwrap(); + let TileKernelSpec::Gemm { inner_block, .. } = &run.kernel else { + continue; + }; + assert!( + inner.physical_end - inner.start <= *inner_block, + "case {case}" + ); + } + } + } +} + +#[test] +fn randomized_deferred_views_materialize_arbitrary_regions() { + let mut random = fastrand::Rng::with_seed(0x7669_6577_5f72_6567); + for case in 0..CASES { + let batch = random.u32(1..=4); + let heads = random.u32(2..=6); + let rows = random.u32(1..=2) * AMP_INNER_BLOCK; + let width = random.u32(1..=2) * AMP_COLUMN_MICRO; + let tiles = u16::try_from(batch * heads).unwrap(); + let mut graph = ComputeGraph::new(); + let input = graph + .host_input("input", [batch, rows, heads * width]) + .unwrap(); + let output = graph.split_heads(input, heads).unwrap(); + graph.set_outputs([output]).unwrap(); + let config = PipelineConfig::new(tiles).with_automatic_input(input, Precision::F16); + let mid = lower(&graph, &config, &Ipu21CostModel).unwrap(); + let operation = mid + .operations + .iter() + .find(|operation| matches!(operation.kind, MidOperationKind::Convert(Some(_), ..))) + .unwrap(); + let MidOperationKind::Convert(Some(transform), ..) = operation.kind else { + unreachable!() + }; + let (source, result) = (operation.inputs[0], operation.results[0]); + let mut state = LoweringState::new(&mid, tiles, config.target).unwrap(); + state.deferred_values.insert( + result, + DeferredValue::View(transform, state.value_shards(source).unwrap().to_vec()), + ); + + let stream = random.u32(0..batch * heads); + let row_start = random.u32(0..rows); + let row_end = random.u32(row_start + 1..=rows); + let column_start = random.u32(0..width); + let column_end = random.u32(column_start + 1..=width); + let region = TensorRegion::logical_bounds([ + (stream, stream + 1), + (row_start, row_end), + (column_start, column_end), + ]) + .unwrap(); + let destination = state + .push_matrix_buffer( + random.u16(0..tiles), + row_end - row_start, + row_end - row_start, + column_end - column_start, + column_end - column_start, + StorageOrder::Linear, + ) + .unwrap(); + let mappings = state + .deferred_region_mappings(result, ®ion, destination) + .unwrap_or_else(|error| panic!("case {case}: {error}")); + assert_eq!( + mappings + .iter() + .map(|(source, _)| source.extents.logical_elements()) + .sum::(), + region.logical_elements(), + "case {case}" + ); + for (source, destination) in mappings { + assert_eq!( + [ + source.extents[0].start, + source.extents[1].start - row_start, + source.extents[2].start - (stream % heads) * width - column_start, + u32::try_from(source.extents.logical_elements()).unwrap(), + ], + [ + stream / heads, + destination.extents[0].start, + destination.extents[1].start, + u32::try_from(destination.extents.logical_elements()).unwrap(), + ], + "case {case}" + ); + } + } +} + +#[test] +fn randomized_tile_local_gelu_reorders_without_exchange() { + let mut random = fastrand::Rng::with_seed(0x6765_6c75); + for case in 0..CASES { + let row_partitions = 1_u16 << random.u32(0..=3); + let column_partitions = 1_u16 << random.u32(0..=3); + let tiles = row_partitions * column_partitions; + let rows = u32::from(row_partitions) * random.u32(1..=8); + let columns = u32::from(column_partitions) * 64 * random.u32(1..=4); + let input_format = TensorFormat { + precision: Precision::F16, + layout: Layout::amp_output_replicated_grid(tiles, row_partitions, column_partitions), + }; + let mut graph = ComputeGraph::new(); + let input = graph.host_input("input", [rows, columns]).unwrap(); + let output = graph.gelu(input).unwrap(); + graph.set_outputs([output]).unwrap(); + let config = PipelineConfig::new(tiles) + .with_search_domain(PlannerSearchDomain::default().with_active_tile_counts([tiles])) + .with_input(input, input_format); + + let mid = lower(&graph, &config, &Ipu21CostModel).unwrap(); + let low = lower_to_tiles(&mid, &config).unwrap(); + assert!(low.exchange_phases.is_empty(), "random case {case}"); + for tile in &low.tiles { + for work in low.work(tile) { + let TileWorkRef::Kernel(run) = work else { + continue; + }; + assert_eq!( + low.shards[run.inputs[0].views[0].shard.index() as usize].tile, + tile.tile + ); + assert_eq!( + low.shards[run.output.shard.index() as usize].tile, + tile.tile + ); + } + } + } +} + +#[test] +fn randomized_multiaxis_pointwise_layouts_need_no_exchange_staging() { + let mut random = fastrand::Rng::with_seed(0x6469_7265_6374_7265); + for case in 0..CASES { + let source_rows = 1_u16 << random.u32(0..=3); + let source_columns = 1_u16 << random.u32(0..=3); + let tiles = source_rows * source_columns; + let rows = u32::from(source_rows.max(source_columns)) * random.u32(1..=4); + let columns = u32::from(source_rows.max(source_columns)) * random.u32(1..=4) * 4; + let layout = |row_partitions, column_partitions| Layout { + order: StorageOrder::Linear, + tiling: TensorTiling { + tile_count: tiles, + replicas: 1, + axes: vec![ + AxisTiling::new(TensorAxis::FromEnd(2), row_partitions, 1, Padding::Reject), + AxisTiling::new( + TensorAxis::FromEnd(1), + column_partitions, + 4, + Padding::Reject, + ), + ], + }, + memory_class: MemoryClass::Standard, + }; + let input_format = TensorFormat { + precision: Precision::F16, + layout: layout(source_rows, source_columns), + }; + let mut graph = ComputeGraph::new(); + let input = graph.host_input("input", [rows, columns]).unwrap(); + let output = graph.gelu(input).unwrap(); + graph.set_outputs([output]).unwrap(); + let config = PipelineConfig::new(tiles) + .with_search_domain(PlannerSearchDomain::default().with_active_tile_counts([tiles])) + .with_input(input, input_format); + + let mid = lower(&graph, &config, &Ipu21CostModel).unwrap(); + let low = lower_to_tiles(&mid, &config).unwrap(); + assert!(low.exchange_phases.is_empty(), "case {case}"); + assert!( + low.shards + .iter() + .all(|shard| !matches!(shard.definition, ShardDefinition::ExchangeStaging)), + "case {case}" + ); + assert!(low.local_copies.iter().all(|copy| { + low.shards[copy.source.index() as usize].tile + == low.shards[copy.destination.index() as usize].tile + })); + } +} + +#[test] +fn randomized_multiaxis_shards_cover_padded_extents_in_whole_blocks() { + let mut random = fastrand::Rng::with_seed(0x7368_6172); + for case in 0..CASES { + let row_partitions = random.u16(1..=4); + let column_partitions = random.u16(1..=4); + let replicas = random.u16(1..=3); + let row_block = 1_u32 << random.u32(0..=3); + let column_block = 1_u32 << random.u32(0..=3); + let tile_count = row_partitions * column_partitions * replicas; + let layout = Layout { + order: StorageOrder::Linear, + tiling: TensorTiling { + tile_count, + replicas, + axes: vec![ + AxisTiling::new( + TensorAxis::FromEnd(1), + column_partitions, + column_block, + Padding::Zero, + ), + AxisTiling::new( + TensorAxis::FromEnd(2), + row_partitions, + row_block, + Padding::Zero, + ), + ], + }, + memory_class: MemoryClass::Standard, + }; + let tensor_type = TensorType::new( + [ + u32::from(row_partitions) * row_block + random.u32(0..=65), + u32::from(column_partitions) * column_block + random.u32(0..=65), + ], + Precision::F16, + layout.clone(), + ); + let resolved = layout.resolve(&tensor_type.shape).unwrap(); + let padded = resolved.padded_shape(); + let shards = shard_extents(&tensor_type).unwrap(); + assert_eq!(shards.len(), usize::from(tile_count), "case {case}"); + + for (axis, partitions, block) in [ + (0, row_partitions, row_block), + (1, column_partitions, column_block), + ] { + let ranges = shards + .iter() + .map(|(_, extents)| (extents[axis].start, extents[axis].physical_end)) + .collect::>(); + assert_eq!(ranges.len(), usize::from(partitions), "case {case}"); + let mut cursor = 0; + for (start, end) in ranges { + assert_eq!(start, cursor, "case {case}"); + assert_eq!(start % block, 0, "case {case}"); + assert_eq!(end % block, 0, "case {case}"); + cursor = end; + } + assert_eq!(cursor, padded.0[axis], "case {case}"); + } + } +} + +#[test] +fn randomized_partition_padding_preserves_logical_groups() { + let mut random = fastrand::Rng::with_seed(0x6772_6f75_705f_7064); + for case in 0..CASES * 8 { + let groups = random.u16(1..=16); + let group_width = random.u32(1..=127); + let rows = random.u32(1..=16); + let physical_multiple = 1_u32 << random.u32(1..=4); + let physical_width = group_width.div_ceil(physical_multiple) * physical_multiple; + let physical_blocks = physical_width / physical_multiple; + let partitions_per_group = random.u16(1..=u16::try_from(physical_blocks.min(8)).unwrap()); + let partitions = groups * partitions_per_group; + let layout = Layout { + order: StorageOrder::Linear, + tiling: TensorTiling { + tile_count: partitions, + replicas: 1, + axes: vec![ + AxisTiling::new( + TensorAxis::FromEnd(1), + partitions, + physical_multiple, + Padding::Zero, + ) + .with_padding_groups(groups), + ], + }, + memory_class: MemoryClass::Standard, + }; + let tensor = TensorType::new( + [rows, u32::from(groups) * group_width], + Precision::F16, + layout, + ); + let shards = shard_extents(&tensor).unwrap(); + assert_eq!(shards.len(), usize::from(partitions), "case {case}"); + for group in 0..groups { + let group_base = u32::from(group) * group_width; + let group_shards = &shards[usize::from(group * partitions_per_group) + ..usize::from((group + 1) * partitions_per_group)]; + let mut cursor = group_base; + let mut allocated = 0; + for (_, extents) in group_shards { + assert_eq!(extents[1].start, cursor, "case {case}"); + assert!( + extents[1].logical_end <= group_base + group_width, + "case {case}" + ); + cursor = extents[1].logical_end; + allocated += extents[1].physical_end - extents[1].start; + assert_eq!( + crate::shard_storage_bytes(&LowShard { + id: LowShardId(0), + tile: 0, + tensor_type: tensor.clone(), + extents: extents.clone(), + definition: ShardDefinition::Staging, + }) + .unwrap(), + rows * (extents[1].physical_end - extents[1].start) * 2, + "case {case}" + ); + } + assert_eq!(cursor, group_base + group_width, "case {case}"); + assert_eq!(allocated, physical_width, "case {case}"); + } + assert_eq!( + crate::cost::physical_elements(&tensor.shape, &tensor.format.layout), + u64::from(rows) * u64::from(groups) * u64::from(physical_width), + "case {case}" + ); + assert_eq!( + crate::cost::maximum_shard_bytes(&tensor), + u64::from(rows) + * u64::from( + physical_blocks.div_ceil(u32::from(partitions_per_group)) * physical_multiple, + ) + * 2, + "case {case}" + ); + } +} + +#[test] +fn randomized_padded_intersections_do_not_claim_adjacent_groups() { + let mut random = fastrand::Rng::with_seed(0x7064_5f69_6e74_6572); + for case in 0..CASES * 8 { + let width = random.u32(1..=127); + let padding = random.u32(1..=31); + let group = random.u32(0..=30); + let start = group * width; + let owned = ShardExtent { + axis: 0, + start, + logical_end: start + width, + physical_end: start + width + padding, + }; + let matching = intersect_extents_with_shared_padding(&[owned], &[owned]).unwrap(); + assert_eq!(matching, vec![owned], "case {case}"); + + let adjacent = ShardExtent { + axis: 0, + start: start + width, + logical_end: start + width * 2, + physical_end: start + width * 2 + padding, + }; + assert!( + intersect_extents_with_shared_padding(&[owned], &[adjacent]).is_none(), + "case {case}" + ); + + let narrower_padding = random.u32(0..=padding); + let narrower = ShardExtent { + physical_end: owned.logical_end + narrower_padding, + ..owned + }; + let intersection = intersect_extents_with_shared_padding(&[owned], &[narrower]).unwrap(); + assert_eq!( + intersection[0].physical_end, + owned.logical_end + narrower_padding, + "case {case}" + ); + } +} + +#[test] +fn randomized_gemm_grid_orders_align_operands_and_pair_shared_payloads() { + let mut random = fastrand::Rng::with_seed(0x6772_6964_5f6f_7264); + for case in 0..CASES { + let row_partitions = 1_u16 << random.u32(1..=3); + let column_partitions = 1_u16 << random.u32(1..=3); + let tiles = row_partitions * column_partitions; + let rows = u32::from(row_partitions) * random.u32(1..=8); + let columns = u32::from(column_partitions) * 64 * random.u32(1..=3); + let inner = 64 * random.u32(1..=4); + for order in [GridOrder::ColumnsFast, GridOrder::RowsFast] { + let left = TensorType::new( + [rows, inner], + Precision::F16, + Layout::amp_left_grid(64, tiles, row_partitions, column_partitions, order), + ); + let right = TensorType::new( + [inner, columns], + Precision::F16, + Layout::block_major_matrix_grid( + 64, + 64, + tiles, + row_partitions, + column_partitions, + order, + ), + ); + let output = TensorType::new( + [rows, columns], + Precision::F16, + Layout::amp_output_grid( + crate::GemmOrientation::Normal, + 64, + tiles, + row_partitions, + column_partitions, + order, + ), + ); + let left = shard_extents(&left).unwrap(); + let right = shard_extents(&right).unwrap(); + let output = shard_extents(&output).unwrap(); + for tile in 0..usize::from(tiles) { + assert_eq!(left[tile].1[0], output[tile].1[0], "case {case}"); + assert_eq!(right[tile].1[1], output[tile].1[1], "case {case}"); + } + for tile in (0..usize::from(tiles)).step_by(2) { + let shared_axis = match order { + GridOrder::ColumnsFast => 0, + GridOrder::RowsFast => 1, + }; + assert_eq!( + output[tile].1[shared_axis], + output[tile + 1].1[shared_axis], + "case {case}" + ); + } + } + } +} + +#[test] +fn randomized_micro_panel_mappings_carry_word_aligned_row_padding() { + let mut random = fastrand::Rng::with_seed(0x7061_6464_6564_5f72); + for case in 0..CASES * 8 { + let rows = random.u32(1..=AMP_INNER_BLOCK); + let panel_rows = rows.div_ceil(AMP_COLUMN_MICRO) * AMP_COLUMN_MICRO; + let source_rows = if random.bool() { + panel_rows + } else { + AMP_INNER_BLOCK + }; + let source = LowShard { + id: LowShardId(0), + tile: 0, + tensor_type: TensorType::new( + [rows, AMP_COLUMN_MICRO], + Precision::F16, + Layout { + order: StorageOrder::Native(NativeKernelOrder::TransposedLeft), + tiling: TensorTiling::replicated(1), + memory_class: MemoryClass::Standard, + }, + ), + extents: vec![ + ShardExtent { + axis: 0, + start: 0, + logical_end: rows, + physical_end: source_rows, + }, + ShardExtent { + axis: 1, + start: 0, + logical_end: AMP_COLUMN_MICRO, + physical_end: AMP_COLUMN_MICRO, + }, + ] + .into(), + definition: ShardDefinition::ExchangeStaging, + }; + let destination = LowShard { + id: LowShardId(1), + tile: 1, + tensor_type: TensorType::new( + [rows, AMP_COLUMN_MICRO], + Precision::F16, + Layout { + order: StorageOrder::Blocked(BlockedOrder::matrix( + AMP_INNER_BLOCK as u16, + AMP_COLUMN_MICRO as u16, + )), + tiling: TensorTiling::replicated(1), + memory_class: MemoryClass::Standard, + }, + ), + extents: vec![ + ShardExtent { + axis: 0, + start: 0, + logical_end: rows, + physical_end: AMP_INNER_BLOCK, + }, + ShardExtent { + axis: 1, + start: 0, + logical_end: AMP_COLUMN_MICRO, + physical_end: AMP_COLUMN_MICRO, + }, + ] + .into(), + definition: ShardDefinition::ExchangeStaging, + }; + let logical_view = |shard: &LowShard| ShardView { + shard: shard.id, + extents: shard + .extents + .iter() + .copied() + .map(|mut extent| { + extent.physical_end = extent.logical_end; + extent + }) + .collect(), + }; + let mappings = split_mapping_at_panel_boundaries( + &source, + logical_view(&source), + &destination, + logical_view(&destination), + ) + .unwrap_or_else(|error| panic!("case {case}, rows {rows}: {error}")); + let source_bytes = mappings + .iter() + .flat_map(|(view, _)| view_byte_spans(&source, view).unwrap()) + .map(|span| { + assert_eq!(span.offset & 0b11, 0, "case {case}, rows {rows}"); + assert_eq!(span.bytes & 0b11, 0, "case {case}, rows {rows}"); + span.bytes + }) + .sum::(); + let destination_bytes = mappings + .iter() + .flat_map(|(_, view)| view_byte_spans(&destination, view).unwrap()) + .map(|span| span.bytes) + .sum::(); + assert_eq!(source_bytes, panel_rows * AMP_COLUMN_MICRO * 2); + assert_eq!(destination_bytes, source_bytes, "case {case}, rows {rows}"); + } +} + +#[test] +fn randomized_schedules_make_kernel_operands_resident() { + let mut random = fastrand::Rng::with_seed(0x6c6f_7721); + for case in 0..CASES { + let tiles = 1_u16 << random.u32(0..=3); + let rows = u32::from(tiles) * random.u32(1..=8) * 16; + let columns = random.u32(1..=8) * 16; + let mut graph = ComputeGraph::new(); + let left = graph.host_input("left", [rows, columns]).unwrap(); + let right = graph.host_input("right", [rows, columns]).unwrap(); + let output = graph.add(left, right).unwrap(); + graph.set_outputs([output]).unwrap(); + let config = PipelineConfig::new(tiles) + .with_search_domain(PlannerSearchDomain::default().with_active_tile_counts([tiles])) + .with_input(left, format(tiles)) + .with_input(right, format(tiles)); + let mid = lower(&graph, &config, &Ipu21CostModel).unwrap(); + let low = lower_to_tiles(&mid, &config).unwrap(); + + assert_eq!(low.tiles.len(), usize::from(tiles), "case {case}"); + for tile in &low.tiles { + for work in low.work(tile) { + if let TileWorkRef::Kernel(run) = work { + crate::validate_kernel_run(run).unwrap(); + assert_eq!( + low.shards[run.output.shard.index() as usize].tile, + tile.tile + ); + assert!( + run.inputs + .iter() + .flat_map(|operand| &operand.views) + .all(|view| { + low.shards[view.shard.index() as usize].tile == tile.tile + }) + ); + } + } + } + for phase in &low.exchange_phases { + assert!( + low.tiles + .iter() + .all(|tile| contains_phase(&low, tile, phase.id)) + ); + for transfer in &phase.transfers { + assert!(transfer.destinations.iter().all(|destination| matches!( + low.shards[destination.shard.index() as usize].definition, + ShardDefinition::Value(_) + | ShardDefinition::ExchangeStaging + | ShardDefinition::LocalCopy(_) + | ShardDefinition::Staging + ))); + } + } + } +} + +#[test] +fn randomized_broadcast_adds_schedule_remote_singleton_views() { + let mut random = fastrand::Rng::with_seed(0x6272_6463); + for case in 0..CASES { + let tiles = 1_u16 << random.u32(1..=3); + let rows = u32::from(tiles) * random.u32(1..=8); + let columns = random.u32(1..=8) * 16; + let mut graph = ComputeGraph::new(); + let bias = graph.host_input("bias", [1, columns]).unwrap(); + let tensor = graph.host_input("tensor", [rows, columns]).unwrap(); + let output = graph.add(bias, tensor).unwrap(); + graph.set_outputs([output]).unwrap(); + let config = PipelineConfig::new(tiles) + .with_input(bias, format(tiles)) + .with_input(tensor, format(tiles)); + let mid = lower(&graph, &config, &Ipu21CostModel).unwrap(); + let low = lower_to_tiles(&mid, &config).unwrap(); + + assert!( + low.exchange_phases + .iter() + .any(|phase| matches!(phase.provenance.reason, WorkReason::OperatorInputs)), + "case {case}" + ); + for tile in &low.tiles { + let add = low + .work(tile) + .find_map(|work| match work { + TileWorkRef::Kernel(run) if matches!(run.kernel, TileKernelSpec::Add) => { + Some(run) + } + _ => None, + }) + .unwrap(); + assert_eq!(add.inputs[0].views[0].extents[0].logical_end, 1); + assert_eq!( + low.shards[add.inputs[0].views[0].shard.index() as usize].tile, + tile.tile + ); + } + } +} + +#[test] +fn randomized_blocked_gemms_expand_to_tile_kernel_phases() { + let mut random = fastrand::Rng::with_seed(0x6765_6d6d); + for case in 0..CASES { + let tiles = 1_u16 << random.u32(0..=3); + let rows = u32::from(tiles) * random.u32(1..=4) * 8; + let inner_blocks = random.u32(1..=4); + let column_blocks = random.u32(1..=4); + let inner = inner_blocks * 64; + let columns = column_blocks * 64; + let mut graph = ComputeGraph::new(); + let left = graph.host_input("left", [rows, inner]).unwrap(); + let right = graph.parameter("right", [inner, columns]).unwrap(); + let output = graph.gemm(left, right).unwrap(); + graph.set_outputs([output]).unwrap(); + let config = PipelineConfig::new(tiles) + .with_search_domain(PlannerSearchDomain::default().with_active_tile_counts([tiles])) + .with_input(left, format(tiles)) + .with_input(right, format(tiles)); + let mid = lower(&graph, &config, &Ipu21CostModel).unwrap(); + let low = lower_to_tiles(&mid, &config).unwrap(); + + assert!(std::mem::size_of::() <= 8); + let mut metadata = Vec::<&Arc>::new(); + for run in &low.kernel_runs { + if let Some(existing) = metadata + .iter() + .find(|existing| existing.as_ref() == run.metadata.as_ref()) + { + assert!(Arc::ptr_eq(existing, &run.metadata), "case {case}"); + } else { + metadata.push(&run.metadata); + } + } + + let gemm_destinations = low + .exchange_phases + .iter() + .filter(|phase| phase.provenance.reason == WorkReason::OperatorInput { input: 1 }) + .flat_map(|phase| &phase.transfers) + .flat_map(|transfer| &transfer.destinations) + .map(|destination| destination.shard) + .collect::>(); + let unique_gemm_staging = gemm_destinations + .iter() + .copied() + .collect::>(); + let panels_per_phase = column_blocks; + assert!( + unique_gemm_staging.len() + <= usize::from(tiles) + * usize::try_from(column_blocks.min(panels_per_phase)).unwrap(), + "case {case}" + ); + assert!(low.exchange_phases.iter().all(|phase| { + phase.provenance.operation.is_some() + && (phase.provenance.value.is_some() + || phase.provenance.reason == WorkReason::OperatorInputs) + })); + for tile in &low.tiles { + let gemms = low + .work(tile) + .filter_map(|work| match work { + TileWorkRef::Kernel(run) + if matches!(run.kernel, TileKernelSpec::Gemm { .. }) => + { + Some(run) + } + _ => None, + }) + .collect::>(); + let mut initialized_columns = std::collections::BTreeSet::new(); + for run in gemms { + assert_eq!(run.provenance.reason, WorkReason::OperatorKernel); + assert!(run.provenance.operation.is_some()); + assert!(run.provenance.value.is_some()); + let TileKernelSpec::Gemm { + mode, + inner_block: kernel_inner, + output_columns: kernel_columns, + .. + } = run.kernel + else { + unreachable!() + }; + let output_key = run + .output + .extents + .iter() + .map(|extent| (extent.start, extent.physical_end)) + .collect::>(); + assert_eq!( + mode, + if initialized_columns.insert(output_key) { + crate::GemmKernelMode::Initialize + } else { + crate::GemmKernelMode::Accumulate + }, + "case {case}" + ); + assert_eq!(run.inputs.len(), 2); + assert!(run.inputs.iter().all(|operand| operand.views.len() == 1)); + assert!( + run.inputs[0].views[0] + .extents + .iter() + .any(|extent| { extent.physical_end - extent.start == kernel_inner }) + ); + assert!( + run.output + .extents + .iter() + .any(|extent| { extent.physical_end - extent.start == kernel_columns }) + ); + } + } + } +} + +#[test] +fn randomized_odd_capacities_use_nonempty_active_tile_subsets() { + let mut random = fastrand::Rng::with_seed(0x7375_6273_6574); + for case in 0..16 { + let active_tiles = 1_u16 << random.u32(2..=5); + let capacity = active_tiles + random.u16(1..active_tiles); + let rows = u32::from(active_tiles); + let mut graph = ComputeGraph::new(); + let left = graph.host_input("left", [rows, 64]).unwrap(); + let right = graph.parameter("right", [64, 64]).unwrap(); + let output = graph.gemm(left, right).unwrap(); + graph.set_outputs([output]).unwrap(); + let config = PipelineConfig::new(capacity) + .with_automatic_input(left, Precision::F16) + .with_automatic_input(right, Precision::F16); + + let mid = lower(&graph, &config, &Ipu21CostModel).unwrap(); + let result = mid.operations.last().unwrap().results[0]; + let selected_tiles = mid.values[result.index() as usize] + .tensor_type + .format + .layout + .tiling + .tile_count; + assert!(selected_tiles <= capacity, "case {case}"); + + let low = lower_to_tiles(&mid, &config).unwrap(); + assert_eq!(low.tile_count, capacity, "case {case}"); + assert_eq!(low.outputs[0].shards.len(), usize::from(selected_tiles)); + for &shard in &low.outputs[0].shards { + assert!( + low.shards[shard.index() as usize] + .extents + .iter() + .all(|extent| extent.start < extent.logical_end), + "case {case} capacity={capacity} selected={selected_tiles} shard={:?} type={:?}", + low.shards[shard.index() as usize].extents, + mid.values[result.index() as usize].tensor_type, + ); + } + assert!(low.tiles.iter().all(|tile| tile.tile < capacity)); + } +} + +#[test] +fn randomized_resident_blocked_weights_lower_without_panel_copies() { + let mut random = fastrand::Rng::with_seed(0x7265_7369); + for _ in 0..48 { + let tiles = 1_u16 << random.u32(0..=3); + let rows = u32::from(tiles) * random.u32(1..=4); + let inner = 64 * random.u32(2..=4); + let columns = 64 * random.u32(1..=4); + let mut graph = ComputeGraph::new(); + let left = graph.host_input("left", [rows, inner]).unwrap(); + let right = graph.parameter("right", [inner, columns]).unwrap(); + let output = graph.gemm(left, right).unwrap(); + graph.set_outputs([output]).unwrap(); + let config = PipelineConfig::new(tiles) + .with_search_domain( + PlannerSearchDomain::default() + .with_active_tile_counts([tiles]) + .with_operator_precisions(crate::OperatorClass::Gemm, [Precision::F16]) + .with_weight_memory_classes([MemoryClass::Interleaved]), + ) + .with_automatic_input(left, Precision::F16) + .with_automatic_input(right, Precision::F16); + let mid = lower(&graph, &config, &Ipu21CostModel).unwrap(); + let operation = mid + .operations + .iter() + .find(|operation| matches!(operation.kind, MidOperationKind::Operator(_))) + .unwrap(); + let right_type = &mid.values[operation.inputs[1].index() as usize].tensor_type; + assert_eq!( + right_type.format.layout.order, + crate::StorageOrder::Blocked(crate::BlockedOrder::matrix( + 64, + crate::layout::AMP_COLUMN_MICRO as u16, + )) + ); + let low = lower_to_tiles(&mid, &config).unwrap(); + assert!(low.tiles.iter().all(|tile| low.work(tile).all(|work| { + !matches!(work, TileWorkRef::LocalCopy(_)) && !matches!(work, TileWorkRef::Exchange(_)) + }))); + assert!( + low.tiles + .iter() + .flat_map(|tile| low.work(tile)) + .any(|work| { + matches!( + work, + TileWorkRef::Kernel(run) + if matches!(run.kernel, TileKernelSpec::Gemm { + weights: crate::GemmWeightLoad::Interleaved, + .. + }) + ) + }) + ); + } +} + +#[test] +fn randomized_partially_sharded_weight_grids_preserve_storage() { + let mut random = fastrand::Rng::with_seed(0x7374_726d_6765_6d6d); + for _ in 0..32 { + let row_partitions = 1_u16 << random.u32(1..=2); + let inner_partitions = 1_u16 << random.u32(1..=row_partitions.ilog2()); + let column_partitions = 1_u16 << random.u32(0..=2); + let tiles = row_partitions * column_partitions; + let rows = u32::from(row_partitions) * random.u32(1..=4); + let inner_blocks = u32::from(row_partitions) * random.u32(1..=2); + let inner = inner_blocks * 64; + let columns = u32::from(column_partitions) * 64; + let mut graph = ComputeGraph::new(); + let left = graph.host_input("left", [rows, inner]).unwrap(); + let right = graph.parameter("right", [inner, columns]).unwrap(); + let output = graph.gemm(left, right).unwrap(); + graph.set_outputs([output]).unwrap(); + let left_format = TensorFormat { + precision: Precision::F16, + layout: Layout::amp_left_grid( + 64, + tiles, + row_partitions, + column_partitions, + crate::operator::GridOrder::ColumnsFast, + ), + }; + let right_format = TensorFormat { + precision: Precision::F16, + layout: Layout::block_major_matrix_storage( + crate::GemmOrientation::Normal, + 64, + 64, + column_partitions, + inner_partitions, + row_partitions / inner_partitions, + crate::MemoryClass::Standard, + ), + }; + let config = PipelineConfig::new(tiles) + .with_search_domain( + PlannerSearchDomain::default() + .with_active_tile_counts([tiles]) + .with_operator_precisions(crate::OperatorClass::Gemm, [Precision::F16]) + .with_weight_memory_classes([MemoryClass::Standard]), + ) + .with_input(left, left_format.clone()) + .with_input(right, right_format.clone()); + let mid = lower(&graph, &config, &Ipu21CostModel).unwrap(); + let low = lower_to_tiles(&mid, &config).unwrap(); + + let expected_weight_bytes = inner + .div_ceil(u32::from(inner_partitions)) + .saturating_mul(columns.div_ceil(u32::from(column_partitions))) + .saturating_mul(2); + assert!(low.inputs[1].shards.iter().all(|shard| { + crate::shard_storage_bytes(&low.shards[shard.index() as usize]) + == Ok(expected_weight_bytes) + })); + } +} + +#[test] +fn randomized_repeats_remain_structured_per_tile() { + let mut random = fastrand::Rng::with_seed(0x7265_706c); + for case in 0..CASES { + let tiles = 1_u16 << random.u32(0..=3); + let count = random.u32(1..=8); + let width = u32::from(tiles) * random.u32(1..=8); + let mut graph = ComputeGraph::new(); + let carried = graph.host_input("carried", [width, 16]).unwrap(); + let parameters = (0..count) + .map(|index| graph.parameter(format!("parameter.{index}"), [width, 16])) + .collect::, _>>() + .unwrap(); + let sequence = graph + .value_sequence("parameters", parameters.clone()) + .unwrap(); + let result = graph + .repeat(count, [carried], [], [sequence], |body, arguments| { + Ok(vec![body.add(arguments.carried[0], arguments.iterated[0])?]) + }) + .unwrap()[0]; + graph.set_outputs([result]).unwrap(); + let mut config = PipelineConfig::new(tiles).with_input(carried, format(tiles)); + for parameter in parameters { + config.inputs.insert(parameter, format(tiles)); + } + let mid = lower(&graph, &config, &Ipu21CostModel).unwrap(); + let low = lower_to_tiles(&mid, &config).unwrap(); + + for tile in &low.tiles { + let repeats = low + .work(tile) + .filter_map(|work| match work { + TileWorkRef::Repeat(repeat) => Some(repeat), + _ => None, + }) + .collect::>(); + assert_eq!(repeats.len(), 1, "case {case}"); + assert_eq!(repeats[0].count, count); + assert_eq!(repeats[0].iterated[0].inputs.len(), count as usize); + assert!(repeats[0].iterated[0].stride_bytes > 0); + assert!( + repeats[0].iterated[0] + .stride_bytes + .is_multiple_of(repeats[0].iterated[0].alignment) + ); + let carried = &repeats[0].carried[0]; + assert_eq!( + low.shards[carried.argument.index() as usize].definition, + ShardDefinition::Alias(carried.initial) + ); + assert_eq!( + low.shards[carried.yielded.index() as usize].definition, + ShardDefinition::WritableAlias(carried.argument) + ); + assert_eq!( + low.shards[carried.result.index() as usize].definition, + ShardDefinition::Alias(carried.initial) + ); + assert!( + low.work(&repeats[0].body) + .any(|work| matches!(work, TileWorkRef::Kernel(_))) + ); + } + } +} + +#[test] +fn randomized_repeats_alias_fresh_results_after_the_last_carried_use() { + let mut random = fastrand::Rng::with_seed(0x696e_706c); + for case in 0..CASES { + let tiles = 1_u16 << random.u32(0..=3); + let count = random.u32(1..=4); + let rows = u32::from(tiles) * random.u32(1..=4) * 8; + let mut graph = ComputeGraph::new(); + let carried = graph.host_input("carried", [rows, 64]).unwrap(); + let weights = (0..count) + .map(|index| graph.parameter(format!("weight.{index}"), [64, 64])) + .collect::, _>>() + .unwrap(); + let sequence = graph.value_sequence("weights", weights.clone()).unwrap(); + let result = graph + .repeat(count, [carried], [], [sequence], |body, arguments| { + Ok(vec![ + body.gemm(arguments.carried[0], arguments.iterated[0])?, + ]) + }) + .unwrap()[0]; + graph.set_outputs([result]).unwrap(); + let mut config = PipelineConfig::new(tiles).with_input(carried, format(tiles)); + for weight in weights { + config.inputs.insert(weight, format(tiles)); + } + let mid = lower(&graph, &config, &Ipu21CostModel).unwrap(); + let low = lower_to_tiles(&mid, &config).unwrap(); + for tile in &low.tiles { + let repeat = low + .work(tile) + .find_map(|work| match work { + TileWorkRef::Repeat(repeat) => Some(repeat), + _ => None, + }) + .unwrap(); + assert_eq!( + low.shards[repeat.carried[0].yielded.index() as usize].definition, + ShardDefinition::WritableAlias(repeat.carried[0].argument), + "case {case}" + ); + } + } +} + +fn contains_phase(program: &LowProgram, list: &TileWorkList, phase: ExchangePhaseId) -> bool { + program.work(list).any(|work| match work { + TileWorkRef::Exchange(candidate) => candidate == phase, + TileWorkRef::Repeat(repeat) => contains_phase(program, &repeat.body, phase), + TileWorkRef::Kernel(_) | TileWorkRef::LocalCopy(_) | TileWorkRef::Checkpoint(..) => false, + }) +} * Unmerged path crates/ipu-codegen/src/memory.rs diff --git a/crates/ipu-codegen/src/metrics.rs b/crates/ipu-codegen/src/metrics.rs new file mode 100644 index 0000000..8e8b218 --- /dev/null +++ b/crates/ipu-codegen/src/metrics.rs @@ -0,0 +1,279 @@ +//! Shared cycle and per-tile memory metrics used by operator and region planning. + +use crate::layout::MemoryClass; +use ipu_target::hardware::{HardwareMemoryConstraints, HardwareTarget}; + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct CostEstimate { + pub cycles: u64, + pub exchange_cycles: u64, + pub exchange_footprint: ExchangeFootprint, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct ExchangeFootprint { + pub phases: u64, + pub maximum_transfer_chunks_per_tile: u64, +} + +impl CostEstimate { + pub const fn exchange_row_bytes(self, target: HardwareTarget) -> u64 { + self.exchange_footprint.estimated_row_bytes(target) + } + + pub const fn sequence(self, next: Self) -> Self { + Self { + cycles: self.cycles.saturating_add(next.cycles), + exchange_cycles: self.exchange_cycles.saturating_add(next.exchange_cycles), + exchange_footprint: ExchangeFootprint { + phases: self + .exchange_footprint + .phases + .saturating_add(next.exchange_footprint.phases), + maximum_transfer_chunks_per_tile: if self + .exchange_footprint + .maximum_transfer_chunks_per_tile + > next.exchange_footprint.maximum_transfer_chunks_per_tile + { + self.exchange_footprint.maximum_transfer_chunks_per_tile + } else { + next.exchange_footprint.maximum_transfer_chunks_per_tile + }, + }, + } + } + + pub const fn repeated(self, count: u32) -> Self { + Self { + cycles: self.cycles.saturating_mul(count as u64), + exchange_cycles: self.exchange_cycles.saturating_mul(count as u64), + exchange_footprint: ExchangeFootprint { + phases: self.exchange_footprint.phases.saturating_mul(count as u64), + maximum_transfer_chunks_per_tile: self + .exchange_footprint + .maximum_transfer_chunks_per_tile, + }, + } + } +} + +impl ExchangeFootprint { + pub const fn estimated_row_bytes(self, target: HardwareTarget) -> u64 { + target + .exchange() + .estimated_row_bytes(self.phases, self.maximum_transfer_chunks_per_tile) + } +} + +/// Maximum per-tile bytes attributed to each address/load class. The classes +/// share physical tile SRAM, so feasibility must check both the individual +/// interleaved-region limit and their combined size. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct MemoryUsage { + pub standard: u64, + pub interleaved: u64, +} + +impl MemoryUsage { + pub const fn total(self) -> u64 { + self.standard.saturating_add(self.interleaved) + } + + pub(crate) fn add_class(&mut self, class: MemoryClass, bytes: u64) { + let target = match class { + MemoryClass::Standard => &mut self.standard, + MemoryClass::Interleaved => &mut self.interleaved, + }; + *target = target.saturating_add(bytes); + } + + pub(crate) fn saturating_add(self, other: Self) -> Self { + Self { + standard: self.standard.saturating_add(other.standard), + interleaved: self.interleaved.saturating_add(other.interleaved), + } + } + + pub fn fits(self, constraints: HardwareMemoryConstraints) -> bool { + self.interleaved <= constraints.interleaved_bytes && self.total() <= constraints.total_bytes + } +} + +/// Independent class maxima and the maximum simultaneous total. The allocator +/// fixes the interleaved arena boundary for the whole program, so feasibility +/// uses the sum of the class maxima even when they occur in different phases. +/// `total` remains useful for ranking the actual peak live working set. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct MemoryPeaks { + pub standard: u64, + pub interleaved: u64, + pub total: u64, + /// Persistent standard-memory estimate for generated exchange rows. + pub exchange_rows: u64, + pub maximum_standard_allocation: u64, + /// Largest amount by which one standard-addressed allocation exceeded + /// both contiguous ranges left around the interleaved region. + pub standard_contiguous_overflow: u64, +} + +impl MemoryPeaks { + pub(crate) fn observe( + &mut self, + usage: MemoryUsage, + maximum_standard_allocation: u64, + constraints: HardwareMemoryConstraints, + ) { + self.standard = self.standard.max(usage.standard); + self.interleaved = self.interleaved.max(usage.interleaved); + self.total = self.total.max(usage.total()); + self.maximum_standard_allocation = self + .maximum_standard_allocation + .max(maximum_standard_allocation); + let interleaved_boundary = self + .interleaved + .div_ceil(constraints.interleaved_element_bytes) + * constraints.interleaved_element_bytes; + let upper_standard = constraints + .interleaved_bytes + .saturating_sub(interleaved_boundary); + let contiguous_capacity = constraints.standard_fixed_bytes.max(upper_standard); + self.standard_contiguous_overflow = self + .maximum_standard_allocation + .saturating_sub(contiguous_capacity); + } + + pub fn fits(self, constraints: HardwareMemoryConstraints) -> bool { + self.fits_with_budget(constraints, 0, constraints.total_bytes) + } + + pub fn fits_with_budget( + self, + constraints: HardwareMemoryConstraints, + reserved_standard_bytes: u64, + tile_memory_budget_bytes: u64, + ) -> bool { + let partitioned_bytes = self + .standard + .saturating_add(self.aligned_interleaved_bytes(constraints)) + .saturating_add(reserved_standard_bytes); + self.interleaved <= constraints.interleaved_bytes + && partitioned_bytes <= tile_memory_budget_bytes.min(constraints.total_bytes) + && self + .standard_contiguous_overflow_with_reservation(constraints, reserved_standard_bytes) + == 0 + } + + fn aligned_interleaved_bytes(self, constraints: HardwareMemoryConstraints) -> u64 { + self.interleaved + .div_ceil(constraints.interleaved_element_bytes) + .saturating_mul(constraints.interleaved_element_bytes) + } + + pub fn standard_contiguous_overflow_with_reservation( + self, + constraints: HardwareMemoryConstraints, + reserved_standard_bytes: u64, + ) -> u64 { + let interleaved_boundary = self + .interleaved + .div_ceil(constraints.interleaved_element_bytes) + .saturating_mul(constraints.interleaved_element_bytes); + let upper_standard = constraints + .interleaved_bytes + .saturating_sub(interleaved_boundary); + let lower_standard = constraints + .standard_fixed_bytes + .saturating_sub(reserved_standard_bytes.saturating_add(self.exchange_rows)); + self.maximum_standard_allocation + .saturating_sub(lower_standard.max(upper_standard)) + } + + pub(crate) fn conservative_tensor_usage(self) -> MemoryUsage { + MemoryUsage { + standard: self.standard.saturating_sub(self.exchange_rows), + interleaved: self.interleaved, + } + } +} + +/// Storage visible at an operator boundary plus phase-local scratch. Peak is +/// the simultaneous requirement used for candidate feasibility. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct MemoryEstimate { + pub live: MemoryUsage, + pub temporary: MemoryUsage, + pub peak: MemoryUsage, + /// Largest phase-local standard-addressed buffer which must fit in one + /// contiguous standard-memory range. + pub maximum_standard_temporary_allocation: u64, +} + +impl MemoryEstimate { + pub(crate) fn peaks(self, exchange_rows: u64) -> MemoryPeaks { + MemoryPeaks { + standard: self.peak.standard, + interleaved: self.peak.interleaved, + total: self.peak.total(), + exchange_rows, + maximum_standard_allocation: self.maximum_standard_temporary_allocation, + standard_contiguous_overflow: 0, + } + } +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct PlanMetrics { + pub cost: CostEstimate, + pub memory: M, +} + +pub type OperationMetrics = PlanMetrics; +pub type RegionMetrics = PlanMetrics; + +impl PlanMetrics { + fn pareto_dimensions(self) -> [u64; 8] { + [ + self.cost.cycles, + self.cost.exchange_cycles, + self.memory.standard, + self.memory.interleaved, + self.memory.total, + self.memory.maximum_standard_allocation, + self.memory.standard_contiguous_overflow, + self.memory.exchange_rows, + ] + } + + pub(crate) fn dominates(self, other: Self) -> bool { + let left = self.pareto_dimensions(); + let right = other.pareto_dimensions(); + left.iter().zip(right).all(|(left, right)| *left <= right) + && left.iter().zip(right).any(|(left, right)| *left < right) + } +} + +pub(crate) fn pareto_frontier( + candidates: impl IntoIterator, + class: impl Fn(&T) -> K, + objective: impl Fn(&T) -> RegionMetrics, +) -> (Vec, usize) { + let mut frontier = Vec::new(); + let mut dominated = 0; + for candidate in candidates { + let candidate_class = class(&candidate); + let candidate_objective = objective(&candidate); + if frontier.iter().any(|kept| { + class(kept) == candidate_class && objective(kept).dominates(candidate_objective) + }) { + dominated += 1; + continue; + } + let before = frontier.len(); + frontier.retain(|kept| { + class(kept) != candidate_class || !candidate_objective.dominates(objective(kept)) + }); + dominated += before - frontier.len(); + frontier.push(candidate); + } + (frontier, dominated) +} diff --git a/crates/ipu-codegen/src/mid/consumers.rs b/crates/ipu-codegen/src/mid/consumers.rs new file mode 100644 index 0000000..5fb5794 --- /dev/null +++ b/crates/ipu-codegen/src/mid/consumers.rs @@ -0,0 +1,89 @@ +use crate::config::{OperatorClass, PipelineConfig}; +use crate::graph::{Operation, OperationKind, TensorShape, ValueId}; +use crate::layout::{AMP_INNER_BLOCK, Layout}; +use crate::operator::Precision; + +pub(super) fn direct_consumer_layouts( + source: &[Operation], + operation_index: usize, + result: ValueId, + output: &TensorShape, + config: &PipelineConfig, +) -> Vec { + if !config + .search_domain + .permits_precision(OperatorClass::Attention, Precision::F16) + { + return Vec::new(); + } + let Ok(streams) = u16::try_from(output.0.first().copied().unwrap_or(0)) else { + return Vec::new(); + }; + if streams == 0 { + return Vec::new(); + } + let Some(&rows) = output.0.get(1) else { + return Vec::new(); + }; + let query_partitions = u16::try_from(rows) + .unwrap_or(u16::MAX) + .min(config.tile_count / streams); + let key_partitions = u16::try_from(rows.div_ceil(AMP_INNER_BLOCK)) + .unwrap_or(u16::MAX) + .min(config.tile_count / streams); + let mut layouts = Vec::new(); + for consumer in &source[operation_index + 1..] { + for input_index in consumer + .inputs + .iter() + .enumerate() + .filter_map(|(index, &input)| (input == result).then_some(index)) + { + let layout = match (&consumer.kind, input_index) { + (OperationKind::FlashAttention(_), 0) if query_partitions != 0 => { + Some(Layout::attention_query(streams, query_partitions)) + } + (OperationKind::FlashAttention(_), 1) if key_partitions != 0 => { + Some(Layout::attention_key(streams, key_partitions)) + } + (OperationKind::FlashAttention(_), 2) if key_partitions != 0 => Some( + Layout::attention_block_major_key_value(streams, key_partitions), + ), + _ => None, + }; + if let Some(layout) = layout + && !layouts.contains(&layout) + { + layouts.push(layout); + } + } + } + layouts +} + +pub(super) fn operator_accepts_input_layout( + operation: &OperationKind, + input_index: usize, + config: &PipelineConfig, +) -> bool { + match operation { + OperationKind::Gelu => { + input_index == 0 + && !config + .search_domain + .precisions(OperatorClass::Gelu) + .is_empty() + } + OperationKind::Add(_) => { + input_index < 2 + && !config + .search_domain + .precisions(OperatorClass::Add) + .is_empty() + } + OperationKind::SplitHeads(_) => input_index == 0, + OperationKind::Gemm(_) | OperationKind::FlashAttention(_) | OperationKind::Repeat(_) => { + false + } + } +} diff --git a/crates/ipu-codegen/src/mid/gemm.rs b/crates/ipu-codegen/src/mid/gemm.rs new file mode 100644 index 0000000..6f6bf31 --- /dev/null +++ b/crates/ipu-codegen/src/mid/gemm.rs @@ -0,0 +1,1237 @@ +use super::{StorageOrderCompatibility, storage_order_compatibility}; +use crate::config::{OperatorClass, PipelineConfig, PlannerSearchDomain}; +use crate::cost::{CostModel, operator_memory_estimate, parallel_reduction_preselection_metrics}; +use crate::graph::{GemmOptions, Operation, OperationKind, TensorShape, ValueId}; +use crate::layout::{ + AMP_COLUMN_MICRO, AMP_INNER_BLOCK, AMP_NARROW_OUTPUT_COLUMN_BLOCK, AMP_OUTPUT_COLUMN_BLOCK, + AMP_WIDE_OUTPUT_COLUMN_BLOCK, Layout, MemoryClass, Padding, StorageOrder, TensorAxis, + TensorFormat, TensorType, +}; +use crate::metrics::{CostEstimate, RegionMetrics, pareto_frontier}; +use crate::operator::*; +use crate::{GemmMap, OperatorSchedule, ScheduleStep, ScheduleValue}; +use std::collections::{BTreeMap, BTreeSet}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) struct GroupedOutputLayout { + pub groups: u16, + pub physical_lane_multiple: u32, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +struct ParallelGridCandidate { + metrics: RegionMetrics, + grid: GemmGrid, + physical_column_groups: u16, + grouped: bool, +} + +pub(super) fn gemm_seed_plans_for_tile_count( + options: GemmOptions, + tile_count: u16, + domain: &PlannerSearchDomain, +) -> Vec { + let candidates = (1..=tile_count) + .rev() + .filter(|columns| tile_count.is_multiple_of(*columns)) + .flat_map(|columns| { + let rows = tile_count / columns; + let geometry = GemmGeometry { + block: GemmBlockShape { + inner: AMP_INNER_BLOCK, + output_columns: AMP_OUTPUT_COLUMN_BLOCK, + }, + orientation: GemmOrientation::Normal, + compute: GemmGrid { + rows, + columns, + inner: 1, + }, + result: GemmResultGrid { rows, columns }, + order: GridOrder::ColumnsFast, + }; + let mut grid = Vec::new(); + let mut placements = Vec::new(); + for &precision in domain.precisions(OperatorClass::Gemm) { + let Some(left_tail) = gemm_left_access_tail(precision) else { + continue; + }; + for &memory_class in &domain.weight_memory_classes { + if gemm_supports_weight_memory(precision, memory_class) { + placements.push(( + precision, + left_tail, + AmpWeightPlacement::resident(memory_class), + )); + } + } + } + if rows > 1 { + for &precision in domain.precisions(OperatorClass::Gemm) { + let Some(left_tail) = gemm_left_access_tail(precision) else { + continue; + }; + for &memory_class in &domain.weight_memory_classes { + if gemm_supports_weight_memory(precision, memory_class) { + placements.push(( + precision, + left_tail, + AmpWeightPlacement::sharded(rows, memory_class), + )); + } + } + } + } + if rows > 2 + && rows.is_multiple_of(2) + && domain.permits_precision(OperatorClass::Gemm, Precision::F16) + && domain.permits_weight_memory(MemoryClass::Interleaved) + { + placements.push(( + Precision::F16, + 16, + AmpWeightPlacement::sharded(2, MemoryClass::Interleaved), + )); + } + for (precision, left_tail, weights) in placements { + for &output_columns in amp_output_column_blocks(precision) { + if output_columns < AMP_OUTPUT_COLUMN_BLOCK + && !(weights.inner_partitions == 1 + && weights.memory_class == MemoryClass::Interleaved) + { + continue; + } + let mut geometry = geometry; + geometry.block.output_columns = output_columns; + let candidate = amp_grid_gemm_plan( + options, + precision, + left_tail, + geometry, + output_columns, + weights, + ); + grid.push(candidate.clone()); + if columns == 1 + && rows == tile_count + && output_columns > AMP_OUTPUT_COLUMN_BLOCK + && weights == AmpWeightPlacement::resident(MemoryClass::Standard) + { + grid.push(amp_grid_gemm_plan( + options, + precision, + left_tail, + geometry, + AMP_OUTPUT_COLUMN_BLOCK, + weights, + )); + } + if precision == Precision::F16 && weights.memory_class == MemoryClass::Standard + { + let mut staged = candidate; + staged.requirements.inputs[1].local_staging = + LocalOperandStaging::Staged(MemoryClass::Interleaved); + grid.push(staged); + } + } + } + grid + }) + .collect::>(); + let mut unique = Vec::with_capacity(candidates.len()); + for candidate in candidates { + if !unique.contains(&candidate) { + unique.push(candidate); + } + } + unique +} + +const fn gemm_left_access_tail(precision: Precision) -> Option { + match precision { + Precision::F16 => Some(16), + Precision::F32 => Some(32), + Precision::F8F143 { .. } => None, + } +} + +const fn gemm_supports_weight_memory(precision: Precision, memory_class: MemoryClass) -> bool { + matches!( + (precision, memory_class), + (Precision::F16, _) | (Precision::F32, MemoryClass::Standard) + ) +} + +pub(super) fn amp_grid_gemm_plan( + options: GemmOptions, + precision: Precision, + left_tail: u32, + geometry: GemmGeometry, + storage_column_block: u32, + weights: AmpWeightPlacement, +) -> OperatorSchedule { + let inner = u16::try_from(geometry.block.inner).unwrap_or(0); + let grid = geometry.result; + let right_layout = match (weights.inner_partitions, weights.memory_class) { + (1, MemoryClass::Standard) => Layout::block_major_matrix_grid( + inner, + storage_column_block, + grid.tile_count(), + grid.rows, + grid.columns, + geometry.order, + ), + (inner_partitions, memory_class) => Layout::block_major_matrix_storage( + GemmOrientation::Normal, + inner, + storage_column_block, + grid.columns, + inner_partitions, + grid.rows / inner_partitions, + memory_class, + ), + }; + let operator = MidOperator::Gemm { + options, + multiply: precision, + accumulate: gemm_accumulation_precision(precision), + }; + let gemm = gemm_map(operator, geometry); + let output = if geometry.compute.inner > 1 { + ScheduleValue::Temporary(0) + } else { + ScheduleValue::Output + }; + let mut steps = vec![ScheduleStep::Gemm(GemmMap { output, ..gemm })]; + if geometry.compute.inner > 1 { + steps.push(ScheduleStep::Reduce { + input: output, + output: ScheduleValue::Output, + staging: ReductionStaging::Complete, + }); + } + OperatorSchedule { + operator, + steps, + requirements: OperatorRequirements { + inputs: vec![ + OperandRequirement::new( + TensorFormat { + precision, + layout: Layout::amp_left_grid( + inner, + grid.tile_count(), + grid.rows, + grid.columns, + geometry.order, + ), + }, + 32, + ) + .with_access_tail(left_tail) + .with_materialization(OperandMaterialization::DispatchSlices), + OperandRequirement::new( + TensorFormat { + precision, + layout: right_layout, + }, + 32, + ) + .with_local_staging(LocalOperandStaging::Direct( + if precision == Precision::F16 && weights.inner_partitions > 1 { + MemoryClass::Interleaved + } else { + weights.memory_class + }, + )), + ], + output: OperandRequirement::new( + TensorFormat { + precision, + layout: if precision == Precision::F16 { + Layout::amp_left_result_grid( + GemmOrientation::Normal, + storage_column_block, + grid.tile_count(), + grid.rows, + grid.columns, + geometry.order, + ) + } else { + Layout::amp_output_grid( + GemmOrientation::Normal, + storage_column_block, + grid.tile_count(), + grid.rows, + grid.columns, + geometry.order, + ) + }, + }, + 32, + ), + output_aliasing: OutputAliasing::Fresh, + memory_space: MemorySpaceRequirements::default() + .with_distinct_elements([MemoryOperand::Output, MemoryOperand::Input(0)]), + }, + } +} + +pub(super) fn grouped_output_layout( + source: &[Operation], + operation_index: usize, + operation: &Operation, + output: &TensorShape, + value_uses: &BTreeMap, +) -> Option { + let result = *operation.results.first()?; + if value_uses.get(&result).copied() != Some(1) { + return None; + } + let consumer = source[operation_index + 1..] + .iter() + .find(|candidate| candidate.inputs.contains(&result))?; + let OperationKind::SplitHeads(options) = consumer.kind else { + return None; + }; + let groups = u16::try_from(options.heads).ok()?; + let width = *output.0.last()?; + (groups != 0 && width.is_multiple_of(u32::from(groups))).then_some(GroupedOutputLayout { + groups, + physical_lane_multiple: AMP_COLUMN_MICRO, + }) +} + +pub(super) fn gemm_map(operator: MidOperator, geometry: GemmGeometry) -> GemmMap { + let MidOperator::Gemm { + multiply, + accumulate, + .. + } = operator + else { + unreachable!("blocked GEMM schedule requires a GEMM operator") + }; + GemmMap { + inputs: [ScheduleValue::Input(0), ScheduleValue::Input(1)], + output: ScheduleValue::Output, + kernel: GemmKernelFamily { + multiply, + accumulate, + weights: GemmWeightLoad::Standard, + }, + geometry, + } +} + +#[derive(Clone, Copy, PartialEq, Eq)] +pub(super) struct AmpWeightPlacement { + pub inner_partitions: u16, + pub memory_class: MemoryClass, +} + +impl AmpWeightPlacement { + pub const fn resident(memory_class: MemoryClass) -> Self { + Self::sharded(1, memory_class) + } + + pub const fn sharded(inner_partitions: u16, memory_class: MemoryClass) -> Self { + Self { + inner_partitions, + memory_class, + } + } +} + +fn amp_output_column_blocks(precision: Precision) -> &'static [u32] { + match precision { + Precision::F16 => &[ + AMP_OUTPUT_COLUMN_BLOCK, + AMP_WIDE_OUTPUT_COLUMN_BLOCK, + AMP_NARROW_OUTPUT_COLUMN_BLOCK, + ], + Precision::F32 | Precision::F8F143 { .. } => &[AMP_OUTPUT_COLUMN_BLOCK], + } +} + +pub(super) const fn gemm_accumulation_precision(precision: Precision) -> AccumulationPrecision { + match precision { + Precision::F16 | Precision::F8F143 { .. } => AccumulationPrecision::F16, + Precision::F32 => AccumulationPrecision::F32, + } +} + +fn balance_parallel_gemm_columns(layout: &mut Layout, axis: TensorAxis) { + if let Some(columns) = layout + .tiling + .axes + .iter_mut() + .find(|tiling| tiling.axis == axis) + { + columns.block_size = AMP_COLUMN_MICRO; + columns.padding_multiple = AMP_COLUMN_MICRO; + columns.padding = Padding::Zero; + } +} + +fn apply_grouped_output_layout( + candidate: &mut OperatorSchedule, + grouping: GroupedOutputLayout, +) -> bool { + if candidate.requirements.output.format.precision != Precision::F16 + || grouping.groups == 0 + || grouping.physical_lane_multiple == 0 + { + return false; + } + let configure = |layout: &mut Layout| { + let Some(axis) = layout + .tiling + .axes + .iter_mut() + .find(|axis| axis.axis == TensorAxis::FromEnd(1)) + else { + return false; + }; + if !axis.partitions.is_multiple_of(grouping.groups) { + return false; + } + axis.block_size = grouping.physical_lane_multiple; + axis.padding_multiple = grouping.physical_lane_multiple; + axis.padding_groups = grouping.groups; + axis.shard_padding_multiple = 1; + axis.padding = Padding::Zero; + true + }; + configure(&mut candidate.requirements.inputs[1].format.layout) + && configure(&mut candidate.requirements.output.format.layout) +} + +fn pad_axis_to_f16_exchange_word(layout: &mut Layout, axis: TensorAxis) { + if let Some(tiling) = layout + .tiling + .axes + .iter_mut() + .find(|tiling| tiling.axis == axis) + { + tiling.block_size = tiling.block_size.div_ceil(2) * 2; + tiling.padding_multiple = tiling.padding_multiple.div_ceil(2) * 2; + tiling.padding = Padding::Zero; + } +} + +#[allow(clippy::too_many_arguments)] +pub(super) fn gemm_plans( + options: GemmOptions, + inputs: &[TensorType], + parameter_inputs: &[bool], + output: &TensorShape, + config: &PipelineConfig, + costs: &impl CostModel, + distributed_result_is_useful: bool, + constraint: Option<&GemmPlanConstraint>, + grouped_output: Option, +) -> Vec { + if options != GemmOptions::default() { + return Vec::new(); + } + let mut plans = Vec::new(); + for &tile_count in &config.resolved_active_tile_counts { + for seed in gemm_seed_plans_for_tile_count(options, tile_count, &config.search_domain) { + let mut variants = vec![seed.clone()]; + variants.extend(parallel_reduction_plans( + &seed, + inputs, + output, + config, + costs, + distributed_result_is_useful, + constraint, + grouped_output, + )); + for (input_index, _) in parameter_inputs + .iter() + .enumerate() + .filter(|(_, parameter)| **parameter) + { + let additions = variants + .iter() + .flat_map(|variant| { + independent_parameter_storage(variant, inputs, input_index, config) + }) + .filter(|independent| !variants.contains(independent)) + .collect::>(); + variants.extend(additions); + } + for plan in variants { + if !plan.supports(inputs, output) { + continue; + } + if !plans.contains(&plan) { + plans.push(plan); + } + } + } + } + plans +} + +pub(super) fn gemm_plan_matches( + constraint: &GemmPlanConstraint, + schedule: &OperatorSchedule, + inputs: &[OperandRequirement], +) -> bool { + let Some(ScheduleStep::Gemm(plan)) = schedule.steps.first() else { + return false; + }; + if plan.geometry.compute.inner < 2 { + return false; + } + let weight_index = plan.geometry.orientation.physical_right_input(); + let Some(weight) = inputs.get(weight_index) else { + return false; + }; + let reduction_staging = schedule.steps.get(1).and_then(|step| match step { + ScheduleStep::Reduce { staging, .. } => Some(*staging), + _ => None, + }); + plan.geometry == constraint.geometry + && reduction_staging == constraint.reduction_staging + && weight.format.layout.memory_class == constraint.weight_memory_class + && weight.local_staging == constraint.local_weight_staging +} + +pub(super) fn independent_parameter_storage( + candidate: &OperatorSchedule, + inputs: &[TensorType], + input_index: usize, + config: &PipelineConfig, +) -> Vec { + let Some(ScheduleStep::Gemm(gemm)) = candidate.steps.first() else { + return Vec::new(); + }; + let Some(requirement) = candidate.requirements.inputs.get(input_index) else { + return Vec::new(); + }; + let StorageOrder::Blocked(order) = requirement.format.layout.order else { + return Vec::new(); + }; + if !order.is_matrix() { + return Vec::new(); + } + let inner_block = order.block_shape[0]; + let Some(input) = inputs.get(input_index) else { + return Vec::new(); + }; + let rank = input.shape.0.len(); + let Some(inner_axis) = rank.checked_sub(2) else { + return Vec::new(); + }; + let Some(&inner) = input.shape.0.get(inner_axis) else { + return Vec::new(); + }; + let Some(&columns) = input.shape.0.last() else { + return Vec::new(); + }; + let inner_blocks = inner.div_ceil(u32::from(inner_block)); + let output_column_block = gemm.geometry.block.output_columns; + if output_column_block < AMP_OUTPUT_COLUMN_BLOCK { + return Vec::new(); + } + let column_blocks = columns.div_ceil(output_column_block); + let mut storage_grids = (1..=inner_blocks.min(u32::from(config.tile_count))) + .flat_map(|inner_partitions| { + let maximum_columns = (u32::from(config.tile_count) / inner_partitions) + .min(column_blocks) + .min(u32::from(u16::MAX)); + (1..=maximum_columns).map(move |column_partitions| { + let panels_per_shard = inner_blocks + .div_ceil(inner_partitions) + .saturating_mul(column_blocks.div_ceil(column_partitions)); + let used = inner_partitions.saturating_mul(column_partitions); + ( + panels_per_shard, + u32::MAX - used, + column_partitions, + inner_partitions, + ) + }) + }) + .collect::>(); + storage_grids.sort_unstable(); + storage_grids + .first() + .and_then(|&(_, _, column_partitions, inner_partitions)| { + Some(( + u16::try_from(column_partitions).ok()?, + u16::try_from(inner_partitions).ok()?, + )) + }) + .into_iter() + .map(|(column_partitions, inner_partitions)| { + let mut independent = candidate.clone(); + independent.requirements.inputs[input_index].format.layout = + Layout::block_major_matrix_storage( + GemmOrientation::Normal, + inner_block, + output_column_block, + column_partitions, + inner_partitions, + 1, + requirement.format.layout.memory_class, + ); + independent + }) + .collect() +} + +pub(super) fn parallel_reduction_plans( + candidate: &OperatorSchedule, + inputs: &[TensorType], + output: &TensorShape, + config: &PipelineConfig, + costs: &impl CostModel, + distributed_result_is_useful: bool, + constraint: Option<&GemmPlanConstraint>, + grouped_output: Option, +) -> Vec { + [GemmOrientation::Normal, GemmOrientation::Swapped] + .into_iter() + .flat_map(|orientation| { + parallel_reduction_plans_for_orientation( + &candidate, + inputs, + output, + config, + costs, + orientation, + distributed_result_is_useful, + constraint, + grouped_output, + ) + }) + .collect() +} + +fn parallel_reduction_plans_for_orientation( + candidate: &OperatorSchedule, + inputs: &[TensorType], + output: &TensorShape, + config: &PipelineConfig, + costs: &impl CostModel, + orientation: GemmOrientation, + distributed_result_is_useful: bool, + constraint: Option<&GemmPlanConstraint>, + grouped_output: Option, +) -> Vec { + let Some(ScheduleStep::Gemm(plan)) = candidate.steps.first() else { + return Vec::new(); + }; + let plan = *plan; + if plan.geometry.compute.inner != 1 { + return Vec::new(); + } + let output_column_block = plan.geometry.block.output_columns; + if !matches!( + candidate.operator, + MidOperator::Gemm { + multiply: Precision::F16, + .. + } + ) || output_column_block != AMP_OUTPUT_COLUMN_BLOCK + { + return Vec::new(); + } + let [left, right] = inputs else { + return Vec::new(); + }; + let rank = left.shape.0.len(); + if rank < 2 || right.shape.0.len() < 2 { + return Vec::new(); + } + let Some(&inner) = left.shape.0.last() else { + return Vec::new(); + }; + let Some(&normal_columns) = right.shape.0.last() else { + return Vec::new(); + }; + let normal_rows = left.shape.0[left.shape.0.len() - 2]; + let [rows, columns] = orientation.physical_order([normal_rows, normal_columns]); + // Generate the shape-specialized family once from the ordinary C64 seed. + // Each grid chooses the exact padded local K and C extents, so one tile + // call traverses all of its AMP micro-groups without fixed K64/C64 + // boundaries. + let tile_count = candidate + .requirements + .output + .format + .layout + .tiling + .tile_count; + let column_groups = columns.div_ceil(AMP_COLUMN_MICRO); + let inner_groups = inner.div_ceil(AMP_COLUMN_MICRO); + let Ok(inner_groups) = u16::try_from(inner_groups) else { + return Vec::new(); + }; + let Ok(column_groups) = u16::try_from(column_groups) else { + return Vec::new(); + }; + let grouped_column_groups = grouped_output.and_then(|grouping| { + let groups = u32::from(grouping.groups); + (groups != 0 && columns.is_multiple_of(groups)).then(|| { + let columns_per_group = columns / groups; + columns_per_group + .div_ceil(grouping.physical_lane_multiple) + .saturating_mul(groups) + }) + }); + let grouped_column_groups = grouped_column_groups + .and_then(|groups| u16::try_from(groups).ok()) + .filter(|groups| *groups >= column_groups); + let output_seed_partitions = candidate + .requirements + .output + .format + .layout + .tiling + .axes + .iter() + .find(|axis| axis.axis == TensorAxis::FromEnd(1)) + .map(|axis| axis.partitions); + if output_seed_partitions != Some(tile_count) + || candidate.requirements.inputs[1].format.layout.memory_class != MemoryClass::Standard + || candidate.requirements.inputs[1] + .local_staging + .stages_local() + { + return Vec::new(); + } + let mut grids = Vec::new(); + for inner_partitions in 2..=inner_groups.min(tile_count) { + let maximum_columns = grouped_column_groups + .unwrap_or(column_groups) + .min(tile_count / inner_partitions); + for column_partitions in 1..=maximum_columns { + let grouped_options = [ + (column_partitions <= column_groups).then_some((false, column_groups)), + grouped_output.and_then(|grouping| { + let physical = grouped_column_groups?; + column_partitions + .is_multiple_of(grouping.groups) + .then_some((true, physical)) + }), + ]; + for (grouped, physical_column_groups) in grouped_options.into_iter().flatten() { + let row_partitions = (tile_count / inner_partitions / column_partitions) + .min(u16::try_from(rows).unwrap_or(u16::MAX)); + let used_tiles = row_partitions + .saturating_mul(column_partitions) + .saturating_mul(inner_partitions); + if used_tiles < tile_count.div_ceil(2) || u32::from(row_partitions) > rows { + continue; + } + let local_columns = + u32::from(physical_column_groups).div_ceil(u32::from(column_partitions)); + let local_inner = u32::from(inner_groups).div_ceil(u32::from(inner_partitions)); + if u32::from(inner_partitions - 1).saturating_mul(local_inner) + >= u32::from(inner_groups) + { + continue; + } + let grid = GemmGrid { + rows: row_partitions, + columns: column_partitions, + inner: inner_partitions, + }; + let block = GemmBlockShape { + inner: local_inner.saturating_mul(AMP_COLUMN_MICRO), + output_columns: local_columns.saturating_mul(AMP_COLUMN_MICRO), + }; + let Some(metrics) = parallel_reduction_preselection_metrics( + config.target, + block, + grid, + orientation, + inputs, + candidate.requirements.output.format.precision, + ) else { + continue; + }; + let constraints = config.target.memory_constraints(); + if !metrics.memory.fits(constraints) { + continue; + } + grids.push(ParallelGridCandidate { + metrics, + grid, + physical_column_groups, + grouped, + }); + } + } + } + let generated_grids = grids.len(); + let grids = if let Some(constraint) = constraint { + grids + .into_iter() + .filter(|grid| { + orientation == constraint.geometry.orientation + && grid.grid == constraint.geometry.compute + }) + .collect::>() + } else { + // Inner partitioning and grouped outputs select different lowering + // families. Within each family the shared metrics vocabulary retains + // cycle, exchange, and memory tradeoffs for precise evaluation below. + let (mut frontier, _) = pareto_frontier( + grids, + |grid| (grid.grid.inner, grid.grouped), + |grid| grid.metrics, + ); + frontier.sort_by_key(|grid| { + ( + grid.metrics.cost.cycles, + grid.metrics.cost.exchange_cycles, + grid.metrics.memory.total, + grid.grid, + grid.grouped, + ) + }); + frontier + }; + let proxy_frontier_grids = grids.len(); + let mut variants = Vec::new(); + for grid in grids { + let ParallelGridCandidate { + grid: + GemmGrid { + rows: row_partitions, + columns: column_partitions, + inner: inner_partitions, + }, + physical_column_groups, + grouped, + .. + } = grid; + let used_tiles = row_partitions + .saturating_mul(column_partitions) + .saturating_mul(inner_partitions); + let kernel_inner_block = u32::from(inner_groups) + .div_ceil(u32::from(inner_partitions)) + .saturating_mul(AMP_COLUMN_MICRO); + let kernel_output_columns = u32::from(physical_column_groups) + .div_ceil(u32::from(column_partitions)) + .saturating_mul(AMP_COLUMN_MICRO); + let Ok(kernel_inner_block_u16) = u16::try_from(kernel_inner_block) else { + continue; + }; + for &memory_class in &config.search_domain.weight_memory_classes { + let mut variant = candidate.clone(); + match orientation { + GemmOrientation::Normal => { + variant.requirements.inputs[0].format.layout = Layout::amp_left_parallel_grid( + orientation, + kernel_inner_block_u16, + used_tiles, + row_partitions, + column_partitions, + inner_partitions, + ); + variant.requirements.inputs[1].format.layout = + Layout::block_major_matrix_storage( + orientation, + kernel_inner_block_u16, + kernel_output_columns, + column_partitions, + inner_partitions, + 1, + memory_class, + ); + balance_parallel_gemm_columns( + &mut variant.requirements.inputs[1].format.layout, + TensorAxis::FromEnd(1), + ); + variant.requirements.output.format.layout = Layout::amp_left_result_grid( + orientation, + kernel_output_columns, + row_partitions.saturating_mul(column_partitions), + row_partitions, + column_partitions, + GridOrder::ColumnsFast, + ); + balance_parallel_gemm_columns( + &mut variant.requirements.output.format.layout, + TensorAxis::FromEnd(1), + ); + } + GemmOrientation::Swapped => { + let mut physical_left = variant.requirements.inputs[1].clone(); + physical_left.format.layout = Layout::amp_left_parallel_grid( + orientation, + kernel_inner_block_u16, + used_tiles, + row_partitions, + column_partitions, + inner_partitions, + ); + physical_left.materialization = OperandMaterialization::DispatchSlices; + let mut physical_right = variant.requirements.inputs[0].clone(); + physical_right.format.layout = Layout::block_major_matrix_storage( + orientation, + kernel_inner_block_u16, + kernel_output_columns, + column_partitions, + inner_partitions, + row_partitions, + memory_class, + ); + balance_parallel_gemm_columns( + &mut physical_right.format.layout, + TensorAxis::FromEnd(2), + ); + physical_right.materialization = OperandMaterialization::Complete; + variant.requirements.inputs = vec![physical_right, physical_left]; + variant.requirements.output.format.layout = Layout::amp_left_result_grid( + orientation, + kernel_output_columns, + row_partitions.saturating_mul(column_partitions), + row_partitions, + column_partitions, + GridOrder::ColumnsFast, + ); + balance_parallel_gemm_columns( + &mut variant.requirements.output.format.layout, + TensorAxis::FromEnd(2), + ); + variant.requirements.memory_space = MemorySpaceRequirements::default() + .with_distinct_elements([MemoryOperand::Output, MemoryOperand::Input(1)]); + } + } + if let Some(ScheduleStep::Gemm(plan)) = variant.steps.first_mut() { + plan.kernel.weights = if memory_class == MemoryClass::Interleaved { + GemmWeightLoad::Interleaved + } else { + GemmWeightLoad::Standard + }; + plan.geometry = GemmGeometry { + block: GemmBlockShape { + inner: kernel_inner_block, + output_columns: kernel_output_columns, + }, + orientation, + compute: GemmGrid { + rows: row_partitions, + columns: column_partitions, + inner: inner_partitions, + }, + result: GemmResultGrid { + rows: row_partitions, + columns: column_partitions, + }, + order: GridOrder::ColumnsFast, + }; + plan.output = ScheduleValue::Temporary(0); + } + variant.steps.push(ScheduleStep::Reduce { + input: ScheduleValue::Temporary(0), + output: ScheduleValue::Output, + staging: ReductionStaging::Complete, + }); + let physical_right_index = orientation.physical_right_input(); + let local_staging_options: &[_] = match orientation { + GemmOrientation::Normal => &[LocalOperandStaging::Direct(MemoryClass::Interleaved)], + GemmOrientation::Swapped => &[ + LocalOperandStaging::Direct(MemoryClass::Interleaved), + LocalOperandStaging::Staged(MemoryClass::Interleaved), + ], + }; + let mut result_layout_variants = Vec::new(); + let maximum_result_rows = u16::try_from(rows / u32::from(row_partitions)) + .unwrap_or(u16::MAX) + .min(inner_partitions); + let maximum_result_columns = u16::try_from( + columns + .div_ceil(AMP_COLUMN_MICRO) + .checked_div(u32::from(column_partitions)) + .unwrap_or(0), + ) + .unwrap_or(u16::MAX) + .min(inner_partitions); + let mut result_partition_options = vec![(1, 1)]; + if distributed_result_is_useful { + if inner_partitions <= maximum_result_rows { + result_partition_options.push((inner_partitions, 1)); + } else if inner_partitions <= maximum_result_columns { + result_partition_options.push((1, inner_partitions)); + } + } + for (result_row_partitions, result_column_partitions) in result_partition_options { + let result_rows = row_partitions.saturating_mul(result_row_partitions); + let result_columns = column_partitions.saturating_mul(result_column_partitions); + let result_column_block = if result_column_partitions > 1 { + AMP_COLUMN_MICRO + } else { + kernel_output_columns + }; + for grid_order in [GridOrder::ColumnsFast, GridOrder::RowsFast] { + if grid_order == GridOrder::RowsFast + && ((result_row_partitions, result_column_partitions) != (1, 1) + || result_rows == 1 + || result_columns == 1) + { + continue; + } + let mut result_variant = variant.clone(); + if let Some(ScheduleStep::Gemm(plan)) = result_variant.steps.first_mut() + && plan.geometry.compute.inner > 1 + { + plan.geometry.result.rows = result_rows; + plan.geometry.result.columns = result_columns; + plan.geometry.order = grid_order; + } + let mut result_layout = Layout::amp_left_result_grid( + orientation, + result_column_block, + result_rows.saturating_mul(result_columns), + result_rows, + result_columns, + grid_order, + ); + let physical_column_axis = orientation.column_axis(); + balance_parallel_gemm_columns(&mut result_layout, physical_column_axis); + result_variant.requirements.output.format.layout = result_layout; + result_layout_variants.push(result_variant); + } + } + let mut layout_variants = Vec::new(); + for mut result_layout in result_layout_variants { + let physical_row_axis = orientation.row_axis(); + let physical_rows = + [normal_rows, normal_columns][orientation.physical_left_input()]; + let physical_left_index = orientation.physical_left_input(); + let result_rows = result_layout + .requirements + .output + .format + .layout + .tiling + .axes + .iter() + .find(|axis| axis.axis == physical_row_axis) + .map_or(row_partitions, |axis| axis.partitions); + // Exchange moves whole 32-bit words. Give every F16 shard an + // even physical-row grain rather than allowing a later layout + // conversion to discover an unsendable two-byte tail. + if u32::from(result_rows) > physical_rows.div_ceil(2) { + continue; + } + pad_axis_to_f16_exchange_word( + &mut result_layout.requirements.inputs[physical_left_index] + .format + .layout, + physical_row_axis, + ); + pad_axis_to_f16_exchange_word( + &mut result_layout.requirements.output.format.layout, + physical_row_axis, + ); + layout_variants.push(result_layout); + } + for mut layout_variant in layout_variants { + if grouped + && !grouped_output.is_some_and(|grouping| { + apply_grouped_output_layout(&mut layout_variant, grouping) + }) + { + continue; + } + for &local_staging in local_staging_options { + let mut staged = layout_variant.clone(); + staged.requirements.inputs[physical_right_index].local_staging = local_staging; + variants.push(staged.clone()); + if let Some(ScheduleStep::Reduce { staging, .. }) = staged.steps.get_mut(1) { + *staging = ReductionStaging::Streamed; + } + variants.push(staged); + } + } + } + } + let generated_variants = variants.len(); + let generated_grouped_variants = variants + .iter() + .filter(|candidate| { + candidate + .requirements + .output + .format + .layout + .tiling + .axes + .iter() + .any(|axis| axis.padding_groups > 1) + }) + .count(); + let retained = if let Some(constraint) = constraint { + variants + .into_iter() + .filter(|candidate| { + gemm_plan_matches(constraint, candidate, &candidate.requirements.inputs) + }) + .collect::>() + } else { + retain_precise_gemm_plans( + variants, + inputs, + output, + costs, + config.planning_beam_width.max(1), + ) + }; + tracing::debug!( + ?orientation, + generated_grids, + proxy_frontier_grids, + generated_variants, + generated_grouped_variants, + retained_grouped_variants = retained + .iter() + .filter(|candidate| candidate + .requirements + .output + .format + .layout + .tiling + .axes + .iter() + .any(|axis| axis.padding_groups > 1)) + .count(), + retained_variants = retained.len(), + "retained parallel GEMM candidates" + ); + retained +} + +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)] +struct GemmPlanCompatibility { + orientation: Option, + reduction_staging: Option, + inputs: Vec<(StorageOrderCompatibility, MemoryClass, LocalOperandStaging)>, + output: ( + StorageOrderCompatibility, + MemoryClass, + Vec<(TensorAxis, u16, u32)>, + ), +} + +fn gemm_plan_compatibility(candidate: &OperatorSchedule) -> GemmPlanCompatibility { + let (orientation, reduction_staging) = match candidate.steps.first() { + Some(ScheduleStep::Gemm(plan)) => ( + Some(plan.geometry.orientation), + candidate.steps.get(1).and_then(|step| match step { + ScheduleStep::Reduce { staging, .. } => Some(*staging), + _ => None, + }), + ), + _ => (None, None), + }; + GemmPlanCompatibility { + orientation, + reduction_staging, + inputs: candidate + .requirements + .inputs + .iter() + .map(|input| { + ( + storage_order_compatibility(input.format.layout.order), + input.format.layout.memory_class, + input.local_staging, + ) + }) + .collect(), + output: ( + storage_order_compatibility(candidate.requirements.output.format.layout.order), + candidate.requirements.output.format.layout.memory_class, + candidate + .requirements + .output + .format + .layout + .tiling + .axes + .iter() + .map(|axis| (axis.axis, axis.padding_groups, axis.shard_padding_multiple)) + .collect(), + ), + } +} + +fn retain_precise_gemm_plans( + candidates: Vec, + inputs: &[TensorType], + output: &TensorShape, + costs: &impl CostModel, + width: usize, +) -> Vec { + let ranked = candidates + .into_iter() + .map(|candidate| { + let planned_inputs = inputs + .iter() + .zip(&candidate.requirements.inputs) + .map(|(input, requirement)| TensorType { + shape: input.shape.clone(), + format: requirement.format.clone(), + }) + .collect::>(); + let planned_output = TensorType { + shape: output.clone(), + format: candidate.requirements.output.format.clone(), + }; + let memory = operator_memory_estimate(&candidate, &planned_inputs, &planned_output); + let exchange = + costs.operator_exchange_footprint(&candidate, &planned_inputs, &planned_output); + let objective = RegionMetrics { + cost: CostEstimate { + cycles: costs.operator_cycles(&candidate, &planned_inputs, &planned_output), + exchange_footprint: exchange, + ..CostEstimate::default() + }, + memory: memory.peaks(exchange.estimated_row_bytes(costs.target())), + }; + let compatibility = gemm_plan_compatibility(&candidate); + (candidate, objective, compatibility) + }) + .collect::>(); + let (mut ranked, _) = pareto_frontier( + ranked, + |(_, _, compatibility)| compatibility.clone(), + |(_, objective, _)| *objective, + ); + ranked.sort_by_key(|(_, objective, _)| { + ( + objective.cost.cycles, + objective.memory.total, + objective.memory.interleaved, + objective.memory.exchange_rows, + ) + }); + let mut selected = BTreeSet::new(); + let mut represented = BTreeSet::new(); + for (index, (_, _, compatibility)) in ranked.iter().enumerate() { + if represented.insert(compatibility.clone()) { + selected.insert(index); + } + } + for index in 0..ranked.len() { + if selected.len() == width { + break; + } + selected.insert(index); + } + ranked + .into_iter() + .enumerate() + .filter_map(|(index, (candidate, _, _))| selected.contains(&index).then_some(candidate)) + .collect() +} diff --git a/crates/ipu-codegen/src/mid/ir.rs b/crates/ipu-codegen/src/mid/ir.rs new file mode 100644 index 0000000..522bd5c --- /dev/null +++ b/crates/ipu-codegen/src/mid/ir.rs @@ -0,0 +1,128 @@ +//! Layout-aware mid-level graph records. + +use crate::OperatorSchedule; +use crate::conversion::{ConversionMapping, ConversionStrategy, DeferredTransform}; +use crate::graph::{GraphInputKind, OperationId, ValueId}; +use crate::layout::TensorType; +use crate::metrics::{OperationMetrics, RegionMetrics}; +use crate::operator::{MidOperator, OperandMaterialization}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct MidValueId(u32); + +impl MidValueId { + pub const fn index(self) -> u32 { + self.0 + } + + pub const fn from_index(index: u32) -> Self { + Self(index) + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct MidValue { + pub id: MidValueId, + pub tensor_type: TensorType, + /// Semantic value represented by this value; conversions retain the same + /// origin. Region arguments also refer to their high-level argument ID. + pub origin: ValueId, + /// Values in the same group use the same logical-to-physical tile mapping. + /// Structured iteration uses this to keep successive parameter blocks + /// addressable by a single advancing base pointer. + pub storage_group: MidValueId, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum MidOperationKind { + Operator(OperatorSchedule), + Convert( + Option, + ConversionStrategy, + OperandMaterialization, + Vec, + ), + CastPrecision, + Repeat(MidRepeat), +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct MidOperation { + pub source: Option, + pub inputs: Vec, + pub results: Vec, + pub kind: MidOperationKind, + pub metrics: OperationMetrics, +} + +impl MidOperation { + pub fn conversion(&self) -> Option<(ConversionStrategy, OperandMaterialization)> { + match &self.kind { + MidOperationKind::Convert(_, strategy, materialization, _) => { + Some((*strategy, *materialization)) + } + MidOperationKind::CastPrecision => Some(( + ConversionStrategy::LocalKernel, + OperandMaterialization::Complete, + )), + MidOperationKind::Operator(_) | MidOperationKind::Repeat(_) => None, + } + } + + pub fn operator_plan(&self) -> Option<&OperatorSchedule> { + match &self.kind { + MidOperationKind::Operator(plan) => Some(plan), + MidOperationKind::Convert(..) + | MidOperationKind::CastPrecision + | MidOperationKind::Repeat(_) => None, + } + } + + pub fn operator_plan_mut(&mut self) -> Option<&mut OperatorSchedule> { + match &mut self.kind { + MidOperationKind::Operator(plan) => Some(plan), + MidOperationKind::Convert(..) + | MidOperationKind::CastPrecision + | MidOperationKind::Repeat(_) => None, + } + } + + pub fn operator(&self) -> Option { + self.operator_plan().map(|plan| plan.operator) + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct MidRegion { + pub arguments: Vec, + pub operations: Vec, + pub yields: Vec, + pub metrics: RegionMetrics, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct MidRepeat { + pub count: u32, + pub carried_inputs: usize, + pub invariant_inputs: usize, + /// One normalized value list for each iterated body argument. Keeping the + /// lists on the structured operation avoids unrolling layer parameters. + pub iterated_inputs: Vec>, + pub body: MidRegion, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct MidInput { + pub name: String, + pub kind: GraphInputKind, + pub value: MidValueId, +} + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct MidGraph { + pub inputs: Vec, + pub values: Vec, + pub operations: Vec, + pub outputs: Vec, + pub metrics: RegionMetrics, +} * Unmerged path crates/ipu-codegen/src/mid/mod.rs diff --git a/crates/ipu-codegen/src/operator.rs b/crates/ipu-codegen/src/operator.rs new file mode 100644 index 0000000..bd4c17f --- /dev/null +++ b/crates/ipu-codegen/src/operator.rs @@ -0,0 +1,857 @@ +//! Whole-device operator plans and tile-kernel specifications. + +use crate::graph::{AddOptions, AttentionOptions, GemmOptions, TensorShape}; +use crate::layout::{ + AMP_COLUMN_MICRO, AMP_INNER_BLOCK, Layout, MemoryClass, NativeKernelOrder, StorageOrder, + TensorAxis, TensorFormat, TensorType, +}; +use crate::schedule::{AttentionBlocking, AttentionMap}; + +/// In-memory representation of one tensor element. +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum Precision { + /// F143 values scaled by a tensor-wide power of two. + F8F143 { + scale_exponent: i8, + }, + F16, + F32, +} + +impl Precision { + pub const fn bytes(self) -> u64 { + match self { + Self::F8F143 { .. } => 1, + Self::F16 => 2, + Self::F32 => 4, + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum AccumulationPrecision { + F16, + F32, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum MidOperator { + Gemm { + options: GemmOptions, + multiply: Precision, + accumulate: AccumulationPrecision, + }, + Gelu, + Add(AddOptions), + FlashAttention { + options: AttentionOptions, + accumulate: AccumulationPrecision, + }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)] +pub enum GemmKernelMode { + Initialize, + Accumulate, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum GemmWeightLoad { + Standard, + Interleaved, +} + +/// Linearization of a GEMM's logical tile grid. +/// +/// The order is part of the operand and output layouts because it determines +/// which tensor coordinates occupy adjacent logical (and therefore paired +/// physical) tiles. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum GridOrder { + #[default] + ColumnsFast, + RowsFast, +} + +/// Physical matrix orientation used by a blocked GEMM implementation. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum GemmOrientation { + #[default] + Normal, + /// Compute `(rightᵀ × leftᵀ)ᵀ`. This preserves GEMM semantics while + /// exchanging the physical row and output-column traversal dimensions. + Swapped, +} + +impl GemmOrientation { + pub const fn physical_left_input(self) -> usize { + match self { + Self::Normal => 0, + Self::Swapped => 1, + } + } + + pub const fn physical_right_input(self) -> usize { + 1 - self.physical_left_input() + } + + pub const fn row_axis(self) -> TensorAxis { + match self { + Self::Normal => TensorAxis::FromEnd(2), + Self::Swapped => TensorAxis::FromEnd(1), + } + } + + pub const fn column_axis(self) -> TensorAxis { + match self { + Self::Normal => TensorAxis::FromEnd(1), + Self::Swapped => TensorAxis::FromEnd(2), + } + } + + pub(crate) fn physical_order(self, values: [T; 2]) -> [T; 2] { + match self { + Self::Normal => values, + Self::Swapped => { + let [left, right] = values; + [right, left] + } + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct GemmGrid { + pub rows: u16, + pub columns: u16, + pub inner: u16, +} + +impl GemmGrid { + pub const fn tile_count(self) -> u16 { + self.rows + .saturating_mul(self.columns) + .saturating_mul(self.inner) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct GemmResultGrid { + pub rows: u16, + pub columns: u16, +} + +impl GemmResultGrid { + pub const fn tile_count(self) -> u16 { + self.rows.saturating_mul(self.columns) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct GemmBlockShape { + pub inner: u32, + pub output_columns: u32, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct GemmKernelFamily { + pub multiply: Precision, + pub accumulate: AccumulationPrecision, + pub weights: GemmWeightLoad, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct GemmGeometry { + pub block: GemmBlockShape, + pub orientation: GemmOrientation, + /// Spatial row, output-column, and K partitions which invoke kernels. + pub compute: GemmGrid, + /// Spatial ownership of the final result. Parallel reductions may spread + /// roots over former K-partition tiles. + pub result: GemmResultGrid, + pub order: GridOrder, +} + +/// Exact blocked-GEMM geometry retained for planner diagnosis. Constraints +/// are keyed by the source graph operation and bypass beam pruning and +/// conservative whole-graph memory rejection. Concrete placement remains the +/// final authority on whether the resulting package fits. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct GemmPlanConstraint { + pub source_operation: u32, + pub geometry: GemmGeometry, + pub reduction_staging: Option, + pub weight_memory_class: MemoryClass, + pub local_weight_staging: LocalOperandStaging, +} + +/// Lifetime policy for partials reduced across a GEMM's K partitions. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord)] +pub enum ReductionStaging { + /// Receive every remote partial into one packed buffer, then reduce once. + #[default] + Complete, + /// Receive and accumulate one remote partial at a time. This minimizes + /// temporary SRAM at the expense of additional exchange epochs and kernel + /// launches. + Streamed, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum EmptyOutputShardPolicy { + Skip, + Reject, +} + +impl crate::OperatorSchedule { + fn empty_output_shard_policy(&self) -> EmptyOutputShardPolicy { + match self.steps.first() { + Some(crate::ScheduleStep::KernelMap(_)) => EmptyOutputShardPolicy::Skip, + Some(crate::ScheduleStep::Gemm(_) | crate::ScheduleStep::Attention(_)) => { + EmptyOutputShardPolicy::Reject + } + Some(crate::ScheduleStep::Reduce { .. }) | None => EmptyOutputShardPolicy::Reject, + } + } +} + +pub(crate) fn layout_has_empty_shards(layout: &Layout, shape: &TensorShape) -> bool { + layout + .resolve(shape) + .map_or(true, |resolved| resolved.has_empty_shards()) +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct OperandRequirement { + pub format: TensorFormat, + pub allocation: AllocationRequirements, + /// How a locally resident operand should be consumed when other tiles use + /// an operator-local staging buffer for the same operand. + pub local_staging: LocalOperandStaging, + /// Whether a schedule may populate and consume bounded operand slices + /// instead of materializing the complete required format first. + pub materialization: OperandMaterialization, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum MemoryElementRequirement { + #[default] + Any, + Distinct, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct AllocationRequirements { + pub alignment: u32, + /// Bytes the kernel may access beyond the logical tensor payload. + pub access_tail_bytes: u32, + pub memory_element: MemoryElementRequirement, +} + +impl AllocationRequirements { + pub fn merge(&mut self, other: Self) { + self.alignment = self.alignment.max(other.alignment); + self.access_tail_bytes = self.access_tail_bytes.max(other.access_tail_bytes); + if other.memory_element == MemoryElementRequirement::Distinct { + self.memory_element = MemoryElementRequirement::Distinct; + } + } + + pub fn require_distinct_element(&mut self) { + self.memory_element = MemoryElementRequirement::Distinct; + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)] +pub enum LocalOperandStaging { + Direct(MemoryClass), + Staged(MemoryClass), +} + +impl Default for LocalOperandStaging { + fn default() -> Self { + Self::Direct(MemoryClass::Standard) + } +} + +impl LocalOperandStaging { + pub(crate) const fn memory_class(self) -> MemoryClass { + match self { + Self::Direct(memory_class) | Self::Staged(memory_class) => memory_class, + } + } + + pub(crate) const fn stages_local(self) -> bool { + matches!(self, Self::Staged(_)) + } +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum OperandMaterialization { + #[default] + Complete, + DispatchSlices, +} + +impl OperandRequirement { + pub fn new(format: TensorFormat, alignment: u32) -> Self { + Self { + format, + allocation: AllocationRequirements { + alignment, + ..AllocationRequirements::default() + }, + local_staging: LocalOperandStaging::default(), + materialization: OperandMaterialization::Complete, + } + } + + pub fn with_access_tail(mut self, bytes: u32) -> Self { + self.allocation.access_tail_bytes = bytes; + self + } + + pub fn with_local_staging(mut self, staging: LocalOperandStaging) -> Self { + self.local_staging = staging; + self + } + + pub fn with_materialization(mut self, materialization: OperandMaterialization) -> Self { + self.materialization = materialization; + self + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum OutputAliasing { + Fresh, + MayAliasInputs(Vec), + MustAliasInput(u16), +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum MemoryOperand { + Output, + Input(u16), +} + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct MemorySpaceRequirements { + /// Each group names operand ranges which must occupy distinct effective + /// tile-memory elements. + pub distinct_element_groups: Vec>, +} + +impl MemorySpaceRequirements { + pub fn with_distinct_elements( + mut self, + operands: impl IntoIterator, + ) -> Self { + self.distinct_element_groups + .push(operands.into_iter().collect()); + self + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct OperatorRequirements { + pub inputs: Vec, + pub output: OperandRequirement, + pub output_aliasing: OutputAliasing, + pub memory_space: MemorySpaceRequirements, +} + +#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)] +pub enum OperatorPlanError { + #[error("operator plan operand arity does not match its requirements")] + OperandArity, + #[error("operator plan schedule does not match the selected operator")] + DispatchMismatch, + #[error("operator plan uses zero or incompatible block dimensions")] + InvalidBlocking, + #[error("operator plan requires corresponding activation and output tile groups")] + IncompatibleTileGroups, + #[error("operator schedule does not support empty output shards")] + EmptyOutputShard, + #[error("blocked GEMM currently requires non-transposed AMP left/right/output formats")] + UnsupportedGemmLayout, +} + +fn alias_compatible( + index: usize, + requirements: &[OperandRequirement], + inputs: &[TensorType], + output_requirement: &OperandRequirement, + output_shape: &TensorShape, +) -> bool { + requirements + .get(index) + .zip(inputs.get(index)) + .is_some_and(|(requirement, input)| { + input.shape == *output_shape && requirement.format == output_requirement.format + }) +} + +fn valid_requirement(requirement: &OperandRequirement, shape: &TensorShape) -> bool { + requirement.allocation.alignment.is_power_of_two() + && requirement.format.layout.resolve(shape).is_ok() +} + +fn valid_memory_operand(operand: MemoryOperand, input_count: usize) -> bool { + match operand { + MemoryOperand::Output => true, + MemoryOperand::Input(index) => usize::from(index) < input_count, + } +} + +impl crate::OperatorSchedule { + pub(crate) fn supports(&self, inputs: &[TensorType], output: &TensorShape) -> bool { + if self.requirements.inputs.len() != inputs.len() + || !valid_requirement(&self.requirements.output, output) + || !self + .requirements + .inputs + .iter() + .zip(inputs) + .all(|(requirement, input)| valid_requirement(requirement, &input.shape)) + { + return false; + } + let alias_valid = match &self.requirements.output_aliasing { + OutputAliasing::Fresh => true, + OutputAliasing::MayAliasInputs(indices) => { + !indices.is_empty() + && indices.iter().any(|index| { + alias_compatible( + usize::from(*index), + &self.requirements.inputs, + inputs, + &self.requirements.output, + output, + ) + }) + } + OutputAliasing::MustAliasInput(index) => alias_compatible( + usize::from(*index), + &self.requirements.inputs, + inputs, + &self.requirements.output, + output, + ), + }; + if !alias_valid + || !self + .requirements + .memory_space + .distinct_element_groups + .iter() + .all(|operands| { + operands.len() >= 2 + && operands + .iter() + .all(|operand| valid_memory_operand(*operand, inputs.len())) + && operands.iter().enumerate().all(|(index, operand)| { + !operands[..index].iter().any(|previous| previous == operand) + }) + }) + { + return false; + } + let planned_inputs = inputs + .iter() + .zip(&self.requirements.inputs) + .map(|(input, requirement)| TensorType { + shape: input.shape.clone(), + format: requirement.format.clone(), + }) + .collect::>(); + let planned_output = TensorType { + shape: output.clone(), + format: self.requirements.output.format.clone(), + }; + self.validate(&planned_inputs, &planned_output).is_ok() + } + + pub fn validate( + &self, + inputs: &[TensorType], + output: &TensorType, + ) -> Result<(), OperatorPlanError> { + if inputs.len() != self.requirements.inputs.len() { + return Err(OperatorPlanError::OperandArity); + } + if self.empty_output_shard_policy() == EmptyOutputShardPolicy::Reject + && layout_has_empty_shards(&output.format.layout, &output.shape) + { + return Err(OperatorPlanError::EmptyOutputShard); + } + match (&self.operator, self.steps.as_slice()) { + ( + MidOperator::Gemm { + options, multiply, .. + }, + [crate::ScheduleStep::Gemm(plan), ..], + ) => { + if plan.inputs + != [ + crate::ScheduleValue::Input(0), + crate::ScheduleValue::Input(1), + ] + { + return Err(OperatorPlanError::DispatchMismatch); + } + let reduction = self.steps.get(1).and_then(|step| match step { + crate::ScheduleStep::Reduce { + input, + output, + staging, + } => Some((input, output, staging)), + _ => None, + }); + if plan.geometry.compute.inner > 1 { + if plan.output != crate::ScheduleValue::Temporary(0) + || reduction.is_none_or(|(input, output, _)| { + *input != plan.output || *output != crate::ScheduleValue::Output + }) + { + return Err(OperatorPlanError::DispatchMismatch); + } + } else if plan.output != crate::ScheduleValue::Output || reduction.is_some() { + return Err(OperatorPlanError::DispatchMismatch); + } + let inner_block = &plan.geometry.block.inner; + let output_column_block = &plan.geometry.block.output_columns; + let compute = plan.geometry.compute; + let orientation = &plan.geometry.orientation; + let [left, right] = inputs else { + return Err(OperatorPlanError::OperandArity); + }; + if compute.inner == 1 + && left.format.layout.tiling.tile_count + != output.format.layout.tiling.tile_count + { + return Err(OperatorPlanError::IncompatibleTileGroups); + } + let formats_match_orientation = match orientation { + GemmOrientation::Normal => { + matches!( + left.format.layout.order, + StorageOrder::Native(NativeKernelOrder::Left) + ) && matches!( + right.format.layout.order, + StorageOrder::Blocked(order) if order.is_matrix() + ) && output.format.layout.order + == StorageOrder::Native(if *multiply == Precision::F16 { + NativeKernelOrder::Left + } else { + NativeKernelOrder::Output + }) + } + GemmOrientation::Swapped => { + matches!( + left.format.layout.order, + StorageOrder::Blocked(order) if order.is_transposed_matrix() + ) && right.format.layout.order + == StorageOrder::Native(NativeKernelOrder::TransposedLeft) + && output.format.layout.order + == StorageOrder::Native(if *multiply == Precision::F16 { + NativeKernelOrder::TransposedLeft + } else { + NativeKernelOrder::TransposedOutput + }) + } + }; + if options.transpose_left || options.transpose_right || !formats_match_orientation { + return Err(OperatorPlanError::UnsupportedGemmLayout); + } + let MidOperator::Gemm { + multiply, + accumulate, + .. + } = &self.operator + else { + return Err(OperatorPlanError::DispatchMismatch); + }; + if plan.kernel.multiply != *multiply || plan.kernel.accumulate != *accumulate { + return Err(OperatorPlanError::DispatchMismatch); + } + if *inner_block == 0 + || *output_column_block == 0 + || left.shape.0.len() < 2 + || output.shape.0.len() < 2 + { + return Err(OperatorPlanError::InvalidBlocking); + } + let row_axis = orientation.row_axis(); + let column_axis = orientation.column_axis(); + let axis_partitions = |axis| { + output + .format + .layout + .tiling + .axes + .iter() + .find(|tiling| tiling.axis == axis) + .map_or(1, |tiling| tiling.partitions) + }; + if output.format.layout.tiling.tile_count != plan.geometry.result.tile_count() + || axis_partitions(row_axis) != plan.geometry.result.rows + || axis_partitions(column_axis) != plan.geometry.result.columns + { + return Err(OperatorPlanError::InvalidBlocking); + } + if compute.inner > 1 { + let result_rows = plan.geometry.result.rows; + let result_columns = plan.geometry.result.columns; + let result_row_partitions = result_rows.checked_div(compute.rows).unwrap_or(0); + let result_column_partitions = + result_columns.checked_div(compute.columns).unwrap_or(0); + if compute.rows == 0 + || compute.columns == 0 + || result_rows == 0 + || result_columns == 0 + || !result_rows.is_multiple_of(compute.rows) + || !result_columns.is_multiple_of(compute.columns) + || result_row_partitions.saturating_mul(result_column_partitions) + > compute.inner + || reduction.is_none() + || axis_partitions(row_axis) != result_rows + || axis_partitions(column_axis) != result_columns + { + return Err(OperatorPlanError::InvalidBlocking); + } + } else if compute.rows != plan.geometry.result.rows + || compute.columns != plan.geometry.result.columns + || reduction.is_some() + { + return Err(OperatorPlanError::InvalidBlocking); + } + let [physical_left, physical_right] = orientation.physical_order([left, right]); + let left_layout = physical_left + .format + .layout + .resolve(&physical_left.shape) + .map_err(|_| OperatorPlanError::InvalidBlocking)?; + let output_layout = output + .format + .layout + .resolve(&output.shape) + .map_err(|_| OperatorPlanError::InvalidBlocking)?; + let left_padded = left_layout.padded_shape(); + let output_padded = output_layout.padded_shape(); + let output_column_axis = orientation + .column_axis() + .resolve(output_padded.0.len()) + .map_err(|_| OperatorPlanError::InvalidBlocking)?; + let columns_per_output_shard = output_layout + .maximum_axis_extent(output_column_axis) + .ok_or(OperatorPlanError::InvalidBlocking)?; + let right_layout = physical_right + .format + .layout + .resolve(&physical_right.shape) + .map_err(|_| OperatorPlanError::InvalidBlocking)?; + let right_padded = right_layout.padded_shape(); + let right_column_axis = orientation + .column_axis() + .resolve(right_padded.0.len()) + .map_err(|_| OperatorPlanError::InvalidBlocking)?; + let columns_per_right_shard = right_layout + .maximum_axis_extent(right_column_axis) + .ok_or(OperatorPlanError::InvalidBlocking)?; + let grid_plan = left.format.layout.tiling.replicas > 1 + || right.format.layout.tiling.replicas > 1 + || right + .format + .layout + .tiling + .axes + .iter() + .any(|axis| axis.axis == TensorAxis::FromEnd(2) && axis.partitions > 1); + if grid_plan + && [left, right, output] + .into_iter() + .any(|tensor| !layout_shards_are_nonempty(tensor)) + { + return Err(OperatorPlanError::InvalidBlocking); + } + let physical_left_inner_axis = orientation + .column_axis() + .resolve(left_padded.0.len()) + .map_err(|_| OperatorPlanError::InvalidBlocking)?; + let balanced_output_columns = compute.inner > 1; + let output_shard_alignment = if balanced_output_columns { + AMP_COLUMN_MICRO + } else { + *output_column_block + }; + if !left_padded.0[physical_left_inner_axis].is_multiple_of(*inner_block) + || !output_padded.0[output_column_axis].is_multiple_of(output_shard_alignment) + || !columns_per_output_shard.is_multiple_of(output_shard_alignment) + || (balanced_output_columns && columns_per_output_shard > *output_column_block) + || columns_per_right_shard < *output_column_block + { + return Err(OperatorPlanError::InvalidBlocking); + } + Ok(()) + } + (MidOperator::Gelu | MidOperator::Add(_), [crate::ScheduleStep::KernelMap(map)]) => { + if map.output != crate::ScheduleValue::Output + || map.inputs.len() != inputs.len() + || !map.inputs.iter().enumerate().all(|(index, (value, _))| { + *value == crate::ScheduleValue::Input(index as u16) + }) + { + return Err(OperatorPlanError::DispatchMismatch); + } + let output_tiles = output.format.layout.tiling.tile_count; + if inputs + .iter() + .any(|input| input.format.layout.tiling.tile_count != output_tiles) + { + Err(OperatorPlanError::IncompatibleTileGroups) + } else { + Ok(()) + } + } + ( + MidOperator::FlashAttention { + options, + accumulate, + }, + [ + crate::ScheduleStep::Attention(AttentionMap { + inputs: scheduled_inputs, + output: scheduled_output, + kernel, + blocking: + AttentionBlocking::Flash { + query_rows, + key_rows, + }, + query_dimension, + value_dimension, + }), + ], + ) => { + if *scheduled_inputs + != [ + crate::ScheduleValue::Input(0), + crate::ScheduleValue::Input(1), + crate::ScheduleValue::Input(2), + ] + || *scheduled_output != crate::ScheduleValue::Output + { + return Err(OperatorPlanError::DispatchMismatch); + } + let [query, key, value] = inputs else { + return Err(OperatorPlanError::OperandArity); + }; + if options.causal + || *accumulate != AccumulationPrecision::F32 + || *query_rows == 0 + || *key_rows != AMP_INNER_BLOCK + || *query_dimension == 0 + || *value_dimension == 0 + || !matches!( + query.format.layout.order, + StorageOrder::Native(NativeKernelOrder::Left) + ) + || !matches!( + key.format.layout.order, + StorageOrder::Native(NativeKernelOrder::TransposedRight) + ) + || !matches!( + value.format.layout.order, + StorageOrder::Blocked(order) if order.is_matrix() + ) + || output.format.layout.order != StorageOrder::Linear + || query.format.layout.tiling.tile_count + != output.format.layout.tiling.tile_count + || key.format.layout.tiling.tile_count != value.format.layout.tiling.tile_count + || kernel.multiply != Precision::F16 + || kernel.accumulate != *accumulate + || kernel.weights != GemmWeightLoad::Standard + { + Err(OperatorPlanError::InvalidBlocking) + } else { + Ok(()) + } + } + ( + MidOperator::FlashAttention { + options, + accumulate, + }, + [ + crate::ScheduleStep::Attention(AttentionMap { + inputs: scheduled_inputs, + output: scheduled_output, + kernel, + blocking: + AttentionBlocking::Materialized { + query_rows, + padded_key_rows, + }, + query_dimension, + value_dimension, + }), + ], + ) => { + if *scheduled_inputs + != [ + crate::ScheduleValue::Input(0), + crate::ScheduleValue::Input(1), + crate::ScheduleValue::Input(2), + ] + || *scheduled_output != crate::ScheduleValue::Output + { + return Err(OperatorPlanError::DispatchMismatch); + } + let [query, key, value] = inputs else { + return Err(OperatorPlanError::OperandArity); + }; + if options.causal + || *accumulate != AccumulationPrecision::F32 + || *query_rows == 0 + || *padded_key_rows == 0 + || !padded_key_rows.is_multiple_of(AMP_INNER_BLOCK) + || *query_dimension == 0 + || *value_dimension == 0 + || !matches!( + query.format.layout.order, + StorageOrder::Native(NativeKernelOrder::Left) + ) + || !matches!( + key.format.layout.order, + StorageOrder::Native(NativeKernelOrder::TransposedRight) + ) + || !matches!( + value.format.layout.order, + StorageOrder::Blocked(order) if order.is_matrix() + ) + || output.format.layout.order != StorageOrder::Linear + || query.format.layout.tiling.tile_count + != output.format.layout.tiling.tile_count + || key.format.layout.tiling.tile_count != value.format.layout.tiling.tile_count + || kernel.multiply != Precision::F16 + || kernel.accumulate != *accumulate + || kernel.weights != GemmWeightLoad::Standard + { + Err(OperatorPlanError::InvalidBlocking) + } else { + Ok(()) + } + } + _ => Err(OperatorPlanError::DispatchMismatch), + } + } +} + +fn layout_shards_are_nonempty(tensor: &TensorType) -> bool { + tensor + .format + .layout + .resolve(&tensor.shape) + .is_ok_and(|resolved| !resolved.has_empty_shards()) +} * Unmerged path crates/ipu-codegen/src/package.rs * Unmerged path crates/ipu-codegen/src/place.rs diff --git a/crates/ipu-codegen/src/schedule.rs b/crates/ipu-codegen/src/schedule.rs new file mode 100644 index 0000000..76fd93b --- /dev/null +++ b/crates/ipu-codegen/src/schedule.rs @@ -0,0 +1,182 @@ +//! Parametric whole-device work selected by mid-level planning. + +use crate::{ + GemmGeometry, GemmKernelFamily, GemmOrientation, GridOrder, Layout, MidOperator, + NativeKernelOrder, OperatorRequirements, ReductionStaging, StorageOrder, TensorFormat, + TensorType, TileKernelSpec, +}; + +/// A value consumed or produced by an operator schedule. +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum ScheduleValue { + Input(u16), + Temporary(u16), + Output, +} + +/// How a mapped kernel obtains one input for each output shard. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ScheduleAccess { + /// Select every logical intersection, including singleton broadcasting. + LogicalOverlap, + /// Consume the corresponding shard already resident on the kernel tile. + TileLocal, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct KernelMap { + pub kernel: TileKernelSpec, + pub inputs: Vec<(ScheduleValue, ScheduleAccess)>, + pub output: ScheduleValue, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct GemmMap { + pub inputs: [ScheduleValue; 2], + pub output: ScheduleValue, + pub kernel: GemmKernelFamily, + pub geometry: GemmGeometry, +} + +impl GemmMap { + /// Tensor produced by the compute grid before a parallel K reduction. + pub(crate) fn partial_tensor(&self, output: &TensorType) -> Option { + if self.geometry.compute.inner < 2 { + return Some(output.clone()); + } + let output_rank = output.shape.0.len(); + let output_column_axis = self + .geometry + .orientation + .column_axis() + .resolve(output_rank) + .ok()?; + let column_tiling = *output + .format + .layout + .tiling + .axes + .iter() + .find(|axis| axis.axis.resolve(output_rank).ok() == Some(output_column_axis))?; + let rows = self.geometry.compute.rows; + let columns = self.geometry.compute.columns; + let tiles = rows.checked_mul(columns)?; + let left_order = match self.geometry.orientation { + GemmOrientation::Normal => NativeKernelOrder::Left, + GemmOrientation::Swapped => NativeKernelOrder::TransposedLeft, + }; + let constructor = if output.format.layout.order == StorageOrder::Native(left_order) { + Layout::amp_left_result_grid + } else { + Layout::amp_output_grid + }; + let mut layout = constructor( + self.geometry.orientation, + self.geometry.block.output_columns, + tiles, + rows, + columns, + GridOrder::ColumnsFast, + ); + let axis = layout + .tiling + .axes + .iter_mut() + .find(|axis| axis.axis.resolve(output_rank).ok() == Some(output_column_axis))?; + axis.block_size = column_tiling.block_size; + axis.padding_multiple = column_tiling.block_size; + if column_tiling.partitions == columns { + axis.block_size = column_tiling.block_size; + axis.padding_multiple = column_tiling.padding_multiple; + axis.shard_padding_multiple = column_tiling.shard_padding_multiple; + } + let partial = TensorType { + shape: output.shape.clone(), + format: TensorFormat { + precision: output.format.precision, + layout, + }, + }; + partial.format.layout.resolve(&partial.shape).ok()?; + Some(partial) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum AttentionBlocking { + Flash { + query_rows: u32, + key_rows: u32, + }, + Materialized { + query_rows: u32, + padded_key_rows: u32, + }, +} + +impl AttentionBlocking { + pub const fn query_rows(self) -> u32 { + match self { + Self::Flash { query_rows, .. } | Self::Materialized { query_rows, .. } => query_rows, + } + } + + pub const fn key_block_rows(self) -> u32 { + match self { + Self::Flash { key_rows, .. } => key_rows, + Self::Materialized { .. } => crate::AMP_INNER_BLOCK, + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct AttentionMap { + pub inputs: [ScheduleValue; 3], + pub output: ScheduleValue, + pub kernel: GemmKernelFamily, + pub blocking: AttentionBlocking, + pub query_dimension: u32, + pub value_dimension: u32, +} + +impl AttentionMap { + pub fn gemm_blocks(&self) -> [crate::GemmBlockShape; 2] { + let key_columns = match self.blocking { + AttentionBlocking::Materialized { + padded_key_rows, .. + } => padded_key_rows, + blocking => blocking.key_block_rows(), + }; + [ + crate::GemmBlockShape { + inner: self.query_dimension, + output_columns: key_columns, + }, + crate::GemmBlockShape { + inner: key_columns, + output_columns: self.value_dimension, + }, + ] + } +} + +/// Ordered, symbolic whole-device work. Placement binds values to concrete +/// shards and tiles; low lowering only materializes the selected steps. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum ScheduleStep { + KernelMap(KernelMap), + Gemm(GemmMap), + Reduce { + input: ScheduleValue, + output: ScheduleValue, + staging: ReductionStaging, + }, + Attention(AttentionMap), +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct OperatorSchedule { + pub operator: MidOperator, + pub steps: Vec, + pub requirements: OperatorRequirements, +} * Unmerged path crates/ipu-codegen/src/storage.rs * Unmerged path crates/ipu-codegen/src/tile.rs diff --git a/crates/ipu-driver/Cargo.toml b/crates/ipu-driver/Cargo.toml index 8907e13..ca18789 100644 --- a/crates/ipu-driver/Cargo.toml +++ b/crates/ipu-driver/Cargo.toml @@ -7,6 +7,7 @@ license.workspace = true [dependencies] fastrand.workspace = true ipu-package = { path = "../ipu-package" } +ipu-target = { path = "../ipu-target" } libc.workspace = true object.workspace = true thiserror.workspace = true diff --git a/crates/ipu-driver/src/lib.rs b/crates/ipu-driver/src/lib.rs index ab87742..e88b6b3 100644 --- a/crates/ipu-driver/src/lib.rs +++ b/crates/ipu-driver/src/lib.rs @@ -1,4 +1,5 @@ -use ipu_package::{Application, HostCall, HostExchange, TILE_MEMORY_BASE}; +use ipu_package::{Application, HostCall, HostExchange}; +use ipu_target::memory::{TILE_MEMORY_BASE, TILE_MEMORY_SIZE}; use object::{Object, ObjectSegment}; use std::collections::HashMap; use std::ffi::CString; @@ -12,7 +13,6 @@ use std::time::{Duration, Instant}; use tracing::{debug, info, trace}; pub const CONFIG_BAR_SIZE: usize = 0x80000; -pub const TILE_MEMORY_SIZE: usize = 624 * 1024; // The secondary loader installs framed application payload at the SDK image's // launch slot. Applications reserve that word and enter at the following word. pub const APPLICATION_LOAD_BASE: u32 = TILE_MEMORY_BASE + 0x10; @@ -34,7 +34,7 @@ pub const SECONDARY_LOADER_MIN_PAYLOAD_SIZE: usize = 0x4134; pub const SECONDARY_LOADER_MAX_FRAMES: usize = 0x283; /// Exclusive upper address that can be represented by that bootloader when /// loading an application from [`APPLICATION_LOAD_BASE`]. -pub const APPLICATION_LOAD_LIMIT: u32 = ipu_package::IPU21_APPLICATION_MEMORY_LIMIT; +pub const APPLICATION_LOAD_LIMIT: u32 = ipu_target::memory::IPU21_APPLICATION_MEMORY_LIMIT; const _: () = assert!( APPLICATION_LOAD_LIMIT == APPLICATION_LOAD_BASE + (SECONDARY_LOADER_MAX_FRAMES * FRAME_PAYLOAD_SIZE) as u32 diff --git a/crates/ipu-exchange/src/lib.rs b/crates/ipu-exchange/src/lib.rs deleted file mode 100644 index 251c5ac..0000000 --- a/crates/ipu-exchange/src/lib.rs +++ /dev/null @@ -1,4421 +0,0 @@ -use serde::{Deserialize, Serialize}; -use std::collections::HashSet; -use tracing::debug; - -pub mod diagnostic; - -pub const PLAN_WORDS: usize = 9; -pub const MAX_TRANSFER_WORDS: u32 = 4148; -/// Largest scheduled delay encodable by one exchange delay instruction. -pub const MAX_PLAN_OFFSET_CYCLES: u32 = 0x8_0000; -pub const EXCHANGE_WINDOW_BASE: u32 = 0x50000; -pub const EXCHANGE_WINDOW_BYTES: u32 = 0x8000; -pub const HOST_SHORT_MAX_BYTES: u32 = 60; -pub const HOST_LONG_MAX_BYTES: u32 = 1024; -pub const TILE_TO_HOST_MAX_BYTES: u32 = 256; -pub const HOST_PAGE_BYTES: u32 = 4096; -pub const HOST_TO_TILE_WINDOW_BYTES: u32 = 0x4000; -pub const TILE_MUX_HOST: u32 = 0x600; -pub const TILE_MUX_EXCHANGE: u32 = 0x640; -const XREQ_BITMAP0_BITS: u32 = 24; - -const OPCODE_MASK: u32 = 0xfc00_0000; -const LONG_OPCODE_MASK: u32 = 0xf800_0000; -const DELAY_OPCODE_MASK: u32 = 0xfff8_0000; -const DELAY_OPCODE: u32 = 0x40a0_0000; -const DELAY_PIC_OPCODE: u32 = 0x6000_0000; -const DELAY_XPIC_OPCODE: u32 = 0x6400_0000; -const PIC_RECEIVE_ADDRESS_MASK: u32 = 0x3ffff; -const SEND_OPCODE: u32 = 0x7800_0000; -const SEND_ADDRESS_MASK: u32 = 0x001f_fff8; -const SEND_OFF_OPCODE: u32 = 0x7000_0000; -// See docs/EXCHANGE_INSTRUCTION_REFERENCE.md. SENDPICP is an aligned two-word -// supervisor instruction whose following word is inline PIC/XPIC payload. -const SEND_PIC_OPCODE: u32 = 0x7010_0000; -const SEND_PICP_OPCODE: u32 = 0xf000_0000; -const SEND_CONTROL_OPCODE_MASK: u32 = 0xf810_0000; -const SEND_PICP_OPCODE_MASK: u32 = 0xf000_0000; -const SEND_COUNT_MASK: u32 = 0x07e0_0000; -const SYNC_OPCODE: u32 = 0x4180_0000; -const SANS_OPCODE: u32 = 0x40c0_0000; -const BR_M_OPCODE: u32 = 0x4300_0000; -const CALL_M_IMMEDIATE_OPCODE: u32 = 0x1800_0000; -const SETZI_M_OPCODE: u32 = 0x1900_0000; -const PUT_SPECIAL_M_OPCODE: u32 = 0x4300_8000; -const LD32_M_IMMEDIATE_OPCODE: u32 = 0x0100_0000; -const ST32_M_IMMEDIATE_OPCODE: u32 = 0x4f00_0000; -const ADD_M_IMMEDIATE_OPCODE: u32 = 0x2200_0000; -const AND_M_IMMEDIATE_OPCODE: u32 = 0x4200_0000; -const SHL_M_IMMEDIATE_OPCODE: u32 = 0x4200_a000; -const BRZ_M_IMMEDIATE_OPCODE: u32 = 0x1300_0000; -const INCOMING_MUX_REGISTER: u8 = 0xa0; -const INCOMING_DCOUNT_REGISTER: u8 = 0xa6; -// The host hierarchy reserves eighteen exchange events for each tile-to-host -// payload. Short payloads must be padded before the next packet header (or the -// closing zero-byte read); longer payloads provide the interval themselves. -const TILE_TO_HOST_MIN_PAYLOAD_EVENTS: u32 = 18; -const HOST_TO_TILE_STREAM_END_BITS: u32 = 0x0c00_0000; -// Time reserved by the SDK supervisor schedule between receiving a host -// command and injecting that command into the device-side dispatch path. -const HOST_COMMAND_ROUTE_CYCLES: u32 = 73; - -pub const SANS_INACTIVE_INSTRUCTION: u32 = sans(0); -pub const SYNC_RECEIVE_INSTRUCTION: u32 = sync(0); -pub const SYNC_ANS_INSTRUCTION: u32 = sync(1); -pub const SYNC_SUPERVISOR_INSTRUCTION: u32 = sync(3); -pub const SYNC_ALL_INSTRUCTION: u32 = sync(7); -pub const SYNC_HOST_INSTRUCTION: u32 = sync(15); -pub const RETURN_M10_INSTRUCTION: u32 = br_m(10); - -pub const fn sans(selector: u8) -> u32 { - SANS_OPCODE | selector as u32 -} - -pub const fn sync(selector: u8) -> u32 { - SYNC_OPCODE | selector as u32 -} - -pub const fn br_m(register: u8) -> u32 { - BR_M_OPCODE | ((register as u32) << 20) -} - -pub fn encode_br_m(register: u8) -> Result { - if register >= 16 { - return Err(ExchangeError::Schedule("branch register")); - } - Ok(br_m(register)) -} - -pub fn encode_call_m_immediate( - return_register: u8, - target_address: u32, -) -> Result { - if return_register >= 16 || target_address & 3 != 0 || target_address >= 1 << 21 { - return Err(ExchangeError::Schedule("call operand")); - } - Ok(CALL_M_IMMEDIATE_OPCODE | (u32::from(return_register) << 20) | (target_address >> 2)) -} - -pub fn encode_setzi_m(register: u8, immediate: u32) -> Result { - if register >= 16 || immediate >= 1 << 20 { - return Err(ExchangeError::Schedule("setzi operand")); - } - Ok(setzi_m(register, immediate)) -} - -pub fn encode_put_special_m(special: u8, register: u8) -> Result { - if register >= 16 { - return Err(ExchangeError::Schedule("put source register")); - } - Ok(PUT_SPECIAL_M_OPCODE | (u32::from(register) << 20) | u32::from(special)) -} - -pub fn encode_ld32_m_immediate( - destination: u8, - base: u8, - delta: u8, - word_offset: u16, -) -> Result { - if destination >= 16 || base >= 16 || delta >= 16 || word_offset >= 1 << 12 { - return Err(ExchangeError::Schedule("ld32 operand")); - } - Ok(LD32_M_IMMEDIATE_OPCODE - | (u32::from(base) << 20) - | (u32::from(destination) << 16) - | (u32::from(delta) << 12) - | u32::from(word_offset)) -} - -pub fn encode_st32_m_immediate( - source: u8, - base: u8, - delta: u8, - word_offset: u16, -) -> Result { - if source >= 16 || base >= 16 || delta >= 16 || word_offset >= 1 << 12 { - return Err(ExchangeError::Schedule("st32 operand")); - } - Ok(ST32_M_IMMEDIATE_OPCODE - | (u32::from(base) << 20) - | (u32::from(source) << 16) - | (u32::from(delta) << 12) - | u32::from(word_offset)) -} - -pub fn encode_add_m_immediate( - destination: u8, - source: u8, - immediate: i32, -) -> Result { - let immediate = - i16::try_from(immediate).map_err(|_| ExchangeError::Schedule("add immediate operand"))?; - if destination >= 16 || source >= 16 { - return Err(ExchangeError::Schedule("add register operand")); - } - Ok(ADD_M_IMMEDIATE_OPCODE - | (u32::from(source) << 20) - | (u32::from(destination) << 16) - | u32::from(immediate as u16)) -} - -pub fn encode_and_m_immediate( - destination: u8, - source: u8, - immediate: u16, -) -> Result { - if destination >= 16 || source >= 16 || immediate >= 1 << 12 { - return Err(ExchangeError::Schedule("and operand")); - } - Ok(AND_M_IMMEDIATE_OPCODE - | (u32::from(source) << 20) - | (u32::from(destination) << 16) - | u32::from(immediate)) -} - -pub fn encode_shl_m_immediate( - destination: u8, - source: u8, - immediate: u16, -) -> Result { - if destination >= 16 || source >= 16 || immediate >= 1 << 12 { - return Err(ExchangeError::Schedule("shift-left operand")); - } - Ok(SHL_M_IMMEDIATE_OPCODE - | (u32::from(source) << 20) - | (u32::from(destination) << 16) - | u32::from(immediate)) -} - -pub fn encode_brz_m_immediate(register: u8, target_address: u32) -> Result { - if register >= 16 || target_address & 3 != 0 || target_address >= 1 << 21 { - return Err(ExchangeError::Schedule("brz operand")); - } - Ok(BRZ_M_IMMEDIATE_OPCODE | (u32::from(register) << 20) | (target_address >> 2)) -} - -/// Encodes a processor delay of `cycles` cycles. -pub fn encode_delay_m(cycles: u32) -> Result { - if !(1..=MAX_PLAN_OFFSET_CYCLES).contains(&cycles) { - return Err(ExchangeError::Schedule("processor delay range")); - } - Ok(delay(cycles - 1)) -} - -const fn setzi_m(register: u8, immediate: u32) -> u32 { - SETZI_M_OPCODE | ((register as u32) << 20) | immediate -} - -const fn put_special_from_m8(register: u8) -> u32 { - PUT_SPECIAL_M_OPCODE | (8 << 20) | register as u32 -} - -pub type PlanRow = [u32; PLAN_WORDS]; - -/// Returns the plan event horizon measured from the entry synchronization. -/// -/// Delay immediates advance to the event `N + 1` cycles later. Send -/// instructions occupy one event per transferred word. -pub fn plan_event_cycles(row: &[u32]) -> Result { - let mut cycles = 0u32; - let mut cursor = 0; - while cursor < row.len() { - let instruction = row[cursor]; - let advance = instruction_advance(instruction); - cycles = cycles - .checked_add(advance) - .ok_or(ExchangeError::Schedule("plan event horizon overflow"))?; - cursor += if is_send_control_pair(instruction) { - 2 - } else { - 1 - }; - } - Ok(cycles) -} - -#[derive(Clone, Debug, Default)] -struct TileProgramSchedule { - senders: Vec, - /// Borrowed transmit lane; no instruction or SRAM access on this tile. - reserved_sender_end: u32, - receive_events: Vec, - event_cycles: u32, - receive_stream: Option, -} - -#[derive(Clone, Debug)] -struct ReceiveStream { - mode: ReceiveMode, - paired_format: Option, - format_end_cycles: Option, - source_end_cycles: u32, - pointer_end_cycles: u32, - next_address: Option, -} - -#[derive(Clone, Debug)] -struct ScheduledSenderRow { - row: PlanRow, - start_cycles: u32, - end_cycles: u32, -} - -#[derive(Clone, Debug)] -pub struct PhaseProgramBuilder { - tile_states: Vec, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct PhaseTransferTiming { - pub payload_start: u32, - pub payload_end: u32, - pub sender_horizon: u32, - pub receiver_payload_starts: Vec, - pub receiver_payload_ends: Vec, - pub receiver_horizons: Vec, - pub horizon: u32, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct PhasePrograms { - pub programs: Vec>>, - pub tile_event_cycles: Vec, - pub event_cycles: u32, -} - -impl PhaseProgramBuilder { - pub fn new(tile_count: u16) -> Self { - Self { - tile_states: vec![TileProgramSchedule::default(); usize::from(tile_count)], - } - } - - pub fn tile_count(&self) -> u16 { - u16::try_from(self.tile_states.len()).expect("phase tile count was supplied as u16") - } - - pub fn tile_event_cycles(&self, tile: u16) -> Result { - self.tile_states - .get(usize::from(tile)) - .map(|schedule| schedule.event_cycles) - .ok_or(ExchangeError::Tile(tile)) - } - - pub fn event_cycles(&self) -> u32 { - self.tile_states - .iter() - .map(|schedule| schedule.event_cycles) - .max() - .unwrap_or(0) - } - - pub fn active_tile_count(&self) -> usize { - self.tile_states - .iter() - .filter(|schedule| schedule.event_cycles != 0) - .count() - } - - /// Finds one absolute offset at which all endpoints can perform the - /// transfer. Compatibility is evaluated across the complete phase state; - /// tile programs are not encoded until the phase is finished. - pub fn earliest_transfer_offset( - &self, - source: u16, - reserved_tiles: &[u16], - receivers: &[u16], - plan: &MulticastPlan, - words: u32, - requested: u32, - ) -> Result { - self.earliest_transfer_offset_impl( - source, - reserved_tiles, - receivers, - plan, - words, - requested, - true, - ) - } - - /// Finds an endpoint-compatible offset while deferring whole-row encoding - /// validation until the complete phase is available. Callers must encode - /// the completed phase and retry with [`Self::earliest_transfer_offset`] - /// if instruction alignment is not representable. - pub fn earliest_transfer_offset_deferred( - &self, - source: u16, - reserved_tiles: &[u16], - receivers: &[u16], - plan: &MulticastPlan, - words: u32, - requested: u32, - ) -> Result { - self.earliest_transfer_offset_impl( - source, - reserved_tiles, - receivers, - plan, - words, - requested, - false, - ) - } - - #[allow(clippy::too_many_arguments)] - fn earliest_transfer_offset_impl( - &self, - source: u16, - reserved_tiles: &[u16], - receivers: &[u16], - plan: &MulticastPlan, - words: u32, - requested: u32, - validate_encoding: bool, - ) -> Result { - if receivers.len() != plan.receivers.len() { - return Err(ExchangeError::Schedule("receiver row count")); - } - let source_schedule = self - .tile_states - .get(usize::from(source)) - .ok_or(ExchangeError::Tile(source))?; - let mut offset = requested; - loop { - let previous = offset; - offset = source_schedule.earliest_sender_offset(&plan.sender, offset)?; - for &tile in reserved_tiles { - let schedule = self - .tile_states - .get(usize::from(tile)) - .ok_or(ExchangeError::Tile(tile))?; - offset = offset.max( - schedule - .senders - .iter() - .map(|row| row.end_cycles) - .fold(schedule.reserved_sender_end, u32::max), - ); - } - for (&receiver, row) in receivers.iter().zip(&plan.receivers) { - let schedule = self - .tile_states - .get(usize::from(receiver)) - .ok_or(ExchangeError::Tile(receiver))?; - offset = schedule.earliest_receiver_offset(row, words, offset)?; - } - if offset == previous { - if !validate_encoding { - return Ok(offset); - } - match self.transfer_is_encodable_at(source, receivers, plan, offset, words) { - Ok(()) => return Ok(offset), - Err(ExchangeError::Schedule("SENDPICP instruction alignment")) => { - offset = offset - .checked_add(1) - .ok_or(ExchangeError::Schedule("receive offset overflow"))?; - } - Err(error) => return Err(error), - } - } - } - } - - fn transfer_is_encodable_at( - &self, - source: u16, - receivers: &[u16], - plan: &MulticastPlan, - schedule_offset: u32, - words: u32, - ) -> Result<(), ExchangeError> { - let mut source_schedule = self - .tile_states - .get(usize::from(source)) - .ok_or(ExchangeError::Tile(source))? - .clone(); - source_schedule.append_sender_at(&plan.sender, schedule_offset)?; - source_schedule.finish()?; - for (&receiver, row) in receivers.iter().zip(&plan.receivers) { - let mut receiver_schedule = self - .tile_states - .get(usize::from(receiver)) - .ok_or(ExchangeError::Tile(receiver))? - .clone(); - receiver_schedule.append_receiver_at(row, schedule_offset, words)?; - receiver_schedule.finish()?; - } - Ok(()) - } - - /// Adds one transfer to the phase schedule. The transfer's sender and all - /// receivers are recorded declaratively and compiled together by - /// [`Self::finish`]. - pub fn append_transfer_at( - &mut self, - source: u16, - reserved_tiles: &[u16], - receivers: &[u16], - plan: &MulticastPlan, - schedule_offset: u32, - words: u32, - ) -> Result { - let transfer_timing = - self.transfer_timing_at(source, receivers, plan, schedule_offset, words)?; - if receivers.len() != plan.receivers.len() { - return Err(ExchangeError::Schedule("receiver row count")); - } - let mut updates = Vec::with_capacity(receivers.len() + reserved_tiles.len() + 1); - let mut source_schedule = self - .tile_states - .get(usize::from(source)) - .ok_or(ExchangeError::Tile(source))? - .clone(); - source_schedule.append_sender_at(&plan.sender, schedule_offset)?; - updates.push((source, source_schedule)); - - for &tile in reserved_tiles { - if tile == source - || receivers.contains(&tile) - || updates.iter().any(|(updated, _)| *updated == tile) - { - return Err(ExchangeError::DuplicateTile); - } - let mut schedule = self - .tile_states - .get(usize::from(tile)) - .ok_or(ExchangeError::Tile(tile))? - .clone(); - schedule.reserved_sender_end = schedule - .reserved_sender_end - .max(transfer_timing.sender_horizon); - schedule.event_cycles = schedule.event_cycles.max(transfer_timing.sender_horizon); - updates.push((tile, schedule)); - } - - for (&receiver, row) in receivers.iter().zip(&plan.receivers) { - if receiver == source || updates.iter().any(|(tile, _)| *tile == receiver) { - return Err(ExchangeError::DuplicateTile); - } - let mut receiver_schedule = self - .tile_states - .get(usize::from(receiver)) - .ok_or(ExchangeError::Tile(receiver))? - .clone(); - receiver_schedule.append_receiver_at(row, schedule_offset, words)?; - updates.push((receiver, receiver_schedule)); - } - for (tile, schedule) in updates { - self.tile_states[usize::from(tile)] = schedule; - } - Ok(transfer_timing) - } - - pub fn transfer_timing_at( - &self, - source: u16, - receivers: &[u16], - plan: &MulticastPlan, - schedule_offset: u32, - words: u32, - ) -> Result { - if receivers.len() != plan.receivers.len() { - return Err(ExchangeError::Schedule("receiver row count")); - } - self.tile_states - .get(usize::from(source)) - .ok_or(ExchangeError::Tile(source))?; - let sender = scheduled_sender_timing(&plan.sender, schedule_offset)?; - let receiver_timings = receivers - .iter() - .zip(&plan.receivers) - .map(|(&receiver, row)| { - if receiver == source { - return Err(ExchangeError::DuplicateTile); - } - let schedule = self - .tile_states - .get(usize::from(receiver)) - .ok_or(ExchangeError::Tile(receiver))?; - let base = receive_row_timing(row, 0)?; - scheduled_receive_window( - &base, - schedule_offset, - words, - schedule.receive_stream.as_ref(), - ) - }) - .collect::, _>>()?; - let horizon = receiver_timings - .iter() - .map(|timing| timing.horizon) - .chain(std::iter::once(sender.horizon)) - .max() - .unwrap_or(schedule_offset); - Ok(PhaseTransferTiming { - payload_start: sender.payload_start, - payload_end: sender.payload_end, - sender_horizon: sender.horizon, - receiver_payload_starts: receiver_timings - .iter() - .map(|timing| timing.payload_start) - .collect(), - receiver_payload_ends: receiver_timings - .iter() - .map(|timing| timing.payload_end) - .collect(), - receiver_horizons: receiver_timings - .iter() - .map(|timing| timing.horizon) - .collect(), - horizon, - }) - } - - pub fn finish(self) -> Result { - // `append_transfer_at` has already merged the authoritative - // phase-wide schedule into these per-tile sender and receive-event - // timelines. Encoding them directly avoids replaying every transfer - // and performing the same conflict checks a second time. - let tile_states = self.tile_states; - let event_cycles = tile_states - .iter() - .map(|schedule| schedule.event_cycles) - .max() - .unwrap_or(0); - let tile_event_cycles = tile_states - .iter() - .map(|schedule| schedule.event_cycles) - .collect::>(); - let programs = tile_states - .iter() - .enumerate() - .map(|(tile, schedule)| { - if schedule.event_cycles == 0 { - return Ok(None); - } - let program = schedule.finish()?; - diagnostic::validate_tile_program(tile, schedule, &program)?; - Ok(Some(program)) - }) - .collect::, _>>()?; - Ok(PhasePrograms { - programs, - tile_event_cycles, - event_cycles, - }) - } -} - -impl TileProgramSchedule { - #[cfg(test)] - fn event_cycles(&self) -> u32 { - self.event_cycles - } - - /// Advances a requested transfer offset until its outgoing message does - /// not overlap another outgoing message on this tile. Receive controls may - /// be represented by the ISA's composite send/control encodings; this is a - /// single supervisor instruction, not dual issue from independent lanes. - fn earliest_sender_offset(&self, row: &PlanRow, requested: u32) -> Result { - let base = sender_row_timing(row, 0)?; - let mut offset = requested.max(self.reserved_sender_end.saturating_sub(base.start_cycles)); - loop { - let start = base - .start_cycles - .checked_add(offset) - .ok_or(ExchangeError::Schedule("send offset overflow"))?; - let end = base - .end_cycles - .checked_add(offset) - .ok_or(ExchangeError::Schedule("send offset overflow"))?; - let conflicting_sender = self - .senders - .iter() - .find(|sender| start < sender.end_cycles && sender.start_cycles < end); - if let Some(sender) = conflicting_sender { - offset = sender - .end_cycles - .checked_sub(base.start_cycles) - .ok_or(ExchangeError::Schedule("send offset order"))?; - continue; - } - // A receive control cannot be encoded before the first outgoing - // word. Leave at least one continuation word so SENDPIC can carry - // a control at the following event. - if self - .receive_events - .iter() - .any(|event| event.cycles == start || event.cycles == start.saturating_add(1)) - { - offset = offset - .checked_add(1) - .ok_or(ExchangeError::Schedule("send offset overflow"))?; - continue; - } - return Ok(offset); - } - } - - /// Advances a requested transfer offset while preserving independent - /// source-selection and local receive-address streams. - fn earliest_receiver_offset( - &self, - row: &PlanRow, - received_words: u32, - requested: u32, - ) -> Result { - let base = receive_row_timing(row, 0)?; - let mut offset = requested; - if let Some(stream) = &self.receive_stream { - let source_cycles = base - .source_cycles - .or(base.format_start_cycles) - .ok_or(ExchangeError::Schedule("receive source timing"))?; - offset = offset.max(stream.source_end_cycles.saturating_sub(source_cycles)); - // Paired format activates two events before its first SRAM write. - // It must not reinterpret the tail of an ordinary receive that is - // still travelling from XPIC to the local payload stream. - if stream.mode == ReceiveMode::Ordinary - && let Some(format_start) = base.format_start_cycles - { - offset = offset.max(stream.pointer_end_cycles.saturating_sub(format_start)); - } - let pointer_cycles = base - .pointer_cycles - .ok_or(ExchangeError::Schedule("receive pointer event"))?; - offset = offset.max(stream.pointer_end_cycles.saturating_sub(pointer_cycles)); - } - loop { - let timing = scheduled_receive_window( - &base, - offset, - received_words, - self.receive_stream.as_ref(), - )?; - let collision = timing.events.iter().any(|new| { - self.receive_events.iter().any(|existing| { - new.cycles == existing.cycles - && !self.receive_stream.as_ref().is_some_and(|previous| { - replaces_receive_event(*existing, base.mode, &timing, previous) - }) - && !receive_events_can_share_instruction(*new, *existing) - }) - }); - let sender_boundary = timing.events.iter().any(|event| { - self.senders.iter().any(|sender| { - event.cycles == sender.start_cycles - || event.cycles == sender.start_cycles.saturating_add(1) - }) - }); - if !collision && !sender_boundary { - return Ok(offset); - } - offset = offset - .checked_add(1) - .ok_or(ExchangeError::Schedule("receive offset overflow"))?; - } - } - - /// Appends a primitive row at an arbitrary absolute phase offset without - /// requiring that offset to fit in the primitive row's spare words. - fn append_sender_at( - &mut self, - row: &PlanRow, - schedule_offset: u32, - ) -> Result<(), ExchangeError> { - let timing = sender_row_timing(row, schedule_offset)?; - if self.senders.iter().any(|sender| { - timing.start_cycles < sender.end_cycles && sender.start_cycles < timing.end_cycles - }) { - return Err(ExchangeError::Schedule("overlapping outgoing messages")); - } - if self.receive_events.iter().any(|event| { - event.cycles == timing.start_cycles - || event.cycles == timing.start_cycles.saturating_add(1) - }) { - return Err(ExchangeError::Schedule("unencodable initial send control")); - } - self.event_cycles = self.event_cycles.max(timing.horizon_cycles); - self.senders.push(ScheduledSenderRow { - row: *row, - start_cycles: timing.start_cycles, - end_cycles: timing.end_cycles, - }); - Ok(()) - } - - /// Appends a receive row, merging its timed control writes with the - /// current receive stream and replacing an immediately preceding neutral - /// mux selection with a direct source cutover. - fn append_receiver_at( - &mut self, - row: &PlanRow, - schedule_offset: u32, - received_words: u32, - ) -> Result { - let base = receive_row_timing(row, 0)?; - let timing = scheduled_receive_window( - &base, - schedule_offset, - received_words, - self.receive_stream.as_ref(), - )?; - let next_address = timing - .pointer_address - .map(|address| { - received_words - .checked_mul(match base.mode { - ReceiveMode::Ordinary => 4, - ReceiveMode::Paired64 => 8, - }) - .and_then(|bytes| address.checked_add(bytes)) - .ok_or(ExchangeError::Schedule("receive address overflow")) - }) - .transpose()?; - let previous = self.receive_stream.take(); - if let Some(stream) = &previous { - if timing.source_start < stream.source_end_cycles - || timing.payload_start < stream.pointer_end_cycles - { - return Err(ExchangeError::Schedule("overlapping receive streams")); - } - self.receive_events - .retain(|event| !replaces_receive_event(*event, base.mode, &timing, stream)); - } - - // `earliest_receiver_offset` checked the new events against the full - // existing stream, while `scheduled_receive_window` validated the new - // group internally. Source and pointer controls are independent and - // may be inserted on opposite sides of an older teardown event, so - // keep insertion order here and sort once when encoding the row. - self.receive_events.extend(timing.events.iter().copied()); - self.event_cycles = self.event_cycles.max(timing.horizon); - self.receive_stream = Some(ReceiveStream { - mode: base.mode, - paired_format: base.paired_format(), - format_end_cycles: base.format_end_cycles.map(|end| end + schedule_offset), - source_end_cycles: timing.source_end, - pointer_end_cycles: timing.payload_end, - next_address, - }); - Ok(timing) - } - - pub fn finish(&self) -> Result, ExchangeError> { - build_scheduled_program(&self.senders, &self.receive_events, self.event_cycles) - } -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum ReceiveEventKind { - OrdinarySource, - OrdinaryNeutral, - PairedSource, - PairedNeutral, - PairedPointer, - Pointer, - Format, -} - -impl ReceiveEventKind { - fn is_xpic(self) -> bool { - matches!( - self, - Self::OrdinarySource | Self::OrdinaryNeutral | Self::PairedSource | Self::PairedNeutral - ) - } - - fn is_pic(self) -> bool { - matches!(self, Self::Pointer | Self::PairedPointer | Self::Format) - } -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum ReceiveMode { - Ordinary, - Paired64, -} - -#[derive(Clone, Copy, Debug)] -struct ReceiveEvent { - cycles: u32, - instruction: u32, - kind: ReceiveEventKind, -} - -fn receive_events_can_share_instruction(left: ReceiveEvent, right: ReceiveEvent) -> bool { - // Only ordinary pointer/mux controls are supported in SENDPICP. A paired - // pointer cannot share with an ordinary XPIC from a subsequent transfer. Combining - // paired XPIC with an ordinary pointer also faults (not just format/XPIC), - // including when the controls belong to different overlapping transfers. - let ordinary_pair = |pic: ReceiveEventKind, xpic: ReceiveEventKind| { - pic == ReceiveEventKind::Pointer - && matches!( - xpic, - ReceiveEventKind::OrdinarySource | ReceiveEventKind::OrdinaryNeutral - ) - }; - left.cycles == right.cycles - && (ordinary_pair(left.kind, right.kind) || ordinary_pair(right.kind, left.kind)) -} - -fn validate_receive_events(events: &[ReceiveEvent]) -> Result<(), ExchangeError> { - let mut cursor = 0; - while cursor < events.len() { - let end = cursor - + events[cursor..].partition_point(|event| event.cycles == events[cursor].cycles); - let group = &events[cursor..end]; - if group.len() > 2 - || (group.len() == 2 && !receive_events_can_share_instruction(group[0], group[1])) - { - return Err(ExchangeError::Schedule( - "incompatible simultaneous receive controls", - )); - } - cursor = end; - } - Ok(()) -} - -struct ReceiveRowTiming { - mode: ReceiveMode, - events: Vec, - neutral_cycles: Option, - source_cycles: Option, - format_start_cycles: Option, - format_end_cycles: Option, - pointer_cycles: Option, - horizon_cycles: u32, - pointer_address: Option, -} - -#[derive(Clone, Debug)] -struct ScheduledReceiverWindow { - continues_format: bool, - events: Vec, - source_start: u32, - source_end: u32, - payload_start: u32, - payload_end: u32, - horizon: u32, - pointer_address: Option, -} - -impl ReceiveRowTiming { - fn paired_format(&self) -> Option { - self.events - .iter() - .find(|event| { - event.kind == ReceiveEventKind::Format - && Some(event.cycles) == self.format_start_cycles - }) - .map(|event| event.instruction & PIC_RECEIVE_ADDRESS_MASK) - } -} - -fn replaces_receive_event( - event: ReceiveEvent, - mode: ReceiveMode, - timing: &ScheduledReceiverWindow, - previous: &ReceiveStream, -) -> bool { - if timing.continues_format - && event.kind == ReceiveEventKind::Format - && Some(event.cycles) == previous.format_end_cycles - { - return true; - } - timing.source_start == previous.source_end_cycles - && event.cycles == previous.source_end_cycles - && match (previous.mode, mode, event.kind) { - (ReceiveMode::Ordinary, ReceiveMode::Ordinary, ReceiveEventKind::OrdinaryNeutral) => { - true - } - (ReceiveMode::Paired64, ReceiveMode::Paired64, ReceiveEventKind::PairedNeutral) => { - timing.continues_format - } - _ => false, - } -} - -fn scheduled_receive_window( - base: &ReceiveRowTiming, - schedule_offset: u32, - received_words: u32, - previous: Option<&ReceiveStream>, -) -> Result { - let timing = receive_row_timing_from_base(base, schedule_offset)?; - // The incoming source is occupied for exactly one event per word. The - // pointer stream has a physical-row phase of its own and must not extend - // source ownership: SDK full-duplex rows switch XPIC at this boundary - // while their PIC update and outgoing SEND stream continue independently. - let source_start = timing - .source_cycles - .or(timing.format_start_cycles) - .ok_or(ExchangeError::Schedule("receive source timing"))?; - let source_end = source_start - .checked_add(received_words) - .ok_or(ExchangeError::Schedule("receive source timing overflow"))?; - let continues_format = timing.paired_format().is_some_and(|format| { - previous.is_some_and(|stream| { - stream.paired_format == Some(format) - && stream.format_end_cycles == timing.format_start_cycles - }) - }); - let carries_pointer = timing.mode == ReceiveMode::Ordinary - && previous.is_some_and(|stream| { - stream.mode == ReceiveMode::Ordinary && timing.pointer_address == stream.next_address - }); - // PIC may retain a contiguous destination address, but the payload still - // follows this transfer's XPIC source cutover through the route pipeline. - // The primitive row's PIC event gives that arrival offset even when the - // redundant address write itself can be omitted. - let payload_start = timing - .pointer_cycles - .ok_or(ExchangeError::Schedule("receive pointer event"))?; - let payload_end = payload_start - .checked_add(received_words) - .ok_or(ExchangeError::Schedule("receive payload timing overflow"))?; - - // Primitive rows serialize the independent XPIC and PIC control streams. - // Their final delay preserves a small guard after both the remote source - // window and the local SRAM write window have completed. Retain that - // guard after the two streams are interleaved with other transfers. - let guard = base - .horizon_cycles - .saturating_sub(source_end.max(payload_end) - schedule_offset); - let horizon = source_end - .max(payload_end) - .checked_add(guard) - .ok_or(ExchangeError::Schedule("receive horizon overflow"))?; - - let mut events = timing - .events - .into_iter() - .filter_map(|mut event| match event.kind { - ReceiveEventKind::OrdinarySource => Some(event), - ReceiveEventKind::PairedPointer => Some(event), - ReceiveEventKind::Pointer if !carries_pointer => Some(event), - ReceiveEventKind::Pointer => None, - ReceiveEventKind::OrdinaryNeutral => { - event.cycles = source_end; - Some(event) - } - ReceiveEventKind::Format - if continues_format && Some(event.cycles) == timing.format_start_cycles => - { - None - } - ReceiveEventKind::PairedSource - | ReceiveEventKind::PairedNeutral - | ReceiveEventKind::Format => Some(event), - }) - .collect::>(); - events.sort_by_key(|event| event.cycles); - validate_receive_events(&events)?; - Ok(ScheduledReceiverWindow { - continues_format, - events, - source_start, - source_end, - payload_start, - payload_end, - horizon, - pointer_address: timing.pointer_address, - }) -} - -fn receive_row_timing( - row: &PlanRow, - schedule_offset: u32, -) -> Result { - if row[0] != SYNC_SUPERVISOR_INSTRUCTION { - return Err(ExchangeError::Schedule("receive row entry")); - } - let end = row - .iter() - .position(|instruction| *instruction == RETURN_M10_INSTRUCTION) - .ok_or(ExchangeError::Schedule("receive row return"))?; - let mut cycles = schedule_offset; - let mut events = Vec::new(); - let mut source_cycles = None; - let mut neutral_cycles = None; - let mut format_start_cycles = None; - let mut format_end_cycles = None; - let mut pointer_cycles = None; - let mut pointer_address = None; - let mut cursor = 1; - while cursor < end { - let instruction = row[cursor]; - if is_send_control_pair(instruction) { - let payload = *row - .get(cursor + 1) - .ok_or(ExchangeError::Schedule("truncated SENDPICP payload"))?; - let control_cycles = cycles - .checked_add(1) - .ok_or(ExchangeError::Schedule("receive event horizon overflow"))?; - let xpic_value = payload >> 18; - let paired = xpic_value & (1 << 13) != 0; - let xpic_kind = if xpic_value & 0x1fff == TILE_MUX_EXCHANGE { - if paired { - ReceiveEventKind::PairedNeutral - } else { - ReceiveEventKind::OrdinaryNeutral - } - } else if paired { - ReceiveEventKind::PairedSource - } else { - ReceiveEventKind::OrdinarySource - }; - events.push(ReceiveEvent { - cycles: control_cycles, - instruction: delay_xpic(0, u32::from(paired), xpic_value & 0x1fff), - kind: xpic_kind, - }); - let pic_selector = (instruction >> 27) & 1; - let pic_value = payload & PIC_RECEIVE_ADDRESS_MASK; - events.push(ReceiveEvent { - cycles: control_cycles, - instruction: delay_pic(0, pic_selector, pic_value), - kind: if pic_selector == 0 { - ReceiveEventKind::Pointer - } else { - ReceiveEventKind::Format - }, - }); - cycles = cycles - .checked_add(instruction_advance(instruction)) - .ok_or(ExchangeError::Schedule("receive event horizon overflow"))?; - cursor += 2; - continue; - } - - cycles = cycles - .checked_add(instruction_advance(instruction)) - .ok_or(ExchangeError::Schedule("receive event horizon overflow"))?; - let kind = if instruction & OPCODE_MASK == DELAY_PIC_OPCODE { - Some(if instruction & (1 << 18) == 0 { - ReceiveEventKind::Pointer - } else { - ReceiveEventKind::Format - }) - } else if instruction & OPCODE_MASK == DELAY_XPIC_OPCODE { - let paired = instruction & (1 << 13) != 0; - Some(if instruction & 0x1fff == TILE_MUX_EXCHANGE { - if paired { - ReceiveEventKind::PairedNeutral - } else { - ReceiveEventKind::OrdinaryNeutral - } - } else if paired { - ReceiveEventKind::PairedSource - } else { - ReceiveEventKind::OrdinarySource - }) - } else { - None - }; - if let Some(kind) = kind { - events.push(ReceiveEvent { - cycles, - instruction, - kind, - }); - } - cursor += 1; - } - - for event in &events { - match event.kind { - ReceiveEventKind::OrdinarySource | ReceiveEventKind::PairedSource => { - if source_cycles.replace(event.cycles).is_some() { - return Err(ExchangeError::Schedule("multiple receive sources")); - } - } - ReceiveEventKind::OrdinaryNeutral | ReceiveEventKind::PairedNeutral => { - if neutral_cycles.replace(event.cycles).is_some() { - return Err(ExchangeError::Schedule("multiple receive teardowns")); - } - } - ReceiveEventKind::Pointer | ReceiveEventKind::PairedPointer => { - pointer_cycles = Some(event.cycles); - if pointer_address - .replace((event.instruction & PIC_RECEIVE_ADDRESS_MASK) << 2) - .is_some() - { - return Err(ExchangeError::Schedule("multiple receive pointers")); - } - } - ReceiveEventKind::Format => { - let value = event.instruction & PIC_RECEIVE_ADDRESS_MASK; - if matches!(value, 1 | 2) { - if format_start_cycles.replace(event.cycles).is_some() { - return Err(ExchangeError::Schedule("multiple paired format starts")); - } - } else if value == 0 { - if format_end_cycles.replace(event.cycles).is_some() { - return Err(ExchangeError::Schedule("multiple paired format ends")); - } - } else { - return Err(ExchangeError::Schedule("paired receive format value")); - } - } - } - } - let mode = if format_start_cycles.is_some() - || events.iter().any(|event| { - matches!( - event.kind, - ReceiveEventKind::PairedSource | ReceiveEventKind::PairedNeutral - ) - }) { - if format_start_cycles.is_none() || format_end_cycles.is_none() { - return Err(ExchangeError::Schedule("incomplete paired receive format")); - } - ReceiveMode::Paired64 - } else { - if source_cycles.is_none() || neutral_cycles.is_none() { - return Err(ExchangeError::Schedule("incomplete ordinary receive mux")); - } - ReceiveMode::Ordinary - }; - if mode == ReceiveMode::Paired64 { - for event in &mut events { - if event.kind == ReceiveEventKind::Pointer { - event.kind = ReceiveEventKind::PairedPointer; - } - } - } - Ok(ReceiveRowTiming { - mode, - events, - neutral_cycles, - source_cycles, - format_start_cycles, - format_end_cycles, - pointer_cycles, - horizon_cycles: cycles, - pointer_address, - }) -} - -struct SenderRowTiming { - start_cycles: u32, - end_cycles: u32, - horizon_cycles: u32, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct ScheduledSenderTiming { - pub payload_start: u32, - pub payload_end: u32, - pub horizon: u32, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct ScheduledReceiverTiming { - pub source_event: u32, - pub pointer_event: Option, - pub source_teardown: u32, - pub horizon: u32, -} - -pub fn scheduled_sender_timing( - row: &PlanRow, - schedule_offset: u32, -) -> Result { - let timing = sender_row_timing(row, schedule_offset)?; - Ok(ScheduledSenderTiming { - payload_start: timing.start_cycles, - payload_end: timing.end_cycles, - horizon: timing.horizon_cycles, - }) -} - -pub fn scheduled_receiver_timing( - row: &PlanRow, - schedule_offset: u32, -) -> Result { - let timing = receive_row_timing(row, schedule_offset)?; - Ok(ScheduledReceiverTiming { - source_event: timing - .source_cycles - .or(timing.format_start_cycles) - .ok_or(ExchangeError::Schedule("receive source timing"))?, - pointer_event: timing.pointer_cycles, - source_teardown: timing - .neutral_cycles - .or(timing.format_end_cycles) - .ok_or(ExchangeError::Schedule("receive teardown timing"))?, - horizon: timing.horizon_cycles, - }) -} - -fn sender_row_timing( - row: &PlanRow, - schedule_offset: u32, -) -> Result { - if row[0] != SYNC_SUPERVISOR_INSTRUCTION { - return Err(ExchangeError::Schedule("sender row entry")); - } - let end = row - .iter() - .position(|instruction| *instruction == RETURN_M10_INSTRUCTION) - .ok_or(ExchangeError::Schedule("sender row return"))?; - let mut cycles = schedule_offset; - let mut start_cycles = None; - let mut end_cycles = None; - for &instruction in &row[1..end] { - if is_payload_send(instruction) { - start_cycles.get_or_insert(cycles); - cycles = cycles - .checked_add(instruction_advance(instruction)) - .ok_or(ExchangeError::Schedule("sender event horizon overflow"))?; - end_cycles = Some(cycles); - } else { - cycles = cycles - .checked_add(instruction_advance(instruction)) - .ok_or(ExchangeError::Schedule("sender event horizon overflow"))?; - } - } - Ok(SenderRowTiming { - start_cycles: start_cycles.ok_or(ExchangeError::Schedule("sender payload"))?, - end_cycles: end_cycles.ok_or(ExchangeError::Schedule("sender payload"))?, - horizon_cycles: cycles, - }) -} - -fn receive_row_timing_from_base( - base: &ReceiveRowTiming, - offset: u32, -) -> Result { - let shift = |cycles: u32| { - cycles - .checked_add(offset) - .ok_or(ExchangeError::Schedule("receive offset overflow")) - }; - Ok(ReceiveRowTiming { - mode: base.mode, - events: base - .events - .iter() - .map(|event| { - Ok(ReceiveEvent { - cycles: shift(event.cycles)?, - ..*event - }) - }) - .collect::>()?, - neutral_cycles: base.neutral_cycles.map(shift).transpose()?, - source_cycles: base.source_cycles.map(shift).transpose()?, - format_start_cycles: base.format_start_cycles.map(shift).transpose()?, - format_end_cycles: base.format_end_cycles.map(shift).transpose()?, - pointer_cycles: base.pointer_cycles.map(shift).transpose()?, - horizon_cycles: shift(base.horizon_cycles)?, - pointer_address: base.pointer_address, - }) -} - -fn build_scheduled_program( - senders: &[ScheduledSenderRow], - receive_events: &[ReceiveEvent], - horizon_cycles: u32, -) -> Result, ExchangeError> { - let mut senders = senders.iter().collect::>(); - senders.sort_by_key(|sender| sender.start_cycles); - if senders - .windows(2) - .any(|pair| pair[0].end_cycles > pair[1].start_cycles) - { - return Err(ExchangeError::Schedule("overlapping outgoing messages")); - } - let mut events = receive_events.to_vec(); - events.sort_by_key(|event| event.cycles); - validate_receive_events(&events)?; - - let mut words = Vec::new(); - let mut event_cycles = 0; - let mut event_index = 0; - for sender in senders { - let before_start = - events[event_index..].partition_point(|event| event.cycles < sender.start_cycles); - let split = event_index + before_start; - append_receive_events( - &mut words, - &mut event_cycles, - &events[event_index..split], - sender.start_cycles, - true, - )?; - event_index = split; - let through_end = - events[event_index..].partition_point(|event| event.cycles <= sender.end_cycles); - let split = event_index + through_end; - append_sender_message( - &mut words, - &mut event_cycles, - sender, - &events[event_index..split], - )?; - event_index = split; - } - append_receive_events( - &mut words, - &mut event_cycles, - &events[event_index..], - horizon_cycles, - true, - )?; - words.push(RETURN_M10_INSTRUCTION); - debug_assert_eq!(plan_event_cycles(&words)?, horizon_cycles); - Ok(words) -} - -fn append_sender_message( - words: &mut Vec, - event_cycles: &mut u32, - sender: &ScheduledSenderRow, - controls: &[ReceiveEvent], -) -> Result<(), ExchangeError> { - append_plain_delay(words, event_cycles, sender.start_cycles)?; - let (initial_instruction, payload_words) = sender_payload(&sender.row)?; - let direction = initial_instruction & 7; - let initial_source = (initial_instruction & SEND_ADDRESS_MASK) >> 3; - let mut remaining = payload_words; - let mut sent = 0u32; - let mut started = false; - - let mut cursor = 0; - while cursor < controls.len() { - let end = cursor - + controls[cursor..].partition_point(|event| event.cycles == controls[cursor].cycles); - let group = &controls[cursor..end]; - let control_start = group[0] - .cycles - .checked_sub(1) - .ok_or(ExchangeError::Schedule("send control at phase entry"))?; - let distance = control_start - .checked_sub(*event_cycles) - .ok_or(ExchangeError::Schedule("send control order"))?; - if distance >= remaining { - return Err(ExchangeError::Schedule("send control beyond payload")); - } - emit_sender_words( - words, - event_cycles, - initial_instruction, - direction, - &mut remaining, - &mut sent, - &mut started, - distance, - (group.len() == 2).then_some(0), - )?; - - let next_boundary = controls - .get(end) - .map_or(sender.end_cycles, |next| next.cycles.saturating_sub(1)); - let mut available = next_boundary - .checked_sub(*event_cycles) - .ok_or(ExchangeError::Schedule("send control order"))?; - let next_is_pair = controls.get(end + 1).is_some_and(|next| { - controls - .get(end) - .is_some_and(|first| next.cycles == first.cycles) - }); - let words_after_control = words.len() + group.len(); - if next_is_pair && words_after_control & 1 != 0 && available > 1 { - // Leave one outgoing word for a SENDOFF instruction that aligns - // the following two-word SENDPICP without changing its event. - available -= 1; - } - let chunk = available.min(64).min(remaining); - if chunk == 0 { - return Err(ExchangeError::Schedule( - "empty receive-control continuation", - )); - } - if group.len() == 2 { - let source = initial_source - .checked_add(sent) - .ok_or(ExchangeError::Schedule("SENDPICP source overflow"))?; - let (instruction, payload) = - encode_send_control_pair(chunk - 1, source, direction, group)?; - words.extend([instruction, payload]); - } else { - words.push(encode_send_control(chunk - 1, group[0])?); - } - *event_cycles += chunk; - remaining -= chunk; - sent += chunk; - cursor = end; - } - - if !started { - let first = remaining.min(64); - words.push(resize_send(initial_instruction, first)?); - *event_cycles += first; - remaining -= first; - } - while remaining != 0 { - let chunk = remaining.min(4096); - words.push(send_off(chunk - 1, direction, 0)); - *event_cycles += chunk; - remaining -= chunk; - } - if *event_cycles != sender.end_cycles { - return Err(ExchangeError::Schedule("sender payload horizon")); - } - Ok(()) -} - -#[allow(clippy::too_many_arguments)] -fn emit_sender_words( - words: &mut Vec, - event_cycles: &mut u32, - initial_instruction: u32, - direction: u32, - remaining: &mut u32, - sent: &mut u32, - started: &mut bool, - count: u32, - final_word_parity: Option, -) -> Result<(), ExchangeError> { - if count > *remaining { - return Err(ExchangeError::Schedule("sender payload overflow")); - } - if count == 0 { - if !*started { - return Err(ExchangeError::Schedule("unencodable initial send control")); - } - if final_word_parity.is_some_and(|parity| words.len() & 1 != parity) { - return Err(ExchangeError::Schedule("SENDPICP instruction alignment")); - } - return Ok(()); - } - - let first_limit = if *started { 4096 } else { 64 }; - let mut chunks = vec![count.min(first_limit)]; - let mut unscheduled = count - chunks[0]; - while unscheduled != 0 { - let chunk = unscheduled.min(4096); - chunks.push(chunk); - unscheduled -= chunk; - } - if final_word_parity.is_some_and(|parity| (words.len() + chunks.len()) & 1 != parity) { - let split = chunks - .iter() - .rposition(|chunk| *chunk >= 2) - .ok_or(ExchangeError::Schedule("SENDPICP instruction alignment"))?; - let tail = chunks[split] - 1; - chunks[split] = 1; - chunks.insert(split + 1, tail); - } - - for chunk in chunks { - if !*started { - words.push(resize_send(initial_instruction, chunk)?); - *started = true; - } else { - words.push(send_off(chunk - 1, direction, 0)); - } - *event_cycles += chunk; - *remaining -= chunk; - *sent += chunk; - } - if final_word_parity.is_some_and(|parity| words.len() & 1 != parity) { - return Err(ExchangeError::Schedule("SENDPICP instruction alignment")); - } - Ok(()) -} - -fn sender_payload(row: &PlanRow) -> Result<(u32, u32), ExchangeError> { - let mut initial = None; - let mut words = 0u32; - for &instruction in row { - if instruction & LONG_OPCODE_MASK == SEND_OPCODE { - if initial.replace(instruction).is_some() { - return Err(ExchangeError::Schedule("multiple initial sends")); - } - words = words - .checked_add(instruction_advance(instruction)) - .ok_or(ExchangeError::Schedule("sender payload overflow"))?; - } else if is_send_off(instruction) { - words = words - .checked_add(instruction_advance(instruction)) - .ok_or(ExchangeError::Schedule("sender payload overflow"))?; - } - } - Ok(( - initial.ok_or(ExchangeError::Schedule("sender payload"))?, - words, - )) -} - -fn resize_send(instruction: u32, words: u32) -> Result { - if instruction & LONG_OPCODE_MASK != SEND_OPCODE || !(1..=64).contains(&words) { - return Err(ExchangeError::Schedule("initial send size")); - } - Ok((instruction & !SEND_COUNT_MASK) | ((words - 1) << 21)) -} - -fn encode_send_control(count_minus_one: u32, event: ReceiveEvent) -> Result { - if count_minus_one > 63 { - return Err(ExchangeError::Schedule("SENDPIC count")); - } - let (selector, operand) = match event.kind { - ReceiveEventKind::OrdinarySource - | ReceiveEventKind::OrdinaryNeutral - | ReceiveEventKind::PairedSource - | ReceiveEventKind::PairedNeutral => { - ((event.instruction >> 13) & 1, event.instruction & 0x1fff) - } - ReceiveEventKind::Pointer | ReceiveEventKind::PairedPointer | ReceiveEventKind::Format => ( - 2 + ((event.instruction >> 18) & 1), - event.instruction & PIC_RECEIVE_ADDRESS_MASK, - ), - }; - Ok(SEND_PIC_OPCODE | (count_minus_one << 21) | (selector << 18) | operand) -} - -fn append_receive_events( - words: &mut Vec, - event_cycles: &mut u32, - events: &[ReceiveEvent], - horizon_cycles: u32, - align_control_pairs: bool, -) -> Result<(), ExchangeError> { - let mut cursor = 0; - while cursor < events.len() { - let end = cursor - + events[cursor..].partition_point(|event| event.cycles == events[cursor].cycles); - let group = &events[cursor..end]; - if group.len() == 2 { - let instruction_start = group[0] - .cycles - .checked_sub(1) - .ok_or(ExchangeError::Schedule("receive control at phase entry"))?; - if align_control_pairs { - append_plain_delay_aligned(words, event_cycles, instruction_start, 0)?; - } else { - // Primitive rows are a declarative timing representation and - // are rebuilt phase-wide before execution. Avoid consuming a - // spare word merely to align an instruction which is never - // executed at this intermediate address. - append_plain_delay(words, event_cycles, instruction_start)?; - } - let next_start = events - .get(end) - .map_or(horizon_cycles, |next| next.cycles.saturating_sub(1)); - let advance = next_start - .checked_sub(*event_cycles) - .ok_or(ExchangeError::Schedule("receive control order"))? - .min(64); - if advance == 0 { - return Err(ExchangeError::Schedule("empty SENDPICP interval")); - } - let (instruction, payload) = encode_send_control_pair(advance - 1, 0, 0, group)?; - words.extend([instruction, payload]); - *event_cycles += advance; - } else { - let event = group[0]; - let advance = event - .cycles - .checked_sub(*event_cycles) - .ok_or(ExchangeError::Schedule("receive event order"))?; - if advance == 0 { - return Err(ExchangeError::Schedule("receive control issue interval")); - } - let mut instruction = event.instruction; - let maximum_advance = maximum_timed_instruction_advance(instruction) - .ok_or(ExchangeError::Schedule("receive timed event"))?; - if advance > maximum_advance { - append_plain_delay(words, event_cycles, event.cycles - maximum_advance)?; - } - set_instruction_advance(&mut instruction, event.cycles - *event_cycles)?; - words.push(instruction); - *event_cycles = event.cycles; - } - cursor = end; - } - append_plain_delay(words, event_cycles, horizon_cycles) -} - -fn encode_send_control_pair( - count_minus_one: u32, - source_word_address: u32, - send_control: u32, - events: &[ReceiveEvent], -) -> Result<(u32, u32), ExchangeError> { - if count_minus_one > 63 - || source_word_address > SEND_ADDRESS_MASK >> 3 - || send_control > 7 - || events.len() != 2 - { - return Err(ExchangeError::Schedule("SENDPICP operand")); - } - let pic = events - .iter() - .find(|event| event.kind.is_pic()) - .ok_or(ExchangeError::Schedule("SENDPICP PIC control"))?; - let xpic = events - .iter() - .find(|event| event.kind.is_xpic()) - .ok_or(ExchangeError::Schedule("SENDPICP XPIC control"))?; - if pic.cycles != xpic.cycles { - return Err(ExchangeError::Schedule("SENDPICP event time")); - } - // SENDPICP uses the same outgoing fields as SEND: an absolute source word - // address followed by the three-bit SCTL field. Its fourth operand carries - // the high PIC selector bit; the inline word holds all fourteen XPIC bits - // and the remaining eighteen PIC bits. - let pointer_selector = (pic.instruction >> 18) & 1; - let instruction = SEND_PICP_OPCODE - | (pointer_selector << 27) - | (count_minus_one << 21) - | ((source_word_address << 3) & SEND_ADDRESS_MASK) - | send_control; - let payload = - ((xpic.instruction & 0x3fff) << 18) | (pic.instruction & PIC_RECEIVE_ADDRESS_MASK); - Ok((instruction, payload)) -} - -fn append_plain_delay( - words: &mut Vec, - event_cycles: &mut u32, - target_cycles: u32, -) -> Result<(), ExchangeError> { - let mut remaining = target_cycles - .checked_sub(*event_cycles) - .ok_or(ExchangeError::Schedule("plan delay order"))?; - while remaining != 0 { - let chunk = remaining.min(MAX_PLAN_OFFSET_CYCLES); - words.push(delay(chunk - 1)); - *event_cycles += chunk; - remaining -= chunk; - } - Ok(()) -} - -/// Advances with plain delays while arranging the next word at the requested -/// two-word-instruction alignment. Exchange rows are placed at eight-byte -/// boundaries, so `word_parity == 0` aligns a following SENDPICP and payload. -fn append_plain_delay_aligned( - words: &mut Vec, - event_cycles: &mut u32, - target_cycles: u32, - word_parity: usize, -) -> Result<(), ExchangeError> { - let distance = target_cycles - .checked_sub(*event_cycles) - .ok_or(ExchangeError::Schedule("plan delay order"))?; - if distance == 0 { - return (words.len() & 1 == word_parity) - .then_some(()) - .ok_or(ExchangeError::Schedule("SENDPICP instruction alignment")); - } - - let minimum_words = distance.div_ceil(MAX_PLAN_OFFSET_CYCLES); - let mut delay_words = minimum_words; - if ((words.len() + delay_words as usize) & 1) != word_parity { - delay_words += 1; - } - if delay_words > distance { - return Err(ExchangeError::Schedule("SENDPICP instruction alignment")); - } - let mut remaining = distance; - for remaining_words in (1..=delay_words).rev() { - let chunk = remaining - .saturating_sub(remaining_words - 1) - .min(MAX_PLAN_OFFSET_CYCLES); - words.push(delay(chunk - 1)); - *event_cycles += chunk; - remaining -= chunk; - } - debug_assert_eq!(remaining, 0); - debug_assert_eq!(words.len() & 1, word_parity); - Ok(()) -} - -fn maximum_timed_instruction_advance(instruction: u32) -> Option { - match instruction & OPCODE_MASK { - DELAY_PIC_OPCODE => Some(0x80), - DELAY_XPIC_OPCODE => Some(0x1000), - _ => None, - } -} - -#[cfg(test)] -fn is_neutral_mux_teardown(instruction: u32) -> bool { - instruction & OPCODE_MASK == DELAY_XPIC_OPCODE && instruction & 0x1fff == TILE_MUX_EXCHANGE -} - -fn instruction_advance(instruction: u32) -> u32 { - if instruction & DELAY_OPCODE_MASK == DELAY_OPCODE { - (instruction & 0x7_ffff) + 1 - } else { - match instruction & OPCODE_MASK { - DELAY_PIC_OPCODE => ((instruction >> 19) & 0x7f) + 1, - DELAY_XPIC_OPCODE => ((instruction >> 14) & 0xfff) + 1, - _ if instruction & LONG_OPCODE_MASK == SEND_OPCODE => ((instruction >> 21) & 0x3f) + 1, - _ if is_send_control_pair(instruction) => ((instruction >> 21) & 0x3f) + 1, - _ if is_send_control(instruction) => ((instruction >> 21) & 0x3f) + 1, - _ if is_send_off(instruction) => { - (((instruction >> 21) & 0x3f) | (((instruction >> 14) & 0x3f) << 6)) + 1 - } - _ => 0, - } - } -} - -fn is_send_control(instruction: u32) -> bool { - instruction & SEND_CONTROL_OPCODE_MASK == SEND_PIC_OPCODE -} - -fn is_send_control_pair(instruction: u32) -> bool { - instruction & SEND_PICP_OPCODE_MASK == SEND_PICP_OPCODE -} - -fn is_send_off(instruction: u32) -> bool { - instruction & SEND_CONTROL_OPCODE_MASK == SEND_OFF_OPCODE -} - -fn is_payload_send(instruction: u32) -> bool { - instruction & LONG_OPCODE_MASK == SEND_OPCODE || is_send_off(instruction) -} - -fn set_instruction_advance(instruction: &mut u32, advance: u32) -> Result<(), ExchangeError> { - if advance == 0 { - return Err(ExchangeError::Schedule("zero event advance")); - } - let immediate = advance - 1; - if *instruction & DELAY_OPCODE_MASK == DELAY_OPCODE { - if immediate > 0x7_ffff { - return Err(ExchangeError::Schedule("delay advance overflow")); - } - *instruction = (*instruction & !0x7_ffff) | immediate; - } else if *instruction & OPCODE_MASK == DELAY_PIC_OPCODE { - if immediate > 0x7f { - return Err(ExchangeError::Schedule("PIC delay advance overflow")); - } - *instruction = (*instruction & !0x03f8_0000) | (immediate << 19); - } else if *instruction & OPCODE_MASK == DELAY_XPIC_OPCODE { - if immediate > 0xfff { - return Err(ExchangeError::Schedule("XPIC delay advance overflow")); - } - *instruction = (*instruction & !0x03ff_c000) | (immediate << 14); - } else { - return Err(ExchangeError::Schedule( - if is_payload_send(*instruction) || is_send_control(*instruction) { - "scheduled offset truncates a send" - } else { - "first scheduled event is not a delay" - }, - )); - } - Ok(()) -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] -pub struct HostPacketHeader { - pub word0: u32, - pub word1: u32, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct HostTransferChunk { - pub tile_address: u32, - pub host_offset: u32, - pub bytes: u32, - pub header: HostPacketHeader, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct TileToHostProgram { - pub instructions: Vec, - pub packet_words: Vec, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct HostHierarchy { - pub xreq_physical_tile: u16, - pub target_physical_tile: u16, -} - -pub fn host_hierarchy(target_physical_tile: u16) -> Result { - validate_host_tile(target_physical_tile)?; - Ok(HostHierarchy { - xreq_physical_tile: target_physical_tile & 0x3d, - target_physical_tile, - }) -} - -pub fn assemble_host_xreq_program( - target_physical_tile: u16, - packet_address: u32, -) -> Result { - assemble_host_xreq_program_for_targets(&[target_physical_tile], packet_address) -} - -pub fn assemble_host_xreq_program_for_targets( - target_physical_tiles: &[u16], - packet_address: u32, -) -> Result { - if target_physical_tiles.is_empty() { - return Err(ExchangeError::HostPacket); - } - if packet_address & 7 != 0 { - return Err(ExchangeError::HostPacket); - } - let mut bitmap = [0u32; 2]; - for &target_physical_tile in target_physical_tiles { - validate_host_tile(target_physical_tile)?; - let bitmap_index = - u32::from(target_physical_tile / 64) * 2 + u32::from((target_physical_tile >> 1) & 1); - if bitmap_index < XREQ_BITMAP0_BITS { - bitmap[0] |= 1 << bitmap_index; - } else { - bitmap[1] |= 1 << (bitmap_index - XREQ_BITMAP0_BITS); - } - } - Ok(TileToHostProgram { - instructions: vec![ - encode_send(1, 3, packet_address >> 2)?, - RETURN_M10_INSTRUCTION, - ], - packet_words: bitmap.to_vec(), - }) -} - -pub fn assemble_host_command_read_program( - packet_address: u32, - destination_address: u32, - host_offset: u32, -) -> Result { - if packet_address & 7 != 0 { - return Err(ExchangeError::HostPacket); - } - let request = host_to_tile_packet(0, destination_address, host_offset, 4)?; - let mut instructions = vec![ - setzi_m(8, TILE_MUX_HOST), - put_special_from_m8(INCOMING_MUX_REGISTER), - SYNC_HOST_INSTRUCTION, - setzi_m(8, 1), - put_special_from_m8(INCOMING_DCOUNT_REGISTER), - encode_send(1, 3, packet_address >> 2)?, - encode_send(1, 3, (packet_address + 8) >> 2)?, - SYNC_RECEIVE_INSTRUCTION, - ]; - append_local_host_completion(&mut instructions); - instructions.extend([ - SYNC_SUPERVISOR_INSTRUCTION, - delay(HOST_COMMAND_ROUTE_CYCLES - 1), - encode_send(0, 3, destination_address >> 2)?, - ]); - instructions.push(RETURN_M10_INSTRUCTION); - Ok(TileToHostProgram { - instructions, - packet_words: vec![1, 0, request.word0, request.word1], - }) -} - -pub fn assemble_host_to_tile_target_program( - physical_tile: u16, - tile_address: u32, - host_offset: u32, - bytes: u32, - packet_address: u32, -) -> Result { - let chunks = plan_host_to_tile(physical_tile, tile_address, host_offset, bytes)?; - if packet_address & 7 != 0 { - return Err(ExchangeError::HostPacket); - } - let mut instructions = vec![ - setzi_m(8, bytes / 4), - put_special_from_m8(INCOMING_DCOUNT_REGISTER), - encode_send(1, 3, packet_address >> 2)?, - ]; - for _ in 1..chunks.len() { - instructions.push(send_off(1, 3, 0)); - } - instructions.push(SYNC_RECEIVE_INSTRUCTION); - instructions.push(RETURN_M10_INSTRUCTION); - let packet_words = chunks - .iter() - .enumerate() - .flat_map(|(index, chunk)| { - let word0 = if index + 1 == chunks.len() { - chunk.header.word0 - } else { - chunk.header.word0 & !HOST_TO_TILE_STREAM_END_BITS - }; - [word0, chunk.header.word1] - }) - .collect(); - Ok(TileToHostProgram { - instructions, - packet_words, - }) -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum HostPacketSize { - Short, - Long, -} - -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] -pub struct Plan { - pub sender: PlanRow, - pub receiver: PlanRow, -} - -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] -pub struct MulticastPlan { - pub sender: PlanRow, - pub receivers: Vec, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct Topology { - logical_to_physical: Vec, -} - -#[derive(Debug, thiserror::Error, PartialEq, Eq)] -pub enum ExchangeError { - #[error("logical tile {0} is out of range")] - Tile(u16), - #[error("exchange endpoints must be distinct")] - DuplicateTile, - #[error("exchange count {0} is outside 1..={MAX_TRANSFER_WORDS}")] - Count(u32), - #[error("invalid multicast receiver set")] - ReceiverSet, - #[error("exchange schedule is not encodable: {0}")] - Schedule(&'static str), - #[error( - "exchange event offset {cycles} cycles from the phase-entry sync exceeds the encodable maximum of {maximum} cycles" - )] - PlanOffsetRange { cycles: u32, maximum: u32 }, - #[error("tile address 0x{0:x} is not encodable")] - Address(u32), - #[error("host exchange address or length is not encodable")] - HostPacket, -} - -pub fn host_to_tile_packet( - physical_tile: u16, - tile_address: u32, - host_offset: u32, - bytes: u32, -) -> Result { - validate_host_tile(physical_tile)?; - if tile_address < EXCHANGE_WINDOW_BASE || tile_address & 31 != 0 { - return Err(ExchangeError::HostPacket); - } - let exchange_address = (tile_address - EXCHANGE_WINDOW_BASE) >> 5; - if tile_address >= EXCHANGE_WINDOW_BASE + HOST_TO_TILE_WINDOW_BYTES { - return Err(ExchangeError::HostPacket); - } - let size = host_packet_size(host_offset, bytes)?; - let opcode = match size { - HostPacketSize::Short => 0xcc00_0200, - HostPacketSize::Long => 0xec00_0200, - }; - Ok(HostPacketHeader { - word0: opcode | host_route_word0(physical_tile) | exchange_address, - word1: host_route_word1(physical_tile) | host_address_length(host_offset, bytes, size)?, - }) -} - -pub fn tile_to_host_packet( - physical_tile: u16, - host_offset: u32, - bytes: u32, -) -> Result { - validate_host_tile(physical_tile)?; - let size = host_packet_size(host_offset, bytes)?; - let opcode = match size { - HostPacketSize::Short => 0x8000_0000, - HostPacketSize::Long => 0xa000_0000, - }; - Ok(HostPacketHeader { - word0: opcode | host_route_word0(physical_tile), - word1: host_route_word1(physical_tile) | host_address_length(host_offset, bytes, size)?, - }) -} - -pub fn zero_byte_read_packet( - physical_tile: u16, - dummy_tile_address: u32, -) -> Result { - validate_host_tile(physical_tile)?; - if dummy_tile_address < EXCHANGE_WINDOW_BASE || dummy_tile_address & 31 != 0 { - return Err(ExchangeError::HostPacket); - } - let exchange_address = (dummy_tile_address - EXCHANGE_WINDOW_BASE) >> 5; - if dummy_tile_address >= EXCHANGE_WINDOW_BASE + HOST_TO_TILE_WINDOW_BYTES { - return Err(ExchangeError::HostPacket); - } - Ok(HostPacketHeader { - word0: 0xcc00_0200 | host_route_word0(physical_tile) | exchange_address, - word1: host_route_word1(physical_tile), - }) -} - -pub fn plan_host_to_tile( - physical_tile: u16, - tile_address: u32, - host_offset: u32, - bytes: u32, -) -> Result, ExchangeError> { - if bytes == 0 || bytes & 3 != 0 || tile_address & 31 != 0 || host_offset & 3 != 0 { - return Err(ExchangeError::HostPacket); - } - plan_host_transfer( - tile_address, - host_offset, - bytes, - HOST_LONG_MAX_BYTES, - |tile, host, count| host_to_tile_packet(physical_tile, tile, host, count), - ) -} - -pub fn plan_tile_to_host( - physical_tile: u16, - tile_address: u32, - host_offset: u32, - bytes: u32, -) -> Result, ExchangeError> { - if bytes == 0 || bytes & 3 != 0 || tile_address & 3 != 0 || host_offset & 3 != 0 { - return Err(ExchangeError::HostPacket); - } - plan_host_transfer( - tile_address, - host_offset, - bytes, - TILE_TO_HOST_MAX_BYTES, - |_tile, host, count| tile_to_host_packet(physical_tile, host, count), - ) -} - -pub fn assemble_tile_to_host_target_program( - physical_tile: u16, - tile_address: u32, - host_offset: u32, - bytes: u32, - packet_address: u32, - command_address: u32, -) -> Result { - let chunks = plan_tile_to_host(physical_tile, tile_address, host_offset, bytes)?; - if packet_address & 7 != 0 { - return Err(ExchangeError::HostPacket); - } - let mut packet_words: Vec = chunks - .iter() - .flat_map(|chunk| [chunk.header.word0, chunk.header.word1]) - .collect(); - let close_address = packet_address - .checked_add(u32::try_from(packet_words.len() * 4).map_err(|_| ExchangeError::HostPacket)?) - .ok_or(ExchangeError::HostPacket)?; - let close = zero_byte_read_packet(physical_tile, command_address)?; - packet_words.extend([close.word0, close.word1]); - Ok(TileToHostProgram { - instructions: tile_to_host_target_instructions(&chunks, packet_address, close_address)?, - packet_words, - }) -} - -fn tile_to_host_target_instructions( - chunks: &[HostTransferChunk], - packet_address: u32, - close_address: u32, -) -> Result, ExchangeError> { - if chunks.is_empty() || packet_address & 7 != 0 { - return Err(ExchangeError::HostPacket); - } - let header_base = packet_address; - let mut instructions = vec![setzi_m(8, 1), put_special_from_m8(INCOMING_DCOUNT_REGISTER)]; - for (index, chunk) in chunks.iter().enumerate() { - instructions.push(encode_send(1, 3, (header_base + index as u32 * 8) >> 2)?); - instructions.push(encode_send( - chunk.bytes / 4 - 1, - 3, - chunk.tile_address >> 2, - )?); - let payload_events = chunk.bytes / 4; - if payload_events < TILE_TO_HOST_MIN_PAYLOAD_EVENTS { - instructions.push(delay(TILE_TO_HOST_MIN_PAYLOAD_EVENTS - payload_events - 1)); - } - } - instructions.push(encode_send(1, 3, close_address >> 2)?); - instructions.push(SYNC_RECEIVE_INSTRUCTION); - instructions.push(RETURN_M10_INSTRUCTION); - Ok(instructions) -} - -fn append_local_host_completion(instructions: &mut Vec) { - instructions.extend([ - SYNC_ALL_INSTRUCTION, - setzi_m(8, TILE_MUX_EXCHANGE), - put_special_from_m8(INCOMING_MUX_REGISTER), - ]); -} - -fn wrap_host_operation( - physical_tile: u16, - operation: &[u32], - entry_sync: u32, -) -> Result, ExchangeError> { - let Some((&RETURN_M10_INSTRUCTION, body)) = operation.split_last() else { - return Err(ExchangeError::Schedule("host target operation return")); - }; - let mut instructions = vec![ - setzi_m(8, host_mux_for_tile(physical_tile)?), - put_special_from_m8(INCOMING_MUX_REGISTER), - entry_sync, - ]; - instructions.extend_from_slice(body); - instructions.extend([ - SYNC_ALL_INSTRUCTION, - setzi_m(8, TILE_MUX_EXCHANGE), - put_special_from_m8(INCOMING_MUX_REGISTER), - RETURN_M10_INSTRUCTION, - ]); - Ok(instructions) -} - -pub fn wrap_host_xreq_operation( - physical_tile: u16, - operation: &[u32], -) -> Result, ExchangeError> { - wrap_host_operation(physical_tile, operation, SYNC_HOST_INSTRUCTION) -} - -pub fn wrap_host_target_operation( - physical_tile: u16, - operation: &[u32], -) -> Result, ExchangeError> { - wrap_host_operation(physical_tile, operation, SYNC_ALL_INSTRUCTION) -} - -pub fn wrap_combined_host_operation( - physical_tile: u16, - operation: &[u32], - xreq_packet_address: u32, -) -> Result, ExchangeError> { - let Some((&RETURN_M10_INSTRUCTION, body)) = operation.split_last() else { - return Err(ExchangeError::Schedule("local host operation return")); - }; - if body.len() < 2 || xreq_packet_address & 7 != 0 { - return Err(ExchangeError::Schedule("local host operation prefix")); - } - let mut instructions = vec![ - setzi_m(8, host_mux_for_tile(physical_tile)?), - put_special_from_m8(INCOMING_MUX_REGISTER), - SYNC_HOST_INSTRUCTION, - ]; - instructions.extend_from_slice(&body[..2]); - instructions.push(encode_send(1, 3, xreq_packet_address >> 2)?); - instructions.extend_from_slice(&body[2..]); - instructions.extend([ - SYNC_ALL_INSTRUCTION, - setzi_m(8, TILE_MUX_EXCHANGE), - put_special_from_m8(INCOMING_MUX_REGISTER), - RETURN_M10_INSTRUCTION, - ]); - Ok(instructions) -} - -fn plan_host_transfer( - mut tile_address: u32, - mut host_offset: u32, - mut bytes: u32, - long_max_bytes: u32, - packet: impl Fn(u32, u32, u32) -> Result, -) -> Result, ExchangeError> { - let mut chunks = Vec::new(); - while bytes != 0 { - let page_bytes = HOST_PAGE_BYTES - host_offset % HOST_PAGE_BYTES; - let available = bytes.min(page_bytes); - let count = if host_offset & 63 == 0 && available >= 64 { - available.min(long_max_bytes) & !63 - } else if available <= HOST_SHORT_MAX_BYTES { - available - } else { - // Keeping intermediate short packets at 32 bytes also preserves - // the destination alignment required by host-to-tile requests. - 32 - }; - let header = packet(tile_address, host_offset, count)?; - chunks.push(HostTransferChunk { - tile_address, - host_offset, - bytes: count, - header, - }); - tile_address = tile_address - .checked_add(count) - .ok_or(ExchangeError::HostPacket)?; - host_offset = host_offset - .checked_add(count) - .ok_or(ExchangeError::HostPacket)?; - bytes -= count; - } - Ok(chunks) -} - -fn validate_host_tile(physical_tile: u16) -> Result<(), ExchangeError> { - if physical_tile > 0xfff { - return Err(ExchangeError::HostPacket); - } - Ok(()) -} - -pub fn host_mux_for_tile(physical_tile: u16) -> Result { - validate_host_tile(physical_tile)?; - Ok(TILE_MUX_HOST + u32::from((physical_tile & 0x3f) & !2)) -} - -fn host_packet_size(host_offset: u32, bytes: u32) -> Result { - if (4..=HOST_SHORT_MAX_BYTES).contains(&bytes) && host_offset & 3 == 0 && bytes & 3 == 0 { - return Ok(HostPacketSize::Short); - } - if (64..=HOST_LONG_MAX_BYTES).contains(&bytes) && host_offset & 63 == 0 && bytes & 63 == 0 { - return Ok(HostPacketSize::Long); - } - Err(ExchangeError::HostPacket) -} - -fn host_address_length( - host_offset: u32, - bytes: u32, - size: HostPacketSize, -) -> Result { - let shift = match size { - HostPacketSize::Short => 2, - HostPacketSize::Long => 6, - }; - let units = bytes >> shift; - let length = if size == HostPacketSize::Long && bytes == HOST_LONG_MAX_BYTES { - 0 - } else { - units - }; - (u64::from(host_offset >> shift) << 4 | u64::from(length)) - .try_into() - .ok() - .filter(|encoded: &u32| *encoded <= 0x7fff_ffff) - .ok_or(ExchangeError::HostPacket) -} - -fn host_route_word0(physical_tile: u16) -> u32 { - let tile = u32::from(physical_tile); - ((tile >> 1) << 16) | ((tile & 1) << 15) -} - -fn host_route_word1(physical_tile: u16) -> u32 { - u32::from(physical_tile & 1) << 31 -} - -impl Topology { - pub fn new(logical_to_physical: Vec) -> Result { - let mut physical = HashSet::new(); - if logical_to_physical.is_empty() - || logical_to_physical - .iter() - .any(|tile| !physical.insert(*tile)) - { - return Err(ExchangeError::ReceiverSet); - } - Ok(Self { - logical_to_physical, - }) - } - - pub fn c600() -> Self { - Self { - logical_to_physical: (0..1472).map(c600_logical_to_physical).collect(), - } - } - - pub fn tile_count(&self) -> usize { - self.logical_to_physical.len() - } - - pub fn physical(&self, logical: u16) -> Result { - self.logical_to_physical - .get(usize::from(logical)) - .copied() - .ok_or(ExchangeError::Tile(logical)) - } - - /// Logical tile that shares this tile's double-width exchange resources. - pub fn paired_logical(&self, logical: u16) -> Result { - let paired_physical = self.physical(logical)? ^ 2; - self.logical_to_physical - .iter() - .position(|physical| *physical == paired_physical) - .map(|paired| u16::try_from(paired).expect("logical tile count fits u16")) - .ok_or(ExchangeError::Tile(logical)) - } - - /// Physical source selected by `INCOMING_MUXPAIR` for a 64-bit send. - pub fn paired_source_mux(&self, sender_logical: u16) -> Result { - Ok(self.physical(sender_logical)? ^ 2) - } - - pub fn is_pair_primary(&self, logical: u16) -> Result { - Ok(self.physical(logical)? & 2 == 0) - } - - /// Direction and width control for one 64-bit route. Both borrowed lanes - /// follow the same fabric direction; the two members of a receiving pair - /// are not an ordinary two-direction multicast. - pub fn paired_send_control( - &self, - sender_logical: u16, - receiver_logical: u16, - ) -> Result { - let sender = u32::from(self.physical(sender_logical)?); - let receiver = u32::from(self.physical(receiver_logical)?); - Ok(u8::try_from(direction(sender, receiver) | 4).expect("send control is three bits")) - } - - /// Whether this member of a double-width receiving pair owns the paired - /// XPIC source-selection stream. This matches the SDK architecture - /// helper `TPair_RxIsEarly`. - pub fn paired_receiver_is_early( - &self, - receiver_logical: u16, - sender_logical: u16, - ) -> Result { - let receiver = u32::from(self.physical(receiver_logical)?); - let sender = u32::from(self.physical(sender_logical)?); - let local = time_to_mux(sender, receiver); - let borrowed = time_to_mux(sender, receiver ^ 2); - Ok(local < borrowed) - } - - /// Builds an SDK-compatible double-width multicast. Receivers must be - /// complete physical tile pairs; both members consume the same 64-bit - /// item stream. - pub fn paired_multicast( - &self, - sender_logical: u16, - receivers: &[u16], - count: u32, - ) -> Result { - validate_count(count)?; - if count < 64 { - return Err(ExchangeError::Count(count)); - } - if receivers.is_empty() || receivers.len() & 1 != 0 { - return Err(ExchangeError::ReceiverSet); - } - let receiver_set = receivers.iter().copied().collect::>(); - if receiver_set.len() != receivers.len() - || receiver_set.contains(&sender_logical) - || receiver_set.contains(&self.paired_logical(sender_logical)?) - || receivers.iter().any(|&receiver| { - self.paired_logical(receiver) - .map_or(true, |paired| !receiver_set.contains(&paired)) - }) - { - return Err(ExchangeError::ReceiverSet); - } - - let mut plan = self.multicast(sender_logical, receivers, count, 0)?; - let sender_physical = u32::from(self.physical(sender_logical)?); - let send_control = if receivers.len() == 2 { - u8::try_from(direction(sender_physical, u32::from(self.physical(receivers[0])?)) | 4) - .expect("send control is three bits") - } else { - 7 - }; - set_sender_control(&mut plan.sender, send_control)?; - let receiver_physical = receivers - .iter() - .map(|&receiver| self.physical(receiver).map(u32::from)) - .collect::, _>>()?; - let minimum_double_mux = receiver_physical - .iter() - .map(|&receiver| paired_time_to_mux(sender_physical, receiver)) - .min() - .expect("paired multicast has receivers"); - plan.receivers = receivers - .iter() - .map(|&receiver| { - self.paired_receiver_row(sender_logical, receiver, count, minimum_double_mux) - }) - .collect::, _>>()?; - Ok(plan) - } - - fn paired_receiver_row( - &self, - sender_logical: u16, - receiver_logical: u16, - count: u32, - minimum_double_mux: i32, - ) -> Result { - let sender = u32::from(self.physical(sender_logical)?); - let receiver = u32::from(self.physical(receiver_logical)?); - // Route times are already expressed in the exchange epoch's event - // clock when every receiver is on the positive side of the timing - // origin. Shift the whole multicast only when its earliest pair would - // otherwise configure the mux before event one. - let epoch_shift = (1 - minimum_double_mux).max(0); - let source_event = u32::try_from(paired_time_to_mux(sender, receiver) + epoch_shift) - .map_err(|_| ExchangeError::Schedule("paired source timing"))?; - // MXP is 59 plus twice the physical row. The SDK selects paired mode - // MXP-8 cycles after this pair's route-specific XPIC source event. - // Adjacent logical pairs can therefore have distinct format windows - // when the logical-to-physical mapping turns into another column. - let format_start = source_event + 51 + 2 * (receiver >> 6); - let mut events = Vec::with_capacity(5); - if self.paired_receiver_is_early(receiver_logical, sender_logical)? { - events.push(ReceiveEvent { - cycles: source_event, - instruction: delay_xpic(0, 1, sender ^ (receiver & 2)), - kind: ReceiveEventKind::PairedSource, - }); - events.push(ReceiveEvent { - cycles: source_event + count, - instruction: delay_xpic(0, 1, TILE_MUX_EXCHANGE), - kind: ReceiveEventKind::PairedNeutral, - }); - } - events.extend([ - ReceiveEvent { - cycles: format_start, - instruction: delay_pic(0, 1, direction(sender, receiver)), - kind: ReceiveEventKind::Format, - }, - ReceiveEvent { - cycles: format_start + 2, - instruction: delay_pic(0, 0, 0), - kind: ReceiveEventKind::PairedPointer, - }, - ReceiveEvent { - cycles: format_start + count, - instruction: delay_pic(0, 1, 0), - kind: ReceiveEventKind::Format, - }, - ]); - events.sort_by_key(|event| event.cycles); - // Directionless SENDPICP can apply ordinary PIC and XPIC controls on - // one event, but the paired receive path faults if format activation - // coincides with its source selection or teardown. Reject that row so - // placement-aware lowering can retain a Word32 transfer instead. - if events - .windows(2) - .any(|pair| pair[0].cycles == pair[1].cycles) - { - return Err(ExchangeError::Schedule( - "paired receive coincident controls", - )); - } - validate_receive_events(&events)?; - - let horizon = events - .last() - .expect("paired receive always has format events") - .cycles - .checked_add(7) - .ok_or(ExchangeError::Schedule("paired receive horizon overflow"))?; - let mut words = vec![SYNC_SUPERVISOR_INSTRUCTION]; - let mut cycles = 0; - append_receive_events(&mut words, &mut cycles, &events, horizon, false)?; - words.push(RETURN_M10_INSTRUCTION); - if words.len() > PLAN_WORDS { - return Err(ExchangeError::Schedule("paired receive row capacity")); - } - let mut row = [0; PLAN_WORDS]; - row[..words.len()].copy_from_slice(&words); - Ok(row) - } - - pub fn point_to_point( - &self, - sender_logical: u16, - receiver_logical: u16, - count: u32, - ) -> Result { - validate_count(count)?; - if sender_logical == receiver_logical { - return Err(ExchangeError::DuplicateTile); - } - let sender = u32::from(self.physical(sender_logical)?); - let receiver = u32::from(self.physical(receiver_logical)?); - let direction = direction(sender, receiver); - let mux_time = time_to_mux(sender, receiver); - let receiver_phase = 2 * (receiver >> 6); - let sender_delay = 111 - mux_time; - if !(-1..=0x7ffff).contains(&sender_delay) { - return Err(ExchangeError::Schedule("sender delay")); - } - - let mut sender_row = [0; PLAN_WORDS]; - sender_row[0] = SYNC_SUPERVISOR_INSTRUCTION; - let mut cursor = 1; - if sender_delay >= 0 { - sender_row[cursor] = delay(sender_delay as u32); - cursor += 1; - } - let first_packet = count.min(64); - sender_row[cursor] = encode_send(first_packet - 1, direction, 0)?; - cursor += 1; - if count > 64 { - sender_row[cursor] = send_off(count - 65, direction, 0); - cursor += 1; - } - let trailing_delay = 4 - sender_delay - count as i32; - if trailing_delay >= 0 { - sender_row[cursor] = delay(trailing_delay as u32); - cursor += 1; - } - sender_row[cursor] = RETURN_M10_INSTRUCTION; - - let mut receiver_row = [0; PLAN_WORDS]; - receiver_row[0] = 1; - receiver_row[1] = SYNC_SUPERVISOR_INSTRUCTION; - receiver_row[2] = delay_xpic(112, 0, 0); - if count <= 51 { - receiver_row[3] = delay_xpic(count - 1, 0, TILE_MUX_EXCHANGE); - receiver_row[4] = delay_pic(51 - count + receiver_phase, 0, 0); - receiver_row[5] = delay(count + 4); - receiver_row[6] = RETURN_M10_INSTRUCTION; - } else if count == 52 { - // Keep PIC at the first payload arrival, as in the other rows. - // Programming it early works alone but understates the write - // window when this primitive is composed with another receive. - receiver_row[3] = delay_pic(51 + receiver_phase, 0, 0); - receiver_row[4] = delay_xpic(0, 0, TILE_MUX_EXCHANGE); - receiver_row[5] = delay(56); - receiver_row[6] = RETURN_M10_INSTRUCTION; - } else { - receiver_row[3] = delay_pic(51 + receiver_phase, 0, 0); - receiver_row[4] = delay_xpic(count - 53, 0, TILE_MUX_EXCHANGE); - receiver_row[5] = delay(56); - receiver_row[6] = RETURN_M10_INSTRUCTION; - } - debug!( - sender_logical, - receiver_logical, count, "assembled point-to-point exchange" - ); - Ok(Plan { - sender: sender_row, - receiver: receiver_row, - }) - } - - pub fn multicast( - &self, - sender_logical: u16, - receiver_logical: &[u16], - count: u32, - schedule_offset: u32, - ) -> Result { - validate_count(count)?; - let source_physical = u32::from(self.physical(sender_logical)?); - let mut used = HashSet::from([sender_logical]); - if receiver_logical.is_empty() - || receiver_logical - .iter() - .any(|receiver| !used.insert(*receiver) || self.physical(*receiver).is_err()) - { - return Err(ExchangeError::ReceiverSet); - } - let mux_times: Vec<_> = receiver_logical - .iter() - .map(|receiver| { - self.physical(*receiver) - .map(|physical| time_to_mux(source_physical, u32::from(physical))) - }) - .collect::>()?; - let minimum_mux = *mux_times.iter().min().ok_or(ExchangeError::ReceiverSet)?; - let natural_start = (-minimum_mux).max(0) as u32; - let start_cycle = natural_start - .checked_add(schedule_offset) - .filter(|cycle| *cycle <= 4095) - .ok_or(ExchangeError::Schedule("multicast start cycle"))?; - let sender_delay = start_cycle as i32 - 1; - - let mut sender = [0; PLAN_WORDS]; - let mut cursor = 0; - sender[cursor] = SYNC_SUPERVISOR_INSTRUCTION; - cursor += 1; - if sender_delay >= 0 { - sender[cursor] = delay(sender_delay as u32); - cursor += 1; - } - let send_direction = if receiver_logical.len() == 1 { - direction( - source_physical, - u32::from(self.physical(receiver_logical[0])?), - ) - } else { - 3 - }; - sender[cursor] = encode_send(count.min(64) - 1, send_direction, 0)?; - cursor += 1; - if count > 64 { - sender[cursor] = send_off(count - 65, send_direction, 0); - cursor += 1; - } - let trailing_delay = 4 - sender_delay - count as i32; - if trailing_delay >= 0 { - sender[cursor] = delay(trailing_delay as u32); - cursor += 1; - } - sender[cursor] = RETURN_M10_INSTRUCTION; - - let mut receivers = Vec::with_capacity(receiver_logical.len()); - for (logical, mux_time) in receiver_logical.iter().zip(mux_times) { - let physical = u32::from(self.physical(*logical)?); - let receive_cycle = start_cycle as i32 + mux_time; - if !(0..=4095).contains(&receive_cycle) { - return Err(ExchangeError::Schedule("multicast receive cycle")); - } - let receiver_phase = 2 * (physical >> 6); - let mut row = [0; PLAN_WORDS]; - row[0] = SYNC_SUPERVISOR_INSTRUCTION; - row[1] = delay_xpic(receive_cycle as u32, 0, source_physical); - if count <= 51 { - // The one-word case still needs this event: without it the - // tile remains connected to the source after the phase ends. - row[2] = delay_xpic(count - 1, 0, TILE_MUX_EXCHANGE); - row[3] = delay_pic(51 - count + receiver_phase, 0, 0) | 0x0001_4000; - row[4] = delay(count + 4); - row[5] = RETURN_M10_INSTRUCTION; - } else if count == 52 { - row[2] = delay_pic(51 + receiver_phase, 0, 0) | 0x0001_4000; - row[3] = delay_xpic(0, 0, TILE_MUX_EXCHANGE); - row[4] = delay(56); - row[5] = RETURN_M10_INSTRUCTION; - } else { - row[2] = delay_pic(51 + receiver_phase, 0, 0) | 0x0001_4000; - row[3] = delay_xpic(count - 53, 0, TILE_MUX_EXCHANGE); - row[4] = delay(56); - row[5] = RETURN_M10_INSTRUCTION; - } - receivers.push(row); - } - debug!( - sender_logical, - receiver_logical = ?receiver_logical, - count, - schedule_offset, - "assembled multicast exchange" - ); - Ok(MulticastPlan { sender, receivers }) - } -} - -/// Selects double-width items in every outgoing instruction in a primitive -/// sender row. Counts and absolute source operands both become 64-bit-item -/// units; use [`patch_sender_address_64`] after selecting this control. -pub fn set_sender_control(row: &mut PlanRow, send_control: u8) -> Result<(), ExchangeError> { - if !(1..=7).contains(&send_control) { - return Err(ExchangeError::Schedule("send control")); - } - let mut found = false; - for instruction in row { - if *instruction & LONG_OPCODE_MASK == SEND_OPCODE - || is_send_off(*instruction) - || (is_send_control_pair(*instruction) && *instruction & 3 != 0) - { - *instruction = (*instruction & !7) | u32::from(send_control); - found = true; - } - } - found - .then_some(()) - .ok_or(ExchangeError::Schedule("sender payload")) -} - -/// Removes standalone source-mux writes while retaining their exact event -/// advances. The secondary member of a paired 64-bit receiver uses this when -/// the primary member owns both `INCOMING_MUX` selections. -pub fn replace_xpic_controls_with_delays(program: &mut [u32]) -> Result<(), ExchangeError> { - for instruction in program { - if *instruction & OPCODE_MASK == DELAY_XPIC_OPCODE { - *instruction = delay(instruction_advance(*instruction) - 1); - } else if is_send_control(*instruction) || is_send_control_pair(*instruction) { - return Err(ExchangeError::Schedule( - "paired receiver has fused XPIC control", - )); - } - } - Ok(()) -} - -/// Selects the paired XPIC stream for every standalone receive-mux write. -/// Double-width receivers use stream one on the primary tile; stream zero is -/// supplied by the paired incoming-mux register. -pub fn select_paired_xpic_stream(program: &mut [u32]) -> Result<(), ExchangeError> { - let mut found = false; - for instruction in program { - if *instruction & OPCODE_MASK == DELAY_XPIC_OPCODE { - *instruction |= 1 << 13; - found = true; - } else if is_send_control(*instruction) || is_send_control_pair(*instruction) { - return Err(ExchangeError::Schedule( - "paired receiver has fused XPIC control", - )); - } - } - found - .then_some(()) - .ok_or(ExchangeError::Schedule("paired receiver XPIC control")) -} - -/// Replaces the selected physical source in standalone XPIC controls. -pub fn patch_xpic_source(program: &mut [u32], source_physical: u16) -> Result<(), ExchangeError> { - if u32::from(source_physical) > 0x1fff { - return Err(ExchangeError::Tile(source_physical)); - } - let instruction = program - .iter_mut() - .find(|instruction| { - **instruction & OPCODE_MASK == DELAY_XPIC_OPCODE - && **instruction & 0x1fff != TILE_MUX_EXCHANGE - }) - .ok_or(ExchangeError::Schedule("receiver source control"))?; - *instruction = (*instruction & !0x1fff) | u32::from(source_physical); - Ok(()) -} - -/// Executable exchange row that reserves a borrowed tile resource without -/// sending or receiving payload data. -pub fn timed_idle_program(cycles: u32) -> Result, ExchangeError> { - let mut words = Vec::new(); - let mut event_cycles = 0; - append_plain_delay(&mut words, &mut event_cycles, cycles)?; - words.push(RETURN_M10_INSTRUCTION); - Ok(words) -} - -pub fn c600_logical_to_physical(logical: u16) -> u16 { - let pair = logical / 2; - let lane = logical & 1; - let block = pair / 23; - let mut row = pair % 23; - if block & 1 != 0 { - row = 22 - row; - } - let column = (block / 2) * 4 + (block & 1); - row * 64 + column + lane * 2 -} - -pub fn patch_sender_address(row: &mut PlanRow, byte_address: u32) -> Result<(), ExchangeError> { - let instruction = row - .iter_mut() - .find(|instruction| **instruction & LONG_OPCODE_MASK == SEND_OPCODE) - .ok_or(ExchangeError::Address(byte_address))?; - patch_sender_instruction(instruction, byte_address) -} - -/// Replaces the address field of one tile-to-tile SEND instruction. -pub fn patch_sender_instruction( - instruction: &mut u32, - byte_address: u32, -) -> Result<(), ExchangeError> { - if *instruction & LONG_OPCODE_MASK != SEND_OPCODE - && !(is_send_control_pair(*instruction) && *instruction & 7 != 0) - { - return Err(ExchangeError::Address(byte_address)); - } - let item_shift = if *instruction & 4 != 0 { 3 } else { 2 }; - if byte_address & ((1 << item_shift) - 1) != 0 - || byte_address >> item_shift > SEND_ADDRESS_MASK >> 3 - { - return Err(ExchangeError::Address(byte_address)); - } - let item_address = byte_address >> item_shift; - *instruction = (*instruction & !SEND_ADDRESS_MASK) | ((item_address << 3) & SEND_ADDRESS_MASK); - Ok(()) -} - -/// Address-bearing instructions for each outgoing message, in execution -/// order. Each entry is `(word offset, byte offset from the message source)`. -/// SENDPICP restarts the outgoing source stream explicitly after its inline -/// control word, so repeat relocation must patch it as well as the first SEND. -pub fn sender_address_instruction_groups( - row: &[u32], -) -> Result>, ExchangeError> { - let mut groups = Vec::>::new(); - let mut sent_words = None; - let mut cursor = 0; - while cursor < row.len() { - let instruction = row[cursor]; - if instruction & LONG_OPCODE_MASK == SEND_OPCODE { - groups.push(vec![(cursor, 0)]); - sent_words = Some(instruction_advance(instruction)); - } else if is_send_control_pair(instruction) && instruction & 7 != 0 { - let sent = sent_words.ok_or(ExchangeError::Schedule( - "SENDPICP precedes initial outgoing SEND", - ))?; - groups - .last_mut() - .ok_or(ExchangeError::Schedule("SENDPICP outgoing group"))? - .push(( - cursor, - sent.checked_mul(if instruction & 4 != 0 { 8 } else { 4 }) - .ok_or(ExchangeError::Schedule("sender byte offset overflow"))?, - )); - sent_words = Some( - sent.checked_add(instruction_advance(instruction)) - .ok_or(ExchangeError::Schedule("sender word offset overflow"))?, - ); - } else if (is_send_off(instruction) || is_send_control(instruction)) && sent_words.is_some() - { - sent_words = Some( - sent_words - .unwrap() - .checked_add(instruction_advance(instruction)) - .ok_or(ExchangeError::Schedule("sender word offset overflow"))?, - ); - } - cursor += if is_send_control_pair(instruction) { - 2 - } else { - 1 - }; - } - Ok(groups) -} - -/// Removes tile-memory address fields while retaining exchange roles, routes, -/// transfer sizes, and event timing. Rows with the same result can share one -/// executable slot and restore their addresses before invocation. -pub fn normalized_exchange_address_words(row: &[u32]) -> Vec { - let mut normalized = row.to_vec(); - let mut cursor = 0; - while cursor < normalized.len() { - let instruction = normalized[cursor]; - if is_send_control_pair(instruction) { - if instruction & 7 != 0 { - normalized[cursor] &= !SEND_ADDRESS_MASK; - } - if instruction & (1 << 27) == 0 - && let Some(payload) = normalized.get_mut(cursor + 1) - { - *payload &= !PIC_RECEIVE_ADDRESS_MASK; - } - cursor += 2; - continue; - } - normalized[cursor] = if instruction & LONG_OPCODE_MASK == SEND_OPCODE { - instruction & !SEND_ADDRESS_MASK - } else if (is_send_control(instruction) && (instruction >> 18) & 3 == 2) - || (instruction & OPCODE_MASK == DELAY_PIC_OPCODE && instruction & (1 << 18) == 0) - { - instruction & !PIC_RECEIVE_ADDRESS_MASK - } else { - instruction - }; - cursor += 1; - } - normalized -} - -pub fn patch_receiver_address(row: &mut PlanRow, byte_address: u32) -> Result<(), ExchangeError> { - if byte_address & 3 != 0 || byte_address >> 2 > PIC_RECEIVE_ADDRESS_MASK { - return Err(ExchangeError::Address(byte_address)); - } - let word_address = byte_address >> 2; - let mut cursor = 0; - while cursor < row.len() { - let instruction = row[cursor]; - if is_send_control_pair(instruction) { - if instruction & (1 << 27) == 0 { - let payload = row - .get_mut(cursor + 1) - .ok_or(ExchangeError::Schedule("truncated SENDPICP payload"))?; - *payload = (*payload & !PIC_RECEIVE_ADDRESS_MASK) | word_address; - return Ok(()); - } - cursor += 2; - continue; - } - if instruction & OPCODE_MASK == DELAY_PIC_OPCODE && instruction & (1 << 18) == 0 { - row[cursor] = (instruction & !PIC_RECEIVE_ADDRESS_MASK) | word_address; - return Ok(()); - } - cursor += 1; - } - Err(ExchangeError::Address(byte_address)) -} - -/// Delays every timed event in a plan row while preserving route-relative timing. -pub fn offset_plan(row: &mut PlanRow, cycles: u32) -> Result<(), ExchangeError> { - if cycles == 0 { - return Ok(()); - } - if row[0] != SYNC_SUPERVISOR_INSTRUCTION { - return Err(ExchangeError::Schedule("plan offset entry")); - } - let end = row - .iter() - .position(|instruction| *instruction == RETURN_M10_INSTRUCTION) - .ok_or(ExchangeError::Schedule("plan offset return"))?; - let available = row.len().saturating_sub(end + 1); - let delay_count = usize::try_from(cycles.div_ceil(MAX_PLAN_OFFSET_CYCLES)) - .map_err(|_| ExchangeError::Schedule("plan offset instruction count"))?; - if delay_count > available { - let maximum = u32::try_from(available) - .unwrap_or(u32::MAX) - .saturating_mul(MAX_PLAN_OFFSET_CYCLES); - return Err(ExchangeError::PlanOffsetRange { cycles, maximum }); - } - if delay_count == 0 { - return Err(ExchangeError::Schedule("plan offset instruction capacity")); - } - row.copy_within(1..=end, 1 + delay_count); - let mut remaining = cycles; - for instruction in &mut row[1..=delay_count] { - let chunk = remaining.min(MAX_PLAN_OFFSET_CYCLES); - *instruction = delay(chunk - 1); - remaining -= chunk; - } - Ok(()) -} - -pub fn finalize_point_receiver( - row: &PlanRow, - source_physical: u16, -) -> Result { - let patch_index = row[0] as usize; - if patch_index >= PLAN_WORDS - 1 || u32::from(source_physical) > 0x1fff { - return Err(ExchangeError::Schedule("point receiver patch index")); - } - let mut executable = [0; PLAN_WORDS]; - executable[..PLAN_WORDS - 1].copy_from_slice(&row[1..]); - executable[patch_index] = (executable[patch_index] & !0x1fff) | u32::from(source_physical); - Ok(executable) -} - -fn validate_count(count: u32) -> Result<(), ExchangeError> { - if (1..=MAX_TRANSFER_WORDS).contains(&count) { - Ok(()) - } else { - Err(ExchangeError::Count(count)) - } -} - -fn route_displacement(source: u32, destination: u32) -> i32 { - let source_raw = ((source >> 2) & 15) as i32; - let destination_raw = ((destination >> 2) & 15) as i32; - let source_column = if source_raw > 7 { - source_raw ^ 15 - } else { - source_raw - }; - let destination_column = if destination_raw > 7 { - destination_raw ^ 15 - } else { - destination_raw - }; - let source_mux = source_column + ((source_raw >> 3) ^ (source & 1) as i32); - let base = (destination_column - source_mux) * 6; - let destination_lane = destination & 3; - let destination_half = destination_raw >> 3; - if destination_lane > 1 { - base + if destination_half == (destination & 1) as i32 { - 2 - } else { - 4 - } - } else { - base + if destination_half == destination_lane as i32 { - 1 - } else { - 5 - } - } -} - -fn direction(source: u32, destination: u32) -> u32 { - if route_displacement(source, destination) < 1 { - 2 - } else { - 1 - } -} - -fn time_to_mux(source: u32, destination: u32) -> i32 { - let source_raw = ((source >> 2) & 15) as i32; - let destination_raw = ((destination >> 2) & 15) as i32; - let source_low = ((source >> 2) & 7) as i32; - let displacement = route_displacement(source, destination); - let source_edge = if source_raw > 7 { - (source_raw * 4) ^ 60 - } else { - source_raw * 4 - }; - let destination_edge = if destination_raw > 7 { - (destination_raw * 4) ^ 60 - } else { - destination_raw * 4 - }; - let local = ((source >> 2) & 8) as i32 | ((source >> 3) & 3) as i32; - let crossing = local - destination_raw + ((source_low >> 1) ^ 3); - let same_region = (source ^ destination) & 0x20 == 0; - let turn = if same_region { - source_low + 1 - } else { - 16 - source_low - }; - let group_delta = (((source >> 6) & 31) as i32 - ((destination >> 6) & 31) as i32) * 2; - crossing + source_edge + turn - destination_edge + group_delta + displacement.abs() - 34 -} - -/// Route timing when a receiving tile borrows its physical partner's lane. -/// It is monotonic in the opposite direction to the SDK's -/// `ColossusBNET_TTDBL`, but differences between receiver pairs are in event -/// cycles and drive the same paired-source schedule. -fn paired_time_to_mux(source: u32, destination: u32) -> i32 { - time_to_mux(source, destination).max(time_to_mux(source, destination ^ 2)) -} - -pub const fn encode_exchange_delay(cycles: u32) -> u32 { - 0x40a0_0000 | (cycles & 0x7ffff) -} - -pub const fn encode_exchange_delay_pic(a: u32, b: u32, c: u32) -> u32 { - 0x6000_0000 | ((a << 19) & 0x03f8_0000) | ((b << 18) & 0x0004_0000) | (c & 0x3ffff) -} - -pub const fn encode_exchange_delay_xpic(a: u32, b: u32, c: u32) -> u32 { - 0x6400_0000 | ((a << 14) & 0x03ff_c000) | ((b << 13) & 0x0000_2000) | (c & 0x1fff) -} - -const fn delay(cycles: u32) -> u32 { - encode_exchange_delay(cycles) -} - -const fn delay_pic(a: u32, b: u32, c: u32) -> u32 { - encode_exchange_delay_pic(a, b, c) -} - -const fn delay_xpic(a: u32, b: u32, c: u32) -> u32 { - encode_exchange_delay_xpic(a, b, c) -} - -pub fn encode_send( - count_minus_one: u32, - direction: u32, - base_word: u32, -) -> Result { - if count_minus_one > 63 || direction > 7 || base_word > 0x3_ffff { - return Err(ExchangeError::Schedule("send instruction operand")); - } - Ok(0x7800_0000 - | ((count_minus_one << 21) & 0x07e0_0000) - | ((base_word << 3) & 0x001f_fff8) - | direction) -} - -fn send_off(count_minus_one: u32, direction: u32, base_word: u32) -> u32 { - 0x7000_0000 - | ((count_minus_one << 21) & 0x07e0_0000) - | (((count_minus_one >> 6) << 14) & 0x000f_c000) - | ((base_word << 3) & 0x0000_3ff8) - | (direction & 7) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn scalar_instruction_encoders_preserve_operands_and_reject_overflow() { - let setzi = encode_setzi_m(15, (1 << 20) - 1).unwrap(); - assert_eq!((setzi >> 20) & 0xf, 15); - assert_eq!(setzi & ((1 << 20) - 1), (1 << 20) - 1); - - let put = encode_put_special_m(0xa6, 8).unwrap(); - assert_eq!((put >> 20) & 0xf, 8); - assert_eq!(put & 0xff, 0xa6); - assert_eq!((encode_br_m(10).unwrap() >> 20) & 0xf, 10); - - let call = encode_call_m_immediate(10, 0x4c100).unwrap(); - assert_eq!((call >> 20) & 0xf, 10); - assert_eq!((call & 0x7ffff) << 2, 0x4c100); - - assert!(encode_setzi_m(16, 0).is_err()); - assert!(encode_setzi_m(0, 1 << 20).is_err()); - assert!(encode_put_special_m(0, 16).is_err()); - assert!(encode_br_m(16).is_err()); - assert!(encode_call_m_immediate(16, 0).is_err()); - assert!(encode_call_m_immediate(0, 2).is_err()); - assert!(encode_call_m_immediate(0, 1 << 21).is_err()); - - assert_eq!(encode_ld32_m_immediate(8, 11, 15, 1).unwrap(), 0x01b8_f001); - assert_eq!(encode_st32_m_immediate(2, 11, 15, 0).unwrap(), 0x4fb2_f000); - assert_eq!(encode_add_m_immediate(11, 11, -32).unwrap(), 0x22bb_ffe0); - assert_eq!(encode_and_m_immediate(0, 8, 1).unwrap(), 0x4280_0001); - assert_eq!(encode_shl_m_immediate(10, 7, 2).unwrap(), 0x427a_a002); - assert!(encode_shl_m_immediate(0, 0, 1 << 12).is_err()); - assert_eq!(encode_brz_m_immediate(0, 0x4c100).unwrap(), 0x1301_3040); - } - - #[test] - fn normalizes_all_sender_address_fields() { - let row = [ - SYNC_SUPERVISOR_INSTRUCTION, - encode_send(1, 3, 0x1a048).unwrap(), - SEND_PICP_OPCODE | (7 << 21) | (0x1a04a << 3) | 3, - 0x1901_5000, - RETURN_M10_INSTRUCTION, - ]; - let normalized = normalized_exchange_address_words(&row); - assert_eq!(normalized[1] ^ row[1], row[1] & SEND_ADDRESS_MASK); - assert_eq!(normalized[2] ^ row[2], row[2] & SEND_ADDRESS_MASK); - assert_eq!(normalized[0], row[0]); - assert_eq!(normalized[4], row[4]); - } - - #[test] - fn paired_send_matches_sdk_full_duplex_oracle() { - let events = [ - ReceiveEvent { - cycles: 189, - instruction: delay_xpic(0, 0, TILE_MUX_EXCHANGE), - kind: ReceiveEventKind::OrdinaryNeutral, - }, - ReceiveEvent { - cycles: 189, - instruction: delay_pic(0, 0, 0x15000), - kind: ReceiveEventKind::Pointer, - }, - ]; - assert_eq!( - encode_send_control_pair(42, 0x14021, 1, &events).unwrap(), - (0xf54a_0109, 0x1901_5000) - ); - } - - #[test] - fn paired_multicast64_matches_near_reverse_and_far_sdk_rows() { - let topology = Topology::c600(); - let cases = [(0, [2, 3], 2), (2, [0, 1], 0), (0, [46, 47], 47)]; - for (source, receivers, early) in cases { - let mut plan = topology.paired_multicast(source, &receivers, 64).unwrap(); - patch_sender_address(&mut plan.sender, 0x50000).unwrap(); - for row in &mut plan.receivers { - patch_receiver_address(row, 0x60000).unwrap(); - } - if source == 0 && receivers == [2, 3] { - assert_eq!( - plan.sender, - [ - SYNC_SUPERVISOR_INSTRUCTION, - delay(30), - encode_send(63, 5, 0x50000 >> 3).unwrap(), - RETURN_M10_INSTRUCTION, - 0, - 0, - 0, - 0, - 0, - ] - ); - assert_eq!( - plan.receivers[0], - [ - SYNC_SUPERVISOR_INSTRUCTION, - delay_xpic(0, 1, 0), - delay_pic(52, 1, 1), - delay_pic(1, 0, 0x60000 >> 2), - delay_xpic(8, 1, TILE_MUX_EXCHANGE), - delay_pic(52, 1, 0), - delay(6), - RETURN_M10_INSTRUCTION, - 0, - ] - ); - assert_eq!( - plan.receivers[1], - [ - SYNC_SUPERVISOR_INSTRUCTION, - delay_pic(53, 1, 1), - delay_pic(1, 0, 0x60000 >> 2), - delay_pic(61, 1, 0), - delay(6), - RETURN_M10_INSTRUCTION, - 0, - 0, - 0, - ] - ); - } - assert!(topology.paired_receiver_is_early(early, source).unwrap()); - assert!( - !topology - .paired_receiver_is_early(early ^ 1, source) - .unwrap() - ); - - let mut builder = PhaseProgramBuilder::new(1472); - let source_pair = topology.paired_logical(source).unwrap(); - let offset = builder - .earliest_transfer_offset(source, &[source_pair], &receivers, &plan, 64, 0) - .unwrap(); - builder - .append_transfer_at(source, &[source_pair], &receivers, &plan, offset, 64) - .unwrap(); - let programs = builder.finish().unwrap(); - assert!(programs.programs[usize::from(source)].is_some()); - assert!(programs.programs[usize::from(source_pair)].is_some()); - assert!( - receivers - .iter() - .all(|receiver| programs.programs[usize::from(*receiver)].is_some()) - ); - } - } - - #[test] - fn randomized_paired_multicasts_preserve_route_relative_timing() { - let topology = Topology::c600(); - let tile_count = topology.tile_count() as u16; - let mut random = fastrand::Rng::with_seed(0x7061_6972_6564_3634); - let all_pairs = (0..tile_count) - .filter_map(|tile| { - let paired = topology.paired_logical(tile).ok()?; - (tile < paired).then_some([tile, paired]) - }) - .collect::>(); - - for prefer_positive_routes in [false, true] { - for _ in 0..64 { - let source = random.u16(0..tile_count); - let source_pair = topology.paired_logical(source).unwrap(); - let source_physical = u32::from(topology.physical(source).unwrap()); - let mut candidates = all_pairs - .iter() - .copied() - .filter(|pair| !pair.contains(&source) && !pair.contains(&source_pair)) - .filter(|pair| { - let receiver = u32::from(topology.physical(pair[0]).unwrap()); - let route_time = paired_time_to_mux(source_physical, receiver); - (route_time > 1) == prefer_positive_routes - }) - .collect::>(); - random.shuffle(&mut candidates); - if candidates.is_empty() { - continue; - } - let pair_count = random.usize(1..=candidates.len().min(12)); - let receivers = candidates[..pair_count] - .iter() - .flatten() - .copied() - .collect::>(); - let count = 512; - let plan = topology - .paired_multicast(source, &receivers, count) - .unwrap(); - let minimum_route_time = receivers - .iter() - .map(|&receiver| { - paired_time_to_mux( - source_physical, - u32::from(topology.physical(receiver).unwrap()), - ) - }) - .min() - .unwrap(); - let epoch_shift = (1 - minimum_route_time).max(0); - - for (&receiver, row) in receivers.iter().zip(&plan.receivers) { - let receiver_physical = u32::from(topology.physical(receiver).unwrap()); - let source_event = u32::try_from( - paired_time_to_mux(source_physical, receiver_physical) + epoch_shift, - ) - .unwrap(); - let format_start = source_event + 51 + 2 * (receiver_physical >> 6); - let timing = receive_row_timing(row, 0).unwrap(); - assert_eq!(timing.mode, ReceiveMode::Paired64); - assert_eq!(timing.format_start_cycles, Some(format_start)); - assert_eq!(timing.pointer_cycles, Some(format_start + 2)); - assert_eq!(timing.format_end_cycles, Some(format_start + count)); - if topology.paired_receiver_is_early(receiver, source).unwrap() { - assert_eq!(timing.source_cycles, Some(source_event)); - assert_eq!(timing.neutral_cycles, Some(source_event + count)); - } else { - assert_eq!(timing.source_cycles, None); - assert_eq!(timing.neutral_cycles, None); - } - } - } - } - } - - #[test] - fn randomized_coincident_paired_controls_are_rejected() { - let topology = Topology::c600(); - let tile_count = topology.tile_count() as u16; - let mut random = fastrand::Rng::with_seed(0x636f_696e_6369_6465); - let mut checked = 0; - for _ in 0..10_000 { - let source = random.u16(0..tile_count); - let receiver = random.u16(0..tile_count); - let paired = topology.paired_logical(receiver).unwrap(); - let source_pair = topology.paired_logical(source).unwrap(); - if receiver > paired - || [source, source_pair] - .iter() - .any(|tile| *tile == receiver || *tile == paired) - { - continue; - } - let receivers = [receiver, paired]; - let source_physical = u32::from(topology.physical(source).unwrap()); - let receiver = receivers - .iter() - .copied() - .find(|receiver| { - topology - .paired_receiver_is_early(*receiver, source) - .unwrap() - }) - .unwrap(); - let receiver_physical = u32::from(topology.physical(receiver).unwrap()); - let source_event = u32::try_from( - paired_time_to_mux(source_physical, receiver_physical) - + (1 - paired_time_to_mux(source_physical, receiver_physical)).max(0), - ) - .unwrap(); - let format_start = source_event + 51 + 2 * (receiver_physical >> 6); - let count = format_start - source_event; - if count < 64 { - continue; - } - assert!(matches!( - topology.paired_multicast(source, &receivers, count), - Err(ExchangeError::Schedule( - "paired receive coincident controls" - )) - )); - checked += 1; - if checked == 64 { - break; - } - } - assert_eq!(checked, 64); - } - - #[test] - fn c600_mapping_is_a_permutation() { - let topology = Topology::c600(); - let physical: HashSet<_> = (0..topology.tile_count() as u16) - .map(|logical| topology.physical(logical).unwrap()) - .collect(); - assert_eq!(physical.len(), 1472); - assert_eq!(topology.physical(46).unwrap(), 1409); - assert!(physical.iter().all(|tile| *tile < 1472)); - } - - #[test] - fn point_to_point_primitive_encodings() { - let topology = Topology::c600(); - let cases = [ - ( - 0, - 1286, - 3, - [ - 0x41800003, 0x40a00086, 0x78400001, 0x43a00000, 0, 0, 0, 0, 0, - ], - [ - 1, 0x41800003, 0x641c0000, 0x64008640, 0x61800000, 0x40a00007, 0x43a00000, 0, 0, - ], - ), - ( - 736, - 100, - 52, - [ - 0x41800003, 0x40a0003e, 0x7e600002, 0x43a00000, 0, 0, 0, 0, 0, - ], - [ - 1, 0x41800003, 0x641c0000, 0x61d80000, 0x64000640, 0x40a00038, 0x43a00000, 0, 0, - ], - ), - ( - 1286, - 1471, - 65, - [ - 0x41800003, 0x40a0007d, 0x7fe00002, 0x70000002, 0x43a00000, 0, 0, 0, 0, - ], - [ - 1, 0x41800003, 0x641c0000, 0x61980000, 0x64030640, 0x40a00038, 0x43a00000, 0, 0, - ], - ), - ( - 32, - 0, - 4148, - [ - 0x41800003, 0x40a0006c, 0x7fe00001, 0x766fc001, 0x43a00000, 0, 0, 0, 0, - ], - [ - 1, 0x41800003, 0x641c0000, 0x61980000, 0x67ffc640, 0x40a00038, 0x43a00000, 0, 0, - ], - ), - ]; - for (sender, receiver, count, expected_sender, expected_receiver) in cases { - let plan = topology.point_to_point(sender, receiver, count).unwrap(); - assert_eq!(&plan.sender[..expected_sender.len()], &expected_sender); - assert_eq!( - &plan.receiver[..expected_receiver.len()], - &expected_receiver - ); - assert!( - plan.sender[expected_sender.len()..] - .iter() - .all(|word| *word == 0) - ); - assert!( - plan.receiver[expected_receiver.len()..] - .iter() - .all(|word| *word == 0) - ); - } - } - - #[test] - fn multicast_primitive_encodings() { - let topology = Topology::c600(); - let plan = topology.multicast(0, &[274, 1286], 3, 0).unwrap(); - assert_eq!(plan.sender[0], 0x41800003); - assert_eq!(plan.sender[1], 0x40a00016); - assert_eq!(plan.sender[2], 0x78400003); - assert_eq!(plan.receivers[0][1], 0x64000000); - assert_eq!(plan.receivers[0][2], 0x64008640); - assert_eq!(plan.receivers[0][3], 0x61814000); - - let boundary = topology.multicast(736, &[100, 900], 52, 0).unwrap(); - assert_eq!(boundary.receivers[0][2], 0x61d94000); - assert_eq!(boundary.receivers[0][3], 0x64000640); - } - - #[test] - fn single_receiver_uses_the_directional_route_for_every_payload_instruction() { - let topology = Topology::c600(); - let mut random = fastrand::Rng::with_seed(0x0072_6f75_7469_6e67); - for _ in 0..128 { - let source = random.u16(0..topology.tile_count() as u16); - let mut destination = random.u16(0..topology.tile_count() as u16); - while destination == source { - destination = random.u16(0..topology.tile_count() as u16); - } - let words = random.u32(65..=MAX_TRANSFER_WORDS); - let expected_direction = direction( - u32::from(topology.physical(source).unwrap()), - u32::from(topology.physical(destination).unwrap()), - ); - let plan = topology - .multicast(source, &[destination], words, 0) - .unwrap(); - let payload_directions = plan - .sender - .iter() - .filter(|instruction| { - matches!( - **instruction & LONG_OPCODE_MASK, - SEND_OPCODE | SEND_OFF_OPCODE - ) - }) - .map(|instruction| instruction & 7) - .collect::>(); - assert_eq!(payload_directions, [expected_direction, expected_direction]); - } - } - - #[test] - fn randomized_internal_receives_leave_the_neutral_mux_selected() { - let topology = Topology::c600(); - let mut random = fastrand::Rng::with_seed(0x006e_6575_7472_616c); - for _ in 0..128 { - let source = random.u16(0..topology.tile_count() as u16); - let receiver_count = random.usize(1..=4); - let mut receivers = Vec::with_capacity(receiver_count); - while receivers.len() < receiver_count { - let receiver = random.u16(0..topology.tile_count() as u16); - if receiver != source && !receivers.contains(&receiver) { - receivers.push(receiver); - } - } - let plan = topology - .multicast(source, &receivers, random.u32(1..=MAX_TRANSFER_WORDS), 0) - .unwrap(); - for receiver in plan.receivers { - assert!(receiver.iter().skip(2).any(|instruction| { - instruction & OPCODE_MASK == DELAY_XPIC_OPCODE - && instruction & 0x1fff == TILE_MUX_EXCHANGE - })); - } - } - } - - #[test] - fn ordinary_pointer_setup_matches_payload_arrival_at_count_boundaries() { - let topology = Topology::c600(); - for receiver in [2, 46, 100, 736, 1286] { - for count in [1, 51, 52, 53, 64, 65, 4148] { - let plan = topology.multicast(0, &[receiver], count, 0).unwrap(); - let timing = receive_row_timing(&plan.receivers[0], 0).unwrap(); - assert_eq!( - timing.pointer_cycles.unwrap() - timing.source_cycles.unwrap(), - 52 + 2 * u32::from(topology.physical(receiver).unwrap() >> 6), - "receiver={receiver} count={count}" - ); - } - } - } - - #[test] - fn paired_format_waits_for_the_previous_ordinary_payload_to_drain() { - let topology = Topology::c600(); - for receiver in [2, 3, 46, 47] { - let mut schedule = TileProgramSchedule::default(); - let ordinary = topology - .multicast(0, &[receiver], 972, 0) - .unwrap() - .receivers[0]; - let previous = schedule.append_receiver_at(&ordinary, 0, 972).unwrap(); - let pair = [receiver & !1, receiver | 1]; - let row = topology.paired_multicast(4, &pair, 176).unwrap().receivers - [usize::from(receiver & 1)]; - let offset = schedule.earliest_receiver_offset(&row, 176, 0).unwrap(); - let timing = receive_row_timing(&row, offset).unwrap(); - assert!(timing.format_start_cycles.unwrap() >= previous.payload_end); - } - } - - #[test] - fn paired_source_setup_does_not_merge_with_an_ordinary_pointer() { - let topology = Topology::c600(); - let mut schedule = TileProgramSchedule::default(); - let ordinary = topology.multicast(0, &[2], 53, 0).unwrap().receivers[0]; - schedule.append_receiver_at(&ordinary, 0, 53).unwrap(); - let paired = topology - .paired_multicast(4, &[2, 3], 176) - .unwrap() - .receivers[0]; - let offset = schedule.earliest_receiver_offset(&paired, 176, 0).unwrap(); - let timing = receive_row_timing(&paired, offset).unwrap(); - assert_eq!(timing.source_cycles, Some(56)); - schedule.append_receiver_at(&paired, offset, 176).unwrap(); - schedule.finish().unwrap(); - } - - #[test] - fn ordinary_source_setup_does_not_merge_with_a_paired_pointer() { - // MLP redistribution: paired activation multicast followed by ordinary - // weights. Fusing their controls in SENDPICP loses activation words on - // hardware even though XPIC itself selects an ordinary source. - let topology = Topology::c600(); - let mut schedule = TileProgramSchedule::default(); - let receivers = (1206..1224).collect::>(); - let paired = topology - .paired_multicast(966, &receivers, 72) - .unwrap() - .receivers[10]; - let previous = schedule.append_receiver_at(&paired, 0, 72).unwrap(); - let ordinary = topology - .multicast(612, &[244, 730, 1216], 4148, 0) - .unwrap() - .receivers[2]; - let offset = schedule - .earliest_receiver_offset(&ordinary, 4148, 0) - .unwrap(); - let timing = receive_row_timing(&ordinary, offset).unwrap(); - assert_ne!(timing.source_cycles, Some(previous.payload_start)); - schedule - .append_receiver_at(&ordinary, offset, 4148) - .unwrap(); - schedule.finish().unwrap(); - } - - #[test] - fn consecutive_paired_receives_match_the_sdks_continuous_format_stream() { - let topology = Topology::c600(); - let mut schedule = TileProgramSchedule::default(); - for (index, source) in [0, 4, 6, 8].into_iter().enumerate() { - let mut row = topology - .paired_multicast(source, &[2, 3], 176) - .unwrap() - .receivers[0]; - patch_receiver_address(&mut row, 0x90000 + index as u32 * 0x1000).unwrap(); - let offset = schedule.earliest_receiver_offset(&row, 176, 0).unwrap(); - schedule.append_receiver_at(&row, offset, 176).unwrap(); - } - let cycles = |kind| { - schedule - .receive_events - .iter() - .filter(|event| event.kind == kind) - .map(|event| event.cycles) - .collect::>() - }; - assert_eq!(cycles(ReceiveEventKind::PairedSource), [1, 177, 353, 529]); - assert_eq!(cycles(ReceiveEventKind::PairedPointer), [56, 232, 408, 584]); - assert_eq!(cycles(ReceiveEventKind::Format), [54, 758]); - assert_eq!(cycles(ReceiveEventKind::PairedNeutral), [705]); - schedule.finish().unwrap(); - } - - #[test] - fn consecutive_paired_receives_separate_format_teardown_from_next_source() { - let topology = Topology::c600(); - let mut schedule = TileProgramSchedule::default(); - for (index, source) in [0, 4].into_iter().enumerate() { - let mut row = topology - .paired_multicast(source, &[2, 3], 176) - .unwrap() - .receivers[0]; - patch_receiver_address(&mut row, 0x90000 + index as u32 * 0x1000).unwrap(); - let offset = schedule.earliest_receiver_offset(&row, 176, 0).unwrap(); - schedule.append_receiver_at(&row, offset, 176).unwrap(); - } - for format in schedule - .receive_events - .iter() - .filter(|event| event.kind == ReceiveEventKind::Format) - { - assert!( - !schedule - .receive_events - .iter() - .any(|event| event.kind.is_xpic() && event.cycles == format.cycles) - ); - } - schedule.finish().unwrap(); - } - - #[test] - fn randomized_receiver_streams_schedule_source_and_pointer_cutovers() { - let topology = Topology::c600(); - let mut random = fastrand::Rng::with_seed(0x6d75_785f_6375_746f); - for _ in 0..128 { - let receiver = random.u16(0..topology.tile_count() as u16); - let transfer_count = random.usize(2..=12); - let words = random.u32(53..=512); - let mut address = 0x50000 + random.u32(0..=0x2000) * 4; - let mut pointer_writes = 1; - let mut builder = TileProgramSchedule::default(); - for index in 0..transfer_count { - let source = loop { - let candidate = random.u16(0..topology.tile_count() as u16); - if candidate != receiver { - break candidate; - } - }; - let has_address_gap = index != 0 && random.bool(); - if has_address_gap { - address += random.u32(1..=16) * 4; - pointer_writes += 1; - } - let mut row = topology - .multicast(source, &[receiver], words, 0) - .unwrap() - .receivers[0]; - patch_receiver_address(&mut row, address).unwrap(); - let offset = builder.earliest_receiver_offset(&row, words, 0).unwrap(); - builder.append_receiver_at(&row, offset, words).unwrap(); - address += words * 4; - } - let expected_cycles = builder.event_cycles(); - let program = builder.finish().unwrap(); - assert_eq!(plan_event_cycles(&program).unwrap(), expected_cycles); - assert_eq!( - program - .iter() - .filter(|instruction| is_neutral_mux_teardown(**instruction)) - .count(), - 1 - ); - assert_eq!( - program - .iter() - .filter(|instruction| **instruction & OPCODE_MASK == DELAY_XPIC_OPCODE) - .count(), - transfer_count + 1 - ); - assert_eq!( - program - .iter() - .filter(|instruction| **instruction & OPCODE_MASK == DELAY_PIC_OPCODE) - .count(), - pointer_writes - ); - } - } - - #[test] - fn randomized_mixed_role_programs_fuse_receive_controls_into_sends() { - let topology = Topology::c600(); - let mut random = fastrand::Rng::with_seed(0x0073_656e_6470_6963); - let mut fused_programs = 0; - for _ in 0..256 { - let tile = random.u16(0..topology.tile_count() as u16); - let incoming_source = loop { - let candidate = random.u16(0..topology.tile_count() as u16); - if candidate != tile { - break candidate; - } - }; - let outgoing_destination = loop { - let candidate = random.u16(0..topology.tile_count() as u16); - if candidate != tile { - break candidate; - } - }; - let words = random.u32(2..=256); - let mut incoming = topology - .multicast(incoming_source, &[tile], words, 0) - .unwrap() - .receivers[0]; - patch_receiver_address(&mut incoming, 0x50000 + random.u32(0..0x1000) * 4).unwrap(); - let outgoing = topology - .multicast(tile, &[outgoing_destination], words, 0) - .unwrap() - .sender; - - let mut builder = TileProgramSchedule::default(); - builder.append_receiver_at(&incoming, 0, words).unwrap(); - let sender_offset = builder.earliest_sender_offset(&outgoing, 0).unwrap(); - builder.append_sender_at(&outgoing, sender_offset).unwrap(); - let expected_horizon = builder.event_cycles(); - let program = builder.finish().unwrap(); - assert_eq!(plan_event_cycles(&program).unwrap(), expected_horizon); - let outgoing_timing = sender_row_timing(&outgoing, sender_offset).unwrap(); - - let base = receive_row_timing(&incoming, 0).unwrap(); - let expected = scheduled_receive_window(&base, 0, words, None) - .unwrap() - .events - .into_iter() - .map(|event| (event.cycles, receive_control_signature(event.instruction))) - .collect::>(); - let mut actual = Vec::new(); - let mut cycles = 0; - let mut sent_words = 0; - let mut cursor = 0; - while cursor < program.len() { - let instruction = program[cursor]; - let before = cycles; - let advance = instruction_advance(instruction); - cycles += advance; - if instruction & OPCODE_MASK == DELAY_PIC_OPCODE - || instruction & OPCODE_MASK == DELAY_XPIC_OPCODE - { - actual.push((cycles, receive_control_signature(instruction))); - } else if is_send_control(instruction) { - fused_programs += 1; - actual.push((before + 1, send_control_signature(instruction))); - if before >= outgoing_timing.start_cycles && before < outgoing_timing.end_cycles - { - sent_words += advance; - } - } else if is_send_control_pair(instruction) { - fused_programs += 1; - let payload = program[cursor + 1]; - let pointer = payload & PIC_RECEIVE_ADDRESS_MASK; - let source = payload >> 18; - actual.push((before + 1, (ReceiveEventKind::Pointer, pointer))); - actual.push(( - before + 1, - ( - if source == TILE_MUX_EXCHANGE { - ReceiveEventKind::OrdinaryNeutral - } else { - ReceiveEventKind::OrdinarySource - }, - source, - ), - )); - if before >= outgoing_timing.start_cycles && before < outgoing_timing.end_cycles - { - sent_words += advance; - } - cursor += 1; - } else if is_payload_send(instruction) { - sent_words += advance; - } - cursor += 1; - } - assert_eq!(actual, expected); - assert_eq!(sent_words, words); - assert_eq!( - program - .iter() - .filter(|instruction| **instruction & LONG_OPCODE_MASK == SEND_OPCODE) - .count(), - 1 - ); - } - assert!(fused_programs > 0); - } - - fn receive_control_signature(instruction: u32) -> (ReceiveEventKind, u32) { - if instruction & OPCODE_MASK == DELAY_PIC_OPCODE { - ( - ReceiveEventKind::Pointer, - instruction & PIC_RECEIVE_ADDRESS_MASK, - ) - } else { - let operand = instruction & 0x1fff; - ( - if operand == TILE_MUX_EXCHANGE { - ReceiveEventKind::OrdinaryNeutral - } else { - ReceiveEventKind::OrdinarySource - }, - operand, - ) - } - } - - fn send_control_signature(instruction: u32) -> (ReceiveEventKind, u32) { - let selector = (instruction >> 18) & 3; - let operand = if selector >= 2 { - instruction & PIC_RECEIVE_ADDRESS_MASK - } else { - instruction & 0x1fff - }; - ( - if selector >= 2 { - ReceiveEventKind::Pointer - } else if operand == TILE_MUX_EXCHANGE { - ReceiveEventKind::OrdinaryNeutral - } else { - ReceiveEventKind::OrdinarySource - }, - operand, - ) - } - - #[test] - fn event_horizon_tracks_transfer_size_and_route() { - let topology = Topology::c600(); - let short = topology.multicast(0, &[736, 1286], 1, 0).unwrap(); - let long = topology.multicast(0, &[736, 1286], 1024, 0).unwrap(); - let horizon = |plan: &MulticastPlan| { - std::iter::once(&plan.sender) - .chain(plan.receivers.iter()) - .map(|row| plan_event_cycles(row).unwrap()) - .max() - .unwrap() - }; - - assert!(horizon(&short) > 0); - assert!(horizon(&long) > horizon(&short)); - assert_ne!( - plan_event_cycles(&short.receivers[0]).unwrap(), - plan_event_cycles(&short.receivers[1]).unwrap() - ); - } - - #[test] - fn encoder_places_receive_then_send_on_one_event_timeline() { - let topology = Topology::c600(); - let first = topology.multicast(0, &[736], 64, 0).unwrap(); - let first_horizon = std::iter::once(&first.sender) - .chain(first.receivers.iter()) - .map(|row| plan_event_cycles(row).unwrap()) - .max() - .unwrap(); - let second = topology - .multicast(736, &[1286], 64, first_horizon + 1) - .unwrap(); - let horizon = std::iter::once(&second.sender) - .chain(second.receivers.iter()) - .map(|row| plan_event_cycles(row).unwrap()) - .max() - .unwrap(); - - let mut relay = TileProgramSchedule::default(); - relay - .append_receiver_at(&first.receivers[0], 0, 64) - .unwrap(); - let offset = relay.earliest_sender_offset(&second.sender, 0).unwrap(); - relay.append_sender_at(&second.sender, offset).unwrap(); - let relay_horizon = relay.event_cycles(); - let relay = relay.finish().unwrap(); - - assert_eq!(relay.last(), Some(&RETURN_M10_INSTRUCTION)); - assert_eq!( - relay - .iter() - .filter(|instruction| **instruction == SYNC_SUPERVISOR_INSTRUCTION) - .count(), - 0 - ); - assert_eq!(plan_event_cycles(&relay).unwrap(), relay_horizon); - assert!(relay_horizon <= horizon); - } - - #[test] - fn validates_limits_and_patches_addresses() { - let topology = Topology::c600(); - assert_eq!( - topology.point_to_point(0, 1, 0), - Err(ExchangeError::Count(0)) - ); - assert_eq!( - topology.multicast(0, &[1, 1], 1, 0), - Err(ExchangeError::ReceiverSet) - ); - let mut plan = topology.multicast(0, &[274], 65, 0).unwrap(); - patch_sender_address(&mut plan.sender, 0x52040).unwrap(); - patch_receiver_address(&mut plan.receivers[0], 0x53080).unwrap(); - assert_eq!( - plan.sender[2] & 0x001f_fff8, - ((0x52040 >> 2) << 3) & 0x001f_fff8 - ); - assert_eq!(plan.receivers[0][2] & 0x3ffff, 0x53080 >> 2); - - patch_receiver_address(&mut plan.receivers[0], 0x8f000).unwrap(); - assert_eq!(plan.receivers[0][2] & 0x3ffff, 0x8f000 >> 2); - assert_eq!( - patch_receiver_address(&mut plan.receivers[0], 0x10_0000), - Err(ExchangeError::Address(0x10_0000)) - ); - } - - #[test] - fn encodes_supervisor_send_fields() { - assert_eq!(encode_send(1, 3, 82_041).unwrap(), 0x782a_03cb); - assert_eq!(encode_send(1, 3, 82_043).unwrap(), 0x782a_03db); - assert!(encode_send(64, 3, 0).is_err()); - assert!(encode_send(1, 8, 0).is_err()); - assert!(encode_send(1, 3, 0x4_0000).is_err()); - } - - #[test] - fn host_packets_match_recovered_sdk_vectors() { - assert_eq!( - tile_to_host_packet(0, 0x40, 64).unwrap(), - HostPacketHeader { - word0: 0xa000_0000, - word1: 0x0000_0011, - } - ); - assert_eq!( - host_to_tile_packet(0, 0x50120, 0x40, 64).unwrap(), - HostPacketHeader { - word0: 0xec00_0209, - word1: 0x0000_0011, - } - ); - assert_eq!( - zero_byte_read_packet(2, 0x50180).unwrap(), - HostPacketHeader { - word0: 0xcc01_020c, - word1: 0, - } - ); - assert_eq!( - tile_to_host_packet(1409, 0x40, 64).unwrap(), - HostPacketHeader { - word0: 0xa2c0_8000, - word1: 0x8000_0011, - } - ); - } - - #[test] - fn host_packets_validate_both_size_classes() { - assert!(tile_to_host_packet(0, 4, 4).is_ok()); - assert!(tile_to_host_packet(0, 0x400, 1024).is_ok()); - assert!(tile_to_host_packet(0, 2, 4).is_err()); - assert!(tile_to_host_packet(0, 0, 0).is_err()); - assert!(tile_to_host_packet(0, 0, 1028).is_err()); - assert!(host_to_tile_packet(0, 0x50124, 0x40, 64).is_err()); - assert!(host_to_tile_packet(0, 0x54000, 0x40, 64).is_err()); - assert!(tile_to_host_packet(0x1000, 0, 4).is_err()); - } - - #[test] - fn host_transfer_planner_covers_unaligned_and_large_ranges() { - let d2h = plan_tile_to_host(2, 0x60004, 4, 2200).unwrap(); - assert_eq!(d2h.first().unwrap().host_offset, 4); - assert_eq!(d2h.iter().map(|chunk| chunk.bytes).sum::(), 2200); - assert!(d2h.iter().all(|chunk| chunk.bytes <= 1024)); - assert!( - d2h.windows(2) - .all(|pair| pair[0].tile_address + pair[0].bytes == pair[1].tile_address) - ); - assert!(d2h.iter().all(|chunk| { - chunk.host_offset / HOST_PAGE_BYTES - == (chunk.host_offset + chunk.bytes - 1) / HOST_PAGE_BYTES - })); - - let paged = plan_tile_to_host(0, 0x60000, 64, 4096).unwrap(); - assert_eq!(paged.iter().map(|chunk| chunk.bytes).sum::(), 4096); - assert!(paged.iter().all(|chunk| { - chunk.host_offset / HOST_PAGE_BYTES - == (chunk.host_offset + chunk.bytes - 1) / HOST_PAGE_BYTES - })); - - let h2d = plan_host_to_tile(1409, 0x50000, 4, 100).unwrap(); - assert_eq!( - h2d.iter().map(|chunk| chunk.bytes).collect::>(), - [32, 32, 36] - ); - assert_eq!(h2d.last().unwrap().tile_address + 36, 0x50064); - assert!(plan_host_to_tile(0, 0x50004, 0, 4).is_err()); - } - - #[test] - fn plan_offsets_extend_beyond_route_timing_fields() { - let topology = Topology::c600(); - let mut plan = topology.multicast(0, &[1, 2], 4096, 0).unwrap(); - let sender_cycles = plan_event_cycles(&plan.sender).unwrap(); - let receiver_cycles = plan_event_cycles(&plan.receivers[0]).unwrap(); - let offset = MAX_PLAN_OFFSET_CYCLES + 70; - - offset_plan(&mut plan.sender, offset).unwrap(); - offset_plan(&mut plan.receivers[0], offset).unwrap(); - - assert_eq!( - plan_event_cycles(&plan.sender).unwrap(), - sender_cycles + offset - ); - assert_eq!( - plan_event_cycles(&plan.receivers[0]).unwrap(), - receiver_cycles + offset - ); - } - - #[test] - fn tile_to_host_target_preserves_packet_and_payload_addresses() { - let plan = - assemble_tile_to_host_target_program(2, 0x50120, 0x40, 64, 0x50160, 0x501a0).unwrap(); - assert_eq!( - &plan.packet_words[..2], - &host_packet_words(tile_to_host_packet(2, 0x40, 64).unwrap()) - ); - assert_eq!(plan.packet_words.len(), 4); - let sends = plan - .instructions - .iter() - .copied() - .filter(|word| word & LONG_OPCODE_MASK == SEND_OPCODE) - .collect::>(); - assert_eq!(sends.len(), 3); - assert_eq!(send_address(sends[0]), 0x50160); - assert_eq!(send_address(sends[1]), 0x50120); - assert_eq!(send_address(sends[2]), 0x50168); - assert_eq!(instruction_advance(sends[1]), 16); - assert_eq!(plan.instructions[0], setzi_m(8, 1)); - let payload = plan - .instructions - .iter() - .position(|instruction| *instruction == sends[1]) - .unwrap(); - let close = plan - .instructions - .iter() - .position(|instruction| *instruction == sends[2]) - .unwrap(); - assert_eq!(close, payload + 2); - assert_eq!( - plan.instructions[payload + 1], - delay(TILE_TO_HOST_MIN_PAYLOAD_EVENTS - 64 / 4 - 1) - ); - } - - #[test] - fn randomized_tile_to_host_packets_observe_the_payload_interval() { - let mut random = fastrand::Rng::with_seed(0x686f_7374_5f70_6164); - for _ in 0..256 { - let host_offset = random.u32(0..256) * 4; - let bytes = random.u32(1..=1024) * 4; - let chunks = plan_tile_to_host(2, 0x52000, host_offset, bytes).unwrap(); - let target = assemble_tile_to_host_target_program( - 2, - 0x52000, - host_offset, - bytes, - 0x50160, - 0x501a0, - ) - .unwrap(); - - let mut cursor = 2; - for _ in 0..chunks.len() { - cursor += 1; // packet header - let payload_events = instruction_advance(target.instructions[cursor]); - cursor += 1; - let padding_events = if payload_events < TILE_TO_HOST_MIN_PAYLOAD_EVENTS { - let events = instruction_advance(target.instructions[cursor]); - cursor += 1; - events - } else { - 0 - }; - assert_eq!( - payload_events + padding_events, - payload_events.max(TILE_TO_HOST_MIN_PAYLOAD_EVENTS) - ); - } - assert_eq!(cursor + 3, target.instructions.len()); - } - } - - #[test] - fn tile_to_host_target_has_no_controller_xreq_or_sync_wrapper() { - let target = - assemble_tile_to_host_target_program(2, 0x50120, 0x40, 64, 0x50160, 0x501a0).unwrap(); - - assert_eq!(target.packet_words.len(), 4); - assert_eq!( - &target.packet_words[..2], - &host_packet_words(tile_to_host_packet(2, 0x40, 64).unwrap()) - ); - assert!(!target.instructions.contains(&SYNC_HOST_INSTRUCTION)); - assert!(!target.instructions.contains(&SYNC_ALL_INSTRUCTION)); - assert_eq!( - target.instructions[target.instructions.len() - 2], - SYNC_RECEIVE_INSTRUCTION - ); - } - - #[test] - fn groups_multi_packet_tile_to_host_payloads() { - let chunks = plan_tile_to_host(2, 0x52000, 0x40, 2048).unwrap(); - assert_eq!(chunks.len(), 8); - assert!( - chunks - .iter() - .all(|chunk| chunk.bytes == TILE_TO_HOST_MAX_BYTES) - ); - - let target = - assemble_tile_to_host_target_program(2, 0x52000, 0x40, 2048, 0x50160, 0x501a0).unwrap(); - - assert_eq!(target.packet_words.len(), chunks.len() * 2 + 2); - assert_eq!(target.instructions[0], setzi_m(8, 1)); - assert!( - target - .instructions - .iter() - .filter(|word| **word & LONG_OPCODE_MASK == SEND_OPCODE) - .skip(1) - .step_by(2) - .all(|word| instruction_advance(*word) == TILE_TO_HOST_MAX_BYTES / 4) - ); - assert_eq!( - target - .instructions - .iter() - .filter(|word| **word & LONG_OPCODE_MASK == SEND_OPCODE) - .count(), - chunks.len() * 2 + 1 - ); - } - - #[test] - fn host_to_tile_target_preserves_packet_and_request_addresses() { - let plan = assemble_host_to_tile_target_program(2, 0x50120, 0x40, 64, 0x50160).unwrap(); - assert_eq!( - &plan.packet_words[..], - &host_packet_words(host_to_tile_packet(2, 0x50120, 0x40, 64).unwrap()) - ); - let sends = plan - .instructions - .iter() - .copied() - .filter(|word| word & LONG_OPCODE_MASK == SEND_OPCODE) - .collect::>(); - assert_eq!(sends.len(), 1); - assert_eq!(send_address(sends[0]), 0x50160); - assert_eq!( - plan.instructions[plan.instructions.len() - 2], - SYNC_RECEIVE_INSTRUCTION - ); - assert!(!plan.instructions.contains(&SYNC_HOST_INSTRUCTION)); - assert!(!plan.instructions.contains(&SYNC_ALL_INSTRUCTION)); - } - - #[test] - fn groups_multi_packet_host_to_tile_as_one_stream_copy() { - let chunks = plan_host_to_tile(63, 0x50000, 0x40, 4096).unwrap(); - let plan = assemble_host_to_tile_target_program(63, 0x50000, 0x40, 4096, 0x54000).unwrap(); - assert_eq!(plan.packet_words.len(), chunks.len() * 2); - assert!( - plan.packet_words[..plan.packet_words.len() - 2] - .chunks_exact(2) - .all(|header| header[0] & HOST_TO_TILE_STREAM_END_BITS == 0) - ); - assert_eq!( - plan.packet_words[plan.packet_words.len() - 2] & HOST_TO_TILE_STREAM_END_BITS, - HOST_TO_TILE_STREAM_END_BITS - ); - assert_eq!( - plan.instructions - .iter() - .filter(|instruction| **instruction & LONG_OPCODE_MASK == SEND_OFF_OPCODE) - .count(), - chunks.len() - 1 - ); - } - - #[test] - fn host_command_read_encoder_preserves_recovered_packet_and_addresses() { - let plan = assemble_host_command_read_program(0x50160, 0x50180, 0x1000).unwrap(); - assert_eq!(plan.packet_words, [1, 0, 0xcc00_020c, 0x4001]); - assert_eq!(send_address(plan.instructions[5]), 0x50160); - assert_eq!(send_address(plan.instructions[6]), 0x50168); - let command_send = plan.instructions[plan.instructions.len() - 2]; - assert_eq!(send_address(command_send), 0x50180); - } - - fn host_packet_words(header: HostPacketHeader) -> [u32; 2] { - [header.word0, header.word1] - } - - #[test] - fn host_mux_uses_the_physical_tile_row_endpoint() { - for (physical_tile, mux) in [(116, 0x634), (582, 0x604), (1173, 0x615)] { - assert_eq!(host_mux_for_tile(physical_tile).unwrap(), mux); - } - } - - #[test] - fn host_xreq_combines_target_endpoint_bits() { - let targets = [31, 81, 768, 1471]; - let combined = assemble_host_xreq_program_for_targets(&targets, 0x50120).unwrap(); - let expected = targets - .into_iter() - .map(|target| { - assemble_host_xreq_program(target, 0x50120) - .unwrap() - .packet_words - }) - .fold([0u32; 2], |mut bitmap, words| { - bitmap[0] |= words[0]; - bitmap[1] |= words[1]; - bitmap - }); - assert_eq!(combined.packet_words, expected); - assert_eq!( - combined.instructions, - assemble_host_xreq_program(31, 0x50120) - .unwrap() - .instructions - ); - } - - #[test] - fn target_operations_match_sdk_logical_tile_100_oracle() { - let hierarchy = host_hierarchy(260).unwrap(); - assert_eq!(hierarchy.xreq_physical_tile, 4); - let xreq = assemble_host_xreq_program(260, 0x50120).unwrap(); - assert_eq!(xreq.instructions, [0x782a_0243, 0x43a0_0000]); - assert_eq!(xreq.packet_words, [0x100, 0]); - let first_group = host_hierarchy(31).unwrap(); - assert_eq!(first_group.xreq_physical_tile, 29); - assert_eq!( - assemble_host_xreq_program(31, 0x50120) - .unwrap() - .packet_words, - [2, 0] - ); - assert_eq!( - assemble_host_xreq_program(81, 0x50120) - .unwrap() - .packet_words, - [4, 0] - ); - assert_eq!( - assemble_host_xreq_program(768, 0x50120) - .unwrap() - .packet_words, - [0, 1] - ); - assert_eq!( - assemble_host_xreq_program(1471, 0x50120) - .unwrap() - .packet_words, - [0, 1 << 21] - ); - let wrapped_xreq = - wrap_host_xreq_operation(hierarchy.xreq_physical_tile, &xreq.instructions).unwrap(); - assert_eq!(&wrapped_xreq[..3], &[0x1980_0604, 0x4380_80a0, 0x4180_000f]); - - let d2h = - assemble_tile_to_host_target_program(260, 0x50120, 0x40, 64, 0x50160, 0x50180).unwrap(); - assert_eq!( - d2h.instructions, - [ - 0x1980_0001, - 0x4380_80a6, - 0x782a_02c3, - 0x79ea_0243, - 0x40a0_0001, - 0x782a_02d3, - 0x4180_0000, - 0x43a0_0000, - ] - ); - assert_eq!(d2h.packet_words, [0xa082_0000, 0x0000_0011, 0xcc82_020c, 0]); - - let h2d = assemble_host_to_tile_target_program(260, 0x50120, 0x40, 64, 0x50170).unwrap(); - assert_eq!( - h2d.instructions, - [ - 0x1980_0010, - 0x4380_80a6, - 0x782a_02e3, - 0x4180_0000, - 0x43a0_0000, - ] - ); - assert_eq!(h2d.packet_words, [0xec82_0209, 0x0000_0011]); - - let wrapped = wrap_host_target_operation(260, &d2h.instructions).unwrap(); - assert_eq!(&wrapped[..3], &[0x1980_0604, 0x4380_80a0, 0x4180_0007]); - assert_eq!( - &wrapped[wrapped.len() - 4..], - &[0x4180_0007, 0x1980_0640, 0x4380_80a0, 0x43a0_0000] - ); - - let local = assemble_host_to_tile_target_program(0, 0x50120, 0x40, 64, 0x50168).unwrap(); - let wrapped_local = wrap_combined_host_operation(0, &local.instructions, 0x50160).unwrap(); - assert_eq!( - &wrapped_local[..3], - &[0x1980_0600, 0x4380_80a0, 0x4180_000f] - ); - - let wrapped_tile_nine = - wrap_combined_host_operation(9, &local.instructions, 0x50160).unwrap(); - assert_eq!(wrapped_tile_nine[0], 0x1980_0609); - assert_eq!(send_address(wrapped_local[5]), 0x50160); - } - - fn send_address(instruction: u32) -> u32 { - ((instruction & 0x001f_fff8) >> 3) * 4 - } - - #[test] - fn finalizes_point_receiver_for_direct_execution() { - let topology = Topology::c600(); - let plan = topology.point_to_point(274, 1286, 64).unwrap(); - let row = finalize_point_receiver(&plan.receiver, topology.physical(274).unwrap()).unwrap(); - assert_eq!(row[0], SYNC_SUPERVISOR_INSTRUCTION); - assert_eq!(row[1] & 0x1fff, 9); - assert_eq!(row[5], RETURN_M10_INSTRUCTION); - } -} diff --git a/crates/ipu-package/Cargo.toml b/crates/ipu-package/Cargo.toml index 2bb95e9..5c468e3 100644 --- a/crates/ipu-package/Cargo.toml +++ b/crates/ipu-package/Cargo.toml @@ -7,6 +7,7 @@ build = "build.rs" [dependencies] capnp.workspace = true +ipu-target = { path = "../ipu-target" } thiserror.workspace = true tracing.workspace = true diff --git a/crates/ipu-package/build.rs b/crates/ipu-package/build.rs index 5816ddc..09951ec 100644 --- a/crates/ipu-package/build.rs +++ b/crates/ipu-package/build.rs @@ -1,10 +1,12 @@ fn main() { println!("cargo:rerun-if-changed=../../schemas/application.capnp"); println!("cargo:rerun-if-changed=../../schemas/profile.capnp"); + println!("cargo:rerun-if-changed=../../schemas/profile_common.capnp"); capnpc::CompilerCommand::new() .src_prefix("../../schemas") .file("../../schemas/application.capnp") .file("../../schemas/profile.capnp") + .file("../../schemas/profile_common.capnp") .run() .expect("compile application.capnp"); } * Unmerged path crates/ipu-package/src/lib.rs diff --git a/crates/ipu-profile/src/lib.rs b/crates/ipu-profile/src/lib.rs index b7edd17..66e00f6 100644 --- a/crates/ipu-profile/src/lib.rs +++ b/crates/ipu-profile/src/lib.rs @@ -1,6 +1,5 @@ use ipu_package::{ - CycleSample, ProfileExchangeActivity, ProfileExchangeActivityKind, ProfileReport, - ProfileStepKind, + CycleSample, ExchangeActivityKind, ProfileExchangeActivity, ProfileReport, ProfileStepKind, }; use serde::{Deserialize, Serialize}; use std::collections::{BTreeMap, BTreeSet, HashMap}; @@ -16,6 +15,35 @@ pub enum GroupBy { Metadata, } +impl std::fmt::Display for GroupBy { + fn fmt(&self, output: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + output.write_str(match self { + Self::Kind => "kind", + Self::Kernel => "kernel", + Self::Operation => "operation", + Self::Phase => "phase", + Self::Tile => "tile", + Self::Metadata => "metadata", + }) + } +} + +impl std::str::FromStr for GroupBy { + type Err = &'static str; + + fn from_str(value: &str) -> Result { + match value { + "kind" => Ok(Self::Kind), + "kernel" => Ok(Self::Kernel), + "operation" => Ok(Self::Operation), + "phase" => Ok(Self::Phase), + "tile" => Ok(Self::Tile), + "metadata" => Ok(Self::Metadata), + _ => Err("expected kind, kernel, operation, phase, tile, or metadata"), + } + } +} + #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] pub enum SortBy { #[default] @@ -26,12 +54,31 @@ pub enum SortBy { Name, } -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum StepKind { - Exchange, - Compute, - Synchronization, - Idle, +impl std::fmt::Display for SortBy { + fn fmt(&self, output: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + output.write_str(match self { + Self::PhaseCycles => "phase-cycles", + Self::WorkCycles => "work-cycles", + Self::MaximumCycles => "maximum-cycles", + Self::Samples => "samples", + Self::Name => "name", + }) + } +} + +impl std::str::FromStr for SortBy { + type Err = &'static str; + + fn from_str(value: &str) -> Result { + match value { + "phase-cycles" => Ok(Self::PhaseCycles), + "work-cycles" => Ok(Self::WorkCycles), + "maximum-cycles" => Ok(Self::MaximumCycles), + "samples" => Ok(Self::Samples), + "name" => Ok(Self::Name), + _ => Err("expected phase-cycles, work-cycles, maximum-cycles, samples, or name"), + } + } } #[derive(Clone, Debug, PartialEq, Eq)] @@ -44,7 +91,7 @@ pub struct MetadataFilter { pub struct Query { pub group_by: GroupBy, pub sort_by: SortBy, - pub kind: Option, + pub kind: Option, pub kernel: Option, pub operation_contains: Option, pub tiles: BTreeSet, @@ -351,13 +398,13 @@ pub fn exchange_activity_summary(report: &ProfileReport) -> ExchangeActivitySumm summary.described_samples += 1; for activity in &sample.step.exchange_activities { match activity.kind { - ProfileExchangeActivityKind::Send => { + ExchangeActivityKind::Send => { summary.send_intervals += 1; } - ProfileExchangeActivityKind::Receive => { + ExchangeActivityKind::Receive => { summary.receive_intervals += 1; } - ProfileExchangeActivityKind::PartnerBusy => { + ExchangeActivityKind::PartnerBusy => { summary.partner_busy_intervals += 1; } } @@ -403,9 +450,9 @@ fn exchange_role_cycles( continue; } let (send, receive, partner_busy) = match activity.kind { - ProfileExchangeActivityKind::Send => (1i32, 0i32, 0i32), - ProfileExchangeActivityKind::Receive => (0, 1, 0), - ProfileExchangeActivityKind::PartnerBusy => (0, 0, 1), + ExchangeActivityKind::Send => (1i32, 0i32, 0i32), + ExchangeActivityKind::Receive => (0, 1, 0), + ExchangeActivityKind::PartnerBusy => (0, 0, 1), }; events.push((start, send, receive, partner_busy)); events.push((end, -send, -receive, -partner_busy)); @@ -687,21 +734,16 @@ fn group_key(query: &Query, tile: u32, sample: &CycleSample) -> (String, BTreeMa (name, dimensions) } -fn kind_of(sample: &CycleSample) -> StepKind { - match sample.step.kind { - ProfileStepKind::Exchange => StepKind::Exchange, - ProfileStepKind::Compute => StepKind::Compute, - ProfileStepKind::Synchronization => StepKind::Synchronization, - ProfileStepKind::Idle => StepKind::Idle, - } +fn kind_of(sample: &CycleSample) -> ProfileStepKind { + sample.step.kind } fn kind_name(sample: &CycleSample) -> &'static str { match kind_of(sample) { - StepKind::Exchange => "exchange", - StepKind::Compute => "compute", - StepKind::Synchronization => "synchronization", - StepKind::Idle => "idle", + ProfileStepKind::Exchange => "exchange", + ProfileStepKind::Compute => "compute", + ProfileStepKind::Synchronization => "synchronization", + ProfileStepKind::Idle => "idle", } } @@ -998,7 +1040,7 @@ mod tests { let result = query( &report, &Query { - kind: Some(StepKind::Compute), + kind: Some(ProfileStepKind::Compute), sample_limit: 2, shared_clock: true, ..Query::default() @@ -1022,9 +1064,9 @@ mod tests { fanout: 0, paired: false, kind: match random.u8(0..3) { - 0 => ProfileExchangeActivityKind::Send, - 1 => ProfileExchangeActivityKind::Receive, - _ => ProfileExchangeActivityKind::PartnerBusy, + 0 => ExchangeActivityKind::Send, + 1 => ExchangeActivityKind::Receive, + _ => ExchangeActivityKind::PartnerBusy, }, start_cycle, end_cycle: random.u32(start_cycle + 1..=event_cycles), @@ -1034,15 +1076,15 @@ mod tests { let expected = (0..event_cycles).fold(ExchangeRoleCycles::default(), |mut totals, cycle| { let send = activities.iter().any(|activity| { - activity.kind == ProfileExchangeActivityKind::Send + activity.kind == ExchangeActivityKind::Send && (activity.start_cycle..activity.end_cycle).contains(&cycle) }); let receive = activities.iter().any(|activity| { - activity.kind == ProfileExchangeActivityKind::Receive + activity.kind == ExchangeActivityKind::Receive && (activity.start_cycle..activity.end_cycle).contains(&cycle) }); let partner_busy = activities.iter().any(|activity| { - activity.kind == ProfileExchangeActivityKind::PartnerBusy + activity.kind == ExchangeActivityKind::PartnerBusy && (activity.start_cycle..activity.end_cycle).contains(&cycle) }); totals.send += u64::from(send); diff --git a/crates/ipu-runtime/Cargo.toml b/crates/ipu-runtime/Cargo.toml index 4a61aeb..951eaa5 100644 --- a/crates/ipu-runtime/Cargo.toml +++ b/crates/ipu-runtime/Cargo.toml @@ -7,5 +7,4 @@ license.workspace = true [dependencies] ipu-driver = { path = "../ipu-driver" } ipu-package = { path = "../ipu-package" } -thiserror.workspace = true tracing-subscriber.workspace = true diff --git a/crates/ipu-runtime/src/lib.rs b/crates/ipu-runtime/src/lib.rs index 65309e6..e4fd6ed 100644 --- a/crates/ipu-runtime/src/lib.rs +++ b/crates/ipu-runtime/src/lib.rs @@ -2,13 +2,7 @@ use ipu_driver::{Device, DriverError, HostSession, Loader, block_device_interrup use ipu_package::Application; use tracing_subscriber::EnvFilter; -#[derive(Debug, thiserror::Error)] -pub enum RuntimeError { - #[error("driver error: {0}")] - Driver(#[from] DriverError), -} - -pub type Result = std::result::Result; +pub type Result = std::result::Result; /// Thin ownership wrapper around an initialized device. /// diff --git a/crates/ipu-exchange/Cargo.toml b/crates/ipu-target/Cargo.toml similarity index 90% rename from crates/ipu-exchange/Cargo.toml rename to crates/ipu-target/Cargo.toml index 552e405..cf289b6 100644 --- a/crates/ipu-exchange/Cargo.toml +++ b/crates/ipu-target/Cargo.toml @@ -1,5 +1,5 @@ [package] -name = "ipu-exchange" +name = "ipu-target" version.workspace = true edition.workspace = true license.workspace = true diff --git a/crates/ipu-target/src/cost.rs b/crates/ipu-target/src/cost.rs new file mode 100644 index 0000000..69837b5 --- /dev/null +++ b/crates/ipu-target/src/cost.rs @@ -0,0 +1,62 @@ +//! Empirical and architectural cycle costs for supported targets. + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct HardwareCosts { + pub exchange_bytes_per_cycle: u64, + pub standard_load_bytes_per_cycle: u64, + pub interleaved_load_bytes_per_cycle: u64, + pub local_copy_bytes_per_cycle: u64, + pub reduction_output_bytes_per_cycle: u64, + pub local_copy_call_cycles: u64, + pub exchange_phase_cycles: u64, + pub kernel_launch_cycles: u64, + pub logical_fragment_cycles: u64, + pub amp_call_cycles: u64, + pub amp_grid_search_setup_cycles: u64, + pub amp_column_group_width: u64, + pub amp_interleaved_column_group_cycles: u64, + pub amp_standard_column_group_cycles: u64, + pub indexed_f16_transform_cycles_per_element: u64, + pub amp_left_pack_cycles_per_element: u64, + pub contiguous_panel_pack_cycles_per_element: u64, + pub block_major_pack_startup_cycles: u64, + pub block_major_pack_cycles_per_element: u64, +} + +/// IPU21 architectural costs and measurements used by analytical planning. +pub(crate) const IPU21_TARGET_COSTS: HardwareCosts = HardwareCosts { + // Target::getExchangeBytesPerCycle. + exchange_bytes_per_cycle: 4, + // Target::getMemcpyBytesPerCycle. Interleaved reads use both memory + // elements, while an ordinary read or local copy uses one data path. + standard_load_bytes_per_cycle: 8, + interleaved_load_bytes_per_cycle: 16, + local_copy_bytes_per_cycle: 8, + // Reduction-add reads two partials and writes one. Current profiles + // sustain roughly one output byte per cycle after the three interleaved + // streams and worker imbalance are included. + reduction_output_bytes_per_cycle: 1, + // Finalized six-worker local copy, including both rendezvous. + local_copy_call_cycles: 288, + // Target::getGlobalSyncCycles. + exchange_phase_cycles: 600, + // popops::internal::basicOpSupervisorOverhead(false). + kernel_launch_cycles: 11, + // Endpoint and receive-pointer cutovers in fragmented logical exchange. + logical_fragment_cycles: 160, + // Generated AMP GEMM kernel: one 16-column group and one 64-element K + // block, with the remaining group cost dominated by weight delivery. + amp_call_cycles: 294, + // Coarse grid search excludes retained kernel state and separately models + // the residual supervisor, weight-feed, and worker setup per microblock. + amp_grid_search_setup_cycles: 160, + amp_column_group_width: 16, + amp_interleaved_column_group_cycles: 940, + amp_standard_column_group_cycles: 1_063, + // Generated F16 layout-transform measurements. + indexed_f16_transform_cycles_per_element: 10, + amp_left_pack_cycles_per_element: 4, + contiguous_panel_pack_cycles_per_element: 3, + block_major_pack_startup_cycles: 4_096, + block_major_pack_cycles_per_element: 4, +}; diff --git a/crates/ipu-target/src/emit.rs b/crates/ipu-target/src/emit.rs new file mode 100644 index 0000000..b69800b --- /dev/null +++ b/crates/ipu-target/src/emit.rs @@ -0,0 +1,966 @@ +//! Encoding of finalized per-tile programs into IPU21 machine code. + +use crate::hardware::HardwareTarget; +use crate::instruction::{ + SANS_INACTIVE_INSTRUCTION, SYNC_SUPERVISOR_INSTRUCTION, encode_add_m_immediate, encode_br_m, + encode_brz_m_immediate, encode_call_m_immediate, encode_ld32_m_immediate, encode_put_special_m, + encode_setzi_m, encode_shl_m_immediate, encode_st32_m_immediate, +}; +use crate::program::{ + ComputeStep, ExchangeSetupPatch, ExchangeStep, HostPhase, HostProgram, PlacedExchangeRow, + RepeatStep, StepProfile, TileAddress, TileProgram, TileStep, +}; +#[cfg(test)] +use crate::program::{ExchangePatch, RepeatPointer}; +use std::collections::BTreeMap; + +const FIRST_INPUT_REGISTER: u8 = 3; +const LAST_VALUE_REGISTER: u8 = 9; + +pub const WORKER_BARRIER_SYMBOL: &str = "ipu_stack_static_worker_barrier"; +pub const COMPLETE_SYMBOL: &str = "ipu_stack_static_complete"; +pub const HOST_RUN_SYMBOL: &str = "ipu_stack_static_host_run"; +pub const REPEAT_CALL_SYMBOL: &str = "ipu_stack_static_repeat_call"; +pub const SAMPLE_CYCLE_SYMBOL: &str = "ipu_stack_static_sample_cycle"; +pub const COPY_U16_SYMBOL: &str = "ipu_stack_static_copy_u16"; +pub const COPY_U32_SYMBOL: &str = "ipu_stack_static_copy_u32"; +pub const COPY_U64_SYMBOL: &str = "ipu_stack_copy_u64"; +pub const COPY_STRIDED_U64_SYMBOL: &str = "ipu_stack_copy_strided_u64"; +pub const FILL_ZERO_U64_SYMBOL: &str = "ipu_stack_fill_zero_u64"; +pub const PATCH_WORD_SYMBOL: &str = "ipu_stack_static_patch_word"; +pub const PATCH_ROW_SYMBOL: &str = "ipu_stack_static_patch_row"; +pub const RUNTIME_ENTRY_SYMBOL: &str = "ipu_stack_static_start"; +pub const PROGRAM_ADDRESS_SYMBOL: &str = "ipu_stack_static_program"; +pub const WORKER_SYNC_CONTEXT_SYMBOL: &str = "ipu_stack_static_worker_sync_context"; +pub const WORKER_STACK_BASE_SYMBOL: &str = "ipu_stack_static_worker_stack_base"; +pub const PRNG_SEED_SYMBOL: &str = "ipu_stack_static_prng_seed"; +pub const HOST_STAGING_SYMBOL: &str = "ipu_stack_static_host_staging"; +pub const COMPLETION_ADDRESS_SYMBOL: &str = "ipu_stack_static_completion"; +const PATCHED_BREAKPOINT_TRAP_BASE: u32 = 0x4180_1000; + +#[derive(Debug, thiserror::Error)] +pub enum CodegenError { + #[error("exchange encoding failed: {0}")] + Exchange(#[from] crate::exchange::ExchangeError), + #[error("invalid tile program: {0}")] + Invalid(String), +} + +pub type Result = std::result::Result; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct CodegenOptions { + pub target: HardwareTarget, + /// Address where the first emitted byte will be placed. + pub code_address: u32, + pub invocations: u32, + pub initial_profile_address: Option, + pub final_profile_address: Option, +} + +impl Default for CodegenOptions { + fn default() -> Self { + Self { + target: HardwareTarget::Ipu21, + code_address: 0, + invocations: 1, + initial_profile_address: None, + final_profile_address: None, + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct GeneratedProgram { + pub bytes: Vec, + /// Exchange data retained verbatim for explicit package placement. + pub exchange_rows: Vec, +} + +pub fn emit( + program: &TileProgram, + symbols: &BTreeMap, + host: &HostProgram, + options: &CodegenOptions, +) -> Result { + if options.invocations == 0 { + return Err(invalid("invocation count must be nonzero")); + } + validate(program)?; + + let complete = symbol(symbols, COMPLETE_SYMBOL)?; + let mut code = TileCode { + target: options.target, + ..TileCode::default() + }; + emit_host_phases(&mut code, symbols, &host.initialize)?; + + if options.invocations > 1 { + code.add_immediate(11, 11, -8)?; + code.setzi(0, options.invocations)?; + code.st32(0, 11, 15, 0)?; + } + let invocation_start = code.address(options.code_address)?; + emit_host_phases(&mut code, symbols, &host.inputs)?; + + if let Some(address) = options.initial_profile_address { + emit_cycle_sample(&mut code, symbols, address)?; + } + + let worker_barrier = program + .steps + .iter() + .any(active_exchange) + .then(|| symbol(symbols, WORKER_BARRIER_SYMBOL)) + .transpose()?; + let mut exchange_rows = Vec::new(); + emit_steps( + &mut code, + program.tile, + &program.steps, + symbols, + worker_barrier, + &mut exchange_rows, + None, + None, + options.code_address, + )?; + + if let Some(address) = options.final_profile_address { + emit_cycle_sample(&mut code, symbols, address)?; + } + emit_host_phases(&mut code, symbols, &host.outputs)?; + if options.invocations > 1 { + code.ld32(0, 11, 15, 0)?; + code.add_immediate(0, 0, -1)?; + code.st32(0, 11, 15, 0)?; + let done_branch = code.words.len(); + code.brz(0, 0)?; + code.jump(invocation_start)?; + let done = code.address(options.code_address)?; + code.words[done_branch] = encode_brz_m_immediate(0, done)?; + code.add_immediate(11, 11, 8)?; + } + code.jump(complete)?; + + let mut unique_exchange_rows = BTreeMap::new(); + for row in exchange_rows { + if unique_exchange_rows + .insert(row.address, row.words.clone()) + .is_some_and(|existing| existing != row.words) + { + return Err(invalid("different exchange rows share an address")); + } + } + Ok(GeneratedProgram { + bytes: code.words.into_iter().flat_map(u32::to_le_bytes).collect(), + exchange_rows: unique_exchange_rows + .into_iter() + .map(|(address, words)| PlacedExchangeRow { address, words }) + .collect(), + }) +} + +#[allow(clippy::too_many_arguments)] +fn emit_steps( + code: &mut TileCode, + tile: u16, + steps: &[TileStep], + symbols: &BTreeMap, + worker_barrier: Option, + exchange_rows: &mut Vec, + repeat_pointer_count: Option, + repeat_count: Option, + code_address: u32, +) -> Result<()> { + let mut index = 0; + while index < steps.len() { + let step = &steps[index]; + if let TileStep::Compute(compute) = step + && compute.symbol == COPY_U16_SYMBOL + && let Some((source, destination)) = absolute_u16_copy(compute) + { + if let Some(address) = compute.profile.before { + emit_cycle_sample(code, symbols, address)?; + } + let mut copies = vec![(source, destination)]; + let mut end = index + 1; + while end < steps.len() + && step_compute_profile(&steps[end - 1]) + .is_none_or(|profile| profile.after.is_none()) + { + let TileStep::Compute(next) = &steps[end] else { + break; + }; + if next.symbol != COPY_U16_SYMBOL || next.profile.before.is_some() { + break; + } + let Some(copy) = absolute_u16_copy(next) else { + break; + }; + copies.push(copy); + end += 1; + if next.profile.after.is_some() { + break; + } + } + code.setzi( + 4, + u32::try_from(copies.len()) + .map_err(|_| invalid("halfword copy table is too large"))?, + )?; + code.call(symbol(symbols, COPY_U16_SYMBOL)?, 10)?; + for (source, destination) in copies { + code.instruction(source); + code.instruction(destination); + } + if let Some(address) = step_compute_profile(&steps[end - 1]).and_then(|p| p.after) { + emit_cycle_sample(code, symbols, address)?; + } + index = end; + continue; + } + match step { + TileStep::Exchange(exchange) => { + let target = *code.target.exchange(); + if let Some(address) = exchange.profile.before { + emit_cycle_sample(code, symbols, address)?; + } + if let Some(patch) = &exchange.setup_patch { + emit_exchange_setup_patch(code, exchange, patch, symbols)?; + } + if !exchange.repeat_patches.is_empty() { + emit_exchange_patches( + code, + exchange, + repeat_count.ok_or_else(|| invalid("exchange patches outside repeat"))?, + symbols, + )?; + } + if !exchange.preserve_base_registers { + code.setzi(8, exchange.incoming_base)?; + code.put_special(target.incoming_base, 8)?; + } + if let Some(source) = exchange.incoming_mux { + code.setzi(8, u32::from(source))?; + code.put_special(target.incoming_mux, 8)?; + } + if exchange.incoming_format != 0 { + code.setzi(8, u32::from(exchange.incoming_format))?; + code.put_special(target.incoming_format, 8)?; + } + if let Some(source) = exchange.incoming_mux_pair { + code.setzi(8, u32::from(source))?; + code.put_special(target.incoming_mux_pair, 8)?; + } + if !exchange.preserve_base_registers { + code.put_special(target.outgoing_base, 15)?; + } + if exchange.active { + code.call( + worker_barrier.expect("active exchange phase has worker barrier"), + 7, + )?; + if exchange.incoming_dcount.is_some() || !exchange.sync_in_program { + code.setzi( + 8, + exchange + .incoming_dcount + .unwrap_or(target.internal_exchange_dcount), + )?; + code.put_special(target.incoming_dcount, 8)?; + } + } + if exchange.active && !exchange.sync_in_program { + code.instruction(SYNC_SUPERVISOR_INSTRUCTION); + } else if !exchange.active { + code.instruction(SANS_INACTIVE_INSTRUCTION); + code.instruction(crate::instruction::SYNC_ANS_INSTRUCTION); + } + code.call(exchange.program.address, 10)?; + if let Some(address) = exchange.profile.after { + emit_cycle_sample(code, symbols, address)?; + } + exchange_rows.push(exchange.program.clone()); + if let Some(patch) = &exchange.setup_patch { + exchange_rows.push(patch.offsets.clone()); + exchange_rows.push(patch.values.clone()); + } + exchange_rows.extend( + exchange + .repeat_patches + .iter() + .map(|patch| patch.values.clone()), + ); + } + TileStep::Compute(compute) => { + if let Some(address) = compute.profile.before { + emit_cycle_sample(code, symbols, address)?; + } + emit_compute(code, tile, compute, symbols, repeat_pointer_count)?; + if let Some(address) = compute.profile.after { + emit_cycle_sample(code, symbols, address)?; + } + } + TileStep::Repeat(repeat) => { + if let Some(address) = repeat.profile.before { + emit_cycle_sample(code, symbols, address)?; + } + emit_repeat( + code, + tile, + repeat, + symbols, + worker_barrier, + exchange_rows, + code_address, + )?; + if let Some(address) = repeat.profile.after { + emit_cycle_sample(code, symbols, address)?; + } + } + TileStep::Checkpoint(checkpoint) => { + code.instruction(PATCHED_BREAKPOINT_TRAP_BASE | u32::from(checkpoint.breakpoint)) + } + } + index += 1; + } + Ok(()) +} + +fn absolute_u16_copy(compute: &ComputeStep) -> Option<(u32, u32)> { + let [TileAddress::Absolute(source)] = compute.input_addresses.as_slice() else { + return None; + }; + let TileAddress::Absolute(destination) = compute.output_address else { + return None; + }; + (compute.arguments.as_slice() == [1]).then_some((*source, destination)) +} + +fn step_compute_profile(step: &TileStep) -> Option<&StepProfile> { + match step { + TileStep::Compute(compute) => Some(&compute.profile), + TileStep::Exchange(_) | TileStep::Repeat(_) | TileStep::Checkpoint(_) => None, + } +} + +fn validate(program: &TileProgram) -> Result<()> { + validate_steps(&program.steps, None, None) +} + +fn validate_steps( + steps: &[TileStep], + repeat_pointer_count: Option, + repeat_count: Option, +) -> Result<()> { + for step in steps { + match step { + TileStep::Exchange(exchange) => { + validate_exchange_program(exchange)?; + if exchange.setup_patch.as_ref().is_some_and(|patch| { + patch.offsets.words.is_empty() + || patch.offsets.words.len() != patch.values.words.len() + }) { + return Err(invalid("exchange setup patch has an invalid shape")); + } + for patch in &exchange.repeat_patches { + if repeat_count.is_none_or(|count| patch.values.words.len() != count as usize) + || patch.word_offset as usize >= exchange.program.words.len() + || patch.values.address & 0b11 != 0 + { + return Err(invalid("exchange patch has invalid shape or address")); + } + } + } + TileStep::Compute(compute) => { + if compute.symbol.is_empty() { + return Err(invalid("compute symbol is empty")); + } + let values = compute.input_addresses.len() + compute.arguments.len(); + let available = usize::from(LAST_VALUE_REGISTER - FIRST_INPUT_REGISTER + 1); + if values == 0 || values > available { + return Err(invalid(format!( + "kernel {} needs {values} input/argument registers; 1..={available} are supported", + compute.symbol + ))); + } + validate_address(compute.output_address, repeat_pointer_count)?; + for &address in &compute.input_addresses { + validate_address(address, repeat_pointer_count)?; + } + } + TileStep::Repeat(repeat) => { + if repeat_pointer_count.is_some() { + return Err(invalid("nested finalized repeats are not yet supported")); + } + if repeat.count == 0 { + return Err(invalid("repeat count must be nonzero")); + } + validate_steps( + &repeat.body, + Some(repeat.iterated_pointers.len()), + Some(repeat.count), + )?; + } + TileStep::Checkpoint(checkpoint) => { + if checkpoint.breakpoint > 1 { + return Err(invalid("checkpoint breakpoint must be zero or one")); + } + } + } + } + Ok(()) +} + +fn validate_exchange_program(exchange: &ExchangeStep) -> Result<()> { + let embedded_sync = exchange + .program + .words + .first() + .is_some_and(|word| *word == SYNC_SUPERVISOR_INSTRUCTION); + if exchange.program.address & 0b11 != 0 + || exchange.program.words.last() != Some(&crate::instruction::RETURN_M10_INSTRUCTION) + || embedded_sync != exchange.sync_in_program + || exchange + .program + .words + .iter() + .skip(usize::from(embedded_sync)) + .any(|word| { + matches!( + *word, + SANS_INACTIVE_INSTRUCTION | SYNC_SUPERVISOR_INSTRUCTION + ) + }) + { + return Err(invalid( + "exchange phase has an invalid boundary or timed program", + )); + } + if exchange.active != (exchange.program.words.len() > 1 + usize::from(embedded_sync)) { + return Err(invalid( + "exchange participation does not match timed program", + )); + } + Ok(()) +} + +fn validate_address(address: TileAddress, repeat_pointer_count: Option) -> Result<()> { + if let TileAddress::RepeatPointer { index, .. } = address + && repeat_pointer_count.is_none_or(|count| usize::from(index) >= count) + { + return Err(invalid( + "compute address refers to an unavailable repeat pointer", + )); + } + Ok(()) +} + +fn active_exchange(step: &TileStep) -> bool { + match step { + TileStep::Exchange(exchange) => exchange.active, + TileStep::Repeat(repeat) => repeat.body.iter().any(active_exchange), + TileStep::Compute(_) | TileStep::Checkpoint(_) => false, + } +} + +fn emit_exchange_patches( + code: &mut TileCode, + exchange: &ExchangeStep, + repeat_count: u32, + symbols: &BTreeMap, +) -> Result<()> { + let helper = symbol(symbols, PATCH_WORD_SYMBOL)?; + for patch in &exchange.repeat_patches { + let byte_offset = patch + .word_offset + .checked_mul(4) + .ok_or_else(|| invalid("exchange patch offset overflow"))?; + code.setzi( + 2, + exchange + .program + .address + .checked_add(byte_offset) + .ok_or_else(|| invalid("exchange patch address overflow"))?, + )?; + code.setzi(3, patch.values.address)?; + code.ld32(4, 11, 15, 0)?; + code.setzi(5, repeat_count)?; + code.call(helper, 9)?; + } + Ok(()) +} + +fn emit_exchange_setup_patch( + code: &mut TileCode, + exchange: &ExchangeStep, + patch: &ExchangeSetupPatch, + symbols: &BTreeMap, +) -> Result<()> { + code.setzi(2, exchange.program.address)?; + code.setzi(3, patch.offsets.address)?; + code.setzi(4, patch.values.address)?; + code.setzi( + 5, + u32::try_from(patch.values.words.len()) + .map_err(|_| invalid("exchange setup patch is too large"))?, + )?; + code.call(symbol(symbols, PATCH_ROW_SYMBOL)?, 9) +} + +fn emit_compute( + code: &mut TileCode, + tile: u16, + compute: &ComputeStep, + symbols: &BTreeMap, + repeat_pointer_count: Option, +) -> Result<()> { + let argument_base = FIRST_INPUT_REGISTER + .checked_add( + u8::try_from(compute.input_addresses.len()) + .map_err(|_| invalid("kernel input count exceeds u8"))?, + ) + .ok_or_else(|| invalid("kernel input register overflow"))?; + emit_address(code, 2, compute.output_address, repeat_pointer_count)?; + for (index, &address) in compute.input_addresses.iter().enumerate() { + emit_address( + code, + FIRST_INPUT_REGISTER + + u8::try_from(index).map_err(|_| invalid("kernel input count exceeds u8"))?, + address, + repeat_pointer_count, + )?; + } + for (index, &argument) in compute.arguments.iter().enumerate() { + code.setzi( + argument_base + + u8::try_from(index).map_err(|_| invalid("kernel argument count exceeds u8"))?, + argument, + )?; + } + let kernel = symbols.get(&compute.symbol).copied().ok_or_else(|| { + invalid(format!( + "tile {tile} references missing kernel symbol {}", + compute.symbol + )) + })?; + code.call(kernel, 10) +} + +fn emit_address( + code: &mut TileCode, + register: u8, + address: TileAddress, + repeat_pointer_count: Option, +) -> Result<()> { + match address { + TileAddress::Absolute(address) => code.setzi(register, address), + TileAddress::RepeatPointer { index, offset } => { + let count = repeat_pointer_count + .ok_or_else(|| invalid("repeat pointer used outside repeat body"))?; + if usize::from(index) >= count { + return Err(invalid("repeat pointer index is out of range")); + } + code.ld32(register, 11, 15, index + 1)?; + code.add_unsigned(register, offset) + } + } +} + +#[allow(clippy::too_many_arguments)] +fn emit_repeat( + code: &mut TileCode, + tile: u16, + repeat: &RepeatStep, + symbols: &BTreeMap, + worker_barrier: Option, + exchange_rows: &mut Vec, + code_address: u32, +) -> Result<()> { + let words = repeat + .iterated_pointers + .len() + .checked_add(1) + .ok_or_else(|| invalid("repeat frame size overflow"))?; + let frame_bytes = i32::try_from((words * 4).next_multiple_of(8)) + .map_err(|_| invalid("repeat frame is too large"))?; + code.add_immediate(11, 11, -frame_bytes)?; + code.setzi(0, repeat.count)?; + code.st32(0, 11, 15, 0)?; + for (index, pointer) in repeat.iterated_pointers.iter().enumerate() { + code.setzi(0, pointer.initial_address)?; + code.st32( + 0, + 11, + 15, + u16::try_from(index + 1).map_err(|_| invalid("too many repeat pointers"))?, + )?; + } + let loop_start = code.address(code_address)?; + emit_steps( + code, + tile, + &repeat.body, + symbols, + worker_barrier, + exchange_rows, + Some(repeat.iterated_pointers.len()), + Some(repeat.count), + code_address, + )?; + for (index, pointer) in repeat.iterated_pointers.iter().enumerate() { + let slot = u16::try_from(index + 1).map_err(|_| invalid("too many repeat pointers"))?; + code.ld32(0, 11, 15, slot)?; + code.add_unsigned(0, pointer.stride_bytes)?; + code.st32(0, 11, 15, slot)?; + } + code.ld32(0, 11, 15, 0)?; + code.add_immediate(0, 0, -1)?; + code.st32(0, 11, 15, 0)?; + let done_branch = code.words.len(); + code.brz(0, 0)?; + code.jump(loop_start)?; + let done = code.address(code_address)?; + code.words[done_branch] = encode_brz_m_immediate(0, done)?; + code.add_immediate(11, 11, frame_bytes) +} + +fn emit_host_phases( + code: &mut TileCode, + symbols: &BTreeMap, + phases: &[HostPhase], +) -> Result<()> { + if phases.is_empty() { + return Ok(()); + } + let repeat_call = phases + .iter() + .any(|phase| !phase.active) + .then(|| symbol(symbols, REPEAT_CALL_SYMBOL)) + .transpose()?; + let host_run = phases + .iter() + .any(|phase| phase.active) + .then(|| symbol(symbols, HOST_RUN_SYMBOL)) + .transpose()?; + let mut index = 0; + while index < phases.len() { + let start = index; + if phases[start].active { + while index < phases.len() + && phases[index].active + && phases[index].address == phases[start].address + { + index += 1; + } + code.setzi( + 2, + u32::try_from(index - start).map_err(|_| invalid("host run overflow"))?, + )?; + code.setzi( + 3, + phases[start] + .run_table + .ok_or_else(|| invalid("active host phase has no run table"))?, + )?; + code.setzi(4, phases[start].address)?; + code.call(host_run.expect("active host phase has host runner"), 9)?; + } else { + while index < phases.len() && !phases[index].active { + index += 1; + } + code.setzi( + 2, + u32::try_from(index - start).map_err(|_| invalid("host run overflow"))?, + )?; + code.setzi(3, phases[start].address)?; + code.call( + repeat_call.expect("inactive host phase has repeat helper"), + 9, + )?; + } + } + Ok(()) +} + +fn emit_cycle_sample( + code: &mut TileCode, + symbols: &BTreeMap, + address: u32, +) -> Result<()> { + code.setzi(2, address)?; + code.call(symbol(symbols, SAMPLE_CYCLE_SYMBOL)?, 10) +} + +fn symbol(symbols: &BTreeMap, name: &str) -> Result { + symbols + .get(name) + .copied() + .ok_or_else(|| invalid(format!("missing runtime symbol {name}"))) +} + +fn invalid(message: impl Into) -> CodegenError { + CodegenError::Invalid(message.into()) +} + +#[derive(Default)] +struct TileCode { + target: HardwareTarget, + words: Vec, +} + +impl TileCode { + fn address(&self, base: u32) -> Result { + base.checked_add( + u32::try_from(self.words.len()) + .map_err(|_| invalid("generated code exceeds u32"))? + .checked_mul(4) + .ok_or_else(|| invalid("generated code size overflow"))?, + ) + .ok_or_else(|| invalid("generated code address overflow")) + } + + fn setzi(&mut self, register: u8, immediate: u32) -> Result<()> { + if immediate < 1 << 20 { + self.words.push(encode_setzi_m(register, immediate)?); + } else { + self.words.push(encode_setzi_m(register, immediate >> 12)?); + self.words + .push(encode_shl_m_immediate(register, register, 12)?); + self.words.push(encode_add_m_immediate( + register, + register, + i32::from((immediate & 0xfff) as u16), + )?); + } + Ok(()) + } + + fn instruction(&mut self, instruction: u32) { + self.words.push(instruction); + } + + fn ld32(&mut self, destination: u8, base: u8, delta: u8, offset: u16) -> Result<()> { + self.words + .push(encode_ld32_m_immediate(destination, base, delta, offset)?); + Ok(()) + } + + fn st32(&mut self, source: u8, base: u8, delta: u8, offset: u16) -> Result<()> { + self.words + .push(encode_st32_m_immediate(source, base, delta, offset)?); + Ok(()) + } + + fn add_immediate(&mut self, destination: u8, source: u8, immediate: i32) -> Result<()> { + self.words + .push(encode_add_m_immediate(destination, source, immediate)?); + Ok(()) + } + + fn add_unsigned(&mut self, register: u8, mut immediate: u32) -> Result<()> { + while immediate != 0 { + let part = immediate.min(i16::MAX as u32); + self.add_immediate(register, register, part as i32)?; + immediate -= part; + } + Ok(()) + } + + fn put_special(&mut self, special: u8, register: u8) -> Result<()> { + self.words.push(encode_put_special_m(special, register)?); + Ok(()) + } + + fn call(&mut self, target: u32, return_register: u8) -> Result<()> { + self.words + .push(encode_call_m_immediate(return_register, target)?); + Ok(()) + } + + fn brz(&mut self, register: u8, target: u32) -> Result<()> { + self.words.push(encode_brz_m_immediate(register, target)?); + Ok(()) + } + + fn jump(&mut self, target: u32) -> Result<()> { + self.setzi(0, target)?; + self.words.push(encode_br_m(0)?); + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn symbols() -> BTreeMap { + [ + (WORKER_BARRIER_SYMBOL.into(), 0x50000), + (COMPLETE_SYMBOL.into(), 0x50004), + (HOST_RUN_SYMBOL.into(), 0x50008), + (REPEAT_CALL_SYMBOL.into(), 0x5000c), + (SAMPLE_CYCLE_SYMBOL.into(), 0x50010), + (PATCH_WORD_SYMBOL.into(), 0x50014), + ("gemm".into(), 0x51000), + ] + .into_iter() + .collect() + } + + #[test] + fn emits_resolved_exchange_and_compute_steps() { + let program = TileProgram { + tile: 7, + steps: vec![ + TileStep::Exchange(ExchangeStep { + active: false, + incoming_base: 0, + preserve_base_registers: false, + incoming_mux: None, + incoming_format: 0, + incoming_mux_pair: None, + incoming_dcount: None, + sync_in_program: false, + program: PlacedExchangeRow { + address: 0x60000, + words: vec![crate::instruction::RETURN_M10_INSTRUCTION], + }, + setup_patch: None, + repeat_patches: Vec::new(), + profile: StepProfile::default(), + }), + TileStep::Compute(ComputeStep { + symbol: "gemm".into(), + output_address: TileAddress::Absolute(0x70000), + input_addresses: vec![ + TileAddress::Absolute(0x71000), + TileAddress::Absolute(0x72000), + ], + arguments: vec![64], + profile: StepProfile::default(), + }), + ], + }; + let generated = emit( + &program, + &symbols(), + &HostProgram::default(), + &CodegenOptions { + code_address: 0x52000, + ..CodegenOptions::default() + }, + ) + .unwrap(); + assert!(!generated.bytes.is_empty()); + assert_eq!(generated.exchange_rows.len(), 1); + assert_eq!(generated.exchange_rows[0].address, 0x60000); + } + + #[test] + fn rejects_unresolved_or_malformed_inputs() { + let program = TileProgram { + tile: 0, + steps: vec![TileStep::Exchange(ExchangeStep { + active: false, + incoming_base: 0, + preserve_base_registers: false, + incoming_mux: None, + incoming_format: 0, + incoming_mux_pair: None, + incoming_dcount: None, + sync_in_program: false, + program: PlacedExchangeRow { + address: 3, + words: Vec::new(), + }, + setup_patch: None, + repeat_patches: Vec::new(), + profile: StepProfile::default(), + })], + }; + assert!(matches!( + emit( + &program, + &symbols(), + &HostProgram::default(), + &CodegenOptions::default() + ), + Err(CodegenError::Invalid(_)) + )); + } + + #[test] + fn randomized_repeat_patch_code_is_independent_of_iteration_count() { + let mut random = fastrand::Rng::with_seed(0x7061_7463_685f_7265); + let mut code_bytes = None; + for _ in 0..64 { + let count = random.u32(2..=128); + let values = (0..count).map(|_| random.u32(..)).collect::>(); + let program = TileProgram { + tile: 0, + steps: vec![TileStep::Repeat(RepeatStep { + count, + iterated_pointers: vec![RepeatPointer { + initial_address: 0x70000, + stride_bytes: 64, + }], + body: vec![TileStep::Exchange(ExchangeStep { + active: true, + incoming_base: 0x70000, + preserve_base_registers: false, + incoming_mux: None, + incoming_format: 0, + incoming_mux_pair: None, + incoming_dcount: None, + sync_in_program: false, + program: PlacedExchangeRow { + address: 0x60000, + words: vec![0, crate::instruction::RETURN_M10_INSTRUCTION], + }, + setup_patch: None, + repeat_patches: vec![ExchangePatch { + word_offset: 0, + values: PlacedExchangeRow { + address: 0x61000, + words: values.clone(), + }, + }], + profile: StepProfile::default(), + })], + profile: StepProfile::default(), + })], + }; + let generated = emit( + &program, + &symbols(), + &HostProgram::default(), + &CodegenOptions { + code_address: 0x52000, + ..CodegenOptions::default() + }, + ) + .unwrap(); + assert_eq!(generated.exchange_rows.len(), 2); + assert_eq!(generated.exchange_rows[1].words, values); + let emitted_words = generated + .bytes + .chunks_exact(4) + .map(|word| u32::from_le_bytes(word.try_into().unwrap())) + .collect::>(); + assert_eq!( + emitted_words + .iter() + .filter(|word| **word == SYNC_SUPERVISOR_INSTRUCTION) + .count(), + 1 + ); + assert_eq!( + *code_bytes.get_or_insert(generated.bytes.len()), + generated.bytes.len() + ); + } + } +} * Unmerged path crates/ipu-target/src/exchange.rs diff --git a/crates/ipu-exchange/src/diagnostic.rs b/crates/ipu-target/src/exchange/parse.rs similarity index 97% rename from crates/ipu-exchange/src/diagnostic.rs rename to crates/ipu-target/src/exchange/parse.rs index 41e5f37..70166ff 100644 --- a/crates/ipu-exchange/src/diagnostic.rs +++ b/crates/ipu-target/src/exchange/parse.rs @@ -5,7 +5,14 @@ //! two-word supervisor instruction: the first word carries the send fields and //! the second is inline PIC/XPIC payload, not an independently executed word. -use super::*; +use serde::{Deserialize, Serialize}; + +use super::{ExchangeError, ReceiveEventKind, TileProgramSchedule}; +use crate::instruction::{ + DELAY_OPCODE, DELAY_OPCODE_MASK, DELAY_PIC_OPCODE, DELAY_XPIC_OPCODE, LONG_OPCODE_MASK, + OPCODE_MASK, PIC_RECEIVE_ADDRESS_MASK, RETURN_M10_INSTRUCTION, SEND_ADDRESS_MASK, SEND_OPCODE, + SYNC_OPCODE, is_send_control, is_send_control_pair, is_send_off, +}; #[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] pub enum IncomingControlStream { @@ -422,6 +429,8 @@ fn control_key(control: IncomingControl) -> (u8, u32) { #[cfg(test)] mod tests { use super::*; + use crate::exchange::{ReceiveEvent, encode_send_control_pair}; + use crate::instruction::{delay_pic, delay_xpic}; #[test] fn sdk_receiver_row_decodes_two_word_controls_as_single_instructions() { diff --git a/crates/ipu-target/src/hardware.rs b/crates/ipu-target/src/hardware.rs new file mode 100644 index 0000000..f0d2321 --- /dev/null +++ b/crates/ipu-target/src/hardware.rs @@ -0,0 +1,56 @@ +//! Supported hardware targets and their resource limits. + +use crate::cost::{HardwareCosts, IPU21_TARGET_COSTS}; +use crate::exchange::{ExchangeConstants, IPU21_EXCHANGE_CONSTANTS}; +use crate::memory::{ + IPU21_DEFAULT_SUPPORT_RESERVATION_BYTES, IPU21_INTERLEAVED_ELEMENT_SIZE, + IPU21_INTERLEAVED_REGION_BYTES, IPU21_PLANNED_DATA_BYTES, IPU21_STANDARD_FIXED_BYTES, +}; +use crate::topology::Topology; + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum HardwareTarget { + #[default] + Ipu21, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct HardwareMemoryConstraints { + pub standard_fixed_bytes: u64, + pub interleaved_bytes: u64, + pub interleaved_element_bytes: u64, + pub total_bytes: u64, + pub default_standard_reservation_bytes: u64, +} + +impl HardwareTarget { + pub const fn costs(self) -> &'static HardwareCosts { + match self { + Self::Ipu21 => &IPU21_TARGET_COSTS, + } + } + + pub const fn exchange(self) -> &'static ExchangeConstants { + match self { + Self::Ipu21 => &IPU21_EXCHANGE_CONSTANTS, + } + } + + pub fn topology(self) -> Topology { + match self { + Self::Ipu21 => Topology::c600(), + } + } + + pub const fn memory_constraints(self) -> HardwareMemoryConstraints { + match self { + Self::Ipu21 => HardwareMemoryConstraints { + standard_fixed_bytes: IPU21_STANDARD_FIXED_BYTES as u64, + interleaved_bytes: IPU21_INTERLEAVED_REGION_BYTES as u64, + interleaved_element_bytes: IPU21_INTERLEAVED_ELEMENT_SIZE as u64, + total_bytes: IPU21_PLANNED_DATA_BYTES as u64, + default_standard_reservation_bytes: IPU21_DEFAULT_SUPPORT_RESERVATION_BYTES as u64, + }, + } + } +} diff --git a/crates/ipu-target/src/instruction.rs b/crates/ipu-target/src/instruction.rs new file mode 100644 index 0000000..6be5d84 --- /dev/null +++ b/crates/ipu-target/src/instruction.rs @@ -0,0 +1,249 @@ +//! IPU21 supervisor instruction encoders used by generated tile programs. + +use crate::exchange::ExchangeError; + +pub(crate) const OPCODE_MASK: u32 = 0xfc00_0000; +pub(crate) const LONG_OPCODE_MASK: u32 = 0xf800_0000; +pub(crate) const DELAY_OPCODE_MASK: u32 = 0xfff8_0000; +pub(crate) const DELAY_OPCODE: u32 = 0x40a0_0000; +pub(crate) const DELAY_PIC_OPCODE: u32 = 0x6000_0000; +pub(crate) const DELAY_XPIC_OPCODE: u32 = 0x6400_0000; +pub(crate) const PIC_RECEIVE_ADDRESS_MASK: u32 = 0x3ffff; +pub(crate) const SEND_OPCODE: u32 = 0x7800_0000; +pub(crate) const SEND_ADDRESS_MASK: u32 = 0x001f_fff8; +pub(crate) const SEND_OFF_OPCODE: u32 = 0x7000_0000; +// See docs/EXCHANGE_INSTRUCTION_REFERENCE.md. SENDPICP is an aligned two-word +// supervisor instruction whose following word is inline PIC/XPIC payload. +pub(crate) const SEND_PIC_OPCODE: u32 = 0x7010_0000; +pub(crate) const SEND_PICP_OPCODE: u32 = 0xf000_0000; +pub(crate) const SEND_CONTROL_OPCODE_MASK: u32 = 0xf810_0000; +pub(crate) const SEND_PICP_OPCODE_MASK: u32 = 0xf000_0000; +pub(crate) const SEND_COUNT_MASK: u32 = 0x07e0_0000; + +pub(crate) const SYNC_OPCODE: u32 = 0x4180_0000; +const SANS_OPCODE: u32 = 0x40c0_0000; +const BR_M_OPCODE: u32 = 0x4300_0000; +const CALL_M_IMMEDIATE_OPCODE: u32 = 0x1800_0000; +const SETZI_M_OPCODE: u32 = 0x1900_0000; +const PUT_SPECIAL_M_OPCODE: u32 = 0x4300_8000; +const LD32_M_IMMEDIATE_OPCODE: u32 = 0x0100_0000; +const ST32_M_IMMEDIATE_OPCODE: u32 = 0x4f00_0000; +const ADD_M_IMMEDIATE_OPCODE: u32 = 0x2200_0000; +const AND_M_IMMEDIATE_OPCODE: u32 = 0x4200_0000; +const SHL_M_IMMEDIATE_OPCODE: u32 = 0x4200_a000; +const BRZ_M_IMMEDIATE_OPCODE: u32 = 0x1300_0000; + +/// Largest delay encodable by one processor or exchange delay instruction. +pub const MAX_PLAN_OFFSET_CYCLES: u32 = 0x8_0000; + +pub const SANS_INACTIVE_INSTRUCTION: u32 = sans(0); +pub const SYNC_RECEIVE_INSTRUCTION: u32 = sync(0); +pub const SYNC_ANS_INSTRUCTION: u32 = sync(1); +pub const SYNC_SUPERVISOR_INSTRUCTION: u32 = sync(3); +pub const SYNC_ALL_INSTRUCTION: u32 = sync(7); +pub const SYNC_HOST_INSTRUCTION: u32 = sync(15); +pub const RETURN_M10_INSTRUCTION: u32 = br_m(10); + +pub const fn sans(selector: u8) -> u32 { + SANS_OPCODE | selector as u32 +} + +pub const fn sync(selector: u8) -> u32 { + SYNC_OPCODE | selector as u32 +} + +pub const fn br_m(register: u8) -> u32 { + BR_M_OPCODE | ((register as u32) << 20) +} + +pub fn encode_br_m(register: u8) -> Result { + if register >= 16 { + return Err(ExchangeError::Schedule("branch register")); + } + Ok(br_m(register)) +} + +pub fn encode_call_m_immediate( + return_register: u8, + target_address: u32, +) -> Result { + if return_register >= 16 || target_address & 0b11 != 0 || target_address >= 1 << 21 { + return Err(ExchangeError::Schedule("call operand")); + } + Ok(CALL_M_IMMEDIATE_OPCODE | (u32::from(return_register) << 20) | (target_address >> 2)) +} + +pub fn encode_setzi_m(register: u8, immediate: u32) -> Result { + if register >= 16 || immediate >= 1 << 20 { + return Err(ExchangeError::Schedule("setzi operand")); + } + Ok(setzi_m(register, immediate)) +} + +pub fn encode_put_special_m(special: u8, register: u8) -> Result { + if register >= 16 { + return Err(ExchangeError::Schedule("put source register")); + } + Ok(PUT_SPECIAL_M_OPCODE | (u32::from(register) << 20) | u32::from(special)) +} + +pub fn encode_ld32_m_immediate( + destination: u8, + base: u8, + delta: u8, + word_offset: u16, +) -> Result { + if destination >= 16 || base >= 16 || delta >= 16 || word_offset >= 1 << 12 { + return Err(ExchangeError::Schedule("ld32 operand")); + } + Ok(LD32_M_IMMEDIATE_OPCODE + | (u32::from(base) << 20) + | (u32::from(destination) << 16) + | (u32::from(delta) << 12) + | u32::from(word_offset)) +} + +pub fn encode_st32_m_immediate( + source: u8, + base: u8, + delta: u8, + word_offset: u16, +) -> Result { + if source >= 16 || base >= 16 || delta >= 16 || word_offset >= 1 << 12 { + return Err(ExchangeError::Schedule("st32 operand")); + } + Ok(ST32_M_IMMEDIATE_OPCODE + | (u32::from(base) << 20) + | (u32::from(source) << 16) + | (u32::from(delta) << 12) + | u32::from(word_offset)) +} + +pub fn encode_add_m_immediate( + destination: u8, + source: u8, + immediate: i32, +) -> Result { + let immediate = + i16::try_from(immediate).map_err(|_| ExchangeError::Schedule("add immediate operand"))?; + if destination >= 16 || source >= 16 { + return Err(ExchangeError::Schedule("add register operand")); + } + Ok(ADD_M_IMMEDIATE_OPCODE + | (u32::from(source) << 20) + | (u32::from(destination) << 16) + | u32::from(immediate as u16)) +} + +pub fn encode_and_m_immediate( + destination: u8, + source: u8, + immediate: u16, +) -> Result { + if destination >= 16 || source >= 16 || immediate >= 1 << 12 { + return Err(ExchangeError::Schedule("and operand")); + } + Ok(AND_M_IMMEDIATE_OPCODE + | (u32::from(source) << 20) + | (u32::from(destination) << 16) + | u32::from(immediate)) +} + +pub fn encode_shl_m_immediate( + destination: u8, + source: u8, + immediate: u16, +) -> Result { + if destination >= 16 || source >= 16 || immediate >= 1 << 12 { + return Err(ExchangeError::Schedule("shift-left operand")); + } + Ok(SHL_M_IMMEDIATE_OPCODE + | (u32::from(source) << 20) + | (u32::from(destination) << 16) + | u32::from(immediate)) +} + +pub fn encode_brz_m_immediate(register: u8, target_address: u32) -> Result { + if register >= 16 || target_address & 0b11 != 0 || target_address >= 1 << 21 { + return Err(ExchangeError::Schedule("brz operand")); + } + Ok(BRZ_M_IMMEDIATE_OPCODE | (u32::from(register) << 20) | (target_address >> 2)) +} + +/// Encodes a processor delay of `cycles` cycles. +pub fn encode_delay_m(cycles: u32) -> Result { + if !(1..=MAX_PLAN_OFFSET_CYCLES).contains(&cycles) { + return Err(ExchangeError::Schedule("processor delay range")); + } + Ok(delay(cycles - 1)) +} + +pub(crate) const fn setzi_m(register: u8, immediate: u32) -> u32 { + SETZI_M_OPCODE | ((register as u32) << 20) | immediate +} + +pub(crate) const fn put_special_from_m8(register: u8) -> u32 { + PUT_SPECIAL_M_OPCODE | (8 << 20) | register as u32 +} + +pub const fn encode_exchange_delay(cycles: u32) -> u32 { + DELAY_OPCODE | (cycles & 0x7ffff) +} + +pub const fn encode_exchange_delay_pic(a: u32, b: u32, c: u32) -> u32 { + DELAY_PIC_OPCODE | ((a << 19) & 0x03f8_0000) | ((b << 18) & 0x0004_0000) | (c & 0x3ffff) +} + +pub const fn encode_exchange_delay_xpic(a: u32, b: u32, c: u32) -> u32 { + DELAY_XPIC_OPCODE | ((a << 14) & 0x03ff_c000) | ((b << 13) & 0x0000_2000) | (c & 0x1fff) +} + +pub(crate) const fn delay(cycles: u32) -> u32 { + encode_exchange_delay(cycles) +} + +pub(crate) const fn delay_pic(a: u32, b: u32, c: u32) -> u32 { + encode_exchange_delay_pic(a, b, c) +} + +pub(crate) const fn delay_xpic(a: u32, b: u32, c: u32) -> u32 { + encode_exchange_delay_xpic(a, b, c) +} + +pub fn encode_send( + count_minus_one: u32, + direction: u32, + base_word: u32, +) -> Result { + if count_minus_one > 63 || direction > 7 || base_word > 0x3_ffff { + return Err(ExchangeError::Schedule("send instruction operand")); + } + Ok(SEND_OPCODE + | ((count_minus_one << 21) & SEND_COUNT_MASK) + | ((base_word << 3) & SEND_ADDRESS_MASK) + | direction) +} + +pub(crate) const fn send_off(count_minus_one: u32, direction: u32, base_word: u32) -> u32 { + SEND_OFF_OPCODE + | ((count_minus_one << 21) & SEND_COUNT_MASK) + | (((count_minus_one >> 6) << 14) & 0x000f_c000) + | ((base_word << 3) & 0x0000_3ff8) + | (direction & 7) +} + +pub(crate) fn is_send_control(instruction: u32) -> bool { + instruction & SEND_CONTROL_OPCODE_MASK == SEND_PIC_OPCODE +} + +pub(crate) fn is_send_control_pair(instruction: u32) -> bool { + instruction & SEND_PICP_OPCODE_MASK == SEND_PICP_OPCODE +} + +pub(crate) fn is_send_off(instruction: u32) -> bool { + instruction & SEND_CONTROL_OPCODE_MASK == SEND_OFF_OPCODE +} + +pub(crate) fn is_payload_send(instruction: u32) -> bool { + instruction & LONG_OPCODE_MASK == SEND_OPCODE || is_send_off(instruction) +} diff --git a/crates/ipu-target/src/lib.rs b/crates/ipu-target/src/lib.rs new file mode 100644 index 0000000..6102ee7 --- /dev/null +++ b/crates/ipu-target/src/lib.rs @@ -0,0 +1,12 @@ +//! IPU21 machine-level definitions and executable program generation. + +pub mod cost; +pub mod emit; +pub mod exchange; +pub mod hardware; +pub mod instruction; +pub mod memory; +pub mod program; +pub mod topology; + +pub use hardware::{HardwareMemoryConstraints, HardwareTarget}; diff --git a/crates/ipu-target/src/memory.rs b/crates/ipu-target/src/memory.rs new file mode 100644 index 0000000..a74b6cb --- /dev/null +++ b/crates/ipu-target/src/memory.rs @@ -0,0 +1,131 @@ +//! IPU21 tile-memory geometry. + +pub const TILE_MEMORY_BASE: u32 = 0x4c000; +pub const TILE_MEMORY_SIZE: u32 = 624 * 1024; +/// IPU21 `TMEM_ELEMSIZE`. Instruction fetch and data access contend at this +/// granularity even when placement policy is supplied by another crate. +pub const TILE_MEMORY_ELEMENT_SIZE: u32 = 0x4000; +/// Maximum supervisor instruction-fetch lookahead used when checking whether +/// executable and data ranges share a memory element. +pub const IPU21_SUPERVISOR_FETCH_LOOKAHEAD: u32 = 8 * 8; +/// End of IPU21 region 0, the only tile-memory region supporting instruction fetch. +pub const IPU21_EXECUTABLE_MEMORY_LIMIT: u32 = 0x80000; +/// First logical address of the commonly used interleaved operand window. +pub const IPU21_INTERLEAVED_MEMORY_BASE: u32 = TILE_MEMORY_BASE + 0x34000; +/// End of the commonly used interleaved operand window. +pub const IPU21_INTERLEAVED_MEMORY_LIMIT: u32 = TILE_MEMORY_BASE + 0x3c000; +/// End of architectural region 1, whose interleave factor is two on IPU21. +pub const IPU21_INTERLEAVED_REGION_LIMIT: u32 = TILE_MEMORY_BASE + TILE_MEMORY_SIZE; +/// Exclusive end of SRAM which the SDK secondary loader can populate. +/// +/// The final 0x450 bytes are architectural tile memory, but lie beyond the +/// loader's 643 frames of 992 payload bytes starting at `TILE_MEMORY_BASE + +/// 0x10`. Runtime-loaded packages must place no segment there. +pub const IPU21_APPLICATION_MEMORY_LIMIT: u32 = 0xe7bb0; +/// Logical bytes covered by a pair of physical elements in interleaved region 1. +pub const IPU21_INTERLEAVED_ELEMENT_SIZE: u32 = 2 * TILE_MEMORY_ELEMENT_SIZE; + +/// Runtime completion word followed by supervisor and worker stack state. +pub const RUNTIME_STATE_BASE: u32 = crate::hardware::HardwareTarget::Ipu21 + .exchange() + .window_base + + crate::hardware::HardwareTarget::Ipu21 + .exchange() + .window_bytes; +pub const WORKER_STACK_HEADROOM: u32 = 0xe0; +pub const WORKER_SYNC_STRIDE: u32 = 0x100; +pub const WORKER_CONTEXTS: u32 = 6; +pub const RUNTIME_STATE_BYTES: u32 = WORKER_STACK_HEADROOM + WORKER_CONTEXTS * WORKER_SYNC_STRIDE; +pub const PROFILE_START_CYCLE: u32 = RUNTIME_STATE_BASE + 4; +pub const PROFILE_END_CYCLE: u32 = RUNTIME_STATE_BASE + 8; + +/// First byte after permanently reserved runtime state. +pub const IPU21_DATA_BASE: u32 = RUNTIME_STATE_BASE + RUNTIME_STATE_BYTES; +/// Loader-populatable region 1 storage available to interleaved data. +pub const IPU21_INTERLEAVED_REGION_BYTES: u32 = + IPU21_APPLICATION_MEMORY_LIMIT - IPU21_INTERLEAVED_MEMORY_BASE; +/// Standard-addressable storage which is not borrowed from region 1. +pub const IPU21_STANDARD_FIXED_BYTES: u32 = IPU21_INTERLEAVED_MEMORY_BASE - IPU21_DATA_BASE; +/// Total tile SRAM available to planned values after permanent runtime state. +pub const IPU21_PLANNED_DATA_BYTES: u32 = + IPU21_STANDARD_FIXED_BYTES + IPU21_INTERLEAVED_REGION_BYTES; +/// Baseline standard-memory reservation for generated package support data. +pub const IPU21_DEFAULT_SUPPORT_RESERVATION_BYTES: u32 = 3 * TILE_MEMORY_ELEMENT_SIZE; + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)] +pub struct MemoryElement { + pub interleaved: bool, + pub index: u32, +} + +/// Iterates the physical SRAM elements touched by a word-addressed span. +pub fn memory_elements_for_words(address: u32, words: u32) -> MemoryElements { + MemoryElements { + cursor: address, + end: address.saturating_add(words.saturating_mul(4)), + } +} + +#[derive(Clone, Debug)] +pub struct MemoryElements { + cursor: u32, + end: u32, +} + +impl Iterator for MemoryElements { + type Item = MemoryElement; + + fn next(&mut self) -> Option { + if self.cursor >= self.end { + return None; + } + let interleaved = self.cursor >= IPU21_INTERLEAVED_MEMORY_BASE; + let (base, size) = if interleaved { + ( + IPU21_INTERLEAVED_MEMORY_BASE, + IPU21_INTERLEAVED_ELEMENT_SIZE, + ) + } else { + (0, TILE_MEMORY_ELEMENT_SIZE) + }; + let index = (self.cursor - base) / size; + let boundary = base.saturating_add((index + 1).saturating_mul(size)); + self.cursor = boundary.min(self.end); + Some(MemoryElement { interleaved, index }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::BTreeSet; + + #[test] + fn randomized_spans_visit_exactly_their_physical_elements() { + let mut random = fastrand::Rng::with_seed(0x6d65_6d6f_7279_5f65); + for _ in 0..512 { + let address = TILE_MEMORY_BASE + random.u32(0..TILE_MEMORY_SIZE / 4) * 4; + let available_words = (TILE_MEMORY_BASE + TILE_MEMORY_SIZE - address) / 4; + let words = random.u32(0..=available_words.min(16_384)); + let actual = memory_elements_for_words(address, words).collect::>(); + let expected = (0..words) + .map(|word| { + let address = address + word * 4; + if address >= IPU21_INTERLEAVED_MEMORY_BASE { + MemoryElement { + interleaved: true, + index: (address - IPU21_INTERLEAVED_MEMORY_BASE) + / IPU21_INTERLEAVED_ELEMENT_SIZE, + } + } else { + MemoryElement { + interleaved: false, + index: address / TILE_MEMORY_ELEMENT_SIZE, + } + } + }) + .collect::>(); + assert_eq!(actual, expected); + } + } +} diff --git a/crates/ipu-target/src/program.rs b/crates/ipu-target/src/program.rs new file mode 100644 index 0000000..996df4d --- /dev/null +++ b/crates/ipu-target/src/program.rs @@ -0,0 +1,150 @@ +//! Address-resolved programs executed by individual tiles and the host. + +use serde::{Deserialize, Serialize}; + +/// A fully resolved program for one logical tile. +#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct TileProgram { + pub tile: u16, + pub steps: Vec, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub enum TileStep { + Exchange(ExchangeStep), + Compute(ComputeStep), + Repeat(RepeatStep), + Checkpoint(CheckpointStep), +} + +/// A debugger-visible operator boundary using alternating PBRK0/PBRK1 traps. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct CheckpointStep { + pub operation: u32, + pub breakpoint: u8, + #[serde(default)] + pub profile: StepProfile, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct RepeatStep { + pub count: u32, + /// Mutable bases used by [`TileAddress::RepeatPointer`] in the body. + pub iterated_pointers: Vec, + pub body: Vec, + #[serde(default)] + pub profile: StepProfile, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct RepeatPointer { + pub initial_address: u32, + pub stride_bytes: u32, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub enum TileAddress { + Absolute(u32), + /// The current base of an enclosing repeat plus a constant byte offset. + RepeatPointer { + index: u16, + offset: u32, + }, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct ExchangeStep { + /// Whether this tile executes a timed send/receive program after the boundary. + pub active: bool, + /// Base address used by point-to-point receive rows. + pub incoming_base: u32, + /// Preserve both exchange base registers on entry. Absolute-address paired + /// rows use the two PIC streams directly and must not reset their state. + #[serde(default)] + pub preserve_base_registers: bool, + /// Ordinary receive source selected outside the timed row when a paired + /// receive uses the neighbouring sender for its waiting half. + #[serde(default)] + pub incoming_mux: Option, + /// IPU21 incoming item format: 0 for 32-bit, 1 for the early half of a + /// paired 64-bit path, and 2 for the waiting half. + #[serde(default)] + pub incoming_format: u8, + /// Fixed source selection for the borrowed half of a paired 64-bit path. + #[serde(default)] + pub incoming_mux_pair: Option, + /// Override the ordinary internal-exchange down-count. Paired 64-bit + /// helper tiles execute mux timing while using zero to ignore the value. + #[serde(default)] + pub incoming_dcount: Option, + /// The exchange row owns its supervisor sync and does not require the + /// generic down-count setup. This is used by paired-width rows whose SDK + /// form treats the sync and the following timing program as one unit. + #[serde(default)] + pub sync_in_program: bool, + /// Synchronization-free timed exchange program. + pub program: PlacedExchangeRow, + /// Address words applied before invoking a structurally shared row. + #[serde(default)] + pub setup_patch: Option, + /// Words rewritten before the timed program is invoked inside a structured repeat. + #[serde(default)] + pub repeat_patches: Vec, + #[serde(default)] + pub profile: StepProfile, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct ExchangeSetupPatch { + /// Byte offsets into the shared executable row, reused by its structural shape. + pub offsets: PlacedExchangeRow, + /// Replacement instruction words for this use of the row. + pub values: PlacedExchangeRow, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct ExchangePatch { + pub word_offset: u32, + /// Full replacement instruction words, indexed by repeat iteration. + pub values: PlacedExchangeRow, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct ComputeStep { + /// Exact linked kernel symbol; no naming convention is applied. + pub symbol: String, + pub output_address: TileAddress, + pub input_addresses: Vec, + pub arguments: Vec, + #[serde(default)] + pub profile: StepProfile, +} + +/// Optional explicit cycle-counter destinations around a step. +/// +/// The addresses belong to caller-managed tile memory. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct StepProfile { + pub before: Option, + pub after: Option, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct HostPhase { + pub address: u32, + pub active: bool, + pub run_table: Option, +} + +#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct HostProgram { + pub initialize: Vec, + pub inputs: Vec, + pub outputs: Vec, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct PlacedExchangeRow { + pub address: u32, + pub words: Vec, +} diff --git a/crates/ipu-target/src/topology.rs b/crates/ipu-target/src/topology.rs new file mode 100644 index 0000000..0b29212 --- /dev/null +++ b/crates/ipu-target/src/topology.rs @@ -0,0 +1,193 @@ +//! IPU21 logical/physical tile mapping and exchange-fabric geometry. + +use std::collections::HashSet; + +use crate::exchange::ExchangeError; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Topology { + logical_to_physical: Vec, +} + +impl Topology { + pub fn new(logical_to_physical: Vec) -> Result { + let mut physical = HashSet::new(); + if logical_to_physical.is_empty() + || logical_to_physical + .iter() + .any(|tile| !physical.insert(*tile)) + { + return Err(ExchangeError::ReceiverSet); + } + Ok(Self { + logical_to_physical, + }) + } + + pub(crate) fn c600() -> Self { + Self { + logical_to_physical: (0..1472).map(c600_logical_to_physical).collect(), + } + } + + pub fn tile_count(&self) -> usize { + self.logical_to_physical.len() + } + + pub fn prefix(&self, tile_count: u16) -> Result { + Self::new( + self.logical_to_physical + .iter() + .copied() + .take(usize::from(tile_count)) + .collect(), + ) + } + + pub fn physical(&self, logical: u16) -> Result { + self.logical_to_physical + .get(usize::from(logical)) + .copied() + .ok_or(ExchangeError::Tile(logical)) + } + + /// Logical tile that shares this tile's double-width exchange resources. + pub fn paired_logical(&self, logical: u16) -> Result { + let paired_physical = self.physical(logical)? ^ 2; + self.logical_to_physical + .iter() + .position(|physical| *physical == paired_physical) + .map(|paired| u16::try_from(paired).expect("logical tile count fits u16")) + .ok_or(ExchangeError::Tile(logical)) + } + + /// Physical source selected by `INCOMING_MUXPAIR` for a 64-bit send. + pub fn paired_source_mux(&self, sender_logical: u16) -> Result { + Ok(self.physical(sender_logical)? ^ 2) + } + + pub fn is_pair_primary(&self, logical: u16) -> Result { + Ok(self.physical(logical)? & 2 == 0) + } + + /// Direction and width control for one 64-bit route. + pub fn paired_send_control( + &self, + sender_logical: u16, + receiver_logical: u16, + ) -> Result { + let sender = u32::from(self.physical(sender_logical)?); + let receiver = u32::from(self.physical(receiver_logical)?); + Ok(u8::try_from(direction(sender, receiver) | 4).expect("send control is three bits")) + } + + /// Whether this member of a double-width receiving pair owns the paired + /// XPIC source-selection stream. + pub fn paired_receiver_is_early( + &self, + receiver_logical: u16, + sender_logical: u16, + ) -> Result { + let receiver = u32::from(self.physical(receiver_logical)?); + let sender = u32::from(self.physical(sender_logical)?); + let local = time_to_mux(sender, receiver); + let borrowed = time_to_mux(sender, receiver ^ 2); + Ok(local < borrowed) + } +} + +pub(crate) fn c600_logical_to_physical(logical: u16) -> u16 { + let pair = logical / 2; + let lane = logical & 1; + let block = pair / 23; + let mut row = pair % 23; + if block & 1 != 0 { + row = 22 - row; + } + let column = (block / 2) * 4 + (block & 1); + row * 64 + column + lane * 2 +} + +fn route_displacement(source: u32, destination: u32) -> i32 { + let source_raw = ((source >> 2) & 15) as i32; + let destination_raw = ((destination >> 2) & 15) as i32; + let source_column = if source_raw > 7 { + source_raw ^ 15 + } else { + source_raw + }; + let destination_column = if destination_raw > 7 { + destination_raw ^ 15 + } else { + destination_raw + }; + let source_mux = source_column + ((source_raw >> 3) ^ (source & 1) as i32); + let base = (destination_column - source_mux) * 6; + let destination_lane = destination & 3; + let destination_half = destination_raw >> 3; + if destination_lane > 1 { + base + if destination_half == (destination & 1) as i32 { + 2 + } else { + 4 + } + } else { + base + if destination_half == destination_lane as i32 { + 1 + } else { + 5 + } + } +} + +pub(crate) fn direction(source: u32, destination: u32) -> u32 { + if route_displacement(source, destination) < 1 { + 2 + } else { + 1 + } +} + +pub(crate) fn time_to_mux(source: u32, destination: u32) -> i32 { + let source_raw = ((source >> 2) & 15) as i32; + let destination_raw = ((destination >> 2) & 15) as i32; + let source_low = ((source >> 2) & 7) as i32; + let displacement = route_displacement(source, destination); + let source_edge = if source_raw > 7 { + (source_raw * 4) ^ 60 + } else { + source_raw * 4 + }; + let destination_edge = if destination_raw > 7 { + (destination_raw * 4) ^ 60 + } else { + destination_raw * 4 + }; + let local = ((source >> 2) & 8) as i32 | ((source >> 3) & 3) as i32; + let crossing = local - destination_raw + ((source_low >> 1) ^ 3); + let same_region = (source ^ destination) & 0x20 == 0; + let turn = if same_region { + source_low + 1 + } else { + 16 - source_low + }; + let group_delta = (((source >> 6) & 31) as i32 - ((destination >> 6) & 31) as i32) * 2; + crossing + source_edge + turn - destination_edge + group_delta + displacement.abs() - 34 +} + +pub(crate) fn paired_time_to_mux(source: u32, destination: u32) -> i32 { + time_to_mux(source, destination).max(time_to_mux(source, destination ^ 2)) +} + +#[cfg(test)] +mod tests { + #[test] + fn c600_mapping_is_a_permutation() { + let topology = crate::hardware::HardwareTarget::Ipu21.topology(); + let mut physical = (0..topology.tile_count()) + .map(|logical| topology.physical(logical as u16).unwrap()) + .collect::>(); + physical.sort_unstable(); + assert_eq!(physical, (0..1472).collect::>()); + } +} diff --git a/crates/ipu-tests/Cargo.toml b/crates/ipu-tests/Cargo.toml index f64d776..72469a5 100644 --- a/crates/ipu-tests/Cargo.toml +++ b/crates/ipu-tests/Cargo.toml @@ -8,10 +8,6 @@ license.workspace = true name = "ipu-trivial-test" path = "src/main.rs" -[[bin]] -name = "ipu-exchange-schedule-bench" -path = "src/bin/exchange_schedule_bench.rs" - [[bin]] name = "ipu-exchange-live-state" path = "src/bin/exchange_live_state.rs" @@ -23,7 +19,7 @@ dotenvy.workspace = true ipu-codegen = { path = "../ipu-codegen" } ipu-driver = { path = "../ipu-driver" } ipu-elf = { path = "../ipu-elf" } -ipu-exchange = { path = "../ipu-exchange" } +ipu-target = { path = "../ipu-target" } ipu-package = { path = "../ipu-package" } ipu-runtime = { path = "../ipu-runtime" } fastrand.workspace = true diff --git a/crates/ipu-tests/src/bin/exchange_live_state.rs b/crates/ipu-tests/src/bin/exchange_live_state.rs index b613876..1391eca 100644 --- a/crates/ipu-tests/src/bin/exchange_live_state.rs +++ b/crates/ipu-tests/src/bin/exchange_live_state.rs @@ -1,7 +1,6 @@ use anyhow::{Context, Result}; use clap::Parser; use ipu_driver::Device; -use ipu_exchange::Topology; #[derive(Parser)] #[command(about = "Read live IPU21 exchange state without resetting the device")] @@ -20,7 +19,7 @@ struct Arguments { fn main() -> Result<()> { let arguments = Arguments::parse(); let device = Device::open(&arguments.device)?; - let topology = Topology::c600(); + let topology = ipu_target::hardware::HardwareTarget::Ipu21.topology(); let requested = (!arguments.tile.is_empty()).then_some(arguments.tile.as_slice()); for logical in 0..u16::try_from(topology.tile_count())? { if requested.is_some_and(|tiles| !tiles.contains(&logical)) { diff --git a/crates/ipu-tests/src/bin/exchange_schedule_bench.rs b/crates/ipu-tests/src/bin/exchange_schedule_bench.rs deleted file mode 100644 index a3f0a52..0000000 --- a/crates/ipu-tests/src/bin/exchange_schedule_bench.rs +++ /dev/null @@ -1,190 +0,0 @@ -use anyhow::{Context, Result, bail}; -use clap::Parser; -use ipu_codegen::{ - ExchangeScheduleSnapshot, schedule_exchange_problem, validate_exchange_schedule, -}; -use ipu_exchange::diagnostic::diagnose_plan_program; -use std::collections::BTreeSet; -use std::fs::File; -use std::hint::black_box; -use std::io::BufReader; -use std::path::PathBuf; -use std::time::{Duration, Instant}; - -#[derive(Parser)] -#[command( - version, - about = "Replay and validate production exchange scheduling without IPU hardware" -)] -struct Arguments { - /// JSON snapshot written by ipu-trivial-test --export-exchange-schedule. - snapshot: PathBuf, - /// Restrict the benchmark to these physical exchange phase IDs. - #[arg(long = "phase")] - phases: Vec, - /// Untimed scheduler/codegen runs before measurement. - #[arg(long, default_value_t = 0)] - warmup: usize, - /// Timed scheduler/codegen runs per selected phase. - #[arg(long, default_value_t = 1)] - iterations: usize, - /// Ignore sender addresses after the first Repeat iteration. This - /// reproduces the unsafe scheduler behavior used before Repeat-aware - /// memory-element hazard checking. - #[arg(long)] - first_iteration_only: bool, - /// Decode the generated exchange program for this logical tile. - #[arg(long)] - dump_tile: Option, -} - -fn main() -> Result<()> { - let arguments = Arguments::parse(); - if arguments.iterations == 0 { - bail!("--iterations must be nonzero"); - } - let input = File::open(&arguments.snapshot) - .with_context(|| format!("open {}", arguments.snapshot.display()))?; - let mut snapshot: ExchangeScheduleSnapshot = serde_json::from_reader(BufReader::new(input)) - .with_context(|| format!("parse {}", arguments.snapshot.display()))?; - snapshot.validate()?; - if arguments.first_iteration_only { - for transfer in snapshot - .phases - .iter_mut() - .flat_map(|phase| &mut phase.transfers) - { - transfer.source_addresses.truncate(1); - } - } - - let selected = arguments.phases.iter().copied().collect::>(); - if selected.len() != arguments.phases.len() { - bail!("--phase contains a duplicate phase ID"); - } - for &phase in &selected { - if !snapshot.phases.iter().any(|problem| problem.phase == phase) { - bail!("snapshot does not contain phase {phase}"); - } - } - - let problems = snapshot - .phases - .iter() - .filter(|problem| selected.is_empty() || selected.contains(&problem.phase)) - .collect::>(); - println!( - "snapshot={} tiles={} phases={} warmup={} iterations={} repeatAware={}", - arguments.snapshot.display(), - snapshot.tile_count, - problems.len(), - arguments.warmup, - arguments.iterations, - !arguments.first_iteration_only, - ); - - let total_start = Instant::now(); - for problem in problems { - for _ in 0..arguments.warmup { - let run = black_box(schedule_exchange_problem(snapshot.tile_count, problem)?); - validate_exchange_schedule(snapshot.tile_count, problem, &run.phase)?; - } - let mut baseline = None; - let mut durations = Vec::with_capacity(arguments.iterations); - let mut validation_durations = Vec::with_capacity(arguments.iterations); - for _ in 0..arguments.iterations { - let start = Instant::now(); - let run = black_box(schedule_exchange_problem(snapshot.tile_count, problem)?); - durations.push(start.elapsed()); - let validation_start = Instant::now(); - validate_exchange_schedule(snapshot.tile_count, problem, &run.phase)?; - validation_durations.push(validation_start.elapsed()); - if let Some(expected) = &baseline { - if &run.phase != expected { - bail!( - "phase {} scheduler/codegen output changed between identical runs", - problem.phase - ); - } - } else { - baseline = Some(run.phase.clone()); - } - if durations.len() == 1 - && let Some(tile) = arguments.dump_tile - { - let words = run - .phase - .programs - .get(tile) - .with_context(|| format!("logical tile {tile} is out of range"))?; - let activities = run - .phase - .activities - .get(tile) - .with_context(|| format!("logical tile {tile} is out of range"))?; - for activity in activities { - println!( - "phase={} logicalTile={} transfer={} {:?} cycles={}..{} memoryEnd={} address=0x{:x} words={}", - problem.phase, - tile, - activity.transfer, - activity.kind, - activity.start_cycle, - activity.end_cycle, - activity.memory_end_cycle, - activity.address, - activity.words, - ); - } - println!( - "phase={} logicalTile={}\n{}", - problem.phase, - tile, - diagnose_plan_program(words, None)?.render() - ); - } - let destination_count = problem - .transfers - .iter() - .map(|transfer| transfer.destinations.len()) - .sum::(); - let row_words = run.phase.programs.iter().map(Vec::len).sum::(); - let maximum_row_words = run.phase.programs.iter().map(Vec::len).max().unwrap_or(0); - if durations.len() == arguments.iterations { - durations.sort_unstable(); - validation_durations.sort_unstable(); - println!( - "phase={} transfers={} destinations={} initialHorizonCycles={} horizonCycles={} endpointLowerBoundCycles={} lowerBoundGapCycles={} neighborhoodImprovements={} rowWords={} maximumRowWords={} scheduleCodegenMinMs={:.3} scheduleCodegenMedianMs={:.3} scheduleCodegenP95Ms={:.3} scheduleCodegenMaxMs={:.3} validationMedianMs={:.3} invariants=PASS", - problem.phase, - problem.transfers.len(), - destination_count, - run.initial_horizon, - run.phase.event_cycles, - run.endpoint_lower_bound, - run.phase - .event_cycles - .saturating_sub(run.endpoint_lower_bound), - run.neighborhood_improvements, - row_words, - maximum_row_words, - milliseconds(durations[0]), - milliseconds(percentile(&durations, 50)), - milliseconds(percentile(&durations, 95)), - milliseconds(*durations.last().expect("iterations is nonzero")), - milliseconds(percentile(&validation_durations, 50)), - ); - } - } - } - println!("totalMs={:.3}", milliseconds(total_start.elapsed())); - Ok(()) -} - -fn percentile(samples: &[Duration], percentile: usize) -> Duration { - let rank = (samples.len() * percentile).div_ceil(100); - samples[rank.clamp(1, samples.len()) - 1] -} - -fn milliseconds(duration: Duration) -> f64 { - duration.as_secs_f64() * 1_000.0 -} diff --git a/crates/ipu-tests/src/diagnostic.rs b/crates/ipu-tests/src/diagnostic.rs index 6281820..f093162 100644 --- a/crates/ipu-tests/src/diagnostic.rs +++ b/crates/ipu-tests/src/diagnostic.rs @@ -1,8 +1,8 @@ use anyhow::{Context, Result, bail}; use ipu_codegen::{ - AttentionScale, ComputeGraph, DiagnosticPackage, DiagnosticTensor, GemmOptions, Operation, - OperationKind, Precision, Region, Repeat, ShardExtent, ShardView, ValueId, - logical_view_byte_spans, + AttentionScale, CompiledPackage, CompiledTensor, CompiledTensorShard, ComputeGraph, + GemmOptions, Operation, OperationKind, Precision, Region, Repeat, ShardExtent, ShardView, + ValueId, logical_view_byte_spans, }; use ipu_driver::{Device, DriverError, TileException}; use ipu_package::{Application, Binding}; @@ -28,7 +28,7 @@ pub(crate) type PreparedInputs = (BTreeMap, Vec, Vec, next: &mut usize, waiting_for_resume: &mut bool, @@ -200,7 +200,7 @@ fn service_checkpoint( fn compare_tensor( device: &Device, - tensor: &DiagnosticTensor, + tensor: &CompiledTensor, expected: &HostTensor, sample_limit: usize, atol: f32, @@ -267,8 +267,8 @@ fn sample_indices(total: usize, limit: usize) -> BTreeSet { } pub(crate) fn shard_elements( - tensor: &DiagnosticTensor, - shard: &ipu_codegen::DiagnosticShard, + tensor: &CompiledTensor, + shard: &CompiledTensorShard, ) -> Result> { let logical_extents = shard .storage @@ -281,7 +281,7 @@ pub(crate) fn shard_elements( .collect::>(); let view = ShardView { shard: shard.storage.id, - extents: logical_extents.clone(), + extents: logical_extents.clone().into(), }; let element_bytes = u32::try_from(tensor.precision.bytes())?; let offsets = logical_view_byte_spans(&shard.storage, &view)? @@ -324,7 +324,7 @@ fn decode_word(word: u32, byte: u32, precision: Precision) -> Result { pub(crate) fn prepare_inputs( graph: &ComputeGraph, application: &Application, - metadata: &[DiagnosticTensor], + metadata: &[CompiledTensor], ) -> Result { let mut values = BTreeMap::new(); for input in graph.inputs() { @@ -860,7 +860,7 @@ fn quantize(value: f32, precision: Precision) -> f32 { #[cfg(test)] mod tests { use super::*; - use ipu_codegen::{AmpOrder, amp_matrix_coordinates}; + use ipu_codegen::{NativeKernelOrder, amp_matrix_coordinates}; #[test] fn randomized_blas_gemm_matches_scalar_reference() -> Result<()> { @@ -933,14 +933,14 @@ mod tests { let source = base + source_word as u32; let destination = base + destination_word; let source_coordinates = amp_matrix_coordinates( - AmpOrder::Output, + NativeKernelOrder::Output, Precision::F16, rows, columns, source * 2 + lane, )?; let destination_coordinates = amp_matrix_coordinates( - AmpOrder::Left, + NativeKernelOrder::Left, Precision::F16, rows, columns, * Unmerged path crates/ipu-tests/src/exchange_stress.rs * Unmerged path crates/ipu-tests/src/main.rs * Unmerged path device/attention_softmax_f16_wrapper.S * Unmerged path device/attention_stages_f16.S diff --git a/device/flash_attention_online_f16.cpp b/device/flash_attention_online_f16.cpp deleted file mode 100644 index 87a133b..0000000 --- a/device/flash_attention_online_f16.cpp +++ /dev/null @@ -1,144 +0,0 @@ -#include -#include - -#ifndef ATTENTION_VERTEX_NAME -#define ATTENTION_VERTEX_NAME FlashAttentionOnlineF16 -#endif -#ifndef ATTENTION_MATRICES -#define ATTENTION_MATRICES 1 -#endif -#ifndef ATTENTION_QUERY_ROWS -#define ATTENTION_QUERY_ROWS 16 -#endif -#ifndef ATTENTION_KEY_ROWS -#define ATTENTION_KEY_ROWS ATTENTION_QUERY_ROWS -#endif -#ifndef ATTENTION_QUERY_DIMENSION -#define ATTENTION_QUERY_DIMENSION 64 -#endif -#ifndef ATTENTION_VALUE_DIMENSION -#define ATTENTION_VALUE_DIMENSION ATTENTION_QUERY_DIMENSION -#endif -#ifndef ATTENTION_SCALE -#define ATTENTION_SCALE (1.0f / __builtin_sqrtf(float(ATTENTION_QUERY_DIMENSION))) -#endif -#ifndef ATTENTION_KEY_BLOCK_ROWS -#define ATTENTION_KEY_BLOCK_ROWS 32 -#endif - -using namespace poplar; - -static_assert(ATTENTION_MATRICES > 0); -static_assert(ATTENTION_QUERY_ROWS > 0); -static_assert(ATTENTION_KEY_ROWS > 0); -static_assert(ATTENTION_QUERY_DIMENSION > 0); -static_assert(ATTENTION_VALUE_DIMENSION > 0); -static_assert(ATTENTION_KEY_BLOCK_ROWS > 0); - -static __attribute__((always_inline)) float attentionDot(const half *query, - const half *key) { - constexpr unsigned dimension = ATTENTION_QUERY_DIMENSION; - float score = 0.0f; -#if ATTENTION_QUERY_DIMENSION % 2 == 0 - for (unsigned column = 0; column < dimension; column += 2) { - const half2 packedQuery = - *reinterpret_cast(&query[column]); - const half2 packedKey = *reinterpret_cast(&key[column]); - const float2 queryPair = __builtin_convertvector(packedQuery, float2); - const float2 keyPair = __builtin_convertvector(packedKey, float2); - score += queryPair[0] * keyPair[0] + queryPair[1] * keyPair[1]; - } -#else - for (unsigned column = 0; column < dimension; ++column) - score += float(query[column]) * float(key[column]); -#endif - return score; -} - -// Exact, non-causal online softmax attention. Each worker owns complete query -// rows, so the running maximum, denominator, and output vector never need to -// be synchronized. The kernel materializes no QK matrix. -class ATTENTION_VERTEX_NAME : public MultiVertex { -public: - Input> query; - Input> key; - Input> value; - Output> output; - - bool compute(unsigned worker) { - constexpr unsigned matrices = ATTENTION_MATRICES; - constexpr unsigned queryRows = ATTENTION_QUERY_ROWS; - constexpr unsigned keyRows = ATTENTION_KEY_ROWS; - constexpr unsigned queryDimension = ATTENTION_QUERY_DIMENSION; - constexpr unsigned valueDimension = ATTENTION_VALUE_DIMENSION; - constexpr float scale = ATTENTION_SCALE; - - for (unsigned flatQuery = worker; flatQuery < matrices * queryRows; - flatQuery += 6) { - const unsigned matrix = flatQuery / queryRows; - const unsigned queryRow = flatQuery % queryRows; - const half *queryVector = - &query[(matrix * queryRows + queryRow) * queryDimension]; - const half *keys = &key[matrix * keyRows * queryDimension]; - const half *values = &value[matrix * keyRows * valueDimension]; - volatile float *destination = - &output[(matrix * queryRows + queryRow) * valueDimension]; - - for (unsigned column = 0; column < valueDimension; ++column) - destination[column] = 0.0f; - - float maximum = -__builtin_inff(); - float denominator = 0.0f; - for (unsigned keyStart = 0; keyStart < keyRows; - keyStart += ATTENTION_KEY_BLOCK_ROWS) { - const unsigned remainingRows = keyRows - keyStart; - const unsigned blockRows = remainingRows < ATTENTION_KEY_BLOCK_ROWS - ? remainingRows - : ATTENTION_KEY_BLOCK_ROWS; - alignas(8) float scores[ATTENTION_KEY_BLOCK_ROWS]; - float blockMaximum = -__builtin_inff(); - for (unsigned blockRow = 0; blockRow < blockRows; ++blockRow) { - const unsigned keyRow = keyStart + blockRow; - scores[blockRow] = - attentionDot(queryVector, &keys[keyRow * queryDimension]) * scale; - blockMaximum = __builtin_fmaxf(blockMaximum, scores[blockRow]); - } - - const float nextMaximum = __builtin_fmaxf(maximum, blockMaximum); - const float previousScale = maximum == -__builtin_inff() - ? 0.0f - : __builtin_expf(maximum - nextMaximum); - denominator *= previousScale; - for (unsigned column = 0; column < valueDimension; ++column) - destination[column] *= previousScale; - - for (unsigned blockRow = 0; blockRow < blockRows; ++blockRow) { - const unsigned keyRow = keyStart + blockRow; - const float weight = __builtin_expf(scores[blockRow] - nextMaximum); - denominator += weight; - const half *valueVector = &values[keyRow * valueDimension]; -#if ATTENTION_VALUE_DIMENSION % 2 == 0 - for (unsigned column = 0; column < valueDimension; column += 2) { - const half2 packedValue = - *reinterpret_cast(&valueVector[column]); - const float2 valuePair = - __builtin_convertvector(packedValue, float2); - volatile float2 *destinationPair = - reinterpret_cast(&destination[column]); - *destinationPair = *destinationPair + valuePair * weight; - } -#else - for (unsigned column = 0; column < valueDimension; ++column) - destination[column] += weight * float(valueVector[column]); -#endif - } - maximum = nextMaximum; - } - - const float reciprocal = 1.0f / denominator; - for (unsigned column = 0; column < valueDimension; ++column) - destination[column] *= reciprocal; - } - return true; - } -}; diff --git a/device/gemm_f32_64_amp.S b/device/gemm_f32_64_amp.S index 3147ed5..26c8235 100644 --- a/device/gemm_f32_64_amp.S +++ b/device/gemm_f32_64_amp.S @@ -8,6 +8,18 @@ #endif #ifndef GEMM_LARGE_ROWS #define GEMM_LARGE_ROWS GEMM_SMALL_ROWS +#endif +#ifndef GEMM_INIT_SMALL_SYMBOL +#define GEMM_INIT_SMALL_SYMBOL ipu_stack_gemm_f32_init_small_rows +#endif +#ifndef GEMM_INIT_LARGE_SYMBOL +#define GEMM_INIT_LARGE_SYMBOL ipu_stack_gemm_f32_init_large_rows +#endif +#ifndef GEMM_ACCUMULATE_SMALL_SYMBOL +#define GEMM_ACCUMULATE_SMALL_SYMBOL ipu_stack_gemm_f32_accumulate_small_rows +#endif +#ifndef GEMM_ACCUMULATE_LARGE_SYMBOL +#define GEMM_ACCUMULATE_LARGE_SYMBOL ipu_stack_gemm_f32_accumulate_large_rows #endif .text @@ -197,10 +209,15 @@ .size \name, .-\name .endm - GEMM_SPECIALIZATION ipu_stack_gemm_f32_init_small_rows, GEMM_SMALL_ROWS, ipu_stack_gemm_f32_init_common - GEMM_SPECIALIZATION ipu_stack_gemm_f32_init_large_rows, GEMM_LARGE_ROWS, ipu_stack_gemm_f32_init_common - GEMM_SPECIALIZATION ipu_stack_gemm_f32_accumulate_small_rows, GEMM_SMALL_ROWS, ipu_stack_gemm_f32_accumulate_common - GEMM_SPECIALIZATION ipu_stack_gemm_f32_accumulate_large_rows, GEMM_LARGE_ROWS, ipu_stack_gemm_f32_accumulate_common +#ifdef GEMM_SINGLE_ROWS + GEMM_SPECIALIZATION GEMM_INIT_SMALL_SYMBOL, GEMM_SMALL_ROWS, ipu_stack_gemm_f32_init_common + GEMM_SPECIALIZATION GEMM_ACCUMULATE_SMALL_SYMBOL, GEMM_SMALL_ROWS, ipu_stack_gemm_f32_accumulate_common +#else + GEMM_SPECIALIZATION GEMM_INIT_SMALL_SYMBOL, GEMM_SMALL_ROWS, ipu_stack_gemm_f32_init_common + GEMM_SPECIALIZATION GEMM_INIT_LARGE_SYMBOL, GEMM_LARGE_ROWS, ipu_stack_gemm_f32_init_common + GEMM_SPECIALIZATION GEMM_ACCUMULATE_SMALL_SYMBOL, GEMM_SMALL_ROWS, ipu_stack_gemm_f32_accumulate_common + GEMM_SPECIALIZATION GEMM_ACCUMULATE_LARGE_SYMBOL, GEMM_LARGE_ROWS, ipu_stack_gemm_f32_accumulate_common +#endif .worker .p2align 3 diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index aaeae94..da1b47a 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -4,9 +4,10 @@ See [Compiler data flow](COMPILER_DATA_FLOW.md) for the selection, costing and e The package path has four explicit components: -1. `ipu-exchange` produces exchange rows. -2. `ipu-codegen` lowers a `ComputeGraph`, emits supervisor code, and coordinates - package construction according to `PackageConfig`. +1. `ipu-target` defines IPU21 topology and SRAM geometry, generates and parses + exchange rows, encodes instructions, and emits finalized tile programs. +2. `ipu-codegen` lowers a `ComputeGraph` and coordinates package construction + according to `PackageConfig`. 3. `ipu-elf` compiles and links the static runtime and selected kernels. 4. `ipu-package` stores final tile images and host protocol metadata for `ipu-driver` and `ipu-runtime`. @@ -105,9 +106,10 @@ A tile program is an ordered list of: - an exchange row and its final address; or - a kernel symbol, output address, input addresses, and scalar arguments. -The code generator validates only local encoding constraints. It does not check +`ipu_target::emit` validates only local encoding constraints. It does not check lifetimes, search memory, merge repeated regions, repack executable objects, or -derive kernel memory requirements. +derive kernel memory requirements. `ipu-codegen` owns those planning and +placement decisions and supplies a fully resolved `ipu_target::program`. Optional cycle samples name explicit destination addresses. This is a narrow mechanism rather than a profiling layout policy. diff --git a/docs/EXCHANGE_INSTRUCTION_REFERENCE.md b/docs/EXCHANGE_INSTRUCTION_REFERENCE.md index 433d5eb..310bef1 100644 --- a/docs/EXCHANGE_INSTRUCTION_REFERENCE.md +++ b/docs/EXCHANGE_INSTRUCTION_REFERENCE.md @@ -5,7 +5,7 @@ It separates facts verified against the SDK assembler/disassembler and emitted Poplar programs from timing interpretations inferred from those programs. The executable decoder and consistency checker are in -`ipu_exchange::diagnostic`. Every row produced by `PhaseProgramBuilder` is +`ipu_target::diagnostic`. Every row produced by `PhaseProgramBuilder` is decoded and checked against its scheduled send intervals and incoming-control events before codegen can use it. @@ -177,7 +177,7 @@ The encodings above are checked in three independent ways: fields. 2. Small Poplar copy graphs supply complete SDK-generated rows, including the inline payload and scheduling choices. -3. `ipu_exchange::diagnostic` decodes ipu-stack rows and checks them against the +3. `ipu_target::diagnostic` decodes ipu-stack rows and checks them against the phase builder's declarative send/control schedule. For an ipu-stack failure, inspect in this order: diff --git a/docs/PATCHED_BREAKPOINT_RE_FINDINGS.md b/docs/PATCHED_BREAKPOINT_RE_FINDINGS.md index 98a2405..37a5b10 100644 --- a/docs/PATCHED_BREAKPOINT_RE_FINDINGS.md +++ b/docs/PATCHED_BREAKPOINT_RE_FINDINGS.md @@ -30,7 +30,7 @@ put CSR[index], $m1 = 0x43008100 | index It is not `index << 4`. Any experiment using `0x43008730` to address CSR `0x73` wrote a different register. -The previous `ipu-exchange-re` note calling CSR `0x73` `DBG_ECSR` is not +The previous `ipu-target-re` note calling CSR `0x73` `DBG_ECSR` is not conclusive. Its actual diagnostic names the index `debugExceptionControl`, ORs bit zero into it, and enables a separately configured IBRK. That proves the IBRK registers at `0x80/0x81`, but does not distinguish `DBG_CTL.CHAN_EN` from diff --git a/docs/PLANNER_SIMPLIFICATION_AUDIT.md b/docs/PLANNER_SIMPLIFICATION_AUDIT.md index b75ad6d..078995b 100644 --- a/docs/PLANNER_SIMPLIFICATION_AUDIT.md +++ b/docs/PLANNER_SIMPLIFICATION_AUDIT.md @@ -191,18 +191,20 @@ Physical element order remains a separate concern: `ResolvedLayout` describes which physical tensor elements each tile owns, while `storage` maps a resolved view to byte spans for row-major, block-major, and AMP encodings. -### Separate implementation families from the search domain +### Make plan generation operator-directed -The operator-candidate list currently combines kernel availability, precision, -active tile counts, layouts, GEMM grids, memory classes, and staging policy. -`plans()` then adds separate shape-dependent SplitHeads, attention, pointwise, -GEMM, and parameter-storage variants. +Each semantic `OperationKind` should invoke its own typed plan generator. GEMM +generation owns GEMM grids, blocking, parameter placement, and reduction +staging; pointwise generation propagates compatible input layouts and offers +useful ownership transitions; attention generation owns its materialized and +blocked algorithms. A flat cross-operator implementation catalogue obscures +these dependencies and still requires a second dispatch layer. -Replace this with a small internal implementation catalogue and an explicit -search domain. Active tile counts, allowed precisions, memory classes, and -diagnostic restrictions should be planner inputs, not duplicated seed -candidates. Layout-transparent pointwise implementations should be described -once rather than once per tile count. +The generators share an explicit search domain for genuinely global choices: +active tile counts, permitted precisions, weight memory classes, attention +strategy, and diagnostic restrictions. These values are planner inputs rather +than placeholder plans. Shape-dependent layouts and dispatches are emitted +only by the generator for the corresponding high-level operator. ### Normalize plan representations @@ -305,7 +307,8 @@ tradeoffs and should be unified rather than removed: ## Suggested order 1. Add canonical layout resolution and migrate every layout consumer. -2. Remove dead configuration and narrow the public planner API. +2. Factor out the shared search domain and make plan generation + operator-directed. 3. Consolidate candidate and selected-plan representations. 4. Replace deferred SplitHeads machinery with generic views. 5. Normalize GEMM plans and their lowering. diff --git a/docs/REPRESENTATION_DUPLICATION_AUDIT.md b/docs/REPRESENTATION_DUPLICATION_AUDIT.md new file mode 100644 index 0000000..c9926d1 --- /dev/null +++ b/docs/REPRESENTATION_DUPLICATION_AUDIT.md @@ -0,0 +1,102 @@ +# Representation boundaries + +This document records the canonical representations shared across planning, +lowering, packaging, profiling, and diagnostics. A decision is stored once; +later stages either refer to it or derive a representation with additional +resolution. + +## Compiler plans + +- `OperatorPlan` is the selected whole-device operator plan. A + `MidOperationKind::Operator` stores it directly, so semantic operator, + dispatch, operand requirements, deferred inputs, and deferred output have + one owner. +- `GemmGeometry` owns the block shape, orientation, final result grid, grid + order, and distribution. `ParallelReductionPlan` adds only its compute grid + and staging policy. `GemmPlanConstraint` refers to the same geometry type. +- `GemmKernelFamily` contains only shape-independent tile-kernel choices. + Concrete `TileKernelSpec::Gemm` calls are derived from the family and the + canonical block shape. +- `AttentionPlan` owns its algorithm, blocking, padding, and GEMM kernel + family. Attention GEMM calls and `AttentionBufferShape` are derived from it. + `AttentionTask` and prepared panels are transient lowering state. +- `AllocationRequirements` is merged and consumed by both estimation and + placement. Alignment, access tails, and memory-element separation are not + re-expressed in estimator- or allocator-private schemas. + +## Cost and memory + +- `CostEstimate` contains total cycles, exchange cycles, and + `ExchangeFootprint`. +- `PlanMetrics` pairs a cost with the relevant memory record. + `OperationMetrics` and `RegionMetrics` are aliases over this type. +- Pareto retention uses `RegionMetrics::dominates`; compatibility classes may + restrict which candidates are comparable without defining another metric + vocabulary. +- Deferred materialization stores its complete `CostEstimate`, so restoring a + deferred operation cannot lose exchange cost or row-footprint information. + +## Tensor and memory regions + +- `TensorRegion` is a rectangular semantic or padded tensor region composed of + axis-labelled `ShardExtent`s. Layout intersections, replica identities, + deferred slices, and lowering views use it directly. +- `AddressRegion` is a half-open tile-SRAM interval. Placement arenas, package + maps, and host-planning free ranges use this physical type. +- `ByteSpan` remains an offset-plus-length description inside one storage + object. It is not an address interval or a semantic tensor region. +- Cycle intervals remain cycle pairs; they are not memory regions. + +## Exchange + +- `PhysicalTransfer` and `TransferEndpoint` are the address-resolved exchange + core. Schedule snapshots serialize `PhysicalTransfer` directly. +- Codegen `PendingTransfer` wraps the physical transfer only with low-IR + provenance, source-slice information, and paired-resource reservations. +- `ResolvedTransfer` adds topology-derived point-to-point encoding. Encoded + exchange rows are the next resolution boundary. +- `PhaseTransferTiming` is the encoder's detailed timing result. The scheduler + retains only dependency-chain summaries needed by its search; profile + activities and emitted diagnostics are observations of the selected + schedule. +- Exchange replay uses the captured physical transfer list as its source of + truth and verifies activity metadata against it. Stress records attach + expected payloads and requested timing to the same physical transfer. + +## Packages and profiles + +- `CompiledPackage` is the result of ordinary and diagnostic compilation. + Diagnostic builds populate its checkpoint list rather than constructing a + parallel package result. +- `CompiledTensor` owns logical tensor metadata and placed shards. Host + `Binding`s and checkpoint tensor descriptions are derived from it. +- `ProfileStepKind` and `ExchangeActivityKind` are defined by `ipu-package` and + used directly by codegen, profile queries, tests, and the CLI. +- Application profile plans and standalone profile reports share the Cap'n + Proto `profile_common.capnp` step schema and the same Rust reader/writer. + `TileProfilePlan` and `CycleSample` remain distinct containers because one + is a static instrumentation plan and the other is a measured interval. + +## Configuration + +- `HardwareTarget` remains the dispatch point for target cost models and + memory constraints, even while IPU21 is the only implemented target. +- `ProfilingConfig` is the explicit `Disabled`, `Overall`, or `Full` policy. + It is not a Boolean wrapper. + +## Stage boundaries + +The following identities intentionally remain distinct: + +- `Operation`, `MidOperation`, low tile work, and address-resolved `TileStep`; +- `ValueId`, `MidValueId`, and `LowShardId`; +- logical exchanges, physical transfers, resolved transfers, and encoded rows; +- graph, mid-level, low-level, and finalized structured-repeat records; +- symbolic tensor regions, storage byte spans, and physical address regions; +- `LinkedSegment`, which borrows a linked image, and package `Segment`, which + owns serialized bytes and permissions; and +- internal device exchange transfers and host-exchange protocol transfers. + +These boundaries add placement, topology, ownership, serialization, or +measurement semantics. Shared nested value types cross them whenever the +underlying decision is unchanged. * Unmerged path schemas/application.capnp * Unmerged path schemas/profile.capnp diff --git a/schemas/profile_common.capnp b/schemas/profile_common.capnp new file mode 100644 index 0000000..04bd7f1 --- /dev/null +++ b/schemas/profile_common.capnp @@ -0,0 +1,37 @@ +@0xf3c4b874bcfe20c9; + +enum StepKind { + exchange @0; + compute @1; + synchronization @2; + idle @3; +} + +enum ExchangeActivityKind { + send @0; + receive @1; + partnerBusy @2; +} + +struct ExchangeActivity { + kind @0 :ExchangeActivityKind; + startCycle @1 :UInt32; + endCycle @2 :UInt32; +} + +struct Metadata { + name @0 :Text; + value @1 :Text; +} + +struct Step { + localIndex @0 :UInt32; + phase @1 :UInt32; + epoch @2 :UInt32; + operation @3 :Text; + kind @4 :StepKind; + kernel @5 :Text; + metadata @6 :List(Metadata); + exchangeActivities @7 :List(ExchangeActivity); + exchangeEventCycles @8 :UInt32; +} diff --git a/scripts/generate-package-callgraph.py b/scripts/generate-package-callgraph.py new file mode 100644 index 0000000..30ff32e --- /dev/null +++ b/scripts/generate-package-callgraph.py @@ -0,0 +1,1219 @@ +#!/usr/bin/env python3 +"""Generate package build/run call graphs for the ipu-stack workspace. + +The extractor prefers rustc MIR because MIR call terminators resolve associated +functions and trait methods. If a crate does not currently compile, it falls +back to a conservative source-level extractor for that crate. Compiler-made +closure functions are folded into their containing named Rust function. + +The generated full graph contains every reachable workspace-local named +function plus every workspace-local type mentioned by those functions. The +overview is the function-only prefix within four calls of either entry point. +""" + +from __future__ import annotations + +import argparse +import collections +import dataclasses +import html +import os +from pathlib import Path +import re +import subprocess +import sys +from typing import Iterable + + +WORKSPACE_PACKAGES = ( + "ipu-codegen", + "ipu-driver", + "ipu-elf", + "ipu-exchange", + "ipu-package", + "ipu-profile", + "ipu-runtime", +) + +CRATE_COLORS = { + "ipu_codegen": "#d9e8fb", + "ipu_driver": "#f9dfc5", + "ipu_elf": "#eadcf8", + "ipu_exchange": "#d8f0dd", + "ipu_package": "#f7e8ae", + "ipu_profile": "#e2e2e2", + "ipu_runtime": "#f7d8e3", +} + +CONTROL_WORDS = { + "as", + "assert", + "async", + "break", + "const", + "continue", + "drop", + "else", + "fn", + "for", + "if", + "loop", + "match", + "move", + "return", + "sizeof", + "static", + "while", +} + +IGNORED_AMBIGUOUS_TYPES = {"Result", "Error", "IntoIter", "Reader", "Builder"} + + +@dataclasses.dataclass(eq=False) +class RustFunction: + key: str + crate: str + module: str + owner: str | None + name: str + path: Path + line: int + signature: str + body: str + + @property + def display(self) -> str: + owner = f"::{self.owner}" if self.owner else "" + return f"{self.module}{owner}::{self.name}" + + +@dataclasses.dataclass(eq=False) +class RustType: + key: str + crate: str + module: str + name: str + kind: str + path: Path + line: int + + @property + def display(self) -> str: + return f"{self.module}::{self.name}" + + +@dataclasses.dataclass +class Edge: + callers: set[str] = dataclasses.field(default_factory=set) + provenance: set[str] = dataclasses.field(default_factory=set) + + +@dataclasses.dataclass +class GraphPartition: + key: str + title: str + crate: str + functions: set[str] = dataclasses.field(default_factory=set) + types: set[str] = dataclasses.field(default_factory=set) + + @property + def filename(self) -> str: + return f"{self.key}.svg" + + +def crate_name(package: str) -> str: + return package.replace("-", "_") + + +def module_name(root: Path, source: Path, crate: str) -> str: + relative = source.relative_to(root / "crates" / crate.replace("_", "-") / "src") + parts = list(relative.parts) + filename = parts.pop() + stem = Path(filename).stem + if stem not in {"lib", "main", "mod"}: + parts.append(stem) + if parts and parts[0] == "bin": + parts = parts[1:] + return "::".join([crate, *parts]) + + +def sanitize_rust(source: str) -> str: + """Blank comments and literal contents while retaining offsets/newlines.""" + output = list(source) + index = 0 + block_depth = 0 + while index < len(source): + if block_depth: + if source.startswith("/*", index): + output[index : index + 2] = " " + block_depth += 1 + index += 2 + elif source.startswith("*/", index): + output[index : index + 2] = " " + block_depth -= 1 + index += 2 + else: + if source[index] != "\n": + output[index] = " " + index += 1 + continue + if source.startswith("//", index): + end = source.find("\n", index) + if end < 0: + end = len(source) + for cursor in range(index, end): + output[cursor] = " " + index = end + continue + if source.startswith("/*", index): + output[index : index + 2] = " " + block_depth = 1 + index += 2 + continue + if source[index] in {'"', "'"}: + quote = source[index] + # A lifetime such as 'a is not a character literal. + if quote == "'" and index + 1 < len(source) and ( + source[index + 1].isalnum() or source[index + 1] == "_" + ): + # Treat 'a and '_ as lifetimes, but retain ordinary character + # literals such as 'a'. + if index + 2 >= len(source) or source[index + 2] != "'": + index += 1 + continue + cursor = index + 1 + while cursor < len(source): + if source[cursor] == "\\": + if source[cursor] != "\n": + output[cursor] = " " + if cursor + 1 < len(source) and source[cursor + 1] != "\n": + output[cursor + 1] = " " + cursor += 2 + continue + if source[cursor] == quote: + break + if source[cursor] != "\n": + output[cursor] = " " + cursor += 1 + index = min(cursor + 1, len(source)) + continue + index += 1 + return "".join(output) + + +def matching_brace(source: str, opening: int) -> int | None: + depth = 0 + for index in range(opening, len(source)): + if source[index] == "{": + depth += 1 + elif source[index] == "}": + depth -= 1 + if depth == 0: + return index + return None + + +def source_line(source: str, offset: int) -> int: + return source.count("\n", 0, offset) + 1 + + +def impl_owner(header: str) -> str | None: + header = re.sub(r"^\s*impl\s*<[^>{}]*>\s*", "", header.strip()) + header = re.sub(r"^\s*impl\s+", "", header) + header = header.split(" where ", 1)[0].strip() + if " for " in header: + header = header.rsplit(" for ", 1)[1] + match = re.search(r"(?:[A-Za-z_]\w*::)*([A-Za-z_]\w*)", header) + return match.group(1) if match else None + + +def enclosing_owner(impls: list[tuple[int, int, str]], offset: int) -> str | None: + candidates = [(start, owner) for start, end, owner in impls if start < offset < end] + return max(candidates, default=(0, None), key=lambda item: item[0])[1] + + +def discover_source(root: Path) -> tuple[dict[str, RustFunction], dict[str, RustType]]: + functions: dict[str, RustFunction] = {} + types: dict[str, RustType] = {} + for package in WORKSPACE_PACKAGES: + crate = crate_name(package) + source_root = root / "crates" / package / "src" + if not source_root.exists(): + continue + for path in sorted(source_root.rglob("*.rs")): + original = path.read_text(encoding="utf-8") + source = sanitize_rust(original) + module = module_name(root, path, crate) + + impls: list[tuple[int, int, str]] = [] + impl_pattern = re.compile(r"\bimpl(?:\s*<[^>{}]*>)?\s+([^;{]+)\{") + for match in impl_pattern.finditer(source): + opening = source.find("{", match.start(), match.end()) + closing = matching_brace(source, opening) + owner = impl_owner(source[match.start() : opening]) + if closing is not None and owner: + impls.append((opening, closing, owner)) + + type_pattern = re.compile( + r"(?m)^\s*(?:pub(?:\([^)]*\))?\s+)?" + r"(struct|enum|union|trait|type)\s+([A-Za-z_]\w*)" + ) + for match in type_pattern.finditer(source): + kind, name = match.groups() + line = source_line(source, match.start()) + key = f"type:{crate}:{path.relative_to(root)}:{line}:{name}" + types[key] = RustType(key, crate, module, name, kind, path, line) + + function_pattern = re.compile( + r"(?m)^\s*(?:pub(?:\([^)]*\))?\s+)?" + r"(?:(?:async|const|unsafe|extern\s+\"[^\"]+\")\s+)*" + r"fn\s+([A-Za-z_]\w*)\s*(?:<[^;{}]*?>\s*)?\(" + ) + for match in function_pattern.finditer(source): + name = match.group(1) + opening = source.find("{", match.end()) + semicolon = source.find(";", match.end()) + if opening < 0 or (semicolon >= 0 and semicolon < opening): + continue + closing = matching_brace(source, opening) + if closing is None: + continue + owner = enclosing_owner(impls, match.start()) + line = source_line(source, match.start()) + key = f"fn:{crate}:{path.relative_to(root)}:{line}:{name}" + signature = re.sub(r"\s+", " ", original[match.start() : opening].strip()) + functions[key] = RustFunction( + key, + crate, + module, + owner, + name, + path, + line, + signature, + source[opening + 1 : closing], + ) + return functions, types + + +class Resolver: + def __init__(self, functions: dict[str, RustFunction]): + self.functions = functions + self.by_leaf: dict[str, list[str]] = collections.defaultdict(list) + self.by_owner_leaf: dict[str, list[str]] = collections.defaultdict(list) + self.by_display: dict[str, list[str]] = collections.defaultdict(list) + for key, function in functions.items(): + self.by_leaf[function.name].append(key) + self.by_display[function.display].append(key) + if function.owner: + self.by_owner_leaf[f"{function.owner}::{function.name}"].append(key) + + @staticmethod + def strip_generics(target: str) -> str: + result: list[str] = [] + index = 0 + while index < len(target): + if target.startswith("::<", index): + depth = 1 + index += 3 + while index < len(target) and depth: + if target[index] == "<": + depth += 1 + elif target[index] == ">": + depth -= 1 + index += 1 + continue + result.append(target[index]) + index += 1 + return "".join(result) + + @staticmethod + def normalize_target(target: str, caller: RustFunction) -> str: + target = Resolver.strip_generics(target.strip()) + target = re.sub(r"^<([^<> ]+)(?: as [^>]+)?>::", r"\1::", target) + target = target.replace("Self::", f"{caller.owner}::" if caller.owner else "") + target = target.replace("crate::", f"{caller.crate}::") + while target.startswith("super::"): + target = target[7:] + return target + + def resolve(self, target: str, caller: RustFunction, method: bool = False) -> list[str]: + target = self.normalize_target(target, caller) + leaf = target.rsplit("::", 1)[-1] + if leaf in CONTROL_WORDS: + return [] + candidates: list[str] = [] + if "::" in target: + candidates = [ + key + for display, keys in self.by_display.items() + if display == target or display.endswith(f"::{target}") + for key in keys + ] + if not candidates: + candidates = list( + self.by_owner_leaf.get("::".join(target.split("::")[-2:]), []) + ) + else: + candidates = list(self.by_leaf.get(leaf, [])) + if not candidates: + return [] + same_crate = [key for key in candidates if self.functions[key].crate == caller.crate] + same_module = [key for key in same_crate if self.functions[key].module == caller.module] + same_owner = [ + key + for key in same_crate + if caller.owner and self.functions[key].owner == caller.owner + ] + if same_owner: + candidates = same_owner + elif same_module: + candidates = same_module + elif same_crate: + candidates = same_crate + if method and len(candidates) != 1: + return [] + # Direct paths should normally be unique. Retain a bounded over-approximation + # rather than inventing an arbitrary target for duplicate helper names. + return sorted(set(candidates)) if len(candidates) <= 4 else [] + + +def source_calls(function: RustFunction, resolver: Resolver) -> set[str]: + result: set[str] = set() + qualified = re.compile( + r"(?)?\s*\(" + ) + methods = re.compile(r"\.\s*([A-Za-z_]\w*)\s*(?:::\s*<[^;{}()]*>)?\s*\(") + trait_calls = re.compile(r"<\s*([A-Za-z_]\w*)\s+as\s+[^>]+>::([A-Za-z_]\w*)\s*\(") + for match in qualified.finditer(function.body): + result.update(resolver.resolve(match.group(1), function)) + for match in methods.finditer(function.body): + result.update(resolver.resolve(match.group(1), function, method=True)) + for match in trait_calls.finditer(function.body): + result.update(resolver.resolve(f"{match.group(1)}::{match.group(2)}", function)) + result.discard(function.key) + return result + + +def mir_call_target(line: str) -> str | None: + marker = " -> [" + if marker not in line or " = " not in line: + return None + expression = line.split(" = ", 1)[1].split(marker, 1)[0].rstrip() + if not expression.endswith(")"): + return None + depth = 0 + for index in range(len(expression) - 1, -1, -1): + if expression[index] == ")": + depth += 1 + elif expression[index] == "(": + depth -= 1 + if depth == 0: + return expression[:index].strip() + return None + + +def mir_definition(raw_name: str, crate: str, functions: dict[str, RustFunction]) -> str | None: + base = raw_name.split("::{closure#", 1)[0] + method = re.search(r">::([A-Za-z_]\w*)$", base) + if method: + candidates = [ + key + for key, function in functions.items() + if function.crate == crate and function.name == method.group(1) and function.owner + ] + else: + leaf = base.rsplit("::", 1)[-1] + candidates = [ + key + for key, function in functions.items() + if function.crate == crate and function.name == leaf + ] + qualified = base.rsplit("::", 1)[0] if "::" in base else "" + if qualified: + narrowed = [ + key + for key in candidates + if functions[key].owner == qualified or functions[key].module.endswith(f"::{qualified}") + ] + if narrowed: + candidates = narrowed + return candidates[0] if len(candidates) == 1 else None + + +def run_mir(root: Path, package: str) -> tuple[str | None, str | None]: + command = [ + "cargo", + "+nightly", + "rustc", + "-q", + "-p", + package, + "--lib", + "--", + "-Zunpretty=mir", + ] + process = subprocess.run( + command, + cwd=root, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + check=False, + ) + if process.returncode: + summary = next( + (line.strip() for line in process.stderr.splitlines() if line.startswith("error")), + f"cargo exited {process.returncode}", + ) + return None, summary + return process.stdout, None + + +def mir_calls( + mir: str, + crate: str, + functions: dict[str, RustFunction], + resolver: Resolver, +) -> dict[str, set[str]]: + result: dict[str, set[str]] = collections.defaultdict(set) + current_raw: str | None = None + current_key: str | None = None + for line in mir.splitlines(): + if line.startswith("fn "): + current_raw = line[3:].split("(", 1)[0] + current_key = mir_definition(current_raw, crate, functions) + if current_key is not None: + result.setdefault(current_key, set()) + continue + if line.startswith("const ") or (line == "}" and current_raw is not None): + if line == "}": + current_raw = None + current_key = None + continue + if current_key is None: + continue + target = mir_call_target(line) + if target is None: + continue + caller = functions[current_key] + for callee in resolver.resolve(target, caller): + if callee != current_key: + result[current_key].add(callee) + return result + + +def type_uses( + function: RustFunction, + types: dict[str, RustType], +) -> set[str]: + text = f"{function.signature}\n{function.body}" + by_name: dict[str, list[str]] = collections.defaultdict(list) + for key, item in types.items(): + by_name[item.name].append(key) + result: set[str] = set() + for name, candidates in by_name.items(): + if name in IGNORED_AMBIGUOUS_TYPES or not re.search(rf"\b{re.escape(name)}\b", text): + continue + same_module = [key for key in candidates if types[key].module == function.module] + same_crate = [key for key in candidates if types[key].crate == function.crate] + if len(same_module) == 1: + result.add(same_module[0]) + elif len(same_crate) == 1: + result.add(same_crate[0]) + elif len(candidates) == 1: + result.add(candidates[0]) + return result + + +def find_function(functions: dict[str, RustFunction], suffix: str) -> str: + matches = [key for key, function in functions.items() if function.display.endswith(suffix)] + if len(matches) != 1: + raise RuntimeError(f"expected one function ending {suffix!r}, found {len(matches)}") + return matches[0] + + +def reachable( + roots: Iterable[str], + calls: dict[str, set[str]], +) -> tuple[set[str], dict[str, int]]: + seen: set[str] = set() + depth: dict[str, int] = {} + queue = collections.deque((root, 1) for root in roots) + while queue: + function, current_depth = queue.popleft() + if function in seen: + depth[function] = min(depth[function], current_depth) + continue + seen.add(function) + depth[function] = current_depth + for callee in calls.get(function, set()): + if callee not in seen: + queue.append((callee, current_depth + 1)) + return seen, depth + + +def dot_escape(value: str) -> str: + return value.replace("\\", "\\\\").replace('"', '\\"').replace("\n", "\\n") + + +def relative_source(root: Path, path: Path, line: int) -> str: + return f"{path.relative_to(root)}:{line}" + + +def write_dot( + output: Path, + title: str, + root: Path, + functions: dict[str, RustFunction], + types: dict[str, RustType], + selected_functions: set[str], + selected_types: set[str], + calls: dict[str, set[str]], + edge_provenance: dict[tuple[str, str], set[str]], + build_reachable: set[str], + run_reachable: set[str], + type_edges: dict[str, set[str]], + include_types: bool, +) -> None: + function_ids = {key: f"f{index}" for index, key in enumerate(sorted(selected_functions))} + type_ids = {key: f"t{index}" for index, key in enumerate(sorted(selected_types))} + lines = [ + "digraph ipu_stack_package_flow {", + ' graph [rankdir=LR, bgcolor="white", fontname="DejaVu Sans", fontsize=18,', + f' label="{dot_escape(title)}", labelloc=t, labeljust=l, pad=0.25, ' + 'nodesep=0.35, ranksep=0.9, newrank=true, overlap=false, splines=polyline, outputorder=edgesfirst];', + ' node [fontname="DejaVu Sans", fontsize=9, style="rounded,filled", color="#52606d", penwidth=0.8];', + ' edge [fontname="DejaVu Sans", fontsize=7, color="#52606d", arrowsize=0.55, penwidth=0.8];', + ' build_entry [label="package build entry", shape=octagon, fillcolor="#cfe2ff", penwidth=1.4];', + ' run_entry [label="ipu_cli::main\\nCommand::HostRun arm", shape=octagon, fillcolor="#d5f5e3", penwidth=1.4];', + ] + + for crate in sorted({functions[key].crate for key in selected_functions} | {types[key].crate for key in selected_types}): + color = CRATE_COLORS.get(crate, "#eeeeee") + lines.append(f' subgraph cluster_{crate} {{ label="{crate}"; color="{color}"; style="rounded";') + for key in sorted(selected_functions, key=lambda item: functions[item].display): + function = functions[key] + if function.crate != crate: + continue + source = relative_source(root, function.path, function.line) + tooltip = f"{function.signature} — {source}" + if key in build_reachable and key in run_reachable: + border = "#7d3c98" + elif key in build_reachable: + border = "#2166ac" + else: + border = "#238b45" + lines.append( + f' {function_ids[key]} [label="{dot_escape(function.display)}", shape=box, ' + f'fillcolor="{color}", color="{border}", tooltip="{dot_escape(tooltip)}"];' + ) + if include_types: + for key in sorted(selected_types, key=lambda item: types[item].display): + rust_type = types[key] + if rust_type.crate != crate: + continue + source = relative_source(root, rust_type.path, rust_type.line) + lines.append( + f' {type_ids[key]} [label="{dot_escape(rust_type.display)}\\n«{rust_type.kind}»", ' + f'shape=ellipse, style="filled,dashed", fillcolor="{color}", color="#7b8794", ' + f'tooltip="{dot_escape(source)}"];' + ) + lines.append(" }") + + build_roots = [key for key in selected_functions if functions[key].display.endswith("::build_package")] + write_roots = [key for key in selected_functions if functions[key].display.endswith("::Application::write")] + run_suffixes = ( + "::Application::read", + "::Runtime::open", + "::Runtime::load", + "::Runtime::host_session", + "::HostSession::start", + "::HostSession::invoke", + ) + for key in build_roots + write_roots: + lines.append(f" build_entry -> {function_ids[key]} [color=\"#2166ac\", penwidth=1.4];") + for key in selected_functions: + if functions[key].display.endswith(run_suffixes): + lines.append(f" run_entry -> {function_ids[key]} [color=\"#238b45\", penwidth=1.4];") + + for caller in sorted(selected_functions): + for callee in sorted(calls.get(caller, set())): + if callee not in selected_functions: + continue + paths = edge_provenance.get((caller, callee), {"source"}) + approximate = paths == {"source"} + if caller in build_reachable and callee in build_reachable and caller in run_reachable and callee in run_reachable: + color = "#7d3c98" + elif caller in build_reachable and callee in build_reachable: + color = "#2166ac" + elif caller in run_reachable and callee in run_reachable: + color = "#238b45" + else: + color = "#52606d" + style = "dotted" if approximate else "solid" + tooltip = "source fallback (conservative)" if approximate else "rustc MIR call edge" + lines.append( + f' {function_ids[caller]} -> {function_ids[callee]} [color="{color}", ' + f'style="{style}", tooltip="{tooltip}"];' + ) + if include_types: + for type_key in sorted(type_edges.get(caller, set())): + if type_key in selected_types: + lines.append( + f' {function_ids[caller]} -> {type_ids[type_key]} ' + '[style=dashed, color="#9aa5b1", arrowhead=none, constraint=false, tooltip="uses workspace type"];' + ) + + lines.extend( + [ + ' legend [shape=note, style="filled", fillcolor="#ffffff", color="#9aa5b1", fontsize=8,', + ' label="Solid edge: rustc MIR-resolved call\\nDotted edge: source fallback\\nDashed edge: function uses workspace type\\nBlue: build path Green: run path Purple: shared"];', + "}", + ] + ) + output.write_text("\n".join(lines) + "\n", encoding="utf-8") + + +def render(dot_path: Path, svg_path: Path, engine: str) -> None: + subprocess.run([engine, "-Tsvg", str(dot_path), "-o", str(svg_path)], check=True) + + +def file_slug(value: str) -> str: + return re.sub(r"[^a-z0-9]+", "-", value.lower()).strip("-") + + +def make_partitions( + functions: dict[str, RustFunction], + types: dict[str, RustType], + selected_functions: set[str], + selected_types: set[str], + max_primary_nodes: int = 100, +) -> tuple[dict[str, GraphPartition], dict[str, str], dict[str, str]]: + """Partition primary nodes by source module, then cap oversized modules.""" + grouped: dict[tuple[str, str], list[tuple[str, str]]] = collections.defaultdict(list) + for key in selected_functions: + function = functions[key] + grouped[(function.crate, function.module)].append(("function", key)) + for key in selected_types: + rust_type = types[key] + grouped[(rust_type.crate, rust_type.module)].append(("type", key)) + + partitions: dict[str, GraphPartition] = {} + function_partition: dict[str, str] = {} + type_partition: dict[str, str] = {} + for (crate, module), nodes in sorted(grouped.items()): + nodes.sort( + key=lambda item: ( + item[0], + functions[item[1]].display if item[0] == "function" else types[item[1]].display, + ) + ) + chunks = [nodes[index : index + max_primary_nodes] for index in range(0, len(nodes), max_primary_nodes)] + for chunk_index, chunk in enumerate(chunks, start=1): + suffix = "" if len(chunks) == 1 else f"-part-{chunk_index}" + key = f"{file_slug(module)}{suffix}" + title = module if len(chunks) == 1 else f"{module} — part {chunk_index}/{len(chunks)}" + partition = GraphPartition(key=key, title=title, crate=crate) + for kind, node_key in chunk: + if kind == "function": + partition.functions.add(node_key) + function_partition[node_key] = key + else: + partition.types.add(node_key) + type_partition[node_key] = key + partitions[key] = partition + return partitions, function_partition, type_partition + + +def call_edge_attributes( + caller: str, + callee: str, + edge_provenance: dict[tuple[str, str], set[str]], + build_reachable: set[str], + run_reachable: set[str], +) -> tuple[str, str, str]: + paths = edge_provenance.get((caller, callee), {"source"}) + approximate = paths == {"source"} + if caller in build_reachable and callee in build_reachable and caller in run_reachable and callee in run_reachable: + color = "#7d3c98" + elif caller in build_reachable and callee in build_reachable: + color = "#2166ac" + elif caller in run_reachable and callee in run_reachable: + color = "#238b45" + else: + color = "#52606d" + style = "dotted" if approximate else "solid" + tooltip = "source fallback (conservative)" if approximate else "rustc MIR call edge" + return color, style, tooltip + + +def write_partition_dot( + output: Path, + root: Path, + partition: GraphPartition, + functions: dict[str, RustFunction], + types: dict[str, RustType], + all_functions: set[str], + calls: dict[str, set[str]], + edge_provenance: dict[tuple[str, str], set[str]], + build_reachable: set[str], + run_reachable: set[str], + type_edges: dict[str, set[str]], + function_partition: dict[str, str], + type_partition: dict[str, str], +) -> None: + function_ids = {key: f"f{index}" for index, key in enumerate(sorted(partition.functions))} + type_ids = {key: f"t{index}" for index, key in enumerate(sorted(partition.types))} + external_functions = { + callee + for caller in partition.functions + for callee in calls.get(caller, set()) + if callee in all_functions and callee not in partition.functions + } + external_types = { + type_key + for caller in partition.functions + for type_key in type_edges.get(caller, set()) + if type_key not in partition.types + } + external_function_ids = {key: f"xf{index}" for index, key in enumerate(sorted(external_functions))} + external_type_ids = {key: f"xt{index}" for index, key in enumerate(sorted(external_types))} + color = CRATE_COLORS.get(partition.crate, "#eeeeee") + lines = [ + "digraph ipu_stack_package_partition {", + ' graph [rankdir=LR, bgcolor="white", fontname="DejaVu Sans", fontsize=17,', + f' label="{dot_escape(partition.title)} — package build/run detail", labelloc=t, labeljust=l, ' + 'pad=0.2, nodesep=0.3, ranksep=0.75, newrank=true, overlap=false, splines=polyline, outputorder=edgesfirst];', + ' node [fontname="DejaVu Sans", fontsize=9, style="rounded,filled", color="#52606d", penwidth=0.8];', + ' edge [fontname="DejaVu Sans", fontsize=7, color="#52606d", arrowsize=0.55, penwidth=0.8];', + ] + + if any(functions[key].display.endswith(("::build_package", "::Application::write")) for key in partition.functions): + lines.append(' build_entry [label="package build entry", shape=octagon, fillcolor="#cfe2ff", penwidth=1.4];') + run_suffixes = ( + "::Application::read", + "::Runtime::open", + "::Runtime::load", + "::Runtime::host_session", + "::HostSession::start", + "::HostSession::invoke", + ) + if any(functions[key].display.endswith(run_suffixes) for key in partition.functions): + lines.append(' run_entry [label="ipu_cli::main\\nCommand::HostRun arm", shape=octagon, fillcolor="#d5f5e3", penwidth=1.4];') + + lines.append(f' subgraph cluster_primary {{ label="primary nodes"; color="{color}"; style="rounded";') + for key in sorted(partition.functions, key=lambda item: functions[item].display): + function = functions[key] + source = relative_source(root, function.path, function.line) + tooltip = f"{function.signature} — {source}" + if key in build_reachable and key in run_reachable: + border = "#7d3c98" + elif key in build_reachable: + border = "#2166ac" + else: + border = "#238b45" + lines.append( + f' {function_ids[key]} [label="{dot_escape(function.display)}", shape=box, ' + f'fillcolor="{color}", color="{border}", tooltip="{dot_escape(tooltip)}"];' + ) + for key in sorted(partition.types, key=lambda item: types[item].display): + rust_type = types[key] + source = relative_source(root, rust_type.path, rust_type.line) + lines.append( + f' {type_ids[key]} [label="{dot_escape(rust_type.display)}\\n«{rust_type.kind}»", ' + f'shape=ellipse, style="filled,dashed", fillcolor="{color}", color="#7b8794", ' + f'tooltip="{dot_escape(source)}"];' + ) + lines.append(" }") + + if external_functions or external_types: + lines.append(' subgraph cluster_external { label="outgoing references — click to open target partition"; color="#d9d9d9"; style="rounded,dashed";') + for key in sorted(external_functions, key=lambda item: functions[item].display): + function = functions[key] + target = function_partition[key] + lines.append( + f' {external_function_ids[key]} [label="{dot_escape(function.display)}", shape=box, ' + f'fillcolor="#f5f5f5", color="#9aa5b1", style="rounded,filled,dashed", ' + f'URL="{dot_escape(target)}.svg", target="_top", tooltip="open target partition"];' + ) + for key in sorted(external_types, key=lambda item: types[item].display): + rust_type = types[key] + target = type_partition[key] + lines.append( + f' {external_type_ids[key]} [label="{dot_escape(rust_type.display)}\\n«{rust_type.kind}»", shape=ellipse, ' + f'fillcolor="#f5f5f5", color="#9aa5b1", style="filled,dashed", ' + f'URL="{dot_escape(target)}.svg", target="_top", tooltip="open defining partition"];' + ) + lines.append(" }") + + for key in partition.functions: + if functions[key].display.endswith(("::build_package", "::Application::write")): + lines.append(f" build_entry -> {function_ids[key]} [color=\"#2166ac\", penwidth=1.4];") + if functions[key].display.endswith(run_suffixes): + lines.append(f" run_entry -> {function_ids[key]} [color=\"#238b45\", penwidth=1.4];") + + for caller in sorted(partition.functions): + for callee in sorted(calls.get(caller, set())): + if callee not in all_functions: + continue + target_id = function_ids.get(callee, external_function_ids.get(callee)) + if target_id is None: + continue + edge_color, style, tooltip = call_edge_attributes( + caller, callee, edge_provenance, build_reachable, run_reachable + ) + lines.append( + f' {function_ids[caller]} -> {target_id} [color="{edge_color}", style="{style}", ' + f'tooltip="{tooltip}"];' + ) + for type_key in sorted(type_edges.get(caller, set())): + target_id = type_ids.get(type_key, external_type_ids.get(type_key)) + if target_id is not None: + lines.append( + f' {function_ids[caller]} -> {target_id} ' + '[style=dashed, color="#9aa5b1", arrowhead=none, constraint=false, tooltip="uses workspace type"];' + ) + lines.extend( + [ + ' legend [shape=note, style="filled", fillcolor="#ffffff", color="#9aa5b1", fontsize=8,', + ' label="Solid: MIR-resolved call Dotted: source fallback\\nDashed, no arrow: uses type Dashed node: link to another partition\\nBlue: build Green: run Purple: shared"];', + "}", + ] + ) + output.write_text("\n".join(lines) + "\n", encoding="utf-8") + + +def partition_dependencies( + partitions: dict[str, GraphPartition], + function_partition: dict[str, str], + type_partition: dict[str, str], + calls: dict[str, set[str]], + type_edges: dict[str, set[str]], +) -> dict[tuple[str, str], tuple[int, int]]: + counts: dict[tuple[str, str], list[int]] = collections.defaultdict(lambda: [0, 0]) + for source_key, partition in function_partition.items(): + for callee in calls.get(source_key, set()): + target = function_partition.get(callee) + if target is not None and target != partition: + counts[(partition, target)][0] += 1 + for type_key in type_edges.get(source_key, set()): + target = type_partition.get(type_key) + if target is not None and target != partition: + counts[(partition, target)][1] += 1 + return {key: (value[0], value[1]) for key, value in counts.items()} + + +def write_partition_map_dot( + output: Path, + partitions: dict[str, GraphPartition], + dependencies: dict[tuple[str, str], tuple[int, int]], +) -> None: + partition_modules = { + key: re.sub(r" — part \d+/\d+$", "", partition.title) + for key, partition in partitions.items() + } + module_partitions: dict[tuple[str, str], list[str]] = collections.defaultdict(list) + for key, partition in partitions.items(): + module_partitions[(partition.crate, partition_modules[key])].append(key) + module_dependencies: dict[tuple[tuple[str, str], tuple[str, str]], list[int]] = collections.defaultdict( + lambda: [0, 0] + ) + for (source, target), (call_count, type_count) in dependencies.items(): + source_module = (partitions[source].crate, partition_modules[source]) + target_module = (partitions[target].crate, partition_modules[target]) + if source_module != target_module: + module_dependencies[(source_module, target_module)][0] += call_count + module_dependencies[(source_module, target_module)][1] += type_count + node_ids = {key: f"p{index}" for index, key in enumerate(sorted(module_partitions))} + lines = [ + "digraph ipu_stack_package_partition_map {", + ' graph [rankdir=LR, bgcolor="white", fontname="DejaVu Sans", fontsize=18,', + ' label="IPU stack package build/run — full graph partition map", labelloc=t, labeljust=l, pad=0.25, nodesep=0.4, ranksep=1.0, overlap=false, splines=polyline, concentrate=true, outputorder=edgesfirst];', + ' node [shape=box, fontname="DejaVu Sans", fontsize=9, style="rounded,filled", color="#52606d", penwidth=0.8];', + ' edge [fontname="DejaVu Sans", fontsize=7, color="#8a94a0", arrowsize=0.55, penwidth=0.7];', + ] + for crate in sorted({partition.crate for partition in partitions.values()}): + color = CRATE_COLORS.get(crate, "#eeeeee") + lines.append(f' subgraph cluster_{crate} {{ label="{crate}"; color="{color}"; style="rounded";') + for key, member_keys in sorted(module_partitions.items()): + module_crate, module = key + if module_crate != crate: + continue + function_count = sum(len(partitions[member].functions) for member in member_keys) + type_count = sum(len(partitions[member].types) for member in member_keys) + page_count = len(member_keys) + counts = f"{function_count} functions, {type_count} types, {page_count} detail page{'s' if page_count != 1 else ''}" + anchor = f"module-{file_slug(module)}" + lines.append( + f' {node_ids[key]} [label="{dot_escape(module)}\\n{counts}", fillcolor="{color}", ' + f'URL="ipu-stack-package-callgraph/index.html#{anchor}", target="_top", tooltip="open module detail list"];' + ) + lines.append(" }") + for (source, target), (call_count, type_count) in sorted(module_dependencies.items()): + label_parts = [] + if call_count: + label_parts.append(f"{call_count} calls") + if type_count: + label_parts.append(f"{type_count} type uses") + label = ", ".join(label_parts) + lines.append(f' {node_ids[source]} -> {node_ids[target]} [tooltip="{label}"];') + lines.extend( + [ + ' help [shape=note, fillcolor="#ffffff", color="#9aa5b1", label="Each box opens that module in the lightweight index.\\nThe raw monolithic DOT remains available for tools."];', + "}", + ] + ) + output.write_text("\n".join(lines) + "\n", encoding="utf-8") + + +def write_index_html( + output: Path, + partitions: dict[str, GraphPartition], + dependencies: dict[tuple[str, str], tuple[int, int]], + build_count: int, + run_count: int, + function_count: int, + type_count: int, + mir_status: dict[str, str], +) -> None: + incoming: dict[str, int] = collections.Counter(target for _, target in dependencies) + outgoing: dict[str, int] = collections.Counter(source for source, _ in dependencies) + rows = [] + seen_modules: set[str] = set() + for key, partition in sorted(partitions.items(), key=lambda item: (item[1].crate, item[1].title, item[0])): + module = re.sub(r" — part \d+/\d+$", "", partition.title) + row_id = "" + if module not in seen_modules: + row_id = f' id="module-{file_slug(module)}"' + seen_modules.add(module) + rows.append( + f"" + f"{html.escape(partition.crate)}" + f'{html.escape(partition.title)}' + f"{len(partition.functions)}{len(partition.types)}" + f"{outgoing.get(key, 0)}{incoming.get(key, 0)}" + f'DOT' + "" + ) + statuses = "".join( + f"
  • {html.escape(crate)}: {html.escape(status)}
    • " for crate, status in sorted(mir_status.items()) + ) + document = f""" + + + + +IPU stack package build/run call graph + + + +

    IPU stack package build/run call graph

    +

    {function_count} reachable workspace functions and +{type_count} involved workspace types, split into {len(partitions)} bounded SVGs. +Build reaches {build_count} functions; run reaches {run_count}.

    + + + +{''.join(rows)} +
    CrateModule / partitionFunctionsTypesOutgoing partitionsIncoming partitionsSource
    +
    Extraction status
      {statuses}
    + + +""" + output.write_text(document, encoding="utf-8") + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--root", type=Path, default=Path(__file__).resolve().parents[1]) + parser.add_argument("--output-directory", type=Path, default=None) + parser.add_argument("--no-mir", action="store_true", help="use source extraction for every crate") + arguments = parser.parse_args() + root = arguments.root.resolve() + output_directory = (arguments.output_directory or root / "docs").resolve() + output_directory.mkdir(parents=True, exist_ok=True) + + functions, types = discover_source(root) + resolver = Resolver(functions) + calls: dict[str, set[str]] = collections.defaultdict(set) + edge_provenance: dict[tuple[str, str], set[str]] = collections.defaultdict(set) + source_by_crate: dict[str, dict[str, set[str]]] = collections.defaultdict(dict) + for key, function in functions.items(): + source_by_crate[function.crate][key] = source_calls(function, resolver) + + mir_status: dict[str, str] = {} + for package in WORKSPACE_PACKAGES: + crate = crate_name(package) + extracted: dict[str, set[str]] | None = None + if not arguments.no_mir: + mir, error = run_mir(root, package) + if mir is not None: + extracted = mir_calls(mir, crate, functions, resolver) + mir_status[crate] = "MIR" + else: + mir_status[crate] = f"source fallback: {error}" + else: + mir_status[crate] = "source fallback: --no-mir" + crate_edges: dict[str, set[str]] = {} + crate_provenance: dict[str, str] = {} + for caller, source_callees in source_by_crate[crate].items(): + if extracted is not None and caller in extracted: + crate_edges[caller] = extracted[caller] + crate_provenance[caller] = "mir" + else: + crate_edges[caller] = source_callees + crate_provenance[caller] = "source" + for caller, callees in crate_edges.items(): + calls[caller].update(callees) + provenance = crate_provenance[caller] + for callee in callees: + edge_provenance[(caller, callee)].add(provenance) + + build_roots = { + find_function(functions, "ipu_codegen::package::build_package"), + find_function(functions, "ipu_package::Application::write"), + } + run_roots = { + find_function(functions, "ipu_package::Application::read"), + find_function(functions, "ipu_runtime::Runtime::open"), + find_function(functions, "ipu_runtime::Runtime::load"), + find_function(functions, "ipu_runtime::Runtime::host_session"), + find_function(functions, "ipu_driver::HostSession::start"), + find_function(functions, "ipu_driver::HostSession::invoke"), + } + build_functions, build_depth = reachable(build_roots, calls) + run_functions, run_depth = reachable(run_roots, calls) + all_functions = build_functions | run_functions + + all_type_edges = {key: type_uses(functions[key], types) for key in all_functions} + all_types = set().union(*all_type_edges.values()) if all_type_edges else set() + + overview_functions = { + key + for key in all_functions + if min(build_depth.get(key, 999), run_depth.get(key, 999)) <= 4 + } + overview_dot = output_directory / "ipu-stack-package-callgraph-overview.dot" + overview_svg = output_directory / "ipu-stack-package-callgraph-overview.svg" + full_dot = output_directory / "ipu-stack-package-callgraph-full.dot" + full_svg = output_directory / "ipu-stack-package-callgraph-full.svg" + map_dot = output_directory / "ipu-stack-package-callgraph-map.dot" + map_svg = output_directory / "ipu-stack-package-callgraph-map.svg" + partition_directory = output_directory / "ipu-stack-package-callgraph" + partition_directory.mkdir(parents=True, exist_ok=True) + + write_dot( + overview_dot, + f"IPU stack package build/run call graph — overview ({len(overview_functions)} functions)", + root, + functions, + types, + overview_functions, + set(), + calls, + edge_provenance, + build_functions, + run_functions, + {}, + False, + ) + write_dot( + full_dot, + f"IPU stack package build/run call graph — full ({len(all_functions)} functions, {len(all_types)} types)", + root, + functions, + types, + all_functions, + all_types, + calls, + edge_provenance, + build_functions, + run_functions, + all_type_edges, + True, + ) + render(overview_dot, overview_svg, "dot") + + partitions, function_partition, type_partition = make_partitions( + functions, types, all_functions, all_types + ) + expected_partition_files = {"index.html"} + for partition in partitions.values(): + expected_partition_files.update({f"{partition.key}.dot", partition.filename}) + for existing in partition_directory.iterdir(): + if ( + existing.is_file() + and existing.suffix in {".dot", ".svg"} + and existing.name not in expected_partition_files + ): + existing.unlink() + for partition in partitions.values(): + partition_dot = partition_directory / f"{partition.key}.dot" + partition_svg = partition_directory / partition.filename + write_partition_dot( + partition_dot, + root, + partition, + functions, + types, + all_functions, + calls, + edge_provenance, + build_functions, + run_functions, + all_type_edges, + function_partition, + type_partition, + ) + render(partition_dot, partition_svg, "dot") + + dependencies = partition_dependencies( + partitions, function_partition, type_partition, calls, all_type_edges + ) + write_partition_map_dot(map_dot, partitions, dependencies) + render(map_dot, map_svg, "dot") + # Keep the old full-SVG path useful and browser-safe: it is now the small + # clickable partition map. The complete monolithic graph remains as DOT. + render(map_dot, full_svg, "dot") + index_html = partition_directory / "index.html" + write_index_html( + index_html, + partitions, + dependencies, + len(build_functions), + len(run_functions), + len(all_functions), + len(all_types), + mir_status, + ) + + print(f"overview: {overview_dot.relative_to(root)} -> {overview_svg.relative_to(root)}") + print(f"full raw graph: {full_dot.relative_to(root)}") + print(f"partition map: {map_dot.relative_to(root)} -> {map_svg.relative_to(root)}") + print(f"browser index: {index_html.relative_to(root)}") + print(f"detail partitions: {len(partitions)}") + print(f"reachable functions: build={len(build_functions)} run={len(run_functions)} union={len(all_functions)}") + print(f"reachable workspace types: {len(all_types)}") + for crate, status in sorted(mir_status.items()): + print(f"{crate}: {status}") + return 0 + + +if __name__ == "__main__": + sys.exit(main())