diff --cc TODO index 54b62b6,05e34c8..0000000 --- a/TODO +++ b/TODO @@@ -1,4 -1,42 +1,49 @@@ ++<<<<<<< HEAD +- unify view-copy-output? strides. +- cleanup +- do something with attention +- exchange scheduling improvable? ++======= + Correctness: + + - Copy kernels should use high-performance ldst instructions where bank constraints allow. + - Must ask models to review for weird hacks thoroughly and remove them. + - Also duplicate code. + - Should automate finding tuning parameters at some point. + - Refactor code across modules e.g. effective_memory_elements is in ipu-exchange for some reason. + - Adjust placements to minimize delays in exchange/pairing (has not worked ever). + - Keep cost models updated. + - Does Add work? + - Why so much standard memory reliance? Scaling issue. + - Something something "unmaterializable plan". + - Fix attention (works, slow). + - Paired transfer usability (has not worked ever). + - Is exchange sufficiently dense? Cutover delays? + - Factorize cost models. + - Remove redundant formats. + - Move memory classes somewhere. + - Exchange density? + - Put EXECUTABLE_MEMORY_LIMIT and such in the right place. + - Code sharing with deferred things. + - prepare_attention_blocks is also wrong. + - finish_gemm_plan etc is suspicious + - plans_for_operation maybe duplicates something + - ParallelGridProxy?! + - IPU21_PLANNED_DATA_BYTES factor + - objective duplication + - `let column_grain` is really long. Lots of long chains like this → hard to read. + - Generally wrong abstractions. Trace through code? + - Merge topology and target + - Fix copy descriptors + + Completeness: + + - LayerNorm. + - MAP head. + - Other ops? Subtract, but that's pretty trivial. + - Inplace, fusion, or something. + - Positional encoding. + - Input projection (probably this is just batched MLP). + - FP8 support! + - Run full model. ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec diff --cc crates/ipu-cli/src/main.rs index 85d11ae,38f5640..0000000 --- a/crates/ipu-cli/src/main.rs +++ b/crates/ipu-cli/src/main.rs @@@ -1,11 -1,10 +1,15 @@@ use anyhow::{Context, Result, bail}; - use clap::{Parser, Subcommand, ValueEnum}; + use clap::{Parser, Subcommand}; use ipu_driver::{Device, block_device_interrupt_signals}; use ipu_elf::{LinkOptions, Toolchain, inspect_object, link, source_tree_digest}; - use ipu_package::{Application, ProfileExchangeActivityKind, ProfileReport, ProfileStepKind}; + use ipu_package::{Application, ExchangeActivityKind, ProfileReport, ProfileStepKind}; use ipu_profile::{ ++<<<<<<< HEAD + GroupBy, Query, SortBy, StepKind, calibrate_profiles, cycle_origin, exchange_activity_summary, + exchange_boundaries, query, ++======= + GroupBy, Query, SortBy, calibrate_profiles, cycle_origin, exchange_activity_summary, query, ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec }; use ipu_runtime::Runtime; use std::collections::{BTreeSet, HashMap}; diff --cc crates/ipu-codegen/src/exchange.rs index ed80278,cf0f315..0000000 --- a/crates/ipu-codegen/src/exchange.rs +++ b/crates/ipu-codegen/src/exchange.rs @@@ -1,34 -1,26 +1,41 @@@ //! Physical exchange programs generated from logical shard transfers. +mod diagnostic; +mod order; +use diagnostic::PhaseDiagnostics; +pub use diagnostic::diagnose_exchange_tile; +use order::{critical_neighborhood_order, point_to_point_matching_wave_order}; +mod reuse; +mod traffic; +pub(crate) use reuse::ExchangeScheduleCache; +pub(crate) use traffic::MappingTraffic; + use crate::{ ++<<<<<<< HEAD + BlockValueId, ExchangePhaseId, LogicalExchange, LowProgram, Placement, ShardDefinition, + logical_view_byte_spans, view_byte_spans, ++======= + ByteSpan, ExchangeOrder, ExchangePhaseId, LogicalExchange, LowProgram, LowShardId, Placement, + ShardDefinition, logical_view_byte_spans, shard_storage_bytes, view_byte_spans, ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec }; - use ipu_exchange::{ - MAX_TRANSFER_WORDS, MulticastPlan, PhaseProgramBuilder, RETURN_M10_INSTRUCTION, Topology, - finalize_point_receiver, patch_receiver_address, patch_sender_address, - patch_sender_instruction, sender_address_instruction_groups, - }; - use ipu_package::{ - IPU21_INTERLEAVED_ELEMENT_SIZE, IPU21_INTERLEAVED_MEMORY_BASE, TILE_MEMORY_ELEMENT_SIZE, + use ipu_package::ExchangeActivityKind; + use ipu_target::exchange::{ + PhaseProgramBuilder, PhaseTransferTiming, PhysicalTransfer, ResolvedTransfer, TransferEndpoint, + TransferWidth, patch_sender_instruction, sender_address_instruction_groups, }; + use ipu_target::hardware::HardwareTarget; + use ipu_target::instruction::RETURN_M10_INSTRUCTION; + use ipu_target::memory::{MemoryElement, memory_elements_for_words}; + use ipu_target::topology::Topology; use rayon::prelude::*; - use serde::{Deserialize, Serialize}; + use std::borrow::Cow; use std::cmp::Reverse; - use std::collections::{BTreeMap, BTreeSet, BinaryHeap, VecDeque}; + use std::collections::{BTreeMap, BTreeSet, BinaryHeap}; + use std::time::Instant; #[cfg(test)] - use ipu_exchange::plan_event_cycles; + use ipu_target::exchange::plan_event_cycles; #[derive(Clone, Debug, PartialEq, Eq)] pub struct PhysicalExchangePhase { @@@ -70,93 -55,6 +72,96 @@@ pub struct ExchangeActivity pub words: u32, } ++<<<<<<< HEAD +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ExchangeActivityKind { + Send, + Receive, + /// This tile's transmit lane is borrowed by its partner; receiving remains available. + PartnerBusy, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)] +pub struct ExchangeMemoryElement { + pub interleaved: bool, + pub index: u32, +} + +pub const EXCHANGE_SCHEDULE_SNAPSHOT_VERSION: u32 = 3; + +/// Address-resolved transfers captured immediately before physical scheduling. +/// Replaying this data exercises the production scheduler and exchange-row +/// encoder without compiling kernels or loading a device. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct ExchangeScheduleSnapshot { + pub schema_version: u32, + pub tile_count: u16, + pub phases: Vec, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct ExchangeScheduleProblem { + pub phase: u32, + pub transfers: Vec, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct ExchangeScheduleTransfer { + pub source: u16, + /// Address used by each structured-repeat iteration. Ordinary transfers + /// contain exactly one entry. + pub source_addresses: Vec, + pub destinations: Vec, + pub words: u32, + pub width: ExchangeItemWidth, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +pub enum ExchangeItemWidth { + #[default] + Word32, + Paired64, +} + +impl ExchangeItemWidth { + fn item_words(self) -> u32 { + match self { + Self::Word32 => 1, + Self::Paired64 => 2, + } + } + + fn item_count(self, words: u32) -> Result { + let item_words = self.item_words(); + if words == 0 || !words.is_multiple_of(item_words) { + return Err(ExchangeLoweringError::UnalignedPayload); + } + Ok(words / item_words) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct ExchangeScheduleDestination { + pub tile: u16, + pub address: u32, +} + +#[derive(Clone, Debug)] +pub struct LoweredExchanges { + pub phases: Vec, + pub schedule_snapshot: ExchangeScheduleSnapshot, +} + +#[derive(Clone, Debug)] +pub struct ExchangeScheduleRun { + pub phase: PhysicalExchangePhase, + pub initial_horizon: u32, + pub endpoint_lower_bound: u32, + pub neighborhood_improvements: usize, +} + ++======= ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec #[derive(Debug, thiserror::Error, PartialEq, Eq)] pub enum ExchangeLoweringError { #[error(transparent)] @@@ -175,38 -73,16 +180,43 @@@ Overflow, #[error("structured-repeat exchange rows have incompatible shapes")] IncompatibleRepeatRows, - #[error("exchange diagnostic refers to missing tile {0}")] - DiagnosticTile(u16), - #[error("invalid exchange-schedule snapshot: {0}")] - InvalidSnapshot(String), - #[error("exchange-schedule invariant failed: {0}")] - Invariant(String), } ++<<<<<<< HEAD +#[cfg(test)] +pub(crate) fn lower_exchanges( + program: &LowProgram, + placement: &Placement, + topology: &Topology, + enable_diagnostics: bool, +) -> Result { + lower_exchanges_cached( + program, + placement, + topology, + enable_diagnostics, + &mut ExchangeScheduleCache::default(), + ) +} + +pub(crate) fn lower_exchanges_cached( + program: &LowProgram, + placement: &Placement, + topology: &Topology, + enable_diagnostics: bool, + cache: &mut ExchangeScheduleCache, +) -> Result { + let mut repeat_inputs = BTreeMap::>::new(); ++======= + pub(crate) fn lower_exchanges( + program: &LowProgram, + placement: &Placement, + target: HardwareTarget, + ) -> Result, ExchangeLoweringError> { + let topology = target.topology().prefix(program.tile_count)?; + let maximum_transfer_words = target.exchange().maximum_transfer_words; + let mut repeat_inputs = BTreeMap::>::new(); ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec for repeat in &program.repeat_runs { for iterated in &repeat.iterated { match repeat_inputs.entry(iterated.argument) { @@@ -247,140 -125,41 +259,156 @@@ .into_iter() .flatten() .collect(); - let mut pending = coalesce_pending_transfers(pending); + let prepared = started.elapsed(); + let mut pending = coalesce_pending_transfers(pending, maximum_transfer_words); + let coalesced = started.elapsed(); attach_repeat_source_addresses(&mut pending, &repeat_inputs, placement)?; ++<<<<<<< HEAD + let ScheduledPending { + pending, + receive_counts, + incoming_bases, + optimized, + } = cache.select(phase.id, topology, pending, program.tile_count)?; + let schedule_problem = schedule_problem(phase.id.index(), &pending); + let mut destination_multiplicity = BTreeMap::new(); + for transfer in &pending { + for &(tile, address) in &transfer.destinations { + *destination_multiplicity + .entry((tile, address, transfer.words)) + .or_insert(0usize) += 1; + } + } + let maximum_identical_destinations = destination_multiplicity + .values() + .copied() + .max() + .unwrap_or(0); + if pending.len() > 1_000 || maximum_identical_destinations > 1 { + tracing::info!( + phase = phase.id.index(), + transfers = pending.len(), + maximum_identical_destinations, + "prepared large physical exchange phase" + ); + } + let OptimizedSchedule { + schedule, + initial_horizon, + endpoint_lower_bound, + selected_kind, + neighborhood_improvements, + } = optimized; + if enable_diagnostics { + let repeat_iterations = pending + .iter() + .map(|transfer| transfer.source_addresses.len()) + .max() + .unwrap_or(1); + if repeat_iterations > 1 { + let mut unsafe_pending = pending.clone(); + for transfer in &mut unsafe_pending { + transfer.source_addresses.truncate(1); + transfer.refresh_source_elements(); + } + let unsafe_schedule = optimize_pending_schedule( + topology, + &unsafe_pending, + &incoming_bases, + &receive_counts, + program.tile_count, + )?; + tracing::info!( + phase = phase.id.index(), + repeat_iterations, + unsafe_horizon = unsafe_schedule.schedule.horizon, + repeat_safe_horizon = schedule.horizon, + repeat_safety_cost = schedule + .horizon + .saturating_sub(unsafe_schedule.schedule.horizon), + "compared repeat-safe exchange schedule with first-iteration-only baseline" + ); + } + } + if pending.len() > 1_000 { + tracing::info!( + phase = phase.id.index(), + initial_horizon, + selected_horizon = schedule.horizon, + endpoint_lower_bound, + lower_bound_gap = schedule.horizon.saturating_sub(endpoint_lower_bound), + selected_kind, + neighborhood_improvements, + "optimized physical exchange schedule" + ); + } ++======= + let incoming_bases = incoming_bases(&pending, program.tile_count)?; + tracing::info!( + phase = phase.id.index(), + logical_transfers = phase.transfers.len(), + physical_transfers = pending.len(), + prepare_ms = prepared.as_millis(), + coalesce_ms = coalesced.saturating_sub(prepared).as_millis(), + "prepared exchange phase" + ); + let schedule = materialize_greedy_schedule( + &topology, + &pending, + &incoming_bases, + program.tile_count, + )?; + tracing::info!( + phase = phase.id.index(), + physical_transfers = pending.len(), + schedule_ms = started.elapsed().saturating_sub(coalesced).as_millis(), + "scheduled exchange phase" + ); ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec let MaterializedSchedule { builder, - horizon, - tile_availability, activities, scheduled_sends, - order, - timings, .. } = schedule; ++<<<<<<< HEAD + let mut diagnostics = + enable_diagnostics.then(|| PhaseDiagnostics::new(program.tile_count)); + if let Some(diagnostics) = &mut diagnostics { + let mut endpoint_roles = vec![0usize; usize::from(program.tile_count)]; + for tile in pending.iter().flat_map(PendingTransfer::tiles) { + endpoint_roles[usize::from(tile)] += 1; + } + diagnostics.maximum_endpoint_roles = endpoint_roles.into_iter().max().unwrap_or(0); + for &index in &order { + let transfer = &pending[index]; + let timing = timings[index].ok_or(ExchangeLoweringError::Overflow)?; + diagnostics.record( + transfer.source, + transfer.source_address(), + &transfer.destinations, + transfer.words, + timing.start, + timing.end, + timing.blocking_tile, + ); + } + } + if let Some(diagnostics) = diagnostics { + diagnostics.emit( + phase.id.index(), + &phase.provenance, + horizon, + &tile_availability, + &builder, + ); + } ++======= + let horizon = builder.event_cycles(); ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec let phase_programs = builder.finish()?; debug_assert_eq!(phase_programs.event_cycles, horizon); - let tile_event_cycles = phase_programs.tile_event_cycles; - let active = phase_programs - .programs - .iter() - .map(Option::is_some) - .collect::>(); - let programs = phase_programs - .programs - .into_iter() - .map(|program| program.unwrap_or_else(inactive_exchange_program)) - .collect::>(); + let programs = phase_programs.programs; let repeat_patches = programs .iter() .enumerate() @@@ -473,9 -227,18 +476,22 @@@ fn prepare_transfer program: &LowProgram, placement: &Placement, transfer: &LogicalExchange, + maximum_transfer_words: u32, ) -> Result, ExchangeLoweringError> { let source = &program.shards[transfer.source.shard.index() as usize]; ++<<<<<<< HEAD + let logical_order = transfer.span_order(&program.shards) == crate::CopyOrder::Semantic; ++======= + let logical_order = transfer.order == ExchangeOrder::Semantic + && transfer.destinations.iter().any(|view| { + program.shards[view.shard.index() as usize] + .tensor_type + .format + .layout + .order + != source.tensor_type.format.layout.order + }); ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec let source_base = placement .shard_addresses .get(&source.id) @@@ -612,15 -418,10 +671,15 @@@ #[derive(Clone)] struct PendingTransfer { ++<<<<<<< HEAD + source: u16, + source_shard: BlockValueId, ++======= + physical: PhysicalTransfer, + source_shard: LowShardId, ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec source_offset: u32, - destinations: Vec<(u16, u32)>, - source_addresses: Vec, - source_elements: Vec, - words: u32, - width: ExchangeItemWidth, - reserved_source: Option, + source_elements: Vec, } impl PendingTransfer { @@@ -678,220 -466,13 +724,224 @@@ fn attach_repeat_source_addresses Ok(()) } ++<<<<<<< HEAD +fn paired_transfer_alternatives( + pending: &[PendingTransfer], + topology: &Topology, + tile_count: u16, +) -> Result>, ExchangeLoweringError> { + let mut alternatives = Vec::with_capacity(pending.len()); + for transfer in pending { + if transfer.width != ExchangeItemWidth::Word32 + || transfer.words < 128 + || transfer.words & 1 != 0 + || transfer + .source_addresses + .iter() + .any(|address| address & 0b111 != 0) + { + alternatives.push(None); + continue; + } + let source_pair = topology.paired_logical(transfer.source)?; + if source_pair >= tile_count { + alternatives.push(None); + continue; + } + + let mut by_pair = BTreeMap::>::new(); + for &(tile, address) in &transfer.destinations { + by_pair + .entry(topology.physical(tile)? & !2) + .or_default() + .push((tile, address)); + } + let mut paired_destinations = Vec::with_capacity(transfer.destinations.len()); + let all_destinations_pairable = by_pair.into_values().all(|destinations| { + let complete_pair = destinations.len() == 2 + && topology + .paired_logical(destinations[0].0) + .is_ok_and(|paired| paired == destinations[1].0); + let pairable = complete_pair + && destinations + .iter() + .all(|(tile, address)| *tile != source_pair && address & 0b111 == 0) + && destinations[0].1 == destinations[1].1; + if pairable { + paired_destinations.extend(destinations); + } + pairable + }); + if !all_destinations_pairable || paired_destinations.is_empty() { + alternatives.push(None); + continue; + } + + let paired_tiles = paired_destinations + .iter() + .map(|&(tile, _)| tile) + .collect::>(); + if topology + .paired_multicast(transfer.source, &paired_tiles, transfer.words / 2) + .is_err() + { + alternatives.push(None); + continue; + } + let mut paired = transfer.clone(); + paired.destinations = paired_destinations; + paired.width = ExchangeItemWidth::Paired64; + paired.reserved_source = Some(source_pair); + alternatives.push(Some(paired)); + } + Ok(alternatives) +} + +fn schedule_problem(phase: u32, pending: &[PendingTransfer]) -> ExchangeScheduleProblem { + ExchangeScheduleProblem { + phase, + transfers: pending + .iter() + .map(|transfer| ExchangeScheduleTransfer { + source: transfer.source, + source_addresses: transfer.source_addresses.clone(), + destinations: transfer + .destinations + .iter() + .map(|&(tile, address)| ExchangeScheduleDestination { tile, address }) + .collect(), + words: transfer.words, + width: transfer.width, + }) + .collect(), + } +} + +fn pending_from_problem( + tile_count: u16, + problem: &ExchangeScheduleProblem, +) -> Result, ExchangeLoweringError> { + problem + .transfers + .iter() + .enumerate() + .map(|(index, transfer)| { + if transfer.source >= tile_count { + return Err(ExchangeLoweringError::InvalidSnapshot(format!( + "phase {} transfer {index} has source tile {} outside 0..{tile_count}", + problem.phase, transfer.source + ))); + } + if transfer.words == 0 || transfer.words > MAX_TRANSFER_WORDS { + return Err(ExchangeLoweringError::InvalidSnapshot(format!( + "phase {} transfer {index} has invalid word count {}", + problem.phase, transfer.words + ))); + } + if transfer.width == ExchangeItemWidth::Paired64 + && (transfer.words < 128 || transfer.words & 1 != 0) + { + return Err(ExchangeLoweringError::InvalidSnapshot(format!( + "phase {} transfer {index} has invalid {}-word paired payload", + problem.phase, transfer.words + ))); + } + if transfer.source_addresses.is_empty() { + return Err(ExchangeLoweringError::InvalidSnapshot(format!( + "phase {} transfer {index} has no source addresses", + problem.phase + ))); + } + let bytes = transfer + .words + .checked_mul(4) + .ok_or(ExchangeLoweringError::Overflow)?; + for &address in &transfer.source_addresses { + let alignment_mask = match transfer.width { + ExchangeItemWidth::Word32 => 0b11, + ExchangeItemWidth::Paired64 => 0b111, + }; + if address & alignment_mask != 0 { + return Err(ExchangeLoweringError::InvalidSnapshot(format!( + "phase {} transfer {index} has unaligned source address {address:#x}", + problem.phase + ))); + } + address + .checked_add(bytes) + .ok_or(ExchangeLoweringError::Overflow)?; + } + if transfer.destinations.is_empty() { + return Err(ExchangeLoweringError::InvalidSnapshot(format!( + "phase {} transfer {index} has no destinations", + problem.phase + ))); + } + let mut destination_tiles = BTreeSet::new(); + let destinations = transfer + .destinations + .iter() + .map(|destination| { + if destination.tile >= tile_count + || destination.tile == transfer.source + || !destination_tiles.insert(destination.tile) + { + return Err(ExchangeLoweringError::InvalidSnapshot(format!( + "phase {} transfer {index} has invalid destination tile {}", + problem.phase, destination.tile + ))); + } + let alignment_mask = match transfer.width { + ExchangeItemWidth::Word32 => 0b11, + ExchangeItemWidth::Paired64 => 0b111, + }; + if destination.address & alignment_mask != 0 { + return Err(ExchangeLoweringError::InvalidSnapshot(format!( + "phase {} transfer {index} has unaligned destination address {:#x}", + problem.phase, destination.address + ))); + } + destination + .address + .checked_add(bytes) + .ok_or(ExchangeLoweringError::Overflow)?; + Ok((destination.tile, destination.address)) + }) + .collect::, ExchangeLoweringError>>()?; + let mut pending = PendingTransfer { + source: transfer.source, + source_shard: BlockValueId::from_index( + u32::try_from(index).map_err(|_| ExchangeLoweringError::Overflow)?, + ), + source_offset: 0, + destinations, + source_addresses: transfer.source_addresses.clone(), + source_elements: Vec::new(), + words: transfer.words, + width: transfer.width, + reserved_source: match transfer.width { + ExchangeItemWidth::Word32 => None, + ExchangeItemWidth::Paired64 => { + Some(Topology::c600().paired_logical(transfer.source)?) + } + }, + }; + pending.refresh_source_elements(); + Ok(pending) + }) + .collect() +} + +fn receive_configuration( ++======= + fn incoming_bases( ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec pending: &[PendingTransfer], tile_count: u16, - ) -> Result<(Vec, Vec), ExchangeLoweringError> { + ) -> Result>, ExchangeLoweringError> { let mut receive_counts = vec![0usize; usize::from(tile_count)]; for transfer in pending { - for &(tile, _) in &transfer.destinations { + for &TransferEndpoint(tile, _) in &transfer.physical.destinations { let count = receive_counts .get_mut(usize::from(tile)) .ok_or(ExchangeLoweringError::InvalidDestination)?; @@@ -900,522 -481,13 +950,530 @@@ } let mut incoming_bases = vec![None::; usize::from(tile_count)]; for transfer in pending { ++<<<<<<< HEAD + if let [(tile, address)] = transfer.destinations.as_slice() ++======= + if let [TransferEndpoint(tile, address)] = transfer.physical.destinations.as_slice() ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec && receive_counts[usize::from(*tile)] == 1 { incoming_bases[usize::from(*tile)] = Some(*address); } } ++<<<<<<< HEAD + Ok(( + receive_counts, + incoming_bases + .into_iter() + .map(|base| base.unwrap_or(0)) + .collect(), + )) +} + +struct OptimizedSchedule { + schedule: MaterializedSchedule, + initial_horizon: u32, + endpoint_lower_bound: u32, + selected_kind: &'static str, + neighborhood_improvements: usize, +} + +struct ScheduledPending { + pending: Vec, + receive_counts: Vec, + incoming_bases: Vec, + optimized: OptimizedSchedule, +} + +fn optimize_owned_pending( + topology: &Topology, + pending: Vec, + tile_count: u16, +) -> Result { + let (receive_counts, incoming_bases) = receive_configuration(&pending, tile_count)?; + let optimized = optimize_pending_schedule( + topology, + &pending, + &incoming_bases, + &receive_counts, + tile_count, + )?; + Ok(ScheduledPending { + pending, + receive_counts, + incoming_bases, + optimized, + }) +} + +/// Compare complete width choices. A single width change can leave another +/// path tied at the horizon, so individually profitable transfers are not a +/// useful prerequisite for pairing. This bounds search to two optimizations. +fn select_transfer_widths( + phase: u32, + topology: &Topology, + pending: Vec, + tile_count: u16, +) -> Result { + let alternatives = paired_transfer_alternatives(&pending, topology, tile_count)?; + let candidates = alternatives.iter().flatten().count(); + if candidates == 0 { + return optimize_owned_pending(topology, pending, tile_count); + } + let paired = pending + .iter() + .zip(alternatives) + .map(|(ordinary, paired)| paired.unwrap_or_else(|| ordinary.clone())) + .collect(); + let ordinary = optimize_owned_pending(topology, pending, tile_count)?; + let paired = match optimize_owned_pending(topology, paired, tile_count) { + Ok(paired) => paired, + Err(error) => { + tracing::debug!(phase, %error, "paired exchange candidate is not encodable"); + return Ok(ordinary); + } + }; + let ordinary_horizon = ordinary.optimized.schedule.horizon; + let paired_horizon = paired.optimized.schedule.horizon; + let use_paired = paired_horizon < ordinary_horizon; + tracing::info!( + phase, + candidates, + ordinary_horizon, + paired_horizon, + use_paired, + "compared ordinary and paired exchange schedules" + ); + Ok(if use_paired { paired } else { ordinary }) +} + +fn optimize_pending_schedule( + topology: &Topology, + pending: &[PendingTransfer], + incoming_bases: &[u32], + receive_counts: &[usize], + tile_count: u16, +) -> Result { + let schedule = materialize_greedy_schedule( + topology, + pending, + incoming_bases, + receive_counts, + tile_count, + )?; + improve_pending_schedule( + topology, + pending, + incoming_bases, + receive_counts, + tile_count, + schedule, + "full-duplex", + ) +} + +#[allow(clippy::too_many_arguments)] +fn improve_pending_schedule( + topology: &Topology, + pending: &[PendingTransfer], + incoming_bases: &[u32], + receive_counts: &[usize], + tile_count: u16, + mut schedule: MaterializedSchedule, + initial_kind: &'static str, +) -> Result { + let initial_horizon = schedule_score(&schedule); + let endpoint_lower_bound = endpoint_work_lower_bound(pending, tile_count); + let mut selected_kind = initial_kind; + let mut neighborhood_improvements = 0usize; + if let Some(order) = point_to_point_matching_wave_order(pending, tile_count, &schedule.order) { + let matching = materialize_schedule_order( + topology, + pending, + incoming_bases, + receive_counts, + tile_count, + &order, + false, + ); + if let Ok(matching) = matching + && schedule_score(&matching) < schedule_score(&schedule) + { + schedule = matching; + selected_kind = "matching-waves"; + } + } + loop { + let repaired_order = critical_neighborhood_order(pending, tile_count, &schedule); + if repaired_order == schedule.order { + break; + } + let repaired = materialize_schedule_order( + topology, + pending, + incoming_bases, + receive_counts, + tile_count, + &repaired_order, + false, + ); + let Ok(repaired) = repaired else { + break; + }; + if schedule_score(&repaired) >= schedule_score(&schedule) { + break; + } + schedule = repaired; + selected_kind = "critical-neighborhood"; + neighborhood_improvements += 1; + } + Ok(OptimizedSchedule { + schedule, + initial_horizon, + endpoint_lower_bound, + selected_kind, + neighborhood_improvements, + }) +} + +impl ExchangeScheduleSnapshot { + pub fn validate(&self) -> Result<(), ExchangeLoweringError> { + if self.schema_version != EXCHANGE_SCHEDULE_SNAPSHOT_VERSION { + return Err(ExchangeLoweringError::InvalidSnapshot(format!( + "unsupported schema version {} (expected {})", + self.schema_version, EXCHANGE_SCHEDULE_SNAPSHOT_VERSION + ))); + } + if self.tile_count == 0 || usize::from(self.tile_count) > Topology::c600().tile_count() { + return Err(ExchangeLoweringError::InvalidSnapshot(format!( + "tile count {} is outside the C600 topology", + self.tile_count + ))); + } + let mut phases = BTreeSet::new(); + for problem in &self.phases { + if !phases.insert(problem.phase) { + return Err(ExchangeLoweringError::InvalidSnapshot(format!( + "duplicate phase {}", + problem.phase + ))); + } + pending_from_problem(self.tile_count, problem)?; + } + Ok(()) + } +} + +/// Runs the same ordering, timing, full-duplex code generation, and row +/// validation used by package lowering on one captured phase. +pub fn schedule_exchange_problem( + tile_count: u16, + problem: &ExchangeScheduleProblem, +) -> Result { + if tile_count == 0 || usize::from(tile_count) > Topology::c600().tile_count() { + return Err(ExchangeLoweringError::InvalidSnapshot(format!( + "tile count {tile_count} is outside the C600 topology" + ))); + } + let topology = Topology::new( + (0..tile_count) + .map(ipu_exchange::c600_logical_to_physical) + .collect(), + )?; + let pending = pending_from_problem(tile_count, problem)?; + let (receive_counts, incoming_bases) = receive_configuration(&pending, tile_count)?; + let OptimizedSchedule { + schedule, + initial_horizon, + endpoint_lower_bound, + neighborhood_improvements, + .. + } = optimize_pending_schedule( + &topology, + &pending, + &incoming_bases, + &receive_counts, + tile_count, + )?; + let MaterializedSchedule { + builder, + horizon, + activities, + .. + } = schedule; + let phase_programs = builder.finish()?; + if phase_programs.event_cycles != horizon { + return Err(ExchangeLoweringError::Invariant(format!( + "phase {} row horizon {} differs from scheduled horizon {horizon}", + problem.phase, phase_programs.event_cycles + ))); + } + let tile_event_cycles = phase_programs.tile_event_cycles; + let active = phase_programs + .programs + .iter() + .map(Option::is_some) + .collect::>(); + let programs = phase_programs + .programs + .into_iter() + .map(|program| program.unwrap_or_else(inactive_exchange_program)) + .collect::>(); + let phase = PhysicalExchangePhase { + id: ExchangePhaseId::from_index(problem.phase), + active, + programs, + incoming_bases, + tile_event_cycles, + event_cycles: horizon, + activities, + repeat_patches: vec![Vec::new(); usize::from(tile_count)], + }; + Ok(ExchangeScheduleRun { + phase, + initial_horizon, + endpoint_lower_bound, + neighborhood_improvements, + }) +} + +/// Checks that scheduled activities and encoded rows preserve the captured +/// transfer set and obey per-tile bus and SRAM-element hazards. +pub fn validate_exchange_schedule( + tile_count: u16, + problem: &ExchangeScheduleProblem, + phase: &PhysicalExchangePhase, +) -> Result<(), ExchangeLoweringError> { + let fail = |message| ExchangeLoweringError::Invariant(message); + let size = usize::from(tile_count); + if phase.id.index() != problem.phase { + return Err(fail(format!( + "phase id {} differs from snapshot phase {}", + phase.id.index(), + problem.phase + ))); + } + for (name, length) in [ + ("active", phase.active.len()), + ("programs", phase.programs.len()), + ("incoming bases", phase.incoming_bases.len()), + ("tile horizons", phase.tile_event_cycles.len()), + ("activities", phase.activities.len()), + ("repeat patches", phase.repeat_patches.len()), + ] { + if length != size { + return Err(fail(format!( + "phase {} has {length} {name} entries for {tile_count} tiles", + problem.phase + ))); + } + } + if phase + .repeat_patches + .iter() + .any(|patches| !patches.is_empty()) + { + return Err(fail(format!( + "standalone phase {} unexpectedly contains repeat patches", + problem.phase + ))); + } + let maximum_horizon = phase.tile_event_cycles.iter().copied().max().unwrap_or(0); + if phase.event_cycles != maximum_horizon { + return Err(fail(format!( + "phase {} horizon {} differs from maximum tile horizon {maximum_horizon}", + problem.phase, phase.event_cycles + ))); + } + + let mut send_counts = vec![0usize; problem.transfers.len()]; + let mut partner_busy_counts = vec![0usize; problem.transfers.len()]; + let mut receive_counts = problem + .transfers + .iter() + .map(|transfer| vec![0usize; transfer.destinations.len()]) + .collect::>(); + let reserved_paired_sources = problem + .transfers + .iter() + .filter(|transfer| transfer.width == ExchangeItemWidth::Paired64) + .map(|transfer| Topology::c600().paired_logical(transfer.source)) + .collect::, _>>()?; + for tile in 0..size { + let decoded = ipu_exchange::diagnostic::diagnose_plan_program(&phase.programs[tile], None)?; + if decoded.event_cycles != phase.tile_event_cycles[tile] { + return Err(fail(format!( + "phase {} tile {tile} decoded horizon {} differs from {}", + problem.phase, decoded.event_cycles, phase.tile_event_cycles[tile] + ))); + } + let tile_u16 = u16::try_from(tile).map_err(|_| ExchangeLoweringError::Overflow)?; + let expected_active = + !phase.activities[tile].is_empty() || reserved_paired_sources.contains(&tile_u16); + if phase.active[tile] != expected_active + || phase.active[tile] != (phase.tile_event_cycles[tile] != 0) + { + return Err(fail(format!( + "phase {} tile {tile} has inconsistent active state", + problem.phase + ))); + } + for activity in &phase.activities[tile] { + if activity.start_cycle > activity.end_cycle + || activity.end_cycle > activity.memory_end_cycle + || activity.memory_end_cycle > phase.tile_event_cycles[tile] + { + return Err(fail(format!( + "phase {} tile {tile} transfer {} has invalid cycle interval", + problem.phase, activity.transfer + ))); + } + let transfer_index = + usize::try_from(activity.transfer).map_err(|_| ExchangeLoweringError::Overflow)?; + let transfer = problem.transfers.get(transfer_index).ok_or_else(|| { + fail(format!( + "phase {} tile {tile} references missing transfer {}", + problem.phase, activity.transfer + )) + })?; + if activity.words != transfer.words { + return Err(fail(format!( + "phase {} tile {tile} transfer {transfer_index} has wrong word count", + problem.phase + ))); + } + match activity.kind { + ExchangeActivityKind::Send => { + if usize::from(transfer.source) != tile + || activity.address != transfer.source_addresses[0] + { + return Err(fail(format!( + "phase {} transfer {transfer_index} has a mismatched send activity", + problem.phase + ))); + } + send_counts[transfer_index] += 1; + } + ExchangeActivityKind::Receive => { + let destination = transfer + .destinations + .iter() + .position(|destination| { + usize::from(destination.tile) == tile + && destination.address == activity.address + }) + .ok_or_else(|| { + fail(format!( + "phase {} transfer {transfer_index} has an unexpected receive activity on tile {tile}", + problem.phase + )) + })?; + receive_counts[transfer_index][destination] += 1; + } + ExchangeActivityKind::PartnerBusy => { + let expected = (transfer.width == ExchangeItemWidth::Paired64) + .then(|| Topology::c600().paired_logical(transfer.source)) + .transpose()?; + if expected != Some(tile_u16) + || activity.address != transfer.source_addresses[0] + { + return Err(fail(format!( + "phase {} transfer {transfer_index} has a mismatched partner-busy activity", + problem.phase + ))); + } + partner_busy_counts[transfer_index] += 1; + } + } + } + for kind in [ExchangeActivityKind::Send, ExchangeActivityKind::Receive] { + let mut intervals = phase.activities[tile] + .iter() + .filter(|activity| activity.kind == kind) + .map(|activity| (activity.start_cycle, activity.end_cycle)) + .collect::>(); + intervals.sort_unstable(); + if intervals.windows(2).any(|pair| pair[1].0 < pair[0].1) { + return Err(fail(format!( + "phase {} tile {tile} has overlapping {kind:?} bus intervals", + problem.phase + ))); + } + } + let sends = phase.activities[tile] + .iter() + .filter(|activity| activity.kind == ExchangeActivityKind::Send); + for send in sends { + let transfer = &problem.transfers[send.transfer as usize]; + for receive in phase.activities[tile] + .iter() + .filter(|activity| activity.kind == ExchangeActivityKind::Receive) + { + let overlaps = send.start_cycle < receive.memory_end_cycle + && receive.start_cycle < send.memory_end_cycle; + if overlaps + && transfer.source_addresses.iter().any(|&address| { + spans_share_effective_memory_element( + address, + send.words, + receive.address, + receive.words, + ) + }) + { + return Err(fail(format!( + "phase {} tile {tile} overlaps send/receive access to one SRAM element", + problem.phase + ))); + } + } + } + for partner_busy in phase.activities[tile] + .iter() + .filter(|activity| activity.kind == ExchangeActivityKind::PartnerBusy) + { + if phase.activities[tile].iter().any(|activity| { + activity.transfer != partner_busy.transfer + && activity.kind != ExchangeActivityKind::Receive + && activity.start_cycle < partner_busy.end_cycle + && partner_busy.start_cycle < activity.end_cycle + }) { + return Err(fail(format!( + "phase {} tile {tile} overlaps borrowed and local transmit intervals", + problem.phase + ))); + } + } + } + for (index, count) in send_counts.into_iter().enumerate() { + if count != 1 { + return Err(fail(format!( + "phase {} transfer {index} has {count} send activities", + problem.phase + ))); + } + } + for (index, count) in partner_busy_counts.into_iter().enumerate() { + let expected = usize::from(problem.transfers[index].width == ExchangeItemWidth::Paired64); + if count != expected { + return Err(fail(format!( + "phase {} transfer {index} has {count} partner-busy activities, expected {expected}", + problem.phase + ))); + } + } + for (transfer, counts) in receive_counts.into_iter().enumerate() { + if counts.into_iter().any(|count| count != 1) { + return Err(fail(format!( + "phase {} transfer {transfer} does not have exactly one activity per destination", + problem.phase + ))); + } + } + Ok(()) ++======= + Ok(incoming_bases) ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec } /// Combines physically contiguous source and destination spans into one @@@ -1554,27 -648,18 +1634,40 @@@ impl<'a> TransferScheduler<'a> let candidate = self.ready.pop()?; let index = candidate.index.0; let transfer = &self.transfers[index]; ++<<<<<<< HEAD + let earliest_start = std::iter::once(self.dependency_ready[index]) + .chain(std::iter::once( + tile_availability[usize::from(transfer.source)].send, + )) + .chain(transfer.reserved_source.into_iter().map(|tile| { + let availability = tile_availability[usize::from(tile)]; + availability.send + })) + .chain( + transfer + .destinations + .iter() + .map(|&(tile, _)| tile_availability[usize::from(tile)].receive), + ) + .max() + .unwrap_or(0); ++======= + let earliest_start = + std::iter::once(self.dependency_ready[index]) + .chain(std::iter::once( + tile_availability[usize::from(transfer.physical.source)].send, + )) + .chain(transfer.physical.destinations.iter().map( + |&TransferEndpoint(tile, _)| tile_availability[usize::from(tile)].receive, + )) + .max() + .unwrap_or(0); ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec if candidate.earliest_start.0 == earliest_start { - return Some((index, earliest_start)); + // Endpoint availability ranks the ready queue, but is not a + // dependency on payload arrival. The row builder pipelines + // source selection and delivery using their actual timings. + return Some((index, self.dependency_ready[index])); } self.push_ready(index, earliest_start); } @@@ -1731,17 -803,14 +1811,25 @@@ struct MaterializedSchedule tile_availability: Vec, memory_accesses: Vec, activities: Vec>, ++<<<<<<< HEAD + scheduled_sends: Vec>, + order: Vec, + timings: Vec>, +} + +impl MaterializedSchedule { + fn new(tile_count: u16, transfers: &[PendingTransfer]) -> Self { + let transfer_count = transfers.len(); ++======= + scheduled_sends: Vec>, + } + + impl MaterializedSchedule { + fn new(tile_count: u16, validate_encoding: bool) -> Self { ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec Self { builder: PhaseProgramBuilder::new(tile_count), - horizon: 0, + validate_encoding, tile_availability: vec![TileAvailability::default(); usize::from(tile_count)], memory_accesses: (0..tile_count) .map(|_| TileMemorySchedule::default()) @@@ -1757,153 -824,95 +1843,177 @@@ &mut self, topology: &Topology, pending: &[PendingTransfer], - incoming_bases: &[u32], - receive_counts: &[usize], + incoming_bases: &[Option], index: usize, dependency_ready: u32, - validate_encoding: bool, - last_transfer: &mut [TilePredecessor], ) -> Result { let transfer = &pending[index]; ++<<<<<<< HEAD + let (blocking_tile, latest_availability) = std::iter::once(( + transfer.source, + self.tile_availability[usize::from(transfer.source)].send, + )) + .chain(transfer.reserved_source.into_iter().map(|tile| { + let availability = self.tile_availability[usize::from(tile)]; + (tile, availability.send) + })) + .chain( + transfer + .destinations + .iter() + .map(|&(tile, _)| (tile, self.tile_availability[usize::from(tile)].receive)), + ) + .max_by_key(|&(tile, availability)| (availability, Reverse(tile))) + .unwrap_or((transfer.source, 0)); + let blocking_tile = if dependency_ready > latest_availability { + transfer.source + } else { + blocking_tile + }; + let predecessor = if blocking_tile == transfer.source + || transfer.reserved_source == Some(blocking_tile) + { + last_transfer[usize::from(blocking_tile)].send + } else { + last_transfer[usize::from(blocking_tile)].receive + }; ++======= + let physical = &transfer.physical; + let latest_availability = + std::iter::once(self.tile_availability[usize::from(physical.source)].send) + .chain( + physical + .destinations + .iter() + .map(|&TransferEndpoint(tile, _)| { + self.tile_availability[usize::from(tile)].receive + }), + ) + .max() + .unwrap_or(0) + .max(dependency_ready); + let incoming_base = physical + .destinations + .first() + .and_then(|endpoint| incoming_bases[usize::from(endpoint.0)]); + let resolved = physical.resolve(topology, incoming_base)?; ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec let timing = append_transfer( - topology, &self.memory_accesses, ++<<<<<<< HEAD + incoming_bases, + receive_counts, + ScheduledTransfer { + source: transfer.source, + destinations: &transfer.destinations, + source_address: transfer.source_address(), + source_elements: &transfer.source_elements, + words: transfer.words, + width: transfer.width, + // Endpoint constraints are enforced at their actual source, + // payload and control events by the row builder. Only true + // memory dependencies constrain the whole transfer's release. + schedule_offset: dependency_ready, + }, ++======= + physical, + &resolved, + &transfer.source_elements, + latest_availability, ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec &mut self.builder, - validate_encoding, + self.validate_encoding, )?; - let payload_end = timing.sender_end; - self.memory_accesses[usize::from(transfer.source)] + let payload_end = timing.payload_end; + self.memory_accesses[usize::from(physical.source)] .sends .push(MemoryAccess { - start: timing.start, - end: timing.sender_memory_end, + start: timing.payload_start, + end: timing.sender_horizon, elements: transfer.source_elements.clone(), }); - for ((&(tile, address), &start), &memory_end) in transfer + for ((&TransferEndpoint(tile, address), &start), &memory_end) in physical .destinations .iter() - .zip(&timing.receiver_starts) - .zip(&timing.receiver_memory_ends) + .zip(&timing.receiver_payload_starts) + .zip(&timing.receiver_horizons) { self.memory_accesses[usize::from(tile)] .receives .push(MemoryAccess { start, end: memory_end, - elements: effective_memory_elements(address, transfer.words), + elements: memory_elements_for_words(address, physical.words).collect(), }); } - self.scheduled_sends[usize::from(transfer.source)] + self.scheduled_sends[usize::from(physical.source)] .push((transfer.source_shard, transfer.source_offset)); ++<<<<<<< HEAD + self.activities[usize::from(transfer.source)].push(ExchangeActivity { + fanout: transfer.destinations.len() as u16, + paired: transfer.width == ExchangeItemWidth::Paired64, ++======= + self.activities[usize::from(physical.source)].push(ExchangeActivity { ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec transfer: u32::try_from(index).map_err(|_| ExchangeLoweringError::Overflow)?, kind: ExchangeActivityKind::Send, - start_cycle: timing.start, + start_cycle: timing.payload_start, end_cycle: payload_end, - memory_end_cycle: timing.sender_memory_end, - address: transfer.source_address(), - words: transfer.words, + address: physical.source_address(), + words: physical.words, }); ++<<<<<<< HEAD + if let Some(tile) = transfer.reserved_source { + self.activities[usize::from(tile)].push(ExchangeActivity { + fanout: transfer.destinations.len() as u16, + paired: transfer.width == ExchangeItemWidth::Paired64, + transfer: u32::try_from(index).map_err(|_| ExchangeLoweringError::Overflow)?, + kind: ExchangeActivityKind::PartnerBusy, + start_cycle: timing.start, + end_cycle: timing.sender_memory_end, + memory_end_cycle: timing.sender_memory_end, + address: transfer.source_address(), + words: transfer.words, + }); + } + for (((&(tile, address), &start_cycle), &end_cycle), &memory_end_cycle) in transfer ++======= + for ((&TransferEndpoint(tile, address), &start_cycle), &end_cycle) in physical ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec .destinations .iter() - .zip(&timing.receiver_starts) - .zip(&timing.receiver_ends) - .zip(&timing.receiver_memory_ends) + .zip(&timing.receiver_payload_starts) + .zip(&timing.receiver_payload_ends) { self.activities[usize::from(tile)].push(ExchangeActivity { + fanout: transfer.destinations.len() as u16, + paired: transfer.width == ExchangeItemWidth::Paired64, transfer: u32::try_from(index).map_err(|_| ExchangeLoweringError::Overflow)?, kind: ExchangeActivityKind::Receive, start_cycle, end_cycle, - memory_end_cycle, address, - words: transfer.words, + words: physical.words, }); } ++<<<<<<< HEAD + self.tile_availability[usize::from(transfer.source)].send = timing.sender_end; + if let Some(tile) = transfer.reserved_source { + self.tile_availability[usize::from(tile)].send = timing.sender_memory_end; + last_transfer[usize::from(tile)].send = Some(index); + } + for (&(tile, _), &receiver_end) in transfer.destinations.iter().zip(&timing.receiver_ends) { ++======= + self.tile_availability[usize::from(physical.source)].send = timing.payload_end; + for (&TransferEndpoint(tile, _), &receiver_end) in physical + .destinations + .iter() + .zip(&timing.receiver_payload_ends) + { ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec self.tile_availability[usize::from(tile)].receive = receiver_end; } - last_transfer[usize::from(transfer.source)].send = Some(index); - for &(tile, _) in &transfer.destinations { - last_transfer[usize::from(tile)].receive = Some(index); - } - self.order.push(index); - self.timings[index] = Some(MaterializedTiming { - start: timing.start, - end: timing.end, - blocking_tile, - predecessor, - }); - Ok(timing.end) - } - - fn finish_horizon(&mut self) { - self.horizon = self.builder.event_cycles(); + Ok(timing.payload_completion()) } } @@@ -1947,139 -944,27 +2045,128 @@@ fn materialize_greedy_schedule_impl tile_count: u16, validate_encoding: bool, ) -> Result { ++<<<<<<< HEAD + let mut schedule = MaterializedSchedule::new(tile_count, pending); ++======= + let mut schedule = MaterializedSchedule::new(tile_count, validate_encoding); ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec let mut scheduler = TransferScheduler::new(pending, tile_count); - let mut last_transfer = vec![TilePredecessor::default(); usize::from(tile_count)]; while let Some((index, dependency_ready)) = scheduler.next(&schedule.tile_availability) { - let completion = schedule.append( - topology, - pending, - incoming_bases, - receive_counts, - index, - dependency_ready, - validate_encoding, - &mut last_transfer, - )?; + let completion = + schedule.append(topology, pending, incoming_bases, index, dependency_ready)?; scheduler.complete(index, completion); } debug_assert!(scheduler.is_complete()); Ok(schedule) } ++<<<<<<< HEAD +fn materialize_schedule_order( + topology: &Topology, + pending: &[PendingTransfer], + incoming_bases: &[u32], + receive_counts: &[usize], + tile_count: u16, + order: &[usize], + validate_encoding: bool, +) -> Result { + if order.len() != pending.len() { + return Err(ExchangeLoweringError::Overflow); + } + let mut schedule = MaterializedSchedule::new(tile_count, pending); + let mut last_transfer = vec![TilePredecessor::default(); usize::from(tile_count)]; + let dependencies = memory_dependencies(pending, tile_count); + let mut predecessors = vec![Vec::new(); pending.len()]; + for (before, after) in dependencies { + predecessors[after].push(before); + } + let mut completion = vec![None; pending.len()]; + for &index in order { + if index >= pending.len() + || completion[index].is_some() + || predecessors[index] + .iter() + .any(|&before| completion[before].is_none()) + { + return Err(ExchangeLoweringError::Invariant( + "exchange order is not a topological permutation".into(), + )); + } + let dependency_ready = predecessors[index] + .iter() + .filter_map(|predecessor| completion[*predecessor]) + .max() + .unwrap_or(0); + completion[index] = Some(schedule.append( + topology, + pending, + incoming_bases, + receive_counts, + index, + dependency_ready, + validate_encoding, + &mut last_transfer, + )?); + } + schedule.finish_horizon(); + if schedule_encoding_is_valid(&schedule)? { + Ok(schedule) + } else { + Err(ipu_exchange::ExchangeError::Schedule("SENDPICP instruction alignment").into()) + } +} + +fn schedule_encoding_is_valid( + schedule: &MaterializedSchedule, +) -> Result { + match schedule.builder.clone().finish() { + Ok(_) => Ok(true), + Err(ipu_exchange::ExchangeError::Schedule("SENDPICP instruction alignment")) => Ok(false), + Err(error) => Err(error.into()), + } +} + +fn schedule_score(schedule: &MaterializedSchedule) -> u32 { + schedule.horizon +} + +fn endpoint_work_lower_bound(pending: &[PendingTransfer], tile_count: u16) -> u32 { + let mut send_words = vec![0u64; usize::from(tile_count)]; + let mut receive_words = vec![0u64; usize::from(tile_count)]; + for transfer in pending { + let items = u64::from(transfer.item_count().unwrap_or(transfer.words)); + send_words[usize::from(transfer.source)] += items; + if let Some(tile) = transfer.reserved_source { + send_words[usize::from(tile)] += items; + } + for &(tile, _) in &transfer.destinations { + receive_words[usize::from(tile)] += items; + } + } + send_words + .into_iter() + .zip(receive_words) + .map(|(send, receive)| send.max(receive)) + .max() + .unwrap_or(0) + .min(u64::from(u32::MAX)) as u32 +} + +/// Orders a balanced point-to-point phase as maximum-cardinality waves over +/// its send and receive buses. The result remains only a candidate: the exact +/// row builder decides whether it improves the incumbent schedule. ++======= ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec fn append_transfer( - topology: &Topology, memory_accesses: &[TileMemorySchedule], - incoming_bases: &[u32], - receive_counts: &[usize], - transfer: ScheduledTransfer<'_>, + physical: &PhysicalTransfer, + resolved: &ResolvedTransfer, + source_elements: &[MemoryElement], + requested_offset: u32, builder: &mut PhaseProgramBuilder, validate_encoding: bool, - ) -> Result { - let ScheduledTransfer { - source, - destinations, - source_address, - source_elements, - words, - width, - schedule_offset: requested_offset, - } = transfer; + ) -> Result { + let words = physical.words; if words == 0 || source_elements.is_empty() { return Err(ExchangeLoweringError::UnalignedPayload); } @@@ -2235,56 -1046,161 +2248,215 @@@ fn memory_safe_transfer_offset Ok(safe_offset) } ++<<<<<<< HEAD +fn spans_share_effective_memory_element( + left_address: u32, + left_words: u32, + right_address: u32, + right_words: u32, +) -> bool { + effective_memory_elements(left_address, left_words) + .into_iter() + .any(|left| { + effective_memory_elements(right_address, right_words) + .into_iter() + .any(|right| left == right) + }) +} + +pub(crate) fn effective_memory_elements(address: u32, words: u32) -> Vec { + let end = address.saturating_add(words.saturating_mul(4)); + let mut elements = Vec::new(); + let mut cursor = address; + while cursor < end { + let interleaved = 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 = (cursor - base) / size; + elements.push(ExchangeMemoryElement { interleaved, index }); + let boundary = base.saturating_add((index + 1).saturating_mul(size)); + cursor = boundary.min(end); + } + elements +} + +struct ScheduledTransferTiming { + start: u32, + end: u32, + sender_end: u32, + sender_memory_end: u32, + receiver_starts: Vec, + receiver_ends: Vec, + receiver_memory_ends: Vec, +} + ++======= ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec pub fn inactive_exchange_program() -> Vec { vec![RETURN_M10_INSTRUCTION] } #[cfg(test)] ++<<<<<<< HEAD +mod tests; ++======= + mod tests { + use super::*; + use crate::{ + ComputeGraph, Ipu21CostModel, Layout, PipelineConfig, Precision, TensorFormat, lower, + lower_to_tiles, place, + }; + + #[test] + fn randomized_transfer_schedules_preserve_hazards_without_same_role_overlap() { + let mut random = fastrand::Rng::with_seed(0x736c_6f74); + for _ in 0..64 { + let tile_count = random.u16(2..=32); + let transfer_count = random.usize(1..=256); + let transfers = (0..transfer_count) + .map(|_| { + let source = random.u16(0..tile_count); + let receiver_count = random.usize(1..=usize::from(tile_count.min(8) - 1)); + let mut receivers = Vec::with_capacity(receiver_count); + while receivers.len() != receiver_count { + let tile = random.u16(0..tile_count); + if tile != source && !receivers.contains(&tile) { + receivers.push(tile); + } + } + let words = + random.u32(1..=HardwareTarget::Ipu21.exchange().maximum_transfer_words); + PendingTransfer { + physical: PhysicalTransfer { + source, + source_addresses: vec![0], + destinations: receivers + .into_iter() + .map(|tile| TransferEndpoint(tile, 0)) + .collect(), + words, + width: TransferWidth::Word32, + }, + source_shard: LowShardId::from_index(u32::from(source)), + source_offset: 0, + source_elements: memory_elements_for_words(0, words).collect(), + } + }) + .collect::>(); + let dependencies = memory_dependencies(&transfers, tile_count); + let mut scheduler = TransferScheduler::new(&transfers, tile_count); + let mut availability = vec![TileAvailability::default(); usize::from(tile_count)]; + let mut occurrences = vec![0u8; transfers.len()]; + let mut intervals = vec![(0u32, 0u32); transfers.len()]; + while let Some((index, start)) = scheduler.next(&availability) { + occurrences[index] += 1; + let transfer = &transfers[index]; + let end = start.saturating_add(transfers[index].physical.words); + intervals[index] = (start, end); + availability[usize::from(transfer.physical.source)].send = end; + for &TransferEndpoint(tile, _) in &transfer.physical.destinations { + availability[usize::from(tile)].receive = end; + } + scheduler.complete(index, end); + } + assert!(scheduler.is_complete()); + assert!(occurrences.into_iter().all(|count| count == 1)); + for &(before, after) in &dependencies { + assert!(intervals[before].1 <= intervals[after].0); + } + for tile in 0..tile_count { + let mut send_intervals = transfers + .iter() + .enumerate() + .filter(|(_, transfer)| transfer.physical.source == tile) + .map(|(index, _)| intervals[index]) + .collect::>(); + send_intervals.sort_unstable(); + assert!(send_intervals.windows(2).all(|pair| pair[0].1 <= pair[1].0)); + let mut receive_intervals = transfers + .iter() + .enumerate() + .filter(|(_, transfer)| { + transfer + .physical + .destinations + .iter() + .any(|&TransferEndpoint(destination, _)| destination == tile) + }) + .map(|(index, _)| intervals[index]) + .collect::>(); + receive_intervals.sort_unstable(); + assert!( + receive_intervals + .windows(2) + .all(|pair| pair[0].1 <= pair[1].0) + ); + } + } + } + + #[test] + fn randomized_gemm_exchanges_produce_one_executable_row_per_tile() { + let mut random = fastrand::Rng::with_seed(0x6578_6368); + for _ in 0..32 { + let tiles = 1_u16 << random.u32(1..=3); + let rows = u32::from(tiles) * random.u32(1..=8); + let columns = random.u32(1..=2) * 64; + let mut graph = ComputeGraph::new(); + let left = graph.host_input("left", [rows, 64]).unwrap(); + let right = graph.parameter("right", [64, columns]).unwrap(); + let output = graph.gemm(left, right).unwrap(); + graph.set_outputs([output]).unwrap(); + let config = PipelineConfig::new(tiles) + .with_input( + left, + TensorFormat { + precision: Precision::F16, + layout: Layout::amp_left(64, tiles), + }, + ) + .with_input( + right, + TensorFormat { + precision: Precision::F16, + layout: Layout::block_major_matrix(64, tiles), + }, + ); + let mid = lower(&graph, &config, &Ipu21CostModel).unwrap(); + let low = lower_to_tiles(&mid, &config).unwrap(); + let placement = place(&low).unwrap(); + let phases = lower_exchanges(&low, &placement, HardwareTarget::Ipu21).unwrap(); + assert_eq!(phases.len(), low.exchange_phases.len()); + for phase in phases { + assert_eq!(phase.programs.len(), usize::from(tiles)); + assert_eq!(phase.activities.len(), usize::from(tiles)); + assert!(phase.event_cycles != 0); + assert!(phase.activities.iter().flatten().next().is_some()); + for activities in &phase.activities { + for activity in activities { + assert!(activity.start_cycle < activity.end_cycle); + assert!(activity.end_cycle <= phase.event_cycles); + } + } + for program in &phase.programs { + let active = program.is_some(); + let program = program.as_deref().unwrap_or(&[RETURN_M10_INSTRUCTION]); + assert_eq!(program.last(), Some(&RETURN_M10_INSTRUCTION)); + assert_eq!(active, program.len() > 1); + assert_eq!(active, plan_event_cycles(program).unwrap() != 0); + assert!(plan_event_cycles(program).unwrap() <= phase.event_cycles); + assert!( + !program.contains(&ipu_target::instruction::SYNC_SUPERVISOR_INSTRUCTION) + ); + } + } + } + } + } ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec diff --cc crates/ipu-codegen/src/kernel/build.rs index 8493182,9de8ea9..0000000 --- a/crates/ipu-codegen/src/kernel/build.rs +++ b/crates/ipu-codegen/src/kernel/build.rs @@@ -1,135 -1,32 +1,167 @@@ ++<<<<<<< HEAD +//! Device object recipes for the selected kernel specializations. + +use super::*; + ++======= ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec #[derive(Clone, Debug, PartialEq, Eq)] pub struct KernelCompilation { pub source: &'static str, pub name: String, ++<<<<<<< HEAD + pub flags: Vec, + pub retained_symbols: Vec, +} + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct KernelBuildPlan { + pub compilations: Vec, + pub(super) symbols: BTreeMap, +} + +impl KernelBuildPlan { + /// Derives device objects from the finalized schedule, so row variants are + /// compiler specializations rather than a fixed collection of binaries. + pub fn from_program(program: &LowProgram) -> Result { + let mut inventory = KernelInventory::default(); + for tile in &program.tiles { + inventory.collect(program, tile)?; + } + Self::from_inventory(inventory) + } + + pub(super) fn from_inventory(inventory: KernelInventory) -> Result { + let KernelInventory { + rows, + gelu, + reduction_add, + rearrangements, + unpacks, + attention, + attention_stages, + } = inventory; + let mut plan = Self::default(); + for (configuration, rows) in rows { + plan.add_gemm(configuration, rows); + } + if gelu { + plan.compilations.push(KernelCompilation { + source: "gelu_f16.S", + name: "gelu_f16".into(), + flags: Vec::new(), + retained_symbols: vec!["ipu_stack_gelu_tanh_approx_f16".into()], + }); + } + if reduction_add { + plan.compilations.push(KernelCompilation { + source: "reduce_add_f16.S", + name: "reduce_add_f16".into(), + flags: Vec::new(), + retained_symbols: vec!["ipu_stack_reduce_sum_f16".into()], + }); + } + let has_worker_codelets = !rearrangements.is_empty() || !unpacks.is_empty(); + for shape in unpacks { + plan.add_unpack(shape); + } + for shape in rearrangements { + plan.add_rearrangement(shape); + } + if has_worker_codelets || !attention.is_empty() || !attention_stages.is_empty() { + plan.compilations.push(KernelCompilation { + source: "worker_support.S", + name: "worker_support".into(), + flags: Vec::new(), + retained_symbols: Vec::new(), + }); + } + for shape in attention { + plan.add_attention(shape); + } + plan.add_attention_stages(attention_stages)?; + Ok(plan) + } + + /// Marshal supervisor registers into the C++ vertex argument block. + /// The register order follows the vertex fields, independently of call ABI. + pub(super) fn add_worker_wrapper( + &mut self, + name: String, + symbol: &str, + vertex: &str, + registers: &[u8], + ) { + let arguments = registers + .iter() + .map(|register| format!("$m{register}")) + .collect::>() + .join(","); + let frame_bytes = (registers.len() * 4).next_multiple_of(16); + self.compilations.push(KernelCompilation { + source: "worker_call.S", + name, + flags: vec![ + format!("-DWORKER_CALL_SYMBOL={symbol}"), + format!("-DWORKER_CODELET_SYMBOL=__runCodelet_{vertex}"), + format!("-DWORKER_ARGUMENTS={arguments}"), + format!("-DWORKER_FRAME_BYTES={frame_bytes}"), + ], + retained_symbols: vec![symbol.to_owned()], + }); + } + + pub fn call(&self, run: &KernelRun) -> Result { + let abi = validate_kernel_run(run)?; + let kernel = &run.kernel; + if abi.availability != KernelAvailability::Implemented { + return Err(KernelAbiError::Unavailable(kernel.clone())); + } + let symbol = match abi.symbols { + KernelSymbols::Exact(symbol) => symbol.to_owned(), + _ => self + .symbols + .get(&KernelSpecialization::from_run(run)?) + .cloned() + .ok_or(KernelAbiError::RequirementMismatch)?, + }; + Ok(PlannedKernelCall { + symbol, + arguments: scalar_values(run, &abi)?, + }) + } + + pub fn retained_symbols(&self) -> impl Iterator { + self.compilations + .iter() + .flat_map(|compilation| compilation.retained_symbols.iter().map(String::as_str)) ++======= + pub optimization: Option, + pub definitions: Vec<(&'static str, String)>, + pub retained_symbols: Vec, + } + + #[derive(Clone, Copy, Debug, PartialEq, Eq)] + pub enum KernelOptimization { + Size, + Speed, + } + + impl KernelCompilation { + /// Produces textual options only at the toolchain boundary. + pub fn compiler_flags(&self) -> Vec { + self.optimization + .map(|optimization| match optimization { + KernelOptimization::Size => "-Os".to_owned(), + KernelOptimization::Speed => "-O2".to_owned(), + }) + .into_iter() + .chain( + self.definitions + .iter() + .map(|(name, value)| format!("-D{name}={value}")), + ) + .collect() ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec } } diff --cc crates/ipu-codegen/src/kernel/mod.rs index 3cc11b5,5d34c35..0000000 --- a/crates/ipu-codegen/src/kernel/mod.rs +++ b/crates/ipu-codegen/src/kernel/mod.rs @@@ -1,29 -1,88 +1,117 @@@ ++<<<<<<< HEAD +//! Kernel ABI, specialization recipes, and placed call materialization. + +mod abi; +mod attention; +mod build; +pub(crate) mod cost; +mod gemm; +mod rearrange; +mod specialization; +#[cfg(test)] +mod tests; +pub(crate) use abi::*; +pub(crate) use build::*; +use specialization::*; + +use crate::{ + AMP_COLUMN_MICRO, AMP_INNER_BLOCK, AttentionKernelShape, KernelAbiError, attention_shape, + gemm_rows, input_matrix_extent, matrix_count, matrix_extent, +}; +use crate::{ + AmpOrder, BlockMajorOrder, BlockValue, BlockValueId, ComputeStep, ElementOrder, GemmKernelMode, + GemmWeightLoad, KernelRequirements, KernelRun, LowProgram, Precision, StepProfile, + StorageError, TileAddress, TileKernelSpec, TileWorkList, TileWorkRef, view_byte_spans, +}; +use std::collections::{BTreeMap, BTreeSet}; + ++======= + //! Tile-kernel ABI, specialization builds, and placed call materialization. + + mod build; + mod materialize; + mod spec; + + pub use build::*; + pub use materialize::materialize_kernel_run; + pub use spec::*; + + #[cfg(test)] + use crate::MemorySpaceRequirements; + use crate::layout::{AMP_COLUMN_MICRO, AMP_INNER_BLOCK}; + use crate::{ + GemmKernelMode, GemmWeightLoad, KernelRequirements, KernelRun, LowProgram, NativeKernelOrder, + Precision, StorageError, StorageOrder, TileWorkList, TileWorkRef, + }; + use std::collections::{BTreeMap, BTreeSet}; + + pub const OUTPUT_REGISTER: u8 = 2; + pub const FIRST_INPUT_REGISTER: u8 = 3; + pub const RETURN_REGISTER: u8 = 10; + + #[derive(Clone, Debug, Default, PartialEq, Eq)] + pub struct KernelBuildPlan { + pub compilations: Vec, + symbols: BTreeMap, + } + + #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)] + enum RearrangeTarget { + AmpLeft, + AmpTransposedRight, + BlockMajor { row_block: u16, column_block: u16 }, + } + + #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)] + enum UnpackSource { + AmpOutput, + AmpTransposedLeft, + } + + impl UnpackSource { + fn from_order(order: StorageOrder) -> Option { + match order { + StorageOrder::Native(NativeKernelOrder::Output) => Some(Self::AmpOutput), + StorageOrder::Native(NativeKernelOrder::TransposedLeft) => { + Some(Self::AmpTransposedLeft) + } + _ => None, + } + } + + const fn codelet_index(self) -> u32 { + match self { + Self::AmpOutput => 0, + Self::AmpTransposedLeft => 1, + } + } + } + + impl RearrangeTarget { + fn from_order(order: StorageOrder) -> Option { + match order { + StorageOrder::Native(NativeKernelOrder::Left) => Some(Self::AmpLeft), + StorageOrder::Native(NativeKernelOrder::TransposedRight) => { + Some(Self::AmpTransposedRight) + } + StorageOrder::Blocked(order) if order.is_matrix() => Some(Self::BlockMajor { + row_block: order.block_shape[0], + column_block: order.block_shape[1], + }), + _ => None, + } + } + + const fn codelet_index(self) -> u32 { + match self { + Self::AmpLeft => 0, + Self::AmpTransposedRight => 1, + Self::BlockMajor { .. } => 2, + } + } + } + ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec #[derive(Clone, Debug, PartialEq, Eq)] pub struct PlannedKernelCall { pub symbol: String, @@@ -31,6 -90,30 +119,33 @@@ } #[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)] ++<<<<<<< HEAD ++======= + pub enum KernelAbiError { + #[error("kernel requirements do not match the tile-kernel family")] + RequirementMismatch, + #[error("kernel run has {actual} pointer operands, ABI requires {expected}")] + PointerArity { expected: usize, actual: usize }, + #[error("kernel operand {0} is fragmented into multiple views")] + FragmentedOperand(usize), + #[error("kernel {0:?} has no device implementation")] + Unavailable(TileKernelSpec), + #[error("GEMM output view does not have a matrix row axis")] + MissingGemmRows, + #[error("kernel element count overflowed")] + ElementCountOverflow, + #[error("GEMM row count {0} is not present in the compilation plan")] + UnplannedGemmRows(u32), + #[error("kernel {symbol} requires an element count divisible by {divisor}, got {count}")] + UnsupportedElementCount { + symbol: &'static str, + count: u32, + divisor: u32, + }, + } + + #[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)] ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec pub enum KernelMaterializationError { #[error(transparent)] Abi(#[from] KernelAbiError), @@@ -44,88 -127,1295 +159,1383 @@@ AddressOverflow, } ++<<<<<<< HEAD +/// 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: &[BlockValue], + 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() + .or_else(|| { + shard_addresses + .get(&view.shard) + .copied() + .map(TileAddress::Absolute) + }) + .ok_or(KernelMaterializationError::UnplacedShard( + view.shard.index(), + ))?; + add_address_offset(base, span.offset) + }; + let mut output_address = resolve(&run.output)?; + if let TileKernelSpec::FillZero { offset, bytes, .. } = run.kernel { + let output_spans = + view_byte_spans(&shards[run.output.shard.index() as usize], &run.output)?; + let allocation_bytes = output_spans[0].bytes; + if !offset.is_multiple_of(8) + || offset + .checked_add(bytes) + .is_none_or(|end| end > allocation_bytes) + { + return Err(StorageError::InvalidView.into()); + } + output_address = add_address_offset(output_address, offset)?; + } + 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)?, + }, + }) +} ++======= + fn specialized_gemm_symbol( + prefix: &str, + mode: GemmKernelMode, + weight_suffix: &str, + inner_block: u32, + output_columns: u32, + rows: u32, + ) -> String { + let operation = match mode { + GemmKernelMode::Initialize => "init", + GemmKernelMode::Accumulate => "accumulate", + }; + format!( + "ipu_stack_gemm_{prefix}_{operation}{weight_suffix}_k{inner_block}_c{output_columns}_r{rows}" + ) + } + + fn specialized_kernel_symbol(kernel: &TileKernelSpec) -> Result { + match kernel { + TileKernelSpec::Gemm { + multiply, + mode, + weights, + inner_block, + output_columns, + rows, + .. + } => { + let prefix = match multiply { + Precision::F16 => "f16", + Precision::F32 => "f32", + Precision::F8F143 { .. } => return Err(KernelAbiError::RequirementMismatch), + }; + let weights = if *weights == GemmWeightLoad::Interleaved { + "_interleaved" + } else { + "" + }; + Ok(specialized_gemm_symbol( + prefix, + *mode, + weights, + *inner_block, + *output_columns, + *rows, + )) + } + TileKernelSpec::Rearrange { + from, + to, + logical_rows, + physical_rows, + logical_columns, + physical_columns, + .. + } => { + let (prefix, order, logical_rows, physical_rows, logical_columns, physical_columns) = + if from.order == StorageOrder::Linear { + let target = RearrangeTarget::from_order(to.order) + .ok_or(KernelAbiError::RequirementMismatch)?; + let (_, logical_rows, physical_rows, logical_columns, physical_columns) = + rearrangement_specialization( + target, + *logical_rows, + *physical_rows, + *logical_columns, + *physical_columns, + ); + ( + "ipu_stack_rearrange_row_major_to_amp_f16", + target.codelet_index(), + logical_rows, + physical_rows, + logical_columns, + physical_columns, + ) + } else { + let source = UnpackSource::from_order(from.order) + .ok_or(KernelAbiError::RequirementMismatch)?; + ( + "ipu_stack_unpack_amp_to_row_major_f16", + source.codelet_index(), + *logical_rows, + *physical_rows, + *logical_columns, + *physical_columns, + ) + }; + Ok(format!( + "{prefix}_o{order}_r{logical_rows}_p{physical_rows}_c{logical_columns}_p{physical_columns}" + )) + } + _ => Err(KernelAbiError::RequirementMismatch), + } + } + + fn attention_kernel_symbol(kernel: &TileKernelSpec) -> Result { + match kernel { + TileKernelSpec::AttentionSoftmax { + query_rows, + head_dimension, + key_columns, + padded_key_columns, + } => Ok(format!( + "ipu_stack_attention_softmax_f16_q{query_rows}_d{head_dimension}_k{key_columns}_p{padded_key_columns}" + )), + TileKernelSpec::AttentionMerge { + query_rows, + value_dimension, + padded_value_dimension, + key_block_columns, + .. + } => Ok(format!( + "ipu_stack_attention_merge_f16_q{query_rows}_v{value_dimension}_p{padded_value_dimension}_k{key_block_columns}" + )), + _ => Err(KernelAbiError::RequirementMismatch), + } + } + + impl KernelBuildPlan { + /// Derives device objects from the finalized schedule, so row variants are + /// compiler specializations rather than a fixed collection of binaries. + pub fn from_program(program: &LowProgram) -> Result { + let mut rows = BTreeMap::<(Precision, GemmWeightLoad, u32, u32), BTreeSet>::new(); + let mut gelu = false; + let mut reduction_add = false; + let mut rearrangements = BTreeSet::new(); + let mut unpacks = BTreeSet::new(); + let mut attention_stages = Vec::new(); + let mut planned_kernels = BTreeSet::new(); + for tile in &program.tiles { + collect_kernels( + program, + tile, + &mut rows, + &mut gelu, + &mut reduction_add, + &mut rearrangements, + &mut unpacks, + &mut attention_stages, + &mut planned_kernels, + )?; + } + let mut plan = Self::default(); + for ((precision, weights, inner_block, output_columns), values) in rows { + let values = values.into_iter().collect::>(); + let (source, prefix) = match precision { + Precision::F16 => ("gemm_f16_amp.S", "f16"), + Precision::F32 => ("gemm_f32_64_amp.S", "f32"), + Precision::F8F143 { .. } => continue, + }; + let weight_suffix = if weights == GemmWeightLoad::Interleaved { + "_interleaved" + } else { + "" + }; + for pair in values.chunks(2) { + let small = pair[0]; + let large = *pair.last().expect("nonempty GEMM row pair"); + let symbols = [ + (GemmKernelMode::Initialize, small), + (GemmKernelMode::Initialize, large), + (GemmKernelMode::Accumulate, small), + (GemmKernelMode::Accumulate, large), + ] + .map(|(mode, rows)| { + specialized_gemm_symbol( + prefix, + mode, + weight_suffix, + inner_block, + output_columns, + rows, + ) + }); + let single_rows = pair.len() == 1; + let mut definitions = vec![ + ("GEMM_SMALL_ROWS", small.to_string()), + ("GEMM_LARGE_ROWS", large.to_string()), + ("GEMM_OUTPUT_COLUMNS", output_columns.to_string()), + ("GEMM_INNER_BLOCK_DIMENSION", inner_block.to_string()), + ("GEMM_INIT_SMALL_SYMBOL", symbols[0].clone()), + ("GEMM_INIT_LARGE_SYMBOL", symbols[1].clone()), + ("GEMM_ACCUMULATE_SMALL_SYMBOL", symbols[2].clone()), + ("GEMM_ACCUMULATE_LARGE_SYMBOL", symbols[3].clone()), + ]; + if single_rows { + definitions.push(("GEMM_SINGLE_ROWS", "1".into())); + } + if weights == GemmWeightLoad::Interleaved { + definitions.push(("GEMM_INTERLEAVED_WEIGHTS", "1".into())); + } + let retained_symbols = if single_rows { + vec![symbols[0].clone(), symbols[2].clone()] + } else { + symbols.into_iter().collect() + }; + plan.compilations.push(KernelCompilation { + source, + name: format!( + "gemm_{prefix}{weight_suffix}_k{inner_block}_c{output_columns}_r{small}_r{large}" + ), + optimization: None, + definitions, + retained_symbols, + }); + } + } + if gelu { + plan.compilations.push(KernelCompilation { + source: "gelu_f16.S", + name: "gelu_f16".into(), + optimization: None, + definitions: Vec::new(), + retained_symbols: vec!["ipu_stack_gelu_tanh_approx_f16".into()], + }); + } + if reduction_add { + plan.compilations.push(KernelCompilation { + source: "reduce_add_f16.S", + name: "reduce_add_f16".into(), + optimization: None, + definitions: Vec::new(), + retained_symbols: vec!["ipu_stack_reduce_sum_f16".into()], + }); + } + let has_worker_codelets = !rearrangements.is_empty() || !unpacks.is_empty(); + for (order, logical_rows, physical_rows, logical_columns, physical_columns) in unpacks { + let order_index = order.codelet_index(); + let suffix = format!( + "o{order_index}_r{logical_rows}_p{physical_rows}_c{logical_columns}_p{physical_columns}" + ); + let vertex = format!("UnpackAmpToRowMajorF16_{suffix}"); + let codelet = format!("__runCodelet_{vertex}"); + let call = format!("ipu_stack_unpack_amp_to_row_major_f16_{suffix}"); + plan.compilations.push(KernelCompilation { + source: "unpack_amp_f16.cpp", + name: format!("unpack_amp_f16_codelet_{suffix}"), + optimization: Some(KernelOptimization::Speed), + definitions: vec![ + ("UNPACK_SOURCE_ORDER", order_index.to_string()), + ("UNPACK_LOGICAL_ROWS", logical_rows.to_string()), + ("UNPACK_PHYSICAL_ROWS", physical_rows.to_string()), + ("UNPACK_LOGICAL_COLUMNS", logical_columns.to_string()), + ("UNPACK_PHYSICAL_COLUMNS", physical_columns.to_string()), + ("UNPACK_VERTEX_NAME", vertex), + ], + retained_symbols: Vec::new(), + }); + plan.compilations.push(KernelCompilation { + source: "rearrange_f16.S", + name: format!("unpack_amp_f16_wrapper_{suffix}"), + optimization: None, + definitions: vec![ + ("REARRANGE_CALL_SYMBOL", call.clone()), + ("REARRANGE_CODELET_SYMBOL", codelet), + ], + retained_symbols: vec![call.clone()], + }); + } + for (order, logical_rows, physical_rows, logical_columns, physical_columns) in + rearrangements + { + let order_index = order.codelet_index(); + let (row_block, column_block) = match order { + RearrangeTarget::BlockMajor { + row_block, + column_block, + } => (row_block, column_block), + _ => (AMP_INNER_BLOCK as u16, AMP_COLUMN_MICRO as u16), + }; + let suffix = format!( + "o{order_index}_r{logical_rows}_p{physical_rows}_c{logical_columns}_p{physical_columns}" + ); + let vertex = format!("RearrangeRowMajorToAmpF16_{suffix}"); + let codelet = format!("__runCodelet_{vertex}"); + let call = format!("ipu_stack_rearrange_row_major_to_amp_f16_{suffix}"); + if order == RearrangeTarget::AmpLeft + && logical_columns.is_multiple_of(2) + && physical_columns.is_multiple_of(AMP_COLUMN_MICRO) + { + plan.compilations.push(KernelCompilation { + source: "rearrange_amp_left_f16.S", + name: format!("rearrange_amp_left_f16_{suffix}"), + optimization: None, + definitions: vec![ + ("REARRANGE_CALL_SYMBOL", call.clone()), + ("REARRANGE_LOGICAL_ROWS", logical_rows.to_string()), + ("REARRANGE_PHYSICAL_ROWS", physical_rows.to_string()), + ("REARRANGE_LOGICAL_COLUMNS", logical_columns.to_string()), + ("REARRANGE_PHYSICAL_COLUMNS", physical_columns.to_string()), + ], + retained_symbols: vec![call.clone()], + }); + continue; + } + if order + == (RearrangeTarget::BlockMajor { + row_block: AMP_INNER_BLOCK as u16, + column_block: AMP_COLUMN_MICRO as u16, + }) + && physical_rows == AMP_INNER_BLOCK + && logical_columns.is_multiple_of(4) + && physical_columns.is_multiple_of(AMP_COLUMN_MICRO) + { + plan.compilations.push(KernelCompilation { + source: "rearrange_block_major_f16.S", + name: format!("rearrange_block_major_f16_{suffix}"), + optimization: None, + definitions: vec![ + ("REARRANGE_CALL_SYMBOL", call.clone()), + ("REARRANGE_PHYSICAL_COLUMNS", physical_columns.to_string()), + ], + retained_symbols: vec![call.clone()], + }); + continue; + } + if order == RearrangeTarget::AmpTransposedRight + && logical_rows == 64 + && physical_rows == 64 + && logical_columns == 16 + && physical_columns == 16 + { + plan.compilations.push(KernelCompilation { + source: "rearrange_transposed_right_f16.S", + name: format!("rearrange_transposed_right_f16_{suffix}"), + optimization: None, + definitions: vec![("REARRANGE_CALL_SYMBOL", call.clone())], + retained_symbols: vec![call.clone()], + }); + continue; + } + plan.compilations.push(KernelCompilation { + source: "rearrange_f16.cpp", + name: format!("rearrange_f16_codelet_{suffix}"), + optimization: Some(KernelOptimization::Speed), + definitions: vec![ + ("REARRANGE_TARGET_ORDER", order_index.to_string()), + ("REARRANGE_LOGICAL_ROWS", logical_rows.to_string()), + ("REARRANGE_PHYSICAL_ROWS", physical_rows.to_string()), + ("REARRANGE_LOGICAL_COLUMNS", logical_columns.to_string()), + ("REARRANGE_PHYSICAL_COLUMNS", physical_columns.to_string()), + ("REARRANGE_INNER_DIMENSION", AMP_COLUMN_MICRO.to_string()), + ("REARRANGE_ROW_BLOCK", row_block.to_string()), + ("REARRANGE_COLUMN_BLOCK", column_block.to_string()), + ("REARRANGE_VERTEX_NAME", vertex), + ], + retained_symbols: Vec::new(), + }); + plan.compilations.push(KernelCompilation { + source: "rearrange_f16.S", + name: format!("rearrange_f16_wrapper_{suffix}"), + optimization: None, + definitions: vec![ + ("REARRANGE_CALL_SYMBOL", call.clone()), + ("REARRANGE_CODELET_SYMBOL", codelet), + ], + retained_symbols: vec![call.clone()], + }); + } + if has_worker_codelets || !attention_stages.is_empty() { + plan.compilations.push(KernelCompilation { + source: "worker_support.S", + name: "worker_support".into(), + optimization: None, + definitions: Vec::new(), + retained_symbols: Vec::new(), + }); + } + if !attention_stages.is_empty() { + let mut query_rows = attention_stages + .iter() + .map(|(_, rows)| *rows) + .collect::>(); + let small_query = query_rows + .pop_first() + .ok_or(KernelAbiError::RequirementMismatch)?; + let large_query = query_rows.pop_last().unwrap_or(small_query); + let mut key_rows = attention_stages + .iter() + .filter_map(|(kernel, _)| match kernel { + TileKernelSpec::AttentionSoftmax { key_columns, .. } => Some(*key_columns), + _ => None, + }) + .collect::>(); + let small_key = key_rows + .pop_first() + .ok_or(KernelAbiError::RequirementMismatch)?; + let large_key = key_rows.pop_last().unwrap_or(small_key); + let configuration: Option<(u32, u32, u32, u32)> = + attention_stages + .iter() + .fold(None, |configuration, (kernel, _)| match kernel { + TileKernelSpec::AttentionSoftmax { + head_dimension, + padded_key_columns, + .. + } => Some(configuration.unwrap_or(( + *head_dimension, + 0, + 0, + *padded_key_columns, + ))), + TileKernelSpec::AttentionMerge { + value_dimension, + padded_value_dimension, + key_block_columns, + .. + } => { + let mut value = configuration.unwrap_or(( + 0, + *value_dimension, + *padded_value_dimension, + *key_block_columns, + )); + value.1 = *value_dimension; + value.2 = *padded_value_dimension; + value.3 = *key_block_columns; + Some(value) + } + _ => configuration, + }); + let (head_dimension, value_dimension, padded_value_dimension, key_block_columns) = + configuration.ok_or(KernelAbiError::RequirementMismatch)?; + let assembly_softmax_keys = attention_stages + .iter() + .filter_map(|(kernel, _)| match kernel { + TileKernelSpec::AttentionSoftmax { + key_columns, + padded_key_columns, + .. + } if key_columns != padded_key_columns => Some(*key_columns), + _ => None, + }) + .collect::>(); + let mut softmax_cpp_symbols = Vec::new(); + let mut softmax_assembly_symbols = Vec::new(); + let mut merge_symbols = Vec::new(); + for (kernel, _) in attention_stages { + let symbol = attention_kernel_symbol(&kernel)?; + let retained_symbols = match &kernel { + TileKernelSpec::AttentionSoftmax { key_columns, .. } + if assembly_softmax_keys.contains(key_columns) => + { + &mut softmax_assembly_symbols + } + TileKernelSpec::AttentionSoftmax { .. } => &mut softmax_cpp_symbols, + TileKernelSpec::AttentionMerge { .. } => &mut merge_symbols, + _ => return Err(KernelAbiError::RequirementMismatch), + }; + if !retained_symbols.contains(&symbol) { + retained_symbols.push(symbol.clone()); + } + plan.symbols.insert(kernel, symbol); + } + let softmax_symbol = |query_rows, key_columns| { + format!( + "ipu_stack_attention_softmax_f16_q{query_rows}_d{head_dimension}_k{key_columns}_p{key_block_columns}" + ) + }; + let merge_symbol = |query_rows| { + format!( + "ipu_stack_attention_merge_f16_q{query_rows}_v{value_dimension}_p{padded_value_dimension}_k{key_block_columns}" + ) + }; + let small_small = softmax_symbol(small_query, small_key); + let large_small = if large_query == small_query { + format!("{small_small}_alternate_query") + } else { + softmax_symbol(large_query, small_key) + }; + let small_large = if large_key == small_key { + format!("{small_small}_alternate_key") + } else { + softmax_symbol(small_query, large_key) + }; + let large_large = if large_query == small_query || large_key == small_key { + format!("{small_small}_alternate_query_key") + } else { + softmax_symbol(large_query, large_key) + }; + let small_merge = merge_symbol(small_query); + let large_merge = if large_query == small_query { + format!("{small_merge}_alternate_query") + } else { + merge_symbol(large_query) + }; + let entry_definitions = vec![ + ( + "ATTENTION_SOFTMAX_SMALL_QUERY_SMALL_KEY_SYMBOL", + small_small, + ), + ( + "ATTENTION_SOFTMAX_LARGE_QUERY_SMALL_KEY_SYMBOL", + large_small, + ), + ( + "ATTENTION_SOFTMAX_SMALL_QUERY_LARGE_KEY_SYMBOL", + small_large, + ), + ( + "ATTENTION_SOFTMAX_LARGE_QUERY_LARGE_KEY_SYMBOL", + large_large, + ), + ("ATTENTION_MERGE_SMALL_QUERY_SYMBOL", small_merge), + ("ATTENTION_MERGE_LARGE_QUERY_SYMBOL", large_merge), + ]; + let scale_bits = (1.0_f32 / (head_dimension as f32).sqrt()).to_bits(); + let softmax_definitions = vec![ + ("ATTENTION_HEAD_DIMENSION", head_dimension.to_string()), + ("ATTENTION_KEY_BLOCK_COLUMNS", key_block_columns.to_string()), + ("ATTENTION_SMALL_QUERY_ROWS", small_query.to_string()), + ("ATTENTION_LARGE_QUERY_ROWS", large_query.to_string()), + ("ATTENTION_SMALL_KEY_ROWS", small_key.to_string()), + ("ATTENTION_LARGE_KEY_ROWS", large_key.to_string()), + ]; + plan.compilations.push(KernelCompilation { + source: "attention_softmax_f16.cpp", + name: format!("attention_softmax_q{small_query}_q{large_query}_d{head_dimension}"), + optimization: Some(KernelOptimization::Size), + definitions: softmax_definitions, + retained_symbols: Vec::new(), + }); + plan.compilations.push(KernelCompilation { + source: "attention_softmax_f16_wrapper.S", + name: "attention_softmax_wrapper".into(), + optimization: None, + definitions: entry_definitions + .iter() + .take(4) + .cloned() + .chain( + [ + assembly_softmax_keys + .contains(&small_key) + .then(|| ("ATTENTION_USE_ASSEMBLY_SMALL_KEY", "1".into())), + (large_key != small_key && assembly_softmax_keys.contains(&large_key)) + .then(|| ("ATTENTION_USE_ASSEMBLY_LARGE_KEY", "1".into())), + ] + .into_iter() + .flatten(), + ) + .collect(), + retained_symbols: softmax_cpp_symbols, + }); + let mut attention_stage_definitions = vec![ + ("ATTENTION_HEAD_DIMENSION", head_dimension.to_string()), + ("ATTENTION_VALUE_DIMENSION", value_dimension.to_string()), + ( + "ATTENTION_PADDED_VALUE_DIMENSION", + padded_value_dimension.to_string(), + ), + ("ATTENTION_KEY_BLOCK_COLUMNS", key_block_columns.to_string()), + ("ATTENTION_SMALL_QUERY_ROWS", small_query.to_string()), + ("ATTENTION_LARGE_QUERY_ROWS", large_query.to_string()), + ("ATTENTION_SMALL_KEY_ROWS", small_key.to_string()), + ("ATTENTION_LARGE_KEY_ROWS", large_key.to_string()), + ("ATTENTION_SCALE_BITS", format!("0x{scale_bits:08x}")), + ]; + attention_stage_definitions.extend(entry_definitions); + if assembly_softmax_keys.contains(&small_key) { + attention_stage_definitions + .push(("ATTENTION_BUILD_ASSEMBLY_SOFTMAX_SMALL_KEY", "1".into())); + } + if large_key != small_key && assembly_softmax_keys.contains(&large_key) { + attention_stage_definitions + .push(("ATTENTION_BUILD_ASSEMBLY_SOFTMAX_LARGE_KEY", "1".into())); + } + merge_symbols.extend(softmax_assembly_symbols); + plan.compilations.push(KernelCompilation { + source: "attention_stages_f16.S", + name: format!( + "attention_stages_q{small_query}_q{large_query}_d{head_dimension}_v{value_dimension}" + ), + optimization: Some(KernelOptimization::Speed), + definitions: attention_stage_definitions, + retained_symbols: merge_symbols, + }); + } + for kernel in planned_kernels { + if !plan.symbols.contains_key(&kernel) { + plan.symbols + .insert(kernel.clone(), specialized_kernel_symbol(&kernel)?); + } + } + Ok(plan) + } + + pub fn call(&self, run: &KernelRun) -> Result { + let abi = validate_kernel_run(run)?; + let kernel = &run.kernel; + if abi.availability != KernelAvailability::Implemented { + return Err(KernelAbiError::Unavailable(kernel.clone())); + } + let symbol = match abi.symbols { + KernelSymbols::Exact(symbol) => symbol.to_owned(), + KernelSymbols::Planned => self + .symbols + .get(kernel) + .cloned() + .ok_or(KernelAbiError::RequirementMismatch)?, + }; + Ok(PlannedKernelCall { + symbol, + arguments: scalar_values(run, &abi)?, + }) + } + + pub fn retained_symbols(&self) -> impl Iterator { + self.compilations + .iter() + .flat_map(|compilation| compilation.retained_symbols.iter().map(String::as_str)) + } + } + + fn collect_kernels( + program: &LowProgram, + tile: &TileWorkList, + rows: &mut BTreeMap<(Precision, GemmWeightLoad, u32, u32), BTreeSet>, + gelu: &mut bool, + reduction_add: &mut bool, + rearrangements: &mut BTreeSet<(RearrangeTarget, u32, u32, u32, u32)>, + unpacks: &mut BTreeSet<(UnpackSource, u32, u32, u32, u32)>, + attention_stages: &mut Vec<(TileKernelSpec, u32)>, + planned_kernels: &mut BTreeSet, + ) -> Result<(), KernelAbiError> { + for work in program.work(tile) { + match work { + TileWorkRef::Kernel(run) => { + let abi = validate_kernel_run(run)?; + let kernel = &run.kernel; + if abi.availability != KernelAvailability::Implemented { + return Err(KernelAbiError::Unavailable(kernel.clone())); + } + if abi.symbols == KernelSymbols::Planned { + planned_kernels.insert(kernel.clone()); + } + if let TileKernelSpec::Gemm { + multiply, + weights, + inner_block, + output_columns, + rows: kernel_rows, + .. + } = kernel + { + rows.entry((*multiply, *weights, *inner_block, *output_columns)) + .or_default() + .insert(*kernel_rows); + } else if matches!(kernel, TileKernelSpec::Gelu) { + *gelu = true; + } else if matches!(kernel, TileKernelSpec::ReductionSum { .. }) { + *reduction_add = true; + } else if let TileKernelSpec::Rearrange { + from: + crate::Layout { + order: StorageOrder::Linear, + .. + }, + to: crate::Layout { order, .. }, + logical_rows, + physical_rows, + logical_columns, + physical_columns, + .. + } = kernel + && let Some(target) = RearrangeTarget::from_order(*order) + { + rearrangements.insert(rearrangement_specialization( + target, + *logical_rows, + *physical_rows, + *logical_columns, + *physical_columns, + )); + } else if let TileKernelSpec::Rearrange { + from: crate::Layout { order, .. }, + to: + crate::Layout { + order: StorageOrder::Linear, + .. + }, + logical_rows, + physical_rows, + logical_columns, + physical_columns, + .. + } = kernel + && let Some(source) = UnpackSource::from_order(*order) + { + unpacks.insert(( + source, + *logical_rows, + *physical_rows, + *logical_columns, + *physical_columns, + )); + } else if matches!( + kernel, + TileKernelSpec::AttentionSoftmax { .. } | TileKernelSpec::AttentionMerge { .. } + ) { + let query_rows = match kernel { + TileKernelSpec::AttentionSoftmax { query_rows, .. } + | TileKernelSpec::AttentionMerge { query_rows, .. } => *query_rows, + _ => unreachable!(), + }; + let stage = (kernel.clone(), query_rows); + if !attention_stages.contains(&stage) { + attention_stages.push(stage); + } + } + } + TileWorkRef::Repeat(repeat) => collect_kernels( + program, + &repeat.body, + rows, + gelu, + reduction_add, + rearrangements, + unpacks, + attention_stages, + planned_kernels, + )?, + TileWorkRef::Exchange(_) | TileWorkRef::LocalCopy(_) | TileWorkRef::Checkpoint(..) => {} + } + } + Ok(()) + } + + fn rearrangement_specialization( + order: RearrangeTarget, + logical_rows: u32, + physical_rows: u32, + logical_columns: u32, + physical_columns: u32, + ) -> (RearrangeTarget, u32, u32, u32, u32) { + if physical_rows == AMP_INNER_BLOCK + && logical_rows < physical_rows + && matches!( + order, + RearrangeTarget::AmpTransposedRight | RearrangeTarget::BlockMajor { .. } + ) + { + (order, 0, physical_rows, 0, physical_columns) + } else { + ( + order, + logical_rows, + physical_rows, + logical_columns, + physical_columns, + ) + } + } + + fn scalar_values(run: &KernelRun, abi: &KernelAbi) -> Result, KernelAbiError> { + let count = element_count(run)?; + abi.scalar_arguments + .iter() + .map(|argument| match argument.name { + "element_count" => Ok(count), + "num_partials" => match &run.kernel { + TileKernelSpec::ReductionSum { partials } => Ok(u32::from(*partials - 1)), + _ => Err(KernelAbiError::RequirementMismatch), + }, + "scale_exponent" => match &run.kernel { + TileKernelSpec::Gemm { + multiply: Precision::F8F143 { scale_exponent }, + .. + } => Ok(u32::from_ne_bytes(i32::from(*scale_exponent).to_ne_bytes())), + _ => Err(KernelAbiError::RequirementMismatch), + }, + "initial_block" => match &run.kernel { + TileKernelSpec::AttentionMerge { initial, .. } => Ok(u32::from(*initial)), + _ => Err(KernelAbiError::RequirementMismatch), + }, + "final_block" => match &run.kernel { + TileKernelSpec::AttentionMerge { final_block, .. } => Ok(u32::from(*final_block)), + _ => Err(KernelAbiError::RequirementMismatch), + }, + "words_per_worker" => output_byte_count(run).map(|bytes| bytes / 8 / 6), + "remainder_workers" => output_byte_count(run).map(|bytes| bytes / 8 % 6), + "matrices" => match run.kernel { + TileKernelSpec::Rearrange { matrices, .. } => Ok(matrices), + _ => Err(KernelAbiError::RequirementMismatch), + }, + "logical_rows" => match run.kernel { + TileKernelSpec::Rearrange { logical_rows, .. } => Ok(logical_rows), + _ => Err(KernelAbiError::RequirementMismatch), + }, + "physical_rows" => match run.kernel { + TileKernelSpec::Rearrange { physical_rows, .. } => Ok(physical_rows), + _ => Err(KernelAbiError::RequirementMismatch), + }, + "logical_columns" => match run.kernel { + TileKernelSpec::Rearrange { + logical_columns, .. + } => Ok(logical_columns), + _ => Err(KernelAbiError::RequirementMismatch), + }, + "physical_columns" => match run.kernel { + TileKernelSpec::Rearrange { + physical_columns, .. + } => Ok(physical_columns), + _ => Err(KernelAbiError::RequirementMismatch), + }, + "target_order" => match &run.kernel { + TileKernelSpec::Rearrange { + to: crate::Layout { order, .. }, + .. + } => RearrangeTarget::from_order(*order) + .map(RearrangeTarget::codelet_index) + .ok_or(KernelAbiError::RequirementMismatch), + _ => Err(KernelAbiError::RequirementMismatch), + }, + _ => Err(KernelAbiError::RequirementMismatch), + }) + .collect() + } + + fn element_count(run: &KernelRun) -> Result { + run.output.extents.iter().try_fold(1u32, |product, extent| { + product + .checked_mul(extent.physical_end - extent.start) + .ok_or(KernelAbiError::ElementCountOverflow) + }) + } + + fn output_byte_count(run: &KernelRun) -> Result { + let precision = match &run.requirements { + KernelRequirements::Operator(requirements) => requirements.output.format.precision, + KernelRequirements::Conversion { output, .. } => output.format.precision, + }; + element_count(run)? + .checked_mul( + u32::try_from(precision.bytes()).map_err(|_| KernelAbiError::ElementCountOverflow)?, + ) + .ok_or(KernelAbiError::ElementCountOverflow) + } + + pub fn tile_kernel_abi( + kernel: &TileKernelSpec, + requirements: &KernelRequirements, + ) -> Result { + let precision = match requirements { + KernelRequirements::Operator(requirements) => requirements.output.format.precision, + KernelRequirements::Conversion { output, .. } => output.format.precision, + }; + let (symbols, availability, inputs, scalars) = match kernel { + TileKernelSpec::FillZero => ( + KernelSymbols::Exact(ipu_target::emit::FILL_ZERO_U64_SYMBOL), + KernelAvailability::Implemented, + 0, + scalar_arguments(0, &["words_per_worker", "remainder_workers"]), + ), + TileKernelSpec::Gemm { + multiply, weights, .. + } => { + if !matches!(requirements, KernelRequirements::Operator(_)) { + return Err(KernelAbiError::RequirementMismatch); + } + if *weights == GemmWeightLoad::Interleaved && *multiply != Precision::F16 { + return Err(KernelAbiError::RequirementMismatch); + } + let scalars = if matches!(multiply, Precision::F8F143 { .. }) { + scalar_arguments(2, &["scale_exponent"]) + } else { + Vec::new() + }; + ( + KernelSymbols::Planned, + if matches!(multiply, Precision::F8F143 { .. }) { + KernelAvailability::Required + } else { + KernelAvailability::Implemented + }, + 2, + scalars, + ) + } + TileKernelSpec::Gelu => { + let symbol = gelu_symbol(requirements).unwrap_or("ipu_stack_unsupported_gelu"); + ( + KernelSymbols::Exact(symbol), + if symbol == "ipu_stack_unsupported_gelu" { + KernelAvailability::Required + } else { + KernelAvailability::Implemented + }, + 1, + scalar_arguments(1, &["element_count"]), + ) + } + TileKernelSpec::ReductionSum { .. } => { + if precision != Precision::F16 { + return Err(KernelAbiError::RequirementMismatch); + } + ( + KernelSymbols::Exact("ipu_stack_reduce_sum_f16"), + KernelAvailability::Implemented, + 2, + scalar_arguments(2, &["num_partials", "element_count"]), + ) + } + TileKernelSpec::Add => ( + exact_symbol(precision, "ipu_stack_add_f16", "ipu_stack_add_f32"), + KernelAvailability::Required, + 2, + scalar_arguments( + 2, + &[ + "element_count", + "left_broadcast_stride", + "right_broadcast_stride", + ], + ), + ), + TileKernelSpec::AttentionSoftmax { .. } => ( + KernelSymbols::Planned, + KernelAvailability::Implemented, + 1, + Vec::new(), + ), + TileKernelSpec::AttentionMerge { .. } => ( + KernelSymbols::Planned, + KernelAvailability::Implemented, + 2, + scalar_arguments(2, &["initial_block", "final_block"]), + ), + TileKernelSpec::Cast { from, to } => ( + KernelSymbols::Exact(cast_symbol(*from, *to)), + KernelAvailability::Required, + 1, + scalar_arguments(1, &["element_count"]), + ), + TileKernelSpec::Rearrange { from, to, .. } + if precision == Precision::F16 + && UnpackSource::from_order(from.order).is_some() + && to.order == StorageOrder::Linear => + { + ( + KernelSymbols::Planned, + KernelAvailability::Implemented, + 1, + scalar_arguments( + 1, + &[ + "matrices", + "logical_rows", + "physical_rows", + "logical_columns", + "physical_columns", + ], + ), + ) + } + TileKernelSpec::Rearrange { from, to, .. } + if precision == Precision::F16 + && from.order == StorageOrder::Linear + && (matches!( + to.order, + StorageOrder::Native( + NativeKernelOrder::Left | NativeKernelOrder::TransposedRight + ) + ) || matches!(to.order, StorageOrder::Blocked(order) if order.is_matrix())) => + { + ( + KernelSymbols::Planned, + KernelAvailability::Implemented, + 1, + scalar_arguments( + 1, + &[ + "logical_rows", + "physical_rows", + "target_order", + "logical_columns", + "physical_columns", + ], + ), + ) + } + TileKernelSpec::Rearrange { .. } => ( + KernelSymbols::Exact("ipu_stack_rearrange"), + KernelAvailability::Required, + 1, + Vec::new(), + ), + }; + Ok(KernelAbi { + symbols, + availability, + output_register: OUTPUT_REGISTER, + input_registers: (0..inputs) + .map(|index| FIRST_INPUT_REGISTER + index as u8) + .collect(), + scalar_arguments: scalars, + return_register: RETURN_REGISTER, + }) + } + + pub fn validate_kernel_run(run: &KernelRun) -> Result { + let kernel = &run.kernel; + let abi = tile_kernel_abi(kernel, &run.requirements)?; + if run.inputs.len() != abi.input_registers.len() { + return Err(KernelAbiError::PointerArity { + expected: abi.input_registers.len(), + actual: run.inputs.len(), + }); + } + if let Some(index) = run + .inputs + .iter() + .position(|operand| operand.views.len() != 1) + { + return Err(KernelAbiError::FragmentedOperand(index)); + } + if matches!(kernel, TileKernelSpec::Gelu) { + let KernelSymbols::Exact(symbol) = abi.symbols else { + return Err(KernelAbiError::RequirementMismatch); + }; + let divisor = 2; + let count = element_count(run)?; + if !count.is_multiple_of(divisor) { + return Err(KernelAbiError::UnsupportedElementCount { + symbol, + count, + divisor, + }); + } + } + if let TileKernelSpec::ReductionSum { partials } = kernel { + let count = element_count(run)?; + if *partials < 2 || !count.is_multiple_of(8) { + return Err(KernelAbiError::UnsupportedElementCount { + symbol: "ipu_stack_reduce_sum_f16", + count, + divisor: 8, + }); + } + } + if matches!(kernel, TileKernelSpec::FillZero) { + let bytes = output_byte_count(run)?; + if !bytes.is_multiple_of(8) { + return Err(KernelAbiError::UnsupportedElementCount { + symbol: ipu_target::emit::FILL_ZERO_U64_SYMBOL, + count: bytes, + divisor: 8, + }); + } + } + Ok(abi) + } + + fn gelu_symbol(requirements: &KernelRequirements) -> Option<&'static str> { + let KernelRequirements::Operator(requirements) = requirements else { + return None; + }; + let [input] = requirements.inputs.as_slice() else { + return None; + }; + if input.format.precision != Precision::F16 + || requirements.output.format.precision != Precision::F16 + { + return None; + } + let input_layout = &input.format.layout; + let output_layout = &requirements.output.format.layout; + (input_layout == output_layout).then_some("ipu_stack_gelu_tanh_approx_f16") + } + + fn exact_symbol( + precision: Precision, + f16_symbol: &'static str, + f32_symbol: &'static str, + ) -> KernelSymbols { + KernelSymbols::Exact(match precision { + Precision::F16 => f16_symbol, + Precision::F32 => f32_symbol, + Precision::F8F143 { .. } => "ipu_stack_unsupported_f8_kernel", + }) + } + + fn cast_symbol(from: Precision, to: Precision) -> &'static str { + match (from, to) { + (Precision::F16, Precision::F32) => "ipu_stack_cast_f16_f32", + (Precision::F32, Precision::F16) => "ipu_stack_cast_f32_f16", + (Precision::F8F143 { .. }, Precision::F16) => "ipu_stack_cast_f8_f16", + (Precision::F8F143 { .. }, Precision::F32) => "ipu_stack_cast_f8_f32", + (Precision::F16, Precision::F8F143 { .. }) => "ipu_stack_cast_f16_f8", + (Precision::F32, Precision::F8F143 { .. }) => "ipu_stack_cast_f32_f8", + _ => "ipu_stack_cast_identity", + } + } + + fn scalar_arguments(input_count: u8, names: &[&'static str]) -> Vec { + names + .iter() + .enumerate() + .map(|(index, name)| ScalarArgument { + register: FIRST_INPUT_REGISTER + input_count + index as u8, + name, + }) + .collect() + } + + #[cfg(test)] + mod tests { + use super::*; + use crate::{ + AccumulationPrecision, ComputeGraph, Ipu21CostModel, Layout, MemoryClass, + OperandRequirement, OperatorRequirements, OutputAliasing, PipelineConfig, + PlannerSearchDomain, TensorFormat, TensorTiling, lower, lower_to_tiles, + }; + + #[test] + fn randomized_gemm_abis_resolve_to_retained_symbols() { + let mut random = fastrand::Rng::with_seed(0x6162_6921); + for _ in 0..64 { + let precision = if random.bool() { + Precision::F16 + } else { + Precision::F32 + }; + let mode = if random.bool() { + GemmKernelMode::Initialize + } else { + GemmKernelMode::Accumulate + }; + let weights = if precision == Precision::F16 && random.bool() { + GemmWeightLoad::Interleaved + } else { + GemmWeightLoad::Standard + }; + let format = TensorFormat { + precision, + layout: Layout { + order: crate::StorageOrder::Linear, + tiling: TensorTiling::replicated(1), + memory_class: MemoryClass::Standard, + }, + }; + let operand = OperandRequirement::new(format, 8); + let requirements = KernelRequirements::Operator(OperatorRequirements { + inputs: vec![operand.clone(), operand.clone()], + output: operand, + output_aliasing: OutputAliasing::Fresh, + memory_space: MemorySpaceRequirements::default(), + }); + let abi = tile_kernel_abi( + &TileKernelSpec::Gemm { + multiply: precision, + accumulate: AccumulationPrecision::F32, + mode, + weights, + inner_block: 64, + output_columns: [32, 64, 128][random.usize(0..3)], + rows: random.u32(1..=64), + }, + &requirements, + ) + .unwrap(); + assert_eq!(abi.availability, KernelAvailability::Implemented); + assert_eq!(abi.symbols, KernelSymbols::Planned); + assert_eq!(abi.input_registers, [3, 4]); + assert_eq!(abi.return_register, 10); + } + } + + #[test] + fn randomized_gemm_plans_compile_and_select_scheduled_row_specializations() { + let mut random = fastrand::Rng::with_seed(0x7370_6563); + for _ in 0..32 { + let tiles = 1_u16 << random.u32(0..=3); + let rows_per_tile = random.u32(1..=12); + let batch = random.u32(1..=4); + let mut graph = ComputeGraph::new(); + let left = graph + .host_input("left", [batch, u32::from(tiles) * rows_per_tile, 64]) + .unwrap(); + let right = graph.parameter("right", [64, 64]).unwrap(); + let result = graph.gemm(left, right).unwrap(); + graph.set_outputs([result]).unwrap(); + let config = PipelineConfig::new(tiles) + .with_search_domain(PlannerSearchDomain::default().with_active_tile_counts([tiles])) + .with_input( + left, + TensorFormat { + precision: Precision::F16, + layout: Layout::amp_left(64, tiles), + }, + ) + .with_input( + right, + TensorFormat { + precision: Precision::F16, + layout: Layout::block_major_matrix(64, tiles), + }, + ); + let mid = lower(&graph, &config, &Ipu21CostModel).unwrap(); + let low = lower_to_tiles(&mid, &config).unwrap(); + let plan = KernelBuildPlan::from_program(&low).unwrap(); + let addresses = low + .shards + .iter() + .map(|shard| (shard.id, 0x60000 + shard.id.index() * 0x10000)) + .collect::>(); + assert_eq!(plan.compilations.len(), 1); + let planned_rows = plan.compilations[0] + .definitions + .iter() + .find_map(|(name, value)| { + (name == &"GEMM_SMALL_ROWS").then(|| value.parse::().unwrap()) + }) + .unwrap(); + assert!( + plan.compilations[0] + .definitions + .iter() + .any(|(name, value)| name == &"GEMM_SMALL_ROWS" + && value == &planned_rows.to_string()) + ); + assert!( + plan.compilations[0] + .definitions + .iter() + .any(|(name, value)| name == &"GEMM_SINGLE_ROWS" && value == "1") + ); + assert_eq!(plan.compilations[0].retained_symbols.len(), 2); + for run in low + .tiles + .iter() + .flat_map(|tile| low.work(tile)) + .filter_map(|work| { + if let TileWorkRef::Kernel(run) = work { + Some(run) + } else { + None + } + }) + { + let call = plan.call(run).unwrap(); + assert!(plan.retained_symbols().any(|symbol| symbol == call.symbol)); + assert!(call.arguments.is_empty()); + let compute = + materialize_kernel_run(run, &low.shards, &addresses, &plan, &BTreeMap::new()) + .unwrap_or_else(|error| { + panic!("batch={batch} tiles={tiles} run={run:?}: {error}") + }); + assert_eq!(compute.symbol, call.symbol); + assert_eq!(compute.input_addresses.len(), 2); + } + } + } + + #[test] + fn randomized_gelu_abis_select_supported_layout_paths() { + let mut random = fastrand::Rng::with_seed(0x6765_6c75); + for _ in 0..64 { + let tiles = 1_u16 << random.u32(0..=5); + let input_layout = if random.bool() { + Layout::amp_left_result(tiles) + } else { + Layout::row_sharded(tiles) + }; + let output_layout = input_layout.clone(); + let requirement = |layout| { + OperandRequirement::new( + TensorFormat { + precision: Precision::F16, + layout, + }, + 8, + ) + }; + let requirements = KernelRequirements::Operator(OperatorRequirements { + inputs: vec![requirement(input_layout)], + output: requirement(output_layout), + output_aliasing: OutputAliasing::Fresh, + memory_space: MemorySpaceRequirements::default(), + }); + let abi = tile_kernel_abi(&TileKernelSpec::Gelu, &requirements).unwrap(); + assert_eq!(abi.availability, KernelAvailability::Implemented); + assert_eq!(abi.input_registers, [3]); + assert_eq!(abi.scalar_arguments[0].register, 4); + assert_eq!( + abi.symbols, + KernelSymbols::Exact("ipu_stack_gelu_tanh_approx_f16") + ); + } + } + } ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec diff --cc crates/ipu-codegen/src/lib.rs index 7dfdbb8,ccd79a0..0000000 --- a/crates/ipu-codegen/src/lib.rs +++ b/crates/ipu-codegen/src/lib.rs @@@ -1,1159 -1,89 +1,1246 @@@ - use ipu_exchange::{ - 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 serde::{Deserialize, Serialize}; - use std::collections::BTreeMap; + //! Whole-device graph compilation and IPU package construction. + //! + //! [`build_package`] is the primary entry point: it plans a [`ComputeGraph`], + //! lowers it to tile work, compiles and links the required kernels, assigns + //! memory, emits tile programs, and returns a loadable [`CompiledPackage`]. ++<<<<<<< HEAD +mod estimate; +mod exchange; +pub mod graph; +mod host; +mod kernel; +mod low; +mod memory; +mod mid; +mod package; +mod place; +mod storage; +mod tile; +pub(crate) use exchange::*; +pub use exchange::{ + EXCHANGE_SCHEDULE_SNAPSHOT_VERSION, ExchangeActivity, ExchangeActivityKind, + ExchangeScheduleSnapshot, PhysicalExchangePhase, inactive_exchange_program, + schedule_exchange_problem, validate_exchange_schedule, ++======= + mod package; + pub use package::{ + CompiledPackage, CompiledTensor, CompiledTensorShard, DiagnosticCheckpoint, PackageBuildError, + PackageBuildResult, PackageConfig, TileProgramData, build_diagnostic_package, build_package, + build_tile_program_package, + }; + + mod config; + mod conversion; + mod cost; + pub mod exchange; + pub mod graph; + mod host; + pub mod kernel; + mod layout; + pub mod low; + pub mod memory; + mod metrics; + pub mod mid; + mod operator; + pub mod place; + mod schedule; + pub mod storage; + pub mod tile; + pub use config::{ + AttentionStrategy, ConversionStreamingPolicy, OperatorClass, PipelineConfig, + PlannerSearchDomain, ProfilingConfig, + }; + pub use conversion::{ + ConversionGeometryError, ConversionMapping, ConversionStrategy, CopyGeometry, DeferredTransform, + }; + pub(crate) use exchange::lower_exchanges; + pub use exchange::{ + ExchangeActivity, ExchangeLoweringError, PhysicalExchangePhase, inactive_exchange_program, ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec }; +pub(crate) use graph::*; pub use graph::{ ++<<<<<<< HEAD + AttentionOptions, AttentionScale, AxisFactorView, ComputeGraph, GemmOptions, GraphError, + GraphInputKind, Operation, OperationKind, Region, Repeat, ValueId, ++======= + AddOptions, AttentionOptions, AttentionScale, BroadcastMode, ComputeGraph, GemmOptions, + GraphError, GraphInput, GraphInputKind, GraphResult, Operation, OperationId, OperationKind, + Region, RegionBuilder, Repeat, RepeatArguments, SplitHeadsOptions, TensorShape, ValueId, + ValueSequence, ValueSequenceId, + }; + pub use kernel::{ + KernelAbi, KernelAbiError, KernelAvailability, KernelBuildPlan, KernelCompilation, + KernelMaterializationError, KernelOptimization, KernelSymbols, PlannedKernelCall, + ScalarArgument, TileKernelSpec, materialize_kernel_run, tile_kernel_abi, validate_kernel_run, + }; + pub use layout::{ + AMP_COLUMN_MICRO, AMP_INNER_BLOCK, AMP_OUTPUT_COLUMN_BLOCK, AxisTiling, BlockedOrder, Layout, + LayoutError, MemoryClass, NativeKernelOrder, Padding, ShardExtent, StorageOrder, TensorAxis, + TensorFormat, TensorRegion, TensorTiling, TensorType, ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec }; +pub(crate) use kernel::*; +pub(crate) use low::*; pub use low::{ ++<<<<<<< HEAD + BlockValue, BlockValueId, ShardDefinition, ShardView, logical_view_byte_spans, + shard_storage_bytes, ++======= + ExchangeOrder, ExchangePhase, ExchangePhaseId, KernelOperand, KernelRequirements, KernelRun, + KernelRunId, KernelRunMetadata, LocalCopy, LocalCopyId, LocalCopyPattern, LogicalExchange, + LowInput, LowLoweringError, LowLoweringResult, LowProgram, LowShard, LowShardId, LowValue, + RepeatCarried, RepeatInvariant, RepeatIterated, RepeatRun, RepeatRunId, ShardDefinition, + ShardView, TileWork, TileWorkList, TileWorkRef, WorkProvenance, WorkReason, lower_to_tiles, + }; + pub use metrics::{ + CostEstimate, ExchangeFootprint, MemoryEstimate, MemoryPeaks, MemoryUsage, OperationMetrics, + PlanMetrics, RegionMetrics, ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec }; +pub(crate) use memory::*; +pub(crate) use mid::*; pub use mid::{ ++<<<<<<< HEAD + AMP_COLUMN_MICRO, AmpOrder, AttentionStrategy, BlockMajorOrder, ConversionStreamingPolicy, + GemmOrientation, GemmPlanConstraint, GridOrder, Layout, LocalOperandStaging, MemoryClass, + MidOperator, PipelineConfig, Precision, ReductionStaging, ShardExtent, TensorFormat, + TensorType, +}; +pub use package::{ + CompiledPackage, DiagnosticPackage, DiagnosticShard, DiagnosticTensor, PackageConfig, + TileProgramData, build_diagnostic_package, build_package, build_tile_program_package, +}; +pub(crate) use place::*; +pub(crate) use storage::*; +pub use storage::{amp_matrix_coordinates, block_major_matrix_coordinates}; +pub(crate) use tile::*; + +const INCOMING_BASE: u8 = 0xa4; +const INCOMING_DCOUNT: u8 = 0xa6; +const INCOMING_MUX: u8 = 0xa0; +const INCOMING_FORMAT: u8 = 0xa3; +const INCOMING_MUXPAIR: u8 = 0xa1; +// Recovered primitive PIC/XPIC plans arm A6 with one; their payload length is +// encoded in the timed instructions rather than this external-stream counter. +// Consolidated phases currently preserve that primitive-plan setting. +const INTERNAL_EXCHANGE_DCOUNT: u32 = 1; +const OUTGOING_BASE: u8 = 0xa7; +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 COMPLETED_SYMBOL: &str = "ipu_stack_static_completed"; +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] ipu_exchange::ExchangeError), + #[error("invalid tile program: {0}")] + Invalid(String), +} + +pub type Result = std::result::Result; + +/// 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)] +pub struct CodegenOptions { + /// 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 { + 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, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct PlacedExchangeRow { + pub address: u32, + pub words: 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::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) => { + 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(INCOMING_BASE, 8)?; + } + if let Some(source) = exchange.incoming_mux { + code.setzi(8, u32::from(source))?; + code.put_special(INCOMING_MUX, 8)?; + } + if exchange.incoming_format != 0 { + code.setzi(8, u32::from(exchange.incoming_format))?; + code.put_special(INCOMING_FORMAT, 8)?; + } + if let Some(source) = exchange.incoming_mux_pair { + code.setzi(8, u32::from(source))?; + code.put_special(INCOMING_MUXPAIR, 8)?; + } + if !exchange.preserve_base_registers { + code.put_special(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(INTERNAL_EXCHANGE_DCOUNT), + )?; + code.put_special(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(ipu_exchange::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(&ipu_exchange::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, + OUTPUT_REGISTER, + 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, RETURN_REGISTER) +} + +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 { + 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: inactive_exchange_program(), + }, + 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, ipu_exchange::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() + ); + } + } +} ++======= + CostModel, Ipu21CostModel, LoweringError, LoweringResult, MidGraph, MidInput, MidOperation, + MidOperationKind, MidRegion, MidRepeat, MidValue, MidValueId, lower, + }; + pub use operator::{ + AccumulationPrecision, AllocationRequirements, GemmBlockShape, GemmGeometry, GemmGrid, + GemmKernelFamily, GemmKernelMode, GemmOrientation, GemmPlanConstraint, GemmResultGrid, + GemmWeightLoad, GridOrder, LocalOperandStaging, MemoryElementRequirement, MemoryOperand, + MemorySpaceRequirements, MidOperator, OperandMaterialization, OperandRequirement, + OperatorPlanError, OperatorRequirements, OutputAliasing, Precision, ReductionStaging, + }; + pub use place::{Placement, PlacementError, place}; + pub use schedule::{ + AttentionBlocking, AttentionMap, GemmMap, KernelMap, OperatorSchedule, ScheduleAccess, + ScheduleStep, ScheduleValue, + }; + pub use storage::{ + ByteSpan, StorageError, StorageResult, amp_matrix_coordinates, block_major_matrix_coordinates, + logical_view_byte_spans, shard_storage_bytes, view_byte_spans, + }; + pub use tile::{TileLoweringError, TileProgramLowering, compact_exchange_row_address}; ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec diff --cc crates/ipu-codegen/src/low/mod.rs index d000ba3,9ce1006..0000000 --- a/crates/ipu-codegen/src/low/mod.rs +++ b/crates/ipu-codegen/src/low/mod.rs @@@ -1,25 -1,291 +1,315 @@@ ++<<<<<<< HEAD +//! Tile expansion and per-tile work, followed by placement and code generation. + +mod call; +mod copy; +pub(crate) mod expand; +mod graph; +mod initialization; +mod passes; +pub(crate) use call::*; +pub use copy::*; +pub use expand::{ + ExpansionError, ExpansionResult, logical_view_byte_spans, shard_storage_bytes, view_byte_spans, +}; +pub use graph::*; + +use crate::graph::OperationId; +use crate::mid::*; +use std::sync::Arc; ++======= + //! Logical per-tile schedule produced from the layout-aware mid-level IR. + //! + //! Tensor shards have tile identities and rectangular physical extents, and + //! work is ordered per tile. Exchanges still refer to logical shards rather + //! than SRAM addresses; kernel runs still name a selected kernel kind rather + //! than a linked symbol. Placement and final code generation resolve those + //! remaining choices. + + use crate::PipelineConfig; + use crate::conversion::{ConversionStrategy, DeferredTransform}; + use crate::graph::{GraphInputKind, OperationId}; + use crate::kernel::TileKernelSpec; + use crate::layout::{ + AMP_COLUMN_MICRO, AMP_INNER_BLOCK, BlockedOrder, Layout, LayoutError, MemoryClass, + NativeKernelOrder, ShardExtent, StorageOrder, TensorRegion, TensorTiling, TensorType, + }; + use crate::mid::{MidGraph, MidOperation, MidOperationKind, MidRepeat, MidValueId}; + use crate::operator::{ + GemmKernelMode, MemoryOperand, MemorySpaceRequirements, OperandRequirement, + OperatorRequirements, OutputAliasing, Precision, ReductionStaging, + }; + use crate::schedule::{OperatorSchedule, ScheduleAccess, ScheduleStep, ScheduleValue}; + use crate::storage::{ByteSpan, StorageError, logical_view_byte_spans, view_byte_spans}; + use ipu_target::hardware::HardwareTarget; + use std::collections::{BTreeMap, BTreeSet}; + use std::sync::Arc; + use std::time::Instant; + + mod attention; + mod gemm; + mod pointwise; + #[cfg(test)] + mod tests; + + #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] + pub struct LowShardId(u32); + + impl LowShardId { + pub const fn index(self) -> u32 { + self.0 + } + + pub const fn from_index(index: u32) -> Self { + Self(index) + } + } + + #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] + pub struct ExchangePhaseId(u32); + + impl ExchangePhaseId { + pub const fn from_index(index: u32) -> Self { + Self(index) + } + + pub const fn index(self) -> u32 { + self.0 + } + } + + #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] + pub struct KernelRunId(u32); + + impl KernelRunId { + pub const fn index(self) -> u32 { + self.0 + } + } + + #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] + pub struct LocalCopyId(u32); + + impl LocalCopyId { + pub const fn index(self) -> u32 { + self.0 + } + } ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] pub struct RepeatRunId(u32); ++<<<<<<< HEAD ++======= + impl RepeatRunId { + pub const fn index(self) -> u32 { + self.0 + } + } + + #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] + pub struct ShardView { + pub shard: LowShardId, + pub extents: TensorRegion, + } + + #[derive(Clone, Copy, Debug, PartialEq, Eq)] + pub enum ShardDefinition { + Value(MidValueId), + /// Transient storage populated by one or more exchange phases. + ExchangeStaging, + LocalCopy(LowShardId), + /// Persistent scratch allocation populated by local copies or exchanges. + Staging, + Alias(LowShardId), + /// Alias intentionally used as an in-place operation destination. + WritableAlias(LowShardId), + /// Canonical format placeholder replaced by schedule-local staging. + Unmaterialized, + } + + #[derive(Clone, Debug, PartialEq, Eq)] + pub struct LowShard { + pub id: LowShardId, + pub tile: u16, + pub tensor_type: TensorType, + pub extents: TensorRegion, + pub definition: ShardDefinition, + } + + #[derive(Clone, Debug, PartialEq, Eq)] + pub struct LowInput { + pub name: String, + pub kind: GraphInputKind, + pub value: MidValueId, + pub shards: Vec, + } + + #[derive(Clone, Debug, PartialEq, Eq)] + pub struct LowValue { + pub value: MidValueId, + pub shards: Vec, + } + + /// One source view may populate arbitrary corresponding views on several + /// tiles. Sequential phases may reuse transient destinations after consumers + /// have run. + #[derive(Clone, Debug, PartialEq, Eq)] + pub struct LogicalExchange { + pub source: ShardView, + pub destinations: Vec, + pub order: ExchangeOrder, + } + + #[derive(Clone, Debug, Default, PartialEq, Eq)] + pub enum ExchangeOrder { + /// Preserve tensor coordinates, converting between physical layouts. + #[default] + Semantic, + /// Preserve allocation order, treating both views as packed byte spans. + Physical, + /// Use copy geometry selected from the resolved layouts by mid-level + /// conversion planning. + Planned(crate::CopyGeometry), + } + + fn semantic_exchange(source: ShardView) -> (ShardView, ExchangeOrder) { + (source, ExchangeOrder::Semantic) + } + + fn physical_exchange(source: ShardView) -> (ShardView, ExchangeOrder) { + (source, ExchangeOrder::Physical) + } + + #[derive(Clone, Debug, PartialEq, Eq)] + pub struct ExchangePhase { + pub id: ExchangePhaseId, + pub provenance: WorkProvenance, + pub transfers: Vec, + } + + #[derive(Clone, Copy, Debug, PartialEq, Eq)] + pub enum WorkReason { + OperatorKernel, + OperatorInput { input: u16 }, + OperatorInputs, + PrecisionCast, + LayoutRearrangement, + Repeat, + } + + #[derive(Clone, Copy, Debug, PartialEq, Eq)] + pub struct WorkProvenance { + pub operation: Option, + pub value: Option, + pub reason: WorkReason, + } + + #[derive(Clone, Debug, PartialEq, Eq)] + pub enum KernelRequirements { + Operator(OperatorRequirements), + Conversion { + input: OperandRequirement, + output: OperandRequirement, + memory_space: crate::MemorySpaceRequirements, + }, + } + + #[derive(Clone, Debug, PartialEq, Eq)] + pub struct KernelOperand { + /// Views resident on the execution tile which form this ABI operand. + pub views: Vec, + } + + #[derive(Clone, Debug, PartialEq, Eq)] + pub struct KernelRunMetadata { + pub provenance: WorkProvenance, + pub kernel: TileKernelSpec, + pub requirements: KernelRequirements, + } + + #[derive(Clone, Debug, PartialEq, Eq)] + pub struct KernelRun { + metadata: Arc, + pub inputs: Vec, + pub output: ShardView, + } + + impl KernelRun { + pub fn new( + provenance: WorkProvenance, + kernel: TileKernelSpec, + inputs: Vec, + output: ShardView, + requirements: KernelRequirements, + ) -> Self { + Self { + metadata: Arc::new(KernelRunMetadata { + provenance, + kernel, + requirements, + }), + inputs, + output, + } + } + } + + impl std::ops::Deref for KernelRun { + type Target = KernelRunMetadata; + + fn deref(&self) -> &Self::Target { + &self.metadata + } + } + + #[derive(Clone, Debug, PartialEq, Eq)] + pub struct LocalCopy { + pub source: LowShardId, + pub source_offset: u32, + pub destination: LowShardId, + pub destination_offset: u32, + pub bytes: u32, + pub pattern: LocalCopyPattern, + } + + #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] + pub enum LocalCopyPattern { + #[default] + Contiguous, + Strided { + rows: u32, + row_bytes: u32, + source_stride: u32, + destination_stride: u32, + }, + } + + #[derive(Clone, Debug, PartialEq, Eq)] + pub struct RepeatCarried { + pub initial: LowShardId, + pub argument: LowShardId, + pub yielded: LowShardId, + pub result: LowShardId, + } + + #[derive(Clone, Debug, PartialEq, Eq)] + pub struct RepeatInvariant { + pub input: LowShardId, + pub argument: LowShardId, + } + + #[derive(Clone, Debug, PartialEq, Eq)] + pub struct RepeatIterated { + pub inputs: Vec, + pub argument: LowShardId, + /// Placement must assign entries consecutively at this byte stride. + pub stride_bytes: u32, + pub alignment: u32, + } + ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec #[derive(Clone, Debug, PartialEq, Eq)] pub struct RepeatRun { pub provenance: WorkProvenance, @@@ -57,16 -323,20 +347,33 @@@ pub struct TileWorkList #[derive(Clone, Debug, PartialEq, Eq)] pub struct LowProgram { ++<<<<<<< HEAD + pub program: Arc, + pub tiles: Vec, + pub repeat_runs: Vec, +} + +impl std::ops::Deref for LowProgram { + type Target = TileGraph; + fn deref(&self) -> &Self::Target { + &self.program + } ++======= + pub tile_count: u16, + pub shards: Vec, + pub exchange_phases: Vec, + pub inputs: Vec, + /// Compact per-tile ordering. Non-exchange entries index the arenas below. + pub tiles: Vec, + /// Tile-specific kernel operands and outputs, with shared call metadata. + pub kernel_runs: Vec, + pub local_copies: Vec, + pub repeat_runs: Vec, + /// Canonical materialization of every mid-level value that reaches tile + /// lowering. Diagnostic metadata uses this without adding device copies. + pub values: Vec, + pub outputs: Vec, ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec } impl LowProgram { @@@ -88,77 -358,2518 +395,2592 @@@ } } ++<<<<<<< HEAD +/// Project tile work and remove redundant finite-only initialization. +/// Numerical inputs/results must be finite; loading must first initialize SRAM. +pub fn lower_to_tiles(program: &Arc, diagnostic_checkpoints: bool) -> LowProgram { + fn project( + region: &BlockRegion, + program: &TileGraph, + repeats: &mut Vec, + checkpoints: bool, + ) -> Vec { + let mut tiles = (0..program.tile_count) ++======= + #[derive(Debug, thiserror::Error, PartialEq, Eq)] + pub enum LowLoweringError { + #[error("low-level lowering requires a nonzero tile count")] + EmptyTileGroup, + #[error("value {value:?} declares {declared} tiles, but the schedule capacity is {scheduled}")] + TileCountMismatch { + value: MidValueId, + declared: u16, + scheduled: u16, + }, + #[error("value {0:?} does not exist")] + UnknownValue(MidValueId), + #[error("operation must have exactly one result")] + ResultArity, + #[error("operator operation is missing its selected whole-device plan")] + MissingOperatorPlan, + #[error("operator plan is incompatible with its values or block dimensions")] + InvalidOperatorPlan, + #[error("conversion plan is incompatible with its input or output")] + InvalidConversionPlan, + #[error("repeat structure is inconsistent with its inputs, arguments, yields, or results")] + InvalidRepeat, + #[error("repeat carried value {0} cannot alias its body argument")] + RepeatRequiresInPlace(usize), + #[error("repeat iterated input {0} cannot be represented as equal contiguous blocks")] + InvalidIteratedBlocks(usize), + #[error("too many logical shards or exchange phases")] + IdOverflow, + #[error("invalid tensor layout: {0}")] + Layout(#[from] LayoutError), + #[error("invalid tensor storage view: {0}")] + Storage(#[from] StorageError), + } + + pub type LowLoweringResult = Result; + + fn split_gemm_matrices( + run: &KernelRun, + axis: usize, + coordinates: &mut [u32], + runs: &mut Vec, + ) -> LowLoweringResult<()> { + if axis < coordinates.len() { + let extent = run + .output + .extents + .get(axis) + .ok_or(LowLoweringError::InvalidOperatorPlan)?; + if extent.logical_end != extent.physical_end { + return Err(LowLoweringError::InvalidOperatorPlan); + } + for coordinate in extent.start..extent.physical_end { + coordinates[axis] = coordinate; + split_gemm_matrices(run, axis + 1, coordinates, runs)?; + } + return Ok(()); + } + + let mut matrix = run.clone(); + narrow_gemm_matrix_view(&mut matrix.output, coordinates)?; + for operand in &mut matrix.inputs { + for view in &mut operand.views { + narrow_gemm_matrix_view(view, coordinates)?; + } + } + runs.push(matrix); + Ok(()) + } + + fn narrow_gemm_matrix_view( + view: &mut ShardView, + output_coordinates: &[u32], + ) -> LowLoweringResult<()> { + let input_axes = view.extents.len().saturating_sub(2); + if input_axes > output_coordinates.len() { + return Err(LowLoweringError::InvalidOperatorPlan); + } + let output_axis_offset = output_coordinates.len() - input_axes; + for (axis, extent) in view.extents[..input_axes].iter_mut().enumerate() { + if extent.physical_end - extent.start == 1 { + continue; + } + let coordinate = output_coordinates[output_axis_offset + axis]; + if coordinate < extent.start || coordinate >= extent.physical_end { + return Err(LowLoweringError::InvalidOperatorPlan); + } + extent.start = coordinate; + extent.logical_end = coordinate + 1; + extent.physical_end = coordinate + 1; + } + Ok(()) + } + + fn append_checkpoint(tiles: &mut [TileWorkList], operation: OperationId, breakpoint: u8) { + for tile in tiles { + tile.work.push(TileWork::Checkpoint(operation, breakpoint)); + } + } + + /// Produces a logical per-tile schedule by expanding selected operator plans. + /// Conversions without plans still use a conservative gather fallback. + #[tracing::instrument( + name = "ipu_codegen.low.lower_to_tiles", + skip(graph, config), + fields( + tile_count = config.tile_count, + operations = graph.operations.len(), + profiling = ?config.profiling + ) + )] + pub fn lower_to_tiles(graph: &MidGraph, config: &PipelineConfig) -> LowLoweringResult { + if config.tile_count == 0 { + return Err(LowLoweringError::EmptyTileGroup); + } + let mut state = LoweringState::new(graph, config.tile_count, config.target)?; + let tiles = state.lower_region( + &graph.operations, + &graph.outputs, + config.diagnostic_checkpoints, + )?; + let inputs = graph + .inputs + .iter() + .map(|input| { + Ok(LowInput { + name: input.name.clone(), + kind: input.kind, + value: input.value, + shards: state.value_shards(input.value)?.to_vec(), + }) + }) + .collect::>()?; + let outputs = graph + .outputs + .iter() + .map(|value| { + Ok(LowValue { + value: *value, + shards: state.value_shards(*value)?.to_vec(), + }) + }) + .collect::>()?; + let values = graph + .values + .iter() + .filter_map(|value| { + let shards = &state.canonical[value.id.index() as usize]; + (!shards.is_empty()).then(|| LowValue { + value: value.id, + shards: shards.clone(), + }) + }) + .collect(); + tracing::info!( + shards = state.shards.len(), + exchange_phases = state.phases.len(), + "built logical tile schedule" + ); + Ok(LowProgram { + tile_count: config.tile_count, + shards: state.shards, + exchange_phases: state.phases, + inputs, + tiles, + kernel_runs: state.kernel_runs, + local_copies: state.local_copies, + repeat_runs: state.repeat_runs, + values, + outputs, + }) + } + + enum DeferredValue { + Conversion(MidValueId), + View(DeferredTransform, Vec), + } + + struct PreparedDistributedPanel { + panel: u32, + row_major: Option, + packed: LowShardId, + tile: u16, + destinations: Vec, + } + + struct PreparedAttentionBlock { + row_start: u32, + valid_rows: u32, + key_panels: Vec, + value_panels: Vec, + } + + struct AttentionTask { + tile: u16, + head: u32, + query_row_start: u32, + query_rows: u32, + query_dimension: u32, + value_dimension: u32, + query: LowShardId, + query_receive: Option, + output: LowShardId, + scratch: LowShardId, + weights: LowShardId, + key_staging: LowShardId, + value_staging: LowShardId, + } + + #[derive(Clone, Copy)] + struct AttentionBufferShape { + query_block_rows: u32, + panel_rows: u32, + logical_staging_rows: u32, + physical_staging_rows: u32, + scratch_columns: u32, + state_columns: u32, + padded_query_dimension: u32, + padded_value_dimension: u32, + reuse_key_staging_for_state: bool, + } + + fn gemm_kernel_spec( + family: crate::GemmKernelFamily, + mode: crate::GemmKernelMode, + block: crate::GemmBlockShape, + rows: u32, + ) -> TileKernelSpec { + TileKernelSpec::Gemm { + multiply: family.multiply, + accumulate: family.accumulate, + mode, + weights: family.weights, + inner_block: block.inner, + output_columns: block.output_columns, + rows, + } + } + + fn gemm_kernel_rows(output: &ShardView, order: StorageOrder) -> LowLoweringResult { + let rank = output.extents.len(); + let column_axis = rank + .checked_sub( + if matches!( + order, + StorageOrder::Native( + NativeKernelOrder::TransposedOutput | NativeKernelOrder::TransposedLeft + ) + ) { + 2 + } else { + 1 + }, + ) + .ok_or(LowLoweringError::InvalidOperatorPlan)?; + output + .extents + .iter() + .enumerate() + .filter(|(axis, _)| *axis != column_axis) + .try_fold(1u32, |rows, (_, extent)| { + rows.checked_mul(extent.physical_end - extent.start) + }) + .filter(|rows| *rows != 0) + .ok_or(LowLoweringError::IdOverflow) + } + + fn rearrange_kernel_spec( + from: Layout, + to: Layout, + input: &ShardView, + output: &ShardView, + ) -> LowLoweringResult { + let view = if from.order == StorageOrder::Linear { + output + } else { + input + }; + let rank = view.extents.len(); + if rank < 2 { + return Err(LowLoweringError::InvalidConversionPlan); + } + let rows = view.extents[rank - 2]; + let columns = view.extents[rank - 1]; + let matrices = view.extents[..rank - 2] + .iter() + .try_fold(1u32, |product, extent| { + product.checked_mul(extent.physical_end - extent.start) + }) + .ok_or(LowLoweringError::IdOverflow)?; + Ok(TileKernelSpec::Rearrange { + from, + to, + matrices, + logical_rows: rows.logical_end - rows.start, + physical_rows: rows.physical_end - rows.start, + logical_columns: columns.logical_end - columns.start, + physical_columns: columns.physical_end - columns.start, + }) + } + + impl AttentionBufferShape { + fn from_plan(plan: &crate::AttentionMap, key_rows: u32) -> Self { + let padded_query_dimension = plan.query_dimension; + let padded_value_dimension = plan.value_dimension; + match plan.blocking { + crate::AttentionBlocking::Flash { + query_rows, + key_rows, + } => Self { + query_block_rows: query_rows, + panel_rows: key_rows, + logical_staging_rows: key_rows, + physical_staging_rows: key_rows, + scratch_columns: padded_value_dimension.max(key_rows), + state_columns: key_rows + 16, + padded_query_dimension, + padded_value_dimension, + reuse_key_staging_for_state: false, + }, + crate::AttentionBlocking::Materialized { + query_rows, + padded_key_rows, + } => Self { + query_block_rows: query_rows, + panel_rows: AMP_INNER_BLOCK, + logical_staging_rows: key_rows, + physical_staging_rows: padded_key_rows, + scratch_columns: padded_key_rows.max(padded_value_dimension), + state_columns: padded_key_rows + AMP_COLUMN_MICRO, + padded_query_dimension, + padded_value_dimension, + reuse_key_staging_for_state: true, + }, + } + } + } + + #[derive(Clone, Copy)] + enum AttentionOperand { + Key, + Value, + } + + struct LoweringState { + target: HardwareTarget, + tile_count: u16, + shards: Vec, + canonical: Vec>, + phases: Vec, + kernel_runs: Vec, + local_copies: Vec, + repeat_runs: Vec, + kernel_metadata: Vec>, + deferred_values: BTreeMap, + } + + impl LoweringState { + fn storage_root(&self, mut shard: LowShardId) -> LowShardId { + let mut remaining = self.shards.len().saturating_add(1); + while remaining != 0 { + remaining -= 1; + shard = match self.shards[shard.index() as usize].definition { + ShardDefinition::Alias(source) | ShardDefinition::WritableAlias(source) => source, + _ => return shard, + }; + } + shard + } + + fn new(graph: &MidGraph, tile_count: u16, target: HardwareTarget) -> LowLoweringResult { + let mut state = Self { + target, + tile_count, + shards: Vec::new(), + canonical: vec![Vec::new(); graph.values.len()], + phases: Vec::new(), + kernel_runs: Vec::new(), + local_copies: Vec::new(), + repeat_runs: Vec::new(), + kernel_metadata: Vec::new(), + deferred_values: BTreeMap::new(), + }; + let parameter_origins = graph + .inputs + .iter() + .filter(|input| input.kind == GraphInputKind::Parameter) + .map(|input| graph.values[input.value.index() as usize].origin) + .collect::>(); + let parameter_values = graph + .values + .iter() + .filter(|value| parameter_origins.contains(&value.origin)) + .map(|value| value.id) + .collect::>(); + let parameter_groups = parameter_values + .iter() + .map(|value| graph.values[value.index() as usize].storage_group) + .collect::>(); + let mut parameter_bytes = vec![0u64; usize::from(tile_count)]; + let mut parameter_offsets = BTreeMap::::new(); + for value in &graph.values { + let declared_tiles = value.tensor_type.format.layout.tiling.tile_count; + if declared_tiles == 0 || declared_tiles > tile_count { + return Err(LowLoweringError::TileCountMismatch { + value: value.id, + declared: declared_tiles, + scheduled: tile_count, + }); + } + let extents = shard_extents(&value.tensor_type)?; + let is_parameter = parameter_values.contains(&value.id); + let placement_group = value.storage_group; + let rotate_parameter = is_parameter || parameter_groups.contains(&placement_group); + let parameter_shard_bytes = if rotate_parameter { + extents + .iter() + .map(|(_, extents)| { + crate::shard_storage_bytes(&LowShard { + id: LowShardId(0), + tile: 0, + tensor_type: value.tensor_type.clone(), + extents: extents.clone().into(), + definition: ShardDefinition::Value(value.id), + }) + .map(u64::from) + .map_err(LowLoweringError::from) + }) + .collect::>>()? + } else { + Vec::new() + }; + let parameter_offset = if rotate_parameter { + if let Some(&offset) = parameter_offsets.get(&placement_group) { + offset + } else { + let offset = (0..tile_count) + .min_by_key(|&offset| { + let mut loads = parameter_bytes.clone(); + for (logical, &bytes) in parameter_shard_bytes.iter().enumerate() { + let tile = + (logical + usize::from(offset)) % usize::from(tile_count); + loads[tile] = loads[tile].saturating_add(bytes); + } + (loads.into_iter().max().unwrap_or(u64::MAX), offset) + }) + .ok_or(LowLoweringError::EmptyTileGroup)?; + tracing::debug!( + ?placement_group, + offset, + shards = parameter_shard_bytes.len(), + "assigned parameter storage group to tiles" + ); + parameter_offsets.insert(placement_group, offset); + offset + } + } else { + 0 + }; + let mut value_shards = Vec::with_capacity(extents.len()); + for (logical_shard, (owner_tile, extents)) in extents.into_iter().enumerate() { + let mut shard = LowShard { + id: LowShardId(0), + tile: 0, + tensor_type: value.tensor_type.clone(), + extents: extents.into(), + definition: ShardDefinition::Value(value.id), + }; + shard.tile = if rotate_parameter { + let tile = (usize::from(owner_tile) + usize::from(parameter_offset)) + % usize::from(tile_count); + let bytes = parameter_shard_bytes[logical_shard]; + parameter_bytes[tile] = parameter_bytes[tile] + .checked_add(bytes) + .ok_or(LowLoweringError::IdOverflow)?; + u16::try_from(tile).map_err(|_| LowLoweringError::IdOverflow)? + } else { + owner_tile + }; + let id = state.push_shard(shard)?; + value_shards.push(id); + } + state.canonical[value.id.index() as usize] = value_shards; + } + Ok(state) + } + + fn push_shard(&mut self, mut shard: LowShard) -> LowLoweringResult { + let id = + LowShardId(u32::try_from(self.shards.len()).map_err(|_| LowLoweringError::IdOverflow)?); + shard.id = id; + self.shards.push(shard); + Ok(id) + } + + fn right_shards_for_block<'a>( + &'a self, + right_shards: &'a [LowShardId], + column_start: u32, + column_end: u32, + inner_start: u32, + inner_end: u32, + ) -> impl Iterator + 'a { + right_shards.iter().copied().filter(move |shard| { + let extents = &self.shards[shard.index() as usize].extents; + let columns = extents[extents.len() - 1]; + let inner = extents[extents.len() - 2]; + columns.start <= column_start + && columns.physical_end >= column_end + && inner.start <= inner_start + && inner.physical_end >= inner_end + }) + } + + fn matrix_shards_for_block<'a>( + &'a self, + shards: &'a [LowShardId], + column_axis: usize, + inner_axis: usize, + column_start: u32, + column_end: u32, + inner_start: u32, + inner_end: u32, + ) -> impl Iterator + 'a { + shards.iter().copied().filter(move |shard| { + let extents = &self.shards[shard.index() as usize].extents; + let columns = extents[column_axis]; + let inner = extents[inner_axis]; + columns.start <= column_start + && columns.physical_end >= column_end + && inner.start <= inner_start + && inner.physical_end >= inner_end + }) + } + + fn prefer_local_shard(&self, shards: &[LowShardId], tile: u16) -> Option { + shards + .iter() + .copied() + .min_by_key(|shard| u8::from(self.shards[shard.index() as usize].tile != tile)) + } + + fn value_shards(&self, value: MidValueId) -> LowLoweringResult<&[LowShardId]> { + self.canonical + .get(value.index() as usize) + .filter(|shards| !shards.is_empty()) + .map(Vec::as_slice) + .ok_or(LowLoweringError::UnknownValue(value)) + } + + fn deferred_view(&self, mut value: MidValueId) -> Option<&DeferredValue> { + loop { + match self.deferred_values.get(&value)? { + DeferredValue::Conversion(source) => value = *source, + view @ DeferredValue::View(..) => return Some(view), + } + } + } + + fn deferred_supports_physical_exchange( + &self, + value: MidValueId, + destination: LowShardId, + ) -> bool { + let Some(DeferredValue::View(_, shards)) = self.deferred_view(value) else { + return false; + }; + let Some(source) = shards.first() else { + return false; + }; + self.value_shards(value) + .ok() + .and_then(|shards| shards.first()) + .and_then(|shard| { + self.shards[shard.index() as usize] + .tensor_type + .shape + .0 + .last() + }) + .is_some_and(|width| width.is_multiple_of(2)) + && self.shards[source.index() as usize] + .tensor_type + .format + .supports_f16_micro_panel_exchange( + &self.shards[destination.index() as usize].tensor_type.format, + ) + } + + fn local_shard(&self, value: MidValueId, tile: u16) -> LowLoweringResult { + let shards = self.value_shards(value)?; + if let Some(&shard) = shards.get(usize::from(tile)) + && self.shards[shard.index() as usize].tile == tile + { + return Ok(shard); + } + shards + .iter() + .copied() + .find(|shard| self.shards[shard.index() as usize].tile == tile) + .ok_or(LowLoweringError::UnknownValue(value)) + } + + fn intersecting_shard_set( + &self, + sources: &[LowShardId], + target: &[ShardExtent], + local_tile: u16, + ) -> Vec<(Vec, LowShardId)> { + let mut groups = BTreeMap::, Vec>::new(); + for &source in sources { + if let Some(extents) = + intersect_extents(&self.shards[source.index() as usize].extents, target) + { + groups.entry(extents).or_default().push(source); + } + } + groups + .into_iter() + .map(|(extents, candidates)| { + let selected = candidates + .iter() + .copied() + .find(|source| self.shards[source.index() as usize].tile == local_tile) + .unwrap_or(candidates[0]); + (extents, selected) + }) + .collect() + } + + fn lower_region( + &mut self, + operations: &[MidOperation], + retained_values: &[MidValueId], + checkpoints: bool, + ) -> LowLoweringResult> { + let mut tiles = (0..self.tile_count) ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec .map(|tile| TileWorkList { tile, work: Vec::new(), }) .collect::>(); ++<<<<<<< HEAD + for operation in ®ion.operations { + match operation { + BlockOperation::Exchange(id) => { + for tile in &mut tiles { + tile.work.push(TileWork::Exchange(*id)); + } + } + BlockOperation::Copy { tile, copy } => tiles[usize::from(*tile)] + .work + .push(TileWork::LocalCopy(*copy)), + BlockOperation::Compute { tile, run } => { + tiles[usize::from(*tile)].work.push(TileWork::Kernel(*run)) + } + BlockOperation::Checkpoint(operation, breakpoint) if checkpoints => { + for tile in &mut tiles { + tile.work + .push(TileWork::Checkpoint(*operation, *breakpoint)); + } + } + BlockOperation::Checkpoint(..) => {} + BlockOperation::Repeat(repeat) => { + let body = project(&repeat.body, program, repeats, false); + for binding in &repeat.bindings { + let id = RepeatRunId( + u32::try_from(repeats.len()) + .expect("too many projected repeat instances"), + ); + repeats.push(RepeatRun { + provenance: repeat.provenance, + count: repeat.count, + carried: binding.carried.clone(), + invariants: binding.invariants.clone(), + iterated: binding.iterated.clone(), + body: Box::new(body[usize::from(binding.tile)].clone()), + }); + tiles[usize::from(binding.tile)] + .work + .push(TileWork::Repeat(id)); + } + } + } + } + tiles + } + let mut repeat_runs = Vec::new(); + let tiles = project( + &program.body, + program, + &mut repeat_runs, + diagnostic_checkpoints, + ); + let mut low = LowProgram { + program: Arc::clone(program), + tiles, + repeat_runs, + }; + initialization::reuse_finite_padding(&mut low); + low ++======= + let mut checkpoint = 0u8; + for (index, operation) in operations.iter().enumerate() { + let started = Instant::now(); + if self.defer_conversion( + operation, + operations.get(index + 1), + operations, + retained_values, + &mut tiles, + )? { + tracing::info!( + operation = index, + source = ?operation.source.map(OperationId::index), + elapsed_ms = started.elapsed().as_millis() as u64, + "deferred low-level conversion" + ); + continue; + } + let lowered = match &operation.kind { + MidOperationKind::Repeat(repeat) => { + self.lower_repeat(operation, repeat, &mut tiles) + } + MidOperationKind::Operator(_) => self.lower_operator(operation, &mut tiles), + kind => self.lower_conversion(operation, kind, &mut tiles), + }; + if let Err(error) = lowered { + tracing::error!( + operation = index, + source = ?operation.source.map(OperationId::index), + kind = ?operation.kind, + inputs = ?operation.inputs, + results = ?operation.results, + ?error, + "failed to lower mid operation to tile work" + ); + return Err(error); + } + if checkpoints + && matches!( + operation.kind, + MidOperationKind::Operator(_) + | MidOperationKind::Convert(Some(_), ..) + | MidOperationKind::Repeat(_) + ) + && let Some(source) = operation.source + { + append_checkpoint(&mut tiles, source, checkpoint); + checkpoint ^= 1; + } + tracing::info!( + operation = index, + source = ?operation.source.map(OperationId::index), + elapsed_ms = started.elapsed().as_millis() as u64, + shards = self.shards.len(), + exchange_phases = self.phases.len(), + "lowered mid operation to tile work" + ); + } + Ok(tiles) + } + + fn unpack_amp_to_row_major( + &mut self, + source: MidValueId, + provenance: WorkProvenance, + tiles: &mut [TileWorkList], + ) -> LowLoweringResult>> { + let sources = self.value_shards(source)?.to_vec(); + for &source_shard in &sources { + let source = &self.shards[source_shard.index() as usize]; + let compatible = source.extents.len() == 3 + && source.tensor_type.format.precision == Precision::F16 + && match source.tensor_type.format.layout.order { + StorageOrder::Native(NativeKernelOrder::Output) => { + let columns = source.extents[2]; + (columns.physical_end - columns.start).is_multiple_of(AMP_COLUMN_MICRO) + } + StorageOrder::Native(NativeKernelOrder::TransposedLeft) => { + let rows = source.extents[1]; + (rows.physical_end - rows.start).is_multiple_of(AMP_COLUMN_MICRO) + } + _ => false, + }; + if !compatible { + tracing::debug!( + shard = source_shard.index(), + rank = source.extents.len(), + precision = ?source.tensor_type.format.precision, + order = ?source.tensor_type.format.layout.order, + extents = ?source.extents, + "cannot unpack source storage into row-major order" + ); + return Ok(None); + } + } + + let mut staging_shards = Vec::with_capacity(sources.len()); + for source_shard in sources { + let source = self.shards[source_shard.index() as usize].clone(); + let mut staging_type = source.tensor_type.clone(); + staging_type.format.layout = Layout::row_major(TensorTiling::replicated(1)); + let staging = self.push_shard(LowShard { + id: LowShardId(0), + tile: source.tile, + tensor_type: staging_type, + extents: source.extents.clone(), + definition: ShardDefinition::Staging, + })?; + self.append_kernel( + tiles, + source.tile, + KernelRun::new( + provenance, + rearrange_kernel_spec( + source.tensor_type.format.layout.clone(), + self.shards[staging.index() as usize] + .tensor_type + .format + .layout + .clone(), + &self.full_view(source_shard), + &self.full_view(staging), + )?, + vec![KernelOperand { + views: vec![self.full_view(source_shard)], + }], + self.full_view(staging), + KernelRequirements::Conversion { + input: OperandRequirement::new(source.tensor_type.format, 4), + output: OperandRequirement::new( + self.shards[staging.index() as usize] + .tensor_type + .format + .clone(), + 4, + ), + memory_space: MemorySpaceRequirements::default(), + }, + ), + )?; + staging_shards.push(staging); + } + Ok(Some(staging_shards)) + } + + fn defer_conversion( + &mut self, + operation: &MidOperation, + next: Option<&MidOperation>, + operations: &[MidOperation], + retained_values: &[MidValueId], + tiles: &mut [TileWorkList], + ) -> LowLoweringResult { + let (transform, strategy, materialization) = match &operation.kind { + MidOperationKind::Convert(transform, strategy, materialization, _) => { + (*transform, *strategy, *materialization) + } + _ => return Ok(false), + }; + if !strategy.uses_intersections() { + return Ok(false); + } + if materialization != crate::OperandMaterialization::DispatchSlices { + return Ok(false); + } + let ([source], [result]) = (operation.inputs.as_slice(), operation.results.as_slice()) + else { + return Ok(false); + }; + let uses = operations + .iter() + .flat_map(|operation| &operation.inputs) + .chain(retained_values) + .filter(|value| **value == *result) + .count(); + if uses != 1 { + return Ok(false); + } + let next = if transform.is_some() { + operations + .iter() + .find(|candidate| candidate.inputs.contains(result)) + } else { + next + }; + let Some(next) = next else { + return Ok(false); + }; + let Some(input_index) = next.inputs.iter().position(|input| input == result) else { + return Ok(false); + }; + let streamable = next + .operator_plan() + .and_then(|plan| plan.requirements.inputs.get(input_index)) + .is_some_and(|requirement| { + requirement.materialization == crate::OperandMaterialization::DispatchSlices + }); + if !streamable { + return Ok(false); + } + if let Some(transform) = transform { + let source_shards = self.value_shards(*source)?.to_vec(); + let source_format = &self.shards[source_shards[0].index() as usize] + .tensor_type + .format; + let result_format = &self.shards[self.value_shards(*result)?[0].index() as usize] + .tensor_type + .format; + let direct = source_format.supports_f16_micro_panel_exchange(result_format); + let shards = if direct + || !matches!( + source_format.layout.order, + StorageOrder::Native( + NativeKernelOrder::Output | NativeKernelOrder::TransposedLeft + ) + ) { + source_shards + } else { + self.unpack_amp_to_row_major( + *source, + operation_provenance(operation, &operation.kind), + tiles, + )? + .ok_or(LowLoweringError::InvalidConversionPlan)? + }; + self.deferred_values + .insert(*result, DeferredValue::View(transform, shards)); + } else { + self.deferred_values + .insert(*result, DeferredValue::Conversion(*source)); + } + for shard in self.value_shards(*result)?.to_vec() { + self.shards[shard.index() as usize].definition = ShardDefinition::Unmaterialized; + } + Ok(true) + } + + fn lower_conversion( + &mut self, + operation: &MidOperation, + kind: &MidOperationKind, + tiles: &mut [TileWorkList], + ) -> LowLoweringResult<()> { + let (strategy, _) = operation + .conversion() + .ok_or(LowLoweringError::InvalidConversionPlan)?; + match strategy { + ConversionStrategy::LocalKernel => self.lower_local_conversion(operation, kind, tiles), + ConversionStrategy::DirectRetile + | ConversionStrategy::DirectLogical + | ConversionStrategy::StageLogicalThenTransform => { + self.lower_intersection_conversion(operation, kind, strategy, tiles) + } + } + } + + fn lower_local_conversion( + &mut self, + operation: &MidOperation, + kind: &MidOperationKind, + tiles: &mut [TileWorkList], + ) -> LowLoweringResult<()> { + let [input] = operation.inputs.as_slice() else { + return Err(LowLoweringError::InvalidConversionPlan); + }; + let [result] = operation.results.as_slice() else { + return Err(LowLoweringError::ResultArity); + }; + let input_format = self.shards[self.value_shards(*input)?[0].index() as usize] + .tensor_type + .format + .clone(); + let output_format = self.shards[self.value_shards(*result)?[0].index() as usize] + .tensor_type + .format + .clone(); + let static_kernel = match kind { + MidOperationKind::CastPrecision => Some(TileKernelSpec::Cast { + from: input_format.precision, + to: output_format.precision, + }), + MidOperationKind::Convert(..) => None, + MidOperationKind::Operator(_) | MidOperationKind::Repeat(_) => { + return Err(LowLoweringError::InvalidConversionPlan); + } + }; + for output in self.value_shards(*result)?.to_vec() { + let tile = self.shards[output.index() as usize].tile; + let input = self.local_shard(*input, tile)?; + let input_view = self.full_view(input); + let output_view = self.full_view(output); + let kernel = match kind { + MidOperationKind::Convert(..) => rearrange_kernel_spec( + input_format.layout.clone(), + output_format.layout.clone(), + &input_view, + &output_view, + )?, + _ => static_kernel + .clone() + .ok_or(LowLoweringError::InvalidConversionPlan)?, + }; + self.append_kernel( + tiles, + tile, + KernelRun::new( + operation_provenance(operation, kind), + kernel, + vec![KernelOperand { + views: vec![input_view], + }], + output_view, + KernelRequirements::Conversion { + input: OperandRequirement::new(input_format.clone(), 8), + output: OperandRequirement::new(output_format.clone(), 8), + memory_space: MemorySpaceRequirements::default(), + }, + ), + )?; + } + Ok(()) + } + + fn lower_intersection_conversion( + &mut self, + operation: &MidOperation, + kind: &MidOperationKind, + strategy: ConversionStrategy, + tiles: &mut [TileWorkList], + ) -> LowLoweringResult<()> { + let [input] = operation.inputs.as_slice() else { + return Err(LowLoweringError::InvalidConversionPlan); + }; + let [result] = operation.results.as_slice() else { + return Err(LowLoweringError::ResultArity); + }; + let inputs = self.value_shards(*input)?.to_vec(); + let outputs = self.value_shards(*result)?.to_vec(); + let mappings = match kind { + MidOperationKind::Convert(_, _, _, mappings) => mappings, + _ => return Err(LowLoweringError::InvalidConversionPlan), + }; + let staged = strategy == ConversionStrategy::StageLogicalThenTransform; + let staging = if staged { + outputs + .iter() + .map(|&output| self.push_conversion_staging(output)) + .collect::>>()? + } else { + outputs.clone() + }; + if strategy == ConversionStrategy::DirectLogical { + for &output in &outputs { + if self.shard_has_padding(output) { + self.append_fill_zero(tiles, output, operation_provenance(operation, kind))?; + } + } + } + if mappings.iter().any(|mapping| mapping.copies.is_empty()) { + return Err(LowLoweringError::InvalidConversionPlan); + } + let bound = mappings + .iter() + .map(|mapping| { + let source = inputs + .get(mapping.source_shard as usize) + .copied() + .ok_or(LowLoweringError::InvalidConversionPlan)?; + let destination = staging + .get(mapping.destination_shard as usize) + .copied() + .ok_or(LowLoweringError::InvalidConversionPlan)?; + if self.shards[source.index() as usize].extents != mapping.source_storage + || self.shards[destination.index() as usize].extents + != mapping.destination_storage + { + tracing::error!( + source_shard = mapping.source_shard, + destination_shard = mapping.destination_shard, + actual_source = ?self.shards[source.index() as usize].extents, + planned_source = ?mapping.source_storage, + actual_destination = ?self.shards[destination.index() as usize].extents, + planned_destination = ?mapping.destination_storage, + "selected conversion geometry did not bind to low shards" + ); + return Err(LowLoweringError::InvalidConversionPlan); + } + Ok((mapping, source, destination)) + }) + .collect::>>()?; + let mut transfers = BTreeMap::<(ShardView, crate::CopyGeometry), Vec>::new(); + let mut source_copies = Vec::new(); + let mut local_copies = Vec::new(); + let mut destination_copies = Vec::new(); + for (mapping, source, destination) in bound { + let source_tile = self.shards[source.index() as usize].tile; + let destination_tile = self.shards[destination.index() as usize].tile; + let (transfer_source, transfer_destination) = if mapping.source_copies.is_empty() { + (source, destination) + } else { + let bytes = mapping + .copies + .iter() + .try_fold(0_u64, |bytes, geometry| bytes.checked_add(geometry.bytes())) + .and_then(|bytes| u32::try_from(bytes).ok()) + .ok_or(LowLoweringError::IdOverflow)?; + let precision = self.shards[source.index() as usize] + .tensor_type + .format + .precision; + let source_staging = self.push_transfer_staging( + source_tile, + precision, + bytes, + ShardDefinition::Staging, + )?; + let destination_staging = self.push_transfer_staging( + destination_tile, + precision, + bytes, + ShardDefinition::ExchangeStaging, + )?; + self.append_fill_zero( + tiles, + source_staging, + operation_provenance(operation, kind), + )?; + for geometry in &mapping.source_copies { + source_copies.extend( + planned_local_copies(source, source_staging, geometry)? + .into_iter() + .map(|copy| (source_tile, copy)), + ); + } + for geometry in &mapping.destination_copies { + destination_copies.extend( + planned_local_copies(destination_staging, destination, geometry)? + .into_iter() + .map(|copy| (destination_tile, copy)), + ); + } + (source_staging, destination_staging) + }; + let source_view = if transfer_source == source { + ShardView { + shard: source, + extents: mapping.source_region.clone(), + } + } else { + self.full_view(transfer_source) + }; + let destination_view = if transfer_destination == destination { + ShardView { + shard: destination, + extents: mapping.destination_region.clone(), + } + } else { + self.full_view(transfer_destination) + }; + for geometry in &mapping.copies { + if source_tile == destination_tile { + local_copies.extend( + planned_local_copies(transfer_source, transfer_destination, geometry)? + .into_iter() + .map(|copy| (destination_tile, copy)), + ); + } else { + transfers + .entry((source_view.clone(), geometry.clone())) + .or_default() + .push(destination_view.clone()); + } + } + } + for (tile, copy) in source_copies { + self.append_local_copy(tiles, tile, copy)?; + } + self.append_phase( + transfers, + operation_provenance(operation, kind), + |(source, geometry)| (source, ExchangeOrder::Planned(geometry)), + tiles, + )?; + for (tile, copy) in local_copies.into_iter().chain(destination_copies) { + self.append_local_copy(tiles, tile, copy)?; + } + if staged { + for (&staging, &destination) in staging.iter().zip(&outputs) { + let source_format = self.shards[staging.index() as usize] + .tensor_type + .format + .clone(); + let destination_format = self.shards[destination.index() as usize] + .tensor_type + .format + .clone(); + let supported_destination = match destination_format.layout.order { + StorageOrder::Native( + NativeKernelOrder::Left | NativeKernelOrder::TransposedRight, + ) => true, + StorageOrder::Blocked(order) => order.is_matrix(), + StorageOrder::Linear | StorageOrder::Native(_) => false, + }; + if source_format.precision != crate::Precision::F16 + || source_format.layout.order != StorageOrder::Linear + || !supported_destination + { + return Err(LowLoweringError::InvalidConversionPlan); + } + let tile = self.shards[destination.index() as usize].tile; + self.append_kernel( + tiles, + tile, + KernelRun::new( + operation_provenance(operation, kind), + rearrange_kernel_spec( + source_format.layout.clone(), + destination_format.layout.clone(), + &self.full_view(staging), + &self.full_view(destination), + )?, + vec![KernelOperand { + views: vec![self.full_view(staging)], + }], + self.full_view(destination), + KernelRequirements::Conversion { + input: OperandRequirement::new(source_format, 2), + output: OperandRequirement::new(destination_format, 2), + memory_space: MemorySpaceRequirements::default() + .with_distinct_elements([ + MemoryOperand::Input(0), + MemoryOperand::Output, + ]), + }, + ), + )?; + } + } + Ok(()) + } + + fn push_conversion_staging( + &mut self, + destination: LowShardId, + ) -> LowLoweringResult { + let destination = &self.shards[destination.index() as usize]; + let mut extents = destination.extents.clone(); + let tile = destination.tile; + let shape = destination.tensor_type.shape.clone(); + let precision = destination.tensor_type.format.precision; + for extent in &mut extents { + extent.physical_end = extent.logical_end; + } + self.push_shard(LowShard { + id: LowShardId(0), + tile, + tensor_type: TensorType { + shape, + format: crate::TensorFormat { + precision, + layout: Layout { + order: StorageOrder::Linear, + tiling: TensorTiling::replicated(1), + memory_class: MemoryClass::Standard, + }, + }, + }, + extents, + definition: ShardDefinition::Staging, + }) + } + + fn push_transfer_staging( + &mut self, + tile: u16, + precision: Precision, + bytes: u32, + definition: ShardDefinition, + ) -> LowLoweringResult { + let element_bytes = precision.bytes() as u32; + let elements = bytes.div_ceil(element_bytes); + self.push_shard(LowShard { + id: LowShardId(0), + tile, + tensor_type: TensorType::new( + [elements], + precision, + Layout::row_major(TensorTiling::replicated(1)), + ), + extents: vec![ShardExtent { + axis: 0, + start: 0, + logical_end: elements, + physical_end: elements, + }] + .into(), + definition, + }) + } + + fn lower_operator( + &mut self, + operation: &MidOperation, + tiles: &mut [TileWorkList], + ) -> LowLoweringResult<()> { + let plan = operation + .operator_plan() + .ok_or(LowLoweringError::MissingOperatorPlan)?; + match plan.steps.as_slice() { + [ScheduleStep::KernelMap(_)] => { + self.lower_schedule(operation, plan, &plan.requirements, tiles) + } + [ScheduleStep::Gemm(gemm)] => { + self.lower_blocked_gemm(operation, gemm, None, &plan.requirements, tiles) + } + [ + ScheduleStep::Gemm(gemm), + ScheduleStep::Reduce { + input, + output, + staging, + }, + ] => { + if *input != gemm.output || *output != ScheduleValue::Output { + return Err(LowLoweringError::InvalidOperatorPlan); + } + self.lower_blocked_gemm(operation, gemm, Some(*staging), &plan.requirements, tiles) + } + [ScheduleStep::Attention(attention)] => match attention.blocking { + crate::AttentionBlocking::Flash { .. } => { + self.lower_blocked_attention(operation, attention, &plan.requirements, tiles) + } + crate::AttentionBlocking::Materialized { .. } => self.lower_materialized_attention( + operation, + attention, + &plan.requirements, + tiles, + ), + }, + _ => Err(LowLoweringError::InvalidOperatorPlan), + } + } + + /// Splits corresponding views at each allocation's F16 micro-panel + /// boundaries. Within every resulting rectangle the source and + /// destination have identical physical traversal, even when their outer + /// panel sequence and tile ownership differ. + fn f16_micro_panel_mappings( + &self, + mappings: Vec<(ShardView, ShardView)>, + ) -> LowLoweringResult>> { + let mut split = Vec::new(); + for (source, destination) in mappings { + let source_shard = &self.shards[source.shard.index() as usize]; + let destination_shard = &self.shards[destination.shard.index() as usize]; + let pieces = split_mapping_at_panel_boundaries( + source_shard, + source, + destination_shard, + destination, + )?; + for (source, destination) in pieces { + let source_spans = view_byte_spans(source_shard, &source)?; + let destination_spans = view_byte_spans(destination_shard, &destination)?; + let valid_spans = source_spans + .iter() + .chain(&destination_spans) + .all(|span| span.offset & 0b11 == 0 && span.bytes & 0b11 == 0); + let source_bytes = source_spans.iter().map(|span| span.bytes).sum::(); + let destination_bytes = + destination_spans.iter().map(|span| span.bytes).sum::(); + if !valid_spans || source_bytes != destination_bytes { + return Ok(None); + } + split.push((source, destination)); + } + } + Ok(Some(split)) + } + + fn deferred_region_mappings( + &self, + value: MidValueId, + logical_target: &TensorRegion, + destination: LowShardId, + ) -> LowLoweringResult> { + let mut source_value = value; + let (source_shards, target, source_axes) = loop { + match self.deferred_values.get(&source_value) { + Some(DeferredValue::Conversion(source)) => source_value = *source, + Some(DeferredValue::View(transform, shards)) => { + let logical_type = + &self.shards[self.value_shards(value)?[0].index() as usize].tensor_type; + let source_type = &self.shards[shards[0].index() as usize].tensor_type; + let mapping = transform + .map_slice(&source_type.shape, &logical_type.shape, logical_target) + .ok_or(LowLoweringError::InvalidOperatorPlan)?; + break (shards.clone(), mapping.source, mapping.source_axes); + } + None => { + break ( + self.value_shards(source_value)?.to_vec(), + logical_target.clone(), + (0..logical_target.len()).map(Some).collect(), + ); + } + } + }; + let destination_tile = self.shards[destination.index() as usize].tile; + let destination_extents = &self.shards[destination.index() as usize].extents; + if destination_extents.len() != source_axes.iter().flatten().count() { + return Err(LowLoweringError::InvalidOperatorPlan); + } + let mut covered = 0u64; + let mut mappings = Vec::new(); + for (source_extents, source) in + self.intersecting_shard_set(&source_shards, &target, destination_tile) + { + let mapped_extents = source_axes + .iter() + .flatten() + .enumerate() + .map(|(destination_axis, &source_axis)| { + let source = source_extents + .get(source_axis) + .ok_or(LowLoweringError::InvalidOperatorPlan)?; + let base = target + .get(source_axis) + .ok_or(LowLoweringError::InvalidOperatorPlan)? + .start; + let destination_base = destination_extents + .get(destination_axis) + .ok_or(LowLoweringError::InvalidOperatorPlan)? + .start; + Ok(ShardExtent { + axis: u16::try_from(destination_axis) + .map_err(|_| LowLoweringError::IdOverflow)?, + start: destination_base + source.start - base, + logical_end: destination_base + source.logical_end - base, + physical_end: destination_base + source.logical_end - base, + }) + }) + .collect::>>()?; + covered = covered + .saturating_add(TensorRegion::new(source_extents.clone()).logical_elements()); + let source_view = ShardView { + shard: source, + extents: source_extents.into(), + }; + let destination_view = ShardView { + shard: destination, + extents: mapped_extents.into(), + }; + mappings.push((source_view, destination_view)); + } + if covered != target.logical_elements() { + return Err(LowLoweringError::InvalidOperatorPlan); + } + Ok(mappings) + } + + fn mapping_word_exchange_fragments( + &self, + mappings: &[(ShardView, ShardView)], + ) -> LowLoweringResult> { + let maximum_bytes = self + .target + .exchange() + .maximum_transfer_words + .checked_mul(4) + .ok_or(LowLoweringError::IdOverflow)?; + let mut fragments = 0u64; + for (source, destination) in mappings { + let source_spans = + logical_view_byte_spans(&self.shards[source.shard.index() as usize], source)?; + let destination_spans = logical_view_byte_spans( + &self.shards[destination.shard.index() as usize], + destination, + )?; + let aligned = source_spans + .iter() + .chain(&destination_spans) + .all(|span| span.offset & 0b11 == 0 && span.bytes & 0b11 == 0); + let source_bytes = source_spans.iter().map(|span| span.bytes).sum::(); + let destination_bytes = destination_spans.iter().map(|span| span.bytes).sum::(); + if !aligned || source_bytes != destination_bytes { + tracing::trace!( + source = ?source, + destination = ?destination, + source_order = ?self.shards[source.shard.index() as usize] + .tensor_type.format.layout.order, + destination_order = ?self.shards[destination.shard.index() as usize] + .tensor_type.format.layout.order, + source_spans = ?source_spans, + destination_spans = ?destination_spans, + aligned, + source_bytes, + destination_bytes, + "deferred logical fragment cannot be exchanged directly" + ); + return Ok(None); + } + let mut source_index = 0usize; + let mut destination_index = 0usize; + let mut source_offset = 0u32; + let mut destination_offset = 0u32; + while source_index < source_spans.len() && destination_index < destination_spans.len() { + let source_remaining = source_spans[source_index].bytes - source_offset; + let destination_remaining = + destination_spans[destination_index].bytes - destination_offset; + let bytes = source_remaining + .min(destination_remaining) + .min(maximum_bytes); + if bytes == 0 || bytes & 0b11 != 0 { + return Ok(None); + } + fragments = fragments.saturating_add(1); + source_offset += bytes; + destination_offset += bytes; + if source_offset == source_spans[source_index].bytes { + source_index += 1; + source_offset = 0; + } + if destination_offset == destination_spans[destination_index].bytes { + destination_index += 1; + destination_offset = 0; + } + } + if source_index != source_spans.len() + || destination_index != destination_spans.len() + || source_offset != 0 + || destination_offset != 0 + { + return Ok(None); + } + } + Ok(Some(fragments)) + } + + fn mappings_benefit_from_word_exchange( + &self, + mappings: &[(ShardView, ShardView)], + destination: LowShardId, + ) -> LowLoweringResult { + let Some(fragments) = self.mapping_word_exchange_fragments(mappings)? else { + return Ok(false); + }; + let shard = &self.shards[destination.index() as usize]; + let bytes = u64::from(crate::shard_storage_bytes(shard)?); + let elements = bytes.div_ceil(shard.tensor_type.format.precision.bytes().max(1)); + let packed_cycles = crate::cost::row_major_pack_cycles(&shard.tensor_type, elements); + let clear_cycles = if self.shard_has_padding(destination) { + self.target + .costs() + .kernel_launch_cycles + .saturating_add(bytes.div_ceil(8 * 6)) + } else { + 0 + }; + let fragment_cycles = fragments + .saturating_mul(self.target.costs().logical_fragment_cycles) + .saturating_add(clear_cycles); + let direct = fragment_cycles < packed_cycles; + tracing::trace!( + destination = destination.index(), + fragments, + fragment_cycles, + packed_cycles, + direct, + "selected logical conversion materialization" + ); + Ok(direct) + } + + fn materialize_deferred_region( + &self, + value: MidValueId, + region: &TensorRegion, + destination: LowShardId, + order: ExchangeOrder, + transfers: &mut BTreeMap>, + local_copies: &mut Vec<(u16, LocalCopy)>, + ) -> LowLoweringResult<()> { + let destination_tile = self.shards[destination.index() as usize].tile; + let mappings = self.deferred_region_mappings(value, region, destination)?; + for (source_view, destination_view) in mappings { + let mappings = if order == ExchangeOrder::Physical { + self.f16_micro_panel_mappings(vec![(source_view, destination_view)])? + .ok_or(LowLoweringError::InvalidOperatorPlan)? + } else { + vec![(source_view, destination_view)] + }; + for (source_view, destination_view) in mappings { + if self.shards[source_view.shard.index() as usize].tile == destination_tile { + if order == ExchangeOrder::Physical { + append_span_copies( + &self.shards, + &source_view, + &destination_view, + destination_tile, + local_copies, + )?; + } else { + append_logical_span_copies( + &self.shards, + &source_view, + &destination_view, + destination_tile, + local_copies, + )?; + } + } else { + transfers + .entry(source_view) + .or_default() + .push(destination_view); + } + } + } + Ok(()) + } + + #[allow(clippy::too_many_arguments)] + fn shard_has_padding(&self, shard: LowShardId) -> bool { + self.shards[shard.index() as usize] + .extents + .iter() + .any(|extent| extent.logical_end < extent.physical_end) + } + + fn append_fill_zero( + &mut self, + tiles: &mut [TileWorkList], + shard: LowShardId, + provenance: WorkProvenance, + ) -> LowLoweringResult<()> { + let shard_data = &self.shards[shard.index() as usize]; + let tile = shard_data.tile; + let output = OperandRequirement::new(shard_data.tensor_type.format.clone(), 8); + self.append_kernel( + tiles, + tile, + KernelRun::new( + provenance, + TileKernelSpec::FillZero, + Vec::new(), + self.full_view(shard), + KernelRequirements::Operator(OperatorRequirements { + inputs: Vec::new(), + output, + output_aliasing: crate::OutputAliasing::Fresh, + memory_space: MemorySpaceRequirements::default(), + }), + ), + ) + } + + fn append_phase( + &mut self, + transfers: impl IntoIterator)>, + provenance: WorkProvenance, + mut order: impl FnMut(K) -> (ShardView, ExchangeOrder), + tiles: &mut [TileWorkList], + ) -> LowLoweringResult<()> { + let mut transfers = transfers + .into_iter() + .map(|(key, mut destinations)| { + let (source, order) = order(key); + destinations.sort_unstable(); + destinations.dedup(); + LogicalExchange { + source, + destinations, + order, + } + }) + .collect::>(); + if transfers.is_empty() { + return Ok(()); + } + if let Some(previous) = self.phases.last().map(|phase| phase.id) + && self.phases[previous.index() as usize] + .provenance + .operation + .is_some() + && self.phases[previous.index() as usize].provenance.operation == provenance.operation + { + let touched = transfers + .iter() + .flat_map(|transfer| { + std::iter::once(transfer.source.shard) + .chain(transfer.destinations.iter().map(|view| view.shard)) + }) + .map(|shard| self.storage_root(shard)) + .collect::>(); + let previous_touched = self.phases[previous.index() as usize] + .transfers + .iter() + .flat_map(|transfer| { + std::iter::once(transfer.source.shard) + .chain(transfer.destinations.iter().map(|view| view.shard)) + }) + .map(|shard| self.storage_root(shard)) + .collect::>(); + let disjoint_transfers = touched.is_disjoint(&previous_touched); + let only_independent_copies_between = tiles.iter().all(|tile| { + let Some(boundary) = tile + .work + .iter() + .rposition(|work| *work == TileWork::Exchange(previous)) + else { + return false; + }; + tile.work[boundary + 1..].iter().all(|work| { + let TileWork::LocalCopy(copy) = *work else { + return false; + }; + let copy = &self.local_copies[copy.0 as usize]; + !touched.contains(&self.storage_root(copy.source)) + && !touched.contains(&self.storage_root(copy.destination)) + }) + }); + if disjoint_transfers && only_independent_copies_between { + let phase = &mut self.phases[previous.index() as usize]; + phase.transfers.append(&mut transfers); + if phase.provenance != provenance { + phase.provenance = WorkProvenance { + operation: provenance.operation, + value: None, + reason: WorkReason::OperatorInputs, + }; + } + tracing::debug!( + phase = previous.index(), + operation = ?provenance.operation.map(OperationId::index), + "consolidated independent exchange transfers" + ); + return Ok(()); + } + } + let id = ExchangePhaseId( + u32::try_from(self.phases.len()).map_err(|_| LowLoweringError::IdOverflow)?, + ); + self.phases.push(ExchangePhase { + id, + provenance, + transfers, + }); + tracing::debug!( + phase = id.index(), + operation = ?provenance.operation.map(OperationId::index), + value = ?provenance.value.map(MidValueId::index), + reason = ?provenance.reason, + "scheduled exchange phase" + ); + for tile in tiles { + tile.work.push(TileWork::Exchange(id)); + } + Ok(()) + } + + fn append_kernel( + &mut self, + tiles: &mut [TileWorkList], + tile: u16, + run: KernelRun, + ) -> LowLoweringResult<()> { + let output_flattens_outer_rows = self + .shards + .get(run.output.shard.index() as usize) + .is_some_and(|shard| { + matches!( + shard.tensor_type.format.layout.order, + StorageOrder::Native(NativeKernelOrder::Left | NativeKernelOrder::Output) + ) + }); + if matches!(run.kernel, TileKernelSpec::Gemm { .. }) + && run.output.extents.len() > 2 + && !output_flattens_outer_rows + { + let matrix_axes = run.output.extents.len() - 2; + let mut coordinates = vec![0; matrix_axes]; + let mut matrix_runs = Vec::new(); + split_gemm_matrices(&run, 0, &mut coordinates, &mut matrix_runs)?; + if matrix_runs.len() > 1 { + for matrix_run in matrix_runs { + self.append_single_kernel(tiles, tile, matrix_run)?; + } + return Ok(()); + } + } + self.append_single_kernel(tiles, tile, run) + } + + fn append_single_kernel( + &mut self, + tiles: &mut [TileWorkList], + tile: u16, + mut run: KernelRun, + ) -> LowLoweringResult<()> { + if let Some(metadata) = self + .kernel_metadata + .iter() + .find(|metadata| metadata.as_ref() == run.metadata.as_ref()) + { + run.metadata = Arc::clone(metadata); + } else { + self.kernel_metadata.push(Arc::clone(&run.metadata)); + } + let id = KernelRunId( + u32::try_from(self.kernel_runs.len()).map_err(|_| LowLoweringError::IdOverflow)?, + ); + self.kernel_runs.push(run); + tiles[usize::from(tile)].work.push(TileWork::Kernel(id)); + Ok(()) + } + + fn append_local_copy( + &mut self, + tiles: &mut [TileWorkList], + tile: u16, + copy: LocalCopy, + ) -> LowLoweringResult<()> { + let id = LocalCopyId( + u32::try_from(self.local_copies.len()).map_err(|_| LowLoweringError::IdOverflow)?, + ); + self.local_copies.push(copy); + tiles[usize::from(tile)].work.push(TileWork::LocalCopy(id)); + Ok(()) + } + + fn append_repeat( + &mut self, + tiles: &mut [TileWorkList], + tile: u16, + repeat: RepeatRun, + ) -> LowLoweringResult<()> { + let id = RepeatRunId( + u32::try_from(self.repeat_runs.len()).map_err(|_| LowLoweringError::IdOverflow)?, + ); + self.repeat_runs.push(repeat); + tiles[usize::from(tile)].work.push(TileWork::Repeat(id)); + Ok(()) + } + + fn full_view(&self, shard: LowShardId) -> ShardView { + ShardView { + shard, + extents: self.shards[shard.index() as usize].extents.clone(), + } + } + + fn narrow_view( + &self, + shard: LowShardId, + ranges: &[(usize, u32, u32)], + ) -> LowLoweringResult { + let mut view = self.full_view(shard); + for &(axis, start, end) in ranges { + let extent = view + .extents + .get_mut(axis) + .ok_or(LowLoweringError::InvalidOperatorPlan)?; + if start < extent.start || end > extent.physical_end || start >= end { + return Err(LowLoweringError::InvalidOperatorPlan); + } + extent.start = start; + extent.physical_end = end; + extent.logical_end = end.min(extent.logical_end).max(start); + } + Ok(view) + } + + fn lower_repeat( + &mut self, + operation: &MidOperation, + repeat: &MidRepeat, + tiles: &mut [TileWorkList], + ) -> LowLoweringResult<()> { + let expected_inputs = repeat.carried_inputs + repeat.invariant_inputs; + let expected_arguments = expected_inputs + repeat.iterated_inputs.len(); + if operation.inputs.len() != expected_inputs + || operation.results.len() != repeat.carried_inputs + || repeat.body.arguments.len() != expected_arguments + || repeat.body.yields.len() != repeat.carried_inputs + || repeat + .iterated_inputs + .iter() + .any(|values| values.len() != repeat.count as usize) + { + return Err(LowLoweringError::InvalidRepeat); + } + for index in 0..repeat.carried_inputs { + if !repeat_yield_can_alias( + repeat.body.yields[index], + repeat.body.arguments[index], + &repeat.body.operations, + ) { + return Err(LowLoweringError::RepeatRequiresInPlace(index)); + } + } + let iterated_requirements = repeat + .iterated_inputs + .iter() + .enumerate() + .map(|(index, _)| { + body_storage_requirement( + repeat.body.arguments[expected_inputs + index], + &repeat.body.operations, + ) + }) + .collect::>(); + let body = self.lower_region(&repeat.body.operations, &repeat.body.yields, false)?; + for tile in 0..self.tile_count { + let mut carried = Vec::with_capacity(repeat.carried_inputs); + for index in 0..repeat.carried_inputs { + let Some(argument) = self.find_local_shard(repeat.body.arguments[index], tile)? + else { + continue; + }; + let initial = self.corresponding_shard(operation.inputs[index], argument)?; + let yielded = self.corresponding_shard(repeat.body.yields[index], argument)?; + let result = self.corresponding_shard(operation.results[index], argument)?; + self.alias_shard(argument, initial); + if yielded != argument { + self.shards[yielded.index() as usize].definition = + ShardDefinition::WritableAlias(argument); + } + self.alias_shard(result, initial); + carried.push(RepeatCarried { + initial, + argument, + yielded, + result, + }); + } + let invariants = (0..repeat.invariant_inputs) + .filter_map(|index| { + let input_index = repeat.carried_inputs + index; + let argument = + match self.find_local_shard(repeat.body.arguments[input_index], tile) { + Ok(Some(argument)) => argument, + Ok(None) => return None, + Err(error) => return Some(Err(error)), + }; + Some( + self.corresponding_shard(operation.inputs[input_index], argument) + .map(|input| RepeatInvariant { input, argument }), + ) + }) + .collect::>()?; + let iterated = repeat + .iterated_inputs + .iter() + .enumerate() + .filter_map(|(index, values)| { + let argument = match self + .find_local_shard(repeat.body.arguments[expected_inputs + index], tile) + { + Ok(Some(argument)) => argument, + Ok(None) => return None, + Err(error) => return Some(Err(error)), + }; + let inputs = values + .iter() + .map(|value| self.corresponding_shard(*value, argument)) + .collect::>>(); + let inputs = match inputs { + Ok(inputs) => inputs, + Err(error) => return Some(Err(error)), + }; + let (alignment, access_tail) = iterated_requirements[index]; + let strides = inputs + .iter() + .map(|shard| self.shard_stride(*shard, alignment, access_tail)) + .collect::>>(); + let strides = match strides { + Ok(strides) => strides, + Err(error) => return Some(Err(error)), + }; + let Some(&stride_bytes) = strides.first() else { + return Some(Err(LowLoweringError::InvalidIteratedBlocks(index))); + }; + if strides.iter().any(|stride| *stride != stride_bytes) { + return Some(Err(LowLoweringError::InvalidIteratedBlocks(index))); + } + Some(Ok(RepeatIterated { + inputs, + argument, + stride_bytes, + alignment, + })) + }) + .collect::>()?; + self.append_repeat( + tiles, + tile, + RepeatRun { + provenance: WorkProvenance { + operation: operation.source, + value: operation.results.first().copied(), + reason: WorkReason::Repeat, + }, + count: repeat.count, + carried, + invariants, + iterated, + body: Box::new(body[usize::from(tile)].clone()), + }, + )?; + } + Ok(()) + } + + fn alias_shard(&mut self, shard: LowShardId, target: LowShardId) { + self.shards[shard.index() as usize].definition = ShardDefinition::Alias(target); + } + + fn find_local_shard( + &self, + value: MidValueId, + tile: u16, + ) -> LowLoweringResult> { + Ok(self + .value_shards(value)? + .iter() + .copied() + .find(|shard| self.shards[shard.index() as usize].tile == tile)) + } + + fn corresponding_shard( + &self, + value: MidValueId, + target: LowShardId, + ) -> LowLoweringResult { + let target = &self.shards[target.index() as usize]; + self.value_shards(value)? + .iter() + .copied() + .filter(|shard| self.shards[shard.index() as usize].extents == target.extents) + .min_by_key(|shard| u8::from(self.shards[shard.index() as usize].tile != target.tile)) + .ok_or(LowLoweringError::UnknownValue(value)) + } + + fn shard_stride( + &self, + shard: LowShardId, + alignment: u32, + access_tail: u32, + ) -> LowLoweringResult { + let shard = &self.shards[shard.index() as usize]; + let elements = shard + .extents + .iter() + .try_fold(1_u64, |elements, extent| { + elements.checked_mul(u64::from(extent.physical_end - extent.start)) + }) + .ok_or(LowLoweringError::IdOverflow)?; + let bytes = elements + .checked_mul(shard.tensor_type.format.precision.bytes()) + .and_then(|bytes| bytes.checked_add(u64::from(access_tail))) + .ok_or(LowLoweringError::IdOverflow)?; + let alignment = u64::from(alignment.max(1)); + let stride = bytes + .checked_add(alignment - 1) + .map(|bytes| bytes / alignment * alignment) + .ok_or(LowLoweringError::IdOverflow)?; + u32::try_from(stride).map_err(|_| LowLoweringError::IdOverflow) + } + } + + fn append_span_copies( + shards: &[LowShard], + source: &ShardView, + destination: &ShardView, + tile: u16, + copies: &mut Vec<(u16, LocalCopy)>, + ) -> LowLoweringResult<()> { + let source_spans = view_byte_spans(&shards[source.shard.index() as usize], source)?; + let destination_spans = + view_byte_spans(&shards[destination.shard.index() as usize], destination)?; + append_byte_span_copies( + source, + destination, + tile, + &source_spans, + &destination_spans, + copies, + ) + } + + fn append_byte_span_copies( + source: &ShardView, + destination: &ShardView, + tile: u16, + source_spans: &[ByteSpan], + destination_spans: &[ByteSpan], + copies: &mut Vec<(u16, LocalCopy)>, + ) -> LowLoweringResult<()> { + let mut pending = Vec::new(); + let mut source_index = 0usize; + let mut destination_index = 0usize; + let mut source_offset = 0u32; + let mut destination_offset = 0u32; + while source_index < source_spans.len() && destination_index < destination_spans.len() { + let source_span = source_spans[source_index]; + let destination_span = destination_spans[destination_index]; + let bytes = + (source_span.bytes - source_offset).min(destination_span.bytes - destination_offset); + pending.push(LocalCopy { + source: source.shard, + source_offset: source_span.offset + source_offset, + destination: destination.shard, + destination_offset: destination_span.offset + destination_offset, + bytes, + pattern: LocalCopyPattern::Contiguous, + }); + source_offset += bytes; + destination_offset += bytes; + if source_offset == source_span.bytes { + source_index += 1; + source_offset = 0; + } + if destination_offset == destination_span.bytes { + destination_index += 1; + destination_offset = 0; + } + } + if source_index != source_spans.len() || destination_index != destination_spans.len() { + return Err(LowLoweringError::InvalidConversionPlan); + } + copies.extend( + coalesce_local_copies(pending) + .into_iter() + .map(|copy| (tile, copy)), + ); + Ok(()) + } + + const PARALLEL_STRIDED_COPY_MAX_BYTES: u32 = 512; + + fn coalesce_local_copies(copies: Vec) -> Vec { + let mut coalesced = Vec::new(); + let mut index = 0; + while index < copies.len() { + let first = &copies[index]; + let Some(second) = copies.get(index + 1) else { + coalesced.push(first.clone()); + break; + }; + if first.source != second.source + || first.destination != second.destination + || first.bytes != second.bytes + || first.bytes == 0 + || !first.bytes.is_multiple_of(8) + { + coalesced.push(first.clone()); + index += 1; + continue; + } + let source_stride = second.source_offset.saturating_sub(first.source_offset); + let destination_stride = second + .destination_offset + .saturating_sub(first.destination_offset); + if source_stride == 0 || destination_stride == 0 { + coalesced.push(first.clone()); + index += 1; + continue; + } + let mut end = index + 2; + while let Some(copy) = copies.get(end) { + let previous = &copies[end - 1]; + if copy.source != first.source + || copy.destination != first.destination + || copy.bytes != first.bytes + || copy.source_offset.checked_sub(previous.source_offset) != Some(source_stride) + || copy + .destination_offset + .checked_sub(previous.destination_offset) + != Some(destination_stride) + { + break; + } + end += 1; + } + let rows = u32::try_from(end - index).unwrap_or(u32::MAX); + // Larger strided regions are deliberately left as contiguous rows: + // spreading them over workers loses more to bank contention than it + // saves in call overhead on IPU21. + if first.bytes.saturating_mul(rows) > PARALLEL_STRIDED_COPY_MAX_BYTES { + coalesced.extend(copies[index..end].iter().cloned()); + index = end; + continue; + } + if source_stride == first.bytes && destination_stride == first.bytes { + let mut copy = first.clone(); + copy.bytes = copy.bytes.saturating_mul(rows); + coalesced.push(copy); + } else { + let mut copy = first.clone(); + copy.bytes = copy.bytes.saturating_mul(rows); + copy.pattern = LocalCopyPattern::Strided { + rows, + row_bytes: first.bytes, + source_stride, + destination_stride, + }; + coalesced.push(copy); + } + index = end; + } + coalesced + } + + fn append_logical_span_copies( + shards: &[LowShard], + source: &ShardView, + destination: &ShardView, + tile: u16, + copies: &mut Vec<(u16, LocalCopy)>, + ) -> LowLoweringResult<()> { + let source_spans = logical_view_byte_spans(&shards[source.shard.index() as usize], source)?; + let destination_spans = + logical_view_byte_spans(&shards[destination.shard.index() as usize], destination)?; + append_byte_span_copies( + source, + destination, + tile, + &source_spans, + &destination_spans, + copies, + ) + } + + fn planned_local_copies( + source: LowShardId, + destination: LowShardId, + geometry: &crate::CopyGeometry, + ) -> LowLoweringResult> { + let mut inner = geometry.clone(); + let retain_inner = inner.dimensions.first().is_some_and(|dimension| { + inner.contiguous_bytes.is_multiple_of(8) + && u64::from(inner.contiguous_bytes) * u64::from(dimension.count) <= 512 + }); + let outer = inner.dimensions.split_off(usize::from(retain_inner)); + let offsets = crate::CopyGeometry { + dimensions: outer, + ..inner.clone() + } + .offsets() + .ok_or(LowLoweringError::IdOverflow)?; + offsets + .into_iter() + .map(|(source_offset, destination_offset)| { + Ok(LocalCopy { + source, + source_offset, + destination, + destination_offset, + bytes: u32::try_from(inner.bytes()).map_err(|_| LowLoweringError::IdOverflow)?, + pattern: inner.dimensions.first().map_or( + LocalCopyPattern::Contiguous, + |dimension| LocalCopyPattern::Strided { + rows: dimension.count, + row_bytes: inner.contiguous_bytes, + source_stride: dimension.source_stride, + destination_stride: dimension.destination_stride, + }, + ), + }) + }) + .collect() + } + + fn value_can_alias(value: MidValueId, target: MidValueId, operations: &[MidOperation]) -> bool { + if value == target { + return true; + } + let Some(operation) = operations + .iter() + .find(|operation| operation.results.contains(&value)) + else { + return false; + }; + let Some(plan) = operation.operator_plan() else { + return false; + }; + let indices = match &plan.requirements.output_aliasing { + OutputAliasing::Fresh => return false, + OutputAliasing::MayAliasInputs(indices) => indices.as_slice(), + OutputAliasing::MustAliasInput(index) => std::slice::from_ref(index), + }; + indices.iter().any(|index| { + operation + .inputs + .get(usize::from(*index)) + .is_some_and(|input| value_can_alias(*input, target, operations)) + }) + } + + fn repeat_yield_can_alias( + value: MidValueId, + carried: MidValueId, + operations: &[MidOperation], + ) -> bool { + if value_can_alias(value, carried, operations) { + return true; + } + let Some(definition) = operations + .iter() + .position(|operation| operation.results.contains(&value)) + else { + return false; + }; + // A repeat reuses the carried allocation on its next iteration. A fresh + // yield may overwrite it when every read of the previous iteration's + // value has completed before the yielding operation begins. + !operations[definition..] + .iter() + .any(|operation| operation.inputs.contains(&carried)) + } + + fn body_storage_requirement(value: MidValueId, operations: &[MidOperation]) -> (u32, u32) { + let mut alignment = 8; + let mut access_tail = 0; + for operation in operations { + for (index, input) in operation.inputs.iter().enumerate() { + if *input != value { + continue; + } + let requirement = operation + .operator_plan() + .and_then(|plan| plan.requirements.inputs.get(index)); + if let Some(requirement) = requirement { + alignment = alignment.max(requirement.allocation.alignment); + access_tail = access_tail.max(requirement.allocation.access_tail_bytes); + } else if operation.conversion().is_some() { + alignment = alignment.max(8); + } + } + } + (alignment, access_tail) + } + + fn operation_provenance(operation: &MidOperation, kind: &MidOperationKind) -> WorkProvenance { + WorkProvenance { + operation: operation.source, + value: operation.results.first().copied(), + reason: match kind { + MidOperationKind::CastPrecision => WorkReason::PrecisionCast, + MidOperationKind::Convert(..) => WorkReason::LayoutRearrangement, + MidOperationKind::Operator(_) => WorkReason::OperatorKernel, + MidOperationKind::Repeat(_) => WorkReason::Repeat, + }, + } + } + + fn intersect_extents(left: &[ShardExtent], right: &[ShardExtent]) -> Option> { + if left.len() != right.len() { + return None; + } + left.iter() + .zip(right) + .map(|(left, right)| { + 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() + } + + fn intersect_extents_with_shared_padding( + left: &[ShardExtent], + right: &[ShardExtent], + ) -> Option> { + if left.len() != right.len() { + return None; + } + left.iter() + .zip(right) + .map(|(left, right)| { + let start = left.start.max(right.start); + let logical_end = left.logical_end.min(right.logical_end); + (start < logical_end).then(|| { + let shared_tail = + if logical_end == left.logical_end && logical_end == right.logical_end { + left.physical_end + .saturating_sub(left.logical_end) + .min(right.physical_end.saturating_sub(right.logical_end)) + } else { + 0 + }; + ShardExtent { + axis: left.axis, + start, + logical_end, + physical_end: logical_end + shared_tail, + } + }) + }) + .collect() + } + + fn shard_extents(tensor_type: &TensorType) -> LowLoweringResult> { + Ok(tensor_type + .format + .layout + .resolve(&tensor_type.shape)? + .shard_extents() + .into_iter() + .map(|shard| (shard.tile, shard.extents)) + .collect()) + } + + fn split_mapping_at_panel_boundaries( + source_shard: &LowShard, + mut source: ShardView, + destination_shard: &LowShard, + mut destination: ShardView, + ) -> LowLoweringResult> { + let source_rank = source.extents.len(); + let destination_rank = destination.extents.len(); + let outer_elements = |extents: &[ShardExtent]| { + extents[..extents.len().saturating_sub(2)] + .iter() + .try_fold(1_u32, |elements, extent| { + elements.checked_mul(extent.logical_end - extent.start) + }) + }; + if source_rank < 2 + || destination_rank < 2 + || source_shard.extents.len() != source_rank + || destination_shard.extents.len() != destination_rank + || outer_elements(&source.extents) != Some(1) + || outer_elements(&destination.extents) != Some(1) + { + return Err(LowLoweringError::InvalidOperatorPlan); + } + + let aligned_ranges = |source: ShardExtent, + source_shard: ShardExtent, + destination: ShardExtent, + destination_shard: ShardExtent| + -> LowLoweringResult> { + let logical_width = source.logical_end - source.start; + if logical_width != destination.logical_end - destination.start { + return Err(LowLoweringError::InvalidOperatorPlan); + } + let width = source.physical_end - source.start; + if width != destination.physical_end - destination.start { + return Err(LowLoweringError::InvalidOperatorPlan); + } + let mut ranges = Vec::new(); + let mut offset = 0; + while offset < width { + let source_position = source + .start + .checked_sub(source_shard.start) + .and_then(|start| start.checked_add(offset)) + .ok_or(LowLoweringError::InvalidOperatorPlan)?; + let destination_position = destination + .start + .checked_sub(destination_shard.start) + .and_then(|start| start.checked_add(offset)) + .ok_or(LowLoweringError::InvalidOperatorPlan)?; + let source_remaining = AMP_COLUMN_MICRO - source_position % AMP_COLUMN_MICRO; + let destination_remaining = AMP_COLUMN_MICRO - destination_position % AMP_COLUMN_MICRO; + let length = (width - offset) + .min(source_remaining) + .min(destination_remaining); + let source_start = source.start + offset; + let destination_start = destination.start + offset; + ranges.push(( + ShardExtent { + axis: source.axis, + start: source_start, + logical_end: source + .logical_end + .min(source_start + length) + .max(source_start), + physical_end: source_start + length, + }, + ShardExtent { + axis: destination.axis, + start: destination_start, + logical_end: destination + .logical_end + .min(destination_start + length) + .max(destination_start), + physical_end: destination_start + length, + }, + )); + offset += length; + } + Ok(ranges) + }; + + let source_row_axis = source_rank - 2; + let source_column_axis = source_rank - 1; + let destination_row_axis = destination_rank - 2; + let destination_column_axis = destination_rank - 1; + + // The global row tail can finish part-way through a micro-panel while + // both allocations contain padding through the same panel boundary. + // Carry that padding with the useful values so the direct physical + // exchange remains word-aligned. A split head's column tail is not + // extended because the following source columns may belong to another + // head rather than padding. + let source_rows = source.extents[source_row_axis]; + let destination_rows = destination.extents[destination_row_axis]; + if source_rows.logical_end == source_shard.tensor_type.shape.0[source_row_axis] + && destination_rows.logical_end + == destination_shard.tensor_type.shape.0[destination_row_axis] + { + let source_panel_tail = (AMP_COLUMN_MICRO + - (source_rows.logical_end - source_shard.extents[source_row_axis].start) + % AMP_COLUMN_MICRO) + % AMP_COLUMN_MICRO; + let destination_panel_tail = (AMP_COLUMN_MICRO + - (destination_rows.logical_end + - destination_shard.extents[destination_row_axis].start) + % AMP_COLUMN_MICRO) + % AMP_COLUMN_MICRO; + let padding = source_panel_tail + .min(destination_panel_tail) + .min(source_shard.extents[source_row_axis].physical_end - source_rows.logical_end) + .min( + destination_shard.extents[destination_row_axis].physical_end + - destination_rows.logical_end, + ); + source.extents[source_row_axis].physical_end += padding; + destination.extents[destination_row_axis].physical_end += padding; + } + + let rows = aligned_ranges( + source.extents[source_row_axis], + source_shard.extents[source_row_axis], + destination.extents[destination_row_axis], + destination_shard.extents[destination_row_axis], + )?; + let columns = aligned_ranges( + source.extents[source_column_axis], + source_shard.extents[source_column_axis], + destination.extents[destination_column_axis], + destination_shard.extents[destination_column_axis], + )?; + let mut pieces = Vec::with_capacity(rows.len().saturating_mul(columns.len())); + for (source_row, destination_row) in rows { + for &(source_column, destination_column) in &columns { + let mut source_extents = source.extents.clone(); + let mut destination_extents = destination.extents.clone(); + source_extents[source_row_axis] = source_row; + source_extents[source_column_axis] = source_column; + destination_extents[destination_row_axis] = destination_row; + destination_extents[destination_column_axis] = destination_column; + pieces.push(( + ShardView { + shard: source.shard, + extents: source_extents, + }, + ShardView { + shard: destination.shard, + extents: destination_extents, + }, + )); + } + } + Ok(pieces) ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec } diff --cc crates/ipu-codegen/src/memory.rs index 58da982,aad0973..0000000 --- a/crates/ipu-codegen/src/memory.rs +++ b/crates/ipu-codegen/src/memory.rs @@@ -31,9 -3,8 +3,14 @@@ use ipu_package::AddressRegion #[derive(Clone, Debug, PartialEq, Eq)] pub(crate) struct MemoryAllocation { pub name: &'static str, ++<<<<<<< HEAD + pub range: Range, + /// Entire reservation, including end alignment and any trailing guard. + pub reserved: Range, ++======= + pub range: AddressRegion, + reserved: AddressRegion, ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec } #[derive(Clone, Debug)] diff --cc crates/ipu-codegen/src/mid/mod.rs index ba07b78,f2f9507..0000000 --- a/crates/ipu-codegen/src/mid/mod.rs +++ b/crates/ipu-codegen/src/mid/mod.rs @@@ -1,355 -1,61 +1,416 @@@ ++<<<<<<< HEAD +//! Whole-device implementation selection, tensor primitives, and layouts. +//! Tile enumeration occurs only after selection, in low expansion. + +pub(crate) mod implementation; +mod primitive; +pub use primitive::*; +mod candidates; +mod catalogue; +mod layout; +mod operator; +mod ownership; +mod planner; +mod resolved; +mod view; +pub use crate::graph::AxisFactorView; + +use candidates::*; +use catalogue::*; +pub use catalogue::{OperatorCandidate, OperatorFormatPolicy}; +pub use layout::*; +pub use operator::*; +#[cfg(test)] +pub(crate) use planner::lower; +pub(crate) fn lower_finalists( + graph: &ComputeGraph, + config: &PipelineConfig, + costs: &impl CostModel, + count: usize, +) -> LoweringResult> { + planner::plan_finalists(graph, config, costs, count)? + .into_iter() + .map(|program| implementation::resolve(program).ok_or(LoweringError::InvalidImplementation)) + .collect() +} +#[cfg(test)] +pub(crate) fn expand_tiles( + program: &MidProgram, +) -> crate::ExpansionResult> { + let program = implementation::resolve(program.clone()) + .ok_or(crate::ExpansionError::InvalidOperatorPlan)?; + crate::low::expand::expand_tiles(&program) +} +use planner::*; + +use crate::estimate::MemoizedCostModel; +pub use crate::estimate::{ + CostModel, IPU21_TARGET_COSTS, Ipu21CostModel, MemoryPeaks, MemoryUsage, +}; +use crate::estimate::{region_peak_memory, region_peak_memory_with_multiplicity}; +use crate::graph::{ + AttentionOptions, ComputeGraph, GemmOptions, GraphInputKind, Operation, OperationId, + OperationKind, Repeat, TensorShape, ValueId, +}; +use rayon::prelude::*; +use std::collections::{BTreeMap, BTreeSet}; + +/// Exact blocked-GEMM geometry retained for planner diagnosis. Constraints +/// are keyed by the source graph operation and bypass candidate 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 orientation: GemmOrientation, + pub row_partitions: u16, + pub column_partitions: u16, + pub inner_partitions: u16, + pub result_row_partitions: u16, + pub result_column_partitions: u16, + pub output_column_block: u32, + pub weight_memory_class: MemoryClass, + pub reduction_staging: ReductionStaging, + pub local_weight_staging: LocalOperandStaging, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct PipelineConfig { + 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, + /// Signatures available independently to each operation. Earlier entries + /// of the appropriate operation kind win when costs are equal. + pub operator_candidates: Vec, + /// Add near-capacity tile counts derived from graph tensor extents. + pub shape_aware_active_tile_counts: bool, + /// Maximum number of partial format assignments retained after each + /// operation in a straight-line region. + pub planning_beam_width: usize, + /// Number of complete beam finalists to materialize and rank with the + /// physical exchange scheduler. One retains analytical-only selection. + pub exchange_schedule_finalists: usize, + /// Diagnostic constraints which retain only one GEMM plan family for the + /// named source operations. + pub gemm_plan_constraints: Vec, + /// Standard-addressed SRAM retained for exchange tables, profiling data, + /// host commands, and generated tile programs built after planning. + pub standard_memory_reservation_bytes: u64, + /// Maximum SRAM per tile available to planned values and the standard + /// reservation. Lower values emulate a model whose other persistent state + /// occupies the remainder of SRAM. + pub tile_memory_budget_bytes: u64, + pub profiling: bool, + /// Insert all-tile patched-breakpoint stops after semantic operators. + pub diagnostic_checkpoints: bool, + /// Emit exchange-scheduler lower bounds, per-tile role pressure, and + /// critical dependency chains while constructing the final package. + pub exchange_diagnostics: bool, + /// Controls whether one-use layout conversions may be populated as + /// bounded slices immediately before their consuming dispatch. + pub conversion_streaming: ConversionStreamingPolicy, + /// Restricts attention planning to one execution strategy for controlled + /// benchmarking; automatic planning retains both alternatives. + pub attention_strategy: AttentionStrategy, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum ConversionStreamingPolicy { + /// Require complete converted values. + Never, + /// Prefer complete values, retaining streaming when materialization does + /// not fit the target memory budget. + #[default] + WhenRequired, + /// Stream every eligible conversion, primarily for diagnostics and + /// memory-constrained deployment experiments. + Always, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum AttentionStrategy { + #[default] + Automatic, + Flash, + Materialized, +} + +impl PipelineConfig { + pub fn new(tile_count: u16) -> Self { + Self { + tile_count, + inputs: BTreeMap::new(), + automatic_inputs: BTreeMap::new(), + operator_candidates: default_operator_candidates(tile_count), + shape_aware_active_tile_counts: true, + planning_beam_width: 64, + exchange_schedule_finalists: 1, + gemm_plan_constraints: Vec::new(), + standard_memory_reservation_bytes: u64::from( + crate::memory::IPU21_DEFAULT_SUPPORT_RESERVATION_BYTES, + ), + tile_memory_budget_bytes: u64::from(crate::memory::IPU21_PLANNED_DATA_BYTES), + profiling: false, + diagnostic_checkpoints: false, + exchange_diagnostics: false, + conversion_streaming: ConversionStreamingPolicy::WhenRequired, + attention_strategy: AttentionStrategy::Automatic, + } + } + + 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_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 + } + + /// Restrict default operator planning to explicit active tile counts. + /// This is useful when evaluating a fixed occupancy rather than allowing + /// the planner to trade occupancy against communication and memory use. + pub fn with_active_tile_counts(mut self, counts: impl IntoIterator) -> Self { + let mut candidates = Vec::new(); + for count in counts { + if count == 0 || count > self.tile_count { + continue; + } + candidates.extend(operator_candidates_for_tile_count(count)); + } + candidates.dedup(); + self.operator_candidates = candidates; + self.shape_aware_active_tile_counts = false; + 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 + } +} + +#[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, + /// Selected rotation of layout ownership within the graph's tile group. + pub tile_offset: u16, + 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 { + Primitive(Primitive), + Operator { + plan: OperatorPlan, + deferred_inputs: Vec>, + implementation: Option>, + }, + Convert(ConversionPlan), + Repeat(MidRepeat), +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct MidOperation { + pub source: Option, + pub inputs: Vec, + pub results: Vec, + pub kind: MidOperationKind, + pub estimated_cycles: u64, + pub estimated_exchange_cycles: u64, +} + +impl MidOperation { + pub fn operator_plan(&self) -> Option<&OperatorPlan> { + match &self.kind { + MidOperationKind::Operator { plan, .. } => Some(plan), + _ => None, + } + } + + pub fn deferred_inputs(&self) -> &[Option] { + match &self.kind { + MidOperationKind::Operator { + deferred_inputs, .. + } => deferred_inputs, + _ => &[], + } + } + + pub fn conversion_plan(&self) -> Option<&ConversionPlan> { + match &self.kind { + MidOperationKind::Convert(plan) => Some(plan), + _ => None, + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct MidRegion { + pub arguments: Vec, + pub operations: Vec, + pub yields: Vec, + pub estimated_cycles: u64, + pub peak_memory: MemoryPeaks, +} + +#[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, +} + +/// Whole-device tensor program. Search recipes retain compact implementations; +/// final selection inlines those primitives before low enumerates tiles. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct MidProgram { + pub tile_count: u16, + pub inputs: Vec, + pub values: Vec, + pub operations: Vec, + pub outputs: Vec, + pub estimated_cycles: u64, + pub estimated_exchange_cycles: u64, + pub peak_memory: MemoryPeaks, +} + +// Estimation policy is kept in `estimate` so this module remains focused on IR and lowering. + +#[derive(Debug, thiserror::Error, PartialEq, Eq)] +pub enum LoweringError { + #[error("selected operator implementation is invalid")] + InvalidImplementation, + #[error("cannot create planning worker pool: {0}")] + PlanningThreads(String), + #[error(transparent)] + Layout(#[from] LayoutError), + #[error(transparent)] + Storage(#[from] crate::storage::StorageError), ++======= + //! Mid-level, layout-aware representation. + //! + //! This is the boundary between semantic graph operations and scheduling. It + //! records tensor shapes, storage precision, element order, axis tiling, and + //! memory-class requirements, but deliberately does not assign tile addresses + //! or emit exchange rows. [`lower`] tries a set of legal operator plans, + //! prices them with a [`CostModel`], and inserts explicit precision casts and + //! layout rearrangements at format boundaries. + + mod consumers; + mod gemm; + mod ir; + + use consumers::{direct_consumer_layouts, operator_accepts_input_layout}; + #[cfg(test)] + use gemm::{ + AmpWeightPlacement, amp_grid_gemm_plan, gemm_accumulation_precision, gemm_map, + gemm_seed_plans_for_tile_count, independent_parameter_storage, parallel_reduction_plans, + }; + use gemm::{GroupedOutputLayout, gemm_plan_matches, gemm_plans, grouped_output_layout}; + pub use ir::{ + MidGraph, MidInput, MidOperation, MidOperationKind, MidRegion, MidRepeat, MidValue, MidValueId, + }; + + use crate::TileKernelSpec; + #[cfg(test)] + use crate::config::PlannerSearchDomain; + use crate::config::{AttentionStrategy, ConversionStreamingPolicy, OperatorClass, PipelineConfig}; + use crate::conversion::{ + ConversionGeometryError, ConversionStrategy, DeferredTransform, finalize_conversion_plans, + layout_conversion_strategy, + }; + pub use crate::cost::{CostModel, Ipu21CostModel}; + use crate::cost::{ + MemoizedCostModel, conversion_memory_estimate, operator_memory_estimate, region_peak_memory, + region_peak_memory_with_multiplicity, + }; + use crate::graph::{ + ComputeGraph, GraphInputKind, Operation, OperationId, OperationKind, Repeat, TensorShape, + ValueId, + }; + use crate::layout::{ + AMP_COLUMN_MICRO, AMP_INNER_BLOCK, Layout, MemoryClass, NativeKernelOrder, StorageOrder, + TensorAxis, TensorFormat, TensorType, + }; + pub use crate::metrics::{CostEstimate, ExchangeFootprint}; + use crate::metrics::{MemoryEstimate, MemoryPeaks, MemoryUsage, OperationMetrics, RegionMetrics}; + use crate::operator::*; + use crate::schedule::{ + AttentionBlocking, AttentionMap, KernelMap, OperatorSchedule, ScheduleAccess, ScheduleStep, + ScheduleValue, + }; + use ipu_target::hardware::HardwareTarget; + use rayon::prelude::*; + use std::collections::{BTreeMap, BTreeSet}; + + #[derive(Debug, thiserror::Error, PartialEq, Eq)] + pub enum LoweringError { ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec #[error("mid-level lowering requires a nonzero tile count")] EmptyTileGroup, #[error("no tensor type was supplied for graph input {0:?}")] @@@ -375,60 -81,3179 +436,3236 @@@ UnsupportedGemmBatching(OperationId), #[error("internal lowering error: value {0:?} is unavailable")] UnknownValue(ValueId), ++<<<<<<< HEAD ++======= + #[error(transparent)] + ConversionGeometry(#[from] ConversionGeometryError), ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec } pub type LoweringResult = std::result::Result; ++<<<<<<< HEAD +#[cfg(test)] +mod tests; + +#[cfg(test)] +#[test] +#[ignore = "manual full-size planner timing; does not access hardware"] +fn profile_mlp_finalist_expansion() { + let mut graph = ComputeGraph::new(); + let input = graph.host_input("input", [1, 729, 1152]).unwrap(); + let up = graph.parameter("up", [1, 1152, 4304]).unwrap(); + let down = graph.parameter("down", [1, 4304, 1152]).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 config = PipelineConfig::new(1472) + .with_automatic_input(input, Precision::F16) + .with_automatic_input(up, Precision::F16) + .with_automatic_input(down, Precision::F16); + let start = std::time::Instant::now(); + let finalists = planner::plan_finalists(&graph, &config, &crate::Ipu21CostModel, 8).unwrap(); + eprintln!( + "compact planning: {:?}, {} finalists", + start.elapsed(), + finalists.len() + ); + for (index, mid) in finalists.into_iter().enumerate() { + for operation in &mid.operations { + if let Some(plan) = operation.operator_plan() + && let OperatorDispatch::BlockedGemm { orientation, .. } = plan.dispatch + { + eprintln!( + "finalist {index}: source {:?}, {:?}, weight memory {:?}", + operation.source, + plan.dispatch, + plan.requirements.inputs[orientation.operand_indices().1] + .format + .layout + .memory_class, + ); + } + } + let mid = implementation::resolve(mid.clone()).unwrap(); + let start = std::time::Instant::now(); + let expanded = crate::low::expand::expand_tiles(&mid).unwrap(); + eprintln!( + "finalist {index}: expansion {:?}, estimated cycles {}, exchange {}", + start.elapsed(), + expanded.estimated_cycles, + expanded.estimated_exchange_cycles + ); ++======= + #[tracing::instrument( + name = "ipu_codegen.mid.lower", + skip(graph, config, costs), + fields(tile_count = config.tile_count, operations = graph.operations().len()) + )] + pub fn lower( + graph: &ComputeGraph, + config: &PipelineConfig, + costs: &impl CostModel, + ) -> LoweringResult { + Ok(lower_finalists(graph, config, costs, 1)?.remove(0)) + } + + pub(crate) fn lower_finalists( + graph: &ComputeGraph, + config: &PipelineConfig, + costs: &impl CostModel, + finalist_count: usize, + ) -> LoweringResult> { + if config.tile_count == 0 { + return Err(LoweringError::EmptyTileGroup); + } + let active_tile_counts = config + .search_domain + .active_tile_counts(config.tile_count, graph.value_shapes().values()); + let mut resolved_config = config.clone(); + resolved_config.resolved_active_tile_counts = active_tile_counts.clone(); + let config = &resolved_config; + let mut state = LoweringState::default(); + let costs = MemoizedCostModel::new(costs, config.tile_count); + let mut values = BTreeMap::new(); + let mut inputs = Vec::with_capacity(graph.inputs().len()); + for input in graph.inputs() { + let (format, automatic) = if let Some(format) = config.inputs.get(&input.value) { + (format.clone(), false) + } else if let Some(&precision) = config.automatic_inputs.get(&input.value) { + ( + TensorFormat { + precision, + layout: Layout::row_sharded(config.tile_count), + }, + true, + ) + } else { + return Err(LoweringError::MissingInputType(input.value)); + }; + let tensor_type = TensorType { + shape: input.shape.clone(), + format, + }; + let value = state.value(input.value, tensor_type); + if input.kind == GraphInputKind::Parameter { + state.parameter_values.insert(value); + } + if automatic { + state.automatic_inputs.insert(value); + } + values.insert(input.value, value); + inputs.push(MidInput { + name: input.name.clone(), + kind: input.kind, + value, + }); + } + let branches = plan_region_frontier( + graph.operations(), + graph.outputs(), + &mut values, + graph.value_shapes(), + graph, + config, + &costs, + &mut state, + &RegionPlanningConstraints::default(), + )?; + let initial = inputs.iter().map(|input| input.value).collect::>(); + branches + .into_iter() + .take(finalist_count.max(1)) + .enumerate() + .map(|(finalist, branch)| { + let outputs = graph + .outputs() + .iter() + .map(|value| lookup(&branch.values, *value)) + .collect::>>()?; + let estimated_cycles = branch + .operations + .iter() + .map(|operation| operation.metrics.cost.cycles) + .sum::(); + let estimated_exchange_cycles = branch + .operations + .iter() + .map(|operation| operation.metrics.cost.exchange_cycles) + .sum::(); + let cost = branch + .operations + .iter() + .fold(CostEstimate::default(), |cost, operation| { + cost.sequence(operation.metrics.cost) + }); + let peak_memory = region_peak_memory( + &initial, + &branch.operations, + &outputs, + &branch.state.values, + config.target, + ); + tracing::info!( + finalist, + values = branch.state.values.len(), + operations = branch.operations.len(), + estimated_cycles, + estimated_exchange_cycles, + exchange_row_bytes = peak_memory.exchange_rows, + peak_standard_bytes = peak_memory.standard, + peak_interleaved_bytes = peak_memory.interleaved, + peak_total_bytes = peak_memory.total, + maximum_standard_allocation_bytes = peak_memory.maximum_standard_allocation, + active_tile_counts = ?branch.operations + .iter() + .filter_map(|operation| operation.results.first()) + .map(|result| branch.state.values[result.index() as usize] + .tensor_type + .format + .layout + .tiling + .tile_count) + .collect::>(), + padding_group_counts = ?branch.operations + .iter() + .flat_map(|operation| operation.results.iter()) + .flat_map(|result| branch.state.values[result.index() as usize] + .tensor_type + .format + .layout + .tiling + .axes + .iter() + .map(|axis| axis.padding_groups)) + .filter(|groups| *groups > 1) + .collect::>(), + conversion_sources = ?branch.operations + .iter() + .filter(|operation| operation.conversion().is_some()) + .map(|operation| operation.source) + .collect::>(), + "retained operator-plan finalist" + ); + tracing::debug!( + finalist, + plans = ?branch.operations + .iter() + .filter_map(|operation| operation.operator_plan().map(|plan| ( + operation.source, + plan, + plan.requirements.inputs.iter().map(|input| &input.format.layout).collect::>(), + &plan.requirements.output.format.layout, + operation.metrics.cost.cycles, + operation.metrics.cost.exchange_cycles, + ))) + .collect::>(), + conversions = ?branch.operations + .iter() + .filter_map(|operation| Some(( + operation.source, + &branch.state.get(*operation.inputs.first()?).tensor_type.format.layout, + &branch.state.get(*operation.results.first()?).tensor_type.format.layout, + operation.metrics.cost.cycles, + operation.metrics.cost.exchange_cycles, + )).filter(|_| operation.conversion().is_some())) + .collect::>(), + "retained operator-plan details" + ); + let mut graph = MidGraph { + inputs: inputs.clone(), + values: branch.state.values, + operations: branch.operations, + outputs, + metrics: RegionMetrics { + cost, + memory: peak_memory, + }, + }; + finalize_conversion_plans(&mut graph)?; + Ok(graph) + }) + .collect() + } + + #[derive(Clone, Default)] + struct LoweringState { + values: Vec, + automatic_inputs: BTreeSet, + parameter_values: BTreeSet, + } + + impl LoweringState { + fn value(&mut self, origin: ValueId, tensor_type: TensorType) -> MidValueId { + let id = MidValueId::from_index(self.values.len() as u32); + self.values.push(MidValue { + id, + tensor_type, + origin, + storage_group: id, + }); + id + } + + fn value_in_storage_group( + &mut self, + origin: ValueId, + tensor_type: TensorType, + storage_group: MidValueId, + ) -> MidValueId { + let result = self.value(origin, tensor_type); + self.values[result.index() as usize].storage_group = storage_group; + result + } + + fn get(&self, id: MidValueId) -> &MidValue { + &self.values[id.index() as usize] + } + + fn derived_value(&mut self, source: MidValueId, tensor_type: TensorType) -> MidValueId { + let origin = self.get(source).origin; + let storage_group = self.get(source).storage_group; + let result = self.value_in_storage_group(origin, tensor_type, storage_group); + if self.parameter_values.contains(&source) { + self.parameter_values.insert(result); + } + result + } + + fn retarget_automatic_input(&mut self, id: MidValueId, layout: Layout) -> bool { + if !self.automatic_inputs.remove(&id) { + return false; + } + self.values[id.index() as usize].tensor_type.format.layout = layout; + true + } + } + + #[derive(Clone)] + struct BeamBranch { + values: BTreeMap, + state: LoweringState, + operations: Vec, + peak_memory: MemoryPeaks, + } + + #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)] + struct FutureValueState { + origin: ValueId, + tensor_type: TensorType, + automatic_input: bool, + parameter: bool, + allocation_copies: u32, + storage_class: u32, + } + + #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)] + struct FutureBeamState { + values: Vec, + equal_formats_satisfied: Vec<(ValueId, ValueId, bool)>, + } + + struct RankedBeamBranch { + branch: BeamBranch, + objective: RegionMetrics, + compatibility: FutureFormatCompatibility, + order: usize, + } + + type FutureFormatCompatibility = Vec<( + ValueId, + Precision, + StorageOrderCompatibility, + MemoryClass, + Vec<(TensorAxis, u16, u32)>, + )>; + + #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)] + enum StorageOrderCompatibility { + Linear, + Blocked { + axes: [TensorAxis; 2], + permutation: [u8; 2], + }, + Native(NativeKernelOrder), + } + + fn storage_order_compatibility(order: StorageOrder) -> StorageOrderCompatibility { + match order { + StorageOrder::Linear => StorageOrderCompatibility::Linear, + StorageOrder::Blocked(order) => StorageOrderCompatibility::Blocked { + axes: order.axes, + permutation: order.permutation, + }, + StorageOrder::Native(order) => StorageOrderCompatibility::Native(order), + } + } + + fn future_format_compatibility( + branch: &BeamBranch, + future_origins: &BTreeSet, + ) -> FutureFormatCompatibility { + let mut formats = Vec::new(); + for &origin in future_origins { + let Some(&id) = branch.values.get(&origin) else { + continue; + }; + let mut add = |format: &TensorFormat| { + let axes = format + .layout + .tiling + .axes + .iter() + .map(|axis| (axis.axis, axis.padding_groups, axis.shard_padding_multiple)) + .collect(); + formats.push(( + origin, + format.precision, + storage_order_compatibility(format.layout.order), + format.layout.memory_class, + axes, + )); + }; + add(&branch.state.get(id).tensor_type.format); + } + formats + } + + #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)] + struct PlanCacheKey { + input_shapes: Vec, + parameter_inputs: Vec, + format_sensitive_inputs: Vec<(usize, TensorFormat)>, + } + + #[derive(Default)] + struct RegionPlanningConstraints { + /// Number of simultaneously resident blocks represented by a region value. + allocation_copies: BTreeMap, + /// Value pairs whose formats must agree at a structured-region boundary. + required_equal_formats: Vec<(ValueId, ValueId)>, + } + + #[allow(clippy::too_many_arguments)] + fn plan_region_frontier( + source: &[Operation], + required_outputs: &[ValueId], + values: &mut BTreeMap, + shapes: &BTreeMap, + graph: &ComputeGraph, + config: &PipelineConfig, + costs: &impl CostModel, + state: &mut LoweringState, + constraints: &RegionPlanningConstraints, + ) -> LoweringResult> { + if source.is_empty() { + return Ok(vec![BeamBranch { + values: values.clone(), + state: state.clone(), + operations: Vec::new(), + peak_memory: MemoryPeaks::default(), + }]); + } + let relevant_origins = source + .iter() + .flat_map(|operation| operation_graph_inputs(operation, graph)) + .chain(required_outputs.iter().copied()) + .chain( + constraints + .required_equal_formats + .iter() + .flat_map(|pair| [pair.0, pair.1]), + ) + .chain(constraints.allocation_copies.keys().copied()) + .collect::>(); + let initial = relevant_origins + .iter() + .filter_map(|origin| values.get(origin).copied()) + .collect::>(); + let mut value_uses = BTreeMap::::new(); + for value in source + .iter() + .flat_map(|operation| operation_graph_inputs(operation, graph)) + .chain(required_outputs.iter().copied()) + { + *value_uses.entry(value).or_default() += 1; + } + let mut beam = vec![BeamBranch { + values: values.clone(), + state: state.clone(), + operations: Vec::new(), + peak_memory: MemoryPeaks::default(), + }]; + for (operation_index, operation) in source.iter().enumerate() { + let distributed_result_is_useful = operation.results.first().is_some_and(|result| { + required_outputs.contains(result) + || (value_uses.get(result).copied() == Some(1) + && source[operation_index + 1..] + .iter() + .find(|consumer| consumer.inputs.contains(result)) + .is_some_and(|consumer| { + consumer.inputs.iter().enumerate().any(|(index, input)| { + input == result + && operator_accepts_input_layout(&consumer.kind, index, config) + }) + })) + }); + let mut expanded = Vec::new(); + let mut rejected_memory = Vec::new(); + let mut saw_legal_plan = false; + let format_sensitive_indices = operation + .inputs + .iter() + .enumerate() + .filter_map(|(index, _)| { + operator_accepts_input_layout(&operation.kind, index, config).then_some(index) + }) + .collect::>(); + let mut plan_cache = BTreeMap::>::new(); + let mut plan_cache_hits = 0usize; + let mut generated_plan_sets = 0usize; + for branch in beam { + if let OperationKind::Repeat(repeat) = &operation.kind { + saw_legal_plan = true; + let mut next = branch.clone(); + lower_repeat( + operation, + repeat, + &mut next.values, + graph, + config, + costs, + &mut next.state, + &mut next.operations, + )?; + let peak = beam_memory_peak( + &next, + &initial, + source, + operation_index, + required_outputs, + graph, + &constraints.allocation_copies, + config.target, + ); + if peak.fits_with_budget( + config.target.memory_constraints(), + config.standard_memory_reservation_bytes, + config.tile_memory_budget_bytes, + ) { + next.peak_memory = peak; + expanded.push(next); + } else { + tracing::trace!( + operation = operation.id.index(), + standard = peak.standard, + interleaved = peak.interleaved, + total = peak.total, + contiguous_overflow = peak.standard_contiguous_overflow_with_reservation( + config.target.memory_constraints(), + config.standard_memory_reservation_bytes, + ), + plan = ?next.operations.last().and_then(|operation| operation.operator_plan()), + "rejected planning branch for memory" + ); + rejected_memory.push(peak); + } + continue; + } + let input_ids = operation + .inputs + .iter() + .map(|value| lookup(&branch.values, *value)) + .collect::>>()?; + let input_types = input_ids + .iter() + .map(|value| branch.state.get(*value).tensor_type.clone()) + .collect::>(); + if matches!(operation.kind, OperationKind::Gemm(_)) + && input_types.get(1).is_some_and(|right| { + right.shape.0[..right.shape.0.len().saturating_sub(2)] + .iter() + .any(|&extent| extent != 1) + }) + { + return Err(LoweringError::UnsupportedGemmBatching(operation.id)); + } + let output_shape = shapes + .get(&operation.results[0]) + .cloned() + .ok_or(LoweringError::MissingShape(operation.results[0]))?; + let grouped_output = grouped_output_layout( + source, + operation_index, + operation, + &output_shape, + &value_uses, + ); + let direct_consumer_layouts = direct_consumer_layouts( + source, + operation_index, + operation.results[0], + &output_shape, + config, + ); + if matches!(operation.kind, OperationKind::SplitHeads(_)) { + let [source_value] = input_ids.as_slice() else { + continue; + }; + let Some(&rows) = output_shape.0.get(1) else { + continue; + }; + let Ok(streams) = u16::try_from(output_shape.0[0]) else { + continue; + }; + if streams == 0 { + continue; + } + let query_partitions = rows.min(u32::from(config.tile_count / streams)); + let key_partitions = rows + .div_ceil(AMP_INNER_BLOCK) + .min(u32::from(config.tile_count / streams)); + let layouts = if direct_consumer_layouts.is_empty() { + [ + (query_partitions != 0).then(|| { + Layout::attention_query( + streams, + u16::try_from(query_partitions).unwrap_or(u16::MAX), + ) + }), + (key_partitions != 0).then(|| { + Layout::attention_key( + streams, + u16::try_from(key_partitions).unwrap_or(u16::MAX), + ) + }), + (key_partitions != 0).then(|| { + Layout::attention_block_major_key_value( + streams, + u16::try_from(key_partitions).unwrap_or(u16::MAX), + ) + }), + ] + .into_iter() + .flatten() + .collect::>() + } else { + direct_consumer_layouts + }; + saw_legal_plan |= !layouts.is_empty(); + let evaluated = layouts + .into_par_iter() + .filter_map(|layout| { + let mut next = branch.clone(); + apply_selected_view( + operation, + output_shape.clone(), + layout, + *source_value, + costs, + &mut next, + )?; + let peak = beam_memory_peak( + &next, + &initial, + source, + operation_index, + required_outputs, + graph, + &constraints.allocation_copies, + config.target, + ); + Some((next, peak)) + }) + .collect::>(); + for (mut next, peak) in evaluated { + if peak.fits_with_budget( + config.target.memory_constraints(), + config.standard_memory_reservation_bytes, + config.tile_memory_budget_bytes, + ) { + next.peak_memory = peak; + expanded.push(next); + } else { + rejected_memory.push(peak); + } + } + continue; + } + let parameter_inputs = input_ids + .iter() + .map(|id| branch.state.parameter_values.contains(id)) + .collect::>(); + let cache_key = PlanCacheKey { + input_shapes: input_types + .iter() + .map(|input| input.shape.clone()) + .collect(), + parameter_inputs: parameter_inputs.clone(), + format_sensitive_inputs: format_sensitive_indices + .iter() + .filter_map(|&index| { + input_types + .get(index) + .map(|input| (index, input.format.clone())) + }) + .collect(), + }; + let cached = if let Some(cached) = plan_cache.get(&cache_key) { + plan_cache_hits += 1; + cached + } else { + generated_plan_sets += 1; + let generated = plans_for_operation( + operation, + &input_types, + ¶meter_inputs, + &output_shape, + config, + costs, + distributed_result_is_useful, + grouped_output, + ); + plan_cache.entry(cache_key).or_insert(generated) + }; + let available_plans = cached + .iter() + .cloned() + .into_iter() + .filter(|plan| { + input_ids + .iter() + .zip(&plan.requirements.inputs) + .all(|(id, requirement)| { + let current = &branch.state.get(*id).tensor_type.format.layout; + branch.state.automatic_inputs.contains(id) + || current.order == requirement.format.layout.order + || !requirement.format.layout.order.requires_direct_population() + || (current.order == StorageOrder::Linear + && requirement + .format + .layout + .order + .supports_row_major_population()) + }) + }) + .collect::>(); + let available_plans = available_plans + .into_iter() + .flat_map(|plan| { + let mut complete = plan.clone(); + for requirement in &mut complete.requirements.inputs { + requirement.materialization = OperandMaterialization::Complete; + } + match config.conversion_streaming { + ConversionStreamingPolicy::Never => vec![complete], + ConversionStreamingPolicy::Always => vec![plan], + ConversionStreamingPolicy::WhenRequired if complete == plan => { + vec![complete] + } + ConversionStreamingPolicy::WhenRequired => vec![complete, plan], + } + }) + .collect::>(); + saw_legal_plan |= !available_plans.is_empty(); + let evaluated = available_plans + .into_par_iter() + .map(|plan| { + let mut next = branch.clone(); + apply_selected_plan( + operation, + output_shape.clone(), + plan, + &operation + .inputs + .iter() + .map(|value| value_uses.get(value).copied().unwrap_or(0) == 1) + .collect::>(), + costs, + &mut next.values, + &mut next.state, + &mut next.operations, + ); + let peak = beam_memory_peak( + &next, + &initial, + source, + operation_index, + required_outputs, + graph, + &constraints.allocation_copies, + config.target, + ); + (next, peak) + }) + .collect::>(); + for (mut next, peak) in evaluated { + let fits = peak.fits_with_budget( + config.target.memory_constraints(), + config.standard_memory_reservation_bytes, + config.tile_memory_budget_bytes, + ); + if fits || branch_contains_gemm_constraint(&next, config) { + if !fits { + tracing::debug!( + operation = operation.id.index(), + standard = peak.standard, + interleaved = peak.interleaved, + total = peak.total, + "retained constrained GEMM past conservative memory estimate" + ); + } + next.peak_memory = peak; + expanded.push(next); + } else { + tracing::trace!( + operation = operation.id.index(), + standard = peak.standard, + interleaved = peak.interleaved, + total = peak.total, + contiguous_overflow = peak.standard_contiguous_overflow_with_reservation( + config.target.memory_constraints(), + config.standard_memory_reservation_bytes, + ), + plan = ?next.operations.last().and_then(|operation| operation.operator_plan()), + "rejected planning branch for memory" + ); + rejected_memory.push(peak); + } + } + } + if expanded.is_empty() { + if saw_legal_plan + && let Some(peak) = rejected_memory + .into_iter() + .min_by_key(|peak| (peak.total, peak.interleaved, peak.standard)) + { + return Err(LoweringError::InsufficientMemory { + operation: operation.id, + standard: peak.standard, + standard_reservation: config.standard_memory_reservation_bytes, + interleaved: peak.interleaved, + total: peak.total, + standard_contiguous_overflow: peak + .standard_contiguous_overflow_with_reservation( + config.target.memory_constraints(), + config.standard_memory_reservation_bytes, + ), + }); + } + return Err(LoweringError::NoCandidate(operation.id)); + } + let future_origins = source[operation_index + 1..] + .iter() + .flat_map(|operation| operation_graph_inputs(operation, graph)) + .chain(required_outputs.iter().copied()) + .chain( + constraints + .required_equal_formats + .iter() + .flat_map(|pair| [pair.0, pair.1]), + ) + .chain(constraints.allocation_copies.keys().copied()) + .collect::>(); + let expanded_count = expanded.len(); + let (expanded, dominated, equivalent, diversity) = retain_pareto_beam( + expanded, + &future_origins, + constraints, + costs, + config.planning_beam_width.max(1), + ); + tracing::debug!( + operation = operation.id.index(), + retained = expanded.len(), + expanded = expanded_count, + pareto_dominated = dominated, + equivalent, + diversity_representatives = diversity, + best_cycles = branch_cycles(&expanded[0]), + generated_plan_sets, + plan_cache_hits, + "retained planning beam" + ); + beam = expanded; + } + let final_operation = source.len().saturating_sub(1); + let beam = beam + .into_iter() + .filter_map(|mut branch| { + let peak = beam_memory_peak( + &branch, + &initial, + source, + final_operation, + required_outputs, + graph, + &constraints.allocation_copies, + config.target, + ); + (peak.fits_with_budget( + config.target.memory_constraints(), + config.standard_memory_reservation_bytes, + config.tile_memory_budget_bytes, + ) || branch_contains_gemm_constraint(&branch, config)) + .then(|| { + branch.peak_memory = peak; + branch + }) + }) + .collect::>(); + let mut beam = beam; + beam.sort_by_key(|branch| { + branch_cycles(branch).saturating_add(format_equality_cost( + branch, + &constraints.required_equal_formats, + costs, + )) + }); + if beam.is_empty() { + return Err(LoweringError::NoCandidate(source[0].id)); + } + Ok(beam) + } + + fn branch_contains_gemm_constraint(branch: &BeamBranch, config: &PipelineConfig) -> bool { + branch.operations.iter().any(|operation| { + operation.source.is_some_and(|source| { + config + .search_domain + .gemm_plan_constraints + .iter() + .any(|constraint| constraint.source_operation == source.index()) + }) + }) + } + + #[allow(clippy::too_many_arguments)] + fn lower_operations( + source: &[Operation], + required_outputs: &[ValueId], + values: &mut BTreeMap, + shapes: &BTreeMap, + graph: &ComputeGraph, + config: &PipelineConfig, + costs: &impl CostModel, + state: &mut LoweringState, + constraints: &RegionPlanningConstraints, + ) -> LoweringResult> { + let mut branches = plan_region_frontier( + source, + required_outputs, + values, + shapes, + graph, + config, + costs, + state, + constraints, + )?; + let best = branches.remove(0); + *values = best.values; + *state = best.state; + Ok(best.operations) + } + + fn retain_pareto_beam( + branches: Vec, + future_origins: &BTreeSet, + constraints: &RegionPlanningConstraints, + costs: &impl CostModel, + width: usize, + ) -> (Vec, usize, usize, usize) { + let mut groups = BTreeMap::>::new(); + for (order, branch) in branches.into_iter().enumerate() { + let signature = future_beam_state(&branch, future_origins, constraints); + let objective = RegionMetrics { + cost: CostEstimate { + cycles: branch_cycles(&branch).saturating_add(format_equality_cost( + &branch, + &constraints.required_equal_formats, + costs, + )), + ..CostEstimate::default() + }, + memory: branch.peak_memory, + }; + groups.entry(signature).or_default().push(RankedBeamBranch { + compatibility: future_format_compatibility(&branch, future_origins), + branch, + objective, + order, + }); + } + + let mut frontier = Vec::new(); + let mut dominated = 0usize; + let mut equivalent = 0usize; + for (_, candidates) in groups { + let mut group_frontier = Vec::::new(); + for candidate in candidates { + if group_frontier + .iter() + .any(|kept| kept.objective == candidate.objective) + { + equivalent += 1; + continue; + } + if group_frontier + .iter() + .any(|kept| kept.objective.dominates(candidate.objective)) + { + dominated += 1; + continue; + } + let before = group_frontier.len(); + group_frontier.retain(|kept| !candidate.objective.dominates(kept.objective)); + dominated += before - group_frontier.len(); + group_frontier.push(candidate); + } + frontier.extend(group_frontier); + } + frontier.sort_by_key(|candidate| (candidate.objective.cost.cycles, candidate.order)); + + let mut selected = BTreeSet::new(); + let mut diversity = 0usize; + if frontier.len() > width { + // Preserve the cheapest representative of every live format family + // before retaining secondary memory tradeoffs. Partition counts are + // intentionally excluded: they are searched within a family, whereas + // physical order and ownership axes determine which imminent + // consumers can use a value without a qualitatively different + // conversion. + let mut represented = BTreeSet::new(); + for (index, entry) in frontier.iter().enumerate() { + if selected.len() == width { + break; + } + if represented.insert(entry.compatibility.clone()) && selected.insert(index) { + diversity += 1; + } + } + let objectives: [fn(&RankedBeamBranch) -> u64; 6] = [ + |entry: &RankedBeamBranch| entry.objective.memory.standard, + |entry: &RankedBeamBranch| entry.objective.memory.interleaved, + |entry: &RankedBeamBranch| entry.objective.memory.total, + |entry: &RankedBeamBranch| entry.objective.memory.maximum_standard_allocation, + |entry: &RankedBeamBranch| entry.objective.memory.standard_contiguous_overflow, + |entry: &RankedBeamBranch| entry.objective.memory.exchange_rows, + ]; + selected.insert(0); + for objective in objectives { + if selected.len() == width { + break; + } + let index = frontier + .iter() + .enumerate() + .min_by_key(|(index, entry)| { + (objective(entry), entry.objective.cost.cycles, *index) + }) + .map(|(index, _)| index) + .unwrap(); + if selected.insert(index) { + diversity += 1; + } + } + for index in 0..frontier.len() { + if selected.len() == width { + break; + } + selected.insert(index); + } + } else { + selected.extend(0..frontier.len()); + } + let mut retained = frontier + .into_iter() + .enumerate() + .filter_map(|(index, entry)| selected.contains(&index).then_some(entry.branch)) + .collect::>(); + retained.sort_by_cached_key(|branch| { + branch_cycles(branch).saturating_add(format_equality_cost( + branch, + &constraints.required_equal_formats, + costs, + )) + }); + (retained, dominated, equivalent, diversity) + } + + fn future_beam_state( + branch: &BeamBranch, + future_origins: &BTreeSet, + constraints: &RegionPlanningConstraints, + ) -> FutureBeamState { + let mut storage_classes = BTreeMap::::new(); + let mut next_storage_class = 0u32; + let mut storage_class = |id: MidValueId| { + let group = branch.state.get(id).storage_group; + *storage_classes.entry(group).or_insert_with(|| { + let class = next_storage_class; + next_storage_class += 1; + class + }) + }; + let mut values = Vec::new(); + for &origin in future_origins { + let Some(&id) = branch.values.get(&origin) else { + continue; + }; + values.push(FutureValueState { + origin, + tensor_type: branch.state.get(id).tensor_type.clone(), + automatic_input: branch.state.automatic_inputs.contains(&id), + parameter: branch.state.parameter_values.contains(&id), + allocation_copies: constraints + .allocation_copies + .get(&origin) + .copied() + .unwrap_or(1), + storage_class: storage_class(id), + }); + } + let equal_formats_satisfied = constraints + .required_equal_formats + .iter() + .map(|&(left, right)| { + let satisfied = branch + .values + .get(&left) + .zip(branch.values.get(&right)) + .is_some_and(|(&left, &right)| { + branch.state.get(left).tensor_type.format + == branch.state.get(right).tensor_type.format + }); + (left, right, satisfied) + }) + .collect(); + FutureBeamState { + values, + equal_formats_satisfied, + } + } + + fn branch_cycles(branch: &BeamBranch) -> u64 { + branch + .operations + .iter() + .map(|operation| operation.metrics.cost.cycles) + .sum() + } + + fn format_equality_cost( + branch: &BeamBranch, + equalities: &[(ValueId, ValueId)], + costs: &impl CostModel, + ) -> u64 { + equalities.iter().fold(0u64, |total, &(source, target)| { + let Some((&source, &target)) = branch.values.get(&source).zip(branch.values.get(&target)) + else { + return total; + }; + let source = &branch.state.get(source).tensor_type; + let target = &branch.state.get(target).tensor_type; + let cast = (source.format.precision != target.format.precision) + .then(|| costs.cast_cycles(source, target.format.precision)) + .unwrap_or(0); + let rearrange = (source.format.layout != target.format.layout) + .then(|| { + costs + .layout_conversion_cost( + &source.shape, + target.format.precision, + &source.format.layout, + &target.format.layout, + ) + .cycles + }) + .unwrap_or(0); + total.saturating_add(cast).saturating_add(rearrange) + }) + } + + fn operation_graph_inputs(operation: &Operation, graph: &ComputeGraph) -> Vec { + let mut inputs = operation.inputs.clone(); + if let OperationKind::Repeat(repeat) = &operation.kind { + for sequence in &repeat.iterated_inputs { + inputs.extend(&graph.sequences()[sequence.index() as usize].values); + } + } + inputs + } + + fn apply_selected_plan( + operation: &Operation, + output_shape: TensorShape, + plan: OperatorSchedule, + single_use_inputs: &[bool], + costs: &impl CostModel, + values: &mut BTreeMap, + state: &mut LoweringState, + operations: &mut Vec, + ) { + let input_ids = operation + .inputs + .iter() + .map(|value| values[value]) + .collect::>(); + let original_input_ids = input_ids.clone(); + let mut source_types = Vec::with_capacity(input_ids.len()); + let mut converted = Vec::with_capacity(input_ids.len()); + for (value, requirement) in input_ids.into_iter().zip(&plan.requirements.inputs) { + let conversion_start = operations.len(); + let converted_value = ensure_format( + value, + requirement.format.clone(), + requirement.materialization, + operation.id, + costs, + state, + operations, + ); + let streamed_source = operations[conversion_start..] + .last_mut() + .and_then(|conversion| { + let streamed = conversion.conversion().is_some_and(|(_, materialization)| { + materialization == OperandMaterialization::DispatchSlices + }); + if streamed { + conversion.metrics.cost.cycles = 0; + conversion.metrics.cost.exchange_cycles = 0; + conversion.inputs.first().copied() + } else { + None + } + }); + let source_value = streamed_source.unwrap_or(converted_value); + source_types.push(state.get(source_value).tensor_type.clone()); + converted.push(converted_value); + } + let result = state.value( + operation.results[0], + TensorType { + shape: output_shape, + format: plan.requirements.output.format.clone(), + }, + ); + let converted_types = converted + .iter() + .map(|value| state.get(*value).tensor_type.clone()) + .collect::>(); + let mut operator_cycles = costs.operator_transition_cycles( + &plan, + &source_types, + &converted_types, + &state.get(result).tensor_type, + ); + let mut operator_exchange_cycles = costs.operator_transition_exchange_cycles( + &plan, + &source_types, + &converted_types, + &state.get(result).tensor_type, + ); + for (input_index, ((&original, &converted), requirement)) in original_input_ids + .iter() + .zip(&converted) + .zip(&plan.requirements.inputs) + .enumerate() + { + let conversion_is_streamed = original == converted + || operations.iter().any(|candidate| { + candidate.inputs.as_slice() == [original] + && candidate.results.as_slice() == [converted] + && candidate.conversion().is_some_and(|(_, materialization)| { + materialization == OperandMaterialization::DispatchSlices + }) + }); + if !conversion_is_streamed + || !single_use_inputs.get(input_index).copied().unwrap_or(false) + || requirement.materialization != OperandMaterialization::DispatchSlices + { + continue; + } + let Some(producer_index) = operations + .iter() + .position(|candidate| candidate.results.as_slice() == [original]) + else { + continue; + }; + let (source, transform, producer_cycles) = match &operations[producer_index].kind { + MidOperationKind::Convert(Some(transform), ..) => ( + operations[producer_index].inputs[0], + *transform, + operations[producer_index].metrics.cost.cycles, + ), + _ => continue, + }; + let fused_cycles = costs.deferred_input_cycles( + transform, + &state.get(source).tensor_type, + &state.get(original).tensor_type, + &converted_types[input_index], + &plan, + producer_cycles, + ); + operator_cycles = operator_cycles.saturating_add(fused_cycles); + operator_exchange_cycles = + operator_exchange_cycles.saturating_add(costs.deferred_input_exchange_cycles( + transform, + &state.get(source).tensor_type, + &state.get(original).tensor_type, + &converted_types[input_index], + &plan, + producer_cycles, + )); + operations[producer_index].metrics.cost = CostEstimate::default(); + } + tracing::trace!( + source = operation.id.index(), + cycles = operator_cycles, + schedule = ?plan, + input_layouts = ?converted_types + .iter() + .map(|input| &input.format.layout) + .collect::>(), + output_layout = ?state.get(result).tensor_type.format.layout, + "costed operator plan" + ); + let exchange = + costs.operator_exchange_footprint(&plan, &converted_types, &state.get(result).tensor_type); + let memory = operator_memory_estimate(&plan, &converted_types, &state.get(result).tensor_type); + operations.push(MidOperation { + source: Some(operation.id), + inputs: converted, + results: vec![result], + kind: MidOperationKind::Operator(plan), + metrics: OperationMetrics { + cost: CostEstimate { + cycles: operator_cycles, + exchange_cycles: operator_exchange_cycles, + exchange_footprint: exchange, + }, + memory, + }, + }); + values.insert(operation.results[0], result); + } + + fn apply_selected_view( + operation: &Operation, + output_shape: TensorShape, + output_layout: Layout, + source: MidValueId, + costs: &impl CostModel, + branch: &mut BeamBranch, + ) -> Option<()> { + let transform = match operation.kind { + OperationKind::SplitHeads(options) => DeferredTransform::SplitLastAxisIntoLeading { + parts: options.heads, + }, + _ => return None, + }; + let source_type = &branch.state.get(source).tensor_type; + let output_type = TensorType { + shape: output_shape, + format: TensorFormat { + precision: source_type.format.precision, + layout: output_layout, + }, + }; + source_type.format.layout.resolve(&source_type.shape).ok()?; + output_type.format.layout.resolve(&output_type.shape).ok()?; + let strategy = layout_conversion_strategy( + source_type.format.precision, + &source_type.format.layout, + &output_type.format.layout, + ); + let cost = costs.rearrangement_cost(source_type, &output_type, strategy); + let memory = conversion_memory_estimate(source_type, &output_type, strategy); + let result = branch + .state + .value(operation.results[0], output_type.clone()); + branch.operations.push(MidOperation { + source: Some(operation.id), + inputs: vec![source], + results: vec![result], + kind: MidOperationKind::Convert( + Some(transform), + strategy, + OperandMaterialization::DispatchSlices, + Vec::new(), + ), + metrics: OperationMetrics { cost, memory }, + }); + branch.values.insert(operation.results[0], result); + Some(()) + } + + fn beam_memory_peak( + branch: &BeamBranch, + initial: &[MidValueId], + source: &[Operation], + operation_index: usize, + required_outputs: &[ValueId], + graph: &ComputeGraph, + allocation_multiplicity: &BTreeMap, + target: HardwareTarget, + ) -> MemoryPeaks { + let live_origins = source[operation_index + 1..] + .iter() + .flat_map(|operation| operation_graph_inputs(operation, graph)) + .chain(required_outputs.iter().copied()) + .collect::>(); + let live = live_origins + .iter() + .filter_map(|origin| branch.values.get(origin).copied()) + .collect::>(); + let multiplicity = branch + .state + .values + .iter() + .filter_map(|value| { + allocation_multiplicity + .get(&value.origin) + .map(|copies| (value.id, *copies)) + }) + .collect::>(); + region_peak_memory_with_multiplicity( + initial, + &branch.operations, + &live, + &branch.state.values, + &multiplicity, + target, + ) + } + + fn plan_fits_operator_memory( + plan: &OperatorSchedule, + inputs: &[TensorType], + output: &TensorShape, + config: &PipelineConfig, + ) -> bool { + let planned_inputs = inputs + .iter() + .zip(&plan.requirements.inputs) + .map(|(input, requirement)| TensorType { + shape: input.shape.clone(), + format: requirement.format.clone(), + }) + .collect::>(); + let planned_output = TensorType { + shape: output.clone(), + format: plan.requirements.output.format.clone(), + }; + let peak = operator_memory_estimate(plan, &planned_inputs, &planned_output).peak; + let constraints = config.target.memory_constraints(); + peak.interleaved <= constraints.interleaved_bytes + && peak + .total() + .saturating_add(config.standard_memory_reservation_bytes) + <= config.tile_memory_budget_bytes.min(constraints.total_bytes) + } + + fn plans_for_operation( + operation: &Operation, + inputs: &[TensorType], + parameter_inputs: &[bool], + output: &TensorShape, + config: &PipelineConfig, + costs: &impl CostModel, + distributed_result_is_useful: bool, + grouped_output: Option, + ) -> Vec { + let mut plans = Vec::new(); + let gemm_constraint = config + .search_domain + .gemm_plan_constraints + .iter() + .find(|constraint| constraint.source_operation == operation.id.index()); + if let OperationKind::FlashAttention(options) = operation.kind + && config + .search_domain + .permits_precision(OperatorClass::Attention, Precision::F16) + && !options.causal + && let [query, key, value] = inputs + && query.shape.0.len() == 3 + && key.shape.0.len() == 3 + && value.shape.0.len() == 3 + && query.shape.0[0] == key.shape.0[0] + && query.shape.0[0] == value.shape.0[0] + && let Ok(heads) = u16::try_from(query.shape.0[0]) + && heads != 0 + { + let query_rows = query.shape.0[1]; + let query_partitions = u16::try_from(query_rows) + .unwrap_or(u16::MAX) + .min(config.tile_count / heads); + if query_partitions != 0 { + let key_partitions = + u16::try_from(key.shape.0[1].div_ceil(AMP_INNER_BLOCK)).unwrap_or(u16::MAX); + if key_partitions == 0 || heads.saturating_mul(key_partitions) > config.tile_count { + return plans; + } + let padded_query_dimension = + query.shape.0[2].div_ceil(AMP_COLUMN_MICRO) * AMP_COLUMN_MICRO; + let padded_value_dimension = + value.shape.0[2].div_ceil(AMP_COLUMN_MICRO) * AMP_COLUMN_MICRO; + let padded_key_rows = key.shape.0[1].div_ceil(AMP_INNER_BLOCK) * AMP_INNER_BLOCK; + let query_format = TensorFormat { + precision: Precision::F16, + layout: Layout::attention_query(heads, query_partitions), + }; + let key_format = TensorFormat { + precision: Precision::F16, + layout: Layout::attention_key(heads, key_partitions), + }; + let value_format = TensorFormat { + precision: Precision::F16, + layout: Layout::attention_block_major_key_value(heads, key_partitions), + }; + let output_format = TensorFormat { + precision: Precision::F32, + layout: Layout::attention_output(heads, query_partitions), + }; + let operator = MidOperator::FlashAttention { + options, + accumulate: AccumulationPrecision::F32, + }; + let kernel = GemmKernelFamily { + multiply: Precision::F16, + accumulate: AccumulationPrecision::F32, + weights: GemmWeightLoad::Standard, + }; + let requirements = OperatorRequirements { + inputs: [query_format, key_format, value_format] + .into_iter() + .map(|format| { + OperandRequirement::new(format, 8) + .with_materialization(OperandMaterialization::DispatchSlices) + }) + .collect(), + output: OperandRequirement::new(output_format, 8), + output_aliasing: OutputAliasing::Fresh, + memory_space: MemorySpaceRequirements::default(), + }; + let query_rows = query_rows.div_ceil(u32::from(query_partitions)); + let blockings = [ + (config.search_domain.attention_strategy != AttentionStrategy::Materialized) + .then_some(AttentionBlocking::Flash { + query_rows, + key_rows: AMP_INNER_BLOCK, + }), + (config.search_domain.attention_strategy != AttentionStrategy::Flash).then_some( + AttentionBlocking::Materialized { + query_rows, + padded_key_rows, + }, + ), + ]; + plans.extend( + blockings + .into_iter() + .flatten() + .map(|blocking| OperatorSchedule { + operator, + steps: vec![ScheduleStep::Attention(AttentionMap { + inputs: [ + ScheduleValue::Input(0), + ScheduleValue::Input(1), + ScheduleValue::Input(2), + ], + output: ScheduleValue::Output, + kernel, + blocking, + query_dimension: padded_query_dimension, + value_dimension: padded_value_dimension, + })], + requirements: requirements.clone(), + }), + ); + } + } + match operation.kind { + OperationKind::Gemm(options) => plans.extend(gemm_plans( + options, + inputs, + parameter_inputs, + output, + config, + costs, + distributed_result_is_useful, + gemm_constraint, + grouped_output, + )), + OperationKind::Gelu => plans.extend(pointwise_plans( + MidOperator::Gelu, + OperatorClass::Gelu, + inputs, + output, + config, + )), + OperationKind::Add(options) => plans.extend(pointwise_plans( + MidOperator::Add(options), + OperatorClass::Add, + inputs, + output, + config, + )), + OperationKind::SplitHeads(_) + | OperationKind::FlashAttention(_) + | OperationKind::Repeat(_) => {} + } + plans.retain(|plan| { + plan.supports(inputs, output) && plan_fits_operator_memory(plan, inputs, output, config) + }); + if let Some(constraint) = gemm_constraint { + plans.retain(|plan| gemm_plan_matches(constraint, plan, &plan.requirements.inputs)); + tracing::info!( + source_operation = constraint.source_operation, + matching_plans = plans.len(), + ?constraint, + "applied GEMM plan constraint" + ); + } + plans + } + + fn pointwise_plans( + operator: MidOperator, + class: OperatorClass, + inputs: &[TensorType], + output: &TensorShape, + config: &PipelineConfig, + ) -> Vec { + let access = match operator { + MidOperator::Gelu => ScheduleAccess::TileLocal, + MidOperator::Add(_) => ScheduleAccess::LogicalOverlap, + _ => return Vec::new(), + }; + let mut plans = Vec::new(); + for (anchor, input) in inputs.iter().enumerate() { + if input.shape != *output + || !config + .search_domain + .permits_precision(class, input.format.precision) + || (matches!(operator, MidOperator::Gelu) && anchor != 0) + { + continue; + } + let mut formats = Vec::new(); + if config.conversion_streaming != ConversionStreamingPolicy::Always { + formats.extend(pointwise_flat_formats(input, output, config)); + } + formats.push(input.format.clone()); + for format in formats { + let aliasing = OutputAliasing::MayAliasInputs( + inputs + .iter() + .enumerate() + .filter_map(|(index, input)| (input.shape == *output).then_some(index as u16)) + .collect(), + ); + let plan = OperatorSchedule { + operator, + steps: vec![ScheduleStep::KernelMap(KernelMap { + kernel: match operator { + MidOperator::Gelu => TileKernelSpec::Gelu, + MidOperator::Add(_) => TileKernelSpec::Add, + _ => unreachable!(), + }, + inputs: (0..inputs.len()) + .map(|index| (ScheduleValue::Input(index as u16), access)) + .collect(), + output: ScheduleValue::Output, + })], + requirements: OperatorRequirements { + inputs: (0..inputs.len()) + .map(|_| OperandRequirement::new(format.clone(), 8)) + .collect(), + output: OperandRequirement::new(format, 8), + output_aliasing: aliasing, + memory_space: MemorySpaceRequirements::default(), + }, + }; + if !plans.contains(&plan) { + plans.push(plan); + } + } + } + plans + } + + fn pointwise_flat_formats( + input: &TensorType, + output: &TensorShape, + config: &PipelineConfig, + ) -> Vec { + let grain = 8_u32.div_ceil(input.format.precision.bytes() as u32); + let Some(&width) = output.0.last() else { + return Vec::new(); + }; + if grain == 0 || !output.elements().is_multiple_of(u64::from(grain)) { + return Vec::new(); + } + let width = u64::from(width); + let grains = output.elements() / u64::from(grain); + let mut candidates = BTreeMap::new(); + for &tiles in &config.resolved_active_tile_counts { + let splits = (1..tiles) + .filter(|&tile| { + let tile = u64::from(tile); + let offset = (tile * (grains / u64::from(tiles)) + + tile.min(grains % u64::from(tiles))) + * u64::from(grain); + !offset.is_multiple_of(width) + }) + .count(); + candidates.entry(tiles).or_insert(splits); + } + candidates + .iter() + .filter(|(tiles, splits)| { + !candidates.iter().any(|(other_tiles, other_splits)| { + other_tiles >= tiles + && other_splits <= splits + && (other_tiles > tiles || other_splits < splits) + }) + }) + .flat_map(|(&tiles, _)| { + let mut layouts = vec![Layout::logical_linear(tiles, grain)]; + if let Some(retained_grain) = input + .format + .layout + .order + .retained_linear_column_grain(input.format.precision) + .filter(|retained_grain| retained_grain.is_multiple_of(grain)) + .filter(|retained_grain| { + output.elements().is_multiple_of(u64::from(*retained_grain)) + && output + .0 + .last() + .is_some_and(|width| width.is_multiple_of(*retained_grain)) + }) + { + layouts.push( + input + .format + .layout + .with_retained_order_linear_ownership(tiles, retained_grain), + ); + } + layouts.into_iter().map(|layout| TensorFormat { + precision: input.format.precision, + layout, + }) + }) + .collect() + } + + #[allow(clippy::too_many_arguments)] + #[allow(clippy::too_many_arguments)] + fn lower_repeat( + operation: &Operation, + repeat: &Repeat, + values: &mut BTreeMap, + graph: &ComputeGraph, + config: &PipelineConfig, + costs: &impl CostModel, + state: &mut LoweringState, + operations: &mut Vec, + ) -> LoweringResult<()> { + let inputs = operation + .inputs + .iter() + .map(|value| lookup(values, *value)) + .collect::>>()?; + let mut argument_types = inputs + .iter() + .map(|value| state.get(*value).tensor_type.clone()) + .collect::>(); + let mut raw_iterated_inputs = Vec::with_capacity(repeat.iterated_inputs.len()); + let mut iterated_parameters = Vec::with_capacity(repeat.iterated_inputs.len()); + let mut iterated_automatic = Vec::with_capacity(repeat.iterated_inputs.len()); + for sequence_id in &repeat.iterated_inputs { + let sequence = &graph.sequences()[sequence_id.index() as usize]; + let first = lookup(values, sequence.values[0])?; + let first_type = state.get(first).tensor_type.clone(); + let sequence_values: Vec<_> = sequence + .values + .iter() + .map(|value| lookup(values, *value)) + .collect::>>()?; + let storage_group = state.get(first).storage_group; + for &value in &sequence_values { + state.values[value.index() as usize].storage_group = storage_group; + } + iterated_parameters.push( + sequence_values + .iter() + .all(|value| state.parameter_values.contains(value)), + ); + iterated_automatic.push( + sequence_values + .iter() + .all(|value| state.automatic_inputs.contains(value)), + ); + raw_iterated_inputs.push(sequence_values); + argument_types.push(first_type); + } + let mut body_values = BTreeMap::new(); + let mut arguments = Vec::new(); + for (argument_index, (&origin, tensor_type)) in + repeat.body.arguments.iter().zip(argument_types).enumerate() + { + let storage_group = if argument_index < inputs.len() { + state.get(inputs[argument_index]).storage_group + } else { + state + .get(raw_iterated_inputs[argument_index - inputs.len()][0]) + .storage_group + }; + let value = state.value_in_storage_group(origin, tensor_type, storage_group); + if argument_index < inputs.len() { + if state.automatic_inputs.contains(&inputs[argument_index]) { + state.automatic_inputs.insert(value); + } + if state.parameter_values.contains(&inputs[argument_index]) { + state.parameter_values.insert(value); + } + } else { + let iterated_index = argument_index - inputs.len(); + if iterated_automatic[iterated_index] { + state.automatic_inputs.insert(value); + } + if iterated_parameters[iterated_index] { + state.parameter_values.insert(value); + } + } + body_values.insert(origin, value); + arguments.push(value); + } + let body_allocation_copies = repeat + .body + .arguments + .iter() + .skip(inputs.len()) + .copied() + .map(|argument| (argument, repeat.count)) + .collect::>(); + let required_equal_formats = repeat + .body + .yields + .iter() + .copied() + .zip(repeat.body.arguments.iter().copied()) + .take(repeat.carried_inputs) + .collect(); + let body_constraints = RegionPlanningConstraints { + allocation_copies: body_allocation_copies, + required_equal_formats, + }; + let mut body_operations = lower_operations( + &repeat.body.operations, + &repeat.body.yields, + &mut body_values, + &repeat.body.value_shapes, + graph, + config, + costs, + state, + &body_constraints, + )?; + for index in 0..repeat.carried_inputs { + let body_layout = state + .get(arguments[index]) + .tensor_type + .format + .layout + .clone(); + state.retarget_automatic_input(inputs[index], body_layout); + } + let iterated_inputs = raw_iterated_inputs + .into_iter() + .enumerate() + .map(|(index, sequence)| { + let target = state + .get(arguments[inputs.len() + index]) + .tensor_type + .format + .clone(); + sequence + .into_iter() + .map(|value| { + ensure_format( + value, + target.clone(), + OperandMaterialization::Complete, + operation.id, + costs, + state, + operations, + ) + }) + .collect::>() + }) + .collect::>(); + let mut yields = Vec::new(); + for (index, high_yield) in repeat.body.yields.iter().enumerate() { + let value = lookup(&body_values, *high_yield)?; + let target = state.get(inputs[index]).tensor_type.format.clone(); + yields.push(ensure_format( + value, + target, + OperandMaterialization::Complete, + operation.id, + costs, + state, + &mut body_operations, + )); + } + let body_metrics = body_operations + .iter() + .fold(CostEstimate::default(), |cost, operation| { + cost.sequence(operation.metrics.cost) + }); + let body_allocation_multiplicity = arguments + .iter() + .skip(inputs.len()) + .copied() + .map(|argument| (argument, repeat.count)) + .collect::>(); + let body_peak = region_peak_memory_with_multiplicity( + &arguments, + &body_operations, + &yields, + &state.values, + &body_allocation_multiplicity, + config.target, + ); + let mut results = Vec::new(); + for (origin, input) in operation.results.iter().zip(&inputs) { + let tensor_type = state.get(*input).tensor_type.clone(); + let storage_group = state.get(*input).storage_group; + let result = state.value_in_storage_group(*origin, tensor_type, storage_group); + values.insert(*origin, result); + results.push(result); + } + operations.push(MidOperation { + source: Some(operation.id), + inputs, + results, + kind: MidOperationKind::Repeat(MidRepeat { + count: repeat.count, + carried_inputs: repeat.carried_inputs, + invariant_inputs: repeat.invariant_inputs, + iterated_inputs, + body: MidRegion { + arguments, + operations: body_operations, + yields, + metrics: RegionMetrics { + cost: body_metrics, + memory: body_peak, + }, + }, + }), + metrics: OperationMetrics { + cost: body_metrics.repeated(repeat.count), + memory: MemoryEstimate { + live: body_peak.conservative_tensor_usage(), + temporary: MemoryUsage::default(), + peak: body_peak.conservative_tensor_usage(), + maximum_standard_temporary_allocation: 0, + }, + }, + }); + Ok(()) + } + + fn ensure_format( + mut value: MidValueId, + target: TensorFormat, + materialization: OperandMaterialization, + source: OperationId, + costs: &impl CostModel, + state: &mut LoweringState, + operations: &mut Vec, + ) -> MidValueId { + if state.retarget_automatic_input(value, target.layout.clone()) + && state.get(value).tensor_type.format.precision == target.precision + { + return value; + } + let original = state.get(value).clone(); + if original.tensor_type.format.precision != target.precision { + let mut tensor_type = original.tensor_type.clone(); + tensor_type.format.precision = target.precision; + let result = state.derived_value(value, tensor_type.clone()); + let memory = conversion_memory_estimate( + &original.tensor_type, + &tensor_type, + ConversionStrategy::LocalKernel, + ); + operations.push(MidOperation { + source: Some(source), + inputs: vec![value], + results: vec![result], + kind: MidOperationKind::CastPrecision, + metrics: OperationMetrics { + cost: CostEstimate { + cycles: costs.cast_cycles(&original.tensor_type, target.precision), + ..CostEstimate::default() + }, + memory, + }, + }); + value = result; + } + let current = state.get(value).clone(); + if current.tensor_type.format.layout != target.layout { + let mut tensor_type = current.tensor_type.clone(); + let from = tensor_type.format.layout.clone(); + tensor_type.format.layout = target.layout.clone(); + let result = state.derived_value(value, tensor_type.clone()); + let strategy = + layout_conversion_strategy(tensor_type.format.precision, &from, &target.layout); + let rearrangement = costs.rearrangement_cost(¤t.tensor_type, &tensor_type, strategy); + let memory = conversion_memory_estimate(¤t.tensor_type, &tensor_type, strategy); + operations.push(MidOperation { + source: Some(source), + inputs: vec![value], + results: vec![result], + kind: MidOperationKind::Convert(None, strategy, materialization, Vec::new()), + metrics: OperationMetrics { + cost: rearrangement, + memory, + }, + }); + value = result; + } + value + } + + fn lookup(values: &BTreeMap, value: ValueId) -> LoweringResult { + values + .get(&value) + .copied() + .ok_or(LoweringError::UnknownValue(value)) + } + + #[cfg(test)] + mod tests { + use super::*; + use crate::graph::{AddOptions, AttentionOptions, GemmOptions}; + use crate::{AMP_OUTPUT_COLUMN_BLOCK, AxisTiling, LayoutError, Padding, TensorTiling}; + use ipu_target::hardware::HardwareTarget; + + const RANDOM_CASES: usize = 128; + + #[test] + fn randomized_memory_peaks_reserve_disjoint_class_arenas() { + let mut random = fastrand::Rng::with_seed(0x636c_6173_735f_7372); + let constraints = HardwareTarget::Ipu21.memory_constraints(); + let capacity = constraints.total_bytes; + let interleaved_capacity = constraints.interleaved_bytes; + let element = u64::from(ipu_target::memory::IPU21_INTERLEAVED_ELEMENT_SIZE); + let mut rejected_noncoincident_peaks = 0; + for _ in 0..RANDOM_CASES * 16 { + let standard = random.u64(0..=capacity); + let interleaved = random.u64(0..=interleaved_capacity); + let reservation = random.u64(0..=capacity / 4); + let simultaneous = + random.u64(standard.max(interleaved)..=standard.saturating_add(interleaved)); + let peaks = MemoryPeaks { + standard, + interleaved, + total: simultaneous, + maximum_standard_allocation: 0, + ..MemoryPeaks::default() + }; + let aligned_interleaved = interleaved.div_ceil(element) * element; + let static_partition = standard + .saturating_add(aligned_interleaved) + .saturating_add(reservation); + let fits = peaks.fits_with_budget( + HardwareTarget::Ipu21.memory_constraints(), + reservation, + capacity, + ); + assert_eq!(fits, static_partition <= capacity); + if simultaneous.saturating_add(reservation) <= capacity && static_partition > capacity { + rejected_noncoincident_peaks += 1; + assert!(!fits); + } + } + assert!(rejected_noncoincident_peaks > 0); + } + + fn dimension(random: &mut fastrand::Rng) -> u32 { + random.u32(1..=128) + } + + fn small_dimension(random: &mut fastrand::Rng) -> u32 { + random.u32(1..=4) + } + + fn precision(random: &mut fastrand::Rng) -> Precision { + if random.bool() { + Precision::F16 + } else { + Precision::F32 + } + } + + fn format(precision: Precision, layout: Layout) -> TensorFormat { + TensorFormat { precision, layout } + } + + fn random_format(random: &mut fastrand::Rng, tiles: u16) -> TensorFormat { + let tiling = if random.bool() { + TensorTiling::replicated(tiles) + } else { + TensorTiling::sharded(TensorAxis::FromEnd(2), tiles) + }; + let mut layout = Layout::row_major(tiling); + if random.bool() { + layout.memory_class = MemoryClass::Interleaved; + } + format(precision(random), layout) + } + + #[test] + fn randomized_future_state_is_id_independent_but_preserves_aliasing() { + let mut random = fastrand::Rng::with_seed(0x616c_6961_7365_7321); + for _ in 0..RANDOM_CASES { + let mut graph = ComputeGraph::new(); + let first = graph.host_input("first", [1]).unwrap(); + let second = graph.host_input("second", [1]).unwrap(); + let dummy = graph.host_input("dummy", [1]).unwrap(); + let tiles = random.u16(1..=64); + let tensor_type = TensorType { + shape: TensorShape::new([random.u32(1..=128)]), + format: random_format(&mut random, tiles), + }; + let aliases = random.bool(); + let automatic = random.bool(); + let parameter = random.bool(); + + let make_branch = |prepend_dummy: bool, aliases: bool| { + let mut state = LoweringState::default(); + if prepend_dummy { + state.value(dummy, tensor_type.clone()); + } + let first_id = state.value(first, tensor_type.clone()); + let second_id = if aliases { + state.value_in_storage_group(second, tensor_type.clone(), first_id) + } else { + state.value(second, tensor_type.clone()) + }; + if automatic { + state.automatic_inputs.extend([first_id, second_id]); + } + if parameter { + state.parameter_values.extend([first_id, second_id]); + } + BeamBranch { + values: [(first, first_id), (second, second_id)] + .into_iter() + .collect(), + state, + operations: Vec::new(), + peak_memory: MemoryPeaks::default(), + } + }; + let future = [first, second].into_iter().collect(); + let constraints = RegionPlanningConstraints { + allocation_copies: [(first, random.u32(1..=8))].into_iter().collect(), + required_equal_formats: vec![(first, second)], + }; + let baseline = future_beam_state(&make_branch(false, aliases), &future, &constraints); + let renumbered = future_beam_state(&make_branch(true, aliases), &future, &constraints); + let changed_aliasing = + future_beam_state(&make_branch(true, !aliases), &future, &constraints); + assert_eq!(baseline, renumbered); + assert_ne!(baseline, changed_aliasing); + } + } + + #[test] + fn randomized_active_tile_candidates_bound_idle_capacity() { + let mut random = fastrand::Rng::with_seed(0x7469_6c65); + for _ in 0..RANDOM_CASES { + let capacity = random.u16(1..=1472); + let counts = crate::config::candidate_active_tile_counts(capacity); + assert_eq!(counts[0], capacity); + assert!(counts.windows(2).all(|pair| pair[0] > pair[1])); + assert!(counts.iter().all(|&count| count <= capacity)); + assert!(counts[1..].iter().all(|count| count.is_power_of_two())); + assert_eq!(counts.last(), Some(&1)); + } + for exponent in 1..=10 { + let capacity = 1_u16 << exponent; + assert_eq!( + crate::config::candidate_active_tile_counts(capacity).len(), + exponent + 1 + ); + } + } + + #[test] + fn randomized_shape_aware_tile_candidates_follow_graph_extents() { + let mut random = fastrand::Rng::with_seed(0x7368_6170_655f_6772); + for case in 0..RANDOM_CASES { + let capacity = random.u16(16..=1472); + let extent = random.u16(2..=capacity); + let shape = TensorShape(vec![u32::from(extent), random.u32(1..=4096)]); + let counts = crate::config::shape_aware_active_tile_counts(capacity, [&shape]); + let expected = capacity / extent * extent; + if expected >= capacity.div_ceil(2) && expected < capacity { + assert!(counts.contains(&expected), "case {case}"); + } + assert!(counts.iter().all(|&count| { + count < capacity + && count >= capacity.div_ceil(2) + && shape.0.iter().any(|&axis| u32::from(count) % axis == 0) + })); + } + } + + #[test] + fn randomized_search_domains_filter_operator_plan_generation() { + let mut random = fastrand::Rng::with_seed(0x646f_6d61_696e_2121); + for case in 0..RANDOM_CASES { + let capacity = random.u16(2..=64); + let first = random.u16(1..=capacity); + let second = random.u16(1..=capacity); + let precisions = match random.u8(0..3) { + 0 => vec![Precision::F16], + 1 => vec![Precision::F32], + _ => vec![Precision::F16, Precision::F32], + }; + let memory_classes = if random.bool() { + vec![MemoryClass::Standard] + } else { + vec![MemoryClass::Interleaved] + }; + let domain = PlannerSearchDomain::default() + .with_active_tile_counts([first, second, first, 0, capacity.saturating_add(1)]) + .with_operator_precisions(OperatorClass::Gemm, precisions.clone()) + .with_operator_precisions(OperatorClass::Gelu, precisions.clone()) + .with_operator_precisions(OperatorClass::Add, precisions.clone()) + .with_weight_memory_classes(memory_classes.clone()); + let shape = TensorShape::new([random.u32(1..=64), random.u32(1..=64)]); + let active = domain.active_tile_counts(capacity, [&shape]); + let mut expected = vec![first, second]; + expected.dedup(); + assert_eq!(active, expected, "case {case}"); + + let seeds = active + .iter() + .flat_map(|&tiles| { + gemm_seed_plans_for_tile_count(GemmOptions::default(), tiles, &domain) + }) + .collect::>(); + for seed in seeds { + assert!( + precisions.contains(&seed.requirements.inputs[0].format.precision), + "case {case}" + ); + assert!( + memory_classes + .contains(&seed.requirements.inputs[1].format.layout.memory_class), + "case {case}" + ); + assert!(active.contains(&seed.requirements.output.format.layout.tiling.tile_count)); + } + + let selected_precision = precisions[0]; + let tensor = TensorType::new( + [u32::from(capacity), random.u32(1..=64)], + selected_precision, + Layout::row_sharded(capacity), + ); + let mut config = PipelineConfig::new(capacity).with_search_domain(domain); + config.resolved_active_tile_counts = active; + for (operator, class, arity) in [ + (MidOperator::Gelu, OperatorClass::Gelu, 1), + ( + MidOperator::Add(AddOptions::default()), + OperatorClass::Add, + 2, + ), + ] { + let inputs = vec![tensor.clone(); arity]; + let plans = pointwise_plans(operator, class, &inputs, &tensor.shape, &config); + assert!(!plans.is_empty(), "case {case}"); + assert!(plans.iter().all(|plan| { + plan.operator == operator + && plan.requirements.output.format.precision == selected_precision + })); + } + } + } + + fn value(lowered: &MidGraph, id: MidValueId) -> &MidValue { + &lowered.values[id.index() as usize] + } + + #[test] + fn randomized_parallel_reduction_plans_cover_uneven_three_axis_grids() { + let mut random = fastrand::Rng::with_seed(0x7061_7274_6961_6c73); + let mut distributed_result_cases = 0; + for _ in 0..RANDOM_CASES { + let output_columns = AMP_OUTPUT_COLUMN_BLOCK; + let inner_partitions = random.u16(2..=4); + let column_partitions = random.u16(1..=4); + let row_partitions = random.u16(1..=8); + let tiles = row_partitions * column_partitions * inner_partitions; + let k = u32::from(inner_partitions) * 64 + random.u32(0..64); + let n = u32::from(column_partitions) * output_columns + random.u32(0..output_columns); + let m = u32::from(row_partitions) + random.u32(0..=16); + let base = amp_grid_gemm_plan( + GemmOptions::default(), + Precision::F16, + 16, + GemmGeometry { + block: GemmBlockShape { + inner: 64, + output_columns, + }, + orientation: GemmOrientation::Normal, + compute: GemmGrid { + rows: 1, + columns: tiles, + inner: 1, + }, + result: GemmResultGrid { + rows: 1, + columns: tiles, + }, + order: GridOrder::ColumnsFast, + }, + output_columns, + AmpWeightPlacement::resident(MemoryClass::Standard), + ); + let inputs = [ + TensorType::new([m, k], Precision::F16, Layout::row_sharded(tiles)), + TensorType::new([k, n], Precision::F16, Layout::row_sharded(tiles)), + ]; + let config = PipelineConfig::new(tiles).with_planning_beam_width(16); + let candidates = parallel_reduction_plans( + &base, + &inputs, + &TensorShape(vec![m, n]), + &config, + &Ipu21CostModel, + true, + None, + None, + ); + assert!( + !candidates.is_empty(), + "shape={m}x{k}x{n} tiles={tiles} output_columns={output_columns}" + ); + distributed_result_cases += usize::from(candidates.iter().any(|candidate| { + let Some(ScheduleStep::Gemm(plan)) = candidate.steps.first() else { + return false; + }; + if plan.geometry.compute.inner < 2 { + return false; + } + plan.geometry.result + != GemmResultGrid { + rows: plan.geometry.compute.rows, + columns: plan.geometry.compute.columns, + } + })); + for candidate in candidates { + assert!( + candidate.supports(&inputs, &TensorShape(vec![m, n])), + "unsupported candidate: {candidate:?}; shape={m}x{k}x{n}" + ); + let Some(ScheduleStep::Gemm(plan)) = candidate.steps.first() else { + panic!("parallel GEMM candidate has non-GEMM schedule"); + }; + if plan.geometry.compute.inner < 2 { + panic!("parallel GEMM candidate has no reduction plan"); + } + let partial = plan + .partial_tensor(&TensorType { + shape: TensorShape(vec![m, n]), + format: candidate.requirements.output.format.clone(), + }) + .expect("supported parallel plan has a realizable partial layout"); + assert_eq!( + partial.format.layout.tiling.tile_count, + plan.geometry.compute.rows * plan.geometry.compute.columns + ); + assert!(!layout_has_empty_shards( + &partial.format.layout, + &partial.shape + )); + assert!(matches!(candidate.steps.first(), + Some(ScheduleStep::Gemm(crate::GemmMap { + geometry: GemmGeometry { + block: GemmBlockShape { inner: inner_block, output_columns: output_column_block }, + orientation, + compute: GemmGrid { rows: actual_rows, columns: actual_columns, inner: actual }, + .. + }, + .. + })) if actual_rows * actual_columns * actual <= tiles + && actual_rows * actual_columns * actual >= tiles.div_ceil(2) + && u32::from(*actual_rows) + <= [m, n][orientation.physical_left_input()] + && u32::from(*actual_columns) * output_column_block + >= [m, n][orientation.physical_right_input()] + && u32::from(*actual) * inner_block >= k + )); + } + } + assert!(distributed_result_cases > 0); + } + + #[test] + fn randomized_cycle_model_rewards_direct_interleaved_weight_loads() { + let mut random = fastrand::Rng::with_seed(0x6379_636c); + for _ in 0..RANDOM_CASES { + let rows = 1_u16 << random.u32(0..=2); + let columns = 1_u16 << random.u32(0..=2); + let tiles = rows * columns; + let m = u32::from(rows) * random.u32(1..=4); + let k = 64 * random.u32(2..=4); + let n = u32::from(columns) * 64; + let left = TensorType::new( + [m, k], + Precision::F16, + Layout::amp_left_grid(64, tiles, rows, columns, GridOrder::ColumnsFast), + ); + let mut standard_layout = Layout::block_major_matrix_grid( + 64, + 64, + tiles, + rows, + columns, + GridOrder::ColumnsFast, + ); + let mut direct_layout = standard_layout.clone(); + direct_layout.memory_class = MemoryClass::Interleaved; + standard_layout.memory_class = MemoryClass::Standard; + let standard = TensorType::new([k, n], Precision::F16, standard_layout); + let direct = TensorType::new([k, n], Precision::F16, direct_layout); + let output = TensorType::new( + [m, n], + Precision::F16, + Layout::amp_output_grid( + GemmOrientation::Normal, + 64, + tiles, + rows, + columns, + GridOrder::ColumnsFast, + ), + ); + let operator = MidOperator::Gemm { + options: GemmOptions::default(), + multiply: Precision::F16, + accumulate: AccumulationPrecision::F32, + }; + let gemm = gemm_map( + operator, + 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 requirements = OperatorRequirements { + inputs: Vec::new(), + output: OperandRequirement::new(output.format.clone(), 8), + output_aliasing: OutputAliasing::Fresh, + memory_space: MemorySpaceRequirements::default(), + }; + let schedule = OperatorSchedule { + operator, + steps: vec![ScheduleStep::Gemm(gemm)], + requirements: requirements.clone(), + }; + let standard_cost = + Ipu21CostModel.operator_cycles(&schedule, &[left.clone(), standard], &output); + let direct_cost = Ipu21CostModel.operator_cycles(&schedule, &[left, direct], &output); + assert!(direct_cost < standard_cost); + } + } + + #[test] + fn randomized_parameter_storage_balances_one_copy_independently_of_compute_grids() { + let mut random = fastrand::Rng::with_seed(0x6f77_6e65_7273); + for case in 0..RANDOM_CASES { + let row_partitions = 1_u16 << random.u32(1..=4); + let column_partitions = 1_u16 << random.u32(0..=4); + let tiles = row_partitions * column_partitions; + let inner_blocks = u32::from(row_partitions) * random.u32(1..=4); + let inner = inner_blocks * AMP_INNER_BLOCK; + let columns = u32::from(column_partitions) * AMP_OUTPUT_COLUMN_BLOCK; + let geometry = GemmGeometry { + block: GemmBlockShape { + inner: 64, + output_columns: AMP_OUTPUT_COLUMN_BLOCK, + }, + orientation: GemmOrientation::Normal, + compute: GemmGrid { + rows: row_partitions, + columns: column_partitions, + inner: 1, + }, + result: GemmResultGrid { + rows: row_partitions, + columns: column_partitions, + }, + order: GridOrder::ColumnsFast, + }; + let candidate = amp_grid_gemm_plan( + GemmOptions::default(), + Precision::F16, + 16, + geometry, + geometry.block.output_columns, + AmpWeightPlacement::resident(MemoryClass::Interleaved), + ); + let inputs = [ + TensorType::new( + [u32::from(row_partitions), inner], + Precision::F16, + candidate.requirements.inputs[0].format.layout.clone(), + ), + TensorType::new( + [inner, columns], + Precision::F16, + candidate.requirements.inputs[1].format.layout.clone(), + ), + ]; + let variants = + independent_parameter_storage(&candidate, &inputs, 1, &PipelineConfig::new(tiles)); + assert!(!variants.is_empty(), "case {case}"); + for variant in variants { + let tiling = &variant.requirements.inputs[1].format.layout.tiling; + assert_eq!(tiling.replicas, 1, "case {case}"); + assert!(tiling.tile_count <= tiles, "case {case}"); + assert_eq!( + tiling.tile_count, + tiling.replicas + * tiling + .axes + .iter() + .map(|axis| axis.partitions) + .product::(), + "case {case}" + ); + assert!( + variant.requirements.inputs[1] + .format + .layout + .resolve(&inputs[1].shape) + .is_ok(), + "case {case}" + ); + } + } + } + + fn assert_conversions_are_explicit(lowered: &MidGraph, operations: &[MidOperation]) { + for operation in operations { + let [input] = operation.inputs.as_slice() else { + continue; + }; + let [result] = operation.results.as_slice() else { + continue; + }; + let before = &value(lowered, *input).tensor_type; + let after = &value(lowered, *result).tensor_type; + match &operation.kind { + MidOperationKind::CastPrecision => { + assert_ne!(before.format.precision, after.format.precision); + assert_eq!(before.shape, after.shape); + assert_eq!(before.format.layout, after.format.layout); + } + MidOperationKind::Convert(transform, strategy, _, mappings) => { + if transform.is_none() { + assert_ne!(before.format.layout, after.format.layout); + assert_eq!( + *strategy, + layout_conversion_strategy( + before.format.precision, + &before.format.layout, + &after.format.layout, + ) + ); + assert_eq!(before.shape, after.shape); + assert_eq!(before.format.precision, after.format.precision); + } + assert!(!strategy.uses_intersections() || !mappings.is_empty()); + } + MidOperationKind::Operator(_) | MidOperationKind::Repeat(_) => {} + } + } + } + + struct ColumnParityCost; + + impl CostModel for ColumnParityCost { + fn target(&self) -> HardwareTarget { + HardwareTarget::Ipu21 + } + + fn operator_cycles( + &self, + schedule: &OperatorSchedule, + _inputs: &[TensorType], + output: &TensorType, + ) -> u64 { + let preferred = if output.shape.0.last().unwrap().is_multiple_of(2) { + Precision::F16 + } else { + Precision::F32 + }; + match schedule.operator { + MidOperator::Gemm { multiply, .. } if multiply == preferred => 0, + MidOperator::Gemm { .. } => 1, + _ => 0, + } + } + + fn cast_cycles(&self, _input: &TensorType, _to: Precision) -> u64 { + 0 + } + + fn rearrangement_cost( + &self, + _source: &TensorType, + _destination: &TensorType, + _strategy: ConversionStrategy, + ) -> crate::CostEstimate { + crate::CostEstimate::default() + } + } + + #[test] + fn randomized_axis_tiling_applies_or_rejects_padding() { + let mut random = fastrand::Rng::with_seed(0x7469_6c65); + for case in 0..RANDOM_CASES { + let rank = random.usize(1..=6); + let axis = random.usize(0..rank); + let extent = dimension(&mut random); + let block_size = random.u32(1..=32); + let partitions = random.u16(1..=16); + let replicas = random.u16(1..=4); + let padding = if random.bool() { + Padding::Reject + } else { + Padding::Zero + }; + let mut shape = (0..rank) + .map(|_| dimension(&mut random)) + .collect::>(); + shape[axis] = extent; + let layout = Layout::row_major(TensorTiling { + tile_count: partitions * replicas, + replicas, + axes: vec![AxisTiling::new( + TensorAxis::FromStart(axis as u16), + partitions, + block_size, + padding, + )], + }); + + let result = layout.resolve(&TensorShape(shape.clone())); + if padding == Padding::Reject && !extent.is_multiple_of(block_size) { + assert!( + matches!(result, Err(LayoutError::IndivisibleAxis { .. })), + "random case {case}" + ); + } else { + let resolved = result.unwrap(); + let padded = resolved.padded_shape(); + let expected = extent.div_ceil(block_size) * block_size; + assert_eq!(padded.0[axis], expected, "random case {case}"); + for (other, original) in shape.iter().enumerate() { + if other != axis { + assert_eq!(padded.0[other], *original, "random case {case}"); + } + } + } + } + } + + #[test] + fn randomized_gemm_lowering_makes_every_format_boundary_explicit() { + let mut random = fastrand::Rng::with_seed(0x6d69_6467); + for case in 0..RANDOM_CASES { + let tiles = [1, 2, 4, 8, 16][random.usize(0..5)]; + let (rows, inner, columns) = ( + u32::from(tiles) * random.u32(1..=2), + random.u32(1..=2) * 64, + random.u32(1..=2) * 64, + ); + let batches = (0..random.usize(0..=2)).map(|_| 1).collect::>(); + let multiply = precision(&mut random); + let mut left_shape = batches.clone(); + left_shape.extend([rows, inner]); + let mut right_shape = vec![1; batches.len()]; + right_shape.extend([inner, columns]); + + let mut graph = ComputeGraph::new(); + let left = graph.host_input("left", left_shape).unwrap(); + let right = graph.parameter("right", right_shape).unwrap(); + let product = graph.gemm(left, right).unwrap(); + graph.set_outputs([product]).unwrap(); + let config = PipelineConfig::new(tiles) + .with_search_domain( + PlannerSearchDomain::default() + .with_operator_precisions(OperatorClass::Gemm, [multiply]), + ) + .with_automatic_input(left, precision(&mut random)) + .with_automatic_input(right, precision(&mut random)); + + let lowered = lower(&graph, &config, &Ipu21CostModel).unwrap_or_else(|error| { + panic!( + "random case {case}: tiles={tiles} rows={rows} inner={inner} columns={columns} batches={batches:?}: {error:?}" + ) + }); + let operator = lowered + .operations + .iter() + .find(|operation| matches!(operation.kind, MidOperationKind::Operator(_))) + .unwrap(); + let Some(MidOperator::Gemm { + multiply: selected_multiply, + accumulate: selected_accumulate, + .. + }) = operator.operator() + else { + panic!("random case {case}: expected GEMM"); + }; + assert_eq!(selected_multiply, multiply, "random case {case}"); + assert_eq!( + selected_accumulate, + gemm_accumulation_precision(multiply), + "random case {case}" + ); + let output = value(&lowered, lowered.outputs[0]); + let expected_shape = graph.value_shape(product).unwrap().clone(); + assert_eq!( + output.tensor_type.shape, expected_shape, + "random case {case}" + ); + assert_eq!( + output.tensor_type.format.precision, multiply, + "random case {case}" + ); + assert_conversions_are_explicit(&lowered, &lowered.operations); + } + } + + #[test] + fn randomized_beam_finalists_have_consistent_costs_and_memory() { + let mut random = fastrand::Rng::with_seed(0x6265_616d); + for case in 0..RANDOM_CASES { + let tiles = [1, 2, 4, 8][random.usize(0..4)]; + let rows = u32::from(tiles) * random.u32(1..=8); + let inner = random.u32(1..=4) * 64; + let columns = random.u32(1..=4) * 64; + let row = format(Precision::F16, Layout::row_sharded(tiles)); + let right = format( + Precision::F16, + Layout::block_major_matrix_storage( + GemmOrientation::Normal, + 64, + AMP_OUTPUT_COLUMN_BLOCK, + tiles, + 1, + 1, + MemoryClass::Standard, + ), + ); + let mut graph = ComputeGraph::new(); + let activation = graph.host_input("activation", [rows, inner]).unwrap(); + let weights = graph.parameter("weights", [inner, columns]).unwrap(); + let activated = graph.gelu(activation).unwrap(); + let product = graph.gemm(activated, weights).unwrap(); + graph.set_outputs([product]).unwrap(); + + let make_config = |beam_width| { + PipelineConfig::new(tiles) + .with_search_domain( + PlannerSearchDomain::default().with_active_tile_counts([tiles]), + ) + .with_input(activation, row.clone()) + .with_input(weights, right.clone()) + .with_planning_beam_width(beam_width) + }; + let searched_config = make_config(2); + let finalists = lower_finalists(&graph, &searched_config, &Ipu21CostModel, 2).unwrap(); + assert!( + !finalists.is_empty() && finalists.len() <= 2, + "random case {case}" + ); + for finalist in &finalists { + assert_eq!( + finalist.metrics.cost.cycles, + finalist + .operations + .iter() + .map(|operation| operation.metrics.cost.cycles) + .sum::(), + "random case {case}" + ); + assert_eq!( + finalist.metrics.cost.exchange_cycles, + finalist + .operations + .iter() + .map(|operation| operation.metrics.cost.exchange_cycles) + .sum::(), + "random case {case}" + ); + assert!( + finalist.metrics.cost.exchange_cycles <= finalist.metrics.cost.cycles, + "random case {case}" + ); + } + let searched = &finalists[0]; + assert!( + searched.metrics.memory.fits_with_budget( + searched_config.target.memory_constraints(), + searched_config.standard_memory_reservation_bytes, + searched_config.tile_memory_budget_bytes, + ), + "random case {case}" + ); + } + } + + #[test] + fn randomized_gemm_lowering_rejects_per_batch_weights() { + let mut random = fastrand::Rng::with_seed(0x6261_7463); + for _ in 0..RANDOM_CASES { + let batch = random.u32(2..=8); + let rows = random.u32(1..=8); + let mut graph = ComputeGraph::new(); + let left = graph.host_input("left", [batch, rows, 64]).unwrap(); + let right = graph.parameter("right", [batch, 64, 64]).unwrap(); + let output = graph.gemm(left, right).unwrap(); + graph.set_outputs([output]).unwrap(); + let config = PipelineConfig::new(1) + .with_automatic_input(left, Precision::F16) + .with_automatic_input(right, Precision::F16); + assert!(matches!( + lower(&graph, &config, &Ipu21CostModel), + Err(LoweringError::UnsupportedGemmBatching(_)) + )); + } + } + + #[test] + fn randomized_gemms_choose_precision_independently_within_one_graph() { + let mut random = fastrand::Rng::with_seed(0x6d75_6c74); + for case in 0..RANDOM_CASES / 4 { + let tiles = random.u16(1..=64); + let rows = u32::from(tiles) * small_dimension(&mut random); + let inner = random.u32(1..=64); + let even_columns = random.u32(1..=16) * 2; + let odd_columns = random.u32(1..=16) * 2 - 1; + let layout = Layout::row_sharded(tiles); + let mut graph = ComputeGraph::new(); + let left = graph.host_input("left", [rows, inner]).unwrap(); + let even_right = graph.parameter("even", [inner, even_columns]).unwrap(); + let odd_right = graph.parameter("odd", [inner, odd_columns]).unwrap(); + let even = graph.gemm(left, even_right).unwrap(); + let odd = graph.gemm(left, odd_right).unwrap(); + graph.set_outputs([even, odd]).unwrap(); + let input_format = format(precision(&mut random), layout); + let config = PipelineConfig::new(tiles) + .with_input(left, input_format.clone()) + .with_input(even_right, input_format.clone()) + .with_input(odd_right, input_format); + + let lowered = lower(&graph, &config, &ColumnParityCost).unwrap(); + let chosen = lowered + .operations + .iter() + .filter_map(|operation| match operation.operator() { + Some(MidOperator::Gemm { multiply, .. }) => Some(multiply), + _ => None, + }) + .collect::>(); + assert_eq!( + chosen, + vec![Precision::F16, Precision::F32], + "random case {case}" + ); + for operation in lowered + .operations + .iter() + .filter(|operation| matches!(operation.operator(), Some(MidOperator::Gemm { .. }))) + { + let requirements = &operation.operator_plan().unwrap().requirements; + assert!( + requirements + .inputs + .iter() + .chain([&requirements.output]) + .all(|requirement| requirement.allocation.alignment == 32) + ); + assert_eq!( + requirements.output.format.layout.memory_class, + MemoryClass::Interleaved + ); + let orientation = match operation.operator_plan().unwrap().steps.first() { + Some(ScheduleStep::Gemm(gemm)) => gemm.geometry.orientation, + _ => panic!("GEMM operation lacks a GEMM schedule step"), + }; + let physical_left = orientation.physical_left_input(); + assert_eq!( + requirements.memory_space.distinct_element_groups, + [vec![ + MemoryOperand::Output, + MemoryOperand::Input(physical_left as u16), + ]] + ); + let expected_tail = match operation.operator() { + Some(MidOperator::Gemm { + multiply: Precision::F16, + .. + }) => 16, + Some(MidOperator::Gemm { + multiply: Precision::F32, + .. + }) => 32, + _ => unreachable!(), + }; + assert_eq!( + requirements.inputs[0].allocation.access_tail_bytes, + expected_tail + ); + } + } + } + + #[test] + fn randomized_non_gemm_lowering_honors_operator_plans() { + let mut random = fastrand::Rng::with_seed(0x6164_642b); + for case in 0..RANDOM_CASES { + let tiles = random.u16(1..=64); + let batch = random.u32(1..=2); + let query_rows = u32::from(tiles) * random.u32(1..=2); + let key_rows = random.u32(1..=8); + let channels = random.u32(1..=8); + let value_channels = random.u32(1..=8); + let mut graph = ComputeGraph::new(); + let activation = graph + .host_input("activation", [batch, query_rows, channels]) + .unwrap(); + let residual = graph + .host_input("residual", [batch, query_rows, channels]) + .unwrap(); + let query = graph + .host_input("query", [batch, query_rows, channels]) + .unwrap(); + let key = graph + .host_input("key", [batch, key_rows, channels]) + .unwrap(); + let attention_value = graph + .host_input("value", [batch, key_rows, value_channels]) + .unwrap(); + let activated = graph.gelu(activation).unwrap(); + let sum = graph.add(activated, residual).unwrap(); + let attended = graph.flash_attention(query, key, attention_value).unwrap(); + graph.set_outputs([sum, attended]).unwrap(); + + let config = PipelineConfig::new(tiles) + .with_automatic_input(activation, Precision::F16) + .with_automatic_input(residual, Precision::F16) + .with_automatic_input(query, Precision::F16) + .with_automatic_input(key, Precision::F16) + .with_automatic_input(attention_value, Precision::F16); + + let lowered = lower(&graph, &config, &Ipu21CostModel).unwrap(); + let operators = lowered + .operations + .iter() + .filter(|operation| matches!(operation.kind, MidOperationKind::Operator(_))) + .collect::>(); + assert_eq!(operators.len(), 3, "random case {case}"); + let gelu = operators + .iter() + .copied() + .find(|operation| matches!(operation.operator(), Some(MidOperator::Gelu))) + .expect("random graph retains its GeLU"); + let add = operators + .iter() + .copied() + .find(|operation| matches!(operation.operator(), Some(MidOperator::Add(_)))) + .expect("random graph retains its add"); + let attention = operators + .iter() + .copied() + .find(|operation| { + matches!( + operation.operator(), + Some(MidOperator::FlashAttention { .. }) + ) + }) + .expect("random graph retains its attention"); + assert_eq!( + value(&lowered, gelu.results[0]) + .tensor_type + .format + .precision, + Precision::F16 + ); + assert_eq!( + gelu.operator_plan().unwrap().requirements.output_aliasing, + OutputAliasing::MayAliasInputs(vec![0]) + ); + assert_eq!( + value(&lowered, add.results[0]).tensor_type.format.precision, + Precision::F16 + ); + assert_eq!( + add.operator_plan().unwrap().requirements.output_aliasing, + OutputAliasing::MayAliasInputs(vec![0, 1]) + ); + assert!(matches!( + attention.operator(), + Some(MidOperator::FlashAttention { options, .. }) + if options == AttentionOptions::default() + )); + assert_eq!( + value(&lowered, attention.results[0]).tensor_type.shape.0, + vec![batch, query_rows, value_channels], + "random case {case}" + ); + assert_conversions_are_explicit(&lowered, &lowered.operations); + } + } + + #[test] + fn randomized_repeat_lowering_retains_sequences_without_unrolling() { + let mut random = fastrand::Rng::with_seed(0x7265_7065); + for case in 0..RANDOM_CASES { + let tiles = random.u16(1..=64); + let size = u32::from(tiles); + let count = random.u32(1..=12); + let layout = Layout::row_sharded(tiles); + let carried_format = format(precision(&mut random), layout.clone()); + let mut graph = ComputeGraph::new(); + let carried = graph.host_input("state", [size, size]).unwrap(); + let weights = (0..count) + .map(|index| graph.parameter(format!("weight.{index}"), [size, size])) + .collect::, _>>() + .unwrap(); + let sequence = graph.value_sequence("weights", weights.clone()).unwrap(); + let output = graph + .repeat(count, [carried], [], [sequence], |body, arguments| { + Ok(vec![ + body.gemm(arguments.carried[0], arguments.iterated[0])?, + ]) + }) + .unwrap()[0]; + graph.set_outputs([output]).unwrap(); + let mut config = PipelineConfig::new(tiles).with_input(carried, carried_format.clone()); + for weight in weights { + config + .inputs + .insert(weight, format(precision(&mut random), layout.clone())); + } + + let lowered = lower(&graph, &config, &Ipu21CostModel).unwrap(); + let repeat = lowered + .operations + .iter() + .find_map(|operation| match &operation.kind { + MidOperationKind::Repeat(repeat) => Some(repeat), + _ => None, + }) + .unwrap(); + assert_eq!(repeat.count, count, "random case {case}"); + assert_eq!(repeat.iterated_inputs.len(), 1, "random case {case}"); + assert_eq!( + repeat.iterated_inputs[0].len(), + count as usize, + "random case {case}" + ); + let sequence_format = &value(&lowered, repeat.iterated_inputs[0][0]) + .tensor_type + .format; + assert!(repeat.iterated_inputs[0].iter().all(|value_id| { + &value(&lowered, *value_id).tensor_type.format == sequence_format + })); + assert_eq!( + &value(&lowered, repeat.body.yields[0]).tensor_type.format, + &carried_format, + "random case {case}" + ); + assert_eq!( + &value(&lowered, lowered.outputs[0]).tensor_type.format, + &carried_format, + "random case {case}" + ); + assert_conversions_are_explicit(&lowered, &lowered.operations); + assert_conversions_are_explicit(&lowered, &repeat.body.operations); + } + } + + #[test] + fn randomized_single_use_views_are_claimed_by_slice_consumers() { + let mut random = fastrand::Rng::with_seed(0x6465_6665_7272_6564); + for case in 0..RANDOM_CASES / 32 { + let heads = random.u32(2..=6); + let head_width = random.u32(4..=40) * 2; + let tokens = random.u32(1..=3) * AMP_INNER_BLOCK; + let model_width = heads * head_width; + let tiles = u16::try_from(heads * tokens.div_ceil(AMP_INNER_BLOCK)).unwrap(); + let mut graph = ComputeGraph::new(); + let input = graph.host_input("input", [1, tokens, model_width]).unwrap(); + let mut projected = Vec::new(); + let mut parameters = Vec::new(); + for index in 0..3 { + let weights = graph + .parameter(format!("projection.{index}"), [model_width, model_width]) + .unwrap(); + parameters.push(weights); + projected.push(graph.gemm(input, weights).unwrap()); + } + let split = projected + .iter() + .map(|&value| graph.split_heads(value, heads).unwrap()) + .collect::>(); + let output = graph.flash_attention(split[0], split[1], split[2]).unwrap(); + graph.set_outputs([output]).unwrap(); + let mut config = PipelineConfig::new(tiles).with_automatic_input(input, Precision::F16); + for parameter in parameters { + config = config.with_automatic_input(parameter, Precision::F16); + } + config.conversion_streaming = ConversionStreamingPolicy::Always; + + let lowered = lower(&graph, &config, &Ipu21CostModel).unwrap(); + let producers = lowered + .operations + .iter() + .filter(|operation| { + matches!(operation.kind, MidOperationKind::Convert(Some(_), ..)) + }) + .collect::>(); + assert_eq!(producers.len(), split.len(), "random case {case}"); + assert!( + producers.iter().all(|operation| { + operation.metrics.cost.cycles == 0 + && operation + .conversion() + .is_some_and(|(strategy, materialization)| { + materialization == OperandMaterialization::DispatchSlices + && strategy.uses_intersections() + }) + }), + "random case {case}" + ); + let consumer = lowered + .operations + .iter() + .find(|operation| { + matches!( + operation.operator(), + Some(MidOperator::FlashAttention { .. }) + ) + }) + .unwrap(); + assert!( + consumer.metrics.cost.exchange_footprint.phases >= 2, + "random case {case}" + ); + assert!( + consumer.metrics.cost.exchange_row_bytes(config.target) != 0, + "random case {case}" + ); + assert_eq!( + lowered.metrics.cost.cycles, + lowered + .operations + .iter() + .map(|operation| operation.metrics.cost.cycles) + .sum::(), + "random case {case}" + ); + let tiled = crate::low::lower_to_tiles(&lowered, &config) + .unwrap_or_else(|error| { + panic!( + "random case {case}, heads {heads}, width {head_width}, tokens {tokens}: {error}" + ) + }); + crate::KernelBuildPlan::from_program(&tiled) + .unwrap_or_else(|error| panic!("random case {case}: {error}")); + let attention_phases = tiled + .exchange_phases + .iter() + .filter(|phase| phase.provenance.operation == consumer.source) + .count(); + assert!( + attention_phases <= tokens.div_ceil(AMP_INNER_BLOCK) as usize + 2, + "random case {case}: {attention_phases} attention exchange phases" + ); + } + } + + #[test] + fn randomized_unclaimed_views_retain_materialization_cost() { + let mut random = fastrand::Rng::with_seed(0x756e_636c_6169_6d65); + for case in 0..RANDOM_CASES / 8 { + let batch = random.u32(1..=4); + let heads = random.u32(1..=8); + let rows = random.u32(1..=4) * AMP_INNER_BLOCK; + let head_width = random.u32(1..=4) * AMP_COLUMN_MICRO; + let mut graph = ComputeGraph::new(); + let input = graph + .host_input("input", [batch, rows, heads * head_width]) + .unwrap(); + let output = graph.split_heads(input, heads).unwrap(); + graph.set_outputs([output]).unwrap(); + let tiles = u16::try_from(batch * heads).unwrap(); + let config = PipelineConfig::new(tiles).with_automatic_input(input, Precision::F16); + + let lowered = lower(&graph, &config, &Ipu21CostModel).unwrap(); + let operation = lowered + .operations + .iter() + .find(|operation| matches!(operation.kind, MidOperationKind::Convert(Some(_), ..))) + .unwrap(); + assert!(operation.metrics.cost.cycles != 0, "random case {case}"); + assert!( + operation + .conversion() + .is_some_and(|(strategy, _)| strategy.uses_intersections()), + "random case {case}" + ); + let tiled = crate::low::lower_to_tiles(&lowered, &config) + .unwrap_or_else(|error| panic!("random case {case}: {error}")); + crate::KernelBuildPlan::from_program(&tiled) + .unwrap_or_else(|error| panic!("random case {case}: {error}")); + } ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec } } diff --cc crates/ipu-codegen/src/package.rs index c778ac2,7a2a419..0000000 --- a/crates/ipu-codegen/src/package.rs +++ b/crates/ipu-codegen/src/package.rs @@@ -10,27 -1,38 +10,51 @@@ use crate::ValueBlocks use crate::graph::{ComputeGraph, OperationId, ValueId}; use crate::host; use crate::low::LowProgram; ++<<<<<<< HEAD +use crate::memory::{ + MemoryLayoutError, MemoryRequest, PROFILE_END_CYCLE, PROFILE_START_CYCLE, RUNTIME_STATE_BASE, + RUNTIME_STATE_BYTES, TileMemoryMap, WORKER_STACK_HEADROOM, +}; +use crate::{ + COMPLETE_SYMBOL, COMPLETION_ADDRESS_SYMBOL, CodegenOptions, HOST_RUN_SYMBOL, KernelBuildPlan, + PRNG_SEED_SYMBOL, PROGRAM_ADDRESS_SYMBOL, REPEAT_CALL_SYMBOL, RUNTIME_ENTRY_SYMBOL, + SAMPLE_CYCLE_SYMBOL, TileProgram, TileProgramLowering, WORKER_BARRIER_SYMBOL, + WORKER_STACK_BASE_SYMBOL, WORKER_SYNC_CONTEXT_SYMBOL, emit, lower_to_tiles, place, ++======= + use crate::memory::{MemoryLayoutError, MemoryRequest, TileMemoryMap}; + use crate::mid::lower_finalists; + use crate::mid::{MidGraph, MidOperationKind}; + use crate::operator::Precision; + use crate::{ + KernelBuildPlan, PipelineConfig, TileProgramLowering, lower_exchanges, lower_to_tiles, place, ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec shard_storage_bytes, }; +use crate::{Ipu21CostModel, PipelineConfig, Precision, TileGraph, lower_finalists}; use ipu_driver::{APPLICATION_LOAD_BASE, TILES_PER_BATCH}; use ipu_elf::{ElfError, LinkOptions, LinkedImage, Toolchain, link}; - use ipu_exchange::{ExchangeError, Topology, encode_br_m, encode_setzi_m}; use ipu_package::{ - Application, Binding, DEBUG_ALL_TILES, DebugRegion, DebugSymbol, EntryPoint, - PROFILE_CYCLES_BINDING, PackageError, ProfileExchangeActivity, ProfileExchangeActivityKind, - ProfileMetadata, ProfileStep, ProfileStepKind, RegionSlice, SEGMENT_EXECUTE, SEGMENT_READ, - SEGMENT_WRITE, Segment, TILE_MEMORY_BASE, TileImage, TileProfilePlan, + AddressRegion, Application, Binding, DEBUG_ALL_TILES, DebugRegion, DebugSymbol, EntryPoint, + PROFILE_CYCLES_BINDING, PackageError, ProfileExchangeActivity, ProfileMetadata, ProfileStep, + ProfileStepKind, RegionSlice, SEGMENT_EXECUTE, SEGMENT_READ, SEGMENT_WRITE, Segment, TileImage, + TileProfilePlan, + }; + use ipu_target::emit::{ + COMPLETE_SYMBOL, COMPLETION_ADDRESS_SYMBOL, CodegenError, CodegenOptions, GeneratedProgram, + HOST_RUN_SYMBOL, HOST_STAGING_SYMBOL, PATCH_ROW_SYMBOL, PATCH_WORD_SYMBOL, PRNG_SEED_SYMBOL, + PROGRAM_ADDRESS_SYMBOL, REPEAT_CALL_SYMBOL, RUNTIME_ENTRY_SYMBOL, SAMPLE_CYCLE_SYMBOL, + WORKER_BARRIER_SYMBOL, WORKER_STACK_BASE_SYMBOL, WORKER_SYNC_CONTEXT_SYMBOL, emit, + }; + use ipu_target::exchange::ExchangeError; + use ipu_target::hardware::HardwareTarget; + use ipu_target::instruction::{encode_br_m, encode_setzi_m}; + use ipu_target::memory::TILE_MEMORY_BASE; + use ipu_target::memory::{ + IPU21_DATA_BASE, PROFILE_END_CYCLE, PROFILE_START_CYCLE, RUNTIME_STATE_BASE, + RUNTIME_STATE_BYTES, WORKER_STACK_HEADROOM, }; + use ipu_target::program::{StepProfile, TileProgram, TileStep}; + use ipu_target::topology::Topology; use rayon::prelude::*; use std::collections::{BTreeMap, HashMap}; use std::fs; @@@ -155,18 -133,392 +157,401 @@@ pub struct CompiledTensor } #[derive(Clone, Debug)] - pub struct DiagnosticShard { + pub struct CompiledTensorShard { pub physical_tile: u16, pub address: u32, - pub storage: crate::LowShard, + pub storage: crate::BlockValue, } ++<<<<<<< HEAD +struct BuiltApplication { + application: Application, + placement: crate::Placement, + exchange_phases: Vec, + exchange_schedule: crate::ExchangeScheduleSnapshot, + exchange_code_base: u32, ++======= + /// Builds an application from address-resolved tile programs. + /// + /// This is the low-level counterpart to [`build_package`]. It deliberately has + /// no tensor bindings: callers supply initialized tile data and inspect it + /// through driver diagnostics. A zero-payload `run` rendezvous starts execution + /// after loading, so breakpoints in the program cannot race the loader. + pub fn build_tile_program_package( + target: HardwareTarget, + programs: &[TileProgram], + data: &[TileProgramData], + outputs: &[Binding], + toolchain: &Toolchain, + runtime_source: &std::path::Path, + ) -> PackageBuildResult { + let topology = target.topology(); + let execution_tiles = u16::try_from(topology.tile_count())?; + if programs.len() != usize::from(execution_tiles) + || programs + .iter() + .enumerate() + .any(|(tile, program)| usize::from(program.tile) != tile) + { + return Err(invalid( + "finalized tile programs must cover every target logical tile in order", + )); + } + if data + .iter() + .any(|segment| segment.tile >= execution_tiles || segment.data.is_empty()) + { + return Err(invalid( + "tile-program data has an invalid tile or empty payload", + )); + } + + let runtime_artifact = toolchain.compile(runtime_source, "static_runtime", &[])?; + let objects = vec![fs::read(runtime_artifact.object)?]; + let kernels = KernelBuildPlan::default(); + let mut retained_runtime = vec![ + COMPLETE_SYMBOL.into(), + HOST_RUN_SYMBOL.into(), + REPEAT_CALL_SYMBOL.into(), + WORKER_BARRIER_SYMBOL.into(), + ]; + for program in programs { + collect_compute_symbols(&mut retained_runtime, &program.steps); + } + retained_runtime.sort_unstable(); + retained_runtime.dedup(); + let layout = link_runtime( + &objects, + runtime_symbols(0, 0, 0)?, + &kernels, + &retained_runtime, + target, + )?; + let symbols = layout + .symbols + .clone() + .into_iter() + .collect::>(); + let linked_end = linked_end(&layout)?; + let mut memory = TileMemoryMap::new(); + reserve_linked_image(&mut memory, &layout, "linked runtime")?; + memory.reserve( + "host exchange aperture", + AddressRegion::new( + target.exchange().window_base, + target.exchange().window_base + target.exchange().window_bytes, + ), + )?; + memory.reserve( + "runtime state", + AddressRegion::new(RUNTIME_STATE_BASE, RUNTIME_EXECUTABLE_START), + )?; + let mut tile_data = vec![Vec::<(u32, u32)>::new(); usize::from(execution_tiles)]; + for segment in data { + let bytes = u32::try_from(segment.data.len())?; + let end = segment + .address + .checked_add(bytes) + .ok_or_else(|| invalid("tile data range overflow"))?; + tile_data[usize::from(segment.tile)].push((segment.address, end)); + } + let mut tile_rows = vec![Vec::<(u32, u32)>::new(); usize::from(execution_tiles)]; + for program in programs { + let mut rows = BTreeMap::new(); + collect_exchange_rows(&mut rows, &program.steps)?; + tile_rows[usize::from(program.tile)].extend(rows); + } + for tile in 0..execution_tiles { + for &(data_start, data_end) in &tile_data[usize::from(tile)] { + if let Some(&(row_start, row_end)) = tile_rows[usize::from(tile)] + .iter() + .find(|&&(row_start, row_end)| data_start < row_end && row_start < data_end) + { + return Err(invalid(format!( + "tile {tile} data at 0x{data_start:x}..0x{data_end:x} overlaps exchange row 0x{row_start:x}..0x{row_end:x}" + ))); + } + } + } + // Generated and linked code use common addresses on every tile, so choose + // them against the union of tile-local data and row ranges. Data on one + // tile may otherwise legally share an address with a row on another tile. + let mut tile_local_ranges = tile_data + .into_iter() + .chain(tile_rows) + .flatten() + .collect::>(); + tile_local_ranges.sort_unstable(); + let mut merged_tile_local = Vec::<(u32, u32)>::new(); + for (start, end) in tile_local_ranges { + if let Some((_, previous_end)) = merged_tile_local.last_mut() + && start <= *previous_end + { + *previous_end = (*previous_end).max(end); + } else { + merged_tile_local.push((start, end)); + } + } + for (start, end) in merged_tile_local { + memory.reserve( + "tile-local data or exchange rows", + AddressRegion::new(start, end), + )?; + } + + let launch = Binding { + name: "run-gate".into(), + dtype: "u32".into(), + shape: vec![1], + slices: vec![RegionSlice { + tile: u32::from(topology.physical(0)?), + tile_address: COMPLETION_ADDRESS + 4, + file_offset: 0, + size: 4, + }], + }; + let finish = Binding { + name: "run-finish".into(), + dtype: "u32".into(), + shape: vec![1], + slices: vec![RegionSlice { + tile: u32::from(topology.physical(0)?), + tile_address: COMPLETION_ADDRESS + 8, + file_offset: 0, + size: 4, + }], + }; + let mut run_outputs = outputs.to_vec(); + run_outputs.push(finish); + let host_bounds = AddressRegion::new( + IPU21_DATA_BASE, + TILE_MEMORY_BASE + ipu_target::memory::TILE_MEMORY_SIZE, + ); + let sizing_host_base = memory.next_free( + linked_end, + AddressRegion::new( + TILE_MEMORY_BASE, + ipu_target::memory::IPU21_EXECUTABLE_MEMORY_LIMIT, + ), + 8, + "host programs", + )?; + let provisional_ranges = memory.free_ranges(host_bounds.clone()); + let provisional_host = host::plan( + target, + &[], + std::slice::from_ref(&launch), + &run_outputs, + execution_tiles, + sizing_host_base, + &vec![provisional_ranges; usize::from(execution_tiles)], + )?; + let host_code_bytes = provisional_host + .end + .checked_sub(sizing_host_base) + .ok_or_else(|| invalid("host program size underflow"))?; + let host_code = memory.allocate(MemoryRequest { + name: "host programs", + bytes: host_code_bytes, + alignment: 8, + bounds: AddressRegion::new( + linked_end, + ipu_target::memory::IPU21_EXECUTABLE_MEMORY_LIMIT, + ), + end_alignment: 8, + guard_after: 0, + })?; + let host_ranges = memory.free_ranges(host_bounds.clone()); + let host = host::plan( + target, + &[], + std::slice::from_ref(&launch), + &run_outputs, + execution_tiles, + host_code.range.start, + &vec![host_ranges; usize::from(execution_tiles)], + )?; + if host.end - host_code.range.start > host_code_bytes { + return Err(invalid("host program grew after placement")); + } + let mut host_data_ranges = host + .segments + .iter() + .flatten() + .filter(|segment| segment.flags & SEGMENT_EXECUTE == 0) + .map(|segment| (segment.address, segment.address + segment.memory_size)) + .collect::>(); + host_data_ranges.sort_unstable(); + let mut merged_host_data = Vec::<(u32, u32)>::new(); + for (start, end) in host_data_ranges { + if let Some((_, previous_end)) = merged_host_data.last_mut() + && start <= *previous_end + { + *previous_end = (*previous_end).max(end); + } else { + merged_host_data.push((start, end)); + } + } + for (start, end) in merged_host_data { + memory.reserve("host program data", AddressRegion::new(start, end))?; + } + + let sizing_address = memory.next_free( + host_code.range.end, + AddressRegion::new( + TILE_MEMORY_BASE, + ipu_target::memory::IPU21_EXECUTABLE_MEMORY_LIMIT, + ), + 8, + "generated tile programs", + )?; + let maximum_bytes = programs.iter().try_fold(0u32, |maximum, program| { + let physical = topology.physical(program.tile)?; + let generated = emit( + program, + &symbols, + &host.programs[usize::from(physical)], + &CodegenOptions { + code_address: sizing_address, + ..CodegenOptions::default() + }, + )?; + Ok::<_, PackageBuildError>(maximum.max(u32::try_from(generated.bytes.len())?)) + })?; + let code_address = memory + .allocate(MemoryRequest { + name: "generated tile programs", + bytes: maximum_bytes, + alignment: 4, + bounds: AddressRegion::new( + linked_end, + ipu_target::memory::IPU21_EXECUTABLE_MEMORY_LIMIT, + ), + // Supervisor instruction fetch and exchange/paired memory access + // cannot safely use the same standard-memory element. Reserve the + // rest of the element so subsequently placed tensor data cannot + // become the source of an exchange while code executes from it. + end_alignment: ipu_target::memory::TILE_MEMORY_ELEMENT_SIZE, + guard_after: 0, + })? + .range + .start; + let generated = programs + .iter() + .map(|program| { + let physical = topology.physical(program.tile)?; + Ok(emit( + program, + &symbols, + &host.programs[usize::from(physical)], + &CodegenOptions { + code_address, + ..CodegenOptions::default() + }, + )?) + }) + .collect::>>()?; + + let mut segments = vec![Vec::new(); usize::from(execution_tiles)]; + for segment in data { + let physical = topology.physical(segment.tile)?; + segments[usize::from(physical)].push(Segment { + address: segment.address, + memory_size: u32::try_from(segment.data.len())?, + data: segment.data.clone(), + flags: SEGMENT_READ | SEGMENT_WRITE, + }); + } + for (physical, host_segments) in host.segments.iter().enumerate() { + segments[physical].extend(host_segments.iter().cloned()); + } + let context = TileBuildContext { + objects: &objects, + kernel_plan: &kernels, + retained_runtime: &retained_runtime, + code_address, + host_staging_address: host.staging_address, + target, + }; + let mut tiles = Vec::with_capacity(usize::from(execution_tiles)); + for logical in 0..execution_tiles { + let physical = topology.physical(logical)?; + tiles.push(build_tile( + u32::from(physical), + u32::from(logical), + &generated[usize::from(logical)], + &segments[usize::from(physical)], + &context, + )?); + } + tiles.sort_unstable_by_key(|tile| tile.physical_tile); + let mut application = Application { + tiles, + ..Application::default() + }; + add_linked_debug_map(&mut application, &layout)?; + for (logical, program) in generated.iter().enumerate() { + let physical = u32::from(topology.physical(u16::try_from(logical)?)?); + add_generated_debug_map(&mut application, physical, code_address, program)?; + } + application.outputs.push(Binding { + name: "completion".into(), + dtype: "u32".into(), + shape: vec![1], + slices: vec![RegionSlice { + tile: 0, + tile_address: COMPLETION_ADDRESS, + file_offset: 0, + size: 4, + }], + }); + application.outputs.extend(run_outputs); + application.inputs.push(launch); + application.entry_points.push(EntryPoint { + name: "run".into(), + command: 0, + external_syncs: 0, + }); + application.host_exchange = host.protocol; + application.validate()?; + Ok(application) + } + + fn collect_exchange_rows( + rows: &mut BTreeMap, + steps: &[TileStep], + ) -> PackageBuildResult<()> { + for step in steps { + match step { + TileStep::Exchange(exchange) => { + let bytes = u32::try_from(exchange.program.words.len())? + .checked_mul(4) + .ok_or_else(|| invalid("exchange row size overflow"))?; + let end = exchange + .program + .address + .checked_add(bytes) + .ok_or_else(|| invalid("exchange row range overflow"))?; + rows.entry(exchange.program.address) + .and_modify(|existing| *existing = (*existing).max(end)) + .or_insert(end); + } + TileStep::Repeat(repeat) => collect_exchange_rows(rows, &repeat.body)?, + TileStep::Compute(_) | TileStep::Checkpoint(_) => {} + } + } + Ok(()) + } + + fn collect_compute_symbols(symbols: &mut Vec, steps: &[TileStep]) { + for step in steps { + match step { + TileStep::Compute(compute) => symbols.push(compute.symbol.clone()), + TileStep::Repeat(repeat) => collect_compute_symbols(symbols, &repeat.body), + TileStep::Exchange(_) | TileStep::Checkpoint(_) => {} + } + } ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec } /// Compiles and packages a compute graph into a directly loadable IPU21 @@@ -180,33 -532,8 +565,38 @@@ pub fn build_package graph: &ComputeGraph, config: &PackageConfig, ) -> PackageBuildResult { ++<<<<<<< HEAD + let (built, low) = build_package_artifacts(graph, config, false)?; + let topology = active_topology(low.tile_count)?; + let inputs = package_inputs(&low, &built.placement, &topology)?; + let outputs = low + .outputs + .iter() + .enumerate() + .map(|(index, output)| { + diagnostic_tensor( + &low, + &built.placement, + &topology, + output.value, + Some(format!("output.{index}")), + ) + }) + .collect::>>()?; + let precisions = package_precisions(&low); + Ok(CompiledPackage { + application: built.application, + inputs, + outputs, + precisions, + exchange_phases: built.exchange_phases, + exchange_schedule: built.exchange_schedule, + exchange_code_base: built.exchange_code_base, + }) ++======= + let (built, _, _) = build_package_artifacts(graph, config, false)?; + Ok(built) ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec } /// Builds an ordinary optimized package with resumable PBRK0 traps after each @@@ -215,16 -542,26 +605,39 @@@ pub fn build_diagnostic_package( graph: &ComputeGraph, config: &PackageConfig, ++<<<<<<< HEAD +) -> PackageBuildResult { + let (built, low) = build_package_artifacts(graph, config, true)?; + let topology = active_topology(low.tile_count)?; + let inputs = package_inputs(&low, &built.placement, &topology)?; + let mut checkpoints = Vec::new(); + for (source, results) in &low.checkpoints { + let source = *source; + let tensors = results + .iter() + .map(|&value| diagnostic_tensor(&low, &built.placement, &topology, value, None)) ++======= + ) -> PackageBuildResult { + let (mut built, mid, low) = build_package_artifacts(graph, config, true)?; + let topology = active_topology(config.pipeline.target, low.tile_count)?; + let mut checkpoints = Vec::new(); + for operation in &mid.operations { + if !matches!( + operation.kind, + MidOperationKind::Operator(_) + | MidOperationKind::Convert(Some(_), ..) + | MidOperationKind::Repeat(_) + ) { + continue; + } + let Some(source) = operation.source else { + continue; + }; + let tensors = operation + .results + .iter() + .map(|&value| compiled_tensor(&mid, &low, &built.placement, &topology, value, None)) ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec .collect::>>()?; // A fully deferred view operation has no device work or independently // materialized boundary to stop at; its consumer's checkpoint covers @@@ -251,33 -588,29 +664,49 @@@ tensors, }); } ++<<<<<<< HEAD + Ok(DiagnosticPackage { + application: built.application, + inputs, + checkpoints, + precisions: package_precisions(&low), + exchange_phases: built.exchange_phases, + exchange_schedule: built.exchange_schedule, + exchange_code_base: built.exchange_code_base, + }) ++======= + built.checkpoints = checkpoints; + Ok(built) ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec } -fn package_precisions(mid: &MidGraph) -> BTreeMap { - mid.values +fn package_precisions(mid: &TileGraph) -> BTreeMap { + mid.logical_values .iter() .map(|value| (value.origin, value.tensor_type.format.precision)) .collect() } ++<<<<<<< HEAD +fn package_inputs( ++======= + fn compiled_graph_tensors( + mid: &MidGraph, ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec low: &LowProgram, placement: &crate::Placement, topology: &Topology, - ) -> PackageBuildResult> { - low.inputs + ) -> PackageBuildResult<(Vec, Vec)> { + let inputs = low + .inputs .iter() .map(|input| { ++<<<<<<< HEAD + diagnostic_tensor( ++======= + compiled_tensor( + mid, ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec low, placement, topology, @@@ -292,25 -641,30 +737,45 @@@ fn build_package_artifacts graph: &ComputeGraph, config: &PackageConfig, diagnostic_checkpoints: bool, ++<<<<<<< HEAD +) -> PackageBuildResult<(BuiltApplication, LowProgram)> { + validate_tile_count(u32::from(config.pipeline.tile_count))?; + let mut planning = config.pipeline.clone(); + planning.diagnostic_checkpoints = diagnostic_checkpoints; + if diagnostic_checkpoints { + planning.profiling = false; ++======= + ) -> PackageBuildResult<(CompiledPackage, MidGraph, LowProgram)> { + validate_tile_count( + config.pipeline.target, + u32::from(config.pipeline.tile_count), + )?; + let mut planning = config.pipeline.clone(); + planning.diagnostic_checkpoints = diagnostic_checkpoints; + if diagnostic_checkpoints { + planning.profiling = crate::ProfilingConfig::Disabled; ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec } + let cost_model = match planning.target { + HardwareTarget::Ipu21 => crate::Ipu21CostModel, + }; let finalists = build_phase("lower_mid", || { Ok(lower_finalists( graph, &planning, ++<<<<<<< HEAD + &Ipu21CostModel, + planning.exchange_schedule_finalists.max(4), ++======= + &cost_model, + planning.exchange_schedule_finalists, ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec )?) })?; - let (mid, low) = build_phase("select_finalist", || { - select_scheduled_finalist(finalists, &planning) + let mut selected = build_phase("select_finalist", || { + select_scheduled_finalist(finalists, &planning, config.tile_mapping.as_deref()) })?; + let low = &selected.program; tracing::info!( logical_shards = low.shards.len(), exchange_phases = low.exchange_phases.len(), @@@ -334,21 -689,83 +800,101 @@@ } Ok(objects) })?; ++<<<<<<< HEAD + let built = build_package_from_objects(&mut selected, &planning, &objects, &kernel_plan)?; + Ok((built, selected.program)) +} + +fn build_package_from_objects( + selected: &mut ScheduledPlan, + config: &PipelineConfig, + objects: &[Vec], + kernel_plan: &KernelBuildPlan, +) -> PackageBuildResult { + let program = &selected.program; + let provisional_placement = &selected.placement; + let provisional_exchanges = &selected.phases; + let exchange_cache = &mut selected.cache; + let topology = active_topology(program.tile_count)?; ++======= + let mut package_config = config.clone(); + package_config.pipeline = planning; + let built = build_package_from_objects(&mid, &low, &package_config, &objects, &kernel_plan)?; + Ok((built, mid, low)) + } + + fn select_scheduled_finalist( + finalists: Vec, + planning: &PipelineConfig, + ) -> PackageBuildResult<(MidGraph, LowProgram)> { + if finalists.len() == 1 { + let mid = finalists.into_iter().next().unwrap(); + tracing::info!( + estimated_cycles = mid.metrics.cost.cycles, + estimated_exchange_cycles = mid.metrics.cost.exchange_cycles, + "selected analytical operator plan" + ); + for operation in &mid.operations { + tracing::debug!( + source = ?operation.source, + kind = ?operation.kind, + memory = ?operation.metrics.memory, + plan = ?operation.operator_plan(), + "selected mid-level operation" + ); + } + let low = lower_to_tiles(&mid, planning)?; + return Ok((mid, low)); + } + + let mut ranked = Vec::with_capacity(finalists.len()); + for (index, mid) in finalists.into_iter().enumerate() { + let low = lower_to_tiles(&mid, planning)?; + let placement = place(&low)?; + let exchanges = lower_exchanges(&low, &placement, planning.target)?; + let scheduled_exchange_cycles = exchanges + .iter() + .map(|phase| u64::from(phase.event_cycles)) + .sum::() + .saturating_add( + (exchanges.len() as u64) + .saturating_mul(planning.target.costs().exchange_phase_cycles), + ); + let estimated_non_exchange_cycles = mid + .metrics + .cost + .cycles + .saturating_sub(mid.metrics.cost.exchange_cycles); + let refined_cycles = + estimated_non_exchange_cycles.saturating_add(scheduled_exchange_cycles); + tracing::info!( + finalist = index, + analytical_cycles = mid.metrics.cost.cycles, + analytical_exchange_cycles = mid.metrics.cost.exchange_cycles, + scheduled_exchange_cycles, + refined_cycles, + "scheduled operator-plan finalist" + ); + ranked.push((refined_cycles, index, mid, low)); + } + ranked.sort_by_key(|(cycles, index, _, _)| (*cycles, *index)); + let (_, selected, mid, low) = ranked.remove(0); + tracing::info!( + selected, + "selected physically scheduled operator-plan finalist" + ); + Ok((mid, low)) + } + + fn build_package_from_objects( + mid: &MidGraph, + program: &LowProgram, + config: &PackageConfig, + objects: &[Vec], + kernel_plan: &KernelBuildPlan, + ) -> PackageBuildResult { + let topology = active_topology(config.pipeline.target, program.tile_count)?; ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec let retained_runtime = runtime_retained_symbols(program, config); let layout = build_phase("link_runtime", || { link_runtime( @@@ -368,16 -789,24 +918,32 @@@ )?; memory.reserve( "runtime state", - RUNTIME_STATE_BASE..RUNTIME_EXECUTABLE_START, + AddressRegion::new(RUNTIME_STATE_BASE, RUNTIME_EXECUTABLE_START), )?; ++<<<<<<< HEAD + let execution_tile_count = u16::try_from(Topology::c600().tile_count())?; ++======= + let provisional_placement = build_phase("plan_exchange_storage", || Ok(place(program)?))?; + let provisional_exchanges = build_phase("lower_exchanges_provisional", || { + Ok(lower_exchanges( + program, + &provisional_placement, + config.pipeline.target, + )?) + })?; + let execution_tile_count = u16::try_from(config.pipeline.target.topology().tile_count())?; ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec let exchange_table_bytes = crate::tile::compact_exchange_table_bytes( - &provisional_exchanges, + provisional_exchanges, execution_tile_count, program.tile_count, )?; ++<<<<<<< HEAD + let profile_samples = config.profiling.then(|| { ++======= + let profile_samples = config.pipeline.profiling.records_steps().then(|| { ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec program .tiles .iter() @@@ -437,22 -874,20 +1011,37 @@@ let provisional_inputs = program .inputs .iter() ++<<<<<<< HEAD + .filter(|input| input.kind == crate::GraphInputKind::Host) + .map(|input| input_binding(program, provisional_placement, &topology, input)) ++======= + .zip(&provisional_tensors) + .filter(|(input, _)| input.kind == crate::GraphInputKind::Host) + .map(|(_, tensor)| tensor.binding()) ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec .collect::>>()?; let provisional_weights = program .inputs .iter() ++<<<<<<< HEAD + .filter(|input| input.kind == crate::GraphInputKind::Parameter) + .map(|input| input_binding(program, provisional_placement, &topology, input)) ++======= + .zip(&provisional_tensors) + .filter(|(input, _)| input.kind == crate::GraphInputKind::Parameter) + .map(|(_, tensor)| tensor.binding()) ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec .collect::>>()?; - let mut provisional_outputs = program - .outputs + let mut provisional_outputs = provisional_output_tensors .iter() ++<<<<<<< HEAD + .enumerate() + .map(|(index, output)| { + output_binding(program, provisional_placement, &topology, output, index) + }) ++======= + .map(CompiledTensor::binding) ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec .collect::>>()?; if let Some(storage) = &profile_storage { provisional_outputs.push(cycle_binding( @@@ -567,9 -1016,18 +1170,23 @@@ &symbols, host, &CodegenOptions { + target: config.pipeline.target, code_address: sizing_code_address, ++<<<<<<< HEAD + initial_profile_address: config.profiling.then_some(PROFILE_START_CYCLE), + final_profile_address: config.profiling.then_some(PROFILE_END_CYCLE), ++======= + initial_profile_address: config + .pipeline + .profiling + .records_overall_time() + .then_some(PROFILE_START_CYCLE), + final_profile_address: config + .pipeline + .profiling + .records_overall_time() + .then_some(PROFILE_END_CYCLE), ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec ..CodegenOptions::default() }, )?; @@@ -620,38 -1091,14 +1250,46 @@@ )?) })?; let lowered_exchanges = build_phase("lower_exchanges", || { - Ok(lower_exchanges( + Ok(crate::exchange::lower_exchanges_cached( program, &placement, ++<<<<<<< HEAD + &topology, + config.exchange_diagnostics, + exchange_cache, + )?) + })?; + let (placement, lowered_exchanges) = build_phase("optimize_exchange_placement", || { + placement::improve_exchange_placement( + program, + &standard_ranges, + &topology, + placement, + lowered_exchanges, + exchange_rows.as_ref().map_or(0, |storage| { + storage.reserved.end + - ipu_package::IPU21_SUPERVISOR_FETCH_LOOKAHEAD + - storage.range.start + }), + exchange_cache, + ) + })?; + let final_cost = + crate::estimate::scheduled_program_cycles(&program.program, &lowered_exchanges.phases)?; + tracing::info!( + final_cycles = final_cost.total, + final_exchange = final_cost.exchange, + "costed final placed program" + ); + let exchange_schedule = lowered_exchanges.schedule_snapshot; + let exchanges = lowered_exchanges.phases; ++======= + config.pipeline.target, + )?) + })?; + let exchanges = lowered_exchanges; + let (tensors, output_tensors) = compiled_graph_tensors(mid, program, &placement, &topology)?; ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec let inputs = program .inputs .iter() @@@ -661,16 -1109,15 +1300,19 @@@ let weights = program .inputs .iter() - .filter(|input| input.kind == crate::GraphInputKind::Parameter) - .map(|input| input_binding(program, &placement, &topology, input)) + .zip(&tensors) + .filter(|(input, _)| input.kind == crate::GraphInputKind::Parameter) + .map(|(_, tensor)| tensor.binding()) .collect::>>()?; - let mut outputs = program - .outputs + let mut outputs = output_tensors .iter() - .enumerate() - .map(|(index, output)| output_binding(program, &placement, &topology, output, index)) + .map(CompiledTensor::binding) .collect::>>()?; ++<<<<<<< HEAD + if config.profiling { ++======= + if config.pipeline.profiling.records_overall_time() { ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec outputs.push(cycle_binding( "profile.start-cycle", PROFILE_START_CYCLE, @@@ -785,9 -1223,18 +1430,23 @@@ &symbols, host, &CodegenOptions { + target: config.pipeline.target, code_address, ++<<<<<<< HEAD + initial_profile_address: config.profiling.then_some(PROFILE_START_CYCLE), + final_profile_address: config.profiling.then_some(PROFILE_END_CYCLE), ++======= + initial_profile_address: config + .pipeline + .profiling + .records_overall_time() + .then_some(PROFILE_START_CYCLE), + final_profile_address: config + .pipeline + .profiling + .records_overall_time() + .then_some(PROFILE_END_CYCLE), ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec ..CodegenOptions::default() }, )?) @@@ -883,36 -1335,37 +1547,47 @@@ }) } ++<<<<<<< HEAD +fn diagnostic_tensor( ++======= + fn compiled_tensor( + mid: &MidGraph, ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec low: &LowProgram, placement: &crate::Placement, topology: &Topology, value: crate::MidValueId, name: Option, ++<<<<<<< HEAD +) -> PackageBuildResult { + let mid_value = low + .logical_values ++======= + ) -> PackageBuildResult { + let mid_value = mid + .values ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec .get(value.index() as usize) .ok_or_else(|| invalid("diagnostic mid-level value is missing"))?; - let low_value = low.values.iter().find(|candidate| candidate.value == value); - let shards = if let Some(low_value) = low_value { - low_value - .shards - .iter() - .filter(|id| { - low.shards - .get(id.index() as usize) - .is_some_and(|shard| shard.definition != crate::ShardDefinition::Unmaterialized) - }) - .map(|id| { - let storage = low - .shards - .get(id.index() as usize) - .ok_or_else(|| invalid("diagnostic low-level shard is missing"))?; - let address = placement + let shards = low + .values + .iter() + .find(|candidate| candidate.value == value) + .into_iter() + .flat_map(|value| &value.shards) + .filter(|id| { + low.shards + .get(id.index() as usize) + .is_some_and(|shard| shard.definition != crate::ShardDefinition::Unmaterialized) + }) + .map(|id| { + let storage = low + .shards + .get(id.index() as usize) + .ok_or_else(|| invalid("compiled tensor shard is missing"))?; + Ok(CompiledTensorShard { + physical_tile: topology.physical(storage.tile)?, + address: placement .shard_addresses .get(id) .copied() @@@ -1145,95 -1591,77 +1813,160 @@@ fn runtime_retained_symbols(program: &L let mut symbols = vec![COMPLETE_SYMBOL.into()]; if !program.exchange_phases.is_empty() { symbols.push(WORKER_BARRIER_SYMBOL.into()); - symbols.push(crate::PATCH_ROW_SYMBOL.into()); + symbols.push(PATCH_ROW_SYMBOL.into()); if !program.repeat_runs.is_empty() { - symbols.push(crate::PATCH_WORD_SYMBOL.into()); + symbols.push(PATCH_WORD_SYMBOL.into()); } } ++<<<<<<< HEAD + if config.profiling { ++======= + if config.pipeline.profiling.records_overall_time() { ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec symbols.push(SAMPLE_CYCLE_SYMBOL.into()); } if !program.inputs.is_empty() || !program.outputs.is_empty() { - symbols.push(crate::HOST_RUN_SYMBOL.into()); - symbols.push(crate::REPEAT_CALL_SYMBOL.into()); + symbols.push(HOST_RUN_SYMBOL.into()); + symbols.push(REPEAT_CALL_SYMBOL.into()); } ++<<<<<<< HEAD + #[derive(Default)] + struct CopySymbols { + local: bool, + halfword: bool, + zero: bool, + } + let mut copies = CopySymbols::default(); + fn collect(program: &LowProgram, tile: &crate::TileWorkList, copies: &mut CopySymbols) { + for work in program.work(tile) { + match work { + crate::TileWorkRef::LocalCopy(copy) => { + copies.local = true; + copies.halfword |= !match copy.pattern { + crate::CopyPattern::Contiguous => copy.bytes, + crate::CopyPattern::Strided { row_bytes, .. } => row_bytes, + } + .is_multiple_of(4); + } + crate::TileWorkRef::Kernel(run) => { + copies.zero |= matches!(run.kernel, crate::TileKernelSpec::FillZero { .. }); + } + crate::TileWorkRef::Repeat(repeat) => collect(program, &repeat.body, copies), + _ => {} + } + } + } + for tile in &program.tiles { + collect(program, tile, &mut copies); + } + if copies.halfword { + symbols.push(crate::COPY_U16_SYMBOL.into()); + } + if copies.local { + symbols.extend( + [ + crate::COPY_U32_SYMBOL, + crate::COPY_U64_SYMBOL, + crate::COPY_STRIDED_U64_SYMBOL, + ] + .map(String::from), + ); + } + if copies.zero { + symbols.push(crate::FILL_ZERO_U64_SYMBOL.into()); ++======= + if program + .tiles + .iter() + .any(|tile| tile_has_local_copy(program, tile)) + { + if program + .tiles + .iter() + .any(|tile| tile_has_halfword_copy(program, tile)) + { + symbols.push(ipu_target::emit::COPY_U16_SYMBOL.into()); + } + symbols.push(ipu_target::emit::COPY_U32_SYMBOL.into()); + symbols.push(ipu_target::emit::COPY_U64_SYMBOL.into()); + symbols.push(ipu_target::emit::COPY_STRIDED_U64_SYMBOL.into()); + } + if program + .tiles + .iter() + .any(|tile| tile_has_fill_zero(program, tile)) + { + symbols.push(ipu_target::emit::FILL_ZERO_U64_SYMBOL.into()); ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec } symbols } ++<<<<<<< HEAD +fn input_binding( + program: &LowProgram, + placement: &crate::Placement, + topology: &Topology, + input: &crate::ProgramInput, +) -> PackageBuildResult { + binding( + program, + placement, + topology, + input.name.clone(), + &input.shards, + ) +} + +fn output_binding( + program: &LowProgram, + placement: &crate::Placement, + topology: &Topology, + output: &ValueBlocks, + index: usize, +) -> PackageBuildResult { + binding( + program, + placement, + topology, + format!("output.{index}"), + &output.shards, + ) ++======= + fn tile_has_fill_zero(program: &LowProgram, tile: &crate::TileWorkList) -> bool { + program.work(tile).any(|work| match work { + crate::TileWorkRef::Kernel(run) => { + matches!(run.kernel, crate::TileKernelSpec::FillZero) + } + crate::TileWorkRef::Repeat(repeat) => tile_has_fill_zero(program, &repeat.body), + crate::TileWorkRef::Exchange(_) + | crate::TileWorkRef::LocalCopy(_) + | crate::TileWorkRef::Checkpoint(..) => false, + }) + } + + fn tile_has_halfword_copy(program: &LowProgram, tile: &crate::TileWorkList) -> bool { + program.work(tile).any(|work| match work { + crate::TileWorkRef::LocalCopy(copy) => match copy.pattern { + crate::LocalCopyPattern::Contiguous => !copy.bytes.is_multiple_of(4), + crate::LocalCopyPattern::Strided { row_bytes, .. } => !row_bytes.is_multiple_of(4), + }, + crate::TileWorkRef::Repeat(repeat) => tile_has_halfword_copy(program, &repeat.body), + crate::TileWorkRef::Exchange(_) + | crate::TileWorkRef::Kernel(_) + | crate::TileWorkRef::Checkpoint(..) => false, + }) + } + + fn tile_has_local_copy(program: &LowProgram, tile: &crate::TileWorkList) -> bool { + program.work(tile).any(|work| match work { + crate::TileWorkRef::LocalCopy(_) => true, + crate::TileWorkRef::Repeat(repeat) => tile_has_local_copy(program, &repeat.body), + crate::TileWorkRef::Exchange(_) + | crate::TileWorkRef::Kernel(_) + | crate::TileWorkRef::Checkpoint(..) => false, + }) ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec } fn cycle_binding(name: &str, address: u32, tile_count: u16, topology: &Topology) -> Binding { @@@ -1256,41 -1684,42 +1989,77 @@@ } } ++<<<<<<< HEAD +fn binding( + program: &LowProgram, + placement: &crate::Placement, + topology: &Topology, + name: String, + shards: &[crate::BlockValueId], +) -> PackageBuildResult { + let first = shards + .first() + .and_then(|id| program.shards.get(id.index() as usize)) + .ok_or_else(|| invalid("binding has no shards"))?; + let dtype = match first.tensor_type.format.precision { + crate::Precision::F8F143 { .. } => "f8f143", + crate::Precision::F16 => "f16", + crate::Precision::F32 => "f32", + }; + let mut file_offset = 0u64; + let slices = shards + .iter() + .map(|id| { + let shard = &program.shards[id.index() as usize]; + let size = u64::from(shard_storage_bytes(shard)?); + let slice = RegionSlice { + tile: u32::from(topology.physical(shard.tile)?), + tile_address: *placement + .shard_addresses + .get(id) + .ok_or_else(|| invalid("binding shard is not placed"))?, ++======= + fn profile_binding( + program: &LowProgram, + physical_to_logical: &[u16], + address: u32, + ) -> PackageBuildResult { + let mut file_offset = 0u64; + let mut sample_count = 0u32; + let slices = physical_to_logical + .iter() + .enumerate() + .filter_map(|(physical, &logical)| { + let steps = if logical < program.tile_count { + profile_step_count(program, &program.tiles[usize::from(logical)]) + } else { + inactive_profile_work(program).len() + }; + (steps != 0).then_some((physical, steps)) + }) + .map(|(physical, steps)| { + let samples = u32::try_from(steps + 1)?; + let size = u64::from(samples) + .checked_mul(4) + .ok_or_else(|| invalid("profile binding size overflow"))?; + let slice = RegionSlice { + tile: u32::try_from(physical)?, + tile_address: address, ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec file_offset, size, }; file_offset = file_offset .checked_add(size) ++<<<<<<< HEAD + .ok_or_else(|| invalid("binding file offset overflow"))?; + Ok(slice) ++======= + .ok_or_else(|| invalid("profile binding offset overflow"))?; + sample_count = sample_count + .checked_add(samples) + .ok_or_else(|| invalid("profile binding sample count overflow"))?; + Ok(slice) }) .collect::>>()?; Ok(Binding { @@@ -1301,6 -1730,423 +2070,424 @@@ }) } + fn instrument_profile( + program: &LowProgram, + exchanges: &[crate::PhysicalExchangePhase], + logical_tile: u16, + physical_tile: u32, + tile_program: &mut TileProgram, + address: u32, + ) -> PackageBuildResult { + let mut plans = Vec::with_capacity(tile_program.steps.len()); + if logical_tile < program.tile_count { + let schedule = program + .work(&program.tiles[usize::from(logical_tile)]) + .collect::>(); + if schedule.len() != tile_program.steps.len() { + return Err(invalid("tile profile work does not match finalized steps")); + } + for (index, (&work, step)) in schedule.iter().zip(&mut tile_program.steps).enumerate() { + if index != 0 && profile_work_can_merge(schedule[index - 1], work) { + continue; + } + let following = schedule[index + 1..].iter().find_map(|work| match work { + crate::TileWorkRef::Kernel(run) => Some(&run.provenance), + crate::TileWorkRef::Repeat(repeat) => Some(&repeat.provenance), + crate::TileWorkRef::Exchange(_) + | crate::TileWorkRef::LocalCopy(_) + | crate::TileWorkRef::Checkpoint(..) => None, + }); + step_profile(step).before = Some(profile_address(address, plans.len())?); + let mut description = profile_step( + program, + exchanges, + logical_tile, + index, + work, + step, + following, + )?; + let invocations = schedule[index + 1..] + .iter() + .take_while(|&&next| profile_work_can_merge(work, next)) + .count() + + 1; + description.metadata.push(ProfileMetadata { + name: "invocations".into(), + value: invocations.to_string(), + }); + description.local_index = u32::try_from(plans.len())?; + plans.push(description); + } + } else { + let schedule = inactive_profile_work(program); + if schedule.len() != tile_program.steps.len() { + return Err(invalid( + "inactive tile profile does not match finalized steps", + )); + } + for (index, (work, step)) in schedule + .into_iter() + .zip(&mut tile_program.steps) + .enumerate() + { + if let (crate::TileWorkRef::Checkpoint(operation, _), TileStep::Checkpoint(_)) = + (work, &*step) + { + step_profile(step).before = Some(profile_address(address, index)?); + plans.push(ProfileStep { + local_index: u32::try_from(index)?, + phase: u32::try_from(index)?, + epoch: 0, + operation: format!("operation.{}", operation.index()), + kind: ProfileStepKind::Idle, + kernel: "diagnostic-checkpoint".into(), + metadata: Vec::new(), + exchange_activities: Vec::new(), + exchange_event_cycles: 0, + }); + continue; + } + let (phase, provenance) = match (work, &*step) { + (crate::TileWorkRef::Exchange(id), TileStep::Exchange(_)) => { + let phase = &program.exchange_phases[id.index() as usize]; + (0x8000_0000 | id.index(), &phase.provenance) + } + (crate::TileWorkRef::Repeat(repeat), TileStep::Repeat(_)) => { + (u32::try_from(index)?, &repeat.provenance) + } + _ => return Err(invalid("inactive tile contains executable work")), + }; + step_profile(step).before = Some(profile_address(address, index)?); + plans.push(inactive_tile_description(index, phase, provenance)?); + } + } + if let Some(last) = tile_program.steps.last_mut() { + step_profile(last).after = Some(profile_address(address, plans.len())?); + } + Ok(TileProfilePlan { + physical_tile, + steps: plans, + }) + } + + fn inactive_profile_work(program: &LowProgram) -> Vec> { + program + .tiles + .first() + .into_iter() + .flat_map(|tile| program.work(tile)) + .filter(|work| { + matches!( + work, + crate::TileWorkRef::Exchange(_) + | crate::TileWorkRef::Repeat(_) + | crate::TileWorkRef::Checkpoint(..) + ) + }) + .collect() + } + + fn profile_step_count(program: &LowProgram, tile: &crate::TileWorkList) -> usize { + let mut previous = None; + let mut count = 0; + for work in program.work(tile) { + if previous.is_none_or(|previous| !profile_work_can_merge(previous, work)) { + count += 1; + } + previous = Some(work); + } + count + } + + fn profile_work_can_merge( + previous: crate::TileWorkRef<'_>, + current: crate::TileWorkRef<'_>, + ) -> bool { + matches!( + (previous, current), + (crate::TileWorkRef::Kernel(previous), crate::TileWorkRef::Kernel(current)) + if previous.kernel == current.kernel && previous.provenance == current.provenance + ) || matches!( + (previous, current), + ( + crate::TileWorkRef::LocalCopy(previous), + crate::TileWorkRef::LocalCopy(current) + ) if previous.bytes == current.bytes && previous.pattern == current.pattern + ) + } + + #[allow(clippy::too_many_arguments)] + fn profile_step( + program: &LowProgram, + exchanges: &[crate::PhysicalExchangePhase], + logical_tile: u16, + index: usize, + work: crate::TileWorkRef<'_>, + step: &mut TileStep, + following: Option<&crate::WorkProvenance>, + ) -> PackageBuildResult { + match (work, step) { + (crate::TileWorkRef::Exchange(id), TileStep::Exchange(exchange)) => { + let phase = &program.exchange_phases[id.index() as usize]; + if !exchange.active { + exchange_synchronization_description( + index, + 0x8000_0000 | id.index(), + &phase.provenance, + ) + } else { + let mut description = profile_description( + index, + 0x8000_0000 | id.index(), + &phase.provenance, + ProfileStepKind::Exchange, + "exchange", + )?; + let physical = exchanges + .get(id.index() as usize) + .ok_or_else(|| invalid("profile exchange phase is missing"))?; + description.exchange_activities = physical + .activities + .get(usize::from(logical_tile)) + .ok_or_else(|| invalid("profile exchange tile is missing"))? + .iter() + .map(|activity| ProfileExchangeActivity { + kind: activity.kind, + start_cycle: activity.start_cycle, + end_cycle: activity.end_cycle, + }) + .collect(); + description.exchange_event_cycles = physical.event_cycles; + Ok(description) + } + } + (crate::TileWorkRef::Kernel(run), TileStep::Compute(compute)) => { + let mut description = profile_description( + index, + u32::try_from(index)?, + &run.provenance, + ProfileStepKind::Compute, + &compute.symbol, + )?; + description.metadata.push(ProfileMetadata { + name: "kernelSpec".into(), + value: format!("{:?}", run.kernel), + }); + description.metadata.push(ProfileMetadata { + name: "outputElements".into(), + value: view_logical_elements(&run.output).to_string(), + }); + for (operand, input) in run.inputs.iter().enumerate() { + description.metadata.push(ProfileMetadata { + name: format!("input{operand}Elements"), + value: input + .views + .iter() + .map(view_logical_elements) + .sum::() + .to_string(), + }); + } + Ok(description) + } + (crate::TileWorkRef::LocalCopy(copy), TileStep::Compute(compute)) => { + if let Some(provenance) = following { + let mut description = profile_description( + index, + u32::try_from(index)?, + provenance, + ProfileStepKind::Compute, + &compute.symbol, + )?; + description.metadata[0].value = "LocalCopy".into(); + description.metadata.extend([ + ProfileMetadata { + name: "bytes".into(), + value: copy.bytes.to_string(), + }, + ProfileMetadata { + name: "pattern".into(), + value: format!("{:?}", copy.pattern), + }, + ]); + Ok(description) + } else { + Ok(ProfileStep { + local_index: u32::try_from(index)?, + phase: u32::try_from(index)?, + epoch: 0, + operation: String::new(), + kind: ProfileStepKind::Compute, + kernel: compute.symbol.clone(), + metadata: vec![ + ProfileMetadata { + name: "reason".into(), + value: "LocalCopy".into(), + }, + ProfileMetadata { + name: "bytes".into(), + value: copy.bytes.to_string(), + }, + ProfileMetadata { + name: "pattern".into(), + value: format!("{:?}", copy.pattern), + }, + ], + exchange_activities: Vec::new(), + exchange_event_cycles: 0, + }) + } + } + (crate::TileWorkRef::Repeat(repeat), TileStep::Repeat(_)) => profile_description( + index, + u32::try_from(index)?, + &repeat.provenance, + ProfileStepKind::Compute, + "repeat", + ), + (crate::TileWorkRef::Checkpoint(operation, _), TileStep::Checkpoint(_)) => { + Ok(ProfileStep { + local_index: u32::try_from(index)?, + phase: u32::try_from(index)?, + epoch: 0, + operation: format!("operation.{}", operation.index()), + kind: ProfileStepKind::Synchronization, + kernel: "diagnostic-checkpoint".into(), + metadata: Vec::new(), + exchange_activities: Vec::new(), + exchange_event_cycles: 0, + }) + } + _ => Err(invalid( + "tile profile work kind does not match finalized step", + )), + } + } + + fn view_logical_elements(view: &crate::ShardView) -> u64 { + view.extents.iter().fold(1u64, |elements, extent| { + elements.saturating_mul(u64::from(extent.logical_end.saturating_sub(extent.start))) + }) + } + + fn exchange_synchronization_description( + index: usize, + phase: u32, + provenance: &crate::WorkProvenance, + ) -> PackageBuildResult { + let mut description = profile_description( + index, + phase, + provenance, + ProfileStepKind::Synchronization, + "sync", + )?; + description.metadata[0].value = "ExchangeBarrier".into(); + Ok(description) + } + + fn inactive_tile_description( + index: usize, + phase: u32, + provenance: &crate::WorkProvenance, + ) -> PackageBuildResult { + let mut description = + profile_description(index, phase, provenance, ProfileStepKind::Idle, "idle")?; + description.metadata[0].value = "InactiveTile".into(); + Ok(description) + } + + fn profile_description( + index: usize, + phase: u32, + provenance: &crate::WorkProvenance, + kind: ProfileStepKind, + kernel: &str, + ) -> PackageBuildResult { + let mut metadata = vec![ProfileMetadata { + name: "reason".into(), + value: format!("{:?}", provenance.reason), + }]; + if let Some(value) = provenance.value { + metadata.push(ProfileMetadata { + name: "value".into(), + value: value.index().to_string(), + }); + } + Ok(ProfileStep { + local_index: u32::try_from(index)?, + phase, + epoch: 0, + operation: provenance + .operation + .map(|operation| format!("operation.{}", operation.index())) + .unwrap_or_default(), + kind, + kernel: kernel.into(), + metadata, + exchange_activities: Vec::new(), + exchange_event_cycles: 0, + }) + } + + fn step_profile(step: &mut TileStep) -> &mut StepProfile { + match step { + TileStep::Exchange(exchange) => &mut exchange.profile, + TileStep::Compute(compute) => &mut compute.profile, + TileStep::Repeat(repeat) => &mut repeat.profile, + TileStep::Checkpoint(checkpoint) => &mut checkpoint.profile, + } + } + + fn profile_address(base: u32, index: usize) -> PackageBuildResult { + base.checked_add( + u32::try_from(index)? + .checked_mul(4) + .ok_or_else(|| invalid("profile address overflow"))?, + ) + .ok_or_else(|| invalid("profile address overflow")) + } + + impl CompiledTensor { + pub fn binding(&self) -> PackageBuildResult { + let name = self + .name + .clone() + .ok_or_else(|| invalid("unnamed compiled tensor cannot be a host binding"))?; + let dtype = match self.precision { + crate::Precision::F8F143 { .. } => "f8f143", + crate::Precision::F16 => "f16", + crate::Precision::F32 => "f32", + }; + let mut file_offset = 0u64; + let slices = self + .shards + .iter() + .map(|shard| { + let size = u64::from(shard_storage_bytes(&shard.storage)?); + let slice = RegionSlice { + tile: u32::from(shard.physical_tile), + tile_address: shard.address, + file_offset, + size, + }; + file_offset = file_offset + .checked_add(size) + .ok_or_else(|| invalid("binding file offset overflow"))?; + Ok(slice) + }) + .collect::>>()?; + Ok(Binding { + name, + dtype: dtype.into(), + shape: self.shape.0.clone(), + slices, ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec + }) + } + } + fn runtime_symbols( physical_tile: u32, program_address: u32, diff --cc crates/ipu-codegen/src/place.rs index d02704b,efa12d0..0000000 --- a/crates/ipu-codegen/src/place.rs +++ b/crates/ipu-codegen/src/place.rs @@@ -1,24 -1,25 +1,44 @@@ //! Deterministic placement of logical shards in IPU21 tile SRAM. ++<<<<<<< HEAD +mod exchange; +pub(crate) use exchange::ExchangeConflicts; + +use crate::low::{LowProgram, TileWorkList, TileWorkRef}; +use crate::memory::IPU21_DATA_BASE; +use crate::{BlockValueId, ShardDefinition}; +use crate::{MemoryClass, MemoryOperand}; +use crate::{StorageError, shard_storage_bytes}; +use ipu_package::{ + IPU21_APPLICATION_MEMORY_LIMIT, IPU21_INTERLEAVED_ELEMENT_SIZE, IPU21_INTERLEAVED_MEMORY_BASE, + TILE_MEMORY_ELEMENT_SIZE, ++======= + use crate::layout::MemoryClass; + use crate::low::{ + KernelRequirements, LowProgram, LowShardId, ShardDefinition, TileWorkList, TileWorkRef, + }; + use crate::operator::{ + AllocationRequirements, MemoryElementRequirement, MemoryOperand, OperandRequirement, + }; + use crate::storage::{StorageError, shard_storage_bytes}; + use ipu_package::AddressRegion; + use ipu_target::memory::{ + IPU21_APPLICATION_MEMORY_LIMIT, IPU21_DATA_BASE, IPU21_INTERLEAVED_ELEMENT_SIZE, + IPU21_INTERLEAVED_MEMORY_BASE, TILE_MEMORY_ELEMENT_SIZE, ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec }; use rayon::prelude::*; use std::collections::{BTreeMap, BTreeSet}; #[derive(Clone, Debug, PartialEq, Eq)] pub struct Placement { ++<<<<<<< HEAD + pub shard_addresses: BTreeMap, + pub tile_auxiliary_ranges: Vec>, ++======= + pub shard_addresses: BTreeMap, + pub tile_auxiliary_ranges: Vec>, ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec } #[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)] @@@ -83,24 -83,17 +102,28 @@@ pub fn place(program: &LowProgram) -> R pub(crate) fn place_with_standard_ranges( program: &LowProgram, - standard_ranges: &[(u32, u32)], + standard_ranges: &[AddressRegion], ) -> Result { + place_with_offset(program, standard_ranges, 0) +} + +pub(crate) fn place_with_offset( + program: &LowProgram, + standard_ranges: &[(u32, u32)], + interleaved_offset: u32, +) -> Result { + if interleaved_offset >= IPU21_INTERLEAVED_ELEMENT_SIZE { + return Err(PlacementError::Overflow); + } if standard_ranges.is_empty() - || standard_ranges.iter().any(|&(start, end)| { - start < IPU21_DATA_BASE || end > IPU21_INTERLEAVED_MEMORY_BASE || start >= end + || standard_ranges.iter().any(|range| { + range.start < IPU21_DATA_BASE + || range.end > IPU21_INTERLEAVED_MEMORY_BASE + || range.is_empty() }) - || standard_ranges.windows(2).any(|pair| pair[0].1 > pair[1].0) + || standard_ranges + .windows(2) + .any(|pair| pair[0].end > pair[1].start) { return Err(PlacementError::OutOfMemory { tile: 0, @@@ -210,14 -174,13 +231,22 @@@ fn analyze_allocations(program: &LowPro fn place_tile( program: &LowProgram, tile: u16, ++<<<<<<< HEAD + standard_ranges: &[(u32, u32)], + interleaved_offset: u32, ++======= + standard_ranges: &[AddressRegion], ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec iterated: &[IteratedGroup], members: &BTreeMap>, root_of_member: &[usize], - root_requirements: &BTreeMap, + root_requirements: &BTreeMap, root_lifetimes: &BTreeMap, ++<<<<<<< HEAD +) -> Result<(u16, BTreeMap, Vec<(u32, u32)>), PlacementError> { ++======= + ) -> Result<(u16, BTreeMap, Vec), PlacementError> { ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec let mut grouped = BTreeSet::::new(); for group in iterated.iter().filter(|group| group.tile == tile) { let roots = group @@@ -235,8 -198,8 +264,13 @@@ // element, then return every remaining byte to standard allocations. let mut addresses = BTreeMap::new(); let mut interleaved = Arena::new( ++<<<<<<< HEAD + &[( + IPU21_INTERLEAVED_MEMORY_BASE + interleaved_offset, ++======= + &[AddressRegion::new( + IPU21_INTERLEAVED_MEMORY_BASE, ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec IPU21_APPLICATION_MEMORY_LIMIT, )], true, @@@ -468,10 -398,19 +505,26 @@@ fn collect_requirements for work in program.work(tile) { match work { TileWorkRef::Kernel(run) => { ++<<<<<<< HEAD + let inputs = &run.requirements.inputs; + let output = &run.requirements.output; + let distinct_elements = &run.requirements.distinct_elements; + for operands in distinct_elements { ++======= + let (inputs, output, memory_space) = match &run.requirements { + KernelRequirements::Operator(operator_requirements) => ( + &operator_requirements.inputs[..], + &operator_requirements.output, + &operator_requirements.memory_space, + ), + KernelRequirements::Conversion { + input, + output, + memory_space, + } => (std::slice::from_ref(input), output, memory_space), + }; + for operands in &memory_space.distinct_element_groups { ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec for operand in operands { match operand { MemoryOperand::Output => { @@@ -515,9 -454,8 +568,14 @@@ } } ++<<<<<<< HEAD +fn apply_requirement(target: &mut Requirement, requirement: &crate::KernelAccess) { + target.alignment = target.alignment.max(requirement.alignment); + target.access_tail = target.access_tail.max(requirement.access_tail_bytes); ++======= + fn apply_requirement(target: &mut AllocationRequirements, requirement: &OperandRequirement) { + target.merge(requirement.allocation); ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec } fn checked_union( @@@ -625,10 -567,10 +687,10 @@@ fn allocate_tile_class grouped: &BTreeSet, members: &BTreeMap>, root_of_member: &[usize], - root_requirements: &BTreeMap, + root_requirements: &BTreeMap, root_lifetimes: &BTreeMap, arena: &mut Arena, - addresses: &mut BTreeMap, + addresses: &mut BTreeMap, ) -> Result<(), PlacementError> { let mut requests = Vec::::new(); for group in iterated.iter().filter(|group| group.tile == tile) { @@@ -982,27 -887,14 +1047,27 @@@ mod tests }, ); let mid = lower(&graph, &config, &Ipu21CostModel).unwrap(); - let low = lower_to_tiles(&mid, &config).unwrap(); - let placement = place(&low).unwrap(); + let low = lower_to_tiles( + &crate::expand_tiles(&mid).unwrap(), + config.diagnostic_checkpoints, + ); + let placement = place_with_offset( + &low, + &[(IPU21_DATA_BASE, IPU21_INTERLEAVED_MEMORY_BASE)], + random.u32(0..8) * 4096, + ) + .unwrap(); let kernels = KernelBuildPlan::from_program(&low).unwrap(); - assert_eq!(placement.shard_addresses.len(), low.shards.len()); - for shard in &low.shards { + let resident = low + .shards + .iter() + .filter(|shard| !matches!(shard.definition, crate::ShardDefinition::Unmaterialized)) + .collect::>(); + assert_eq!(placement.shard_addresses.len(), resident.len()); + for shard in resident { let address = placement.shard_addresses[&shard.id]; match shard.tensor_type.format.layout.memory_class { - MemoryClass::Ipu21Interleaved => { + MemoryClass::Interleaved => { assert!( (IPU21_INTERLEAVED_MEMORY_BASE..IPU21_APPLICATION_MEMORY_LIMIT) .contains(&address) @@@ -1022,9 -914,8 +1087,14 @@@ &BTreeMap::new(), ) .unwrap(); ++<<<<<<< HEAD + { + let requirements = &run.requirements; + for operands in &requirements.distinct_elements { ++======= + if let KernelRequirements::Operator(requirements) = &run.requirements { + for operands in &requirements.memory_space.distinct_element_groups { ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec let mut ranges = Vec::new(); for operand in operands { let shards = match operand { diff --cc crates/ipu-codegen/src/storage.rs index 712e5d7,57bd048..0000000 --- a/crates/ipu-codegen/src/storage.rs +++ b/crates/ipu-codegen/src/storage.rs @@@ -1,15 -1,10 +1,23 @@@ //! Conversion from logical shard views to physical byte ranges. ++<<<<<<< HEAD +use crate::{ + AMP_COLUMN_MICRO, AmpOrder, BlockMajorOrder, ElementOrder, Precision, ShardExtent, TensorFormat, +}; + +/// Relative storage geometry, independent of low IR identities and placement. +#[derive(Clone, Copy, Debug)] +pub(crate) struct TensorStorage<'a> { + pub format: &'a TensorFormat, + pub extents: &'a [ShardExtent], +} ++======= + use crate::layout::{ + AMP_COLUMN_MICRO, BlockedOrder, LayoutError, NativeKernelOrder, StorageOrder, TensorRegion, + }; + use crate::low::{LowShard, ShardView}; + use crate::operator::Precision; ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub struct ByteSpan { @@@ -192,32 -175,73 +202,95 @@@ fn byte_spans Ok(spans) } ++<<<<<<< HEAD +fn physical_index( + shard: TensorStorage<'_>, ++======= + fn physical_index(shard: &LowShard, widths: &[u32], coordinates: &[u32]) -> StorageResult { + physical_index_for( + shard.tensor_type.format.layout.order, + shard.tensor_type.format.precision, + widths, + coordinates, + ) + } + + /// Resolves one semantic coordinate to a byte offset within an owning tensor + /// region. Conversion planning uses this only at layout block boundaries; + /// enumerating individual tensor elements remains a validation facility. + pub(crate) fn physical_byte_offset( + order: StorageOrder, + precision: Precision, + owner: &TensorRegion, + coordinates: &[u32], + ) -> StorageResult { + if coordinates.len() != owner.len() { + return Err(StorageError::InvalidView); + } + let mut local = Vec::with_capacity(coordinates.len()); + let mut widths = Vec::with_capacity(owner.len()); + for (coordinate, extent) in coordinates.iter().zip(owner.iter()) { + if *coordinate < extent.start || *coordinate >= extent.physical_end { + return Err(StorageError::InvalidView); + } + local.push(*coordinate - extent.start); + widths.push(extent.physical_end - extent.start); + } + let elements = physical_index_for(order, precision, &widths, &local)?; + let element_bytes = precision.bytes(); + u32::try_from(elements) + .ok() + .and_then(|elements| elements.checked_mul(element_bytes as u32)) + .ok_or(StorageError::Overflow) + } + + fn physical_index_for( + order: StorageOrder, + precision: Precision, ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec widths: &[u32], coordinates: &[u32], ) -> StorageResult { let rank = widths.len(); ++<<<<<<< HEAD + match shard.format.layout.order { + ElementOrder::RowMajor => encode_row_major(widths, coordinates), + ElementOrder::BlockMajor(order) => { ++======= + match order { + StorageOrder::Linear => encode_row_major(widths, coordinates), + StorageOrder::Blocked(order) => { ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec if rank < 2 { return Err(StorageError::MatrixRank); } - let rows = widths[rank - 2]; - let columns = widths[rank - 1]; - let outer = encode_row_major(&widths[..rank - 2], &coordinates[..rank - 2])?; - let row = coordinates[rank - 2]; - let column = coordinates[rank - 1]; + let [row_axis, column_axis] = order.physical_axes(rank)?; + let mut outer_widths = Vec::with_capacity(rank - 2); + let mut outer_coordinates = Vec::with_capacity(rank - 2); + for axis in 0..rank { + if axis != row_axis && axis != column_axis { + outer_widths.push(widths[axis]); + outer_coordinates.push(coordinates[axis]); + } + } + let rows = widths[row_axis]; + let columns = widths[column_axis]; + let outer = encode_row_major(&outer_widths, &outer_coordinates)?; + let row = coordinates[row_axis]; + let column = coordinates[column_axis]; let matrix_elements = u64::from(rows) * u64::from(columns); ++<<<<<<< HEAD + let within = block_major_matrix_index( + order, + shard.format.precision, + rows, + columns, + row, + column, + )?; ++======= + let within = block_major_matrix_index(order, precision, rows, columns, row, column)?; ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec outer .checked_mul(matrix_elements) .and_then(|base| base.checked_add(u64::from(within))) @@@ -232,10 -256,9 +305,14 @@@ let outer = encode_row_major(&widths[..rank - 2], &coordinates[..rank - 2])?; let row = coordinates[rank - 2]; let column = coordinates[rank - 1]; - if role == AmpOrder::TransposedRight { + if role == NativeKernelOrder::TransposedRight { let matrix_elements = u64::from(rows) * u64::from(columns); ++<<<<<<< HEAD + let within = + right_matrix_index(shard.format.precision, columns, rows, column, row)?; ++======= + let within = right_matrix_index(precision, columns, rows, column, row)?; ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec outer .checked_mul(matrix_elements) .and_then(|base| base.checked_add(u64::from(within))) @@@ -246,7 -269,7 +323,11 @@@ })?; amp_matrix_index( role, ++<<<<<<< HEAD + shard.format.precision, ++======= + precision, ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec flat_rows, columns, u32::try_from(outer) @@@ -259,8 -282,7 +340,12 @@@ .map(u64::from) } else { let matrix_elements = u64::from(rows) * u64::from(columns); ++<<<<<<< HEAD + let within = + amp_matrix_index(role, shard.format.precision, rows, columns, row, column)?; ++======= + let within = amp_matrix_index(role, precision, rows, columns, row, column)?; ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec outer .checked_mul(matrix_elements) .and_then(|base| base.checked_add(u64::from(within))) @@@ -420,23 -432,15 +495,35 @@@ fn block_major_panel_spans return Ok(None); } let rank = shard.extents.len(); ++<<<<<<< HEAD + let (inner_block, inner_axis, column_axis, column_tensor_axis) = match shard.format.layout.order + { + ElementOrder::BlockMajor(BlockMajorOrder::Matrix { row_block, .. }) => ( + u32::from(row_block), + rank - 2, + rank - 1, + crate::TensorAxis::FromEnd(1), + ), + ElementOrder::BlockMajor(BlockMajorOrder::TransposedMatrix { row_block, .. }) => ( + u32::from(row_block), + rank - 1, + rank - 2, + crate::TensorAxis::FromEnd(2), + ), + _ => return Ok(None), + }; + if shard.extents[..rank - 2] != view[..rank - 2] { ++======= + let StorageOrder::Blocked(order) = shard.tensor_type.format.layout.order else { + return Ok(None); + }; + let [inner_axis, column_axis] = order.physical_axes(rank)?; + let inner_block = u32::from(order.block_shape[0]); + let column_tensor_axis = order.axes[usize::from(order.permutation[1])]; + if (0..rank).any(|axis| { + axis != inner_axis && axis != column_axis && shard.extents[axis] != view.extents[axis] + }) { ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec return Ok(None); } let rows = shard.extents[inner_axis].physical_end - shard.extents[inner_axis].start; @@@ -541,9 -562,9 +628,15 @@@ fn physical_coordinates ) -> StorageResult> { let rank = widths.len(); let mut coordinates = vec![0; rank]; ++<<<<<<< HEAD + match shard.format.layout.order { + ElementOrder::RowMajor => decode_row_major(widths, physical, &mut coordinates), + ElementOrder::BlockMajor(order) => { ++======= + match shard.tensor_type.format.layout.order { + StorageOrder::Linear => decode_row_major(widths, physical, &mut coordinates), + StorageOrder::Blocked(order) => { ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec if rank < 2 { return Err(StorageError::MatrixRank); } @@@ -756,19 -774,16 +841,27 @@@ pub(crate) fn amp_micro_dimension(preci #[cfg(test)] mod tests { use super::*; ++<<<<<<< HEAD + use crate::mid::ShardExtent; + use crate::{ + AMP_COLUMN_MICRO, AMP_INNER_BLOCK, AmpOrder, BlockMajorOrder, ElementOrder, Layout, + MemoryClass, Precision, TensorTiling, TensorType, ++======= + use crate::low::{LowShardId, ShardDefinition}; + use crate::operator::Precision; + use crate::{ + AMP_COLUMN_MICRO, AMP_INNER_BLOCK, BlockedOrder, Layout, MemoryClass, NativeKernelOrder, + ShardExtent, StorageOrder, TensorTiling, TensorType, ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec + }; + use crate::{ + BlockValue, BlockValueId, ShardDefinition, ShardView, logical_view_byte_spans, + shard_storage_bytes, view_byte_spans, }; - fn shard(layout: Layout, dimensions: &[u32]) -> LowShard { - LowShard { - id: LowShardId::from_index(0), + fn shard(layout: Layout, dimensions: &[u32]) -> BlockValue { + BlockValue { + id: BlockValueId::from_index(0), tile: 0, tensor_type: TensorType::new(dimensions.iter().copied(), Precision::F16, layout), extents: dimensions diff --cc crates/ipu-codegen/src/tile.rs index 33c7a2a,0928107..0000000 --- a/crates/ipu-codegen/src/tile.rs +++ b/crates/ipu-codegen/src/tile.rs @@@ -1,10 -1,13 +1,20 @@@ //! Final lowering from logical per-tile work to address-resolved programs. use crate::{ ++<<<<<<< HEAD + BlockValueId, ExchangePatch, ExchangePhaseId, ExchangeSetupPatch, ExchangeStep, + KernelBuildPlan, LowProgram, PhysicalExchangePhase, PlacedExchangeRow, Placement, + RepeatPointer, RepeatRun, RepeatStep, StepProfile, TileAddress, TileProgram, TileStep, + TileWorkList, TileWorkRef, materialize_kernel_run, ++======= + ExchangePhaseId, KernelBuildPlan, LowProgram, LowShardId, PhysicalExchangePhase, Placement, + RepeatRun, TileWorkList, TileWorkRef, materialize_kernel_run, + }; + use ipu_target::instruction::RETURN_M10_INSTRUCTION; + use ipu_target::program::{ + CheckpointStep, ComputeStep, ExchangePatch, ExchangeSetupPatch, ExchangeStep, + PlacedExchangeRow, RepeatPointer, RepeatStep, StepProfile, TileAddress, TileProgram, TileStep, ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec }; use std::collections::BTreeMap; @@@ -701,15 -724,15 +718,24 @@@ mod tests }, ); let mid = lower(&graph, &config, &Ipu21CostModel).unwrap(); - let low = lower_to_tiles(&mid, &config).unwrap(); + let low = lower_to_tiles( + &crate::expand_tiles(&mid).unwrap(), + config.diagnostic_checkpoints, + ); let placement = place(&low).unwrap(); let kernels = KernelBuildPlan::from_program(&low).unwrap(); ++<<<<<<< HEAD + let exchanges = lower_exchanges(&low, &placement, &Topology::c600(), false) + .unwrap() + .phases; ++======= + let exchanges = lower_exchanges( + &low, + &placement, + ipu_target::hardware::HardwareTarget::Ipu21, + ) + .unwrap(); ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec let filler_tiles = random.u16(1..=4); let lowering = TileProgramLowering::new( &low, @@@ -755,15 -778,15 +781,15 @@@ let words = random.u32(1..=4_096); let bytes = words * 4; let copy = crate::LocalCopy { - source: crate::LowShardId::from_index(0), + source: crate::BlockValueId::from_index(0), source_offset: 0, - destination: crate::LowShardId::from_index(1), + destination: crate::BlockValueId::from_index(1), destination_offset: 0, bytes, - pattern: crate::LocalCopyPattern::Contiguous, + pattern: crate::CopyPattern::Contiguous, }; let (symbol, arguments) = local_copy_call(©).unwrap(); - if symbol == crate::COPY_U64_SYMBOL { + if symbol == ipu_target::emit::COPY_U64_SYMBOL { assert!(arguments[0] != 0); assert_eq!((arguments[0] * 6 + arguments[1]) * 8, bytes); assert!(arguments[1] < 6); diff --cc crates/ipu-package/src/lib.rs index e5d76b9,60010b7..0000000 --- a/crates/ipu-package/src/lib.rs +++ b/crates/ipu-package/src/lib.rs @@@ -34,10 -64,7 +64,14 @@@ pub enum ExchangeActivityKind #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub struct ProfileExchangeActivity { ++<<<<<<< HEAD + /// Number of destination tiles; zero means unknown in older profiles. + pub fanout: u16, + pub paired: bool, + pub kind: ProfileExchangeActivityKind, ++======= + pub kind: ExchangeActivityKind, ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec /// Estimated event-cycle offset within the exchange phase. pub start_cycle: u32, /// Estimated event-cycle offset within the exchange phase. @@@ -99,48 -212,7 +219,52 @@@ impl ProfileReport let mut output_sample = samples.reborrow().get(sample_index as u32); output_sample.set_start_cycle(sample.start_cycle); output_sample.set_end_cycle(sample.end_cycle); ++<<<<<<< HEAD + let mut step = output_sample.reborrow().init_step(); + step.set_local_index(sample.step.local_index); + step.set_phase(sample.step.phase); + step.set_epoch(sample.step.epoch); + step.set_operation(&sample.step.operation); + step.set_kind(match sample.step.kind { + ProfileStepKind::Exchange => profile_capnp::StepKind::Exchange, + ProfileStepKind::Compute => profile_capnp::StepKind::Compute, + ProfileStepKind::Synchronization => profile_capnp::StepKind::Synchronization, + ProfileStepKind::Idle => profile_capnp::StepKind::Idle, + }); + step.set_kernel(&sample.step.kernel); + step.set_exchange_event_cycles(sample.step.exchange_event_cycles); + let mut metadata = step + .reborrow() + .init_metadata(sample.step.metadata.len() as u32); + for (index, entry) in sample.step.metadata.iter().enumerate() { + let mut output_entry = metadata.reborrow().get(index as u32); + output_entry.set_name(&entry.name); + output_entry.set_value(&entry.value); + } + let mut activities = step + .reborrow() + .init_exchange_activities(sample.step.exchange_activities.len() as u32); + for (index, activity) in sample.step.exchange_activities.iter().enumerate() { + let mut output_activity = activities.reborrow().get(index as u32); + output_activity.set_kind(match activity.kind { + ProfileExchangeActivityKind::Send => { + profile_capnp::ExchangeActivityKind::Send + } + ProfileExchangeActivityKind::Receive => { + profile_capnp::ExchangeActivityKind::Receive + } + ProfileExchangeActivityKind::PartnerBusy => { + profile_capnp::ExchangeActivityKind::PartnerBusy + } + }); + output_activity.set_start_cycle(activity.start_cycle); + output_activity.set_end_cycle(activity.end_cycle); + output_activity.set_fanout(activity.fanout); + output_activity.set_paired(activity.paired); + } ++======= + write_profile_step(output_sample.reborrow().init_step(), &sample.step); ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec } } serialize::write_message(&mut output, &message)?; @@@ -164,57 -236,8 +288,60 @@@ .get_samples()? .iter() .map(|sample| { - let step = sample.get_step()?; Ok(CycleSample { ++<<<<<<< HEAD + step: ProfileStep { + local_index: step.get_local_index(), + phase: step.get_phase(), + epoch: step.get_epoch(), + operation: step.get_operation()?.to_str()?.into(), + kind: match step.get_kind()? { + profile_capnp::StepKind::Exchange => ProfileStepKind::Exchange, + profile_capnp::StepKind::Compute => ProfileStepKind::Compute, + profile_capnp::StepKind::Synchronization => { + ProfileStepKind::Synchronization + } + profile_capnp::StepKind::Idle => ProfileStepKind::Idle, + }, + kernel: step.get_kernel()?.to_str()?.into(), + metadata: step + .get_metadata()? + .iter() + .map(|entry| { + Ok(ProfileMetadata { + name: entry.get_name()?.to_str()?.into(), + value: entry.get_value()?.to_str()?.into(), + }) + }) + .collect::>()?, + exchange_activities: step + .get_exchange_activities()? + .iter() + .map(|activity| { + Ok(ProfileExchangeActivity { + fanout: activity.get_fanout(), + paired: activity.get_paired(), + kind: match activity.get_kind()? { + profile_capnp::ExchangeActivityKind::Send => { + ProfileExchangeActivityKind::Send + } + profile_capnp::ExchangeActivityKind::Receive => { + ProfileExchangeActivityKind::Receive + } + profile_capnp::ExchangeActivityKind::PartnerBusy => { + ProfileExchangeActivityKind::PartnerBusy + } + }, + start_cycle: activity.get_start_cycle(), + end_cycle: activity.get_end_cycle(), + }) + }) + .collect::>()?, + exchange_event_cycles: step.get_exchange_event_cycles(), + }, ++======= + step: read_profile_step(sample.get_step()?)?, ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec start_cycle: sample.get_start_cycle(), end_cycle: sample.get_end_cycle(), }) @@@ -992,50 -1036,7 +1140,54 @@@ fn write_profile_tiles output_tile.set_physical_tile(tile.physical_tile); let mut steps = output_tile.reborrow().init_steps(tile.steps.len() as u32); for (step_index, step) in tile.steps.iter().enumerate() { ++<<<<<<< HEAD + let mut output_step = steps.reborrow().get(step_index as u32); + output_step.set_local_index(step.local_index); + output_step.set_phase(step.phase); + output_step.set_epoch(step.epoch); + output_step.set_operation(&step.operation); + output_step.set_kind(match step.kind { + ProfileStepKind::Exchange => application_capnp::ProfileStepKind::Exchange, + ProfileStepKind::Compute => application_capnp::ProfileStepKind::Compute, + ProfileStepKind::Synchronization => { + application_capnp::ProfileStepKind::Synchronization + } + ProfileStepKind::Idle => application_capnp::ProfileStepKind::Idle, + }); + output_step.set_kernel(&step.kernel); + output_step.set_exchange_event_cycles(step.exchange_event_cycles); + let mut metadata = output_step + .reborrow() + .init_metadata(step.metadata.len() as u32); + for (metadata_index, entry) in step.metadata.iter().enumerate() { + let mut output_entry = metadata.reborrow().get(metadata_index as u32); + output_entry.set_name(&entry.name); + output_entry.set_value(&entry.value); + } + let mut activities = output_step + .reborrow() + .init_exchange_activities(step.exchange_activities.len() as u32); + for (activity_index, activity) in step.exchange_activities.iter().enumerate() { + let mut output_activity = activities.reborrow().get(activity_index as u32); + output_activity.set_kind(match activity.kind { + ProfileExchangeActivityKind::Send => { + application_capnp::ProfileExchangeActivityKind::Send + } + ProfileExchangeActivityKind::Receive => { + application_capnp::ProfileExchangeActivityKind::Receive + } + ProfileExchangeActivityKind::PartnerBusy => { + application_capnp::ProfileExchangeActivityKind::PartnerBusy + } + }); + output_activity.set_start_cycle(activity.start_cycle); + output_activity.set_end_cycle(activity.end_cycle); + output_activity.set_fanout(activity.fanout); + output_activity.set_paired(activity.paired); + } ++======= + write_profile_step(steps.reborrow().get(step_index as u32), step); ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec } } } @@@ -1051,61 -1052,7 +1203,65 @@@ fn read_profile_tiles steps: tile .get_steps()? .iter() ++<<<<<<< HEAD + .map(|step| { + Ok(ProfileStep { + local_index: step.get_local_index(), + phase: step.get_phase(), + epoch: step.get_epoch(), + operation: step.get_operation()?.to_str()?.into(), + kind: match step.get_kind()? { + application_capnp::ProfileStepKind::Exchange => { + ProfileStepKind::Exchange + } + application_capnp::ProfileStepKind::Compute => { + ProfileStepKind::Compute + } + application_capnp::ProfileStepKind::Synchronization => { + ProfileStepKind::Synchronization + } + application_capnp::ProfileStepKind::Idle => ProfileStepKind::Idle, + }, + kernel: step.get_kernel()?.to_str()?.into(), + metadata: step + .get_metadata()? + .iter() + .map(|entry| { + Ok(ProfileMetadata { + name: entry.get_name()?.to_str()?.into(), + value: entry.get_value()?.to_str()?.into(), + }) + }) + .collect::>()?, + exchange_activities: step + .get_exchange_activities()? + .iter() + .map(|activity| { + Ok(ProfileExchangeActivity { + fanout: activity.get_fanout(), + paired: activity.get_paired(), + kind: match activity.get_kind()? { + application_capnp::ProfileExchangeActivityKind::Send => { + ProfileExchangeActivityKind::Send + } + application_capnp::ProfileExchangeActivityKind::Receive => { + ProfileExchangeActivityKind::Receive + } + application_capnp::ProfileExchangeActivityKind::PartnerBusy => { + ProfileExchangeActivityKind::PartnerBusy + } + }, + start_cycle: activity.get_start_cycle(), + end_cycle: activity.get_end_cycle(), + }) + }) + .collect::>()?, + exchange_event_cycles: step.get_exchange_event_cycles(), + }) + }) ++======= + .map(read_profile_step) ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec .collect::>()?, }) }) @@@ -1262,16 -1209,12 +1418,24 @@@ mod tests }], exchange_activities: vec![ ProfileExchangeActivity { ++<<<<<<< HEAD + fanout: 2, + paired: true, + kind: ProfileExchangeActivityKind::Receive, ++======= + kind: ExchangeActivityKind::Receive, ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec start_cycle: 3, end_cycle: 11, }, ProfileExchangeActivity { ++<<<<<<< HEAD + fanout: 2, + paired: true, + kind: ProfileExchangeActivityKind::PartnerBusy, ++======= + kind: ExchangeActivityKind::PartnerBusy, ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec start_cycle: 1, end_cycle: 3, }, @@@ -1504,16 -1447,12 +1668,24 @@@ exchange_activities: if kind == ProfileStepKind::Exchange { vec![ ProfileExchangeActivity { ++<<<<<<< HEAD + fanout: 2, + paired: true, + kind: ProfileExchangeActivityKind::Send, ++======= + kind: ExchangeActivityKind::Send, ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec start_cycle: 2, end_cycle: 6, }, ProfileExchangeActivity { ++<<<<<<< HEAD + fanout: 2, + paired: true, + kind: ProfileExchangeActivityKind::PartnerBusy, ++======= + kind: ExchangeActivityKind::PartnerBusy, ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec start_cycle: 6, end_cycle: 8, }, diff --cc crates/ipu-target/src/exchange.rs index 251c5ac,8e747fc..0000000 --- a/crates/ipu-target/src/exchange.rs +++ b/crates/ipu-target/src/exchange.rs @@@ -390,15 -372,9 +376,15 @@@ impl PhaseProgramBuilder .tile_states .get(usize::from(tile)) .ok_or(ExchangeError::Tile(tile))?; - offset = offset.max(schedule.event_cycles); + 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) { + for (&receiver, row) in transfer.receivers.iter().zip(&transfer.plan.receivers) { let schedule = self .tile_states .get(usize::from(receiver)) @@@ -454,60 -427,28 +437,37 @@@ /// [`Self::finish`]. pub fn append_transfer_at( &mut self, - source: u16, - reserved_tiles: &[u16], - receivers: &[u16], - plan: &MulticastPlan, + transfer: &ResolvedTransfer, 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)); + let transfer_timing = self.transfer_timing_at(transfer, schedule_offset)?; + self.tile_states + .get_mut(usize::from(transfer.source)) + .ok_or(ExchangeError::Tile(transfer.source))? + .append_sender_at(&transfer.plan.sender, schedule_offset)?; - 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 + if let Some(tile) = transfer.reserved_source { + let schedule = self .tile_states ++<<<<<<< HEAD:crates/ipu-exchange/src/lib.rs + .get(usize::from(tile)) + .ok_or(ExchangeError::Tile(tile))? + .clone(); + schedule.reserved_sender_end = schedule + .reserved_sender_end + .max(transfer_timing.sender_horizon); ++======= + .get_mut(usize::from(tile)) + .ok_or(ExchangeError::Tile(tile))?; ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec:crates/ipu-target/src/exchange.rs 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)) + for (&receiver, row) in transfer.receivers.iter().zip(&transfer.plan.receivers) { + self.tile_states + .get_mut(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; + .append_receiver_at(row, schedule_offset, transfer.words)?; } Ok(transfer_timing) } @@@ -694,17 -624,30 +649,36 @@@ impl TileProgramSchedule received_words, self.receive_stream.as_ref(), )?; - let replaces_neutral = self - .receive_stream - .as_ref() - .is_some_and(|stream| timing.source_start == stream.source_end_cycles); let collision = timing.events.iter().any(|new| { ++<<<<<<< HEAD:crates/ipu-exchange/src/lib.rs + 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 first = self + .receive_events + .partition_point(|existing| existing.cycles < new.cycles); + self.receive_events[first..] + .iter() + .take_while(|existing| existing.cycles == new.cycles) + .any(|existing| { + new.cycles == existing.cycles + && !(replaces_neutral + && new.kind == ReceiveEventKind::OrdinarySource + && existing.kind == ReceiveEventKind::OrdinaryNeutral) + && !receive_events_can_share_instruction(*new, *existing) + }) ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec:crates/ipu-target/src/exchange.rs }); let sender_boundary = timing.events.iter().any(|event| { - self.senders.iter().any(|sender| { + let first = self + .senders + .partition_point(|sender| sender.start_cycles < event.cycles.saturating_sub(1)); + self.senders[first..].iter().take(2).any(|sender| { event.cycles == sender.start_cycles || event.cycles == sender.start_cycles.saturating_add(1) }) @@@ -781,16 -740,29 +771,42 @@@ { return Err(ExchangeError::Schedule("overlapping receive streams")); } ++<<<<<<< HEAD:crates/ipu-exchange/src/lib.rs + 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()); ++======= + if timing.source_start == stream.source_end_cycles { + let first = self + .receive_events + .partition_point(|event| event.cycles < stream.source_end_cycles); + let neutral = self.receive_events[first..] + .iter() + .take_while(|event| event.cycles == stream.source_end_cycles) + .position(|event| event.kind == ReceiveEventKind::OrdinaryNeutral) + .map(|position| first + position) + .filter(|_| { + stream.mode == ReceiveMode::Ordinary && base.mode == ReceiveMode::Ordinary + }); + if let Some(neutral) = neutral { + self.receive_events.remove(neutral); + } + } + } + for event in timing.events.iter().copied() { + let insertion = self + .receive_events + .partition_point(|existing| existing.cycles <= event.cycles); + self.receive_events.insert(insertion, event); + } ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec:crates/ipu-target/src/exchange.rs self.event_cycles = self.event_cycles.max(timing.horizon); self.receive_stream = Some(ReceiveStream { mode: base.mode, @@@ -3412,18 -3119,7 +3198,22 @@@ mod tests } #[test] ++<<<<<<< HEAD:crates/ipu-exchange/src/lib.rs + 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() { ++======= + fn point_to_point_matches_cpp_oracle_vectors() { ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec:crates/ipu-target/src/exchange.rs let topology = Topology::c600(); let cases = [ ( diff --cc crates/ipu-tests/src/exchange_stress.rs index 29ba6a1,19b308d..0000000 --- a/crates/ipu-tests/src/exchange_stress.rs +++ b/crates/ipu-tests/src/exchange_stress.rs @@@ -798,571 -748,9 +748,574 @@@ pub(crate) fn build }) } ++<<<<<<< HEAD +pub(crate) fn build_phase_replay( + compiled: &CompiledPackage, + phase_index: usize, + toolchain: &Toolchain, + runtime_source: &Path, +) -> Result { + let phase = compiled + .exchange_phases + .get(phase_index) + .with_context(|| format!("exchange phase {phase_index} is out of range"))?; + build_physical_phase_replay(phase, phase_index, toolchain, runtime_source) +} + +pub(crate) fn build_schedule_phase_replay( + snapshot: &ipu_codegen::ExchangeScheduleSnapshot, + phase_index: usize, + first_transfer: usize, + transfer_limit: Option, + toolchain: &Toolchain, + runtime_source: &Path, +) -> Result { + let mut problem = snapshot + .phases + .get(phase_index) + .with_context(|| format!("exchange phase {phase_index} is out of range"))? + .clone(); + if first_transfer > problem.transfers.len() { + bail!("--exchange-replay-first-transfer is beyond the selected phase"); + } + problem.transfers.drain(..first_transfer); + if let Some(limit) = transfer_limit { + if limit == 0 { + bail!("--exchange-replay-transfer-limit must be nonzero"); + } + problem.transfers.truncate(limit); + } + let scheduled = ipu_codegen::schedule_exchange_problem(snapshot.tile_count, &problem)?; + build_physical_phase_replay(&scheduled.phase, phase_index, toolchain, runtime_source) +} + +fn build_physical_phase_replay( + phase: &ipu_codegen::PhysicalExchangePhase, + phase_index: usize, + toolchain: &Toolchain, + runtime_source: &Path, +) -> Result { + let topology = Topology::c600(); + let execution_tiles = u16::try_from(topology.tile_count())?; + let scheduled_tiles = u16::try_from(phase.programs.len())?; + if phase.active.len() != usize::from(scheduled_tiles) + || phase.incoming_bases.len() != usize::from(scheduled_tiles) + || phase.activities.len() != usize::from(scheduled_tiles) + { + bail!("exchange phase {phase_index} has inconsistent per-tile metadata"); + } + if phase + .repeat_patches + .iter() + .any(|patches| !patches.is_empty()) + { + bail!( + "exchange phase {phase_index} uses repeat patches; replay a concrete iteration instead" + ); + } + + let programs = (0..execution_tiles) + .map(|tile| { + let scheduled = tile < scheduled_tiles; + let words = if scheduled { + phase.programs[usize::from(tile)].clone() + } else { + inactive_exchange_program() + }; + let row_address = ROW_BASE; + Ok(TileProgram { + tile, + steps: vec![ + TileStep::Exchange(ExchangeStep { + active: scheduled && phase.active[usize::from(tile)], + incoming_base: if scheduled { + phase.incoming_bases[usize::from(tile)] + } else { + 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: row_address, + words, + }, + setup_patch: None, + repeat_patches: Vec::new(), + profile: StepProfile::default(), + }), + // A patched breakpoint immediately following the final + // internal-exchange dispatch is not durable on IPU21. + // Cross one ordinary worker-call boundary before trapping; + // this occurs after the exchange epoch under test. + TileStep::Compute(ComputeStep { + symbol: "ipu_stack_static_worker_delay".into(), + output_address: TileAddress::Absolute(DATA_BASE), + input_addresses: vec![TileAddress::Absolute(DATA_BASE)], + arguments: vec![1], + profile: StepProfile::default(), + }), + TileStep::Checkpoint(CheckpointStep { + operation: u32::try_from(phase_index)?, + breakpoint: 0, + profile: StepProfile::default(), + }), + ], + }) + }) + .collect::>>()?; + + let mut transfers = BTreeMap::< + u32, + ( + Option<(u16, ExchangeActivity)>, + Vec<(u16, ExchangeActivity)>, + ), + >::new(); + let mut initial = vec![BTreeMap::::new(); usize::from(scheduled_tiles)]; + for (tile, activities) in phase.activities.iter().enumerate() { + let tile = u16::try_from(tile)?; + for &activity in activities { + if activity.address & 0b11 != 0 { + bail!( + "phase {phase_index} transfer {} has unaligned address 0x{:x}", + activity.transfer, + activity.address + ); + } + let transfer = transfers.entry(activity.transfer).or_default(); + match activity.kind { + ExchangeActivityKind::Send => { + if transfer.0.replace((tile, activity)).is_some() { + bail!( + "phase {phase_index} transfer {} has multiple senders", + activity.transfer + ); + } + } + ExchangeActivityKind::Receive => transfer.1.push((tile, activity)), + ExchangeActivityKind::PartnerBusy => {} + } + for word in 0..activity.words { + let address = activity + .address + .checked_add(word.checked_mul(4).context("activity offset overflow")?) + .context("activity address overflow")?; + initial[usize::from(tile)] + .entry(address) + .or_insert_with(|| replay_word(tile, address)); + } + } + } + + let transfers = transfers + .into_iter() + .map(|(id, (send, receives))| { + let (source, send) = + send.with_context(|| format!("phase {phase_index} transfer {id} has no sender"))?; + if receives.is_empty() { + bail!("phase {phase_index} transfer {id} has no receivers"); + } + if receives + .iter() + .any(|(_, receive)| receive.words != send.words) + { + bail!("phase {phase_index} transfer {id} has inconsistent word counts"); + } + Ok((id, source, ReplayTransfer { send, receives })) + }) + .collect::>>()?; + let mut events = Vec::new(); + for &(transfer, source, ref replay) in &transfers { + events.push(( + replay.send.start_cycle, + 1u8, + ReplayEvent::Send { + transfer, + tile: source, + address: replay.send.address, + words: replay.send.words, + }, + )); + for &(tile, receive) in &replay.receives { + events.push(( + receive.memory_end_cycle, + 0u8, + ReplayEvent::Receive { + transfer, + tile, + address: receive.address, + }, + )); + } + } + events.sort_unstable_by_key(|&(cycle, order, _)| (cycle, order)); + + let mut expected_memory = initial.clone(); + let mut payloads = BTreeMap::>::new(); + for (_, _, event) in events { + match event { + ReplayEvent::Send { + transfer, + tile, + address, + words, + } => { + let memory = &expected_memory[usize::from(tile)]; + let payload = (0..words) + .map(|word| { + let address = address + word * 4; + memory.get(&address).copied().with_context(|| { + format!( + "phase {phase_index} transfer {transfer} reads untracked tile {tile} address 0x{address:x}" + ) + }) + }) + .collect::>>()?; + payloads.insert(transfer, payload); + } + ReplayEvent::Receive { + transfer, + tile, + address, + } => { + let payload = payloads.get(&transfer).with_context(|| { + format!( + "phase {phase_index} transfer {transfer} completes a receive before its send snapshot" + ) + })?; + let memory = &mut expected_memory[usize::from(tile)]; + for (word, &value) in payload.iter().enumerate() { + memory.insert(address + u32::try_from(word)? * 4, value); + } + } + } + } + + let data = memory_spans(&initial) + .into_iter() + .map(|span| TileProgramData { + tile: span.tile, + address: span.address, + data: span.words.into_iter().flat_map(u32::to_le_bytes).collect(), + }) + .collect::>(); + let expected = memory_spans(&expected_memory); + let mut initial_origins = BTreeMap::>::new(); + for (tile, memory) in initial.iter().enumerate() { + for (&address, &word) in memory { + initial_origins + .entry(word) + .or_default() + .push((u16::try_from(tile)?, address)); + } + } + let application = build_tile_program_package(&programs, &data, &[], toolchain, runtime_source)?; + let overlapping_tiles = phase + .activities + .iter() + .filter(|activities| { + activities.iter().any(|send| { + send.kind == ExchangeActivityKind::Send + && activities.iter().any(|receive| { + receive.kind == ExchangeActivityKind::Receive + && send.start_cycle < receive.end_cycle + && receive.start_cycle < send.end_cycle + }) + }) + }) + .count(); + eprintln!( + "exchangeReplay phase={phase_index} transfers={} activeTiles={} overlappingTiles={} eventCycles={} touchedSpans={}", + transfers.len(), + phase.active.iter().filter(|&&active| active).count(), + overlapping_tiles, + phase.event_cycles, + expected.len(), + ); + Ok(PhaseReplayPackage { + application, + phase: phase_index, + expected, + activities: phase.activities.clone(), + initial_origins, + }) +} + +impl PhaseReplayPackage { + pub(crate) fn service_readback( + &self, + device: &Device, + sample_limit: usize, + tiles: &[u16], + serviced: &mut bool, + ) -> Result<()> { + if *serviced { + return Ok(()); + } + let topology = Topology::c600(); + for tile in &self.application.tiles { + if device.tile_context_state(u16::try_from(tile.physical_tile)?, 0)? != 2 { + return Ok(()); + } + } + for tile in &self.application.tiles { + let physical = u16::try_from(tile.physical_tile)?; + let status = device.read_tile_context_status(physical, 0)?; + let exception = TileException::from_status(status); + if exception != TileException::PatchedBreak0 { + bail!( + "exchange replay phase {} tile {physical} stopped with {exception} (status {status:#x})", + self.phase, + ); + } + } + for tile in tiles { + if !self.expected.iter().any(|span| span.tile == *tile) { + bail!( + "exchange replay phase {} has no touched words on logical tile {tile}", + self.phase + ); + } + } + let expected = self + .expected + .iter() + .filter(|span| tiles.is_empty() || tiles.contains(&span.tile)) + .collect::>(); + let samples = replay_samples(&expected, sample_limit); + let mut checked = 0usize; + for (tile, samples) in samples { + let physical = topology.physical(tile)?; + let addresses = samples + .iter() + .map(|&(address, _)| address) + .collect::>(); + let actual = device + .read_tile_words_at_addresses_from_inactive_context(physical, 1, &addresses) + .with_context(|| { + format!( + "read replay phase {} logical tile {} physical tile {} at {} sampled addresses", + self.phase, tile, physical, addresses.len(), + ) + })?; + let differences = actual + .iter() + .zip(&samples) + .filter(|(actual, (_, expected))| **actual != *expected) + .take(16) + .map(|(&actual, &(address, expected))| { + let roles = self.activities[usize::from(tile)] + .iter() + .filter(|activity| { + let end = activity.address.saturating_add(activity.words * 4); + (activity.address..end).contains(&address) + }) + .map(|activity| { + ( + activity.transfer, + activity.kind, + activity.start_cycle, + activity.end_cycle, + activity.memory_end_cycle, + ) + }) + .collect::>(); + let window = roles.iter().fold(None, |window, &(_, _, start, end, _)| { + Some(window.map_or((start, end), |(first, last): (u32, u32)| { + (first.min(start), last.max(end)) + })) + }); + let concurrent = window.map_or_else(Vec::new, |(start, end)| { + self.activities[usize::from(tile)] + .iter() + .filter(|activity| { + activity.start_cycle < end && start < activity.end_cycle + }) + .map(|activity| { + ( + activity.transfer, + activity.kind, + activity.start_cycle, + activity.end_cycle, + activity.address, + activity.words, + ) + }) + .collect::>() + }); + let transfer_ids = roles + .iter() + .map(|&(transfer, ..)| transfer) + .collect::>(); + let transfer_ids_ref = &transfer_ids; + let endpoints = self + .activities + .iter() + .enumerate() + .flat_map(|(endpoint_tile, activities)| { + activities.iter().filter_map(move |activity| { + transfer_ids_ref.contains(&activity.transfer).then_some(( + u16::try_from(endpoint_tile).unwrap(), + activity.transfer, + activity.kind, + activity.start_cycle, + activity.end_cycle, + activity.address, + activity.words, + )) + }) + }) + .collect::>(); + let endpoint_concurrent = endpoints + .iter() + .flat_map(|&(endpoint_tile, transfer, kind, start, end, _, _)| { + self.activities[usize::from(endpoint_tile)] + .iter() + .filter(move |activity| { + activity.transfer != transfer + && activity.start_cycle < end + && start < activity.end_cycle + }) + .map(move |activity| { + ( + endpoint_tile, + kind, + transfer, + activity.transfer, + activity.kind, + activity.start_cycle, + activity.end_cycle, + activity.address, + activity.words, + ) + }) + }) + .collect::>(); + ( + address, + expected, + actual, + self.initial_origins.get(&expected), + self.initial_origins.get(&actual), + roles, + concurrent, + endpoints, + endpoint_concurrent, + ) + }) + .collect::>(); + if !differences.is_empty() { + bail!( + "exchange replay phase {} corrupted logical tile {tile}: {differences:?}", + self.phase, + ); + } + checked += samples.len(); + } + eprintln!("exchangeReplay phase={} sampledWords={checked}", self.phase); + const IPU21_NOP_INSTRUCTION: u32 = 0x19e0_0000; + for tile in &self.application.tiles { + let physical = u16::try_from(tile.physical_tile)?; + let pc = device.read_tile_program_counter(physical, 0)?; + device.write_tile_word_from_stopped_context(physical, 0, pc, IPU21_NOP_INSTRUCTION)?; + device.clear_tile_exception(physical, 0)?; + } + *serviced = true; + Ok(()) + } +} + +fn replay_samples(expected: &[&ExpectedSpan], limit: usize) -> BTreeMap> { + let total_words = expected.iter().map(|span| span.words.len()).sum::(); + let wanted = limit.min(total_words); + let mut selected = std::collections::BTreeSet::<(usize, usize)>::new(); + if wanted == 0 || expected.is_empty() { + return BTreeMap::new(); + } + let first_span_samples = wanted.min(expected.len()); + for sample in 0..first_span_samples { + let span = sample * expected.len() / first_span_samples; + selected.insert((span, 0)); + } + if selected.len() < wanted && first_span_samples == expected.len() { + for (span, values) in expected.iter().enumerate() { + if selected.len() == wanted { + break; + } + selected.insert((span, values.words.len() - 1)); + } + } + let mut ends = Vec::with_capacity(expected.len()); + let mut end = 0usize; + for span in expected { + end += span.words.len(); + ends.push(end); + } + let attempts = wanted.saturating_mul(4).max(wanted); + for sample in 0..attempts { + if selected.len() == wanted { + break; + } + let linear = sample * total_words / attempts; + let span = ends.partition_point(|&end| end <= linear); + let start = span.checked_sub(1).map_or(0, |previous| ends[previous]); + selected.insert((span, linear - start)); + } + let mut result = BTreeMap::>::new(); + for (span, word) in selected { + let span = &expected[span]; + result.entry(span.tile).or_default().push(( + span.address + u32::try_from(word).expect("tile word index fits u32") * 4, + span.words[word], + )); + } + for samples in result.values_mut() { + samples.sort_unstable_by_key(|&(address, _)| address); + } + result +} + +fn memory_spans(memory: &[BTreeMap]) -> Vec { + let mut spans = Vec::new(); + for (tile, words) in memory.iter().enumerate() { + let mut current: Option = None; + for (&address, &word) in words { + let contiguous = current.as_ref().is_some_and(|span| { + span.address + u32::try_from(span.words.len()).unwrap() * 4 == address + }); + if contiguous { + current.as_mut().unwrap().words.push(word); + } else { + if let Some(span) = current.take() { + spans.push(span); + } + current = Some(ExpectedSpan { + tile: u16::try_from(tile).expect("tile count was supplied as u16"), + address, + words: vec![word], + }); + } + } + if let Some(span) = current { + spans.push(span); + } + } + spans +} + +fn replay_word(tile: u16, address: u32) -> u32 { + 0xa5a5_5a5a ^ u32::from(tile).wrapping_mul(0x9e37_79b9) ^ (address >> 2).rotate_left(13) +} + ++======= ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec impl StressPackage { pub(crate) fn live_exchange_state(&self, runtime: &Runtime) -> String { - let topology = Topology::c600(); + let topology = ipu_target::hardware::HardwareTarget::Ipu21.topology(); let mut states = Vec::new(); let mut relevant = Vec::new(); for transfer in &self.transfers { diff --cc crates/ipu-tests/src/main.rs index 0d221cb,626bc4e..0000000 --- a/crates/ipu-tests/src/main.rs +++ b/crates/ipu-tests/src/main.rs @@@ -64,37 -63,9 +65,43 @@@ struct Arguments /// Decode every active supervisor row for this exchange-stress case. #[arg(long, requires = "exchange_diagnostics")] exchange_diagnostic_case: Option, ++<<<<<<< HEAD + /// Replace the compiled workload with a tokenized replay of one exact + /// physical exchange phase and verify sampled words after execution. + #[arg(long, conflicts_with_all = ["reuse_package", "diagnostic_run"])] + exchange_replay_phase: Option, + /// Replay a phase from an exported address-resolved schedule without + /// recompiling the model. + #[arg( + long, + requires = "exchange_replay_phase", + conflicts_with = "export_exchange_schedule" + )] + replay_exchange_schedule: Option, + /// Maximum number of systematically distributed words read back by an + /// exact exchange-phase replay. + #[arg(long, default_value_t = 8192)] + exchange_replay_samples: usize, + /// Restrict replay readback to these logical tiles (comma-separated). + #[arg(long, value_delimiter = ',', requires = "exchange_replay_phase")] + exchange_replay_tiles: Vec, + /// Replay only this prefix of the selected phase's transfer list. + #[arg(long, requires = "exchange_replay_phase")] + exchange_replay_transfer_limit: Option, + /// Skip this many transfers before applying the replay transfer limit. + #[arg(long, default_value_t = 0, requires = "exchange_replay_phase")] + exchange_replay_first_transfer: usize, + /// Summarize packaged exchange rows and exit before loading hardware. + #[arg(long)] + inspect_exchanges: bool, + /// Write the address-resolved exchange input, then exit unless a phase replay is requested. + #[arg(long, conflicts_with_all = ["reuse_package", "diagnostic_run"])] + export_exchange_schedule: Option, ++======= + /// Summarize packaged exchange rows and exit before loading hardware. + #[arg(long)] + inspect_exchanges: bool, ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec /// Include complete decoded rows for one physical tile in the inspection. #[arg(long, requires = "inspect_exchanges")] inspect_exchange_tile: Option, @@@ -331,35 -304,32 +343,45 @@@ fn parse_gemm_plan_constraint(value: &s source_operation: operation .parse::() .map_err(|error| error.to_string())?, - orientation: match *orientation { - "normal" => GemmOrientation::Normal, - "swapped" => GemmOrientation::Swapped, - _ => return Err("orientation must be normal or swapped".into()), + geometry: GemmGeometry { + block: GemmBlockShape { + inner: *inner_block, + output_columns: *output_column_block, + }, + orientation, + compute: GemmGrid { + rows: *row_partitions, + columns: *column_partitions, + inner: *inner_partitions, + }, + result: GemmResultGrid { + rows: result_rows, + columns: result_columns, + }, + order: ipu_codegen::GridOrder::ColumnsFast, }, - row_partitions: *row_partitions, - column_partitions: *column_partitions, - inner_partitions: *inner_partitions, - result_row_partitions: *result_row_partitions, - result_column_partitions: *result_column_partitions, - output_column_block, + reduction_staging: Some(staging), weight_memory_class: match *memory { - "standard" => MemoryClass::Ipu21Standard, - "interleaved" => MemoryClass::Ipu21Interleaved, + "standard" => MemoryClass::Standard, + "interleaved" => MemoryClass::Interleaved, _ => return Err("memory must be standard or interleaved".into()), }, ++<<<<<<< HEAD + reduction_staging: match *reduction { + "complete" => ReductionStaging::Complete, + "streamed" => ReductionStaging::Streamed, + value if value.starts_with("batch-") => ReductionStaging::Batched( + value[6..] + .parse() + .map_err(|_| "batch size must be a nonzero u16")?, + ), + _ => return Err("reduction must be complete, streamed, or batch-N".into()), + }, ++======= ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec local_weight_staging: match *local { - "direct" => LocalOperandStaging::Direct, - "match-remote" => LocalOperandStaging::MatchRemote, + "direct" => LocalOperandStaging::Direct(MemoryClass::Interleaved), + "match-remote" => LocalOperandStaging::Staged(MemoryClass::Interleaved), _ => return Err("local staging must be direct or match-remote".into()), }, }) @@@ -672,16 -608,6 +671,19 @@@ fn main() -> Result<()> .with_automatic_input(query_weights, Precision::F16) .with_automatic_input(key_weights, Precision::F16) .with_automatic_input(value_weights, Precision::F16); ++<<<<<<< HEAD + // This benchmark compares the two F16 attention strategies. Keep + // projection precision controlled as batch size changes instead + // of allowing a different GEMM precision to confound the sweep. + pipeline.operator_candidates.retain(|candidate| { + !matches!( + candidate.plan.operator, + MidOperator::Gemm { multiply, .. } if multiply != Precision::F16 + ) + }); + pipeline.profiling = !arguments.no_profile; ++======= ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec } else { let (heads, query_rows, key_rows) = (4, 17, 19); let query_dimension = SIGLIP_ATTENTION_HEAD_DIMENSION; @@@ -719,11 -645,7 +721,11 @@@ layout: Layout::attention_output(heads, key_partitions), }, ); - // Exercise the tiled attention lowering and its explicit input - // conversions rather than the independent whole-head codelet. - pipeline.operator_candidates.clear(); pipeline.conversion_streaming = ipu_codegen::ConversionStreamingPolicy::Never; ++<<<<<<< HEAD + pipeline.profiling = !arguments.no_profile; ++======= ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec } } else if matches!(arguments.workload, Workload::SiglipMlpBenchmark) { validate_mlp_benchmark_shape( @@@ -771,7 -693,6 +773,10 @@@ )?[0] }; graph.set_outputs([output])?; ++<<<<<<< HEAD + pipeline.profiling = !arguments.no_profile; ++======= ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec pipeline = pipeline.with_automatic_input(left, Precision::F16); for weight in right0.into_iter().chain(right1) { pipeline = pipeline.with_automatic_input(weight, Precision::F16); @@@ -789,7 -710,6 +794,10 @@@ )?; let output = graph.gemm(left, right)?; graph.set_outputs([output])?; ++<<<<<<< HEAD + pipeline.profiling = !arguments.no_profile; ++======= ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec pipeline = pipeline .with_automatic_input(left, Precision::F16) .with_automatic_input(right, Precision::F16); @@@ -821,44 -736,6 +829,47 @@@ } None }; ++<<<<<<< HEAD + if let Some(path) = &arguments.export_exchange_schedule { + let compiled = compiled_package + .as_ref() + .context("--export-exchange-schedule requires a newly compiled package")?; + let output = + fs::File::create(path).with_context(|| format!("create {}", path.display()))?; + serde_json::to_writer(std::io::BufWriter::new(output), &compiled.exchange_schedule) + .with_context(|| format!("write {}", path.display()))?; + let transfers = compiled + .exchange_schedule + .phases + .iter() + .map(|phase| phase.transfers.len()) + .sum::(); + println!( + "exchangeSchedule={} tiles={} phases={} transfers={}", + path.display(), + compiled.exchange_schedule.tile_count, + compiled.exchange_schedule.phases.len(), + transfers + ); + if arguments.exchange_replay_phase.is_none() { + return Ok(()); + } + } + if let Some(phase) = arguments.exchange_replay_phase { + let compiled = compiled_package + .as_ref() + .context("--exchange-replay-phase requires a newly compiled package")?; + let replay = exchange_stress::build_phase_replay( + compiled, + phase, + &package_config.toolchain, + &package_config.runtime_source, + )?; + execute_exchange_replay(&arguments, &replay, &bootloader)?; + return Ok(()); + } ++======= ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec let application = Application::read( fs::File::open(&arguments.package) .with_context(|| format!("open {}", arguments.package.display()))?, @@@ -1015,167 -889,157 +1028,281 @@@ fn open_and_load_once Ok(runtime) } -fn retry_after_reset(sdk: &Path, mut attempt: impl FnMut() -> Result) -> Result { - match attempt() { - Ok(value) => Ok(value), - Err(error) - if error - .chain() - .any(|cause| matches!(cause.downcast_ref(), Some(DriverError::Timeout(_)))) => - { - tracing::warn!(%error, "device timed out; resetting and retrying once"); - let reset = sdk.join("bin/gc-reset"); - let status = Command::new(&reset) - .arg("-m") - .status() - .with_context(|| format!("run {} -m", reset.display()))?; - if !status.success() { - bail!("{} -m exited with {status}", reset.display()); - } - attempt().context("hardware execution failed after gc-reset -m") +fn lock_device(arguments: &Arguments) -> Result> { + arguments + .device_lock + .as_ref() + .map(|path| { + let file = fs::OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(path) + .with_context(|| format!("open device lock {}", path.display()))?; + tracing::info!(path = %path.display(), "waiting for hardware ownership"); + file.lock().context("acquire hardware ownership")?; + tracing::info!(path = %path.display(), "acquired hardware ownership"); + Ok(file) + }) + .transpose() +} + ++<<<<<<< HEAD +fn execute_exchange_replay( + arguments: &Arguments, + replay: &exchange_stress::PhaseReplayPackage, + bootloader: &Path, +) -> Result<()> { + write_package(&replay.application, &arguments.package)?; + let configuration = fs::read(&arguments.configuration) + .with_context(|| format!("read {}", arguments.configuration.display()))?; + let bootloader_bytes = + fs::read(bootloader).with_context(|| format!("read {}", bootloader.display()))?; + { + let _device_lock = lock_device(arguments)?; + let runtime = open_and_load_once( + &arguments.device, + &configuration, + &replay.application, + &bootloader_bytes, + replay.application.host_exchange.startup_mark, + )?; + let mut session = runtime.host_session(&replay.application)?; + session.start()?; + let mut serviced = false; + let executed = session + .invoke_streaming_deferred_with_poll("run", &[0; 4], |device| { + replay + .service_readback( + device, + arguments.exchange_replay_samples, + &arguments.exchange_replay_tiles, + &mut serviced, + ) + .map_err(|error| DriverError::Invalid(error.to_string())) + }) + .inspect_err(|error| { + tracing::error!( + %error, + device = %device_failure_diagnostics(&runtime, &replay.application), + "exchange replay failed" + ); + })?; + if !serviced { + bail!("exchange replay completed without reaching its readback trap"); } - Err(error) => Err(error), + let _ = session.collect(&executed)?; + diagnose_completion( + &runtime, + &replay.application, + Duration::from_secs(arguments.timeout_seconds), + )?; } + println!( + "package={} exchangeReplayPhase={} hardwareTest=PASS", + arguments.package.display(), + replay.phase + ); + Ok(()) } ++======= ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec fn run_gemm( runtime: &Runtime, application: &Application, - active_tiles: u16, - batch: u32, + package: &CompiledPackage, timeout_seconds: u64, ) -> Result<()> { ++<<<<<<< HEAD + // Generate logical tensors; placement, padding and element order come from + // the compiler's storage maps, never from the binding slice index. + let mut values = std::collections::BTreeMap::new(); + for tensor in &package.inputs { + let data = (0..tensor.shape.elements()) + .map(|index| { + let index = index as u32; + let bits = if tensor.name.as_deref() == Some("left") { + let row = index / 64; + let batch = row / tensor.shape.0[1]; + let row = row % tensor.shape.0[1]; + if index % 64 == (batch * 7 + row) % 64 { + 0x3c00 + } else { + 0 + } + } else { + gemm_right_value(index / tensor.shape.0[2], index % tensor.shape.0[2]) + }; + half_to_f32(bits) + }) + .collect(); + values.insert( + tensor.value, + diagnostic::HostTensor { + shape: tensor.shape.0.clone(), + values: data, + }, + ); ++======= + let left = application + .inputs + .iter() + .find(|binding| binding.name == "left") + .cloned() + .context("GEMM package has no left input binding")?; + let right = application + .weights + .iter() + .find(|binding| binding.name == "right") + .cloned() + .context("GEMM package has no right weight binding")?; + let left_bytes = packed_binding(&left, |logical_tile, linear, elements| { + let (batch_index, inner) = amp_matrix_coordinates( + NativeKernelOrder::Left, + Precision::F16, + batch, + elements / batch, + linear, + )?; + let selected_inner = (batch_index * 7 + u32::from(logical_tile)) % 64; + Ok(if selected_inner == inner { 0x3c00 } else { 0 }) + })?; + let right_bytes = packed_binding(&right, |logical_tile, linear, elements| { + let (inner, column) = block_major_matrix_coordinates( + BlockedOrder::matrix(64, 16), + Precision::F16, + 64, + elements / 64, + linear, + )?; + Ok(gemm_right_value( + inner, + u32::from(logical_tile) * 64 + column, + )) + })?; + if left.slices.len() != usize::from(active_tiles) + || right.slices.len() != usize::from(active_tiles) + { + bail!("GEMM bindings do not cover every active tile"); ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec } - let output = run_initialized_program( - runtime, - application, - &right_bytes, - &left_bytes, - timeout_seconds, - )?; - verify_gemm_output(application, active_tiles, batch, &output) + let (weights, inputs) = diagnostic::pack_inputs(application, &package.inputs, &values)?; + let bytes = run_initialized_program(runtime, application, &weights, &inputs, timeout_seconds)?; + let tensor = package + .outputs + .iter() + .find(|tensor| tensor.name.as_deref() == Some("output.0")) + .context("GEMM package has no logical output storage map")?; + let rows = tensor.shape.0[1]; + let columns = tensor.shape.0[2]; + let expected = (0..tensor.shape.elements()) + .map(|index| { + let index = index as u32; + let row = index / columns; + half_to_f32(gemm_right_value( + (row / rows * 7 + row % rows) % 64, + index % columns, + )) + }) + .collect::>(); + verify_logical_f16_output(application, tensor, &bytes, &expected, (0.0, 0.0))?; + println!("gemmNumericalChecks={} numericalTest=PASS", expected.len()); + Ok(()) } -fn run_mlp_chain( +fn run_reference( runtime: &Runtime, application: &Application, - active_tiles: u16, + graph: &ComputeGraph, + package: &CompiledPackage, timeout_seconds: u64, ++<<<<<<< HEAD + tolerance: (f32, f32), +) -> Result<(Vec, f32)> { + let (host_inputs, weights, inputs) = + diagnostic::prepare_inputs(graph, application, &package.inputs)?; + let references = diagnostic::evaluate(graph, host_inputs, &package.precisions)?; + let output = run_initialized_program(runtime, application, &weights, &inputs, timeout_seconds)?; + let tensor = package + .outputs + .iter() + .find(|tensor| tensor.name.as_deref() == Some("output.0")) + .context("package has no logical output storage map")?; + let expected = references + .get(&tensor.value) + .context("host reference has no graph output")?; + let maximum_error = + verify_logical_f16_output(application, tensor, &output, &expected.values, tolerance)?; + Ok((output, maximum_error)) ++======= + ) -> Result<()> { + let binding = |name: &str, bindings: &[Binding]| { + bindings + .iter() + .find(|binding| binding.name == name) + .cloned() + .with_context(|| format!("MLP package has no {name} binding")) + }; + let left = binding("left", &application.inputs)?; + let right0 = binding("right.0", &application.weights)?; + let right1 = binding("right.1", &application.weights)?; + let left_bytes = packed_binding(&left, |logical_tile, linear, elements| { + let (_, inner) = + amp_matrix_coordinates(NativeKernelOrder::Left, Precision::F16, 1, elements, linear)?; + Ok(mlp_smoke_value( + MLP_INPUT_SEED, + u64::from(logical_tile) * u64::from(MLP_SMOKE_WIDTH) + u64::from(inner), + MLP_INPUT_STANDARD_DEVIATION, + )) + })?; + let right0_bytes = packed_binding(&right0, |logical_tile, linear, elements| { + let (inner, column) = block_major_matrix_coordinates( + BlockedOrder::matrix(64, 16), + Precision::F16, + 64, + elements / 64, + linear, + )?; + let column = u32::from(logical_tile) * 64 + column; + Ok(if column < 64 { + mlp_smoke_value( + MLP_FIRST_WEIGHT_SEED, + u64::from(inner) * u64::from(MLP_SMOKE_WIDTH) + u64::from(column), + MLP_WEIGHT_STANDARD_DEVIATION, + ) + } else { + 0 + }) + })?; + let right1_bytes = packed_binding(&right1, |logical_tile, linear, elements| { + let (inner, column) = block_major_matrix_coordinates( + BlockedOrder::matrix(64, 16), + Precision::F16, + 64, + elements / 64, + linear, + )?; + let column = u32::from(logical_tile) * 64 + column; + Ok(if column < 64 { + mlp_smoke_value( + MLP_SECOND_WEIGHT_SEED, + u64::from(inner) * u64::from(MLP_SMOKE_WIDTH) + u64::from(column), + MLP_WEIGHT_STANDARD_DEVIATION, + ) + } else { + 0 + }) + })?; + let mut weights = Vec::with_capacity(right0_bytes.len() + right1_bytes.len()); + weights.extend_from_slice(&right0_bytes); + weights.extend_from_slice(&right1_bytes); + + let output = + run_initialized_program(runtime, application, &weights, &left_bytes, timeout_seconds)?; + verify_mlp_output(application, active_tiles, &output) ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec } fn run_projected_attention_benchmark( @@@ -1618,15 -1487,24 +1750,36 @@@ fn run_siglip_mlp_benchmark if clock_hz == 0 { bail!("benchmark clock must be nonzero"); } ++<<<<<<< HEAD + let (output, maximum_absolute_error) = run_reference( + runtime, + application, + graph, + package, + timeout_seconds, + (0.03, 0.05), + )?; + if !profiling_enabled { ++======= + let (host_inputs, weights, left_bytes) = + diagnostic::prepare_inputs(graph, application, &package.inputs)?; + let references = diagnostic::evaluate(graph, host_inputs, &package.precisions)?; + + let output = + run_initialized_program(runtime, application, &weights, &left_bytes, timeout_seconds)?; + + let output_metadata = package + .outputs + .iter() + .find(|tensor| tensor.name.as_deref() == Some("output.0")) + .context("MLP benchmark package has no logical output storage map")?; + let expected = references + .get(&output_metadata.value) + .context("MLP host reference has no graph output")?; + let maximum_absolute_error = + verify_logical_f16_output(application, output_metadata, &output, &expected.values)?; + if !profiling.records_overall_time() { ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec println!( "workload=siglip-mlp-f16-b{batch}-t{tokens}-d{dimension}-h{hidden_dimension}-n{blocks} benchmark=siglip-mlp-f16 batch={batch} tokens={tokens} dimension={dimension} hiddenDimension={hidden_dimension} blocks={blocks} biases=false profiling=false maximumAbsoluteError={maximum_absolute_error:.6}" ); @@@ -1764,10 -1644,9 +1919,10 @@@ fn verify_benchmark_output(application fn verify_logical_f16_output( application: &Application, - tensor: &DiagnosticTensor, + tensor: &CompiledTensor, bytes: &[u8], expected: &[f32], + tolerance: (f32, f32), ) -> Result { let (binding, base) = output_binding(application, "output.0")?; if tensor.precision != Precision::F16 @@@ -1890,6 -1769,67 +2045,70 @@@ fn filled_f16_binding(binding: &Binding Ok(bytes) } ++<<<<<<< HEAD ++======= + fn verify_mlp_output(application: &Application, active_tiles: u16, bytes: &[u8]) -> Result<()> { + let output = application + .outputs + .iter() + .find(|binding| binding.name == "output.0") + .context("MLP package has no output binding")?; + let expected_bytes = output + .slices + .iter() + .map(|slice| slice.file_offset + slice.size) + .max() + .context("MLP output has no slices")?; + if bytes.len() != usize::try_from(expected_bytes)? { + bail!( + "MLP returned {} bytes, expected {expected_bytes}", + bytes.len() + ); + } + let mut maximum_error = 0.0f32; + let mut mismatches = Vec::new(); + let mut checked = 0usize; + if output.slices.len() != usize::from(active_tiles) { + bail!("MLP output does not cover every logical tile"); + } + for (row, slice) in output.slices.iter().enumerate() { + let row = u16::try_from(row)?; + let expected_row = mlp_smoke_reference(row); + let elements = u32::try_from(slice.size / 2)?; + for linear in 0..elements { + let (_, column) = amp_matrix_coordinates( + NativeKernelOrder::Left, + Precision::F16, + 1, + elements, + linear, + )?; + if column >= 64 { + continue; + } + let offset = usize::try_from(slice.file_offset + u64::from(linear) * 2)?; + let actual = half_to_f32(u16::from_le_bytes( + bytes[offset..offset + 2].try_into().unwrap(), + )); + let expected = expected_row[column as usize]; + let error = (actual - expected).abs(); + maximum_error = maximum_error.max(error); + checked += 1; + if error > 0.02 && mismatches.len() < 16 { + mismatches.push((row, column, expected, actual, error)); + } + } + } + if !mismatches.is_empty() { + bail!("MLP numerical verification failed after {checked} checks: {mismatches:?}"); + } + println!( + "mlpNumericalChecks={checked} maximumAbsoluteError={maximum_error:.6} numericalTest=PASS" + ); + Ok(()) + } + ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec fn gelu_reference(value: f32) -> f32 { 0.5 * value * (1.0 + (0.797_884_6 * (value + 0.044_715 * value.powi(3))).tanh()) } @@@ -1940,6 -1922,62 +2159,65 @@@ fn packed_binding Ok(bytes) } ++<<<<<<< HEAD ++======= + fn verify_gemm_output( + application: &Application, + active_tiles: u16, + batch: u32, + bytes: &[u8], + ) -> Result<()> { + let output = application + .outputs + .iter() + .find(|binding| binding.name == "output.0") + .context("GEMM package has no output binding")?; + let expected_bytes = output + .slices + .iter() + .map(|slice| slice.file_offset + slice.size) + .max() + .context("GEMM output has no slices")?; + if bytes.len() != usize::try_from(expected_bytes)? { + bail!( + "GEMM returned {} bytes, expected {expected_bytes}", + bytes.len() + ); + } + let mut mismatches = Vec::new(); + let mut checked = 0usize; + if output.slices.len() != usize::from(active_tiles) { + bail!("GEMM output does not cover every logical tile"); + } + for (row, slice) in output.slices.iter().enumerate() { + let row = u16::try_from(row)?; + let elements = u32::try_from(slice.size / 2)?; + for linear in 0..elements { + let (batch_index, column) = amp_matrix_coordinates( + NativeKernelOrder::Output, + Precision::F16, + batch, + u32::from(active_tiles) * 64, + linear, + )?; + let offset = usize::try_from(slice.file_offset + u64::from(linear) * 2)?; + let actual = u16::from_le_bytes(bytes[offset..offset + 2].try_into().unwrap()); + let selected_inner = (batch_index * 7 + u32::from(row)) % 64; + let expected = gemm_right_value(selected_inner, column); + checked += 1; + if actual != expected && mismatches.len() < 16 { + mismatches.push((row, batch_index, column, expected, actual)); + } + } + } + if !mismatches.is_empty() { + bail!("GEMM numerical verification failed after {checked} checks: {mismatches:?}"); + } + println!("gemmNumericalChecks={checked} numericalTest=PASS"); + Ok(()) + } + ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec fn gemm_right_value(inner: u32, column: u32) -> u16 { const VALUES: [u16; 5] = [0xbc00, 0xb800, 0x0000, 0x3800, 0x3c00]; VALUES[((inner * 3 + column) % VALUES.len() as u32) as usize] diff --cc device/attention_stages_f16.S index b89871c,47a742a..0000000 --- a/device/attention_stages_f16.S +++ b/device/attention_stages_f16.S @@@ -77,7 -110,14 +95,18 @@@ .size \name, .-\name .endm ++<<<<<<< HEAD + ATTENTION_SOFTMAX_ENTRY ATTENTION_SOFTMAX_SYMBOL ++======= + #ifdef ATTENTION_BUILD_ASSEMBLY_SOFTMAX_SMALL_KEY + ATTENTION_SOFTMAX_ENTRY ATTENTION_SOFTMAX_SMALL_QUERY_SMALL_KEY_SYMBOL, ATTENTION_SMALL_QUERY_ROWS, ATTENTION_SMALL_KEY_ROWS + ATTENTION_SOFTMAX_ENTRY ATTENTION_SOFTMAX_LARGE_QUERY_SMALL_KEY_SYMBOL, ATTENTION_LARGE_QUERY_ROWS, ATTENTION_SMALL_KEY_ROWS + #endif + #ifdef ATTENTION_BUILD_ASSEMBLY_SOFTMAX_LARGE_KEY + ATTENTION_SOFTMAX_ENTRY ATTENTION_SOFTMAX_SMALL_QUERY_LARGE_KEY_SYMBOL, ATTENTION_SMALL_QUERY_ROWS, ATTENTION_LARGE_KEY_ROWS + ATTENTION_SOFTMAX_ENTRY ATTENTION_SOFTMAX_LARGE_QUERY_LARGE_KEY_SYMBOL, ATTENTION_LARGE_QUERY_ROWS, ATTENTION_LARGE_KEY_ROWS + #endif ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec .macro ATTENTION_SOFTMAX_MAX_PAIR ld32 $a0, $m3, $m15, 0 @@@ -698,7 -738,8 +727,12 @@@ .size \name, .-\name .endm ++<<<<<<< HEAD + ATTENTION_MERGE_ENTRY ATTENTION_MERGE_SYMBOL ++======= + ATTENTION_MERGE_ENTRY ATTENTION_MERGE_SMALL_QUERY_SYMBOL, ATTENTION_SMALL_QUERY_ROWS + ATTENTION_MERGE_ENTRY ATTENTION_MERGE_LARGE_QUERY_SYMBOL, ATTENTION_LARGE_QUERY_ROWS ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec .macro ATTENTION_MERGE_INITIAL_PAIR ld32 $a0, $m1, $m15, 0 diff --cc schemas/application.capnp index 021d20f,75b10f3..0000000 --- a/schemas/application.capnp +++ b/schemas/application.capnp @@@ -72,47 -74,9 +74,50 @@@ struct DeviceConfigWrite value @1 :UInt32; } ++<<<<<<< HEAD +struct ProfileMetadata { + name @0 :Text; + value @1 :Text; +} + +enum ProfileStepKind { + exchange @0; + compute @1; + synchronization @2; + idle @3; +} + +enum ProfileExchangeActivityKind { + send @0; + receive @1; + partnerBusy @2; +} + +struct ProfileExchangeActivity { + kind @0 :ProfileExchangeActivityKind; + startCycle @1 :UInt32; + endCycle @2 :UInt32; + fanout @3 :UInt16; # Zero means unavailable in older profiles. + paired @4 :Bool; +} + +struct ProfileStepPlan { + localIndex @0 :UInt32; + phase @1 :UInt32; + epoch @2 :UInt32; + operation @3 :Text; + kind @4 :ProfileStepKind; + kernel @5 :Text; + metadata @6 :List(ProfileMetadata); + exchangeActivities @7 :List(ProfileExchangeActivity); + exchangeEventCycles @8 :UInt32; +} + ++======= ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec struct TileProfilePlan { physicalTile @0 :UInt32; - steps @1 :List(ProfileStepPlan); + steps @1 :List(Profile.Step); } struct DebugSymbol { diff --cc schemas/profile.capnp index 6ed9858,6b2f7e0..0000000 --- a/schemas/profile.capnp +++ b/schemas/profile.capnp @@@ -1,45 -1,9 +1,49 @@@ @0xbadb9c24f7721fa3; ++<<<<<<< HEAD +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; + fanout @3 :UInt16; # Zero means unavailable in older profiles. + paired @4 :Bool; +} + +struct ProfileStep { + localIndex @0 :UInt32; + phase @1 :UInt32; + epoch @2 :UInt32; + operation @3 :Text; + kind @4 :StepKind; + kernel @5 :Text; + metadata @6 :List(MetadataEntry); + exchangeActivities @7 :List(ExchangeActivity); + exchangeEventCycles @8 :UInt32; +} + +struct MetadataEntry { + name @0 :Text; + value @1 :Text; +} ++======= + using Common = import "profile_common.capnp"; ++>>>>>>> bf5aaa90ef308b26250ce2474e86806a3bd498ec struct CycleSample { - step @0 :ProfileStep; + step @0 :Common.Step; startCycle @1 :UInt32; endCycle @2 :UInt32; } * Unmerged path crates/ipu-codegen/src/cost/kernel.rs * Unmerged path device/attention_softmax_f16_wrapper.S