From 1b4f070bef51b01bf73b009d2ed2b3a0e5e67ed0 Mon Sep 17 00:00:00 2001 From: NiccoloN Date: Wed, 29 Jul 2026 18:20:44 +0200 Subject: [PATCH] finally fast googlenet with correct latency artifacts for fair comparison --- .../SPATIAL_TARGET_GENERALITY_INVARIANT.md | 73 +++ AGENTS.md | 1 + README.md | 5 +- .../src/lib/instruction_set/isa.rs | 61 ++- .../src/lib/tracing/trace/tracing_isa.rs | 25 +- .../pim/pim-simulator/tests/placeholder.rs | 21 +- .../pim/pim-simulator/tests/simd.rs | 18 +- backend-simulators/pim/pimsim-nn | 2 +- src/PIM/Compiler/PimCodeGen.cpp | 203 ++++---- src/PIM/Compiler/PimCodeGen.hpp | 22 +- src/PIM/Compiler/PimCompilerOptions.cpp | 6 + src/PIM/Compiler/PimCompilerOptions.hpp | 3 + src/PIM/Compiler/PimCompilerUtils.cpp | 266 ++++++++++- src/PIM/Compiler/PimWeightEmitter.cpp | 8 +- .../Conversion/ONNXToSpatial/CMakeLists.txt | 1 - .../Common/RowStripLayoutUtils.cpp | 113 ++++- .../Common/RowStripLayoutUtils.hpp | 10 + .../ONNXToSpatial/LowerSpatialPlansPass.cpp | 172 ++++++- .../Conversion/ONNXToSpatial/ONNXToSpatial.td | 5 - .../ONNXToSpatial/ONNXToSpatialPass.cpp | 9 +- .../ONNXToSpatial/ONNXToSpatialVerifier.cpp | 3 + src/PIM/Conversion/ONNXToSpatial/Patterns.cpp | 1 - src/PIM/Conversion/ONNXToSpatial/Patterns.hpp | 1 - .../Patterns/GeneratedConversion.cpp | 18 - .../ONNXToSpatial/Patterns/Math/Conv.cpp | 421 +++++++---------- .../Patterns/Math/Elementwise.cpp | 7 + .../ONNXToSpatial/Patterns/NN/Pool.cpp | 135 ++++++ .../ONNXToSpatial/Patterns/Tensor/Concat.cpp | 11 + .../Conversion/ONNXToSpatial/PlanLowering.hpp | 8 + .../SpatialLayoutPlanningPass.cpp | 69 +++ .../SpatialToPim/SpatialToPimPass.cpp | 14 +- src/PIM/Dialect/Pim/PimOpsVerify.cpp | 13 +- .../Bufferization/PimBufferizationPass.cpp | 68 +++ .../HostConstantFolding/CMakeLists.txt | 1 + .../HostConstantFoldingPass.cpp | 2 + src/PIM/Dialect/Spatial/Spatial.td | 47 ++ src/PIM/Dialect/Spatial/SpatialOpsAsm.cpp | 9 +- src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp | 73 ++- .../DeferredCommunicationPlanning.cpp | 23 +- .../DeferredCommunicationPlanning.hpp | 2 +- .../DeferredCommunicationRealization.cpp | 196 ++++---- .../DeferredCommunicationRealization.hpp | 7 +- .../MergeComputeNodesPass.cpp | 51 ++- .../ScheduledComputePlan.hpp | 11 +- .../ScheduledComputePlanning.cpp | 8 +- .../ScheduledComputeVerification.cpp | 9 +- .../Scheduling/ComputeGraph.cpp | 432 ++++++++++-------- .../Scheduling/ComputeGraph.hpp | 42 +- .../Scheduling/ComputeInstanceUtils.cpp | 62 +-- .../Scheduling/ComputeInstanceUtils.hpp | 24 +- .../Scheduling/MergeSchedule.hpp | 1 + .../Scheduling/MergeSchedulingAnalysis.cpp | 43 +- .../Scheduling/MergeSchedulingAnalysis.hpp | 4 +- .../Scheduling/PeftScheduler.cpp | 273 ++++++----- .../Scheduling/PeftScheduler.hpp | 20 +- .../Scheduling/SchedulingTarget.hpp | 54 +++ .../TrivialGraphComputeMergePass.cpp | 56 ++- src/PIM/Pass/PIMPasses.h | 7 + src/PIM/PimAccelerator.cpp | 4 +- test/PIM/CMakeLists.txt | 7 + test/PIM/SpatialSchedulingTargetTest.cpp | 58 +++ third_party/PIMCOMP-NN | 2 +- validation/.gitignore | 1 + validation/README.md | 14 +- validation/networks/pimcomp_models/README.md | 65 ++- validation/networks/pimcomp_models/RESULTS.md | 32 -- ...softmax.onnx => googlenet-12-latency.onnx} | Bin 28021922 -> 28021720 bytes .../networks/pimcomp_models/results.csv | 5 + validation/raptor_validation/pimsim_nn.py | 73 +++ validation/raptor_validation/validate_one.py | 29 +- .../{ => pimcomp}/compare_raptor_pimcomp.py | 160 ++++--- .../pimcomp/run_pimcomp_paper_latency.py | 219 +++++++++ validation/tools/run_pimcomp_paper_latency.py | 148 ------ validation/validate.py | 17 +- 74 files changed, 2773 insertions(+), 1311 deletions(-) create mode 100644 .agents/invariants/SPATIAL_TARGET_GENERALITY_INVARIANT.md delete mode 100644 src/PIM/Conversion/ONNXToSpatial/Patterns/GeneratedConversion.cpp create mode 100644 src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/SchedulingTarget.hpp create mode 100644 test/PIM/SpatialSchedulingTargetTest.cpp delete mode 100644 validation/networks/pimcomp_models/RESULTS.md rename validation/networks/pimcomp_models/googlenet/{googlenet-12-no-softmax.onnx => googlenet-12-latency.onnx} (99%) create mode 100644 validation/networks/pimcomp_models/results.csv create mode 100644 validation/raptor_validation/pimsim_nn.py rename validation/tools/{ => pimcomp}/compare_raptor_pimcomp.py (94%) create mode 100755 validation/tools/pimcomp/run_pimcomp_paper_latency.py delete mode 100755 validation/tools/run_pimcomp_paper_latency.py diff --git a/.agents/invariants/SPATIAL_TARGET_GENERALITY_INVARIANT.md b/.agents/invariants/SPATIAL_TARGET_GENERALITY_INVARIANT.md new file mode 100644 index 0000000..3bfa4e3 --- /dev/null +++ b/.agents/invariants/SPATIAL_TARGET_GENERALITY_INVARIANT.md @@ -0,0 +1,73 @@ +# Spatial Target Generality Invariant + +## Scope + +This invariant applies to: + +- the Spatial dialect and its verifiers; +- ONNX-to-Spatial planning and lowering; +- graph transforms and scheduling over Spatial IR; +- target information consumed while optimizing Spatial IR; +- lowerings from Spatial to target-specific dialects. + +## Invariant + +Spatial represents logical compute, dataflow, layout choices, parallel work, +and target-independent resource requirements. It must remain usable by +different targets, including PIM and future targets such as PULPIM. + +Raptor may ingest target information and use it to choose Spatial layouts, +partitions, placements, or schedules. That information must cross an explicit +target interface and be expressed in target-neutral terms at the Spatial +layer. The Spatial dialect and scheduler must not parse a simulator-specific +configuration, depend on a target dialect, or encode one target's instruction +latencies, memory hierarchy, communication protocol, or resource policy. + +Target adapters own translating a target configuration into the neutral +information consumed by Spatial. Target-specific dialects and their lowerings +own instruction semantics, physical memory details, communication mechanisms, +and final legality. + +## Ownership boundary + +- Spatial IR owns logical and physical planning concepts shared across targets. +- A target adapter owns configuration parsing and cost-model construction. +- The scheduler consumes an injected cost/resource model; it does not infer a + target from global PIM options or hardcoded constants. +- Spatial-to-target lowering makes the selected representation explicit. +- Target dialect verifiers reject target-specific illegal states. + +Target-neutral information may include processor topology, operation and +transfer cost queries, available parallel capacity, and opaque resource +requirements. Names and APIs at this boundary must describe those concepts, +not a particular simulator or target implementation. + +## Forbidden coupling + +Do not: + +- include PIM or PULPIM dialect headers in the Spatial dialect or scheduler; +- read PIM compiler globals directly from generic scheduling algorithms; +- parse `pimsim-nn`, PIMCOMP, or another simulator's schema in Spatial code; +- hardcode Arch-A timing, mesh, crossbar, memory, or vector constants in + Spatial cost calculations; +- add target-named Spatial operations when an existing logical/layout concept + expresses the invariant; +- repair target-specific legality in generic Spatial cleanup passes. + +## Required proof + +Changes that use target information in Spatial must show: + +- the target boundary or injected interface used; +- that Spatial IR remains valid without target-specific attributes; +- that an unknown or unsupported target fails clearly rather than silently + using PIM defaults; +- unit coverage with at least two distinct target profiles when scheduling or + cost decisions change; +- target-specific validation after lowering for every implemented target + affected by the change. + +If only one target implementation exists, keep the interface narrow and test +it with two profiles. Do not add speculative target operations or a framework +for unimplemented targets. diff --git a/AGENTS.md b/AGENTS.md index 7d2047d..7fc40a0 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -6,6 +6,7 @@ Before modifying the relevant subsystem, read: * `.agents/invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md` * `.agents/invariants/PERFORMANCE_OPTIMIZATION_INVARIANT.md` +* `.agents/invariants/SPATIAL_TARGET_GENERALITY_INVARIANT.md` * Build commands: * `cmake --build ./build_release` * `cmake --build ./build_debug` diff --git a/README.md b/README.md index 5cad246..5519851 100644 --- a/README.md +++ b/README.md @@ -100,6 +100,9 @@ options; `onnx-mlir --help` lists the inherited ONNX-MLIR options. - `--core-count=` - required positive core count for PIM compilation. - `--crossbar-size=` - crossbar width/height. Default in code is `128`. - `--crossbar-count=` - crossbars per core. Default in code is `64`. +- `--pim-target-config=` - optional PIM target configuration used by the + target adapter to construct the target-neutral Spatial scheduling cost and + topology model. Resource values must match the explicit core/crossbar flags. - `--pim-memory-report=` - emit the concise combined memory report under `reports/memory_report.txt`, or disable it. Default is `summary`. - `--pim-only-codegen` - assume input is already bufferized PIM IR and only run @@ -155,7 +158,7 @@ This writes PIM artifacts under `/tmp/raptor/pim/`. ## Validation Functional validation compiles ONNX models, compares native ONNX-MLIR and PIM -simulator outputs, and optionally reports latency and power. See +simulator outputs, and optionally reports latency, power, and energy. See [`validation/README.md`](validation/README.md) for prerequisites, usage, options, artifacts, and results. diff --git a/backend-simulators/pim/pim-simulator/src/lib/instruction_set/isa.rs b/backend-simulators/pim/pim-simulator/src/lib/instruction_set/isa.rs index 2d46907..1f78916 100644 --- a/backend-simulators/pim/pim-simulator/src/lib/instruction_set/isa.rs +++ b/backend-simulators/pim/pim-simulator/src/lib/instruction_set/isa.rs @@ -16,7 +16,7 @@ use anyhow::{Context, Result, ensure}; use rayon::prelude::*; use paste::paste; -use std::{borrow::Cow, cell::OnceCell, collections::HashMap }; +use std::{borrow::Cow, cell::OnceCell, collections::HashMap, mem::size_of}; use std::{collections::HashSet, sync::LazyLock}; macro_rules! add_name { @@ -170,8 +170,6 @@ macro_rules! add_simd_to_map { tmp.insert((32_usize,64_usize), ([<$id _impl>]:: as InstructionType)); tmp.insert((64_usize,32_usize), ([<$id _impl>]:: as InstructionType)); tmp.insert((64_usize,64_usize), ([<$id _impl>]:: as InstructionType)); - //TODO WTF WHY - tmp.insert((8_usize,8_usize), ([<$id _impl>]:: as InstructionType)); $storage.insert($id as *const () as usize, tmp); } @@ -224,6 +222,14 @@ pub fn is_setbw(functor: InstructionType) -> bool { functor as usize == setbw as *const () as usize } +fn vector_lengths(imm_len: i32) -> Result<(usize, usize)> { + let element_count: usize = imm_len.try_into().context("imm_len can not be negative")?; + let byte_len = element_count + .checked_mul(size_of::()) + .context("vector byte length overflow")?; + Ok((element_count, byte_len)) +} + #[inline(never)] pub fn setbw(cores: &mut CPU, data: InstructionData) -> Result { panic!("You are calling a placeholder, this instruction is resolved in the construction phase"); @@ -361,11 +367,11 @@ where let r1_val = add_offset_r1(r1_val, offset_select, offset_value); let r2_val = add_offset_r2(r2_val, offset_select, offset_value); let rd_val = add_offset_rd(rd_val, offset_select, offset_value); - let imm_len: usize = imm_len.try_into().context("imm_len can not be negative")?; + let (element_count, byte_len) = vector_lengths::(imm_len)?; let loads = core - .reserve_load(r1_val, imm_len)? - .reserve_load(r2_val, imm_len)? + .reserve_load(r1_val, byte_len)? + .reserve_load(r2_val, byte_len)? .execute_load::()?; let (load1, load2) = (loads[0], loads[1]); let res: Vec = load1 @@ -374,7 +380,7 @@ where .map(|(&a, &b)| a + b) .collect(); ensure!( - imm_len / size_of::() == res.len(), + element_count == res.len(), "vvadd generate a vector bigger thant it's requested elements" ); let res_up: Cow<[T]> = res.as_slice().up(); @@ -405,11 +411,11 @@ where let r1_val = add_offset_r1(r1_val, offset_select, offset_value); let r2_val = add_offset_r2(r2_val, offset_select, offset_value); let rd_val = add_offset_rd(rd_val, offset_select, offset_value); - let imm_len: usize = imm_len.try_into().context("imm_len can not be negative")?; + let (element_count, byte_len) = vector_lengths::(imm_len)?; let loads = core - .reserve_load(r1_val, imm_len)? - .reserve_load(r2_val, imm_len)? + .reserve_load(r1_val, byte_len)? + .reserve_load(r2_val, byte_len)? .execute_load::()?; let (load1, load2) = (loads[0], loads[1]); let res: Vec = load1 @@ -418,7 +424,7 @@ where .map(|(&a, &b)| a - b) .collect(); ensure!( - imm_len / size_of::() == res.len(), + element_count == res.len(), "vvadd generate a vector bigger thant it's requested elements" ); let res_up: Cow<[T]> = res.as_slice().up(); @@ -447,10 +453,10 @@ where let r1_val = add_offset_r1(r1_val, offset_select, offset_value); let r2_val = add_offset_r2(r2_val, offset_select, offset_value); let rd_val = add_offset_rd(rd_val, offset_select, offset_value); - let imm_len: usize = imm_len.try_into().context("imm_len can not be negative")?; + let (element_count, byte_len) = vector_lengths::(imm_len)?; let loads = core - .reserve_load(r1_val, imm_len)? - .reserve_load(r2_val, imm_len)? + .reserve_load(r1_val, byte_len)? + .reserve_load(r2_val, byte_len)? .execute_load::()?; let (load1, load2) = (loads[0], loads[1]); let res: Vec = load1 @@ -459,7 +465,7 @@ where .map(|(&a, &b)| a * b) .collect(); ensure!( - imm_len / size_of::() == res.len(), + element_count == res.len(), "vvadd generate a vector bigger thant it's requested elements" ); let res_up: Cow<[T]> = res.as_slice().up(); @@ -488,9 +494,10 @@ where let rd_val = core.register(rd); let r1_val = add_offset_r1(r1_val, offset_select, offset_value); let r2_val = add_offset_r2(r2_val, offset_select, offset_value); + let (_, byte_len) = vector_lengths::(imm_len)?; let loads = core - .reserve_load(r1_val, imm_len)? - .reserve_load(r2_val, imm_len)? + .reserve_load(r1_val, byte_len)? + .reserve_load(r2_val, byte_len)? .execute_load::()?; let (load1, load2) = (loads[0], loads[1]); let res: [F; 1] = [load1 @@ -527,10 +534,11 @@ where let r1_val = add_offset_r1(r1_val, offset_select, offset_value); let r2_val = add_offset_r2(r2_val, offset_select, offset_value); let rd_val = add_offset_rd(rd_val, offset_select, offset_value); + let (_, byte_len) = vector_lengths::(imm_len)?; let loads = core - .reserve_load(r1_val, imm_len)? - .reserve_load(r2_val, imm_len)? + .reserve_load(r1_val, byte_len)? + .reserve_load(r2_val, byte_len)? .execute_load::()?; let (load1, load2) = (loads[0], loads[1]); let res: Vec = load1 @@ -583,7 +591,8 @@ where "Offset select cannot be different from 1" ); let r1_val = add_offset_r1(r1_val, offset_select, offset_value); - let loads = core.reserve_load(r1_val, imm_len)?.execute_load::()?; + let (_, byte_len) = vector_lengths::(imm_len)?; + let loads = core.reserve_load(r1_val, byte_len)?.execute_load::()?; let load1 = loads[0]; let len = load1.len(); let res: [F; _] = @@ -613,7 +622,8 @@ where let rd_val = core.register(rd); let r1_val = add_offset_r1(r1_val, offset_select, offset_value); let rd_val = add_offset_rd(rd_val, offset_select, offset_value); - let loads = core.reserve_load(r1_val, imm_len)?.execute_load::()?; + let (_, byte_len) = vector_lengths::(imm_len)?; + let loads = core.reserve_load(r1_val, byte_len)?.execute_load::()?; let load1 = loads[0]; let res: Vec = load1 .iter() @@ -646,7 +656,8 @@ where let r1_val = add_offset_r1(r1_val, offset_select, offset_value); let rd_val = add_offset_rd(rd_val, offset_select, offset_value); - let loads = core.reserve_load(r1_val, imm_len)?.execute_load::()?; + let (_, byte_len) = vector_lengths::(imm_len)?; + let loads = core.reserve_load(r1_val, byte_len)?.execute_load::()?; let load1 = loads[0]; let res: Vec = load1.iter().map(|&a| a.tanh()).collect(); @@ -675,7 +686,8 @@ where let rd_val = core.register(rd); let r1_val = add_offset_r1(r1_val, offset_select, offset_value); let rd_val = add_offset_rd(rd_val, offset_select, offset_value); - let loads = core.reserve_load(r1_val, imm_len)?.execute_load::()?; + let (_, byte_len) = vector_lengths::(imm_len)?; + let loads = core.reserve_load(r1_val, byte_len)?.execute_load::()?; let load1 = loads[0]; let res: Vec = load1.iter().map(|&a| a.sigm()).collect(); let res_up: Cow<[T]> = res.as_slice().up(); @@ -706,7 +718,8 @@ where let rd_val = core.register(rd); let r1_val = add_offset_r1(r1_val, offset_select, offset_value); let rd_val = add_offset_rd(rd_val, offset_select, offset_value); - let loads = core.reserve_load(r1_val, imm_len)?.execute_load::()?; + let (_, byte_len) = vector_lengths::(imm_len)?; + let loads = core.reserve_load(r1_val, byte_len)?.execute_load::()?; let load1 = loads[0]; ensure!(!load1.is_empty(), "vsoftmax does not support empty vectors"); let max_val = load1 diff --git a/backend-simulators/pim/pim-simulator/src/lib/tracing/trace/tracing_isa.rs b/backend-simulators/pim/pim-simulator/src/lib/tracing/trace/tracing_isa.rs index b11f3e0..e4acd80 100644 --- a/backend-simulators/pim/pim-simulator/src/lib/tracing/trace/tracing_isa.rs +++ b/backend-simulators/pim/pim-simulator/src/lib/tracing/trace/tracing_isa.rs @@ -1,5 +1,5 @@ use crate::{tracing::trace::pretty_print, utility::add_offset_r2}; -use std::fs::File; +use std::{fs::File, mem::size_of}; use crate::{ cpu::CPU, @@ -381,7 +381,15 @@ impl Trace { writeln!(file, "\trs1({}): {}", r1, r1_val); writeln!(file, "\trs2({}): {}", r2, r2_val); writeln!(file, "{} Immediate:", prefix); - writeln!(file, "\tLoad Length: {} bytes", imm_len); + let element_count: usize = imm_len.try_into().expect("imm_len can not be negative"); + let byte_len = element_count + .checked_mul(size_of::()) + .expect("vector byte length overflow"); + writeln!( + file, + "\tLoad Length: {} elements ({} bytes)", + element_count, byte_len + ); writeln!( file, "\toffset_select: {} offset_value: {}", @@ -390,22 +398,21 @@ impl Trace { let r1_final = add_offset_r1(r1_val, offset_select, offset_value); let r2_final = add_offset_r2(r2_val, offset_select, offset_value); let rd_final = add_offset_rd(rd_val, offset_select, offset_value); - let imm_len: usize = imm_len.try_into().expect("imm_len can not be negative"); let loads = core - .reserve_load(r1_final, imm_len) + .reserve_load(r1_final, byte_len) .unwrap() - .reserve_load(r2_final, imm_len) + .reserve_load(r2_final, byte_len) .unwrap() - .reserve_load(rd_final, imm_len) + .reserve_load(rd_final, byte_len) .unwrap() .execute_load::() .unwrap(); writeln!(file, "{} Memory:", prefix); - write!(file, "\tLocal[{}:{}](A): ", r1_final, r1_final + imm_len); + write!(file, "\tLocal[{}:{}](A): ", r1_final, r1_final + byte_len); pretty_print::print_slice::<_,f32>(file, loads[0], 30); - write!(file, "\tLocal[{}:{}](B): ", r2_final , r2_final + imm_len); + write!(file, "\tLocal[{}:{}](B): ", r2_final , r2_final + byte_len); pretty_print::print_slice::<_,f32>(file, loads[1], 30); - write!(file, "\tLocal[{}:{}](out): ", rd_final, rd_final+ imm_len); + write!(file, "\tLocal[{}:{}](out): ", rd_final, rd_final+ byte_len); pretty_print::print_slice::<_,f32>(file, loads[2], 30); if prefix == "Post" { writeln!(file, "\n###############################################\n"); diff --git a/backend-simulators/pim/pim-simulator/tests/placeholder.rs b/backend-simulators/pim/pim-simulator/tests/placeholder.rs index 9b335ad..576ce8f 100644 --- a/backend-simulators/pim/pim-simulator/tests/placeholder.rs +++ b/backend-simulators/pim/pim-simulator/tests/placeholder.rs @@ -25,7 +25,7 @@ fn wrong_size_place_holder() { vvadd, idata_build .set_rdr1r2(3, 1, 2) - .set_imm_len(8 * size_of::() as i32) + .set_imm_len(8) .build(), ); let core_instruction = vec![inst_builder.build().into()]; @@ -33,6 +33,25 @@ fn wrong_size_place_holder() { executable.execute(); } +#[test] +#[should_panic(expected = "Function not found for the requested size input:8 output:8")] +fn unsupported_8_bit_vectors_do_not_alias_f32() { + let mut inst_builder = InstructionsBuilder::new(); + let mut idata_build = InstructionDataBuilder::new(); + idata_build.set_core_indx(0).fix_core_indx(); + inst_builder.make_inst( + setbw, + idata_build.set_ibiw_obiw(8, 8).build(), + ); + inst_builder.make_inst( + vvadd, + idata_build + .set_rdr1r2(3, 1, 2) + .set_imm_len(8) + .build(), + ); +} + fn place_holder(inst : InstructionType) { diff --git a/backend-simulators/pim/pim-simulator/tests/simd.rs b/backend-simulators/pim/pim-simulator/tests/simd.rs index 5a3ea04..8f7c3fd 100644 --- a/backend-simulators/pim/pim-simulator/tests/simd.rs +++ b/backend-simulators/pim/pim-simulator/tests/simd.rs @@ -55,7 +55,7 @@ where vvadd, idata_build .set_rdr1r2(3, 1, 2) - .set_imm_len(8 * size_of::() as i32) + .set_imm_len(8) .build(), ); let core_instruction = vec![inst_builder.build().into()]; @@ -159,7 +159,7 @@ where vvsub, idata_build .set_rdr1r2(3, 1, 2) - .set_imm_len(8 * size_of::() as i32) + .set_imm_len(8) .build(), ); let core_instruction = vec![inst_builder.build().into()]; @@ -263,7 +263,7 @@ where vvmul, idata_build .set_rdr1r2(3, 1, 2) - .set_imm_len(8 * size_of::() as i32) + .set_imm_len(8) .build(), ); let core_instruction = vec![inst_builder.build().into()]; @@ -367,7 +367,7 @@ where vvdmul, idata_build .set_rdr1r2(3, 1, 2) - .set_imm_len(8 * size_of::() as i32) + .set_imm_len(8) .build(), ); let core_instruction = vec![inst_builder.build().into()]; @@ -464,7 +464,7 @@ where vvmax, idata_build .set_rdr1r2(3, 1, 2) - .set_imm_len(8 * size_of::() as i32) + .set_imm_len(8) .build(), ); let core_instruction = vec![inst_builder.build().into()]; @@ -565,7 +565,7 @@ where idata_build .set_rdr1r2(3, 1, 1) .set_offset_select(1) - .set_imm_len(8 * size_of::() as i32) + .set_imm_len(8) .build(), ); let core_instruction = vec![inst_builder.build().into()]; @@ -658,7 +658,7 @@ where vrelu, idata_build .set_rdr1r2(3, 1, 1) - .set_imm_len(8 * size_of::() as i32) + .set_imm_len(8) .build(), ); let core_instruction = vec![inst_builder.build().into()]; @@ -758,7 +758,7 @@ where vtanh, idata_build .set_rdr1r2(3, 1, 1) - .set_imm_len(8 * size_of::() as i32) + .set_imm_len(8) .build(), ); let core_instruction = vec![inst_builder.build().into()]; @@ -860,7 +860,7 @@ where vsigm, idata_build .set_rdr1r2(3, 1, 1) - .set_imm_len(8 * size_of::() as i32) + .set_imm_len(8) .build(), ); let core_instruction = vec![inst_builder.build().into()]; diff --git a/backend-simulators/pim/pimsim-nn b/backend-simulators/pim/pimsim-nn index b405399..0d03316 160000 --- a/backend-simulators/pim/pimsim-nn +++ b/backend-simulators/pim/pimsim-nn @@ -1 +1 @@ -Subproject commit b405399512eb5e84ca555225251910d06ff40178 +Subproject commit 0d03316df4734eab19d8d3178fd84afdd3ee2349 diff --git a/src/PIM/Compiler/PimCodeGen.cpp b/src/PIM/Compiler/PimCodeGen.cpp index 4e2a0dc..9806a59 100644 --- a/src/PIM/Compiler/PimCodeGen.cpp +++ b/src/PIM/Compiler/PimCodeGen.cpp @@ -92,11 +92,12 @@ static MemoryReportKind classifyMemoryReportKind(mlir::Value value) { return MemoryReportKind::None; } -static int32_t getVectorByteSizeOrCrash(ShapedType type) { - auto byteSize = pim::getCheckedShapedTypeSizeInBytes(type, UnknownLoc::get(type.getContext()), "vector byte size"); - if (failed(byteSize)) - llvm_unreachable("Failed to compute checked vector byte size"); - return pim::checkedI32OrCrash(*byteSize, "vector byte size"); +static int32_t getVectorElementCountOrCrash(ShapedType type) { + return pim::checkedI32OrCrash(type.getNumElements(), "vector element count"); +} + +static int32_t getVectorElementBitwidthOrCrash(ShapedType type) { + return pim::checkedI32OrCrash(static_cast(type.getElementTypeBitWidth()), "vector element bitwidth"); } static Operation* getDiagnosticAnchor(mlir::Value value) { @@ -165,8 +166,7 @@ size_t PimMemory::allocateAddress(size_t size, const MemoryValueKey& key) { checkedAlignedEnd = checkedAlignTo(*checkedEnd, minAlignment, anchor, "local memory alignment"); if (address > kPimLocalMemoryAddressLimit || failed(checkedEnd) || *checkedEnd > kPimLocalMemoryAddressLimit || failed(checkedAlignedEnd) || *checkedAlignedEnd > kPimLocalMemoryAddressLimit) { - printMemoryOverflowDiagnostic( - key, + printMemoryOverflowDiagnostic(key, size, firstAvailableAddress, succeeded(checkedAlignedEnd) ? *checkedAlignedEnd : kPimLocalMemoryAddressLimit); @@ -208,9 +208,7 @@ void PimMemory::allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memE switch (reportKind) { case MemoryReportKind::Alloca: case MemoryReportKind::Global: - case MemoryReportKind::Input: - ++reportRow.hostObjectCount; - break; + case MemoryReportKind::Input: ++reportRow.hostObjectCount; break; case MemoryReportKind::None: break; } } @@ -259,8 +257,7 @@ void PimMemory::allocateCore(const CompiledCoreMemoryPlan& plan, std::optional PimAcceleratorMemory::getIndexValue(mlir::Value value, PimAcceleratorMemory::PimAcceleratorMemory() : hostMem(memEntriesMap), fileReport(openMemoryReport(pimMemoryReport == PimMemoryReportSummary)) {} -PimAcceleratorMemory::PimAcceleratorMemory( - const llvm::SmallDenseMap& initialMemEntries, bool enableReport) +PimAcceleratorMemory::PimAcceleratorMemory(const llvm::SmallDenseMap& initialMemEntries, + bool enableReport) : memEntriesMap(initialMemEntries), hostMem(memEntriesMap), fileReport(enableReport ? openMemoryReport(true) : std::fstream()) {} @@ -367,10 +364,8 @@ PimAcceleratorMemory::PimAcceleratorMemory( void PimAcceleratorMemory::reportHost() { hostReportRow = hostMem.getReportRow(); } void PimAcceleratorMemory::recordCoreReport(size_t coreId, const MemoryReportRow& row) { - reportEntries.push_back({MemoryReportEntry::Kind::Core, - coreId, - {pim::checkedI32OrCrash(coreId, "memory report core id")}, - row}); + reportEntries.push_back( + {MemoryReportEntry::Kind::Core, coreId, {pim::checkedI32OrCrash(coreId, "memory report core id")}, row}); } void PimAcceleratorMemory::recordBatchReport(uint64_t batchId, @@ -441,8 +436,7 @@ void PimAcceleratorMemory::flushReport() { os << " Weights memory: " << formatReportMemory(totalWeightBytes) << "\n"; os << " Local memory before reuse: " << formatReportMemory(logicalBytes) << "\n"; os << " Local memory after reuse: " << formatReportMemory(physicalBytes) << "\n"; - os << " Saved local memory: " << formatReportMemory(savedBytes) << " (" - << formatv("{0:F1}%", savedPercent) << ")\n"; + os << " Saved local memory: " << formatReportMemory(savedBytes) << " (" << formatv("{0:F1}%", savedPercent) << ")\n"; os << " Largest core local memory: " << formatReportMemory(largest) << "\n"; if (!groups.empty()) { os << " "; @@ -456,20 +450,15 @@ void PimAcceleratorMemory::flushReport() { printLabel(group); os << "\n"; uint64_t groupSaved = group.row.logicalLocalBytes - group.row.physicalLocalBytes; - double groupPercent = group.row.logicalLocalBytes == 0 - ? 0.0 - : 100.0 * groupSaved / group.row.logicalLocalBytes; + double groupPercent = group.row.logicalLocalBytes == 0 ? 0.0 : 100.0 * groupSaved / group.row.logicalLocalBytes; if (group.coreIds.size() == 1) { os << " Local memory: " << formatReportMemory(group.row.logicalLocalBytes) << " → " - << formatReportMemory(group.row.physicalLocalBytes) << " (" << formatv("{0:F1}% saved", groupPercent) - << ")\n"; + << formatReportMemory(group.row.physicalLocalBytes) << " (" << formatv("{0:F1}% saved", groupPercent) << ")\n"; } else { os << " Per core: " << formatReportMemory(group.row.logicalLocalBytes) << " → " - << formatReportMemory(group.row.physicalLocalBytes) << " (" << formatv("{0:F1}% saved", groupPercent) - << ")\n"; - os << " Total after reuse: " << formatReportMemory(group.row.physicalLocalBytes * group.coreIds.size()) - << "\n"; + << formatReportMemory(group.row.physicalLocalBytes) << " (" << formatv("{0:F1}% saved", groupPercent) << ")\n"; + os << " Total after reuse: " << formatReportMemory(group.row.physicalLocalBytes * group.coreIds.size()) << "\n"; } } if (groups.size() > kGroupLimit) @@ -479,12 +468,6 @@ void PimAcceleratorMemory::flushReport() { fileReport.close(); } -size_t PimCodeGen::remapCoreId(size_t coreId) const { - auto it = emittedCoreIds.find(coreId); - assert(it != emittedCoreIds.end() && "Missing emitted core id remapping"); - return it->second; -} - void PimCodeGen::emitInstruction(const pim_binary::InstructionRecord& instruction) const { if (failed(instructionWriter.append(instruction))) return; @@ -493,6 +476,19 @@ void PimCodeGen::emitInstruction(const pim_binary::InstructionRecord& instructio updateScalarRegisterCache(instruction); } +void PimCodeGen::ensureVectorBitwidth(int32_t inputBitwidth, int32_t outputBitwidth) const { + std::array requested = {inputBitwidth, outputBitwidth}; + if (vectorBitwidths == requested) + return; + + pim_binary::InstructionRecord instruction; + instruction.opcode = pim_binary::Opcode::setbw; + instruction.generic1 = inputBitwidth; + instruction.generic2 = outputBitwidth; + emitInstruction(instruction); + vectorBitwidths = requested; +} + void PimCodeGen::updateScalarRegisterCache(const pim_binary::InstructionRecord& instruction) const { switch (instruction.opcode) { case pim_binary::Opcode::sldi: scalarRegisterValues[instruction.rd] = instruction.r2OrImm; break; @@ -563,7 +559,7 @@ void PimCodeGen::emitCommunicationOp(pim_binary::Opcode opcode, size_t bufferAdd pim_binary::InstructionRecord instruction; instruction.opcode = opcode; instruction.rd = 0; - instruction.r2OrImm = pim::checkedI32OrCrash(remapCoreId(coreId), "communication core id"); + instruction.r2OrImm = pim::checkedI32OrCrash(coreId, "physical communication core id"); instruction.generic1 = 0; instruction.generic2 = 0; instruction.generic3 = pim::checkedI32OrCrash(size, "communication byte size"); @@ -679,6 +675,8 @@ void PimCodeGen::codeGenMVMLikeOp(size_t mvmId, MVMTy mvmLikeOp, bool transposeMatrix, const StaticValueKnowledge& knowledge) { + ensureVectorBitwidth(getVectorElementBitwidthOrCrash(cast(mvmLikeOp.getInput().getType())), + getVectorElementBitwidthOrCrash(cast(mvmLikeOp.getOutputBuffer().getType()))); emitMvmOp(mvmId, addressOf(mvmLikeOp.getOutputBuffer(), knowledge), 0, addressOf(mvmLikeOp.getInput(), knowledge), 0); // TODO: save weights somewhere (if transposeMatrix=true, transpose the weight matrix) @@ -688,25 +686,29 @@ void PimCodeGen::emitBinaryVectorOp(pim_binary::Opcode opcode, mlir::Value output, mlir::Value lhs, mlir::Value rhs, - size_t byteSize, const StaticValueKnowledge& knowledge) const { + auto inputType = cast(lhs.getType()); + ensureVectorBitwidth(getVectorElementBitwidthOrCrash(inputType), + getVectorElementBitwidthOrCrash(cast(output.getType()))); setupRdRs1Rs2(addressOf(output, knowledge), 0, addressOf(lhs, knowledge), 0, addressOf(rhs, knowledge), 0); pim_binary::InstructionRecord instruction; instruction.opcode = opcode; instruction.rd = 0; instruction.r1 = 1; instruction.r2OrImm = 2; - instruction.generic3 = pim::checkedI32OrCrash(byteSize, "vector byte size"); + instruction.generic3 = getVectorElementCountOrCrash(inputType); emitInstruction(instruction); } void PimCodeGen::emitUnaryVectorOp(pim_binary::Opcode opcode, mlir::Value output, mlir::Value input, - size_t byteSize, const StaticValueKnowledge& knowledge, int32_t r2OrImm, int32_t generic1) const { + auto inputType = cast(input.getType()); + ensureVectorBitwidth(getVectorElementBitwidthOrCrash(inputType), + getVectorElementBitwidthOrCrash(cast(output.getType()))); setupRdRs1(addressOf(output, knowledge), 0, addressOf(input, knowledge), 0); pim_binary::InstructionRecord instruction; instruction.opcode = opcode; @@ -714,7 +716,7 @@ void PimCodeGen::emitUnaryVectorOp(pim_binary::Opcode opcode, instruction.r1 = 1; instruction.r2OrImm = r2OrImm; instruction.generic1 = generic1; - instruction.generic3 = pim::checkedI32OrCrash(byteSize, "vector byte size"); + instruction.generic3 = getVectorElementCountOrCrash(inputType); emitInstruction(instruction); } @@ -812,8 +814,7 @@ static SmallVector collectTopLevelCoreLikeOps(func::FuncOp funcOp) { static FailureOr compileCoreMemoryPlan(Operation* coreLikeOp) { CompiledCoreMemoryPlan plan; auto arenaAttr = coreLikeOp->getAttrOfType(kLocalMemorySizeAttrName); - if (!arenaAttr || arenaAttr.getInt() < 0 - || static_cast(arenaAttr.getInt()) > kPimLocalMemoryAddressLimit) { + if (!arenaAttr || arenaAttr.getInt() < 0 || static_cast(arenaAttr.getInt()) > kPimLocalMemoryAddressLimit) { coreLikeOp->emitError("requires a valid pim.local_memory_size attribute before codegen"); return failure(); } @@ -841,7 +842,9 @@ static FailureOr compileCoreMemoryPlan(Operation* coreLi hasFailure = true; return; } - plan.entries.push_back({allocOp.getResult(), {address, static_cast(*checkedSize)}}); + plan.entries.push_back({ + allocOp.getResult(), {address, static_cast(*checkedSize)} + }); auto logicalBytes = pim::checkedAdd( static_cast(plan.logicalBytes), static_cast(*checkedSize), allocOp, "logical local bytes"); if (failed(logicalBytes)) { @@ -995,21 +998,10 @@ static LogicalResult executeCompiledCorePlan( } auto emitBinary = [&](auto op, pim_binary::Opcode opcode) { - coreCodeGen.emitBinaryVectorOp(opcode, - op.getOutputBuffer(), - op.getLhs(), - op.getRhs(), - getVectorByteSizeOrCrash(cast(op.getLhs().getType())), - knowledge); + coreCodeGen.emitBinaryVectorOp(opcode, op.getOutputBuffer(), op.getLhs(), op.getRhs(), knowledge); }; auto emitUnary = [&](auto op, pim_binary::Opcode opcode, int32_t r2OrImm, int32_t generic1) { - coreCodeGen.emitUnaryVectorOp(opcode, - op.getOutputBuffer(), - op.getInput(), - getVectorByteSizeOrCrash(cast(op.getInput().getType())), - knowledge, - r2OrImm, - generic1); + coreCodeGen.emitUnaryVectorOp(opcode, op.getOutputBuffer(), op.getInput(), knowledge, r2OrImm, generic1); }; switch (node.opKind) { @@ -1092,8 +1084,8 @@ static void aliasMaterializedHostGlobals(CoreLikeOpTy coreLikeOp, }); } -static OnnxMlirCompilerErrorCodes emitEmptyCoreArtifacts(StringRef outputDirPath, size_t emittedCoreId) { - std::string outputCorePath = (outputDirPath + "/core_" + std::to_string(emittedCoreId) + ".pim").str(); +static OnnxMlirCompilerErrorCodes emitEmptyCoreArtifacts(StringRef outputDirPath, size_t physicalCoreId) { + std::string outputCorePath = (outputDirPath + "/core_" + std::to_string(physicalCoreId) + ".pim").str(); std::error_code errorCode; raw_fd_ostream coreBinaryStream(outputCorePath, errorCode, sys::fs::OF_None); if (errorCode) { @@ -1117,7 +1109,7 @@ static OnnxMlirCompilerErrorCodes emitEmptyCoreArtifacts(StringRef outputDirPath if (!pimEmitJson.getValue()) return CompilerSuccess; - std::string outputCoreJsonPath = (outputDirPath + "/core_" + std::to_string(emittedCoreId) + ".json").str(); + std::string outputCoreJsonPath = (outputDirPath + "/core_" + std::to_string(physicalCoreId) + ".json").str(); errorCode = std::error_code(); raw_fd_ostream coreJsonStream(outputCoreJsonPath, errorCode); if (errorCode) { @@ -1179,36 +1171,15 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: return it->second.get(); }; - llvm::DenseMap emittedCoreIds; - size_t nextEmittedCoreId = 0; - - for (Operation* op : coreLikeOps) { - if (auto coreOp = dyn_cast(op)) { - size_t originalCoreId = static_cast(coreOp.getCoreId()); - if (!emittedCoreIds.contains(originalCoreId)) - emittedCoreIds[originalCoreId] = nextEmittedCoreId++; - continue; - } - - auto coreBatchOp = cast(op); - auto batchCoreIds = getBatchCoreIds(coreBatchOp); - for (unsigned lane = 0; lane < static_cast(coreBatchOp.getLaneCount()); ++lane) { - size_t originalCoreId = static_cast(batchCoreIds[lane]); - if (!emittedCoreIds.contains(originalCoreId)) - emittedCoreIds[originalCoreId] = nextEmittedCoreId++; - } - } - SmallVector jobs; SmallVector> batchJobIndices; for (Operation* op : coreLikeOps) { if (auto coreOp = dyn_cast(op)) { - size_t originalCoreId = static_cast(coreOp.getCoreId()); CoreEmissionJob job; job.coreLikeOp = coreOp; job.program = getCompiledProgram(op); job.memoryPlan = getMemoryPlan(op); - job.emittedCoreId = emittedCoreIds.lookup(originalCoreId); + job.physicalCoreId = static_cast(coreOp.getCoreId()); jobs.push_back(std::move(job)); continue; } @@ -1220,16 +1191,15 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: lanesByCoreId[static_cast(batchCoreIds[lane])].push_back(lane); SmallVector jobIndices; - SmallVector orderedOriginalCoreIds = llvm::to_vector(lanesByCoreId.keys()); - llvm::sort(orderedOriginalCoreIds, - [&](size_t lhs, size_t rhs) { return emittedCoreIds.lookup(lhs) < emittedCoreIds.lookup(rhs); }); - for (size_t originalCoreId : orderedOriginalCoreIds) { + SmallVector physicalCoreIds = llvm::to_vector(lanesByCoreId.keys()); + llvm::sort(physicalCoreIds); + for (size_t physicalCoreId : physicalCoreIds) { CoreEmissionJob job; job.coreLikeOp = coreBatchOp; job.program = getCompiledProgram(op); job.memoryPlan = getMemoryPlan(op); - job.emittedCoreId = emittedCoreIds.lookup(originalCoreId); - job.lanes = lanesByCoreId.lookup(originalCoreId); + job.physicalCoreId = physicalCoreId; + job.lanes = lanesByCoreId.lookup(physicalCoreId); job.batchReportId = nextBatchReportId; jobIndices.push_back(jobs.size()); jobs.push_back(std::move(job)); @@ -1238,8 +1208,11 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: ++nextBatchReportId; } - auto linkCoreWeights = - [&](size_t coreId, ArrayRef weightFiles, json::Array& xbarsPerGroup) -> OnnxMlirCompilerErrorCodes { + auto linkCoreWeights = [&](size_t coreId, + ArrayRef weightFiles, + ArrayRef weights, + json::Array& xbarsPerGroup) -> OnnxMlirCompilerErrorCodes { + assert(weightFiles.size() == weights.size() && "weight files must match resolved weight views"); auto coreWeightsDirPath = outputDirPath + "/core_" + std::to_string(coreId); if (auto error = sys::fs::create_directory(coreWeightsDirPath); error && error != std::errc::file_exists) { errs() << "Error creating core directory: " << coreWeightsDirPath << ": " << error.message() << '\n'; @@ -1247,7 +1220,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: } for (auto [slot, fileName] : llvm::enumerate(weightFiles)) { - xbarsPerGroup.push_back(1); + xbarsPerGroup.push_back(weights[slot].shape[1] / static_cast(crossbarSize)); std::string sourcePath = outputDirPath + "/weights/" + fileName; std::string targetPath = coreWeightsDirPath + "/crossbar_" + std::to_string(slot) + ".bin"; sys::fs::remove(targetPath); @@ -1295,7 +1268,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: }; std::error_code errorCode; - auto outputCorePath = outputDirPath + "/core_" + std::to_string(job.emittedCoreId) + ".pim"; + auto outputCorePath = outputDirPath + "/core_" + std::to_string(job.physicalCoreId) + ".pim"; raw_fd_ostream coreBinaryStream(outputCorePath, errorCode, sys::fs::OF_None); if (errorCode) { errs() << "Error while opening core file `" << outputCorePath << "`: " << errorCode.message() << '\n'; @@ -1305,7 +1278,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: std::unique_ptr coreJsonStream; if (pimEmitJson.getValue()) { - std::string outputCoreJsonPath = outputDirPath + "/core_" + std::to_string(job.emittedCoreId) + ".json"; + std::string outputCoreJsonPath = outputDirPath + "/core_" + std::to_string(job.physicalCoreId) + ".json"; errorCode = std::error_code(); coreJsonStream = std::make_unique(outputCoreJsonPath, errorCode); if (errorCode) { @@ -1317,7 +1290,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: } PimInstructionWriter instructionWriter(coreBinaryStream); - PimCodeGen coreCodeGen(jobMemory, instructionWriter, coreJsonStream.get(), emittedCoreIds); + PimCodeGen coreCodeGen(jobMemory, instructionWriter, coreJsonStream.get()); auto finalizeInstructions = [&]() { bool succeeded = mlir::succeeded(instructionWriter.finalize()); coreBinaryStream.close(); @@ -1330,7 +1303,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: if (auto coreOp = dyn_cast(job.coreLikeOp)) { aliasMaterializedHostGlobals(coreOp, moduleOp, materializedHostGlobals, jobMemory); - auto& deviceMemory = jobMemory.getOrCreateDeviceMem(job.emittedCoreId); + auto& deviceMemory = jobMemory.getOrCreateDeviceMem(job.physicalCoreId); deviceMemory.allocateCore(*job.memoryPlan); StaticValueKnowledge knowledge = seedCoreCodegenKnowledge(coreOp); @@ -1345,7 +1318,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: else { auto coreBatchOp = cast(job.coreLikeOp); aliasMaterializedHostGlobals(coreBatchOp, moduleOp, materializedHostGlobals, jobMemory); - auto& deviceMemory = jobMemory.getOrCreateDeviceMem(job.emittedCoreId); + auto& deviceMemory = jobMemory.getOrCreateDeviceMem(job.physicalCoreId); for (unsigned lane : job.lanes) { StaticValueKnowledge knowledge = seedCoreBatchCodegenKnowledge(coreBatchOp, lane); @@ -1398,18 +1371,29 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: if (jobResults[jobIndex].status != CompilerSuccess) return jobResults[jobIndex].status; - if (jobs.empty()) { - if (auto err = emitEmptyCoreArtifacts(outputDirPath, 0)) + size_t maxPhysicalCoreId = 0; + for (const CoreEmissionJob& job : jobs) + maxPhysicalCoreId = std::max(maxPhysicalCoreId, job.physicalCoreId); + std::vector activePhysicalCores(maxPhysicalCoreId + 1); + for (const CoreEmissionJob& job : jobs) + activePhysicalCores[job.physicalCoreId] = true; + for (size_t physicalCoreId = 0; + physicalCoreId < activePhysicalCores.size(); ++physicalCoreId) { + if (activePhysicalCores[physicalCoreId]) + continue; + if (auto err = + emitEmptyCoreArtifacts(outputDirPath, physicalCoreId)) return err; - xbarsPerArrayGroup["core0"] = json::Array {}; - memory.recordCoreReport(0, MemoryReportRow {}); + xbarsPerArrayGroup["core" + std::to_string(physicalCoreId)] = + json::Array {}; + memory.recordCoreReport(physicalCoreId, MemoryReportRow {}); } llvm::SmallVector weightRequests; weightRequests.reserve(jobs.size()); for (size_t jobIndex = 0; jobIndex < jobs.size(); ++jobIndex) { WeightFileRequest request; - request.coreId = jobs[jobIndex].emittedCoreId; + request.coreId = jobs[jobIndex].physicalCoreId; request.weights = jobResults[jobIndex].usedWeights; weightRequests.push_back(std::move(request)); } @@ -1422,10 +1406,11 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: json::Array xbarsPerGroup; if (auto coreOp = dyn_cast(job.coreLikeOp)) { - if (auto err = linkCoreWeights(job.emittedCoreId, mapCoreWeightToFileName[job.emittedCoreId], xbarsPerGroup)) + if (auto err = linkCoreWeights( + job.physicalCoreId, mapCoreWeightToFileName[job.physicalCoreId], result.usedWeights, xbarsPerGroup)) return err; - xbarsPerArrayGroup["core" + std::to_string(job.emittedCoreId)] = std::move(xbarsPerGroup); - memory.recordCoreReport(job.emittedCoreId, result.reportRow); + xbarsPerArrayGroup["core" + std::to_string(job.physicalCoreId)] = std::move(xbarsPerGroup); + memory.recordCoreReport(job.physicalCoreId, result.reportRow); continue; } } @@ -1438,10 +1423,11 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: const CoreEmissionJob& job = jobs[jobIndex]; const CoreEmissionResult& result = jobResults[jobIndex]; json::Array xbarsPerGroup; - if (auto err = linkCoreWeights(job.emittedCoreId, mapCoreWeightToFileName[job.emittedCoreId], xbarsPerGroup)) + if (auto err = linkCoreWeights( + job.physicalCoreId, mapCoreWeightToFileName[job.physicalCoreId], result.usedWeights, xbarsPerGroup)) return err; - xbarsPerArrayGroup["core" + std::to_string(job.emittedCoreId)] = std::move(xbarsPerGroup); - reportedCoreIds.push_back(pim::checkedI32OrCrash(job.emittedCoreId, "batch report core id")); + xbarsPerArrayGroup["core" + std::to_string(job.physicalCoreId)] = std::move(xbarsPerGroup); + reportedCoreIds.push_back(pim::checkedI32OrCrash(job.physicalCoreId, "batch report physical core id")); if (!batchPerCoreRow) batchPerCoreRow = result.reportRow; else if (!(*batchPerCoreRow == result.reportRow)) @@ -1449,11 +1435,10 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: } uint64_t batchReportId = jobs[group.front()].batchReportId.value_or(0); - memory.recordBatchReport( - batchReportId, reportedCoreIds, batchPerCoreRow.value_or(MemoryReportRow {})); + memory.recordBatchReport(batchReportId, reportedCoreIds, batchPerCoreRow.value_or(MemoryReportRow {})); } - maxCoreId = nextEmittedCoreId == 0 ? 0 : nextEmittedCoreId - 1; + maxCoreId = maxPhysicalCoreId; memory.flushReport(); return writeConfigJson(funcOp, memory, maxCoreId, std::move(xbarsPerArrayGroup), outputDirPath); diff --git a/src/PIM/Compiler/PimCodeGen.hpp b/src/PIM/Compiler/PimCodeGen.hpp index 41bc893..ab6c2fb 100644 --- a/src/PIM/Compiler/PimCodeGen.hpp +++ b/src/PIM/Compiler/PimCodeGen.hpp @@ -134,8 +134,7 @@ private: public: PimAcceleratorMemory(); - PimAcceleratorMemory( - const llvm::SmallDenseMap& initialMemEntries, bool enableReport); + PimAcceleratorMemory(const llvm::SmallDenseMap& initialMemEntries, bool enableReport); PimMemory& getOrCreateDeviceMem(size_t id); @@ -145,8 +144,7 @@ public: llvm::FailureOr getIndexValue(mlir::Value value, const StaticValueKnowledge& knowledge = {}) const; void reportHost(); void recordCoreReport(size_t coreId, const MemoryReportRow& row); - void recordBatchReport( - uint64_t batchId, llvm::ArrayRef coreIds, const MemoryReportRow& perCoreRow); + void recordBatchReport(uint64_t batchId, llvm::ArrayRef coreIds, const MemoryReportRow& perCoreRow); void setTotalWeightBytes(uint64_t bytes) { totalWeightBytes = bytes; } void flushReport(); }; @@ -155,7 +153,7 @@ struct CoreEmissionJob { mlir::Operation* coreLikeOp = nullptr; const CompiledCoreProgram* program = nullptr; const CompiledCoreMemoryPlan* memoryPlan = nullptr; - size_t emittedCoreId = 0; + size_t physicalCoreId = 0; llvm::SmallVector lanes; std::optional batchReportId; }; @@ -164,17 +162,16 @@ class PimCodeGen { PimAcceleratorMemory& memory; PimInstructionWriter& instructionWriter; llvm::raw_fd_ostream* coreJsonStream; - const llvm::DenseMap& emittedCoreIds; std::optional batchLane; mutable std::array, 256> scalarRegisterValues = {}; + mutable std::optional> vectorBitwidths; size_t addressOf(mlir::Value value, const StaticValueKnowledge& knowledge) const { return memory.getValueAddress(value, knowledge, batchLane); } - size_t remapCoreId(size_t coreId) const; - void emitInstruction(const pim_binary::InstructionRecord& instruction) const; void updateScalarRegisterCache(const pim_binary::InstructionRecord& instruction) const; + void ensureVectorBitwidth(int32_t inputBitwidth, int32_t outputBitwidth) const; void genSetRegisterImmediate(uint8_t registerNumber, int32_t immediate) const; void genSetRegisterImmediateUnsigned(size_t registerNumber, size_t immediate) const; @@ -198,21 +195,16 @@ public: mlir::Value output, mlir::Value lhs, mlir::Value rhs, - size_t byteSize, const StaticValueKnowledge& knowledge) const; void emitUnaryVectorOp(pim_binary::Opcode opcode, mlir::Value output, mlir::Value input, - size_t byteSize, const StaticValueKnowledge& knowledge, int32_t r2OrImm = 0, int32_t generic1 = 0) const; - PimCodeGen(PimAcceleratorMemory& memory, - PimInstructionWriter& instructionWriter, - llvm::raw_fd_ostream* coreJson, - const llvm::DenseMap& emittedCoreIds) - : memory(memory), instructionWriter(instructionWriter), coreJsonStream(coreJson), emittedCoreIds(emittedCoreIds) {} + PimCodeGen(PimAcceleratorMemory& memory, PimInstructionWriter& instructionWriter, llvm::raw_fd_ostream* coreJson) + : memory(memory), instructionWriter(instructionWriter), coreJsonStream(coreJson) {} void setBatchLane(std::optional lane) { batchLane = lane; } llvm::FailureOr indexOf(mlir::Value value, const StaticValueKnowledge& knowledge) const { diff --git a/src/PIM/Compiler/PimCompilerOptions.cpp b/src/PIM/Compiler/PimCompilerOptions.cpp index 5e50aa7..dfa5e6c 100644 --- a/src/PIM/Compiler/PimCompilerOptions.cpp +++ b/src/PIM/Compiler/PimCompilerOptions.cpp @@ -125,6 +125,12 @@ llvm::cl::opt coresCount("core-count", llvm::cl::desc("Number of cores in the chip. Required for PIM compilation."), llvm::cl::init(-1)); +llvm::cl::opt pimTargetConfig( + "pim-target-config", + llvm::cl::desc("PIM target configuration used to construct the Spatial scheduling cost model"), + llvm::cl::init(""), + llvm::cl::cat(OnnxMlirOptions)); + llvm::cl::opt ignoreConcatError("ignore-concat-error", llvm::cl::desc("Ignore ConcatOp corner case: do not assert and do a simplification"), diff --git a/src/PIM/Compiler/PimCompilerOptions.hpp b/src/PIM/Compiler/PimCompilerOptions.hpp index 5e58c23..00c2472 100644 --- a/src/PIM/Compiler/PimCompilerOptions.hpp +++ b/src/PIM/Compiler/PimCompilerOptions.hpp @@ -2,6 +2,8 @@ #include "llvm/Support/CommandLine.h" +#include + #define INSTRUMENTSTAGE_ENUM_PIM #define INSTRUMENTSTAGE_CL_ENUM_PIM @@ -63,6 +65,7 @@ extern llvm::cl::opt pimTraceCommunicationMaterialization; extern llvm::cl::opt crossbarSize; extern llvm::cl::opt crossbarCountInCore; extern llvm::cl::opt coresCount; +extern llvm::cl::opt pimTargetConfig; extern llvm::cl::opt pimConvIm2colMaxElements; extern llvm::cl::opt pimConvStreamChunkPositions; diff --git a/src/PIM/Compiler/PimCompilerUtils.cpp b/src/PIM/Compiler/PimCompilerUtils.cpp index ad48d03..9fc8580 100644 --- a/src/PIM/Compiler/PimCompilerUtils.cpp +++ b/src/PIM/Compiler/PimCompilerUtils.cpp @@ -1,9 +1,21 @@ #include "mlir/Conversion/AffineToStandard/AffineToStandard.h" #include "mlir/Transforms/Passes.h" +#include "llvm/Support/Error.h" +#include "llvm/Support/ErrorHandling.h" +#include "llvm/Support/JSON.h" +#include "llvm/Support/MemoryBuffer.h" +#include "llvm/Support/Path.h" + +#include "llvm/ADT/SmallString.h" +#include +#include +#include + #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp" #include "src/Accelerators/PIM/Compiler/PimCompilerUtils.hpp" #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/SchedulingTarget.hpp" #include "src/Accelerators/PIM/Pass/PIMPasses.h" #include "src/Compiler/CompilerPasses.hpp" @@ -14,6 +26,254 @@ using namespace onnx_mlir; namespace onnx_mlir { +namespace { + +void setDefaultPimInterProcessorLatencies( + spatial::SchedulingTarget& target) { + size_t rows = static_cast( + std::sqrt(static_cast(target.processorCount))); + while (rows > 1 && target.processorCount % rows != 0) + --rows; + size_t columns = (target.processorCount + rows - 1) / rows; + + target.interProcessorLatencyNs.assign( + target.processorCount * target.processorCount, 0); + Cost latencySum = 0; + size_t pairCount = 0; + for (size_t source = 0; source < target.processorCount; ++source) { + for (size_t destination = 0; + destination < target.processorCount; ++destination) { + if (source == destination) + continue; + size_t sourceRow = source / columns; + size_t sourceColumn = source % columns; + size_t destinationRow = destination / columns; + size_t destinationColumn = destination % columns; + size_t rowDistance = sourceRow > destinationRow + ? sourceRow - destinationRow + : destinationRow - sourceRow; + size_t columnDistance = sourceColumn > destinationColumn + ? sourceColumn - destinationColumn + : destinationColumn - sourceColumn; + Cost latency = static_cast(2 + rowDistance + columnDistance); + target.interProcessorLatencyNs[ + source * target.processorCount + destination] = latency; + latencySum = checkedAdd(latencySum, latency); + ++pairCount; + } + } + target.averageInterProcessorLatencyNs = + pairCount == 0 + ? 0 + : (latencySum + static_cast(pairCount) - 1) + / static_cast(pairCount); +} + +spatial::SchedulingTarget getDefaultPimSchedulingTarget() { + spatial::SchedulingTarget target; + target.processorCount = static_cast(coresCount.getValue()); + target.residentWeightCapacity = crossbarCountInCore.getValue(); + target.matrixRows = crossbarSize.getValue(); + target.matrixColumns = crossbarSize.getValue(); + + setDefaultPimInterProcessorLatencies(target); + return target; +} + +const llvm::json::Object& requireObject(const llvm::json::Object& object, + llvm::StringRef key, + llvm::StringRef path) { + const llvm::json::Object* nested = object.getObject(key); + if (!nested) + llvm::report_fatal_error("PIM target config is missing object '" + path + "." + key + "'"); + return *nested; +} + +Cost getConfigCost(const llvm::json::Object& object, + llvm::StringRef key, + Cost fallback, + bool allowZero = false) { + std::optional number = object.getNumber(key); + if (!number) + return fallback; + if (!std::isfinite(*number) || *number < 0.0 || (!allowZero && *number == 0.0) + || *number > static_cast(std::numeric_limits::max())) + llvm::report_fatal_error("PIM target config field '" + key + "' must be a valid positive number"); + return static_cast(std::ceil(*number)); +} + +std::pair getConfigPair(const llvm::json::Object& object, + llvm::StringRef key) { + const llvm::json::Array* values = object.getArray(key); + if (!values || values->size() != 2) + llvm::report_fatal_error("PIM target config field '" + key + "' must contain two integers"); + std::optional first = (*values)[0].getAsInteger(); + std::optional second = (*values)[1].getAsInteger(); + if (!first || !second || *first <= 0 || *second <= 0) + llvm::report_fatal_error("PIM target config field '" + key + "' must contain two positive integers"); + return {static_cast(*first), static_cast(*second)}; +} + +void loadPimInterProcessorLatencies( + spatial::SchedulingTarget& target, + const llvm::json::Object& network) { + std::optional filename = + network.getString("net_config_file_path"); + if (!filename) + llvm::report_fatal_error( + "PIM target config is missing network latency file path"); + + llvm::SmallString<256> networkPath(*filename); + if (!llvm::sys::path::is_absolute(networkPath)) { + llvm::SmallString<256> configDirectory(pimTargetConfig.getValue()); + llvm::sys::path::remove_filename(configDirectory); + llvm::sys::path::append(configDirectory, networkPath); + networkPath = configDirectory; + } + + auto buffer = llvm::MemoryBuffer::getFile(networkPath); + if (!buffer) + llvm::report_fatal_error( + llvm::Twine("failed to read PIM network config '") + + networkPath + "': " + buffer.getError().message()); + auto parsed = llvm::json::parse(buffer.get()->getBuffer()); + if (!parsed) + llvm::report_fatal_error( + llvm::Twine("failed to parse PIM network config '") + + networkPath + "': " + llvm::toString(parsed.takeError())); + const llvm::json::Object* root = parsed->getAsObject(); + const llvm::json::Object* latencies = + root ? root->getObject("latency") : nullptr; + if (!latencies) + llvm::report_fatal_error( + "PIM network config is missing its latency matrix"); + + target.interProcessorLatencyNs.assign( + target.processorCount * target.processorCount, 0); + Cost latencySum = 0; + size_t pairCount = 0; + for (size_t source = 0; source < target.processorCount; ++source) { + std::string sourceKey = std::to_string(source); + const llvm::json::Object* row = latencies->getObject(sourceKey); + if (!row) + llvm::report_fatal_error( + llvm::Twine("PIM network config is missing latency row ") + + sourceKey); + for (size_t destination = 0; + destination < target.processorCount; ++destination) { + if (source == destination) + continue; + std::string destinationKey = std::to_string(destination); + std::optional latency = row->getNumber(destinationKey); + if (!latency || !std::isfinite(*latency) || *latency <= 0.0) + llvm::report_fatal_error( + llvm::Twine("PIM network config is missing latency ") + + sourceKey + " -> " + destinationKey); + Cost roundedLatency = static_cast(std::ceil(*latency)); + target.interProcessorLatencyNs[ + source * target.processorCount + destination] = roundedLatency; + latencySum = checkedAdd(latencySum, roundedLatency); + ++pairCount; + } + } + target.averageInterProcessorLatencyNs = + pairCount == 0 + ? 0 + : (latencySum + static_cast(pairCount) - 1) + / static_cast(pairCount); +} + +spatial::SchedulingTarget getPimSchedulingTarget() { + spatial::SchedulingTarget target = getDefaultPimSchedulingTarget(); + if (pimTargetConfig.empty()) + return target; + + auto buffer = llvm::MemoryBuffer::getFile(pimTargetConfig); + if (!buffer) + llvm::report_fatal_error( + llvm::Twine("failed to read PIM target config '") + + pimTargetConfig.getValue() + "': " + buffer.getError().message()); + auto parsed = llvm::json::parse(buffer.get()->getBuffer()); + if (!parsed) + llvm::report_fatal_error( + llvm::Twine("failed to parse PIM target config '") + + pimTargetConfig.getValue() + "': " + + llvm::toString(parsed.takeError())); + const llvm::json::Object* root = parsed->getAsObject(); + if (!root) + llvm::report_fatal_error("PIM target config must contain a JSON object"); + + const llvm::json::Object& chip = requireObject(*root, "chip_config", "root"); + const llvm::json::Object& core = requireObject(chip, "core_config", "chip_config"); + const llvm::json::Object& matrix = + requireObject(core, "matrix_config", "chip_config.core_config"); + const llvm::json::Object& localMemory = + requireObject(core, "local_memory_config", "chip_config.core_config"); + const llvm::json::Object& network = + requireObject(chip, "network_config", "chip_config"); + + std::optional coreCount = chip.getInteger("core_cnt"); + if (!coreCount || *coreCount <= 0) + llvm::report_fatal_error("PIM target config field 'core_cnt' must be a positive integer"); + target.processorCount = static_cast(*coreCount); + target.residentWeightCapacity = + getConfigCost(matrix, "xbar_array_count", target.residentWeightCapacity); + std::tie(target.matrixRows, target.matrixColumns) = + getConfigPair(matrix, "xbar_size"); + + if (target.processorCount != static_cast(coresCount.getValue()) + || target.residentWeightCapacity != crossbarCountInCore.getValue() + || target.matrixRows != crossbarSize.getValue() + || target.matrixColumns != crossbarSize.getValue()) + llvm::report_fatal_error("PIM target config resources do not match --core-count, " + "--crossbar-count, and --crossbar-size"); + loadPimInterProcessorLatencies(target, network); + + target.processorPeriodNs = + getConfigCost(core, "period", target.processorPeriodNs); + target.localMemoryWidthBytes = + getConfigCost(localMemory, "data_width", target.localMemoryWidthBytes); + target.localMemoryReadLatencyCycles = + getConfigCost(localMemory, "read_latency_cycle", target.localMemoryReadLatencyCycles); + target.localMemoryWriteLatencyCycles = + getConfigCost(localMemory, "write_latency_cycle", target.localMemoryWriteLatencyCycles); + target.transferWidthBytes = + getConfigCost(network, "bus_width", target.transferWidthBytes); + target.vectorWidth = getConfigCost(core, "vector_width", target.vectorWidth); + target.vectorLatencyCycles = + getConfigCost(core, "vector_latency_cycle", target.vectorLatencyCycles); + + target.matrixPeriodNs = + getConfigCost(matrix, "period", target.matrixPeriodNs); + target.matrixInputResolutionBits = + getConfigCost(matrix, "dac_resolution", target.matrixInputResolutionBits); + target.matrixInputLatencyCycles = + getConfigCost(matrix, "dac_latency_cycle", target.matrixInputLatencyCycles); + target.matrixInputParallelism = + getConfigCost(matrix, "dac_count", target.matrixInputParallelism); + target.matrixReadLatencyNs = + getConfigCost(matrix, "xbar_latency", target.matrixReadLatencyNs); + target.matrixSampleLatencyCycles = + getConfigCost(matrix, "sample_hold_latency_cycle", target.matrixSampleLatencyCycles); + target.matrixOutputLatencyCycles = + getConfigCost(matrix, "adc_latency_cycle", target.matrixOutputLatencyCycles); + target.matrixOutputParallelism = + getConfigCost(matrix, "adc_count", target.matrixOutputParallelism); + target.matrixShiftLatencyCycles = + getConfigCost(matrix, "shift_adder_latency_cycle", target.matrixShiftLatencyCycles); + target.matrixBufferLatencyCycles = + getConfigCost(matrix, "output_buffer_latency_cycle", target.matrixBufferLatencyCycles); + target.matrixInputBufferLatencyCycles = + getConfigCost(matrix, + "input_buffer_latency_cycle", + target.matrixInputBufferLatencyCycles, + /*allowZero=*/true); + target.matrixPipeline = matrix.getBoolean("pipeline_mode").value_or(target.matrixPipeline); + return target; +} + +} // namespace + void addPassesPim(OwningOpRef& module, PassManager& pm, EmissionTargetType& emissionTarget, @@ -31,11 +291,13 @@ void addPassesPim(OwningOpRef& module, addONNXToMLIRPasses(pm, /*target CPU*/ false); if (pimEmissionTarget >= EmitSpatial) { + spatial::SchedulingTarget schedulingTarget = getPimSchedulingTarget(); pm.addPass(createONNXToSpatialPass()); pm.addPass(createSpatialLayoutPlanningPass()); pm.addPass(createLowerSpatialPlansPass()); - pm.addPass(createTrivialGraphComputeMergePass()); - pm.addPass(createMergeComputeNodesPass()); + pm.addPass(createTrivialGraphComputeMergePass( + schedulingTarget.residentWeightCapacity)); + pm.addPass(createMergeComputeNodesPass(schedulingTarget)); pm.addPass(createMessagePass("Onnx lowered to Spatial")); } diff --git a/src/PIM/Compiler/PimWeightEmitter.cpp b/src/PIM/Compiler/PimWeightEmitter.cpp index f85bce7..7255166 100644 --- a/src/PIM/Compiler/PimWeightEmitter.cpp +++ b/src/PIM/Compiler/PimWeightEmitter.cpp @@ -42,7 +42,9 @@ WeightEmissionResult createAndPopulateWeightFolder(ArrayRef r assert(isMatrixShape(shape) && "Weight matrix must be 2-dimensional"); int64_t numRows = shape[0]; int64_t numCols = shape[1]; - assert(numRows <= xbarSize && numCols <= xbarSize && "Weight dimensions must not exceed crossbar size"); + assert(numRows <= xbarSize && numCols % xbarSize == 0 + && numCols / xbarSize <= static_cast(crossbarCountInCore) + && "Weight dimensions must fit in one array group"); size_t elementByteWidth = getElementTypeSizeInBytes(denseAttr.getElementType()); @@ -57,7 +59,7 @@ WeightEmissionResult createAndPopulateWeightFolder(ArrayRef r uint64_t zero = 0; for (int64_t row = 0; row < xbarSize; row++) { - for (int64_t col = 0; col < xbarSize; col++) { + for (int64_t col = 0; col < numCols; col++) { if (row < numRows && col < numCols) { int64_t elementIndex = weightView.offset + row * weightView.strides[0] + col * weightView.strides[1]; APInt bits = denseAttr.getValues()[elementIndex].bitcastToAPInt(); @@ -73,7 +75,7 @@ WeightEmissionResult createAndPopulateWeightFolder(ArrayRef r weightFileStream.close(); materializedWeights.push_back({weightView, newFileName}); uint64_t weightBytes = pim::checkedMulOrCrash( - pim::checkedMulOrCrash(static_cast(xbarSize), static_cast(xbarSize), "weight element count"), + pim::checkedMulOrCrash(static_cast(xbarSize), static_cast(numCols), "weight element count"), elementByteWidth, "weight byte size"); result.totalWeightBytes = pim::checkedAddOrCrash(result.totalWeightBytes, weightBytes, "total weight bytes"); diff --git a/src/PIM/Conversion/ONNXToSpatial/CMakeLists.txt b/src/PIM/Conversion/ONNXToSpatial/CMakeLists.txt index 60e9e44..6d9395c 100644 --- a/src/PIM/Conversion/ONNXToSpatial/CMakeLists.txt +++ b/src/PIM/Conversion/ONNXToSpatial/CMakeLists.txt @@ -8,7 +8,6 @@ add_pim_library(OMONNXToSpatial ONNXToSpatialVerifier.cpp Patterns/Pre.cpp Patterns/Post.cpp - Patterns/GeneratedConversion.cpp Patterns/Math/Conv.cpp Patterns/Math/ConvGeometry.cpp Patterns/Math/Elementwise.cpp diff --git a/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp b/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp index ccaf398..e198b77 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp @@ -19,9 +19,11 @@ FailureOr describeRowStripPhysicalValue(Value storage, Ra || storageType.getRank() != 5 || logicalType.getRank() != 4 || logicalType.getDimSize(0) != 1 || storageType.getElementType() != logicalType.getElementType() || storageType.getDimSize(1) != 1 || storageType.getDimSize(2) != 1 - || storageType.getDimSize(3) != logicalType.getDimSize(3) || storageType.getDimSize(4) <= 0) + || storageType.getDimSize(3) != logicalType.getDimSize(3) + || storageType.getDimSize(4) <= 0) return failure(); - const int64_t tilesPerRow = ceilIntegerDivide(logicalType.getDimSize(1), storageType.getDimSize(4)); + const int64_t tilesPerRow = + ceilIntegerDivide(logicalType.getDimSize(1), storageType.getDimSize(4)); if (storageType.getDimSize(0) != logicalType.getDimSize(2) * tilesPerRow) return failure(); return RowStripPhysicalValue {storage, logicalType, @@ -249,4 +251,111 @@ FailureOr applyRowStripBiasAdd(const RowStripPhysicalValue& value, return batchOp->getResult(0); } +FailureOr applyRowStripAdd(const RowStripPhysicalValue& lhs, + const RowStripPhysicalValue& rhs, + PatternRewriter& rewriter, + Location loc) { + if (lhs.logicalType != rhs.logicalType || lhs.fragmentType != rhs.fragmentType + || lhs.storage.getType() != rhs.storage.getType() || lhs.tilesPerRow != rhs.tilesPerRow) + return failure(); + auto storageType = cast(lhs.storage.getType()); + const int64_t laneCount = storageType.getDimSize(0); + auto batch = createSpatComputeBatch( + rewriter, + loc, + TypeRange {storageType}, + laneCount, + {}, + ValueRange {lhs.storage, rhs.storage}, + [&](detail::SpatComputeBatchBodyArgs args) { + FailureOr lhsFragment = + extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[0], args.lane, lhs.fragmentType); + FailureOr rhsFragment = + extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[1], args.lane, rhs.fragmentType); + if (failed(lhsFragment) || failed(rhsFragment)) + return failure(); + Value added = spatial::SpatVAddOp::create(rewriter, loc, lhs.fragmentType, *lhsFragment, *rhsFragment); + publishGraphBatchPhysicalFragment(rewriter, loc, added, args.outputs.front(), args.lane); + return success(); + }); + if (failed(batch)) + return failure(); + return batch->getResult(0); +} + +FailureOr applyRowStripConcat(ArrayRef inputs, + RankedTensorType outputType, + PatternRewriter& rewriter, + Location loc) { + if (inputs.empty() || !outputType || !outputType.hasStaticShape() || outputType.getRank() != 4 + || outputType.getDimSize(0) != 1) + return failure(); + int64_t channels = 0; + for (const RowStripPhysicalValue& input : inputs) { + if (input.logicalType.getElementType() != outputType.getElementType() + || input.logicalType.getDimSize(0) != outputType.getDimSize(0) + || input.logicalType.getDimSize(2) != outputType.getDimSize(2) + || input.logicalType.getDimSize(3) != outputType.getDimSize(3)) + return failure(); + channels += input.logicalType.getDimSize(1); + } + if (channels != outputType.getDimSize(1)) + return failure(); + + SmallVector storages; + llvm::transform( + inputs, std::back_inserter(storages), [](const RowStripPhysicalValue& input) { return input.storage; }); + const int64_t tileWidth = outputType.getDimSize(3); + auto fragmentType = getRowStripFragmentType(outputType); + auto storageType = getRowStripStorageType(outputType); + auto batch = createSpatComputeBatch( + rewriter, + loc, + TypeRange {storageType}, + outputType.getDimSize(2), + {}, + storages, + [&](detail::SpatComputeBatchBodyArgs args) { + Operation* anchor = rewriter.getInsertionBlock()->getParentOp(); + SmallVector fragments; + for (auto [inputIndex, input] : llvm::enumerate(inputs)) { + Value tileStart = affineMulConst( + rewriter, loc, args.lane, input.tilesPerRow, anchor); + for (int64_t tile = 0; tile < input.tilesPerRow; ++tile) { + Value slot = + affineAddConst(rewriter, loc, tileStart, tile, anchor); + FailureOr fragment = + extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[inputIndex], slot, input.fragmentType); + if (failed(fragment)) + return failure(); + int64_t channelOffset = tile * input.fragmentType.getDimSize(3); + int64_t validChannels = + std::min(input.fragmentType.getDimSize(3), input.logicalType.getDimSize(1) - channelOffset); + auto validType = + RankedTensorType::get( + {1, 1, tileWidth, validChannels}, + outputType.getElementType()); + MixedSliceGeometry slice; + slice.offsets.assign(4, rewriter.getIndexAttr(0)); + slice.sizes = {rewriter.getIndexAttr(1), + rewriter.getIndexAttr(1), + rewriter.getIndexAttr(tileWidth), + rewriter.getIndexAttr(validChannels)}; + slice.strides.assign(4, rewriter.getIndexAttr(1)); + Value valid = extractMixedSliceOrIdentity(rewriter, loc, *fragment, validType, slice); + if (!valid) + return failure(); + fragments.push_back(valid); + } + } + Value concatenated = + spatial::SpatConcatOp::create(rewriter, loc, fragmentType, rewriter.getI64IntegerAttr(3), fragments); + publishGraphBatchPhysicalFragment(rewriter, loc, concatenated, args.outputs.front(), args.lane); + return success(); + }); + if (failed(batch)) + return failure(); + return batch->getResult(0); +} + } // namespace onnx_mlir diff --git a/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp b/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp index e8c24a6..d66cbc8 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp +++ b/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp @@ -66,4 +66,14 @@ mlir::FailureOr applyRowStripBiasAdd(const RowStripPhysicalValue& v mlir::PatternRewriter& rewriter, mlir::Location loc); +mlir::FailureOr applyRowStripAdd(const RowStripPhysicalValue& lhs, + const RowStripPhysicalValue& rhs, + mlir::PatternRewriter& rewriter, + mlir::Location loc); + +mlir::FailureOr applyRowStripConcat(llvm::ArrayRef inputs, + mlir::RankedTensorType outputType, + mlir::PatternRewriter& rewriter, + mlir::Location loc); + } // namespace onnx_mlir diff --git a/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp b/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp index 12f6727..5c8795b 100644 --- a/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp @@ -64,6 +64,22 @@ static FailureOr lowerRowStripBiasAdd(const RowStripPhysicalValue& input, return applyRowStripBiasAdd(input, planOp.getBias(), rewriter, planOp.getLoc()); } +static FailureOr lowerRowStripAdd(const RowStripPhysicalValue& lhs, + const RowStripPhysicalValue& rhs, + spatial::SpatAddPlanOp planOp, + PatternRewriter& rewriter) { + return applyRowStripAdd(lhs, rhs, rewriter, planOp.getLoc()); +} + +static FailureOr lowerRowStripConcat(ArrayRef inputs, + spatial::SpatConcatPlanOp planOp, + PatternRewriter& rewriter) { + auto outputType = dyn_cast(planOp.getOutput().getType()); + if (!outputType) + return failure(); + return applyRowStripConcat(inputs, outputType, rewriter, planOp.getLoc()); +} + static FailureOr materializeRowStripToDense(const RowStripPhysicalValue& rowStripValue, Location loc, PatternRewriter& rewriter) { if (rowStripValue.logicalType.getRank() != 4 || !rowStripValue.logicalType.hasStaticShape()) @@ -122,6 +138,99 @@ static FailureOr lowerDenseBatchBiasAdd(Value input, Value bias, RankedTe return batch->getResult(0); } +static LogicalResult lowerAddPlan(spatial::SpatAddPlanOp planOp, + llvm::DenseMap& rowStripValues, + llvm::SmallPtrSetImpl& eraseAfterLowering, + PatternRewriter& rewriter) { + FailureOr lhs = getRowStripValue(rowStripValues, planOp.getLhs()); + FailureOr rhs = getRowStripValue(rowStripValues, planOp.getRhs()); + if (succeeded(lhs) && succeeded(rhs)) { + auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) { + auto blueprint = dyn_cast(user); + return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout; + }); + if (outputBlueprint == planOp.getResult().getUsers().end()) + return planOp.emitOpError("row-strip add plan requires a row-strip blueprint result"); + + rewriter.setInsertionPoint(planOp); + FailureOr lowered = lowerRowStripAdd(*lhs, *rhs, planOp, rewriter); + if (failed(lowered)) + return planOp.emitOpError("failed to lower selected row-strip Spatial add plan"); + auto blueprint = cast(*outputBlueprint); + FailureOr output = buildRowStripValue(blueprint, *lowered); + if (failed(output)) + return failure(); + rowStripValues[blueprint.getResult()] = *output; + eraseAfterLowering.insert(planOp); + eraseAfterLowering.insert(blueprint); + return success(); + } + + rewriter.setInsertionPoint(planOp); + auto compute = createSpatCompute<2>(rewriter, + planOp.getLoc(), + planOp.getOutput().getType(), + {}, + ValueRange {planOp.getLhs(), planOp.getRhs()}, + [&](Value lhsValue, Value rhsValue) { + Value added = spatial::SpatVAddOp::create( + rewriter, planOp.getLoc(), planOp.getOutput().getType(), lhsValue, rhsValue); + spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), added); + }); + rewriter.replaceOp(planOp, compute.getResults()); + return success(); +} + +static LogicalResult lowerConcatPlan(spatial::SpatConcatPlanOp planOp, + llvm::DenseMap& rowStripValues, + llvm::SmallPtrSetImpl& eraseAfterLowering, + PatternRewriter& rewriter) { + SmallVector inputs; + for (Value input : planOp.getInputs()) { + FailureOr physical = getRowStripValue(rowStripValues, input); + if (failed(physical)) { + inputs.clear(); + break; + } + inputs.push_back(*physical); + } + if (inputs.size() == planOp.getInputs().size()) { + auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) { + auto blueprint = dyn_cast(user); + return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout; + }); + if (outputBlueprint == planOp.getResult().getUsers().end()) + return planOp.emitOpError("row-strip concat plan requires a row-strip blueprint result"); + rewriter.setInsertionPoint(planOp); + FailureOr lowered = lowerRowStripConcat(inputs, planOp, rewriter); + if (failed(lowered)) + return planOp.emitOpError("failed to lower selected row-strip Spatial concat plan"); + auto blueprint = cast(*outputBlueprint); + FailureOr output = buildRowStripValue(blueprint, *lowered); + if (failed(output)) + return failure(); + rowStripValues[blueprint.getResult()] = *output; + eraseAfterLowering.insert(planOp); + eraseAfterLowering.insert(blueprint); + return success(); + } + + rewriter.setInsertionPoint(planOp); + auto compute = createSpatCompute( + rewriter, + planOp.getLoc(), + TypeRange {planOp.getOutput().getType()}, + {}, + planOp.getInputs(), + [&](ValueRange values) { + Value concatenated = spatial::SpatConcatOp::create( + rewriter, planOp.getLoc(), planOp.getOutput().getType(), rewriter.getI64IntegerAttr(planOp.getAxis()), values); + spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), concatenated); + }); + rewriter.replaceOp(planOp, compute.getResults()); + return success(); +} + struct LowerSpatialPlansPass final : PassWrapper> { MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerSpatialPlansPass) @@ -274,6 +383,40 @@ struct LowerSpatialPlansPass final : PassWrapper(&op)) { + auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) { + auto blueprint = dyn_cast(user); + return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout; + }); + if (outputBlueprint == planOp.getResult().getUsers().end()) { + planOp.emitOpError("selected global AveragePool plan requires a row-strip blueprint result"); + signalPassFailure(); + return; + } + + FailureOr input = getRowStripValue(rowStripValues, planOp.getInput()); + rewriter.setInsertionPoint(planOp); + std::optional physicalInput; + if (succeeded(input)) + physicalInput = input->storage; + FailureOr lowered = + lowerSelectedGlobalAveragePoolPlan(planOp, physicalInput, rewriter); + if (failed(lowered)) { + planOp.emitOpError("failed to lower selected row-strip Spatial global AveragePool plan"); + signalPassFailure(); + return; + } + auto blueprint = cast(*outputBlueprint); + FailureOr output = buildRowStripValue(blueprint, *lowered); + if (failed(output)) { + signalPassFailure(); + return; + } + rowStripValues[blueprint.getResult()] = *output; + eraseAfterLowering.insert(planOp); + eraseAfterLowering.insert(blueprint); + continue; + } if (auto planOp = dyn_cast(&op)) { if (succeeded(getRowStripValue(rowStripValues, planOp.getInput()))) { auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) { @@ -339,6 +482,20 @@ struct LowerSpatialPlansPass final : PassWrapper(&op)) { + if (failed(lowerAddPlan(planOp, rowStripValues, eraseAfterLowering, rewriter))) { + signalPassFailure(); + return; + } + continue; + } + if (auto planOp = dyn_cast(&op)) { + if (failed(lowerConcatPlan(planOp, rowStripValues, eraseAfterLowering, rewriter))) { + signalPassFailure(); + return; + } + continue; + } if (auto flattenOp = dyn_cast(&op)) { if (flattenOp.getInputs().size() == 1) { FailureOr input = @@ -488,12 +645,15 @@ struct LowerSpatialPlansPass final : PassWrapperemitOpError("planning blueprint must not remain after LowerSpatialPlans"); hasIllegalOps = true; - } else if (isa(op) - || op->getDialect()->getNamespace() == "onnx") { + } + else if (isa(op) + || op->getDialect()->getNamespace() == "onnx") { op->emitOpError("operation must not remain after LowerSpatialPlans"); hasIllegalOps = true; } diff --git a/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatial.td b/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatial.td index 1789e45..ab7d8e9 100644 --- a/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatial.td +++ b/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatial.td @@ -60,11 +60,6 @@ def convAddToConvWithBiasRight : Pat< def replaceWithOperationOfValue : NativeCodeCall<"$0">; -def removeLRN : Pat< - (ONNXLRNOp $A, $_, $_, $_, $_), - (replaceWithOperationOfValue $A) ->; - def HaveSameStaticShape: Constraint< CPred<"onnx_mlir::haveSameStaticShape($0, $1)">, "Two tensors have the same static shape">; diff --git a/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp b/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp index f9a7637..7bc8445 100644 --- a/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp @@ -47,12 +47,17 @@ static void populateEmptyFunction(func::FuncOp funcOp) { SmallVector computeBatches(funcOp.getOps()); SmallVector convPlans(funcOp.getOps()); SmallVector biasAddPlans(funcOp.getOps()); + SmallVector addPlans(funcOp.getOps()); + SmallVector concatPlans(funcOp.getOps()); SmallVector reluPlans(funcOp.getOps()); SmallVector maxPoolPlans(funcOp.getOps()); + SmallVector globalAveragePoolPlans( + funcOp.getOps()); SmallVector blueprints(funcOp.getOps()); SmallVector materializers(funcOp.getOps()); - if (!computes.empty() || !computeBatches.empty() || !convPlans.empty() || !biasAddPlans.empty() || !reluPlans.empty() - || !maxPoolPlans.empty() || !blueprints.empty() || !materializers.empty()) { + if (!computes.empty() || !computeBatches.empty() || !convPlans.empty() || !biasAddPlans.empty() || !addPlans.empty() + || !concatPlans.empty() || !reluPlans.empty() || !maxPoolPlans.empty() || !blueprints.empty() + || !globalAveragePoolPlans.empty() || !materializers.empty()) { return; } diff --git a/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.cpp b/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.cpp index d1230ab..2f60fd8 100644 --- a/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.cpp @@ -146,8 +146,11 @@ void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter spatial::SpatGraphComputeBatch, spatial::SpatConv2DPlanOp, spatial::SpatBiasAddPlanOp, + spatial::SpatAddPlanOp, + spatial::SpatConcatPlanOp, spatial::SpatReluPlanOp, spatial::SpatMaxPool2DPlanOp, + spatial::SpatGlobalAveragePoolPlanOp, spatial::SpatBlueprintOp, spatial::SpatMaterializeLayoutOp>(&op)) { continue; diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns.cpp index 1abe958..b2106fd 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns.cpp @@ -8,7 +8,6 @@ namespace onnx_mlir { void populatePrePatterns(RewritePatternSet& patterns, MLIRContext* ctx) { populateGeneratedPrePatterns(patterns, ctx); } void populateConversionPatterns(RewritePatternSet& patterns, MLIRContext* ctx) { - populateGeneratedConversionPatterns(patterns, ctx); populateElementwisePatterns(patterns, ctx); populateMatMulRewritePatterns(patterns, ctx); populateGemmPatterns(patterns, ctx); diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns.hpp b/src/PIM/Conversion/ONNXToSpatial/Patterns.hpp index da2a7a6..0a0e6aa 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns.hpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns.hpp @@ -13,7 +13,6 @@ void populateConversionPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRCon void populatePostPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx); void populateGeneratedPrePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx); -void populateGeneratedConversionPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx); void populateWeightPromotionPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx); void populateConvPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx); diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/GeneratedConversion.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/GeneratedConversion.cpp deleted file mode 100644 index 5fc96a2..0000000 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/GeneratedConversion.cpp +++ /dev/null @@ -1,18 +0,0 @@ -#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp" -#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp" - -using namespace mlir; - -namespace onnx_mlir { - -namespace { - -#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatial.hpp.inc" - -} // namespace - -void populateGeneratedConversionPatterns(RewritePatternSet& patterns, MLIRContext* ctx) { - patterns.add(ctx); -} - -} // namespace onnx_mlir diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp index 8b9c076..5fcdb35 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp @@ -1242,8 +1242,9 @@ static Value buildPackedWeights(DenseElementsAttr wDenseAttr, const Tiling& tiling, PatternRewriter& rewriter, Location loc) { + const int64_t paddedOutputChannels = static_cast(crossbarSize.getValue()); auto packedWeightType = RankedTensorType::get( - {tiling.numChannelTiles, tiling.tileInputRows, tiling.tileOutputChannels}, wType.getElementType()); + {tiling.numChannelTiles, tiling.tileInputRows, paddedOutputChannels}, wType.getElementType()); SmallVector packedValues(packedWeightType.getNumElements(), cast(rewriter.getZeroAttr(wType.getElementType()))); SmallVector sourceValues(wDenseAttr.getValues()); @@ -1262,7 +1263,7 @@ static Value buildPackedWeights(DenseElementsAttr wDenseAttr, ((globalOutChannel * wType.getDimSize(1) * wType.getDimSize(2)) + kernelH) * wType.getDimSize(3) + kernelW; const int64_t targetCol = localChannel * tiling.outputMultiplier + multiplierIndex; const int64_t targetFlatIndex = - ((tileIndex * tiling.tileInputRows) + targetRow) * tiling.tileOutputChannels + targetCol; + ((tileIndex * tiling.tileInputRows) + targetRow) * paddedOutputChannels + targetCol; packedValues[targetFlatIndex] = sourceValues[sourceFlatIndex]; } } @@ -1353,11 +1354,12 @@ static Value createWeightTile(Value packedWeights, PatternRewriter& rewriter, Location loc) { SmallVector offsets {channelTileIndex, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}; + const int64_t paddedOutputChannels = static_cast(crossbarSize.getValue()); SmallVector sizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(tiling.tileInputRows), - rewriter.getIndexAttr(tiling.tileOutputChannels)}; + rewriter.getIndexAttr(paddedOutputChannels)}; auto collapsedType = - RankedTensorType::get({tiling.tileInputRows, tiling.tileOutputChannels}, packedWeightType.getElementType()); + RankedTensorType::get({tiling.tileInputRows, paddedOutputChannels}, packedWeightType.getElementType()); return extractMixedSliceOrIdentity( rewriter, loc, packedWeights, collapsedType, {offsets, sizes, getUnitStrides(rewriter, 3)}); @@ -1547,6 +1549,8 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter& auto gemmOutType = RankedTensorType::get({tiling->totalPatches, state.outType.getDimSize(1)}, state.outType.getElementType()); auto rowTileType = RankedTensorType::get({1, tiling->tileOutputChannels}, state.outType.getElementType()); + auto paddedRowTileType = RankedTensorType::get( + {1, static_cast(crossbarSize.getValue())}, state.outType.getElementType()); auto piecesType = spatial::getGraphBatchPhysicalResultType( tiling->totalPatches * tiling->numChannelTiles, rowTileType); auto paddedInputType = cast(paddedInput.getType()); @@ -1617,7 +1621,17 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter& *tiling, rewriter, loc); - Value rowTile = spatial::SpatVMMOp::create(rewriter, loc, rowTileType, weightTile, inputTile).getResult(); + Value paddedRowTile = + spatial::SpatVMMOp::create(rewriter, loc, paddedRowTileType, weightTile, inputTile).getResult(); + Value rowTile = tensor::ExtractSliceOp::create( + rewriter, + loc, + rowTileType, + paddedRowTile, + SmallVector {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}, + SmallVector {rewriter.getIndexAttr(1), + rewriter.getIndexAttr(tiling->tileOutputChannels)}, + getUnitStrides(rewriter, 2)); if (args.inputs.size() > 1) { Value biasArg = pickInputByRank(/*rank=*/2); if (!biasArg) { @@ -2493,6 +2507,7 @@ static Value createZeroGemmBias(RankedTensorType gemmResultType, PatternRewriter static bool rowStripOutputTileFitsOneCore(const ConvGeometry& geometry) { return ceilIntegerDivide(geometry.k, geometry.xbarSize) + * ceilIntegerDivide(geometry.c, geometry.xbarSize) <= static_cast(crossbarCountInCore.getValue()); } @@ -2521,28 +2536,6 @@ static bool canConsumePixelMajorRowStripFragments(const ConvLoweringState& state failureReason = "dilation_not_one"; return false; } - const bool pointwise = state.xHeight == 1 && state.xWidth == 1 && state.outHeight == 1 && state.outWidth == 1 - && state.wHeight == 1 && state.wWidth == 1 && state.padHeightBegin == 0 - && state.padHeightEnd == 0 && state.padWidthBegin == 0 && state.padWidthEnd == 0; - if (pointwise) { - if (!getHostConstDenseElementsAttr(state.w)) { - failureReason = "non_constant_weight"; - return false; - } - if (state.hasBias && !isSupportedBiasAddValue(state.b, state.outType)) { - failureReason = "unsupported_bias"; - return false; - } - return true; - } - if (state.wHeight != 3 || state.wWidth != 3) { - failureReason = "kernel_not_3x3"; - return false; - } - if (state.padHeightBegin != 1 || state.padHeightEnd != 1 || state.padWidthBegin != 1 || state.padWidthEnd != 1) { - failureReason = "padding_not_1"; - return false; - } if (state.outHeight != state.xHeight || state.outWidth != state.xWidth) { failureReason = "not_same_spatial_shape"; return false; @@ -3043,57 +3036,154 @@ static FailureOr createConvOutputRow(ValueRange inputTiles, Location loc) { auto elementType = cast(inputTiles.front().getType()).getElementType(); auto rowType = RankedTensorType::get({1, outputChannels}, elementType); - auto tileWeightsType = - RankedTensorType::get({paddedK, xbarDim}, - cast(paddedWeights.getType()).getElementType()); const int64_t outputTileCount = ceilIntegerDivide(outputChannels, xbarDim); - - auto getTileWeights = [&](int64_t outputTile) { - if (outputTileCount == 1) - return paddedWeights; - SmallVector offsets { - rewriter.getIndexAttr(outputTile), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}; - SmallVector sizes { - rewriter.getIndexAttr(1), rewriter.getIndexAttr(paddedK), rewriter.getIndexAttr(xbarDim)}; - return extractStaticSliceOrIdentity( - rewriter, loc, paddedWeights, tileWeightsType, offsets, sizes, getUnitStrides(rewriter, 3)); - }; - - if (outputTileCount == 1) { - FailureOr rowResult = createConvOutputTile( - inputTiles, getTileWeights(0), outputChannels, xbarDim, rewriter, loc); - if (failed(rowResult)) - return failure(); - Value validRow = *rowResult; - if (bias) - validRow = spatial::SpatVAddOp::create(rewriter, loc, rowType, validRow, bias).getResult(); - return validRow; - } - const int64_t paddedOutputChannels = outputTileCount * xbarDim; auto paddedOutputType = RankedTensorType::get({1, paddedOutputChannels}, elementType); - Value paddedOutput = tensor::EmptyOp::create(rewriter, loc, paddedOutputType.getShape(), elementType); - for (int64_t outputTile = 0; outputTile < outputTileCount; ++outputTile) { - FailureOr tileResult = createConvOutputTile( - inputTiles, getTileWeights(outputTile), xbarDim, xbarDim, rewriter, loc); - if (failed(tileResult)) - return failure(); - SmallVector tileOffsets { - rewriter.getIndexAttr(0), rewriter.getIndexAttr(outputTile * xbarDim)}; - SmallVector tileSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim)}; - paddedOutput = tensor::InsertSliceOp::create( - rewriter, loc, *tileResult, paddedOutput, tileOffsets, tileSizes, getUnitStrides(rewriter, 2)); + auto weightSliceType = RankedTensorType::get( + {xbarDim, paddedOutputChannels}, + cast(paddedWeights.getType()).getElementType()); + + Value paddedOutput; + for (auto [kSlice, inputTile] : llvm::enumerate(inputTiles)) { + const int64_t kOffset = static_cast(kSlice) * xbarDim; + Value weightSlice = extractStaticSliceOrIdentity( + rewriter, + loc, + paddedWeights, + weightSliceType, + SmallVector {rewriter.getIndexAttr(kOffset), rewriter.getIndexAttr(0)}, + SmallVector {rewriter.getIndexAttr(xbarDim), + rewriter.getIndexAttr(paddedOutputChannels)}, + getUnitStrides(rewriter, 2)); + Value piece = + spatial::SpatVMMOp::create(rewriter, loc, paddedOutputType, weightSlice, inputTile).getResult(); + paddedOutput = paddedOutput + ? spatial::SpatVAddOp::create( + rewriter, loc, paddedOutputType, paddedOutput, piece).getResult() + : piece; } - SmallVector outputOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}; - SmallVector outputSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(outputChannels)}; - Value validRow = tensor::ExtractSliceOp::create( - rewriter, loc, rowType, paddedOutput, outputOffsets, outputSizes, getUnitStrides(rewriter, 2)); + Value validRow = outputChannels == paddedOutputChannels + ? paddedOutput + : tensor::ExtractSliceOp::create( + rewriter, + loc, + rowType, + paddedOutput, + SmallVector {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}, + SmallVector {rewriter.getIndexAttr(1), + rewriter.getIndexAttr(outputChannels)}, + getUnitStrides(rewriter, 2)) + .getResult(); if (bias) validRow = spatial::SpatVAddOp::create(rewriter, loc, rowType, validRow, bias).getResult(); return validRow; } +static FailureOr createRowStripConvOutput(const ConvLoweringState& state, + Value input, + Value paddedWeights, + Value bias, + int64_t paddedK, + int64_t numKSlices, + int64_t xbarDim, + PatternRewriter& rewriter, + Location loc) { + const int64_t laneCount = state.outHeight; + const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth; + const bool hasPartialInputTile = patchSize % xbarDim != 0; + auto elementType = state.outType.getElementType(); + auto partialInputScratchType = RankedTensorType::get({1, xbarDim}, elementType); + auto outputPixelType = RankedTensorType::get({1, 1, 1, state.numChannelsOut}, elementType); + auto fragmentType = getRowStripFragmentType(state.outType); + auto storageType = getRowStripStorageType(state.outType); + + auto batch = createSpatComputeBatch( + rewriter, + loc, + TypeRange {storageType}, + laneCount, + ValueRange {paddedWeights}, + bias ? ValueRange {input, bias} : ValueRange {input}, + [&](detail::SpatComputeBatchBodyArgs args) { + Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp(); + Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0); + Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1); + Value cOutWidth = getOrCreateIndexConstant(rewriter, anchorOp, state.outWidth); + FailureOr inputWindow = + createConvInputWindow(args.inputs.front(), state, args.lane, rewriter, loc); + if (failed(inputWindow)) + return failure(); + Value fragmentInit = tensor::EmptyOp::create(rewriter, loc, fragmentType.getShape(), elementType); + SmallVector loopInit {fragmentInit}; + if (hasPartialInputTile) + loopInit.push_back(createZeroTensorConstant(partialInputScratchType, rewriter)); + auto loop = buildNormalizedScfFor( + rewriter, + loc, + c0, + cOutWidth, + c1, + loopInit, + [&](OpBuilder&, + Location pixelLoc, + Value localColumn, + ValueRange iterArgs, + SmallVectorImpl& yielded) { + Value partialInputScratch = hasPartialInputTile ? iterArgs[1] : Value(); + FailureOr> inputTiles = createConvInputTiles(*inputWindow, + state, + localColumn, + partialInputScratch, + patchSize, + numKSlices, + xbarDim, + rewriter, + pixelLoc); + if (failed(inputTiles)) + return failure(); + FailureOr output = createConvOutputRow(*inputTiles, + paddedK, + state.numChannelsOut, + args.weights.front(), + bias ? args.inputs[1] : Value(), + xbarDim, + rewriter, + pixelLoc); + if (failed(output)) + return failure(); + Value outputPixel = tensor::ExpandShapeOp::create( + rewriter, pixelLoc, outputPixelType, *output, SmallVector {{0, 1, 2}, {3}}); + Value next = tensor::InsertSliceOp::create( + rewriter, + pixelLoc, + outputPixel, + iterArgs.front(), + SmallVector {rewriter.getIndexAttr(0), + rewriter.getIndexAttr(0), + localColumn, + rewriter.getIndexAttr(0)}, + SmallVector {rewriter.getIndexAttr(1), + rewriter.getIndexAttr(1), + rewriter.getIndexAttr(1), + rewriter.getIndexAttr(state.numChannelsOut)}, + getUnitStrides(rewriter, 4)); + yielded.push_back(next); + if (hasPartialInputTile) + yielded.push_back(partialInputScratch); + return success(); + }); + if (failed(loop)) + return failure(); + publishGraphBatchPhysicalFragment( + rewriter, loc, loop->results.front(), args.outputs.front(), args.lane); + return success(); + }); + if (failed(batch)) + return failure(); + return batch->getResult(0); +} + static FailureOr createOutputChannelTiledRowStripConvOutput(const ConvLoweringState& state, Value input, Value paddedWeights, @@ -3229,21 +3319,11 @@ static FailureOr const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth; const int64_t numKSlices = ceilIntegerDivide(patchSize, xbarDim); const int64_t paddedK = numKSlices * xbarDim; - const bool hasPartialInputTile = patchSize % xbarDim != 0; - auto elementType = state.outType.getElementType(); - auto partialInputScratchType = RankedTensorType::get({1, xbarDim}, elementType); - auto fragmentType = getRowStripFragmentType(state.outType); - auto outputPixelType = RankedTensorType::get({1, 1, 1, state.numChannelsOut}, elementType); - auto outputStorageType = getRowStripStorageType(state.outType); - Value paddedWeights = state.numChannelsOut <= xbarDim - ? standard::createPaddedPixelMajorWeightConstant( - weightDenseAttr, state, paddedK, xbarDim, rewriter) - : standard::createPaddedOutputChannelTiledWeightConstant( - weightDenseAttr, state, paddedK, xbarDim, rewriter); - if (state.numChannelsOut > xbarDim) - return createOutputChannelTiledRowStripConvOutput( - state, state.x, paddedWeights, paddedK, numKSlices, xbarDim, rewriter, loc); + const int64_t paddedOutputChannels = + ceilIntegerDivide(state.numChannelsOut, xbarDim) * xbarDim; + Value paddedWeights = standard::createPaddedPixelMajorWeightConstant( + weightDenseAttr, state, paddedK, paddedOutputChannels, rewriter); FailureOr bias = failure(); if (state.hasBias) @@ -3251,83 +3331,9 @@ static FailureOr if (state.hasBias && failed(bias)) return failure(); - auto batchOp = createSpatComputeBatch( - rewriter, - loc, - TypeRange {outputStorageType}, - state.outHeight, - ValueRange {paddedWeights}, - state.hasBias ? ValueRange {state.x, *bias} : ValueRange {state.x}, - [&](detail::SpatComputeBatchBodyArgs args) { - Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp(); - Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0); - Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1); - Value cOutWidth = getOrCreateIndexConstant(rewriter, anchorOp, state.outWidth); - FailureOr inputWindow = - createConvInputWindow(args.inputs.front(), state, args.lane, rewriter, loc); - if (failed(inputWindow)) - return failure(); - Value fragmentInit = tensor::EmptyOp::create(rewriter, loc, fragmentType.getShape(), elementType); - SmallVector widthLoopInit {fragmentInit}; - if (hasPartialInputTile) - widthLoopInit.push_back(createZeroTensorConstant(partialInputScratchType, rewriter)); - auto widthLoop = buildNormalizedScfFor( - rewriter, - loc, - c0, - cOutWidth, - c1, - widthLoopInit, - [&](OpBuilder&, Location widthLoc, Value widthIndex, ValueRange widthIterArgs, SmallVectorImpl& widthYielded) { - Value partialInputScratch = hasPartialInputTile ? widthIterArgs[1] : Value(); - FailureOr> inputTiles = createConvInputTiles(*inputWindow, - state, - widthIndex, - partialInputScratch, - patchSize, - numKSlices, - xbarDim, - rewriter, - widthLoc); - if (failed(inputTiles)) - return failure(); - FailureOr outputRow = createConvOutputRow(*inputTiles, - paddedK, - state.numChannelsOut, - args.weights.front(), - state.hasBias ? args.inputs[1] : Value(), - xbarDim, - rewriter, - widthLoc); - if (failed(outputRow)) - return failure(); - - Value outputFragment = tensor::ExpandShapeOp::create(rewriter, - widthLoc, - outputPixelType, - *outputRow, - SmallVector {{0, 1, 2}, {3}}); - SmallVector rowOffsets { - rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex, rewriter.getIndexAttr(0)}; - SmallVector rowSizes { - rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), - rewriter.getIndexAttr(state.numChannelsOut)}; - Value nextFragment = tensor::InsertSliceOp::create( - rewriter, widthLoc, outputFragment, widthIterArgs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 4)); - widthYielded.push_back(nextFragment); - if (hasPartialInputTile) - widthYielded.push_back(partialInputScratch); - return success(); - }); - if (failed(widthLoop)) - return failure(); - - insertRowStripFragment(widthLoop->results.front(), args.outputs.front(), state.outType, args.lane, rewriter, loc); - return success(); - }); - if (failed(batchOp)) - return failure(); - return batchOp->getResult(0); + return createRowStripConvOutput( + state, state.x, paddedWeights, state.hasBias ? *bias : Value(), + paddedK, numKSlices, xbarDim, rewriter, loc); } static FailureOr createConvOutputFromPixelMajorRowStripFragments(Value rowStripStorage, @@ -3346,105 +3352,22 @@ static FailureOr createConvOutputFromPixelMajorRowStripFragments(Value ro const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth; const int64_t numKSlices = ceilIntegerDivide(patchSize, xbarDim); const int64_t paddedK = numKSlices * xbarDim; - const bool hasPartialInputTile = patchSize % xbarDim != 0; - auto elementType = state.outType.getElementType(); - auto partialInputScratchType = RankedTensorType::get({1, xbarDim}, elementType); - auto outputPixelType = RankedTensorType::get({1, 1, 1, state.numChannelsOut}, elementType); - auto outputStorageType = getRowStripStorageType(state.outType); auto weightDenseAttr = getHostConstDenseElementsAttr(state.w); if (!weightDenseAttr) return failure(); - Value paddedWeights = state.numChannelsOut <= xbarDim - ? standard::createPaddedPixelMajorWeightConstant( - weightDenseAttr, state, paddedK, xbarDim, rewriter) - : standard::createPaddedOutputChannelTiledWeightConstant( - weightDenseAttr, state, paddedK, xbarDim, rewriter); - if (state.numChannelsOut > xbarDim) - return createOutputChannelTiledRowStripConvOutput( - state, rowStripStorage, paddedWeights, paddedK, numKSlices, xbarDim, rewriter, loc); + const int64_t paddedOutputChannels = + ceilIntegerDivide(state.numChannelsOut, xbarDim) * xbarDim; + Value paddedWeights = standard::createPaddedPixelMajorWeightConstant( + weightDenseAttr, state, paddedK, paddedOutputChannels, rewriter); FailureOr bias = failure(); if (state.hasBias) bias = createBiasRowConstant(state, rewriter); if (state.hasBias && failed(bias)) return failure(); - auto batchOp = createSpatComputeBatch( - rewriter, - loc, - TypeRange {outputStorageType}, - state.outHeight, - ValueRange {paddedWeights}, - state.hasBias ? ValueRange {rowStripStorage, *bias} : ValueRange {rowStripStorage}, - [&](detail::SpatComputeBatchBodyArgs args) { - Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp(); - Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0); - Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1); - Value cOutWidth = getOrCreateIndexConstant(rewriter, anchorOp, state.outWidth); - auto fragmentType = getRowStripFragmentType(state.outType); - FailureOr inputWindow = createConvInputWindow(args.inputs.front(), state, args.lane, rewriter, loc); - if (failed(inputWindow)) - return failure(); - Value fragmentInit = tensor::EmptyOp::create(rewriter, loc, fragmentType.getShape(), elementType); - SmallVector widthLoopInit {fragmentInit}; - if (hasPartialInputTile) - widthLoopInit.push_back(createZeroTensorConstant(partialInputScratchType, rewriter)); - auto widthLoop = buildNormalizedScfFor( - rewriter, - loc, - c0, - cOutWidth, - c1, - widthLoopInit, - [&](OpBuilder&, Location widthLoc, Value widthIndex, ValueRange widthIterArgs, SmallVectorImpl& widthYielded) { - Value partialInputScratch = hasPartialInputTile ? widthIterArgs[1] : Value(); - FailureOr> inputTiles = createConvInputTiles(*inputWindow, - state, - widthIndex, - partialInputScratch, - patchSize, - numKSlices, - xbarDim, - rewriter, - widthLoc); - if (failed(inputTiles)) - return failure(); - FailureOr outputRow = createConvOutputRow(*inputTiles, - paddedK, - state.numChannelsOut, - args.weights.front(), - state.hasBias ? args.inputs[1] : Value(), - xbarDim, - rewriter, - widthLoc); - if (failed(outputRow)) - return failure(); - - Value outputFragment = tensor::ExpandShapeOp::create(rewriter, - widthLoc, - outputPixelType, - *outputRow, - SmallVector {{0, 1, 2}, {3}}); - SmallVector rowOffsets { - rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex, rewriter.getIndexAttr(0)}; - SmallVector rowSizes { - rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), - rewriter.getIndexAttr(state.numChannelsOut)}; - Value nextFragment = tensor::InsertSliceOp::create( - rewriter, widthLoc, outputFragment, widthIterArgs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 4)); - widthYielded.push_back(nextFragment); - if (hasPartialInputTile) - widthYielded.push_back(partialInputScratch); - return success(); - }); - if (failed(widthLoop)) - return failure(); - - insertRowStripFragment(widthLoop->results.front(), args.outputs.front(), state.outType, args.lane, rewriter, loc); - return success(); - }); - if (failed(batchOp)) - return failure(); - return batchOp->getResult(0); + return createRowStripConvOutput( + state, rowStripStorage, paddedWeights, state.hasBias ? *bias : Value(), + paddedK, numKSlices, xbarDim, rewriter, loc); } static FailureOr createPointwiseOutputFromRowStripFragments(Value rowStripStorage, diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Elementwise.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Elementwise.cpp index acb3c54..1172902 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Elementwise.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Elementwise.cpp @@ -193,6 +193,13 @@ struct AddToSpatialCompute : OpConversionPattern { return success(); } + if (resultType.getRank() == 4 && adaptor.getA().getType() == resultType && adaptor.getB().getType() == resultType) { + auto plan = spatial::SpatAddPlanOp::create( + rewriter, op.getLoc(), resultType, adaptor.getA(), adaptor.getB(), rewriter.getStringAttr("nchw")); + rewriter.replaceOp(op, plan.getResult()); + return success(); + } + auto lhs = prepareElementwiseOperand(adaptor.getA(), resultType, rewriter, op.getLoc()); if (failed(lhs)) return failure(); diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/NN/Pool.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/NN/Pool.cpp index 1876149..26692de 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/NN/Pool.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/NN/Pool.cpp @@ -246,6 +246,15 @@ struct PoolToSpatialComputeBase : public OpConversionPattern { return success(); } } + else if (batchSize == 1 && outputHeight == 1 && outputWidth == 1 + && kernelHeight == inputHeight && kernelWidth == inputWidth + && dilationHeight == 1 && dilationWidth == 1 && padTop == 0 + && padLeft == 0 && padBottom == 0 && padRight == 0) { + auto plan = spatial::SpatGlobalAveragePoolPlanOp::create( + rewriter, loc, outType, x, rewriter.getStringAttr("nchw")); + rewriter.replaceOp(poolOp, plan.getResult()); + return success(); + } const int64_t xbarSize = static_cast(crossbarSize.getValue()); const int64_t channelTileCount = (channels + xbarSize - 1) / xbarSize; @@ -676,6 +685,132 @@ FailureOr lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp, return batch->getResult(0); } +LogicalResult canLowerGlobalAveragePoolPlanToRowStrip(spatial::SpatGlobalAveragePoolPlanOp planOp) { + auto inputType = dyn_cast(planOp.getInput().getType()); + auto outputType = dyn_cast(planOp.getOutput().getType()); + if (!inputType || !outputType || !inputType.hasStaticShape() || !outputType.hasStaticShape()) + return failure(); + if (inputType.getRank() != 4 || outputType.getRank() != 4 || inputType.getDimSize(0) != 1 + || outputType.getDimSize(0) != 1 || inputType.getDimSize(1) != outputType.getDimSize(1) + || outputType.getDimSize(2) != 1 || outputType.getDimSize(3) != 1) + return failure(); + return success(); +} + +FailureOr lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp, + std::optional rowStripInput, + PatternRewriter& rewriter) { + if (failed(canLowerGlobalAveragePoolPlanToRowStrip(planOp))) + return failure(); + + Location loc = planOp.getLoc(); + auto inputType = cast(planOp.getInput().getType()); + auto outputType = cast(planOp.getOutput().getType()); + auto elementType = dyn_cast(inputType.getElementType()); + if (!elementType) + return failure(); + + Value input = rowStripInput.value_or(planOp.getInput()); + auto actualInputType = dyn_cast(input.getType()); + FailureOr physicalValue = describeRowStripPhysicalValue(input, inputType); + const bool physicalInput = succeeded(physicalValue); + if (!physicalInput && actualInputType != inputType) + return failure(); + + const int64_t height = inputType.getDimSize(2); + const int64_t width = inputType.getDimSize(3); + const int64_t channels = inputType.getDimSize(1); + const int64_t tilesPerRow = physicalInput ? physicalValue->tilesPerRow : 1; + auto inputFragmentType = + physicalInput ? physicalValue->fragmentType : getRowStripFragmentType(inputType); + auto nchwInputFragmentType = RankedTensorType::get( + {1, channels, 1, width}, inputType.getElementType(), inputType.getEncoding()); + auto outputFragmentType = RankedTensorType::get( + {1, 1, 1, inputFragmentType.getDimSize(3)}, elementType, outputType.getEncoding()); + auto outputStorageType = + spatial::getGraphBatchPhysicalResultType(tilesPerRow, outputFragmentType); + auto zero = getOrCreateConstant( + rewriter, rewriter.getInsertionBlock()->getParentOp(), rewriter.getZeroAttr(outputFragmentType), outputFragmentType); + auto scaleAttr = DenseElementsAttr::get( + outputFragmentType, rewriter.getFloatAttr(elementType, 1.0 / static_cast(height * width))); + auto scale = getOrCreateConstant( + rewriter, rewriter.getInsertionBlock()->getParentOp(), scaleAttr, outputFragmentType); + + auto batch = createSpatComputeBatch( + rewriter, + loc, + TypeRange {outputStorageType}, + tilesPerRow, + ValueRange {zero, scale}, + ValueRange {input}, + [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult { + Value reduced = args.weights[0]; + for (int64_t row = 0; row < height; ++row) { + Value fragment; + if (physicalInput) { + Value sourceSlot = args.lane; + if (row != 0) + sourceSlot = arith::AddIOp::create( + rewriter, + loc, + sourceSlot, + getOrCreateIndexConstant( + rewriter, rewriter.getInsertionBlock()->getParentOp(), row * tilesPerRow)); + FailureOr physicalFragment = extractGraphBatchPhysicalFragment( + rewriter, loc, args.inputs.front(), sourceSlot, inputFragmentType); + if (failed(physicalFragment)) + return failure(); + fragment = *physicalFragment; + } + else { + Value nchw = tensor::ExtractSliceOp::create( + rewriter, + loc, + nchwInputFragmentType, + args.inputs.front(), + SmallVector {rewriter.getIndexAttr(0), + rewriter.getIndexAttr(0), + rewriter.getIndexAttr(row), + rewriter.getIndexAttr(0)}, + SmallVector {rewriter.getIndexAttr(1), + rewriter.getIndexAttr(channels), + rewriter.getIndexAttr(1), + rewriter.getIndexAttr(width)}, + getUnitStrides(rewriter, 4)); + fragment = ONNXTransposeOp::create( + rewriter, loc, inputFragmentType, nchw, rewriter.getI64ArrayAttr({0, 2, 3, 1})); + } + for (int64_t column = 0; column < width; ++column) { + Value point = tensor::ExtractSliceOp::create( + rewriter, + loc, + outputFragmentType, + fragment, + SmallVector {rewriter.getIndexAttr(0), + rewriter.getIndexAttr(0), + rewriter.getIndexAttr(column), + rewriter.getIndexAttr(0)}, + SmallVector {rewriter.getIndexAttr(1), + rewriter.getIndexAttr(1), + rewriter.getIndexAttr(1), + rewriter.getIndexAttr(inputFragmentType.getDimSize(3))}, + getUnitStrides(rewriter, 4)); + point = materializeTileTensor(rewriter, loc, point); + reduced = spatial::SpatVAddOp::create( + rewriter, loc, outputFragmentType, reduced, point); + } + } + reduced = spatial::SpatVMulOp::create( + rewriter, loc, outputFragmentType, reduced, args.weights[1]); + publishGraphBatchPhysicalFragment( + rewriter, loc, reduced, args.outputs.front(), args.lane); + return success(); + }); + if (failed(batch)) + return failure(); + return batch->getResult(0); +} + void populatePoolPatterns(RewritePatternSet& patterns, MLIRContext* ctx) { patterns.insert>(ctx); patterns.insert>(ctx); diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Concat.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Concat.cpp index 86a96a6..a23076e 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Concat.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Concat.cpp @@ -25,6 +25,17 @@ struct Concat : public OpConversionPattern { return success(); } + auto resultType = dyn_cast(maxpoolOp.getResult().getType()); + if (axis == 1 && resultType && resultType.hasStaticShape() && resultType.getRank() == 4 + && llvm::all_of(inputs, [](Value input) { + auto type = dyn_cast(input.getType()); + return type && type.hasStaticShape() && type.getRank() == 4; + })) { + rewriter.replaceOpWithNewOp( + maxpoolOp, resultType, inputs, rewriter.getI64IntegerAttr(axis), rewriter.getStringAttr("nchw")); + return success(); + } + auto computeOp = createSpatCompute( rewriter, maxpoolOp.getLoc(), TypeRange {maxpoolOp.getResult().getType()}, {}, inputs, [&](ValueRange args) { spatial::SpatYieldOp::create( diff --git a/src/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp b/src/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp index f24b421..be2308e 100644 --- a/src/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp +++ b/src/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp @@ -27,6 +27,14 @@ lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp, std::optional rowStripInput, mlir::PatternRewriter& rewriter); +mlir::LogicalResult +canLowerGlobalAveragePoolPlanToRowStrip(spatial::SpatGlobalAveragePoolPlanOp planOp); + +mlir::FailureOr +lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp, + std::optional rowStripInput, + mlir::PatternRewriter& rewriter); + mlir::LogicalResult canLowerFlattenFromRowStrip(spatial::SpatGraphCompute flattenOp); mlir::LogicalResult lowerFlattenFromRowStrip(const RowStripPhysicalValue& input, diff --git a/src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp b/src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp index 57d75e4..f132fae 100644 --- a/src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp @@ -36,10 +36,16 @@ static bool usesSelectedRowStrip(Operation* user, llvm::DenseMap(user)) return getSelectedLayout(layouts, biasAddPlan.getResult()) == SelectedLayout::PixelMajorRowStrip; + if (auto addPlan = dyn_cast(user)) + return getSelectedLayout(layouts, addPlan.getResult()) == SelectedLayout::PixelMajorRowStrip; + if (auto concatPlan = dyn_cast(user)) + return getSelectedLayout(layouts, concatPlan.getResult()) == SelectedLayout::PixelMajorRowStrip; if (auto convPlan = dyn_cast(user)) return getSelectedLayout(layouts, convPlan.getResult()) == SelectedLayout::PixelMajorRowStrip; if (auto maxPoolPlan = dyn_cast(user)) return getSelectedLayout(layouts, maxPoolPlan.getResult()) == SelectedLayout::PixelMajorRowStrip; + if (auto averagePoolPlan = dyn_cast(user)) + return getSelectedLayout(layouts, averagePoolPlan.getResult()) == SelectedLayout::PixelMajorRowStrip; if (auto flattenCompute = dyn_cast(user)) return succeeded(canLowerFlattenFromRowStrip(flattenCompute)); return false; @@ -62,10 +68,16 @@ static bool canConsumeRowStripAsUser(Operation* user) { auto resultType = dyn_cast(biasAddPlan.getOutput().getType()); return resultType && isSupportedBiasAddValue(biasAddPlan.getBias(), resultType); } + if (isa(user)) + return true; + if (isa(user)) + return true; if (auto convPlan = dyn_cast(user)) return succeeded(canConsumeAndProduceRowStrip(convPlan)); if (auto maxPoolPlan = dyn_cast(user)) return succeeded(canLowerMaxPoolPlanToRowStrip(maxPoolPlan)); + if (auto averagePoolPlan = dyn_cast(user)) + return succeeded(canLowerGlobalAveragePoolPlanToRowStrip(averagePoolPlan)); return false; } @@ -118,11 +130,38 @@ static SelectedLayout chooseBiasAddLayout(spatial::SpatBiasAddPlanOp biasAddPlan return SelectedLayout::PixelMajorRowStrip; } +static SelectedLayout chooseAddLayout(spatial::SpatAddPlanOp addPlan, llvm::DenseMap& layouts) { + if (getSelectedLayout(layouts, addPlan.getLhs()) != SelectedLayout::PixelMajorRowStrip + || getSelectedLayout(layouts, addPlan.getRhs()) != SelectedLayout::PixelMajorRowStrip) + return SelectedLayout::DenseNchw; + if (!allUsersCanHandleRowStrip(addPlan.getResult(), layouts)) + return SelectedLayout::DenseNchw; + return SelectedLayout::PixelMajorRowStrip; +} + +static SelectedLayout chooseConcatLayout(spatial::SpatConcatPlanOp concatPlan, + llvm::DenseMap& layouts) { + if (llvm::any_of(concatPlan.getInputs(), [&](Value input) { + return getSelectedLayout(layouts, input) != SelectedLayout::PixelMajorRowStrip; + })) + return SelectedLayout::DenseNchw; + if (!allUsersCanHandleRowStrip(concatPlan.getResult(), layouts)) + return SelectedLayout::DenseNchw; + return SelectedLayout::PixelMajorRowStrip; +} + static SelectedLayout chooseMaxPoolLayout(spatial::SpatMaxPool2DPlanOp maxPoolPlan) { return succeeded(canLowerMaxPoolPlanToRowStrip(maxPoolPlan)) ? SelectedLayout::PixelMajorRowStrip : SelectedLayout::DenseNchw; } +static SelectedLayout chooseGlobalAveragePoolLayout( + spatial::SpatGlobalAveragePoolPlanOp averagePoolPlan) { + return succeeded(canLowerGlobalAveragePoolPlanToRowStrip(averagePoolPlan)) + ? SelectedLayout::PixelMajorRowStrip + : SelectedLayout::DenseNchw; +} + static spatial::SpatBlueprintOp insertRowStripBlueprint(IRRewriter& rewriter, Value value) { auto outputType = cast(value.getType()); auto [offsets, sizes] = buildRowStripMetadata(outputType); @@ -215,6 +254,22 @@ struct SpatialLayoutPlanningPass final : PassWrapper(&op)) { + SelectedLayout selected = chooseAddLayout(addPlan, layouts); + if (layouts[addPlan.getResult()] != selected) { + layouts[addPlan.getResult()] = selected; + changed = true; + } + continue; + } + if (auto concatPlan = dyn_cast(&op)) { + SelectedLayout selected = chooseConcatLayout(concatPlan, layouts); + if (layouts[concatPlan.getResult()] != selected) { + layouts[concatPlan.getResult()] = selected; + changed = true; + } + continue; + } if (auto maxPoolPlan = dyn_cast(&op)) { SelectedLayout selected = chooseMaxPoolLayout(maxPoolPlan); if (layouts[maxPoolPlan.getResult()] != selected) { @@ -223,6 +278,14 @@ struct SpatialLayoutPlanningPass final : PassWrapper(&op)) { + SelectedLayout selected = chooseGlobalAveragePoolLayout(averagePoolPlan); + if (layouts[averagePoolPlan.getResult()] != selected) { + layouts[averagePoolPlan.getResult()] = selected; + changed = true; + } + continue; + } } } @@ -232,10 +295,16 @@ struct SpatialLayoutPlanningPass final : PassWrapper(&op)) producedValue = biasAddPlan.getResult(); + else if (auto addPlan = dyn_cast(&op)) + producedValue = addPlan.getResult(); + else if (auto concatPlan = dyn_cast(&op)) + producedValue = concatPlan.getResult(); else if (auto reluPlan = dyn_cast(&op)) producedValue = reluPlan.getResult(); else if (auto maxPoolPlan = dyn_cast(&op)) producedValue = maxPoolPlan.getResult(); + else if (auto averagePoolPlan = dyn_cast(&op)) + producedValue = averagePoolPlan.getResult(); else continue; diff --git a/src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp b/src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp index 44c0ee6..60f1324 100644 --- a/src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp +++ b/src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp @@ -264,7 +264,13 @@ LogicalResult raptor::SpatialToPimPass::enlargeVMMOutTensorsToCrossbarSize(func: auto outputType = cast(vmmOp.getOutput().getType()); ArrayRef outputShape = outputType.getShape(); assert(isHVectorShape(outputShape) && "expected a horizontal vector output"); - assert(outputShape[1] <= static_cast(crossbarSize) && "output width must fit in one crossbar"); + auto weightType = cast(vmmOp.getWeight().getType()); + const int64_t xbarDim = static_cast(crossbarSize); + const int64_t paddedOutputWidth = ceilIntegerDivide(outputShape[1], xbarDim) * xbarDim; + assert(weightType.getRank() == 2 && weightType.getDimSize(1) == paddedOutputWidth + && "expected VMM weight width to match the padded output width"); + assert(paddedOutputWidth / xbarDim <= static_cast(crossbarCountInCore) + && "output width must fit in one core"); rewriter.setInsertionPoint(vmmOp); auto paddedInput = padHVectorInputToCrossbarSize(rewriter, vmmOp.getLoc(), vmmOp.getInput()); @@ -273,8 +279,8 @@ LogicalResult raptor::SpatialToPimPass::enlargeVMMOutTensorsToCrossbarSize(func: return WalkResult::interrupt(); } auto paddedOutputType = RankedTensorType::get( - {outputShape[0], static_cast(crossbarSize)}, outputType.getElementType(), outputType.getEncoding()); - Value paddedOutputBuffer = outputShape[1] == static_cast(crossbarSize) + {outputShape[0], paddedOutputWidth}, outputType.getElementType(), outputType.getEncoding()); + Value paddedOutputBuffer = outputShape[1] == paddedOutputWidth ? vmmOp.getOutputBuffer() : createEmptyTensorFromShaped(rewriter, vmmOp.getLoc(), paddedOutputType).getResult(); vmmOp.getInputMutable().assign(*paddedInput); @@ -282,7 +288,7 @@ LogicalResult raptor::SpatialToPimPass::enlargeVMMOutTensorsToCrossbarSize(func: vmmOp.getOutput().setType(paddedOutputType); - if (outputShape[1] == static_cast(crossbarSize)) + if (outputShape[1] == paddedOutputWidth) return WalkResult::advance(); SmallVector offsets = {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}; diff --git a/src/PIM/Dialect/Pim/PimOpsVerify.cpp b/src/PIM/Dialect/Pim/PimOpsVerify.cpp index c29eda3..30ca606 100644 --- a/src/PIM/Dialect/Pim/PimOpsVerify.cpp +++ b/src/PIM/Dialect/Pim/PimOpsVerify.cpp @@ -143,18 +143,21 @@ LogicalResult PimVMMOp::verify() { int64_t M = matrixShape[1]; if (N <= 0 || M <= 0) return emitError("matrix shape must be (N, M) with N > 0 and M > 0"); - if (N > static_cast(crossbarSize) || M > static_cast(crossbarSize)) - return emitError("matrix dimensions must fit in one crossbar"); + const int64_t xbarDim = static_cast(crossbarSize); + if (N > xbarDim || M > xbarDim * static_cast(crossbarCountInCore)) + return emitError("matrix dimensions must fit in one array group"); + if (M % xbarDim != 0) + return emitError("matrix output width must be padded to a whole number of crossbars"); int64_t vector1 = vectorShape[0]; int64_t vectorWidth = vectorShape[1]; - if (vector1 != 1 || vectorWidth != static_cast(crossbarSize)) + if (vector1 != 1 || vectorWidth != xbarDim) return emitError("vector shape must be (1, crossbar-size)"); int64_t output1 = outputShape[0]; int64_t outputWidth = outputShape[1]; - if (output1 != 1 || outputWidth != static_cast(crossbarSize)) - return emitError("output shape must be (1, crossbar-size)"); + if (output1 != 1 || outputWidth != M) + return emitError("output shape must match the array-group width"); return success(); } diff --git a/src/PIM/Dialect/Pim/Transforms/Bufferization/PimBufferizationPass.cpp b/src/PIM/Dialect/Pim/Transforms/Bufferization/PimBufferizationPass.cpp index b7977b1..9f95470 100644 --- a/src/PIM/Dialect/Pim/Transforms/Bufferization/PimBufferizationPass.cpp +++ b/src/PIM/Dialect/Pim/Transforms/Bufferization/PimBufferizationPass.cpp @@ -6,7 +6,9 @@ #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/Dominance.h" #include "mlir/IR/PatternMatch.h" +#include "mlir/Interfaces/DestinationStyleOpInterface.h" #include "mlir/Pass/Pass.h" #include "mlir/Transforms/GreedyPatternRewriteDriver.h" @@ -78,6 +80,15 @@ lowerMemRefCopyToPimCopy(memref::CopyOp copyOp, bool sourceIsDevice = isDeviceLocalPimAddress(copyOp.getSource(), knowledge); bool targetIsDevice = isDeviceLocalPimAddress(copyOp.getTarget(), knowledge); + auto sourceAddress = resolveContiguousAddress(copyOp.getSource(), knowledge); + auto targetAddress = resolveContiguousAddress(copyOp.getTarget(), knowledge); + if (succeeded(sourceAddress) && succeeded(targetAddress) + && sourceAddress->base == targetAddress->base + && sourceAddress->byteOffset == targetAddress->byteOffset) { + rewriter.eraseOp(copyOp); + return success(); + } + if (targetIsDevice && sourceIsHost) { pim::PimMemCopyHostToDevOp::create(rewriter, copyOp.getLoc(), @@ -184,6 +195,62 @@ static void forwardSingleConsumerContiguousInputCopies(func::FuncOp funcOp) { } } +static void forwardSingleConsumerPimOutputCopies(func::FuncOp funcOp) { + DominanceInfo dominance(funcOp); + SmallVector copies; + funcOp.walk([&](memref::CopyOp copy) { copies.push_back(copy); }); + + for (memref::CopyOp copy : copies) { + Value copiedSource = copy.getSource(); + auto expand = copiedSource.getDefiningOp(); + Value producerOutput = expand ? expand.getSrc() : copiedSource; + auto result = dyn_cast(producerOutput); + Operation* producer = result ? result.getOwner() : nullptr; + auto dps = dyn_cast_or_null(producer); + if (!producer || producer->getName().getDialectNamespace() + != PimDialect::getDialectNamespace() + || result.getResultNumber() != 0 || !copiedSource.hasOneUse() + || !producerOutput.hasOneUse() || !dps) + continue; + + MutableOperandRange inits = dps.getDpsInitsMutable(); + if (inits.size() != 1) + continue; + OpOperand& init = *inits.begin(); + auto outputAlloc = init.get().getDefiningOp(); + if (!outputAlloc || !init.get().hasOneUse()) + continue; + + Value destination = copy.getTarget(); + Operation* destinationView = destination.getDefiningOp(); + if (!dominance.dominates(destination, producer)) { + if (!destinationView || destinationView->getBlock() != producer->getBlock() + || !destination.hasOneUse() + || !llvm::all_of(destinationView->getOperands(), + [&](Value operand) { + return dominance.dominates(operand, producer); + })) + continue; + destinationView->moveBefore(producer); + } + + OpBuilder builder(producer); + if (expand) + destination = memref::CollapseShapeOp::create( + builder, copy.getLoc(), destination, + expand.getReassociationIndices()); + if (cast(destination.getType()).getShape() + != cast(producerOutput.getType()).getShape() + || cast(destination.getType()).getElementType() + != cast(producerOutput.getType()).getElementType()) + continue; + + init.set(destination); + if (outputAlloc.use_empty()) + outputAlloc.erase(); + } +} + enum class ExpectedPimCopyDirection { HostToDevice, DeviceToHost, DeviceToDevice }; static LogicalResult verifyPimCopyEndpoints(Operation* copy, @@ -340,6 +407,7 @@ void PimBufferizationPass::runOnOperation() { } forwardSingleConsumerContiguousInputCopies(funcOp); + forwardSingleConsumerPimOutputCopies(funcOp); MLIRContext* ctx = moduleOp.getContext(); PatternRewriter rewriter(ctx); diff --git a/src/PIM/Dialect/Pim/Transforms/HostConstantFolding/CMakeLists.txt b/src/PIM/Dialect/Pim/Transforms/HostConstantFolding/CMakeLists.txt index 7b815ab..41c8ad3 100644 --- a/src/PIM/Dialect/Pim/Transforms/HostConstantFolding/CMakeLists.txt +++ b/src/PIM/Dialect/Pim/Transforms/HostConstantFolding/CMakeLists.txt @@ -9,4 +9,5 @@ add_pim_library(OMPimHostConstantFolding LINK_LIBS PUBLIC MLIRLinalgDialect OMPimCommon + OMPimBufferization ) diff --git a/src/PIM/Dialect/Pim/Transforms/HostConstantFolding/HostConstantFoldingPass.cpp b/src/PIM/Dialect/Pim/Transforms/HostConstantFolding/HostConstantFoldingPass.cpp index 133e92b..a7d4cc9 100644 --- a/src/PIM/Dialect/Pim/Transforms/HostConstantFolding/HostConstantFoldingPass.cpp +++ b/src/PIM/Dialect/Pim/Transforms/HostConstantFolding/HostConstantFoldingPass.cpp @@ -5,6 +5,7 @@ #include "Patterns.hpp" #include "src/Accelerators/PIM/Common/PimCommon.hpp" +#include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.hpp" using namespace mlir; @@ -26,6 +27,7 @@ struct HostConstantFoldingPass : PassWrapper(std::move(owningPatterns)); return success(); diff --git a/src/PIM/Dialect/Spatial/Spatial.td b/src/PIM/Dialect/Spatial/Spatial.td index 22ea176..b4ab4be 100644 --- a/src/PIM/Dialect/Spatial/Spatial.td +++ b/src/PIM/Dialect/Spatial/Spatial.td @@ -307,6 +307,21 @@ def SpatMaxPool2DPlanOp : SpatOp<"max_pool2d_plan", []> { let hasVerifier = 1; } +def SpatGlobalAveragePoolPlanOp : SpatOp<"global_average_pool_plan", []> { + let summary = "Layout-aware NCHW global average-pool planning op"; + + let arguments = (ins + SpatTensor:$input, + StrAttr:$logicalLayout + ); + + let results = (outs + SpatTensor:$output + ); + + let hasVerifier = 1; +} + def SpatBiasAddPlanOp : SpatOp<"bias_add_plan", []> { let summary = "Layout-aware Conv-style bias add planning op"; @@ -323,6 +338,38 @@ def SpatBiasAddPlanOp : SpatOp<"bias_add_plan", []> { let hasVerifier = 1; } +def SpatAddPlanOp : SpatOp<"add_plan", []> { + let summary = "Layout-aware elementwise add planning op"; + + let arguments = (ins + SpatTensor:$lhs, + SpatTensor:$rhs, + StrAttr:$logicalLayout + ); + + let results = (outs + SpatTensor:$output + ); + + let hasVerifier = 1; +} + +def SpatConcatPlanOp : SpatOp<"concat_plan", []> { + let summary = "Layout-aware tensor concatenation planning op"; + + let arguments = (ins + Variadic:$inputs, + I64Attr:$axis, + StrAttr:$logicalLayout + ); + + let results = (outs + SpatTensor:$output + ); + + let hasVerifier = 1; +} + def SpatBlueprintOp : SpatOp<"blueprint", []> { let summary = "Blueprint for assembling logical tensors from published fragments"; diff --git a/src/PIM/Dialect/Spatial/SpatialOpsAsm.cpp b/src/PIM/Dialect/Spatial/SpatialOpsAsm.cpp index d987c09..1f1f90e 100644 --- a/src/PIM/Dialect/Spatial/SpatialOpsAsm.cpp +++ b/src/PIM/Dialect/Spatial/SpatialOpsAsm.cpp @@ -149,7 +149,8 @@ void printComputeLikeOp(ComputeOpTy op, OpAsmPrinter& printer) { if (auto coreIdAttr = op->template getAttrOfType(onnx_mlir::kCoreIdAttrName)) printer << " coreId " << coreIdAttr.getInt(); - printer << " crossbarWeights " << collectDistinctCrossbarWeights(op.getOperation()).size(); + printer << " crossbarWeights " + << collectDistinctResidentWeights(op.getOperation()).size(); printer.printOptionalAttrDict(op->getAttrs(), {op.getOperandSegmentSizesAttrName().getValue(), onnx_mlir::kCoreIdAttrName}); @@ -275,7 +276,11 @@ void printComputeBatchLikeOp(ComputeBatchOpTy op, OpAsmPrinter& printer) { printer << " shared_outs"; printBlockArgumentList(printer, outputArgs); } - printer << " crossbarWeights " << getComputeInstanceCrossbarUsage({op.getOperation(), 0, op.getLaneCount()}).size(); + printer << " crossbarWeights " + << getComputeInstanceResidentWeights( + {op.getOperation(), 0, + static_cast(op.getLaneCount())}) + .size(); if (auto coreIdsAttr = op->template getAttrOfType(onnx_mlir::kCoreIdsAttrName)) { printer << " coreIds "; printCompressedIntegerList(printer, coreIdsAttr.asArrayRef()); diff --git a/src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp b/src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp index 1e92546..5b8b4a7 100644 --- a/src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp +++ b/src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp @@ -353,31 +353,30 @@ LogicalResult SpatExtractRowsOp::verify() { return success(); } -LogicalResult SpatConcatOp::verify() { - if (getInputs().empty()) - return emitError("requires at least one input"); +static LogicalResult verifyConcatTypes(Operation* op, ValueRange inputs, Value output, int64_t axis) { + if (inputs.empty()) + return op->emitError("requires at least one input"); - auto outputType = dyn_cast(getOutput().getType()); + auto outputType = dyn_cast(output.getType()); if (!outputType || !outputType.hasRank()) - return emitError("output must be a ranked shaped type"); + return op->emitError("output must be a ranked shaped type"); - int64_t axis = getAxis(); int64_t rank = outputType.getRank(); if (axis < 0 || axis >= rank) - return emitError("axis must be within the output rank"); + return op->emitError("axis must be within the output rank"); int64_t concatenatedDimSize = 0; bool concatenatedDimDynamic = false; Type outputElementType = outputType.getElementType(); - for (Value input : getInputs()) { + for (Value input : inputs) { auto inputType = dyn_cast(input.getType()); if (!inputType || !inputType.hasRank()) - return emitError("inputs must be ranked shaped types"); + return op->emitError("inputs must be ranked shaped types"); if (inputType.getRank() != rank) - return emitError("all inputs must have the same rank as the output"); + return op->emitError("all inputs must have the same rank as the output"); if (inputType.getElementType() != outputElementType) - return emitError("all inputs must have the same element type as the output"); + return op->emitError("all inputs must have the same element type as the output"); for (int64_t dim = 0; dim < rank; ++dim) { if (dim == axis) @@ -385,7 +384,7 @@ LogicalResult SpatConcatOp::verify() { int64_t inputDim = inputType.getDimSize(dim); int64_t outputDim = outputType.getDimSize(dim); if (!ShapedType::isDynamic(inputDim) && !ShapedType::isDynamic(outputDim) && inputDim != outputDim) - return emitError("non-concatenated dimensions must match the output shape"); + return op->emitError("non-concatenated dimensions must match the output shape"); } int64_t inputConcatDim = inputType.getDimSize(axis); @@ -398,11 +397,24 @@ LogicalResult SpatConcatOp::verify() { int64_t outputConcatDim = outputType.getDimSize(axis); if (!concatenatedDimDynamic && !ShapedType::isDynamic(outputConcatDim) && concatenatedDimSize != outputConcatDim) - return emitError("output concatenated dimension must equal the sum of input sizes"); + return op->emitError("output concatenated dimension must equal the sum of input sizes"); return success(); } +LogicalResult SpatConcatOp::verify() { return verifyConcatTypes(getOperation(), getInputs(), getOutput(), getAxis()); } + +LogicalResult SpatConcatPlanOp::verify() { + if (getLogicalLayout() != "nchw") + return emitError("requires logicalLayout = \"nchw\""); + auto outputType = dyn_cast(getOutput().getType()); + if (!outputType || !outputType.hasStaticShape() || outputType.getRank() != 4) + return emitError("requires a static rank-4 output"); + if (getAxis() != 1) + return emitError("only channel-axis concatenation is supported"); + return verifyConcatTypes(getOperation(), getInputs(), getOutput(), getAxis()); +} + static bool isKnownLogicalLayout(StringRef layout) { return layout == "nchw"; } static bool isKnownPhysicalLayout(StringRef layout) { @@ -480,6 +492,24 @@ LogicalResult SpatMaxPool2DPlanOp::verify() { return success(); } +LogicalResult SpatGlobalAveragePoolPlanOp::verify() { + if (failed(verifyPlanTensorTypes( + getOperation(), getInput(), getOutput(), "spat.global_average_pool_plan"))) + return failure(); + auto inputType = dyn_cast(getInput().getType()); + auto outputType = dyn_cast(getOutput().getType()); + if (!inputType.hasStaticShape() || !outputType.hasStaticShape() || inputType.getRank() != 4 + || outputType.getRank() != 4) + return emitError("requires static rank-4 input and output tensors"); + if (getLogicalLayout() != "nchw") + return emitError("requires logical layout \"nchw\""); + if (inputType.getDimSize(0) != 1 || outputType.getDimSize(0) != 1 + || inputType.getDimSize(1) != outputType.getDimSize(1) + || outputType.getDimSize(2) != 1 || outputType.getDimSize(3) != 1) + return emitError("requires batch-one input and matching 1x1 output channels"); + return success(); +} + LogicalResult SpatBiasAddPlanOp::verify() { if (failed(verifyPlanTensorTypes(getOperation(), getInput(), getOutput(), "spat.bias_add_plan"))) return failure(); @@ -513,6 +543,23 @@ LogicalResult SpatBiasAddPlanOp::verify() { return success(); } +LogicalResult SpatAddPlanOp::verify() { + auto lhsType = dyn_cast(getLhs().getType()); + auto rhsType = dyn_cast(getRhs().getType()); + auto outputType = dyn_cast(getOutput().getType()); + if (!lhsType || !rhsType || !outputType) + return emitError("requires ranked tensor operands and output"); + if (!lhsType.hasStaticShape() || !rhsType.hasStaticShape() || !outputType.hasStaticShape()) + return emitError("requires static tensor operands and output"); + if (lhsType != rhsType || lhsType != outputType) + return emitError("requires matching operand and output tensor types"); + if (outputType.getRank() != 4) + return emitError("requires rank-4 operands and output"); + if (getLogicalLayout() != "nchw") + return emitError("requires logical layout \"nchw\""); + return success(); +} + LogicalResult SpatBlueprintOp::verify() { auto modeAttr = getModeAttr(); bool isFragmentAssembly = modeAttr && modeAttr.getValue() == "fragment_assembly"; diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp index 5c77214..b1c3f8f 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp @@ -262,28 +262,30 @@ static void collectClosure(Value value, Block &body, const DeferredInputPlan &pl } // namespace -bool isDeferredFragmentAssemblyInput(Value input) { +bool isDeferredFragmentAssemblyInput(Value input, size_t processorCount) { auto blueprint = input.getDefiningOp(); if (!blueprint || blueprint.getMode() != "fragment_assembly") return false; return llvm::all_of(getBlueprintFragments(blueprint), [&](Value fragment) { - return getProducerValueRef(fragment, nullptr).has_value(); + return getProducerValueRef(fragment, nullptr, processorCount).has_value(); }); } LogicalResult prepareSingleCpuInput(OpBuilder &, Location loc, Value input, BlockArgument graphInput, - const ComputeInstance &consumerInstance, const MergeScheduleResult &, + const ComputeInstance &consumerInstance, + const MergeScheduleResult &schedule, ValueRange scheduledInputs, Block &block, unsigned firstInputArgument, const DenseMap &availableValues, Value graphLane, Value scheduledGraphLane, DeferredInputPlan &plan) { plan = {graphInput, {}, {}, {}, graphLane, scheduledGraphLane, {}, {}, {}, {}, 1, nullptr}; - if (isDeferredFragmentAssemblyInput(input)) { + if (isDeferredFragmentAssemblyInput(input, schedule.processorCount)) { plan.blueprint = input.getDefiningOp(); plan.originalSources = getBlueprintFragments(plan.blueprint); return success(); } - auto producer = getProducerValueRef(input, &consumerInstance); + auto producer = getProducerValueRef( + input, &consumerInstance, schedule.processorCount); if (!producer) { plan.availableValue = getBlockOperand(block, scheduledInputs, input, firstInputArgument); return success(); } ProducerValueKey key {producer->instance, producer->resultIndex}; auto batch = dyn_cast(producer->instance.op); @@ -304,17 +306,19 @@ LogicalResult prepareSingleCpuInput(OpBuilder &, Location loc, Value input, Bloc LogicalResult prepareMultiCpuTupleInput(OpBuilder &, Location loc, Value input, BlockArgument graphInput, const ComputeStepTuple &tuple, const PeftClassPlan &, - const MergeScheduleResult &, ValueRange scheduledInputs, Block &block, + const MergeScheduleResult &schedule, + ValueRange scheduledInputs, Block &block, unsigned firstInputArgument, Value graphLane, Value scheduledGraphLane, Value scheduledLane, DeferredInputPlan &plan) { const ComputeInstance &representative = tuple.instances.front(); plan = {graphInput, {}, {}, {}, graphLane, scheduledGraphLane, scheduledLane, {}, {}, {}, 1, nullptr}; - if (isDeferredFragmentAssemblyInput(input)) { + if (isDeferredFragmentAssemblyInput(input, schedule.processorCount)) { plan.blueprint = input.getDefiningOp(); plan.originalSources = getBlueprintFragments(plan.blueprint); return success(); } - auto producer = getProducerValueRef(input, &representative); + auto producer = getProducerValueRef( + input, &representative, schedule.processorCount); if (!producer) { plan.availableValue = getBlockOperand(block, scheduledInputs, input, firstInputArgument); return success(); } auto inputs = getComputeInstanceInputs(representative); auto it = llvm::find(inputs, input); @@ -323,7 +327,8 @@ LogicalResult prepareMultiCpuTupleInput(OpBuilder &, Location loc, Value input, for (const ComputeInstance &instance : tuple.instances) { auto laneInputs = getComputeInstanceInputs(instance); if (inputIndex >= laneInputs.size()) return emitError(loc) << "scheduled batch step input out of range"; - auto laneProducer = getProducerValueRef(laneInputs[inputIndex], &instance); + auto laneProducer = getProducerValueRef( + laneInputs[inputIndex], &instance, schedule.processorCount); if (!laneProducer) return emitError(loc) << "scheduled batch step mixes host and producer inputs"; auto source = getOriginalProducerValue(*laneProducer); if (failed(source)) return emitError(loc) << "cannot resolve original graph producer value"; diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.hpp index ebf68ee..7f719d3 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.hpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.hpp @@ -22,7 +22,7 @@ struct DeferredInputPlan { Block *scalarizedHoistBlock = nullptr; }; -bool isDeferredFragmentAssemblyInput(Value input); +bool isDeferredFragmentAssemblyInput(Value input, size_t processorCount); LogicalResult prepareSingleCpuInput(OpBuilder &builder, Location loc, Value input, BlockArgument graphInput, diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp index e82dfb5..8af6bf9 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp @@ -1,39 +1,97 @@ -#include "DeferredCommunicationRealization.hpp" +#include "mlir/IR/Dominance.h" #include "DeferredBoundaryPlanning.hpp" #include "DeferredBoundaryRealization.hpp" #include "DeferredCommunicationDeadlock.hpp" +#include "DeferredCommunicationRealization.hpp" #include "DeferredCommunicationScheduling.hpp" #include "DeferredTransferPlanning.hpp" - -#include "mlir/IR/Dominance.h" +#include "Scheduling/PeftScheduler.hpp" +#include "src/Accelerators/PIM/Common/PimCommon.hpp" +#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp" namespace onnx_mlir::spatial { using namespace mlir; namespace { -static LogicalResult replaceFinalGraphPublications( - func::FuncOp funcOp, DeferredTransferPlan &plan) { - for (Operation &op : funcOp.getOps()) { +static LogicalResult placeLogicalProcessorsOnPhysicalCores(DeferredTransferPlan& plan, const SchedulingTarget& target) { + std::vector logicalTrafficFlits(target.processorCount * target.processorCount, 0); + for (const std::unique_ptr& exchange : plan.exchanges) + for (const ExternalTransferFamily& transfer : exchange->external) { + auto fragmentType = dyn_cast(transfer.requirement->publicationFragmentType); + if (!fragmentType || !fragmentType.hasStaticShape()) + return exchange->deferred.emitOpError("physical core placement requires a static transfer fragment"); + auto fragmentBytes = pim::getCheckedShapedTypeSizeInBytes( + fragmentType, exchange->deferred, "physical core placement transfer fragment"); + if (failed(fragmentBytes)) + return failure(); + Cost flits = static_cast(*fragmentBytes) / target.transferWidthBytes + + (*fragmentBytes % target.transferWidthBytes != 0); + for (size_t index = 0; index < transfer.sourceCores.size(); ++index) { + size_t sourceLogicalProcessor = static_cast(transfer.sourceCores.valueAt(index)); + size_t targetLogicalProcessor = static_cast(transfer.targetCores.valueAt(index)); + Cost& traffic = logicalTrafficFlits[sourceLogicalProcessor * target.processorCount + targetLogicalProcessor]; + traffic = checkedAdd(traffic, flits); + } + } + + std::vector physicalCoreForLogicalProcessor = + mapLogicalProcessorsToPhysicalCores(logicalTrafficFlits, target); + auto getPhysicalCore = [&](int64_t logicalProcessor) { + assert(logicalProcessor >= 0 && static_cast(logicalProcessor) < physicalCoreForLogicalProcessor.size() + && "logical processor is outside the scheduling target"); + return static_cast(physicalCoreForLogicalProcessor[logicalProcessor]); + }; + auto remap = [&](StaticIntSequence& logicalProcessors) { + SmallVector physicalCores; + physicalCores.reserve(logicalProcessors.size()); + for (size_t index = 0; index < logicalProcessors.size(); ++index) + physicalCores.push_back(getPhysicalCore(logicalProcessors.valueAt(index))); + logicalProcessors = StaticIntSequence::fromValues(physicalCores); + }; + + for (ScheduledInfo& scheduled : plan.scheduled) { + for (int64_t& logicalProcessor : scheduled.cores) + logicalProcessor = getPhysicalCore(logicalProcessor); + if (isa(scheduled.op)) { + scheduled.op->setAttr( + kCoreIdAttrName, IntegerAttr::get(IntegerType::get(scheduled.op->getContext(), 32), scheduled.cores.front())); + } + else { + SmallVector physicalCores; + physicalCores.reserve(scheduled.cores.size()); + for (int64_t physicalCore : scheduled.cores) + physicalCores.push_back(static_cast(physicalCore)); + scheduled.op->setAttr(kCoreIdsAttrName, DenseI32ArrayAttr::get(scheduled.op->getContext(), physicalCores)); + } + } + for (const std::unique_ptr& produced : plan.producedStorage) + produced->core = getPhysicalCore(produced->core); + for (const std::unique_ptr& exchange : plan.exchanges) + for (ExternalTransferFamily& transfer : exchange->external) { + remap(transfer.sourceCores); + remap(transfer.targetCores); + } + return success(); +} + +static LogicalResult replaceFinalGraphPublications(func::FuncOp funcOp, DeferredTransferPlan& plan) { + for (Operation& op : funcOp.getOps()) { if (!isa(op)) continue; auto graphId = op.getAttrOfType("scheduled.graph_id"); if (!graphId) continue; for (auto [resultIndex, result] : llvm::enumerate(op.getResults())) { - SmallVector externalUses; - for (OpOperand &use : result.getUses()) { - Operation *user = use.getOwner(); - if (isa(user)) + SmallVector externalUses; + for (OpOperand& use : result.getUses()) { + Operation* user = use.getOwner(); + if (isa(user)) continue; if (auto blueprint = dyn_cast(user)) { - bool blueprintEscapes = llvm::any_of( - blueprint.getOutput().getUses(), [](OpOperand &blueprintUse) { - return !isa( - blueprintUse.getOwner()); - }); + bool blueprintEscapes = llvm::any_of(blueprint.getOutput().getUses(), [](OpOperand& blueprintUse) { + return !isa(blueprintUse.getOwner()); + }); if (!blueprintEscapes) continue; } @@ -42,21 +100,16 @@ static LogicalResult replaceFinalGraphPublications( if (externalUses.empty()) continue; SmallVector exact; - for (ProducedValue *produced : - plan.producedByGraph.lookup(graphId.getInt())) - if (produced->resultIndex == resultIndex - && produced->published - && produced->published.getType() == result.getType() - && !llvm::is_contained(exact, produced->published)) + for (ProducedValue* produced : plan.producedByGraph.lookup(graphId.getInt())) + if (produced->resultIndex == resultIndex && produced->published + && produced->published.getType() == result.getType() && !llvm::is_contained(exact, produced->published)) exact.push_back(produced->published); if (exact.size() != 1) - return op.emitOpError( - "phase 2 final publication ownership changed after planning"); - for (OpOperand *use : externalUses) { - Operation *consumer = use->getOwner(); - Operation *producer = exact.front().getDefiningOp(); - if (consumer->getBlock() == producer->getBlock() - && consumer->isBeforeInBlock(producer)) + return op.emitOpError("phase 2 final publication ownership changed after planning"); + for (OpOperand* use : externalUses) { + Operation* consumer = use->getOwner(); + Operation* producer = exact.front().getDefiningOp(); + if (consumer->getBlock() == producer->getBlock() && consumer->isBeforeInBlock(producer)) consumer->moveAfter(producer); use->set(exact.front()); } @@ -65,13 +118,12 @@ static LogicalResult replaceFinalGraphPublications( return success(); } -static LogicalResult eraseOldGraph(func::FuncOp funcOp, - IRRewriter &rewriter) { - SmallVector oldGraph; - for (Operation &op : funcOp.getOps()) +static LogicalResult eraseOldGraph(func::FuncOp funcOp, IRRewriter& rewriter) { + SmallVector oldGraph; + for (Operation& op : funcOp.getOps()) if (isa(op)) oldGraph.push_back(&op); - for (Operation *op : llvm::reverse(oldGraph)) { + for (Operation* op : llvm::reverse(oldGraph)) { if (auto blueprint = dyn_cast(op)) { if (blueprint.getOutput().use_empty()) rewriter.eraseOp(blueprint); @@ -80,10 +132,9 @@ static LogicalResult eraseOldGraph(func::FuncOp funcOp, if (!op->use_empty()) { for (OpResult result : op->getResults()) { if (!result.use_empty()) { - Operation *user = result.use_begin()->getOwner(); - return op->emitOpError() - << "phase 2 cannot erase old graph result " - << result.getResultNumber() << " used by " << user->getName(); + Operation* user = result.use_begin()->getOwner(); + return op->emitOpError() << "phase 2 cannot erase old graph result " << result.getResultNumber() + << " used by " << user->getName(); } } } @@ -92,34 +143,27 @@ static LogicalResult eraseOldGraph(func::FuncOp funcOp, return success(); } -static LogicalResult eraseDeferredSourceSelectors( - func::FuncOp funcOp, IRRewriter &rewriter) { +static LogicalResult eraseDeferredSourceSelectors(func::FuncOp funcOp, IRRewriter& rewriter) { SmallVector selectors; - funcOp.walk([&](SpatDeferredSourceSelectOp selector) { - selectors.push_back(selector); - }); + funcOp.walk([&](SpatDeferredSourceSelectOp selector) { selectors.push_back(selector); }); for (SpatDeferredSourceSelectOp selector : llvm::reverse(selectors)) { if (!selector.getOutput().use_empty()) - return selector.emitOpError( - "phase 2 left a live deferred source selection"); + return selector.emitOpError("phase 2 left a live deferred source selection"); rewriter.eraseOp(selector); } return success(); } -static void eraseUnusedIdentityDeferredCommunications( - func::FuncOp funcOp, IRRewriter &rewriter) { +static void eraseUnusedIdentityDeferredCommunications(func::FuncOp funcOp, IRRewriter& rewriter) { SmallVector unused; funcOp.walk([&](SpatDeferredCommunicationOp deferred) { if (!deferred.getOutput().use_empty() || !deferred.getBody().hasOneBlock()) return; - Block &body = deferred.getBody().front(); + Block& body = deferred.getBody().front(); auto yield = dyn_cast(body.getTerminator()); - auto argument = yield && yield.getOutputs().size() == 1 - ? dyn_cast(yield.getOutputs().front()) - : BlockArgument(); - if (argument && argument.getOwner() == &body - && argument.getArgNumber() < deferred.getSources().size()) + auto argument = + yield && yield.getOutputs().size() == 1 ? dyn_cast(yield.getOutputs().front()) : BlockArgument(); + if (argument && argument.getOwner() == &body && argument.getArgNumber() < deferred.getSources().size()) unused.push_back(deferred); }); for (SpatDeferredCommunicationOp deferred : llvm::reverse(unused)) @@ -128,11 +172,10 @@ static void eraseUnusedIdentityDeferredCommunications( static LogicalResult verifyDominance(func::FuncOp funcOp) { DominanceInfo dominance(funcOp); - WalkResult result = funcOp.walk([&](Operation *op) { + WalkResult result = funcOp.walk([&](Operation* op) { for (auto [index, operand] : llvm::enumerate(op->getOperands())) if (!dominance.dominates(operand, op)) { - op->emitOpError() << "phase 2 produced non-dominating operand " - << index << ": " << operand; + op->emitOpError() << "phase 2 produced non-dominating operand " << index << ": " << operand; return WalkResult::interrupt(); } return WalkResult::advance(); @@ -142,26 +185,23 @@ static LogicalResult verifyDominance(func::FuncOp funcOp) { } // namespace -LogicalResult realizeDeferredCommunication( - func::FuncOp funcOp, - const ScheduledComputeMaterializationResult &materialization) { +LogicalResult realizeDeferredCommunication(func::FuncOp funcOp, + const ScheduledComputeMaterializationResult& materialization, + const SchedulingTarget& target) { IRRewriter rewriter(funcOp.getContext()); eraseUnusedIdentityDeferredCommunications(funcOp, rewriter); auto transfers = buildDeferredTransferPlan(funcOp, materialization); if (failed(transfers)) - return funcOp.emitOpError( - "phase 2 failed to build symbolic transfer families"); + return funcOp.emitOpError("phase 2 failed to build symbolic transfer families"); + if (failed(placeLogicalProcessorsOnPhysicalCores(*transfers, target))) + return failure(); auto schedule = scheduleDeferredCommunication(funcOp, *transfers); - if (failed(schedule) - || failed(verifyPlannedCommunicationDeadlockFree( - funcOp, transfers->stepCounts, *schedule))) - return funcOp.emitOpError( - "phase 2 failed to schedule symbolic communication"); + if (failed(schedule) || failed(verifyPlannedCommunicationDeadlockFree(funcOp, transfers->stepCounts, *schedule))) + return funcOp.emitOpError("phase 2 failed to schedule symbolic communication"); auto boundaries = buildDeferredBoundaryPlan(*transfers, *schedule); if (failed(boundaries)) - return funcOp.emitOpError( - "phase 2 failed to build sparse boundary programs"); + return funcOp.emitOpError("phase 2 failed to build sparse boundary programs"); if (failed(retargetDeferredPublications(funcOp, *transfers)) || failed(replaceFinalGraphPublications(funcOp, *transfers))) @@ -169,28 +209,22 @@ LogicalResult realizeDeferredCommunication( ConstantPool constants(funcOp, rewriter); DeferredEmissionContext context(rewriter, constants); DeferredReplacementMap replacements; - if (failed(realizeDeferredBoundaries( - boundaries->boundaries, boundaries->results, context, replacements))) + if (failed(realizeDeferredBoundaries(boundaries->boundaries, boundaries->results, context, replacements))) return failure(); for (auto [op, replacement] : replacements) { if (op->getResult(0) == replacement) - return op->emitOpError( - "phase 2 cannot replace deferred communication with itself"); + return op->emitOpError("phase 2 cannot replace deferred communication with itself"); op->getResult(0).replaceAllUsesWith(replacement); if (!op->use_empty()) - return op->emitOpError( - "phase 2 cannot erase deferred communication with live uses"); + return op->emitOpError("phase 2 cannot erase deferred communication with live uses"); rewriter.eraseOp(op); } - if (failed(eraseDeferredSourceSelectors(funcOp, rewriter)) - || failed(eraseOldGraph(funcOp, rewriter)) - || failed(verifyDominance(funcOp)) - || failed(verifyRealizedCommunicationDeadlockFree(funcOp, *schedule))) + if (failed(eraseDeferredSourceSelectors(funcOp, rewriter)) || failed(eraseOldGraph(funcOp, rewriter)) + || failed(verifyDominance(funcOp)) || failed(verifyRealizedCommunicationDeadlockFree(funcOp, *schedule))) return failure(); bool deferredRemains = false; funcOp.walk([&](SpatDeferredCommunicationOp deferred) { - deferred.emitOpError( - "phase 2 left an unrealized deferred communication"); + deferred.emitOpError("phase 2 left an unrealized deferred communication"); deferredRemains = true; }); return success(!deferredRemains); diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.hpp index 17a0635..76795e2 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.hpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.hpp @@ -5,9 +5,10 @@ namespace onnx_mlir::spatial { struct ScheduledComputeMaterializationResult; +struct SchedulingTarget; -mlir::LogicalResult realizeDeferredCommunication( - mlir::func::FuncOp funcOp, - const ScheduledComputeMaterializationResult &materialization); +mlir::LogicalResult realizeDeferredCommunication(mlir::func::FuncOp funcOp, + const ScheduledComputeMaterializationResult& materialization, + const SchedulingTarget& target); } // namespace onnx_mlir::spatial diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp index a1890ca..d33a6dc 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp @@ -1,15 +1,14 @@ +#include "mlir/Pass/Pass.h" + +#include "DeferredCommunicationRealization.hpp" #include "ScheduledComputeMaterialization.hpp" #include "ScheduledComputeReport.hpp" #include "ScheduledComputeVerification.hpp" -#include "SpatialDataflowCsvExporter.hpp" -#include "DeferredCommunicationRealization.hpp" - -#include "mlir/Pass/Pass.h" - #include "Scheduling/MergeSchedulingAnalysis.hpp" -#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp" +#include "SpatialDataflowCsvExporter.hpp" #include "src/Accelerators/PIM/Common/PimCommon.hpp" #include "src/Accelerators/PIM/Common/Support/DebugDump.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp" #include "src/Accelerators/PIM/Pass/PIMPasses.h" using namespace mlir; @@ -21,6 +20,10 @@ namespace { struct MergeComputeNodesPass final : PassWrapper> { MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MergeComputeNodesPass) + MergeComputeNodesPass() = default; + explicit MergeComputeNodesPass(const SchedulingTarget& schedulingTarget) + : target(schedulingTarget), hasTarget(true) {} + StringRef getArgument() const override { return "pim-merge-compute-nodes"; } StringRef getDescription() const override { return "Materialize scheduled Spatial compute with deferred communication placeholders."; @@ -28,6 +31,13 @@ struct MergeComputeNodesPass final : PassWrapper 1 && target.averageInterProcessorLatencyNs == 0)) { + moduleOp.emitError("MergeComputeNodes requires an explicit valid Spatial scheduling target"); + signalPassFailure(); + return; + } auto entryFunc = getPimEntryFunc(moduleOp); if (failed(entryFunc)) { moduleOp.emitError("failed to locate the PIM entry function during MergeComputeNodes"); @@ -36,10 +46,10 @@ struct MergeComputeNodesPass final : PassWrapper materialization = - materializeScheduledCompute(funcOp, schedule, rewriter); + materializeScheduledCompute(funcOp, logicalSchedule, rewriter); if (failed(materialization)) { signalPassFailure(); return; @@ -48,7 +58,7 @@ struct MergeComputeNodesPass final : PassWrapperpeftClassPlans, materialization->graphComputeToBlockMap, materialization->materializedSchedules))) { @@ -75,18 +85,14 @@ struct MergeComputeNodesPass final : PassWrappermaterializedSchedules, - "spatial3_scheduled_no_comm", "spatial3"))) { + funcOp, materialization->materializedSchedules, "spatial3_scheduled_no_comm", "spatial3"))) { signalPassFailure(); return; } - dumpScheduledComputeReport(moduleOp, - funcOp, - schedule, - materialization->peftClassPlans, - materialization->materializedSchedules); - if (failed(realizeDeferredCommunication(funcOp, *materialization))) { + dumpScheduledComputeReport( + moduleOp, funcOp, logicalSchedule, materialization->peftClassPlans, materialization->materializedSchedules); + if (failed(realizeDeferredCommunication(funcOp, *materialization, target))) { moduleOp.emitError("MergeComputeNodes phase 2 communication realization failed"); signalPassFailure(); return; @@ -100,11 +106,14 @@ struct MergeComputeNodesPass final : PassWrappermaterializedSchedules, - "spatial4_scheduled", "spatial4"))) { + funcOp, materialization->materializedSchedules, "spatial4_scheduled", "spatial4"))) { signalPassFailure(); } } + +private: + SchedulingTarget target; + bool hasTarget = false; }; } // namespace @@ -112,4 +121,8 @@ struct MergeComputeNodesPass final : PassWrapper createMergeComputeNodesPass() { return std::make_unique(); } +std::unique_ptr createMergeComputeNodesPass(const spatial::SchedulingTarget& target) { + return std::make_unique(target); +} + } // namespace onnx_mlir diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputePlan.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputePlan.hpp index aa177f0..e5d330e 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputePlan.hpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputePlan.hpp @@ -164,7 +164,8 @@ inline size_t getScheduledCpuForComputeInstance(const ComputeInstance &instance, auto batch = dyn_cast(instance.op); assert(batch && instance.laneCount != 0 && "missing scheduled CPU for non-batch compute instance"); assert(instance.laneStart < static_cast(batch.getLaneCount()) && "batch lane start out of range"); - ComputeInstance chunk = getBatchChunkForLane(batch, instance.laneStart); + ComputeInstance chunk = getBatchChunkForLane( + batch, instance.laneStart, schedule.processorCount); auto it = schedule.computeToCpuMap.find(chunk); assert(it != schedule.computeToCpuMap.end() && "missing scheduled CPU for batch chunk"); return it->second; @@ -184,13 +185,17 @@ inline unsigned getScheduledBatchResultArgBase(SpatScheduledComputeBatch schedul return inputArgBase + scheduled.getInputs().size(); } -inline SmallVector collectExpectedGraphComputeBlockKeys(func::FuncOp funcOp) { +inline SmallVector collectExpectedGraphComputeBlockKeys( + func::FuncOp funcOp, size_t processorCount) { SmallVector keys; for (Operation &op : funcOp.getOps()) { if (auto compute = dyn_cast(&op)) keys.push_back(getGraphComputeBlockKey({compute.getOperation(), 0, 1})); else if (auto batch = dyn_cast(&op)) - for (ComputeInstance chunk : getBatchChunksForRange(batch, 0, static_cast(batch.getLaneCount()))) + for (ComputeInstance chunk : + getBatchChunksForRange( + batch, 0, static_cast(batch.getLaneCount()), + processorCount)) keys.push_back(getGraphComputeBlockKey(chunk)); } return keys; diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputePlanning.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputePlanning.cpp index aab4beb..c98f3ae 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputePlanning.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputePlanning.cpp @@ -179,7 +179,9 @@ LogicalResult collectPeftClassOperandsAndResults( for (Value weight : getComputeInstanceWeights(instance)) appendUnique(peftClassPlan.weights, weight); for (Value input : getComputeInstanceInputs(instance)) - if (!getProducerValueRef(input, &instance) && !isDeferredFragmentAssemblyInput(input)) + if (!getProducerValueRef(input, &instance, schedule.processorCount) + && !isDeferredFragmentAssemblyInput( + input, schedule.processorCount)) appendUnique(peftClassPlan.inputs, input); } return success(); @@ -222,7 +224,9 @@ LogicalResult collectPeftClassOperandsAndResults( for (Value weight : getComputeInstanceWeights(instance)) appendUnique(peftClassPlan.weights, weight); for (Value input : getComputeInstanceInputs(instance)) - if (!getProducerValueRef(input, &instance) && !isDeferredFragmentAssemblyInput(input)) + if (!getProducerValueRef(input, &instance, schedule.processorCount) + && !isDeferredFragmentAssemblyInput( + input, schedule.processorCount)) appendUnique(peftClassPlan.inputs, input); } } diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp index be23e03..8f5c959 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp @@ -48,7 +48,8 @@ LogicalResult verifyMaterializedScheduleMapping( } } - for (GraphComputeBlockKey key : collectExpectedGraphComputeBlockKeys(funcOp)) { + for (GraphComputeBlockKey key : + collectExpectedGraphComputeBlockKeys(funcOp, schedule.processorCount)) { if (graphComputeToBlockMap.count(key)) continue; diagnostics.report(key.op, [&](Operation *illegalOp) { @@ -66,10 +67,12 @@ LogicalResult verifyMaterializedScheduleMapping( } } - if (graphComputeToBlockMap.size() != collectExpectedGraphComputeBlockKeys(funcOp).size()) { + const size_t expectedGraphComputeBlockCount = + collectExpectedGraphComputeBlockKeys(funcOp, schedule.processorCount).size(); + if (graphComputeToBlockMap.size() != expectedGraphComputeBlockCount) { diagnostics.report(funcOp.getOperation(), [&](Operation *illegalOp) { illegalOp->emitOpError() << "phase-check expected " - << collectExpectedGraphComputeBlockKeys(funcOp).size() + << expectedGraphComputeBlockCount << " graph compute block mappings but saw " << graphComputeToBlockMap.size(); }); } diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp index 66635a7..b11cda5 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp @@ -12,7 +12,6 @@ #include "llvm/Support/Casting.h" #include -#include #include #include #include @@ -22,7 +21,6 @@ #include "ComputeGraph.hpp" #include "ComputeInstanceUtils.hpp" -#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp" #include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp" #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp" #include "src/Accelerators/PIM/Common/IR/ShapingUtils.hpp" @@ -38,28 +36,25 @@ uint64_t countComputeBodyOperationInstances(Region& body); namespace { -struct PimsimSchedulerCostModel { - static constexpr Cost kDefaultBitwidth = 8; - static constexpr Cost kCorePeriodNs = 1; - static constexpr Cost kLocalMemoryWidthBytes = 64; - static constexpr Cost kLocalMemoryLatencyCycles = 1; - static constexpr Cost kNetworkBusWidthBytes = 8; - static constexpr Cost kNetworkBaseLatencyNs = 2; - static constexpr Cost kNetworkPerHopLatencyNs = 1; - static constexpr Cost kVectorWidth = 16; - static constexpr Cost kVectorLatencyCycles = 4; - static constexpr Cost kDacResolutionBits = 1; - static constexpr Cost kDacLatencyCycles = 1; - static constexpr Cost kDacCount = 128; - static constexpr Cost kXbarReadLatencyNs = 30; - static constexpr Cost kSampleHoldLatencyCycles = 1; - static constexpr Cost kAdcLatencyCycles = 10; - static constexpr Cost kAdcCount = 2; - static constexpr Cost kShiftAdderLatencyCycles = 1; - static constexpr Cost kOutputBufferLatencyCycles = 1; - static constexpr Cost kInputBufferLatencyCycles = 0; +TransferCost addTransferCosts(const TransferCost& lhs, + const TransferCost& rhs) { + return {checkedAdd(lhs.fixed, rhs.fixed), + checkedAdd(lhs.networkFlits, rhs.networkFlits)}; +} + +TransferCost scaleTransferCost(const TransferCost& cost, + Cost numerator, + Cost denominator = 1) { + assert(denominator > 0 && "transfer cost denominator must be positive"); + return {checkedMultiply(cost.fixed, numerator) / denominator, + checkedMultiply(cost.networkFlits, numerator) / denominator}; +} + +struct SchedulerCostModel { static constexpr Cost kFallbackOperationCost = 1; + const SchedulingTarget& target; + static Cost ceilDiv(Cost numerator, Cost denominator) { assert(denominator > 0 && "denominator must be positive"); return (numerator + denominator - 1) / denominator; @@ -72,7 +67,7 @@ struct PimsimSchedulerCostModel { return static_cast(shaped.getNumElements()); } - static Cost getBitwidthOrDefault(Type type) { + Cost getBitwidthOrDefault(Type type) const { if (auto shaped = dyn_cast(type)) type = shaped.getElementType(); if (auto intType = dyn_cast(type)) @@ -81,14 +76,14 @@ struct PimsimSchedulerCostModel { return floatType.getWidth(); if (isa(type)) return 64; - return kDefaultBitwidth; + return target.computeBitwidth; } - static Cost getComputeBitwidth(Type type) { - return std::min(getBitwidthOrDefault(type), kDefaultBitwidth); + Cost getComputeBitwidth(Type type) const { + return std::min(getBitwidthOrDefault(type), target.computeBitwidth); } - static Cost getByteSize(Type type, Cost fallbackBitwidth = kDefaultBitwidth) { + Cost getByteSize(Type type, Cost fallbackBitwidth = 0) const { auto elementCount = getStaticElementCount(type); if (!elementCount) return kFallbackOperationCost; @@ -96,60 +91,40 @@ struct PimsimSchedulerCostModel { if (bitwidth <= 0) bitwidth = getBitwidthOrDefault(type); if (bitwidth <= 0) - bitwidth = kDefaultBitwidth; + bitwidth = target.computeBitwidth; return ceilDiv(checkedMultiply(*elementCount, bitwidth), static_cast(8)); } - static Cost getVectorReadWriteCost(Cost readBytes, Cost writeBytes) { - Cost totalBytes = checkedAdd(readBytes, writeBytes); - return checkedMultiply(ceilDiv(totalBytes, kLocalMemoryWidthBytes), - checkedMultiply(kLocalMemoryLatencyCycles, kCorePeriodNs)); + Cost getVectorReadWriteCost(Cost readBytes, Cost writeBytes) const { + Cost reads = checkedMultiply(ceilDiv(readBytes, target.localMemoryWidthBytes), + target.localMemoryReadLatencyCycles); + Cost writes = checkedMultiply(ceilDiv(writeBytes, target.localMemoryWidthBytes), + target.localMemoryWriteLatencyCycles); + return checkedMultiply(checkedAdd(reads, writes), target.processorPeriodNs); } - static Cost getVectorComputeCost(Cost elementCount) { - return checkedMultiply(ceilDiv(elementCount, kVectorWidth), - checkedMultiply(kVectorLatencyCycles, kCorePeriodNs)); + Cost getVectorComputeCost(Cost elementCount) const { + return checkedMultiply(ceilDiv(elementCount, target.vectorWidth), + checkedMultiply(target.vectorLatencyCycles, target.processorPeriodNs)); } - static Cost getTensorMoveCost(Type type) { + Cost getTensorMoveCost(Type type) const { return getVectorReadWriteCost(getByteSize(type), 0); } - static std::pair estimateMeshShape() { - Cost coreCount = static_cast(std::max(1, coresCount.getValue())); - Cost rows = static_cast(std::sqrt(static_cast(coreCount))); - if (rows == 0) - rows = 1; - while (rows > 1 && coreCount % rows != 0) - --rows; - Cost cols = ceilDiv(coreCount, rows); - return {rows, cols}; + TransferCost getTransferCostFromBytes(Cost bytes) const { + Cost localRead = checkedMultiply(ceilDiv(bytes, target.localMemoryWidthBytes), + checkedMultiply(target.localMemoryReadLatencyCycles, + target.processorPeriodNs)); + Cost localWrite = checkedMultiply(ceilDiv(bytes, target.localMemoryWidthBytes), + checkedMultiply(target.localMemoryWriteLatencyCycles, + target.processorPeriodNs)); + Cost payloadFlits = ceilDiv(bytes, target.transferWidthBytes); + return {checkedAdd(localRead, localWrite), + checkedAdd(static_cast(2), payloadFlits)}; } - static Cost getAverageInterCoreLatencyNs() { - auto [rows, cols] = estimateMeshShape(); - auto averageAxisDistance = [](Cost size) -> Cost { - if (size <= 1) - return 0; - return checkedMultiply(size, size) - 1; - }; - Cost avgRow = averageAxisDistance(rows) / (static_cast(3) * rows); - Cost avgCol = averageAxisDistance(cols) / (static_cast(3) * cols); - return checkedAdd(kNetworkBaseLatencyNs, checkedMultiply(kNetworkPerHopLatencyNs, checkedAdd(avgRow, avgCol))); - } - - static Cost getInterCoreTransferCostFromBytes(Cost bytes) { - Cost localRead = checkedMultiply(ceilDiv(bytes, kLocalMemoryWidthBytes), - checkedMultiply(kLocalMemoryLatencyCycles, kCorePeriodNs)); - Cost localWrite = checkedMultiply(ceilDiv(bytes, kLocalMemoryWidthBytes), - checkedMultiply(kLocalMemoryLatencyCycles, kCorePeriodNs)); - Cost payloadFlits = ceilDiv(bytes, kNetworkBusWidthBytes); - Cost averageNoCLatency = getAverageInterCoreLatencyNs(); - Cost network = checkedMultiply(checkedAdd(static_cast(2), payloadFlits), averageNoCLatency); - return checkedAdd(checkedAdd(localRead, localWrite), network); - } - - static Cost getUnaryVectorCost(Type inputType, Type outputType, bool scalarOutput = false) { + Cost getUnaryVectorCost(Type inputType, Type outputType, bool scalarOutput = false) const { auto maybeElements = getStaticElementCount(inputType); if (!maybeElements) return kFallbackOperationCost; @@ -159,7 +134,7 @@ struct PimsimSchedulerCostModel { return checkedAdd(getVectorReadWriteCost(inputBytes, outputBytes), getVectorComputeCost(*maybeElements)); } - static Cost getBinaryVectorCost(Type lhsType, Type rhsType, Type outputType, bool scalarOutput = false) { + Cost getBinaryVectorCost(Type lhsType, Type rhsType, Type outputType, bool scalarOutput = false) const { auto maybeElements = getStaticElementCount(lhsType); if (!maybeElements) return kFallbackOperationCost; @@ -170,24 +145,35 @@ struct PimsimSchedulerCostModel { return checkedAdd(getVectorReadWriteCost(readBytes, outputBytes), getVectorComputeCost(*maybeElements)); } - static Cost getMatrixComputeLatency(Cost inputBitwidth) { - Cost xbarDim = static_cast(crossbarSize.getValue()); - Cost inputTimes = ceilDiv(inputBitwidth, kDacResolutionBits); - Cost dacTimes = ceilDiv(xbarDim, kDacCount); - Cost adcTimes = ceilDiv(xbarDim, kAdcCount); - Cost frontStage = kInputBufferLatencyCycles + kDacLatencyCycles + kXbarReadLatencyNs + kSampleHoldLatencyCycles; - Cost backPipe = std::max(kAdcLatencyCycles, checkedAdd(kShiftAdderLatencyCycles, kOutputBufferLatencyCycles)); - Cost backStage = checkedAdd(checkedAdd(kAdcLatencyCycles, kShiftAdderLatencyCycles), kOutputBufferLatencyCycles); - backStage = checkedAdd(backStage, checkedMultiply(adcTimes - 1, backPipe)); - Cost totalTimes = checkedMultiply(inputTimes, dacTimes); + Cost getMatrixComputeLatency(Cost inputBitwidth) const { + Cost inputTimes = ceilDiv(inputBitwidth, target.matrixInputResolutionBits); + Cost inputPasses = ceilDiv(target.matrixRows, target.matrixInputParallelism); + Cost outputPasses = ceilDiv(target.matrixColumns, target.matrixOutputParallelism); + Cost readCycles = ceilDiv(target.matrixReadLatencyNs, target.matrixPeriodNs); + Cost frontStage = target.matrixInputBufferLatencyCycles + target.matrixInputLatencyCycles + + readCycles + target.matrixSampleLatencyCycles; + Cost backPipe = std::max(target.matrixOutputLatencyCycles, + checkedAdd(target.matrixShiftLatencyCycles, + target.matrixBufferLatencyCycles)); + Cost backStage = checkedAdd( + checkedAdd(target.matrixOutputLatencyCycles, target.matrixShiftLatencyCycles), + target.matrixBufferLatencyCycles); + backStage = checkedAdd(backStage, checkedMultiply(outputPasses - 1, backPipe)); + Cost totalTimes = checkedMultiply(inputTimes, inputPasses); + if (!target.matrixPipeline) + return checkedMultiply( + checkedMultiply(checkedAdd(frontStage, backStage), totalTimes), + target.matrixPeriodNs); Cost stagePipe = std::max(frontStage, backStage); - return checkedAdd(checkedAdd(frontStage, backStage), - checkedMultiply(totalTimes - 1, stagePipe)); + return checkedMultiply( + checkedAdd(checkedAdd(frontStage, backStage), + checkedMultiply(totalTimes - 1, stagePipe)), + target.matrixPeriodNs); } - static Cost getWvmmCost(Type inputType, Type outputType) { + Cost getWvmmCost(Type inputType, Type outputType) const { Cost inputBitwidth = getComputeBitwidth(inputType); - Cost inputBytes = checkedMultiply(static_cast(crossbarSize.getValue()), + Cost inputBytes = checkedMultiply(target.matrixRows, ceilDiv(inputBitwidth, static_cast(8))); inputBytes = checkedMultiply(inputBytes, static_cast(8)); Cost outputBytes = getByteSize(outputType, getComputeBitwidth(outputType)); @@ -196,22 +182,22 @@ struct PimsimSchedulerCostModel { }; std::optional getStaticTripCount(scf::ForOp loop); -Cost getOperationCost(Operation& op); +Cost getOperationCost(Operation& op, const SchedulerCostModel& costModel); -Cost getRegionCost(Region& body) { +Cost getRegionCost(Region& body, const SchedulerCostModel& costModel) { Cost cost = 0; for (Block& block : body) for (Operation& op : block) - cost = checkedAdd(cost, getOperationCost(op)); + cost = checkedAdd(cost, getOperationCost(op, costModel)); return cost; } -Cost getOperationCost(Operation& op) { +Cost getOperationCost(Operation& op, const SchedulerCostModel& costModel) { if (auto loop = dyn_cast(&op)) { std::optional tripCount = getStaticTripCount(loop); if (!tripCount) - return PimsimSchedulerCostModel::kFallbackOperationCost; - return checkedMultiply(getRegionCost(loop.getRegion()), static_cast(*tripCount)); + return SchedulerCostModel::kFallbackOperationCost; + return checkedMultiply(getRegionCost(loop.getRegion(), costModel), static_cast(*tripCount)); } if (isa(&op)) - return PimsimSchedulerCostModel::getWvmmCost(wvmm.getInput().getType(), wvmm.getOutput().getType()); + return costModel.getWvmmCost(wvmm.getInput().getType(), wvmm.getOutput().getType()); if (auto vvdmul = dyn_cast(&op)) - return PimsimSchedulerCostModel::getBinaryVectorCost( + return costModel.getBinaryVectorCost( vvdmul.getLhs().getType(), vvdmul.getRhs().getType(), vvdmul.getOutput().getType(), /*scalarOutput=*/true); if (auto vadd = dyn_cast(&op)) - return PimsimSchedulerCostModel::getBinaryVectorCost(vadd.getLhs().getType(), vadd.getRhs().getType(), - vadd.getOutput().getType()); + return costModel.getBinaryVectorCost( + vadd.getLhs().getType(), vadd.getRhs().getType(), vadd.getOutput().getType()); if (auto vsub = dyn_cast(&op)) - return PimsimSchedulerCostModel::getBinaryVectorCost(vsub.getLhs().getType(), vsub.getRhs().getType(), - vsub.getOutput().getType()); + return costModel.getBinaryVectorCost( + vsub.getLhs().getType(), vsub.getRhs().getType(), vsub.getOutput().getType()); if (auto vmul = dyn_cast(&op)) - return PimsimSchedulerCostModel::getBinaryVectorCost(vmul.getLhs().getType(), vmul.getRhs().getType(), - vmul.getOutput().getType()); + return costModel.getBinaryVectorCost( + vmul.getLhs().getType(), vmul.getRhs().getType(), vmul.getOutput().getType()); if (auto vmax = dyn_cast(&op)) - return PimsimSchedulerCostModel::getBinaryVectorCost(vmax.getLhs().getType(), vmax.getRhs().getType(), - vmax.getOutput().getType()); + return costModel.getBinaryVectorCost( + vmax.getLhs().getType(), vmax.getRhs().getType(), vmax.getOutput().getType()); if (auto vavg = dyn_cast(&op)) - return PimsimSchedulerCostModel::getUnaryVectorCost(vavg.getInput().getType(), vavg.getOutput().getType(), - /*scalarOutput=*/true); + return costModel.getUnaryVectorCost( + vavg.getInput().getType(), vavg.getOutput().getType(), /*scalarOutput=*/true); if (auto relu = dyn_cast(&op)) - return PimsimSchedulerCostModel::getUnaryVectorCost(relu.getInput().getType(), relu.getOutput().getType()); + return costModel.getUnaryVectorCost(relu.getInput().getType(), relu.getOutput().getType()); if (auto sigm = dyn_cast(&op)) - return PimsimSchedulerCostModel::getUnaryVectorCost(sigm.getInput().getType(), sigm.getOutput().getType()); + return costModel.getUnaryVectorCost(sigm.getInput().getType(), sigm.getOutput().getType()); if (auto softmax = dyn_cast(&op)) { - Cost unary = PimsimSchedulerCostModel::getUnaryVectorCost(softmax.getInput().getType(), softmax.getOutput().getType()); + Cost unary = + costModel.getUnaryVectorCost(softmax.getInput().getType(), softmax.getOutput().getType()); return checkedMultiply(unary, static_cast(4)); } if (auto extract = dyn_cast(&op)) - return PimsimSchedulerCostModel::getTensorMoveCost(extract.getResult().getType()); + return costModel.getTensorMoveCost(extract.getResult().getType()); if (auto insert = dyn_cast(&op)) - return PimsimSchedulerCostModel::getTensorMoveCost(insert.getSource().getType()); + return costModel.getTensorMoveCost(insert.getSource().getType()); Cost nestedCost = 0; for (Region& region : op.getRegions()) - nestedCost = checkedAdd(nestedCost, getRegionCost(region)); - return checkedAdd(PimsimSchedulerCostModel::kFallbackOperationCost, nestedCost); + nestedCost = checkedAdd(nestedCost, getRegionCost(region, costModel)); + return checkedAdd(SchedulerCostModel::kFallbackOperationCost, nestedCost); } std::optional getStaticTripCount(scf::ForOp loop) { @@ -271,8 +258,8 @@ std::optional getStaticTripCount(scf::ForOp loop) { return (distance + stride - 1) / stride; } -Cost getComputeBodyCost(Region& body) { - return getRegionCost(body); +Cost getComputeBodyCost(Region& body, const SchedulerCostModel& costModel) { + return getRegionCost(body, costModel); } uint64_t countOperationInstances(Operation& op) { @@ -348,7 +335,9 @@ std::optional getConstantExtractLane(tensor::ExtractSliceOp extract) { return std::nullopt; } -std::optional getBatchProjectedInputTransferCost(SpatComputeBatch batch, Value input) { +std::optional getBatchProjectedInputTransferCost( + SpatComputeBatch batch, Value input, + const SchedulerCostModel& costModel) { auto inputIt = llvm::find(batch.getInputs(), input); if (inputIt == batch.getInputs().end()) return std::nullopt; @@ -359,7 +348,7 @@ std::optional getBatchProjectedInputTransferCost(SpatComputeBatch batch, V if (!inputArg || !laneArg) return std::nullopt; - Cost projectedCost = 0; + TransferCost projectedCost; for (Operation* user : inputArg->getUsers()) { auto extract = dyn_cast(user); if (!extract || extract.getSource() != *inputArg) @@ -370,11 +359,13 @@ std::optional getBatchProjectedInputTransferCost(SpatComputeBatch batch, V auto resultType = dyn_cast(extract.getResult().getType()); if (!resultType || !resultType.hasStaticShape()) return std::nullopt; - projectedCost = checkedAdd( - projectedCost, PimsimSchedulerCostModel::getInterCoreTransferCostFromBytes(static_cast(getSizeInBytes(resultType)))); + projectedCost = addTransferCosts( + projectedCost, + costModel.getTransferCostFromBytes( + costModel.getByteSize(resultType, costModel.getComputeBitwidth(resultType)))); } - if (projectedCost == 0) + if (projectedCost.fixed == 0 && projectedCost.networkFlits == 0) return std::nullopt; return projectedCost; } @@ -382,7 +373,8 @@ std::optional getBatchProjectedInputTransferCost(SpatComputeBatch batch, V static std::optional> collectProjectedProducerValueRefs(SpatComputeBatch producer, Value input, - const ComputeInstance& consumerInstance) { + const ComputeInstance& consumerInstance, + size_t processorCount) { auto consumer = dyn_cast(consumerInstance.op); if (!consumer) return std::nullopt; @@ -417,7 +409,8 @@ collectProjectedProducerValueRefs(SpatComputeBatch producer, int64_t producerLane = *offset + index * *stride; if (producerLane < 0 || producerLane >= producer.getLaneCount()) return std::nullopt; - ComputeInstance instance = getBatchChunkForLane(producer, static_cast(producerLane)); + ComputeInstance instance = getBatchChunkForLane( + producer, static_cast(producerLane), processorCount); if (llvm::none_of(producers, [&](const ProducerValueRef& ref) { return ref.instance == instance; })) producers.push_back({instance, 0}); } @@ -426,12 +419,16 @@ collectProjectedProducerValueRefs(SpatComputeBatch producer, return producers; } -Cost getInputTransferCost(const ComputeInstance& consumerInstance, Value input) { +TransferCost getInputTransferCost(const ComputeInstance& consumerInstance, + Value input, + const SchedulerCostModel& costModel) { auto inputType = cast(input.getType()); if (auto batch = dyn_cast(consumerInstance.op)) - if (std::optional projectedCost = getBatchProjectedInputTransferCost(batch, input)) + if (std::optional projectedCost = + getBatchProjectedInputTransferCost(batch, input, costModel)) return *projectedCost; - return PimsimSchedulerCostModel::getInterCoreTransferCostFromBytes(static_cast(getSizeInBytes(inputType))); + return costModel.getTransferCostFromBytes( + costModel.getByteSize(inputType, costModel.getComputeBitwidth(inputType))); } uint32_t getLaneOverlapCount(const ComputeInstance& lhs, const ComputeInstance& rhs) { @@ -442,15 +439,20 @@ uint32_t getLaneOverlapCount(const ComputeInstance& lhs, const ComputeInstance& : 0; } -Cost scaleTransferCostByLaneCount(Cost totalCost, uint32_t totalLaneCount, uint32_t fragmentLaneCount) { +TransferCost scaleTransferCostByLaneCount( + const TransferCost& totalCost, uint32_t totalLaneCount, + uint32_t fragmentLaneCount) { assert(totalLaneCount > 0 && "laneCount must be positive"); assert(fragmentLaneCount > 0 && "fragmentLaneCount must be positive"); if (fragmentLaneCount >= totalLaneCount) return totalCost; - return checkedMultiply(totalCost, static_cast(fragmentLaneCount)) / static_cast(totalLaneCount); + return scaleTransferCost(totalCost, static_cast(fragmentLaneCount), + static_cast(totalLaneCount)); } -SmallVector collectProducerValueRefs(Value value, const ComputeInstance& consumerInstance) { +SmallVector collectProducerValueRefs( + Value value, const ComputeInstance& consumerInstance, + size_t processorCount) { SmallVector producers; Operation* op = value.getDefiningOp(); if (!op) @@ -461,13 +463,16 @@ SmallVector collectProducerValueRefs(Value value, const Com auto batch = dyn_cast_or_null(source.getDefiningOp()); if (batch && batch.getNumResults() != 0) { if (std::optional lane = getConstantExtractLane(extract)) { - ComputeInstance instance = getBatchChunkForLane(batch, *lane); + ComputeInstance instance = + getBatchChunkForLane(batch, *lane, processorCount); producers.push_back({instance, 0}); return producers; } for (ComputeInstance instance : - getBatchChunksForRange(batch, 0, static_cast(batch.getLaneCount()))) + getBatchChunksForRange(batch, 0, + static_cast(batch.getLaneCount()), + processorCount)) producers.push_back({instance, 0}); return producers; } @@ -488,19 +493,24 @@ SmallVector collectProducerValueRefs(Value value, const Com if (auto batch = dyn_cast(op)) { if (batch.getNumResults() != 0) { - if (auto projected = collectProjectedProducerValueRefs(batch, value, consumerInstance)) + if (auto projected = collectProjectedProducerValueRefs( + batch, value, consumerInstance, processorCount)) return *projected; - std::optional producer = getProducerValueRef(value, &consumerInstance); + std::optional producer = + getProducerValueRef(value, &consumerInstance, processorCount); if (!producer) return producers; for (ComputeInstance instance : - getBatchChunksForRange(batch, producer->instance.laneStart, producer->instance.laneCount)) + getBatchChunksForRange(batch, producer->instance.laneStart, + producer->instance.laneCount, + processorCount)) producers.push_back({instance, 0}); return producers; } uint32_t lane = cast(value).getResultNumber(); - ComputeInstance instance = getBatchChunkForLane(batch, lane); + ComputeInstance instance = + getBatchChunkForLane(batch, lane, processorCount); producers.push_back({instance, lane - instance.laneStart}); return producers; } @@ -508,18 +518,23 @@ SmallVector collectProducerValueRefs(Value value, const Com return producers; } -Cost getProducerTransferCost(Value input, - const ComputeInstance& consumerInstance, - const ProducerValueRef& producerRef) { - Cost transferCost = getInputTransferCost(consumerInstance, input); +TransferCost getProducerTransferCost( + Value input, const ComputeInstance& consumerInstance, + const ProducerValueRef& producerRef, + const SchedulerCostModel& costModel) { + TransferCost transferCost = + getInputTransferCost(consumerInstance, input, costModel); auto producerBatch = dyn_cast(producerRef.instance.op); if (!producerBatch || producerBatch.getNumResults() == 0) return transferCost; if (auto consumerBatch = dyn_cast(consumerInstance.op)) { - if (std::optional projectedCost = getBatchProjectedInputTransferCost(consumerBatch, input)) { + if (std::optional projectedCost = + getBatchProjectedInputTransferCost(consumerBatch, input, costModel)) { uint32_t overlapLaneCount = getLaneOverlapCount(consumerInstance, producerRef.instance); - return checkedMultiply(*projectedCost, static_cast(std::max(1, overlapLaneCount))); + return scaleTransferCost( + *projectedCost, + static_cast(std::max(1, overlapLaneCount))); } } @@ -527,8 +542,8 @@ Cost getProducerTransferCost(Value input, transferCost, static_cast(producerBatch.getLaneCount()), producerRef.instance.laneCount); } -static CrossbarWeight getOpaqueCrossbarWeight(Value value, std::optional lane) { - CrossbarWeight weight; +static ResidentWeight getOpaqueResidentWeight(Value value, std::optional lane) { + ResidentWeight weight; weight.opaqueValue = value; weight.opaqueLane = lane.value_or(std::numeric_limits::max()); return weight; @@ -619,7 +634,7 @@ static FailureOr> evaluateIndexList(ArrayRef(root)) { if (auto compute = dyn_cast(owner)) { for (auto [index, operand] : llvm::enumerate(compute.getWeights())) @@ -637,11 +652,11 @@ static Value resolveCrossbarWeightRoot(Operation* owner, Value root) { return root; } -static CrossbarWeight completeCrossbarWeight(Value root, +static ResidentWeight completeResidentWeight(Value root, SmallVector offsets, SmallVector sizes, SmallVector strides) { - CrossbarWeight weight; + ResidentWeight weight; weight.root = root; if (auto constant = root.getDefiningOp()) weight.rootAttr = static_cast(constant.getValue()); @@ -651,14 +666,14 @@ static CrossbarWeight completeCrossbarWeight(Value root, return weight; } -static FailureOr getStaticCrossbarWeight(Operation* owner, +static FailureOr getStaticResidentWeight(Operation* owner, Value value, const DenseMap& bindings, std::optional lane, Value laneArg) { if (auto extract = value.getDefiningOp()) { - FailureOr sourceWeight = - getStaticCrossbarWeight(owner, extract.getSource(), bindings, lane, laneArg); + FailureOr sourceWeight = + getStaticResidentWeight(owner, extract.getSource(), bindings, lane, laneArg); auto offsets = evaluateIndexList(extract.getMixedOffsets(), bindings, lane, laneArg); auto sizes = evaluateIndexList(extract.getMixedSizes(), bindings, lane, laneArg); auto strides = evaluateIndexList(extract.getMixedStrides(), bindings, lane, laneArg); @@ -678,7 +693,7 @@ static FailureOr getStaticCrossbarWeight(Operation* owner, return *sourceWeight; } - Value root = resolveCrossbarWeightRoot(owner, value); + Value root = resolveResidentWeightRoot(owner, value); auto type = dyn_cast(root.getType()); if (!type || !type.hasStaticShape()) return failure(); @@ -686,18 +701,18 @@ static FailureOr getStaticCrossbarWeight(Operation* owner, SmallVector offsets(type.getRank(), 0); SmallVector sizes(type.getShape().begin(), type.getShape().end()); SmallVector strides(type.getRank(), 1); - return completeCrossbarWeight(root, std::move(offsets), std::move(sizes), std::move(strides)); + return completeResidentWeight(root, std::move(offsets), std::move(sizes), std::move(strides)); } -static void addCrossbarWeight(CrossbarUsage& usage, CrossbarWeight weight) { - if (!containsCrossbarWeight(usage, weight)) +static void addResidentWeight(ResidentWeightSet& usage, ResidentWeight weight) { + if (!containsResidentWeight(usage, weight)) usage.push_back(std::move(weight)); } -static void collectCrossbarWeightsFromOp(Operation* op, +static void collectResidentWeightsFromOp(Operation* op, Operation* owner, DenseMap& bindings, - CrossbarUsage& usage, + ResidentWeightSet& usage, Value laneArg, std::optional lane) { if (auto loop = dyn_cast(op)) { @@ -710,36 +725,37 @@ static void collectCrossbarWeightsFromOp(Operation* op, for (int64_t iv = *lb; iv < *ub; iv += *step) { bindings[loop.getInductionVar()] = iv; for (Operation& nested : loop.getBody()->without_terminator()) - collectCrossbarWeightsFromOp(&nested, owner, bindings, usage, laneArg, lane); + collectResidentWeightsFromOp(&nested, owner, bindings, usage, laneArg, lane); } bindings.erase(loop.getInductionVar()); return; } if (auto vmm = dyn_cast(op)) { - FailureOr weight = getStaticCrossbarWeight(owner, vmm.getWeight(), bindings, lane, laneArg); + FailureOr weight = getStaticResidentWeight(owner, vmm.getWeight(), bindings, lane, laneArg); if (failed(weight)) { - addCrossbarWeight(usage, getOpaqueCrossbarWeight(vmm.getWeight(), lane)); + addResidentWeight(usage, getOpaqueResidentWeight(vmm.getWeight(), lane)); return; } - addCrossbarWeight(usage, *weight); + addResidentWeight(usage, *weight); return; } for (Region& region : op->getRegions()) for (Block& block : region) for (Operation& nested : block.without_terminator()) - collectCrossbarWeightsFromOp(&nested, owner, bindings, usage, laneArg, lane); + collectResidentWeightsFromOp(&nested, owner, bindings, usage, laneArg, lane); } std::vector aggregateEdges(llvm::ArrayRef edges) { - llvm::DenseMap, Cost> edgeCosts; + llvm::DenseMap, TransferCost> edgeCosts; for (const ComputeGraphEdge& edge : edges) { if (edge.source == edge.target) continue; auto inserted = edgeCosts.try_emplace({edge.source, edge.target}, edge.transferCost); if (!inserted.second) - inserted.first->second = checkedAdd(inserted.first->second, edge.transferCost); + inserted.first->second = + addTransferCosts(inserted.first->second, edge.transferCost); } std::vector aggregatedEdges; @@ -770,8 +786,8 @@ uint64_t countComputeBodyOperationInstances(Region& body) { return instances; } -CrossbarUsage collectDistinctCrossbarWeights(Operation* owner, std::optional lane) { - CrossbarUsage usage; +ResidentWeightSet collectDistinctResidentWeights(Operation* owner, std::optional lane) { + ResidentWeightSet usage; DenseMap bindings; Value laneArg; if (auto batch = dyn_cast(owner)) @@ -781,54 +797,87 @@ CrossbarUsage collectDistinctCrossbarWeights(Operation* owner, std::optionalgetRegions()) for (Block& block : region) for (Operation& op : block.without_terminator()) - collectCrossbarWeightsFromOp(&op, owner, bindings, usage, laneArg, lane); + collectResidentWeightsFromOp(&op, owner, bindings, usage, laneArg, lane); return usage; } -Cost getComputeInstanceCost(const ComputeInstance& instance) { +Cost getComputeInstanceCost(const ComputeInstance& instance, const SchedulingTarget& target) { + SchedulerCostModel costModel {target}; if (auto spatCompute = dyn_cast(instance.op)) - return getComputeBodyCost(spatCompute.getBody()); + return getComputeBodyCost(spatCompute.getBody(), costModel); auto batch = cast(instance.op); - return checkedMultiply(getComputeBodyCost(batch.getBody()), static_cast(instance.laneCount)); + return checkedMultiply( + getComputeBodyCost(batch.getBody(), costModel), static_cast(instance.laneCount)); } -bool containsCrossbarWeight(ArrayRef usage, const CrossbarWeight& weight) { +bool containsResidentWeight(ArrayRef usage, const ResidentWeight& weight) { return llvm::is_contained(usage, weight); } -unsigned countCrossbarOverlap(ArrayRef lhs, ArrayRef rhs) { +unsigned countResidentWeightOverlap(ArrayRef lhs, ArrayRef rhs) { unsigned overlap = 0; - for (const CrossbarWeight& weight : rhs) - if (containsCrossbarWeight(lhs, weight)) + for (const ResidentWeight& weight : rhs) + if (containsResidentWeight(lhs, weight)) ++overlap; return overlap; } -size_t getCrossbarUnionSize(ArrayRef lhs, ArrayRef rhs) { +size_t getResidentWeightUnionSize(ArrayRef lhs, ArrayRef rhs) { size_t size = lhs.size(); - for (const CrossbarWeight& weight : rhs) - if (!containsCrossbarWeight(lhs, weight)) + for (const ResidentWeight& weight : rhs) + if (!containsResidentWeight(lhs, weight)) ++size; return size; } -void insertCrossbarWeights(CrossbarUsage& usage, ArrayRef weights) { - for (const CrossbarWeight& weight : weights) - addCrossbarWeight(usage, weight); +void insertResidentWeights(ResidentWeightSet& usage, ArrayRef weights) { + for (const ResidentWeight& weight : weights) + addResidentWeight(usage, weight); } -CrossbarUsage getComputeInstanceCrossbarUsage(const ComputeInstance& instance) { - CrossbarUsage usage; +ResidentWeightSet getComputeInstanceResidentWeights(const ComputeInstance& instance) { + ResidentWeightSet usage; if (isa(instance.op)) - return collectDistinctCrossbarWeights(instance.op); + return collectDistinctResidentWeights(instance.op); for (uint32_t lane = instance.laneStart; lane < instance.laneStart + instance.laneCount; ++lane) - insertCrossbarWeights(usage, collectDistinctCrossbarWeights(instance.op, lane)); + insertResidentWeights(usage, collectDistinctResidentWeights(instance.op, lane)); return usage; } -ComputeGraph buildComputeGraph(Operation* entryOp) { +ResidentWeightSet getComputeInstanceResidentWeights(const ComputeInstance& instance, + const SchedulingTarget& target) { + ResidentWeightSet tiled; + for (const ResidentWeight& weight : getComputeInstanceResidentWeights(instance)) { + if (weight.opaqueValue || weight.sizes.size() < 2 + || target.matrixRows == 0 || target.matrixColumns == 0) { + addResidentWeight(tiled, weight); + continue; + } + + const size_t rowDim = weight.sizes.size() - 2; + const size_t columnDim = weight.sizes.size() - 1; + for (int64_t row = 0; row < weight.sizes[rowDim]; + row += static_cast(target.matrixRows)) { + for (int64_t column = 0; column < weight.sizes[columnDim]; + column += static_cast(target.matrixColumns)) { + ResidentWeight tile = weight; + tile.offsets[rowDim] += row * tile.strides[rowDim]; + tile.offsets[columnDim] += column * tile.strides[columnDim]; + tile.sizes[rowDim] = + std::min(target.matrixRows, weight.sizes[rowDim] - row); + tile.sizes[columnDim] = + std::min(target.matrixColumns, weight.sizes[columnDim] - column); + addResidentWeight(tiled, std::move(tile)); + } + } + } + return tiled; +} + +ComputeGraph buildComputeGraph(Operation* entryOp, const SchedulingTarget& target) { ComputeGraph graph; + SchedulerCostModel costModel {target}; for (Region& region : entryOp->getRegions()) { for (Block& block : region) { @@ -838,20 +887,26 @@ ComputeGraph buildComputeGraph(Operation* entryOp) { continue; ComputeInstance instance {spatCompute.getOperation(), 0, 1}; size_t index = graph.nodes.size(); - graph.nodes.push_back( - {instance, getComputeInstanceCost(instance), getComputeInstanceCrossbarUsage(instance), index}); + graph.nodes.push_back({instance, + getComputeInstanceCost(instance, target), + getComputeInstanceResidentWeights(instance, target), + index}); graph.instanceToIndex[instance] = index; continue; } if (auto batch = dyn_cast(&op)) { if (isUsedAsWeightOnly(batch.getOperation())) continue; - size_t chunkCount = getBatchChunkTargetCount(batch); + size_t chunkCount = + getBatchChunkTargetCount(batch, target.processorCount); for (size_t chunkIndex = 0; chunkIndex < chunkCount; ++chunkIndex) { - ComputeInstance instance = getBatchChunkForIndex(batch, chunkIndex); + ComputeInstance instance = getBatchChunkForIndex( + batch, chunkIndex, target.processorCount); size_t index = graph.nodes.size(); - graph.nodes.push_back( - {instance, getComputeInstanceCost(instance), getComputeInstanceCrossbarUsage(instance), index}); + graph.nodes.push_back({instance, + getComputeInstanceCost(instance, target), + getComputeInstanceResidentWeights(instance, target), + index}); graph.instanceToIndex[instance] = index; } } @@ -863,12 +918,15 @@ ComputeGraph buildComputeGraph(Operation* entryOp) { for (const auto& [targetIndex, node] : llvm::enumerate(graph.nodes)) { llvm::SmallVector inputs = getComputeInstanceInputs(node.instance); for (Value input : inputs) { - for (const ProducerValueRef& producerRef : collectProducerValueRefs(input, node.instance)) { + for (const ProducerValueRef& producerRef : + collectProducerValueRefs(input, node.instance, + target.processorCount)) { auto producerIt = graph.instanceToIndex.find(producerRef.instance); if (producerIt == graph.instanceToIndex.end()) continue; - rawEdges.push_back( - {producerIt->second, targetIndex, getProducerTransferCost(input, node.instance, producerRef)}); + rawEdges.push_back({producerIt->second, + targetIndex, + getProducerTransferCost(input, node.instance, producerRef, costModel)}); } } } diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.hpp index a7c3c97..4ef76d7 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.hpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.hpp @@ -11,9 +11,10 @@ #include #include "ComputeInstance.hpp" +#include "SchedulingTarget.hpp" #include "Utils.hpp" -struct CrossbarWeight { +struct ResidentWeight { mlir::Value root; mlir::Attribute rootAttr; llvm::SmallVector offsets; @@ -22,14 +23,14 @@ struct CrossbarWeight { mlir::Value opaqueValue; uint32_t opaqueLane = 0; - bool operator==(const CrossbarWeight& other) const { + bool operator==(const ResidentWeight& other) const { bool sameRoot = rootAttr && other.rootAttr ? rootAttr == other.rootAttr : root == other.root; return sameRoot && offsets == other.offsets && sizes == other.sizes && strides == other.strides && opaqueValue == other.opaqueValue && opaqueLane == other.opaqueLane; } }; -using CrossbarUsage = llvm::SmallVector; +using ResidentWeightSet = llvm::SmallVector; namespace onnx_mlir { namespace spatial { @@ -37,36 +38,47 @@ namespace spatial { struct ComputeGraphNode { ComputeInstance instance; Cost cost = 0; - CrossbarUsage crossbarUsage; + ResidentWeightSet residentWeights; size_t originalOrder = 0; }; +struct TransferCost { + Cost fixed = 0; + Cost networkFlits = 0; +}; + struct ComputeGraphEdge { size_t source = 0; size_t target = 0; - Cost transferCost = 0; + TransferCost transferCost; }; struct ComputeGraph { std::vector nodes; std::vector edges; - std::vector>> successors; - std::vector>> predecessors; + std::vector>> successors; + std::vector>> predecessors; llvm::DenseMap instanceToIndex; }; -ComputeGraph buildComputeGraph(mlir::Operation* entryOp); +ComputeGraph buildComputeGraph(mlir::Operation* entryOp, const SchedulingTarget& target); bool verifyAcyclic(const ComputeGraph& graph); uint64_t countComputeBodyInstructions(mlir::Region& body); uint64_t countComputeBodyOperationInstances(mlir::Region& body); -Cost getComputeInstanceCost(const ComputeInstance& instance); -CrossbarUsage collectDistinctCrossbarWeights(mlir::Operation* owner, std::optional lane = std::nullopt); -CrossbarUsage getComputeInstanceCrossbarUsage(const ComputeInstance& instance); -bool containsCrossbarWeight(llvm::ArrayRef usage, const CrossbarWeight& weight); -unsigned countCrossbarOverlap(llvm::ArrayRef lhs, llvm::ArrayRef rhs); -size_t getCrossbarUnionSize(llvm::ArrayRef lhs, llvm::ArrayRef rhs); -void insertCrossbarWeights(CrossbarUsage& usage, llvm::ArrayRef weights); +Cost getComputeInstanceCost(const ComputeInstance& instance, const SchedulingTarget& target); +ResidentWeightSet collectDistinctResidentWeights(mlir::Operation* owner, + std::optional lane = std::nullopt); +ResidentWeightSet getComputeInstanceResidentWeights(const ComputeInstance& instance); +ResidentWeightSet getComputeInstanceResidentWeights(const ComputeInstance& instance, + const SchedulingTarget& target); +bool containsResidentWeight(llvm::ArrayRef usage, const ResidentWeight& weight); +unsigned countResidentWeightOverlap(llvm::ArrayRef lhs, + llvm::ArrayRef rhs); +size_t getResidentWeightUnionSize(llvm::ArrayRef lhs, + llvm::ArrayRef rhs); +void insertResidentWeights(ResidentWeightSet& usage, + llvm::ArrayRef weights); } // namespace spatial } // namespace onnx_mlir diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp index 1e3906e..204402a 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp @@ -2,23 +2,15 @@ #include "mlir/Dialect/Tensor/IR/Tensor.h" #include -#include #include #include "ComputeInstanceUtils.hpp" -#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp" using namespace mlir; namespace onnx_mlir { namespace spatial { -size_t getSchedulingCpuBudget() { - if (coresCount.getValue() > 0) - return static_cast(coresCount.getValue()); - return std::numeric_limits::max(); -} - static BatchChunkRange getBatchChunkRange(int32_t laneCount, size_t chunkCount, size_t chunkIndex) { assert(laneCount > 0 && "laneCount must be positive"); assert(chunkIndex < chunkCount && "chunkIndex out of range"); @@ -33,22 +25,27 @@ static BatchChunkRange getBatchChunkRange(int32_t laneCount, size_t chunkCount, return {static_cast(start), static_cast(count)}; } -size_t getBatchChunkTargetCount(SpatComputeBatch batch) { +size_t getBatchChunkTargetCount(SpatComputeBatch batch, size_t processorCount) { int32_t laneCount = batch.getLaneCount(); assert(laneCount > 0 && "laneCount must be positive"); - return std::min(static_cast(laneCount), getSchedulingCpuBudget()); + assert(processorCount > 0 && "processorCount must be positive"); + return std::min(static_cast(laneCount), processorCount); } -BatchChunkRange getBatchChunkRange(SpatComputeBatch batch, size_t chunkIndex) { - return getBatchChunkRange(batch.getLaneCount(), getBatchChunkTargetCount(batch), chunkIndex); +BatchChunkRange getBatchChunkRange(SpatComputeBatch batch, size_t chunkIndex, + size_t processorCount) { + return getBatchChunkRange(batch.getLaneCount(), + getBatchChunkTargetCount(batch, processorCount), + chunkIndex); } -size_t getBatchChunkIndexForLane(SpatComputeBatch batch, uint32_t lane) { +size_t getBatchChunkIndexForLane(SpatComputeBatch batch, uint32_t lane, + size_t processorCount) { int32_t laneCount = batch.getLaneCount(); assert(laneCount > 0 && "laneCount must be positive"); assert(lane < static_cast(laneCount) && "lane out of range"); - size_t chunkCount = getBatchChunkTargetCount(batch); + size_t chunkCount = getBatchChunkTargetCount(batch, processorCount); size_t laneCountSize = static_cast(laneCount); size_t baseChunkSize = laneCountSize / chunkCount; size_t remainder = laneCountSize % chunkCount; @@ -61,17 +58,22 @@ size_t getBatchChunkIndexForLane(SpatComputeBatch batch, uint32_t lane) { return remainder + ((laneIndex - largerChunkLanes) / baseChunkSize); } -ComputeInstance getBatchChunkForIndex(SpatComputeBatch batch, size_t chunkIndex) { - BatchChunkRange chunk = getBatchChunkRange(batch, chunkIndex); +ComputeInstance getBatchChunkForIndex(SpatComputeBatch batch, size_t chunkIndex, + size_t processorCount) { + BatchChunkRange chunk = getBatchChunkRange(batch, chunkIndex, processorCount); return {batch.getOperation(), chunk.laneStart, chunk.laneCount}; } -ComputeInstance getBatchChunkForLane(SpatComputeBatch batch, uint32_t lane) { - return getBatchChunkForIndex(batch, getBatchChunkIndexForLane(batch, lane)); +ComputeInstance getBatchChunkForLane(SpatComputeBatch batch, uint32_t lane, + size_t processorCount) { + return getBatchChunkForIndex( + batch, getBatchChunkIndexForLane(batch, lane, processorCount), + processorCount); } llvm::SmallVector -getBatchChunksForRange(SpatComputeBatch batch, uint32_t laneStart, uint32_t laneCount) { +getBatchChunksForRange(SpatComputeBatch batch, uint32_t laneStart, + uint32_t laneCount, size_t processorCount) { llvm::SmallVector chunks; if (laneCount == 0) return chunks; @@ -80,11 +82,13 @@ getBatchChunksForRange(SpatComputeBatch batch, uint32_t laneStart, uint32_t lane assert(laneEnd >= laneStart && "lane range overflow"); assert(laneEnd <= static_cast(batch.getLaneCount()) && "lane range out of bounds"); - size_t firstChunk = getBatchChunkIndexForLane(batch, laneStart); - size_t lastChunk = getBatchChunkIndexForLane(batch, laneEnd - 1); + size_t firstChunk = + getBatchChunkIndexForLane(batch, laneStart, processorCount); + size_t lastChunk = + getBatchChunkIndexForLane(batch, laneEnd - 1, processorCount); chunks.reserve(lastChunk - firstChunk + 1); for (size_t chunkIndex = firstChunk; chunkIndex <= lastChunk; ++chunkIndex) - chunks.push_back(getBatchChunkForIndex(batch, chunkIndex)); + chunks.push_back(getBatchChunkForIndex(batch, chunkIndex, processorCount)); return chunks; } @@ -150,7 +154,9 @@ static std::optional getResultfulBatchProducerValueRef(SpatCom }; } -std::optional getProducerValueRef(Value value, const ComputeInstance* consumerInstance) { +std::optional getProducerValueRef( + Value value, const ComputeInstance* consumerInstance, + size_t processorCount) { Operation* op = value.getDefiningOp(); if (!op) return std::nullopt; @@ -187,7 +193,8 @@ std::optional getProducerValueRef(Value value, const ComputeIn if (batch.getNumResults() != 0) return getResultfulBatchProducerValueRef(batch, value, consumerInstance); uint32_t lane = cast(value).getResultNumber(); - ComputeInstance instance = getBatchChunkForLane(batch, lane); + ComputeInstance instance = + getBatchChunkForLane(batch, lane, processorCount); size_t resultIndex = lane - instance.laneStart; return ProducerValueRef {instance, resultIndex}; } @@ -195,8 +202,11 @@ std::optional getProducerValueRef(Value value, const ComputeIn return std::nullopt; } -std::optional getComputeProducerInstance(Value value, const ComputeInstance* consumerInstance) { - if (std::optional producer = getProducerValueRef(value, consumerInstance)) +std::optional getComputeProducerInstance( + Value value, const ComputeInstance* consumerInstance, + size_t processorCount) { + if (std::optional producer = + getProducerValueRef(value, consumerInstance, processorCount)) return producer->instance; return std::nullopt; } diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.hpp index ad556f0..22a23ae 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.hpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.hpp @@ -26,19 +26,25 @@ struct BatchChunkRange { uint32_t laneCount = 0; }; -size_t getSchedulingCpuBudget(); -size_t getBatchChunkTargetCount(SpatComputeBatch batch); -BatchChunkRange getBatchChunkRange(SpatComputeBatch batch, size_t chunkIndex); -size_t getBatchChunkIndexForLane(SpatComputeBatch batch, uint32_t lane); -ComputeInstance getBatchChunkForIndex(SpatComputeBatch batch, size_t chunkIndex); -ComputeInstance getBatchChunkForLane(SpatComputeBatch batch, uint32_t lane); +size_t getBatchChunkTargetCount(SpatComputeBatch batch, size_t processorCount); +BatchChunkRange getBatchChunkRange(SpatComputeBatch batch, size_t chunkIndex, + size_t processorCount); +size_t getBatchChunkIndexForLane(SpatComputeBatch batch, uint32_t lane, + size_t processorCount); +ComputeInstance getBatchChunkForIndex(SpatComputeBatch batch, size_t chunkIndex, + size_t processorCount); +ComputeInstance getBatchChunkForLane(SpatComputeBatch batch, uint32_t lane, + size_t processorCount); llvm::SmallVector -getBatchChunksForRange(SpatComputeBatch batch, uint32_t laneStart, uint32_t laneCount); +getBatchChunksForRange(SpatComputeBatch batch, uint32_t laneStart, + uint32_t laneCount, size_t processorCount); std::optional getProducerValueRef(mlir::Value value, - const ComputeInstance* consumerInstance = nullptr); + const ComputeInstance* consumerInstance, + size_t processorCount); std::optional getComputeProducerInstance(mlir::Value value, - const ComputeInstance* consumerInstance = nullptr); + const ComputeInstance* consumerInstance, + size_t processorCount); llvm::SmallVector getComputeInstanceInputs(const ComputeInstance& instance); llvm::SmallVector getComputeInstanceWeights(const ComputeInstance& instance); diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/MergeSchedule.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/MergeSchedule.hpp index 4990658..19af1cb 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/MergeSchedule.hpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/MergeSchedule.hpp @@ -13,6 +13,7 @@ namespace onnx_mlir { namespace spatial { struct MergeScheduleResult { + size_t processorCount = 0; std::vector dominanceOrderCompute; llvm::DenseMap computeToCpuMap; llvm::DenseMap computeToCpuSlotMap; diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp index cf366f8..e4b6fe5 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp @@ -9,7 +9,6 @@ #include "ComputeGraph.hpp" #include "MergeSchedulingAnalysis.hpp" #include "PeftScheduler.hpp" -#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp" namespace onnx_mlir { namespace spatial { @@ -18,8 +17,7 @@ namespace { void verifySchedule(const ComputeGraph& graph, const MergeScheduleResult& result, - unsigned long crossbarCapacity, - size_t processorCount) { + const SchedulingTarget& target) { llvm::DenseMap>> tasksByCpu; tasksByCpu.reserve(result.cpuToLastComputeMap.size()); @@ -44,13 +42,13 @@ void verifySchedule(const ComputeGraph& graph, return lhs.second < rhs.second; }); - CrossbarUsage usedCrossbars; + ResidentWeightSet residentWeights; for (size_t slot = 0; slot < scheduledTasks.size(); ++slot) { if (scheduledTasks[slot].first != slot) llvm::report_fatal_error("merge scheduling: CPU slots are not contiguous"); - insertCrossbarWeights(usedCrossbars, graph.nodes[scheduledTasks[slot].second].crossbarUsage); - if (usedCrossbars.size() > crossbarCapacity) - llvm::report_fatal_error("merge scheduling: CPU crossbar capacity exceeded"); + insertResidentWeights(residentWeights, graph.nodes[scheduledTasks[slot].second].residentWeights); + if (residentWeights.size() > target.residentWeightCapacity) + llvm::report_fatal_error("merge scheduling: processor resident-weight capacity exceeded"); } const ComputeInstance expectedLast = graph.nodes[scheduledTasks.back().second].instance; @@ -63,20 +61,21 @@ void verifySchedule(const ComputeGraph& graph, for (const ComputeGraphEdge& edge : graph.edges) { const ComputeInstance source = graph.nodes[edge.source].instance; - const ComputeInstance target = graph.nodes[edge.target].instance; + const ComputeInstance destination = graph.nodes[edge.target].instance; const size_t sourceCpu = result.computeToCpuMap.lookup(source); - const size_t targetCpu = result.computeToCpuMap.lookup(target); + const size_t targetCpu = result.computeToCpuMap.lookup(destination); const size_t sourceSlot = result.computeToCpuSlotMap.lookup(source); - const size_t targetSlot = result.computeToCpuSlotMap.lookup(target); + const size_t targetSlot = result.computeToCpuSlotMap.lookup(destination); const Time sourceStart = static_cast