diff --git a/crates/ipu-codegen/src/planner/gemm_tests.rs b/crates/ipu-codegen/src/planner/gemm_tests.rs index 65475f1f..badc0f3d 100644 --- a/crates/ipu-codegen/src/planner/gemm_tests.rs +++ b/crates/ipu-codegen/src/planner/gemm_tests.rs @@ -4,6 +4,84 @@ use super::*; use crate::planner::tests::gelu; use ipu_target::Target; +#[test] +#[ignore] +fn trace_fp16_regression() { + use super::super::search::Search; + let mut high=HighGraph::new(); + let x=high.host_input("x",[1,729,1152]).unwrap(); + let w1=high.parameter("up",[1,1152,4304]).unwrap(); + let w2=high.parameter("down",[1,4304,1152]).unwrap(); + let h=high.gemm(x,w1).unwrap(); + let a=high.gelu(h).unwrap(); + let y=high.gemm(a,w2).unwrap(); + high.set_outputs([y]).unwrap(); + let config=PipelineConfig::new(Target::Ipu21,1472).with_input(x,TensorFormat{precision:Precision::F16,layout:Layout::logical_linear(1472,4)}); + let boundaries=boundary_layouts(&high,&config); + let mut old=Vec::new(); + for (position,k,n,kw,nw,rows,inner,columns) in [(0,1152,4304,144,288,12,8,15),(2,4304,1152,160,192,9,27,6)] { + let op=&high.operations()[position]; + let lhs=TensorType::new([1,729,k],Precision::F16,Layout::logical_linear(1472,4)); + let mut lhs=lhs; + if position==2 {lhs.format.layout.tiling.axes[0].padding=Padding::Zero;} + let rhs=TensorType{shape:TensorShape(vec![1,k,n]),format:parameter_format(&TensorShape(vec![1,k,n]),false,false,Precision::F16,&config)}; + let shape=TensorShape(vec![1,729,n]); + let assignments=choices([&lhs,&rhs],&shape,GemmOptions::default(),1472).unwrap(); + let choice=assignments.iter().find(|c|!c.swapped && c.inner_block==kw && c.output_columns==nw + && c.operands[0].format.layout.tiling.axes[0].partitions==rows + && c.operands[0].format.layout.tiling.axes[0].tile_stride==Some(inner*columns) + && c.operands[0].format.layout.tiling.axes[1].tile_stride==Some(columns) + && c.operands[1].format.layout.memory_class==MemoryClass::Ipu21Interleaved + && c.reduction.as_ref().is_some_and(|r|r.order==ElementOrder::Amp(AmpOrder::Left)&&r.tiling.axes[0].partitions==rows*inner&&r.tiling.axes[1].partitions==columns)).unwrap(); + let operands=LiveValues::from([(op.inputs[0],BoundaryValue{tensor:lhs,owners:OwnerMap::default()}),(op.inputs[1],BoundaryValue{tensor:rhs,owners:OwnerMap::default()})]); + let mut candidate=Candidate::inputs(&high,&operands,1472,position+1); + let mut layout=Layout::logical_linear(if position==0 {1472}else{1469},4); + layout.tiling.axes[0].padding=Padding::Zero; + let output=append(&mut candidate.graph,[candidate.bindings[&op.inputs[0]],candidate.bindings[&op.inputs[1]]],choice,GemmOptions::default(),&shape,op.id,op.results[0],Some(&layout)); + candidate.bindings.insert(op.results[0],output); candidate.graph.outputs=vec![output]; + candidate.graph.compose_copies(); + for (i,op) in candidate.graph.operations.iter().enumerate() { + if let MidOperationKind::Copy{mapping,policy,..}=&op.kind { + let cost=crate::estimate::operation_cost(Target::Ipu21,op,&candidate.graph.values,1472).unwrap().0; + if let Some(t)=crate::estimate::conversion_traffic(Target::Ipu21,&candidate.graph.values[op.inputs[0].index() as usize],&candidate.graph.values[op.results[0].index() as usize],mapping,1472) { + let (bytes,calls)=if *policy==crate::CopyPolicy::DirectRetile {(t.maximum_local_bytes,t.maximum_local_intersections)}else{(2*t.maximum_destination_bytes,t.maximum_intersections)}; + let local=bytes.div_ceil(8)+288*calls; + eprintln!("COPY_COMPONENTS old{position} op{i} total={} exchange={} generic_local={local} transform={} bytes={bytes} calls={calls} actual_local_bytes={} local_intersections={}",cost.total,cost.exchange,cost.total-cost.exchange-local,t.maximum_local_bytes,t.maximum_local_intersections); + } + } + } + old.push((position,candidate)); + } + if std::env::var_os("TRACE_OLD_ONLY").is_some() {return;} + *super::super::search::WATCH.lock().unwrap()=old.iter().map(|(_,c)|c.graph.clone()).collect(); + let mut catalogue=super::super::candidates::catalogue(&high,&boundaries,&config).unwrap(); + for (p,c) in &old { eprintln!("EXACT_RETAINED {p} {}",catalogue[*p].iter().filter(|v|v.graph.values==c.graph.values && v.graph.operations==c.graph.operations).count()); } + if std::env::var_os("TRACE_CATALOGUE_ONLY").is_some() { return; } + for mode in ["normal","injected","old-only"] { + if mode=="injected" {for (p,c) in &old {catalogue[*p].push(c.clone());}} + if mode=="old-only" {for (p,c) in &old {catalogue[*p]=vec![c.clone()];}} + let mut search=Search::new(&high,&boundaries,&config,SearchLimits::default()).unwrap(); + for (position,candidates) in catalogue.iter().enumerate() { + for state in search.take_states(position) {search.extend(position,&state,candidates).unwrap();} + } + let graph=search.finish().unwrap(); + eprintln!("RESULT {mode} estimated={} memory={:?}",graph.estimated_cycles,graph.peak_memory); + for op in &graph.operations { + eprintln!("RESULT_OP {mode} {:?} {:?}",op.kind,crate::estimate::operation_cost(Target::Ipu21,op,&graph.values,1472)); + } + let low=crate::low::expand::expand_tiles(Target::Ipu21,&graph,false).unwrap(); + eprintln!("LOW_COST {mode} {:?}",crate::estimate::program_cycles(Target::Ipu21,&low,None)); + let mut previous=crate::estimate::ProgramCycles::default(); + for end in 1..=graph.operations.len() { + let mut prefix=graph.clone(); prefix.operations.truncate(end); prefix.outputs=prefix.operations[end-1].results.clone(); + let low=crate::low::expand::expand_tiles(Target::Ipu21,&prefix,false).unwrap(); + let cost=crate::estimate::program_cycles(Target::Ipu21,&low,None).unwrap(); + eprintln!("PREFIX {mode} {} delta={} exchange={} {:?}",end-1,cost.total as i64-previous.total as i64,cost.exchange as i64-previous.exchange as i64,graph.operations[end-1].kind); + previous=cost; + } + } +} + #[test] #[ignore = "large catalogue measurement; run explicitly in release mode"] fn report_mlp_boundary_diversity() { diff --git a/crates/ipu-codegen/src/planner/search.rs b/crates/ipu-codegen/src/planner/search.rs index c810e983..b93d3b28 100644 --- a/crates/ipu-codegen/src/planner/search.rs +++ b/crates/ipu-codegen/src/planner/search.rs @@ -14,6 +14,13 @@ use crate::tensor::{OwnerMap, TensorType}; use std::collections::{BTreeMap, HashMap}; use std::rc::Rc; +#[cfg(test)] +pub(super) static WATCH: std::sync::Mutex> = std::sync::Mutex::new(Vec::new()); +#[cfg(test)] +pub(super) fn trace_old(c: &Candidate) -> bool { + WATCH.lock().unwrap().iter().any(|g| g.values==c.graph.values && g.operations==c.graph.operations) +} + /// Heuristic shortlist for one fixed boundary. Keep memory/cycle tradeoffs; /// unlike `prune`, this makes no contextual dominance guarantee. Bounding it /// prevents detailed costing from becoming exhaustive again. @@ -22,11 +29,15 @@ pub(super) fn shortlist( settings: &PipelineConfig, ) -> Vec { let mut frontier: Vec = Vec::new(); - for mut candidate in candidates { + for (ordinal, mut candidate) in candidates.into_iter().enumerate() { + #[cfg(test)] + let watched = trace_old(&candidate); candidate.graph.compose_copies(); let Some((cycles, peak)) = crate::estimate::coarse::analyze(settings.target, &candidate.graph) else { + #[cfg(test)] + if watched { eprintln!("COARSE_UNAVAILABLE {:?}",candidate.graph.operations); } continue; }; candidate.graph.estimated_cycles = cycles.total; @@ -39,8 +50,12 @@ pub(super) fn shortlist( && a.graph.peak_memory.standard <= b.graph.peak_memory.standard }; if frontier.iter().any(|old| dominates(old, &candidate)) { + #[cfg(test)] + if trace_old(&candidate) { let winner=frontier.iter().find(|c|dominates(c,&candidate)).unwrap(); eprintln!("COARSE_DROP old={} {:?} winner={} {:?} {:?}",candidate.graph.estimated_cycles,candidate.graph.peak_memory,winner.graph.estimated_cycles,winner.graph.peak_memory,winner.graph.operations.iter().find(|o|matches!(o.kind,crate::MidOperationKind::Gemm{..})).map(|o|&o.kind)); } continue; } + #[cfg(test)] + for old in &frontier { if trace_old(old) && dominates(&candidate,old) { eprintln!("COARSE_REPLACED old={} {:?} winner={} {:?} {:?}",old.graph.estimated_cycles,old.graph.peak_memory,candidate.graph.estimated_cycles,candidate.graph.peak_memory,candidate.graph.operations.iter().find(|o|matches!(o.kind,crate::MidOperationKind::Gemm{..})).map(|o|&o.kind)); } } frontier.retain(|old| !dominates(&candidate, old)); frontier.push(candidate); if frontier.len() > 32 { @@ -73,6 +88,8 @@ pub(super) fn shortlist( frontier[i + 1].graph.estimated_cycles - frontier[i - 1].graph.estimated_cycles }) .unwrap(); + #[cfg(test)] + if trace_old(&frontier[remove]) { eprintln!("COARSE_CAP {} {:?} ordinal={ordinal} neighbors={:?}",frontier[remove].graph.estimated_cycles,frontier[remove].graph.peak_memory,frontier[remove-1..=remove+1].iter().map(|c|(c.graph.estimated_cycles,c.graph.peak_memory)).collect::>()); } frontier.remove(remove); } }