--- a/crates/ipu-codegen/src/mid/packing.rs +++ b/crates/ipu-codegen/src/mid/packing.rs @@ -175,6 +175,9 @@ #[test] fn distributed_panels_retile_without_unpacking_the_packed_intermediate() { for view in [false, true] { + let mut provenance = crate::ComputeGraph::new(); + let argument = provenance.host_input("x", [1, 16]).unwrap(); + provenance.gelu(argument).unwrap(); let source = TensorType::new( if view { vec![1, 729, 144] @@ -231,8 +234,8 @@ }], outputs: vec![MidValueId(1)], operations: vec![MidOperation { - site: None, - source: None, + site: Some("packing".into()), + source: Some(provenance.operations()[0].id), inputs: vec![MidValueId(0)], results: vec![MidValueId(1)], kind: MidOperationKind::Copy { @@ -254,11 +257,15 @@ .filter_map(|rows| program.with_distributed_packing(rows)) .collect::>(); assert!(!candidates.is_empty() && candidates.len() <= 4); - for packed in candidates { + for (index, packed) in candidates.into_iter().enumerate() { let low = crate::lower_to_tiles( &crate::low::expand::expand_tiles(&packed, false).unwrap(), false, ); + if let Ok(directory) = std::env::var("IPU_REFACTOR_PACKING_CAPTURE") { + std::fs::create_dir_all(&directory).unwrap(); + std::fs::write(std::path::Path::new(&directory).join(format!("view-{view}-{index}.txt")), format!("{low:#?}")).unwrap(); + } let mut packs = 0; for run in &low.kernel_runs { if let TileKernelSpec::Rearrange { from, to } = &run.kernel {