diff --git a/crates/ipu-codegen/src/planner/gemm_tests.rs b/crates/ipu-codegen/src/planner/gemm_tests.rs index 65475f1f..f300c050 100644 --- a/crates/ipu-codegen/src/planner/gemm_tests.rs +++ b/crates/ipu-codegen/src/planner/gemm_tests.rs @@ -4,6 +4,68 @@ 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 catalogue=super::super::candidates::catalogue(&high,&boundaries,&config).unwrap(); + 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(); + let found=catalogue[position].iter().filter(|c|c.graph.values==candidate.graph.values && c.graph.operations==candidate.graph.operations).count(); + eprintln!("OLD_FRAGMENT position={position} retained={found} coarse={:?} detailed={:?}",crate::estimate::coarse::analyze(Target::Ipu21,&candidate.graph),crate::estimate::analyze_mid(Target::Ipu21,&candidate.graph,&Default::default())); + for c in &catalogue[position] { + let gemm=c.graph.operations.iter().find(|o|matches!(o.kind,MidOperationKind::Gemm{..})).unwrap(); + eprintln!("CAT position={position} cost={} peak={:?} kernel={:?} input={:?} output={:?}",c.graph.estimated_cycles,c.graph.peak_memory,gemm.kind,c.graph.values[gemm.inputs[0].index() as usize].tensor_type.format.layout,c.graph.values[c.graph.outputs[0].index() as usize].tensor_type.format.layout); + } + old.push((position,candidate)); + } + 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)); + } +} + #[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..d60cbab4 100644 --- a/crates/ipu-codegen/src/planner/search.rs +++ b/crates/ipu-codegen/src/planner/search.rs @@ -14,6 +14,25 @@ use crate::tensor::{OwnerMap, TensorType}; use std::collections::{BTreeMap, HashMap}; use std::rc::Rc; +#[cfg(test)] +fn trace_old(c: &Candidate) -> bool { + use crate::{MidOperationKind, TensorAxis, ElementOrder, AmpOrder, BlockMajorOrder}; + let Some(op)=c.graph.operations.iter().find(|op|matches!(op.kind,MidOperationKind::Gemm{..})) else{return false;}; + let MidOperationKind::Gemm{axes,inner_block,output_columns,..}=op.kind else{return false;}; + let (rows,inner,columns)=match (inner_block,output_columns){(144,288)=>(12,8,15),(160,192)=>(9,27,6),_=>return false}; + let input=&c.graph.values[op.inputs[0].index() as usize].tensor_type.format.layout; + axes.output_column==TensorAxis::FromEnd(1) + && input.tiling.axes[0].partitions==rows + && input.tiling.axes[0].tile_stride==Some(inner*columns) + && input.tiling.axes[1].tile_stride==Some(columns) + && c.graph.values[c.graph.outputs[0].index() as usize].tensor_type.format.layout.tiling.linear_grain().is_some() + && c.graph.inputs.iter().any(|i| i.kind==GraphInputKind::Parameter && c.graph.values[i.value.index() as usize].tensor_type.format.layout.order==ElementOrder::BlockMajor(BlockMajorOrder::Matrix{row_block:16,column_block:16})) + && c.graph.operations.iter().any(|o| matches!(o.kind,MidOperationKind::ReductionSum{..}) && { + let r=&c.graph.values[o.results[0].index() as usize].tensor_type.format.layout; + r.order==ElementOrder::Amp(AmpOrder::Left) && r.tiling.axes[0].partitions==rows*inner && r.tiling.axes[1].partitions==columns + }) +} + /// 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. @@ -39,8 +58,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 {