finally fast googlenet with correct latency artifacts for fair comparison

.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.

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`

README.md

@@ -100,6 +100,9 @@ options; `onnx-mlir --help` lists the inherited ONNX-MLIR options.
 - `--core-count=<N>` - required positive core count for PIM compilation.
 - `--crossbar-size=<N>` - crossbar width/height. Default in code is `128`.
 - `--crossbar-count=<N>` - crossbars per core. Default in code is `64`.
+- `--pim-target-config=<PATH>` - 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=<summary|none>` - 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.
 

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>]::<f32, f64> as InstructionType));
         tmp.insert((64_usize,32_usize), ([<$id _impl>]::<f64, f32> as InstructionType));
         tmp.insert((64_usize,64_usize), ([<$id _impl>]::<f64, f64> as InstructionType));
-        //TODO WTF WHY
-        tmp.insert((8_usize,8_usize), ([<$id _impl>]::<f32, f32> 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<F>(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::<F>())
+        .context("vector byte length overflow")?;
+    Ok((element_count, byte_len))
+}
+
 #[inline(never)]
 pub fn setbw(cores: &mut CPU, data: InstructionData) -> Result<InstructionStatus> {
     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::<F>(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::<F>()?;
     let (load1, load2) = (loads[0], loads[1]);
     let res: Vec<F> = load1
@@ -374,7 +380,7 @@ where
         .map(|(&a, &b)| a + b)
         .collect();
     ensure!(
-        imm_len / size_of::<F>() == 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::<F>(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::<F>()?;
     let (load1, load2) = (loads[0], loads[1]);
     let res: Vec<F> = load1
@@ -418,7 +424,7 @@ where
         .map(|(&a, &b)| a - b)
         .collect();
     ensure!(
-        imm_len / size_of::<F>() == 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::<F>(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::<F>()?;
     let (load1, load2) = (loads[0], loads[1]);
     let res: Vec<F> = load1
@@ -459,7 +465,7 @@ where
         .map(|(&a, &b)| a * b)
         .collect();
     ensure!(
-        imm_len / size_of::<F>() == 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::<F>(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::<F>()?;
     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::<F>(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::<F>()?;
     let (load1, load2) = (loads[0], loads[1]);
     let res: Vec<F> = 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::<F>()?;
+    let (_, byte_len) = vector_lengths::<F>(imm_len)?;
+    let loads = core.reserve_load(r1_val, byte_len)?.execute_load::<F>()?;
     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::<F>()?;
+    let (_, byte_len) = vector_lengths::<F>(imm_len)?;
+    let loads = core.reserve_load(r1_val, byte_len)?.execute_load::<F>()?;
     let load1 = loads[0];
     let res: Vec<F> = 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::<F>()?;
+    let (_, byte_len) = vector_lengths::<F>(imm_len)?;
+    let loads = core.reserve_load(r1_val, byte_len)?.execute_load::<F>()?;
     let load1 = loads[0];
     let res: Vec<F> = 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::<F>()?;
+    let (_, byte_len) = vector_lengths::<F>(imm_len)?;
+    let loads = core.reserve_load(r1_val, byte_len)?.execute_load::<F>()?;
     let load1 = loads[0];
     let res: Vec<F> = 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::<F>()?;
+    let (_, byte_len) = vector_lengths::<F>(imm_len)?;
+    let loads = core.reserve_load(r1_val, byte_len)?.execute_load::<F>()?;
     let load1 = loads[0];
     ensure!(!load1.is_empty(), "vsoftmax does not support empty vectors");
     let max_val = load1

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::<F>())
+            .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::<F>()
             .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");

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::<f32>() 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) {

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::<F>() 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::<F>() 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::<F>() 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::<F>() 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::<F>() 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::<F>() 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::<F>() 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::<F>() 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::<F>() as i32)
+            .set_imm_len(8)
             .build(),
     );
     let core_instruction = vec![inst_builder.build().into()];

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<uint64_t>(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<u
     reportRow.logicalLocalAllocationCount = plan.logicalAllocationCount;
     reportRow.logicalLocalBytes = plan.logicalBytes;
   }
-  else if (*localArenaSize != plan.arenaSize
-           || reportRow.logicalLocalAllocationCount != plan.logicalAllocationCount
+  else if (*localArenaSize != plan.arenaSize || reportRow.logicalLocalAllocationCount != plan.logicalAllocationCount
            || reportRow.logicalLocalBytes != plan.logicalBytes)
     llvm_unreachable("inconsistent PIM local-memory plan across core-batch lanes");
   for (const CompiledLocalMemoryEntry& entry : plan.entries) {
@@ -358,8 +355,8 @@ llvm::FailureOr<int64_t> PimAcceleratorMemory::getIndexValue(mlir::Value value,
 PimAcceleratorMemory::PimAcceleratorMemory()
 : hostMem(memEntriesMap), fileReport(openMemoryReport(pimMemoryReport == PimMemoryReportSummary)) {}
 
-PimAcceleratorMemory::PimAcceleratorMemory(
-  const llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32>& initialMemEntries, bool enableReport)
+PimAcceleratorMemory::PimAcceleratorMemory(const llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32>& 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<int32_t, 2> 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<ShapedType>(mvmLikeOp.getInput().getType())),
+                       getVectorElementBitwidthOrCrash(cast<ShapedType>(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<ShapedType>(lhs.getType());
+  ensureVectorBitwidth(getVectorElementBitwidthOrCrash(inputType),
+                       getVectorElementBitwidthOrCrash(cast<ShapedType>(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<ShapedType>(input.getType());
+  ensureVectorBitwidth(getVectorElementBitwidthOrCrash(inputType),
+                       getVectorElementBitwidthOrCrash(cast<ShapedType>(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<Operation*> collectTopLevelCoreLikeOps(func::FuncOp funcOp) {
 static FailureOr<CompiledCoreMemoryPlan> compileCoreMemoryPlan(Operation* coreLikeOp) {
   CompiledCoreMemoryPlan plan;
   auto arenaAttr = coreLikeOp->getAttrOfType<IntegerAttr>(kLocalMemorySizeAttrName);
-  if (!arenaAttr || arenaAttr.getInt() < 0
-      || static_cast<uint64_t>(arenaAttr.getInt()) > kPimLocalMemoryAddressLimit) {
+  if (!arenaAttr || arenaAttr.getInt() < 0 || static_cast<uint64_t>(arenaAttr.getInt()) > kPimLocalMemoryAddressLimit) {
     coreLikeOp->emitError("requires a valid pim.local_memory_size attribute before codegen");
     return failure();
   }
@@ -841,7 +842,9 @@ static FailureOr<CompiledCoreMemoryPlan> compileCoreMemoryPlan(Operation* coreLi
       hasFailure = true;
       return;
     }
-    plan.entries.push_back({allocOp.getResult(), {address, static_cast<size_t>(*checkedSize)}});
+    plan.entries.push_back({
+      allocOp.getResult(), {address, static_cast<size_t>(*checkedSize)}
+    });
     auto logicalBytes = pim::checkedAdd(
       static_cast<size_t>(plan.logicalBytes), static_cast<size_t>(*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<ShapedType>(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<ShapedType>(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<size_t, size_t> emittedCoreIds;
-  size_t nextEmittedCoreId = 0;
-
-  for (Operation* op : coreLikeOps) {
-    if (auto coreOp = dyn_cast<pim::PimCoreOp>(op)) {
-      size_t originalCoreId = static_cast<size_t>(coreOp.getCoreId());
-      if (!emittedCoreIds.contains(originalCoreId))
-        emittedCoreIds[originalCoreId] = nextEmittedCoreId++;
-      continue;
-    }
-
-    auto coreBatchOp = cast<pim::PimCoreBatchOp>(op);
-    auto batchCoreIds = getBatchCoreIds(coreBatchOp);
-    for (unsigned lane = 0; lane < static_cast<unsigned>(coreBatchOp.getLaneCount()); ++lane) {
-      size_t originalCoreId = static_cast<size_t>(batchCoreIds[lane]);
-      if (!emittedCoreIds.contains(originalCoreId))
-        emittedCoreIds[originalCoreId] = nextEmittedCoreId++;
-    }
-  }
-
   SmallVector<CoreEmissionJob> jobs;
   SmallVector<SmallVector<size_t>> batchJobIndices;
   for (Operation* op : coreLikeOps) {
     if (auto coreOp = dyn_cast<pim::PimCoreOp>(op)) {
-      size_t originalCoreId = static_cast<size_t>(coreOp.getCoreId());
       CoreEmissionJob job;
       job.coreLikeOp = coreOp;
       job.program = getCompiledProgram(op);
       job.memoryPlan = getMemoryPlan(op);
-      job.emittedCoreId = emittedCoreIds.lookup(originalCoreId);
+      job.physicalCoreId = static_cast<size_t>(coreOp.getCoreId());
       jobs.push_back(std::move(job));
       continue;
     }
@@ -1220,16 +1191,15 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
       lanesByCoreId[static_cast<size_t>(batchCoreIds[lane])].push_back(lane);
 
     SmallVector<size_t> jobIndices;
-    SmallVector<size_t> 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<size_t> 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<std::string> weightFiles, json::Array& xbarsPerGroup) -> OnnxMlirCompilerErrorCodes {
+  auto linkCoreWeights = [&](size_t coreId,
+                             ArrayRef<std::string> weightFiles,
+                             ArrayRef<ResolvedWeightView> 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<int64_t>(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<raw_fd_ostream> 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<raw_fd_ostream>(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<pim::PimCoreOp>(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<pim::PimCoreBatchOp>(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<bool> 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<WeightFileRequest, 8> 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<pim::PimCoreOp>(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);

src/PIM/Compiler/PimCodeGen.hpp

@@ -134,8 +134,7 @@ private:
 
 public:
   PimAcceleratorMemory();
-  PimAcceleratorMemory(
-    const llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32>& initialMemEntries, bool enableReport);
+  PimAcceleratorMemory(const llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32>& initialMemEntries, bool enableReport);
 
   PimMemory& getOrCreateDeviceMem(size_t id);
 
@@ -145,8 +144,7 @@ public:
   llvm::FailureOr<int64_t> 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<int32_t> coreIds, const MemoryReportRow& perCoreRow);
+  void recordBatchReport(uint64_t batchId, llvm::ArrayRef<int32_t> 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<unsigned, 4> lanes;
   std::optional<uint64_t> batchReportId;
 };
@@ -164,17 +162,16 @@ class PimCodeGen {
   PimAcceleratorMemory& memory;
   PimInstructionWriter& instructionWriter;
   llvm::raw_fd_ostream* coreJsonStream;
-  const llvm::DenseMap<size_t, size_t>& emittedCoreIds;
   std::optional<unsigned> batchLane;
   mutable std::array<std::optional<int32_t>, 256> scalarRegisterValues = {};
+  mutable std::optional<std::array<int32_t, 2>> 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<size_t, size_t>& 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<unsigned> lane) { batchLane = lane; }
   llvm::FailureOr<int64_t> indexOf(mlir::Value value, const StaticValueKnowledge& knowledge) const {

src/PIM/Compiler/PimCompilerOptions.cpp

@@ -125,6 +125,12 @@ llvm::cl::opt<long> coresCount("core-count",
                                llvm::cl::desc("Number of cores in the chip. Required for PIM compilation."),
                                llvm::cl::init(-1));
 
+llvm::cl::opt<std::string> 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<bool>
   ignoreConcatError("ignore-concat-error",
                     llvm::cl::desc("Ignore ConcatOp corner case: do not assert and do a simplification"),

src/PIM/Compiler/PimCompilerOptions.hpp

@@ -2,6 +2,8 @@
 
 #include "llvm/Support/CommandLine.h"
 
+#include <string>
+
 #define INSTRUMENTSTAGE_ENUM_PIM
 
 #define INSTRUMENTSTAGE_CL_ENUM_PIM
@@ -63,6 +65,7 @@ extern llvm::cl::opt<bool> pimTraceCommunicationMaterialization;
 extern llvm::cl::opt<size_t> crossbarSize;
 extern llvm::cl::opt<size_t> crossbarCountInCore;
 extern llvm::cl::opt<long> coresCount;
+extern llvm::cl::opt<std::string> pimTargetConfig;
 extern llvm::cl::opt<uint64_t> pimConvIm2colMaxElements;
 extern llvm::cl::opt<uint64_t> pimConvStreamChunkPositions;
 

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 <cmath>
+#include <limits>
+#include <tuple>
+
 #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<size_t>(
+    std::sqrt(static_cast<long double>(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<Cost>(2 + rowDistance + columnDistance);
+      target.interProcessorLatencyNs[
+        source * target.processorCount + destination] = latency;
+      latencySum = checkedAdd(latencySum, latency);
+      ++pairCount;
+    }
+  }
+  target.averageInterProcessorLatencyNs =
+    pairCount == 0
+      ? 0
+      : (latencySum + static_cast<Cost>(pairCount) - 1)
+          / static_cast<Cost>(pairCount);
+}
+
+spatial::SchedulingTarget getDefaultPimSchedulingTarget() {
+  spatial::SchedulingTarget target;
+  target.processorCount = static_cast<size_t>(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<double> number = object.getNumber(key);
+  if (!number)
+    return fallback;
+  if (!std::isfinite(*number) || *number < 0.0 || (!allowZero && *number == 0.0)
+      || *number > static_cast<double>(std::numeric_limits<Cost>::max()))
+    llvm::report_fatal_error("PIM target config field '" + key + "' must be a valid positive number");
+  return static_cast<Cost>(std::ceil(*number));
+}
+
+std::pair<size_t, size_t> 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<int64_t> first = (*values)[0].getAsInteger();
+  std::optional<int64_t> 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<size_t>(*first), static_cast<size_t>(*second)};
+}
+
+void loadPimInterProcessorLatencies(
+    spatial::SchedulingTarget& target,
+    const llvm::json::Object& network) {
+  std::optional<llvm::StringRef> 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<double> 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<Cost>(std::ceil(*latency));
+      target.interProcessorLatencyNs[
+        source * target.processorCount + destination] = roundedLatency;
+      latencySum = checkedAdd(latencySum, roundedLatency);
+      ++pairCount;
+    }
+  }
+  target.averageInterProcessorLatencyNs =
+    pairCount == 0
+      ? 0
+      : (latencySum + static_cast<Cost>(pairCount) - 1)
+          / static_cast<Cost>(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<int64_t> 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<size_t>(*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<size_t>(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<ModuleOp>& module,
                   PassManager& pm,
                   EmissionTargetType& emissionTarget,
@@ -31,11 +291,13 @@ void addPassesPim(OwningOpRef<ModuleOp>& 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"));
   }
 

src/PIM/Compiler/PimWeightEmitter.cpp

@@ -42,7 +42,9 @@ WeightEmissionResult createAndPopulateWeightFolder(ArrayRef<WeightFileRequest> 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<int64_t>(crossbarCountInCore)
+           && "Weight dimensions must fit in one array group");
 
     size_t elementByteWidth = getElementTypeSizeInBytes(denseAttr.getElementType());
 
@@ -57,7 +59,7 @@ WeightEmissionResult createAndPopulateWeightFolder(ArrayRef<WeightFileRequest> 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<APFloat>()[elementIndex].bitcastToAPInt();
@@ -73,7 +75,7 @@ WeightEmissionResult createAndPopulateWeightFolder(ArrayRef<WeightFileRequest> r
     weightFileStream.close();
     materializedWeights.push_back({weightView, newFileName});
     uint64_t weightBytes = pim::checkedMulOrCrash(
-      pim::checkedMulOrCrash(static_cast<size_t>(xbarSize), static_cast<size_t>(xbarSize), "weight element count"),
+      pim::checkedMulOrCrash(static_cast<size_t>(xbarSize), static_cast<size_t>(numCols), "weight element count"),
       elementByteWidth,
       "weight byte size");
     result.totalWeightBytes = pim::checkedAddOrCrash(result.totalWeightBytes, weightBytes, "total weight bytes");

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

src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp

@@ -19,9 +19,11 @@ FailureOr<RowStripPhysicalValue> 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<Value> applyRowStripBiasAdd(const RowStripPhysicalValue& value,
   return batchOp->getResult(0);
 }
 
+FailureOr<Value> 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<RankedTensorType>(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<Value> lhsFragment =
+        extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[0], args.lane, lhs.fragmentType);
+      FailureOr<Value> 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<Value> applyRowStripConcat(ArrayRef<RowStripPhysicalValue> 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<Value> 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<Value> 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<Value> 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

src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp

@@ -66,4 +66,14 @@ mlir::FailureOr<mlir::Value> applyRowStripBiasAdd(const RowStripPhysicalValue& v
                                                   mlir::PatternRewriter& rewriter,
                                                   mlir::Location loc);
 
+mlir::FailureOr<mlir::Value> applyRowStripAdd(const RowStripPhysicalValue& lhs,
+                                              const RowStripPhysicalValue& rhs,
+                                              mlir::PatternRewriter& rewriter,
+                                              mlir::Location loc);
+
+mlir::FailureOr<mlir::Value> applyRowStripConcat(llvm::ArrayRef<RowStripPhysicalValue> inputs,
+                                                 mlir::RankedTensorType outputType,
+                                                 mlir::PatternRewriter& rewriter,
+                                                 mlir::Location loc);
+
 } // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp

@@ -64,6 +64,22 @@ static FailureOr<Value> lowerRowStripBiasAdd(const RowStripPhysicalValue& input,
   return applyRowStripBiasAdd(input, planOp.getBias(), rewriter, planOp.getLoc());
 }
 
+static FailureOr<Value> lowerRowStripAdd(const RowStripPhysicalValue& lhs,
+                                         const RowStripPhysicalValue& rhs,
+                                         spatial::SpatAddPlanOp planOp,
+                                         PatternRewriter& rewriter) {
+  return applyRowStripAdd(lhs, rhs, rewriter, planOp.getLoc());
+}
+
+static FailureOr<Value> lowerRowStripConcat(ArrayRef<RowStripPhysicalValue> inputs,
+                                            spatial::SpatConcatPlanOp planOp,
+                                            PatternRewriter& rewriter) {
+  auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
+  if (!outputType)
+    return failure();
+  return applyRowStripConcat(inputs, outputType, rewriter, planOp.getLoc());
+}
+
 static FailureOr<Value>
 materializeRowStripToDense(const RowStripPhysicalValue& rowStripValue, Location loc, PatternRewriter& rewriter) {
   if (rowStripValue.logicalType.getRank() != 4 || !rowStripValue.logicalType.hasStaticShape())
@@ -122,6 +138,99 @@ static FailureOr<Value> lowerDenseBatchBiasAdd(Value input, Value bias, RankedTe
   return batch->getResult(0);
 }
 
+static LogicalResult lowerAddPlan(spatial::SpatAddPlanOp planOp,
+                                  llvm::DenseMap<Value, RowStripPhysicalValue>& rowStripValues,
+                                  llvm::SmallPtrSetImpl<Operation*>& eraseAfterLowering,
+                                  PatternRewriter& rewriter) {
+  FailureOr<RowStripPhysicalValue> lhs = getRowStripValue(rowStripValues, planOp.getLhs());
+  FailureOr<RowStripPhysicalValue> rhs = getRowStripValue(rowStripValues, planOp.getRhs());
+  if (succeeded(lhs) && succeeded(rhs)) {
+    auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
+      auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(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<Value> lowered = lowerRowStripAdd(*lhs, *rhs, planOp, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected row-strip Spatial add plan");
+    auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
+    FailureOr<RowStripPhysicalValue> 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<Value, RowStripPhysicalValue>& rowStripValues,
+                                     llvm::SmallPtrSetImpl<Operation*>& eraseAfterLowering,
+                                     PatternRewriter& rewriter) {
+  SmallVector<RowStripPhysicalValue> inputs;
+  for (Value input : planOp.getInputs()) {
+    FailureOr<RowStripPhysicalValue> 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<spatial::SpatBlueprintOp>(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<Value> lowered = lowerRowStripConcat(inputs, planOp, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected row-strip Spatial concat plan");
+    auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
+    FailureOr<RowStripPhysicalValue> 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<LowerSpatialPlansPass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerSpatialPlansPass)
 
@@ -274,6 +383,40 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
         eraseAfterLowering.insert(blueprint);
         continue;
       }
+      if (auto planOp = dyn_cast<spatial::SpatGlobalAveragePoolPlanOp>(&op)) {
+        auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
+          auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(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<RowStripPhysicalValue> input = getRowStripValue(rowStripValues, planOp.getInput());
+        rewriter.setInsertionPoint(planOp);
+        std::optional<Value> physicalInput;
+        if (succeeded(input))
+          physicalInput = input->storage;
+        FailureOr<Value> 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<spatial::SpatBlueprintOp>(*outputBlueprint);
+        FailureOr<RowStripPhysicalValue> 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<spatial::SpatBiasAddPlanOp>(&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<LowerSpatialPlansPass, Operatio
         rewriter.replaceOp(planOp, computeOp.getResults());
         continue;
       }
+      if (auto planOp = dyn_cast<spatial::SpatAddPlanOp>(&op)) {
+        if (failed(lowerAddPlan(planOp, rowStripValues, eraseAfterLowering, rewriter))) {
+          signalPassFailure();
+          return;
+        }
+        continue;
+      }
+      if (auto planOp = dyn_cast<spatial::SpatConcatPlanOp>(&op)) {
+        if (failed(lowerConcatPlan(planOp, rowStripValues, eraseAfterLowering, rewriter))) {
+          signalPassFailure();
+          return;
+        }
+        continue;
+      }
       if (auto flattenOp = dyn_cast<spatial::SpatGraphCompute>(&op)) {
         if (flattenOp.getInputs().size() == 1) {
           FailureOr<RowStripPhysicalValue> input =
@@ -488,12 +645,15 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
           return;
         op->emitOpError("planning blueprint must not remain after LowerSpatialPlans");
         hasIllegalOps = true;
-      } else if (isa<spatial::SpatConv2DPlanOp,
-                     spatial::SpatBiasAddPlanOp,
-                     spatial::SpatReluPlanOp,
-                     spatial::SpatMaxPool2DPlanOp,
-                     spatial::SpatMaterializeLayoutOp>(op)
-          || op->getDialect()->getNamespace() == "onnx") {
+      }
+      else if (isa<spatial::SpatConv2DPlanOp,
+                   spatial::SpatBiasAddPlanOp,
+                   spatial::SpatAddPlanOp,
+                   spatial::SpatReluPlanOp,
+                   spatial::SpatMaxPool2DPlanOp,
+                   spatial::SpatGlobalAveragePoolPlanOp,
+                   spatial::SpatMaterializeLayoutOp>(op)
+               || op->getDialect()->getNamespace() == "onnx") {
         op->emitOpError("operation must not remain after LowerSpatialPlans");
         hasIllegalOps = true;
       }

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">;

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp

@@ -47,12 +47,17 @@ static void populateEmptyFunction(func::FuncOp funcOp) {
   SmallVector<spatial::SpatGraphComputeBatch> computeBatches(funcOp.getOps<spatial::SpatGraphComputeBatch>());
   SmallVector<spatial::SpatConv2DPlanOp> convPlans(funcOp.getOps<spatial::SpatConv2DPlanOp>());
   SmallVector<spatial::SpatBiasAddPlanOp> biasAddPlans(funcOp.getOps<spatial::SpatBiasAddPlanOp>());
+  SmallVector<spatial::SpatAddPlanOp> addPlans(funcOp.getOps<spatial::SpatAddPlanOp>());
+  SmallVector<spatial::SpatConcatPlanOp> concatPlans(funcOp.getOps<spatial::SpatConcatPlanOp>());
   SmallVector<spatial::SpatReluPlanOp> reluPlans(funcOp.getOps<spatial::SpatReluPlanOp>());
   SmallVector<spatial::SpatMaxPool2DPlanOp> maxPoolPlans(funcOp.getOps<spatial::SpatMaxPool2DPlanOp>());
+  SmallVector<spatial::SpatGlobalAveragePoolPlanOp> globalAveragePoolPlans(
+    funcOp.getOps<spatial::SpatGlobalAveragePoolPlanOp>());
   SmallVector<spatial::SpatBlueprintOp> blueprints(funcOp.getOps<spatial::SpatBlueprintOp>());
   SmallVector<spatial::SpatMaterializeLayoutOp> materializers(funcOp.getOps<spatial::SpatMaterializeLayoutOp>());
-  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;
   }
 

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;

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);

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);

src/PIM/Conversion/ONNXToSpatial/Patterns/GeneratedConversion.cpp

@@ -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<removeLRN>(ctx);
-}
-
-} // namespace onnx_mlir

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<int64_t>(crossbarSize.getValue());
   auto packedWeightType = RankedTensorType::get(
-    {tiling.numChannelTiles, tiling.tileInputRows, tiling.tileOutputChannels}, wType.getElementType());
+    {tiling.numChannelTiles, tiling.tileInputRows, paddedOutputChannels}, wType.getElementType());
   SmallVector<Attribute> packedValues(packedWeightType.getNumElements(),
                                       cast<Attribute>(rewriter.getZeroAttr(wType.getElementType())));
   SmallVector<Attribute> sourceValues(wDenseAttr.getValues<Attribute>());
@@ -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<OpFoldResult> offsets {channelTileIndex, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
+  const int64_t paddedOutputChannels = static_cast<int64_t>(crossbarSize.getValue());
   SmallVector<OpFoldResult> 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<int64_t>(crossbarSize.getValue())}, state.outType.getElementType());
   auto piecesType = spatial::getGraphBatchPhysicalResultType(
     tiling->totalPatches * tiling->numChannelTiles, rowTileType);
   auto paddedInputType = cast<RankedTensorType>(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<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
+        SmallVector<OpFoldResult> {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<int64_t>(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<Value> createConvOutputRow(ValueRange inputTiles,
                                             Location loc) {
   auto elementType = cast<RankedTensorType>(inputTiles.front().getType()).getElementType();
   auto rowType = RankedTensorType::get({1, outputChannels}, elementType);
-  auto tileWeightsType =
-    RankedTensorType::get({paddedK, xbarDim},
-                          cast<RankedTensorType>(paddedWeights.getType()).getElementType());
   const int64_t outputTileCount = ceilIntegerDivide(outputChannels, xbarDim);
-
-  auto getTileWeights = [&](int64_t outputTile) {
-    if (outputTileCount == 1)
-      return paddedWeights;
-    SmallVector<OpFoldResult> offsets {
-      rewriter.getIndexAttr(outputTile), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
-    SmallVector<OpFoldResult> 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<Value> 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<Value> tileResult = createConvOutputTile(
-      inputTiles, getTileWeights(outputTile), xbarDim, xbarDim, rewriter, loc);
-    if (failed(tileResult))
-      return failure();
-    SmallVector<OpFoldResult> tileOffsets {
-      rewriter.getIndexAttr(0), rewriter.getIndexAttr(outputTile * xbarDim)};
-    SmallVector<OpFoldResult> 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<RankedTensorType>(paddedWeights.getType()).getElementType());
+
+  Value paddedOutput;
+  for (auto [kSlice, inputTile] : llvm::enumerate(inputTiles)) {
+    const int64_t kOffset = static_cast<int64_t>(kSlice) * xbarDim;
+    Value weightSlice = extractStaticSliceOrIdentity(
+      rewriter,
+      loc,
+      paddedWeights,
+      weightSliceType,
+      SmallVector<OpFoldResult> {rewriter.getIndexAttr(kOffset), rewriter.getIndexAttr(0)},
+      SmallVector<OpFoldResult> {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<OpFoldResult> outputOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
-  SmallVector<OpFoldResult> 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<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
+                         SmallVector<OpFoldResult> {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<Value> 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<Value> 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<Value> 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<Value>& yielded) {
+          Value partialInputScratch = hasPartialInputTile ? iterArgs[1] : Value();
+          FailureOr<SmallVector<Value>> inputTiles = createConvInputTiles(*inputWindow,
+                                                                          state,
+                                                                          localColumn,
+                                                                          partialInputScratch,
+                                                                          patchSize,
+                                                                          numKSlices,
+                                                                          xbarDim,
+                                                                          rewriter,
+                                                                          pixelLoc);
+          if (failed(inputTiles))
+            return failure();
+          FailureOr<Value> 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<ReassociationIndices> {{0, 1, 2}, {3}});
+          Value next = tensor::InsertSliceOp::create(
+            rewriter,
+            pixelLoc,
+            outputPixel,
+            iterArgs.front(),
+            SmallVector<OpFoldResult> {rewriter.getIndexAttr(0),
+                                       rewriter.getIndexAttr(0),
+                                       localColumn,
+                                       rewriter.getIndexAttr(0)},
+            SmallVector<OpFoldResult> {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<Value> createOutputChannelTiledRowStripConvOutput(const ConvLoweringState& state,
                                                                    Value input,
                                                                    Value paddedWeights,
@@ -3229,21 +3319,11 @@ static FailureOr<Value>
   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<Value> bias = failure();
   if (state.hasBias)
@@ -3251,83 +3331,9 @@ static FailureOr<Value>
   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<Value> 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<Value> 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<Value>& widthYielded) {
-          Value partialInputScratch = hasPartialInputTile ? widthIterArgs[1] : Value();
-          FailureOr<SmallVector<Value>> inputTiles = createConvInputTiles(*inputWindow,
-                                                                          state,
-                                                                          widthIndex,
-                                                                          partialInputScratch,
-                                                                          patchSize,
-                                                                          numKSlices,
-                                                                          xbarDim,
-                                                                          rewriter,
-                                                                          widthLoc);
-          if (failed(inputTiles))
-            return failure();
-          FailureOr<Value> 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<ReassociationIndices> {{0, 1, 2}, {3}});
-          SmallVector<OpFoldResult> rowOffsets {
-            rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex, rewriter.getIndexAttr(0)};
-          SmallVector<OpFoldResult> 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<Value> createConvOutputFromPixelMajorRowStripFragments(Value rowStripStorage,
@@ -3346,105 +3352,22 @@ static FailureOr<Value> 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<Value> 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<Value> 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<Value> 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<Value>& widthYielded) {
-          Value partialInputScratch = hasPartialInputTile ? widthIterArgs[1] : Value();
-          FailureOr<SmallVector<Value>> inputTiles = createConvInputTiles(*inputWindow,
-                                                                          state,
-                                                                          widthIndex,
-                                                                          partialInputScratch,
-                                                                          patchSize,
-                                                                          numKSlices,
-                                                                          xbarDim,
-                                                                          rewriter,
-                                                                          widthLoc);
-          if (failed(inputTiles))
-            return failure();
-          FailureOr<Value> 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<ReassociationIndices> {{0, 1, 2}, {3}});
-          SmallVector<OpFoldResult> rowOffsets {
-            rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex, rewriter.getIndexAttr(0)};
-          SmallVector<OpFoldResult> 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<Value> createPointwiseOutputFromRowStripFragments(Value rowStripStorage,

src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Elementwise.cpp

@@ -193,6 +193,13 @@ struct AddToSpatialCompute : OpConversionPattern<ONNXAddOp> {
       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();

src/PIM/Conversion/ONNXToSpatial/Patterns/NN/Pool.cpp

@@ -246,6 +246,15 @@ struct PoolToSpatialComputeBase : public OpConversionPattern<PoolOp> {
         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<int64_t>(crossbarSize.getValue());
     const int64_t channelTileCount = (channels + xbarSize - 1) / xbarSize;
@@ -676,6 +685,132 @@ FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
   return batch->getResult(0);
 }
 
+LogicalResult canLowerGlobalAveragePoolPlanToRowStrip(spatial::SpatGlobalAveragePoolPlanOp planOp) {
+  auto inputType = dyn_cast<RankedTensorType>(planOp.getInput().getType());
+  auto outputType = dyn_cast<RankedTensorType>(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<Value> lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp,
+                                                    std::optional<Value> rowStripInput,
+                                                    PatternRewriter& rewriter) {
+  if (failed(canLowerGlobalAveragePoolPlanToRowStrip(planOp)))
+    return failure();
+
+  Location loc = planOp.getLoc();
+  auto inputType = cast<RankedTensorType>(planOp.getInput().getType());
+  auto outputType = cast<RankedTensorType>(planOp.getOutput().getType());
+  auto elementType = dyn_cast<FloatType>(inputType.getElementType());
+  if (!elementType)
+    return failure();
+
+  Value input = rowStripInput.value_or(planOp.getInput());
+  auto actualInputType = dyn_cast<RankedTensorType>(input.getType());
+  FailureOr<RowStripPhysicalValue> 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<double>(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<Value> 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<OpFoldResult> {rewriter.getIndexAttr(0),
+                                      rewriter.getIndexAttr(0),
+                                      rewriter.getIndexAttr(row),
+                                      rewriter.getIndexAttr(0)},
+            SmallVector<OpFoldResult> {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<OpFoldResult> {rewriter.getIndexAttr(0),
+                                      rewriter.getIndexAttr(0),
+                                      rewriter.getIndexAttr(column),
+                                      rewriter.getIndexAttr(0)},
+            SmallVector<OpFoldResult> {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<PoolToSpatialCompute<ONNXMaxPoolSingleOutOp>>(ctx);
   patterns.insert<PoolToSpatialCompute<ONNXAveragePoolOp>>(ctx);

src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Concat.cpp

@@ -25,6 +25,17 @@ struct Concat : public OpConversionPattern<ONNXConcatOp> {
       return success();
     }
 
+    auto resultType = dyn_cast<RankedTensorType>(maxpoolOp.getResult().getType());
+    if (axis == 1 && resultType && resultType.hasStaticShape() && resultType.getRank() == 4
+        && llvm::all_of(inputs, [](Value input) {
+             auto type = dyn_cast<RankedTensorType>(input.getType());
+             return type && type.hasStaticShape() && type.getRank() == 4;
+           })) {
+      rewriter.replaceOpWithNewOp<spatial::SpatConcatPlanOp>(
+        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(

src/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp

@@ -27,6 +27,14 @@ lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
                            std::optional<mlir::Value> rowStripInput,
                            mlir::PatternRewriter& rewriter);
 
+mlir::LogicalResult
+canLowerGlobalAveragePoolPlanToRowStrip(spatial::SpatGlobalAveragePoolPlanOp planOp);
+
+mlir::FailureOr<mlir::Value>
+lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp,
+                                   std::optional<mlir::Value> rowStripInput,
+                                   mlir::PatternRewriter& rewriter);
+
 mlir::LogicalResult canLowerFlattenFromRowStrip(spatial::SpatGraphCompute flattenOp);
 
 mlir::LogicalResult lowerFlattenFromRowStrip(const RowStripPhysicalValue& input,

src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp

@@ -36,10 +36,16 @@ static bool usesSelectedRowStrip(Operation* user, llvm::DenseMap<Value, Selected
     return getSelectedLayout(layouts, reluPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
   if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(user))
     return getSelectedLayout(layouts, biasAddPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
+  if (auto addPlan = dyn_cast<spatial::SpatAddPlanOp>(user))
+    return getSelectedLayout(layouts, addPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
+  if (auto concatPlan = dyn_cast<spatial::SpatConcatPlanOp>(user))
+    return getSelectedLayout(layouts, concatPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
   if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(user))
     return getSelectedLayout(layouts, convPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
   if (auto maxPoolPlan = dyn_cast<spatial::SpatMaxPool2DPlanOp>(user))
     return getSelectedLayout(layouts, maxPoolPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
+  if (auto averagePoolPlan = dyn_cast<spatial::SpatGlobalAveragePoolPlanOp>(user))
+    return getSelectedLayout(layouts, averagePoolPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
   if (auto flattenCompute = dyn_cast<spatial::SpatGraphCompute>(user))
     return succeeded(canLowerFlattenFromRowStrip(flattenCompute));
   return false;
@@ -62,10 +68,16 @@ static bool canConsumeRowStripAsUser(Operation* user) {
     auto resultType = dyn_cast<RankedTensorType>(biasAddPlan.getOutput().getType());
     return resultType && isSupportedBiasAddValue(biasAddPlan.getBias(), resultType);
   }
+  if (isa<spatial::SpatAddPlanOp>(user))
+    return true;
+  if (isa<spatial::SpatConcatPlanOp>(user))
+    return true;
   if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(user))
     return succeeded(canConsumeAndProduceRowStrip(convPlan));
   if (auto maxPoolPlan = dyn_cast<spatial::SpatMaxPool2DPlanOp>(user))
     return succeeded(canLowerMaxPoolPlanToRowStrip(maxPoolPlan));
+  if (auto averagePoolPlan = dyn_cast<spatial::SpatGlobalAveragePoolPlanOp>(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<Value, SelectedLayout>& 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<Value, SelectedLayout>& 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<RankedTensorType>(value.getType());
   auto [offsets, sizes] = buildRowStripMetadata(outputType);
@@ -215,6 +254,22 @@ struct SpatialLayoutPlanningPass final : PassWrapper<SpatialLayoutPlanningPass,
           }
           continue;
         }
+        if (auto addPlan = dyn_cast<spatial::SpatAddPlanOp>(&op)) {
+          SelectedLayout selected = chooseAddLayout(addPlan, layouts);
+          if (layouts[addPlan.getResult()] != selected) {
+            layouts[addPlan.getResult()] = selected;
+            changed = true;
+          }
+          continue;
+        }
+        if (auto concatPlan = dyn_cast<spatial::SpatConcatPlanOp>(&op)) {
+          SelectedLayout selected = chooseConcatLayout(concatPlan, layouts);
+          if (layouts[concatPlan.getResult()] != selected) {
+            layouts[concatPlan.getResult()] = selected;
+            changed = true;
+          }
+          continue;
+        }
         if (auto maxPoolPlan = dyn_cast<spatial::SpatMaxPool2DPlanOp>(&op)) {
           SelectedLayout selected = chooseMaxPoolLayout(maxPoolPlan);
           if (layouts[maxPoolPlan.getResult()] != selected) {
@@ -223,6 +278,14 @@ struct SpatialLayoutPlanningPass final : PassWrapper<SpatialLayoutPlanningPass,
           }
           continue;
         }
+        if (auto averagePoolPlan = dyn_cast<spatial::SpatGlobalAveragePoolPlanOp>(&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<SpatialLayoutPlanningPass,
         producedValue = convPlan.getResult();
       else if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(&op))
         producedValue = biasAddPlan.getResult();
+      else if (auto addPlan = dyn_cast<spatial::SpatAddPlanOp>(&op))
+        producedValue = addPlan.getResult();
+      else if (auto concatPlan = dyn_cast<spatial::SpatConcatPlanOp>(&op))
+        producedValue = concatPlan.getResult();
       else if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(&op))
         producedValue = reluPlan.getResult();
       else if (auto maxPoolPlan = dyn_cast<spatial::SpatMaxPool2DPlanOp>(&op))
         producedValue = maxPoolPlan.getResult();
+      else if (auto averagePoolPlan = dyn_cast<spatial::SpatGlobalAveragePoolPlanOp>(&op))
+        producedValue = averagePoolPlan.getResult();
       else
         continue;
 

src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp

@@ -264,7 +264,13 @@ LogicalResult raptor::SpatialToPimPass::enlargeVMMOutTensorsToCrossbarSize(func:
     auto outputType = cast<RankedTensorType>(vmmOp.getOutput().getType());
     ArrayRef<int64_t> outputShape = outputType.getShape();
     assert(isHVectorShape(outputShape) && "expected a horizontal vector output");
-    assert(outputShape[1] <= static_cast<int64_t>(crossbarSize) && "output width must fit in one crossbar");
+    auto weightType = cast<RankedTensorType>(vmmOp.getWeight().getType());
+    const int64_t xbarDim = static_cast<int64_t>(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<int64_t>(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<int64_t>(crossbarSize)}, outputType.getElementType(), outputType.getEncoding());
-    Value paddedOutputBuffer = outputShape[1] == static_cast<int64_t>(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<int64_t>(crossbarSize))
+    if (outputShape[1] == paddedOutputWidth)
       return WalkResult::advance();
 
     SmallVector<OpFoldResult> offsets = {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};

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<int64_t>(crossbarSize) || M > static_cast<int64_t>(crossbarSize))
-    return emitError("matrix dimensions must fit in one crossbar");
+  const int64_t xbarDim = static_cast<int64_t>(crossbarSize);
+  if (N > xbarDim || M > xbarDim * static_cast<int64_t>(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<int64_t>(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<int64_t>(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();
 }

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<memref::CopyOp> copies;
+  funcOp.walk([&](memref::CopyOp copy) { copies.push_back(copy); });
+
+  for (memref::CopyOp copy : copies) {
+    Value copiedSource = copy.getSource();
+    auto expand = copiedSource.getDefiningOp<memref::ExpandShapeOp>();
+    Value producerOutput = expand ? expand.getSrc() : copiedSource;
+    auto result = dyn_cast<OpResult>(producerOutput);
+    Operation* producer = result ? result.getOwner() : nullptr;
+    auto dps = dyn_cast_or_null<DestinationStyleOpInterface>(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<memref::AllocOp>();
+    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<ShapedType>(destination.getType()).getShape()
+          != cast<ShapedType>(producerOutput.getType()).getShape()
+        || cast<ShapedType>(destination.getType()).getElementType()
+          != cast<ShapedType>(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);

src/PIM/Dialect/Pim/Transforms/HostConstantFolding/CMakeLists.txt

@@ -9,4 +9,5 @@ add_pim_library(OMPimHostConstantFolding
   LINK_LIBS PUBLIC
   MLIRLinalgDialect
   OMPimCommon
+  OMPimBufferization
 )

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<HostConstantFoldingPass, OperationP
 
     populateConstantFoldingConstantPatterns(owningPatterns);
     populateConstantFoldingSubviewPatterns(owningPatterns);
+    pim::populatePimContiguityNormalizationPatterns(owningPatterns);
 
     patterns = std::make_shared<FrozenRewritePatternSet>(std::move(owningPatterns));
     return success();

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<SpatTensor>:$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";
 

src/PIM/Dialect/Spatial/SpatialOpsAsm.cpp

@@ -149,7 +149,8 @@ void printComputeLikeOp(ComputeOpTy op, OpAsmPrinter& printer) {
 
   if (auto coreIdAttr = op->template getAttrOfType<IntegerAttr>(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<uint32_t>(op.getLaneCount())})
+               .size();
   if (auto coreIdsAttr = op->template getAttrOfType<DenseI32ArrayAttr>(onnx_mlir::kCoreIdsAttrName)) {
     printer << " coreIds ";
     printCompressedIntegerList(printer, coreIdsAttr.asArrayRef());

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<ShapedType>(getOutput().getType());
+  auto outputType = dyn_cast<ShapedType>(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<ShapedType>(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<RankedTensorType>(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<RankedTensorType>(getInput().getType());
+  auto outputType = dyn_cast<RankedTensorType>(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<RankedTensorType>(getLhs().getType());
+  auto rhsType = dyn_cast<RankedTensorType>(getRhs().getType());
+  auto outputType = dyn_cast<RankedTensorType>(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";

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<SpatBlueprintOp>();
   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<ProducerValueKey, MaterializedProducerRef> &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<SpatBlueprintOp>();
     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<SpatComputeBatch>(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<SpatBlueprintOp>();
     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";

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,

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<Cost> logicalTrafficFlits(target.processorCount * target.processorCount, 0);
+  for (const std::unique_ptr<DeferredExchangePlan>& exchange : plan.exchanges)
+    for (const ExternalTransferFamily& transfer : exchange->external) {
+      auto fragmentType = dyn_cast<ShapedType>(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<Cost>(*fragmentBytes) / target.transferWidthBytes
+                 + (*fragmentBytes % target.transferWidthBytes != 0);
+      for (size_t index = 0; index < transfer.sourceCores.size(); ++index) {
+        size_t sourceLogicalProcessor = static_cast<size_t>(transfer.sourceCores.valueAt(index));
+        size_t targetLogicalProcessor = static_cast<size_t>(transfer.targetCores.valueAt(index));
+        Cost& traffic = logicalTrafficFlits[sourceLogicalProcessor * target.processorCount + targetLogicalProcessor];
+        traffic = checkedAdd(traffic, flits);
+      }
+    }
+
+  std::vector<size_t> physicalCoreForLogicalProcessor =
+    mapLogicalProcessorsToPhysicalCores(logicalTrafficFlits, target);
+  auto getPhysicalCore = [&](int64_t logicalProcessor) {
+    assert(logicalProcessor >= 0 && static_cast<size_t>(logicalProcessor) < physicalCoreForLogicalProcessor.size()
+           && "logical processor is outside the scheduling target");
+    return static_cast<int64_t>(physicalCoreForLogicalProcessor[logicalProcessor]);
+  };
+  auto remap = [&](StaticIntSequence& logicalProcessors) {
+    SmallVector<int64_t> 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<SpatScheduledCompute>(scheduled.op)) {
+      scheduled.op->setAttr(
+        kCoreIdAttrName, IntegerAttr::get(IntegerType::get(scheduled.op->getContext(), 32), scheduled.cores.front()));
+    }
+    else {
+      SmallVector<int32_t> physicalCores;
+      physicalCores.reserve(scheduled.cores.size());
+      for (int64_t physicalCore : scheduled.cores)
+        physicalCores.push_back(static_cast<int32_t>(physicalCore));
+      scheduled.op->setAttr(kCoreIdsAttrName, DenseI32ArrayAttr::get(scheduled.op->getContext(), physicalCores));
+    }
+  }
+  for (const std::unique_ptr<ProducedValue>& produced : plan.producedStorage)
+    produced->core = getPhysicalCore(produced->core);
+  for (const std::unique_ptr<DeferredExchangePlan>& 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<SpatGraphCompute, SpatGraphComputeBatch>(op))
       continue;
     auto graphId = op.getAttrOfType<IntegerAttr>("scheduled.graph_id");
     if (!graphId)
       continue;
     for (auto [resultIndex, result] : llvm::enumerate(op.getResults())) {
-      SmallVector<OpOperand *> externalUses;
-      for (OpOperand &use : result.getUses()) {
-        Operation *user = use.getOwner();
-        if (isa<SpatGraphCompute, SpatGraphComputeBatch,
-                SpatDeferredCommunicationOp>(user))
+      SmallVector<OpOperand*> externalUses;
+      for (OpOperand& use : result.getUses()) {
+        Operation* user = use.getOwner();
+        if (isa<SpatGraphCompute, SpatGraphComputeBatch, SpatDeferredCommunicationOp>(user))
           continue;
         if (auto blueprint = dyn_cast<SpatBlueprintOp>(user)) {
-          bool blueprintEscapes = llvm::any_of(
-            blueprint.getOutput().getUses(), [](OpOperand &blueprintUse) {
-              return !isa<SpatGraphCompute, SpatGraphComputeBatch,
-                          SpatDeferredCommunicationOp>(
-                blueprintUse.getOwner());
-            });
+          bool blueprintEscapes = llvm::any_of(blueprint.getOutput().getUses(), [](OpOperand& blueprintUse) {
+            return !isa<SpatGraphCompute, SpatGraphComputeBatch, SpatDeferredCommunicationOp>(blueprintUse.getOwner());
+          });
           if (!blueprintEscapes)
             continue;
         }
@@ -42,21 +100,16 @@ static LogicalResult replaceFinalGraphPublications(
       if (externalUses.empty())
         continue;
       SmallVector<Value> 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<Operation *> oldGraph;
-  for (Operation &op : funcOp.getOps())
+static LogicalResult eraseOldGraph(func::FuncOp funcOp, IRRewriter& rewriter) {
+  SmallVector<Operation*> oldGraph;
+  for (Operation& op : funcOp.getOps())
     if (isa<SpatGraphCompute, SpatGraphComputeBatch, SpatBlueprintOp>(op))
       oldGraph.push_back(&op);
-  for (Operation *op : llvm::reverse(oldGraph)) {
+  for (Operation* op : llvm::reverse(oldGraph)) {
     if (auto blueprint = dyn_cast<SpatBlueprintOp>(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<SpatDeferredSourceSelectOp> 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<SpatDeferredCommunicationOp> 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<SpatYieldOp>(body.getTerminator());
-    auto argument = yield && yield.getOutputs().size() == 1
-                      ? dyn_cast<BlockArgument>(yield.getOutputs().front())
-                      : BlockArgument();
-    if (argument && argument.getOwner() == &body
-        && argument.getArgNumber() < deferred.getSources().size())
+    auto argument =
+      yield && yield.getOutputs().size() == 1 ? dyn_cast<BlockArgument>(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);

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

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<MergeComputeNodesPass, OperationPass<ModuleOp>> {
   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<MergeComputeNodesPass, Operatio
 
   void runOnOperation() override {
     ModuleOp moduleOp = getOperation();
+    if (!hasTarget || target.processorCount == 0 || target.residentWeightCapacity == 0 || target.transferWidthBytes == 0
+        || target.interProcessorLatencyNs.size() != target.processorCount * target.processorCount
+        || (target.processorCount > 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<MergeComputeNodesPass, Operatio
     }
 
     func::FuncOp funcOp = *entryFunc;
-    MergeScheduleResult schedule = MergeSchedulingAnalysis(funcOp).getResult();
+    MergeScheduleResult logicalSchedule = MergeSchedulingAnalysis(funcOp, target).getResult();
     PatternRewriter rewriter(moduleOp.getContext());
     FailureOr<ScheduledComputeMaterializationResult> materialization =
-      materializeScheduledCompute(funcOp, schedule, rewriter);
+      materializeScheduledCompute(funcOp, logicalSchedule, rewriter);
     if (failed(materialization)) {
       signalPassFailure();
       return;
@@ -48,7 +58,7 @@ struct MergeComputeNodesPass final : PassWrapper<MergeComputeNodesPass, Operatio
     // payloads must be diagnosed from the producer-owned body.
     dumpModule(moduleOp, "spatial3_scheduled_no_comm", /*assumeVerified=*/true);
     if (failed(verifyMaterializedScheduleMapping(funcOp,
-                                                 schedule,
+                                                 logicalSchedule,
                                                  materialization->peftClassPlans,
                                                  materialization->graphComputeToBlockMap,
                                                  materialization->materializedSchedules))) {
@@ -75,18 +85,14 @@ struct MergeComputeNodesPass final : PassWrapper<MergeComputeNodesPass, Operatio
     SpatialDataflowExportStage exportMode = getSpatialDataflowExportStage();
     if (shouldExportSpatialDataflowStage(exportMode, SpatialDataflowExportStage::Spatial3)
         && failed(exportSpatialDataflowCsvScheduled(
-          funcOp, materialization->materializedSchedules,
-          "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 : PassWrapper<MergeComputeNodesPass, Operatio
     }
     if (shouldExportSpatialDataflowStage(exportMode, SpatialDataflowExportStage::Spatial4)
         && failed(exportSpatialDataflowCsvScheduled(
-          funcOp, materialization->materializedSchedules,
-          "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<MergeComputeNodesPass, Operatio
 
 std::unique_ptr<Pass> createMergeComputeNodesPass() { return std::make_unique<spatial::MergeComputeNodesPass>(); }
 
+std::unique_ptr<Pass> createMergeComputeNodesPass(const spatial::SchedulingTarget& target) {
+  return std::make_unique<spatial::MergeComputeNodesPass>(target);
+}
+
 } // namespace onnx_mlir

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputePlan.hpp

@@ -164,7 +164,8 @@ inline size_t getScheduledCpuForComputeInstance(const ComputeInstance &instance,
   auto batch = dyn_cast<SpatComputeBatch>(instance.op);
   assert(batch && instance.laneCount != 0 && "missing scheduled CPU for non-batch compute instance");
   assert(instance.laneStart < static_cast<uint32_t>(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<GraphComputeBlockKey> collectExpectedGraphComputeBlockKeys(func::FuncOp funcOp) {
+inline SmallVector<GraphComputeBlockKey> collectExpectedGraphComputeBlockKeys(
+    func::FuncOp funcOp, size_t processorCount) {
   SmallVector<GraphComputeBlockKey> keys;
   for (Operation &op : funcOp.getOps()) {
     if (auto compute = dyn_cast<SpatGraphCompute>(&op))
       keys.push_back(getGraphComputeBlockKey({compute.getOperation(), 0, 1}));
     else if (auto batch = dyn_cast<SpatGraphComputeBatch>(&op))
-      for (ComputeInstance chunk : getBatchChunksForRange(batch, 0, static_cast<uint32_t>(batch.getLaneCount())))
+      for (ComputeInstance chunk :
+           getBatchChunksForRange(
+             batch, 0, static_cast<uint32_t>(batch.getLaneCount()),
+             processorCount))
         keys.push_back(getGraphComputeBlockKey(chunk));
   }
   return keys;

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);
     }
   }

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();
     });
   }

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp

@@ -12,7 +12,6 @@
 #include "llvm/Support/Casting.h"
 
 #include <algorithm>
-#include <cmath>
 #include <iterator>
 #include <limits>
 #include <optional>
@@ -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<Cost>(shaped.getNumElements());
   }
 
-  static Cost getBitwidthOrDefault(Type type) {
+  Cost getBitwidthOrDefault(Type type) const {
     if (auto shaped = dyn_cast<ShapedType>(type))
       type = shaped.getElementType();
     if (auto intType = dyn_cast<IntegerType>(type))
@@ -81,14 +76,14 @@ struct PimsimSchedulerCostModel {
       return floatType.getWidth();
     if (isa<IndexType>(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<Cost>(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<Cost, Cost> estimateMeshShape() {
-    Cost coreCount = static_cast<Cost>(std::max<long>(1, coresCount.getValue()));
-    Cost rows = static_cast<Cost>(std::sqrt(static_cast<long double>(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<Cost>(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<Cost>(3) * rows);
-    Cost avgCol = averageAxisDistance(cols) / (static_cast<Cost>(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<Cost>(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<Cost>(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<Cost>(crossbarSize.getValue()),
+    Cost inputBytes = checkedMultiply(target.matrixRows,
                                       ceilDiv(inputBitwidth, static_cast<Cost>(8)));
     inputBytes = checkedMultiply(inputBytes, static_cast<Cost>(8));
     Cost outputBytes = getByteSize(outputType, getComputeBitwidth(outputType));
@@ -196,22 +182,22 @@ struct PimsimSchedulerCostModel {
 };
 
 std::optional<uint64_t> 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<scf::ForOp>(&op)) {
     std::optional<uint64_t> tripCount = getStaticTripCount(loop);
     if (!tripCount)
-      return PimsimSchedulerCostModel::kFallbackOperationCost;
-    return checkedMultiply(getRegionCost(loop.getRegion()), static_cast<Cost>(*tripCount));
+      return SchedulerCostModel::kFallbackOperationCost;
+    return checkedMultiply(getRegionCost(loop.getRegion(), costModel), static_cast<Cost>(*tripCount));
   }
 
   if (isa<SpatYieldOp, SpatInParallelOp, affine::AffineApplyOp, arith::ConstantOp,
@@ -219,42 +205,43 @@ Cost getOperationCost(Operation& op) {
     return 0;
 
   if (auto wvmm = dyn_cast<SpatVMMOp>(&op))
-    return PimsimSchedulerCostModel::getWvmmCost(wvmm.getInput().getType(), wvmm.getOutput().getType());
+    return costModel.getWvmmCost(wvmm.getInput().getType(), wvmm.getOutput().getType());
   if (auto vvdmul = dyn_cast<SpatVVDMulOp>(&op))
-    return PimsimSchedulerCostModel::getBinaryVectorCost(
+    return costModel.getBinaryVectorCost(
       vvdmul.getLhs().getType(), vvdmul.getRhs().getType(), vvdmul.getOutput().getType(), /*scalarOutput=*/true);
   if (auto vadd = dyn_cast<SpatVAddOp>(&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<SpatVSubOp>(&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<SpatVMulOp>(&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<SpatVMaxOp>(&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<SpatVAvgOp>(&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<SpatReluOp>(&op))
-    return PimsimSchedulerCostModel::getUnaryVectorCost(relu.getInput().getType(), relu.getOutput().getType());
+    return costModel.getUnaryVectorCost(relu.getInput().getType(), relu.getOutput().getType());
   if (auto sigm = dyn_cast<SpatSigmoidOp>(&op))
-    return PimsimSchedulerCostModel::getUnaryVectorCost(sigm.getInput().getType(), sigm.getOutput().getType());
+    return costModel.getUnaryVectorCost(sigm.getInput().getType(), sigm.getOutput().getType());
   if (auto softmax = dyn_cast<SpatSoftmaxOp>(&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<Cost>(4));
   }
   if (auto extract = dyn_cast<tensor::ExtractSliceOp>(&op))
-    return PimsimSchedulerCostModel::getTensorMoveCost(extract.getResult().getType());
+    return costModel.getTensorMoveCost(extract.getResult().getType());
   if (auto insert = dyn_cast<tensor::InsertSliceOp>(&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<uint64_t> getStaticTripCount(scf::ForOp loop) {
@@ -271,8 +258,8 @@ std::optional<uint64_t> 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<uint32_t> getConstantExtractLane(tensor::ExtractSliceOp extract) {
   return std::nullopt;
 }
 
-std::optional<Cost> getBatchProjectedInputTransferCost(SpatComputeBatch batch, Value input) {
+std::optional<TransferCost> 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<Cost> getBatchProjectedInputTransferCost(SpatComputeBatch batch, V
   if (!inputArg || !laneArg)
     return std::nullopt;
 
-  Cost projectedCost = 0;
+  TransferCost projectedCost;
   for (Operation* user : inputArg->getUsers()) {
     auto extract = dyn_cast<tensor::ExtractSliceOp>(user);
     if (!extract || extract.getSource() != *inputArg)
@@ -370,11 +359,13 @@ std::optional<Cost> getBatchProjectedInputTransferCost(SpatComputeBatch batch, V
     auto resultType = dyn_cast<ShapedType>(extract.getResult().getType());
     if (!resultType || !resultType.hasStaticShape())
       return std::nullopt;
-    projectedCost = checkedAdd(
-      projectedCost, PimsimSchedulerCostModel::getInterCoreTransferCostFromBytes(static_cast<Cost>(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<Cost> getBatchProjectedInputTransferCost(SpatComputeBatch batch, V
 static std::optional<SmallVector<ProducerValueRef, 4>>
 collectProjectedProducerValueRefs(SpatComputeBatch producer,
                                   Value input,
-                                  const ComputeInstance& consumerInstance) {
+                                  const ComputeInstance& consumerInstance,
+                                  size_t processorCount) {
   auto consumer = dyn_cast<SpatComputeBatch>(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<uint32_t>(producerLane));
+        ComputeInstance instance = getBatchChunkForLane(
+          producer, static_cast<uint32_t>(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<ShapedType>(input.getType());
   if (auto batch = dyn_cast<SpatComputeBatch>(consumerInstance.op))
-    if (std::optional<Cost> projectedCost = getBatchProjectedInputTransferCost(batch, input))
+    if (std::optional<TransferCost> projectedCost =
+          getBatchProjectedInputTransferCost(batch, input, costModel))
       return *projectedCost;
-  return PimsimSchedulerCostModel::getInterCoreTransferCostFromBytes(static_cast<Cost>(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<Cost>(fragmentLaneCount)) / static_cast<Cost>(totalLaneCount);
+  return scaleTransferCost(totalCost, static_cast<Cost>(fragmentLaneCount),
+                           static_cast<Cost>(totalLaneCount));
 }
 
-SmallVector<ProducerValueRef, 4> collectProducerValueRefs(Value value, const ComputeInstance& consumerInstance) {
+SmallVector<ProducerValueRef, 4> collectProducerValueRefs(
+    Value value, const ComputeInstance& consumerInstance,
+    size_t processorCount) {
   SmallVector<ProducerValueRef, 4> producers;
   Operation* op = value.getDefiningOp();
   if (!op)
@@ -461,13 +463,16 @@ SmallVector<ProducerValueRef, 4> collectProducerValueRefs(Value value, const Com
     auto batch = dyn_cast_or_null<SpatComputeBatch>(source.getDefiningOp());
     if (batch && batch.getNumResults() != 0) {
       if (std::optional<uint32_t> 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<uint32_t>(batch.getLaneCount())))
+           getBatchChunksForRange(batch, 0,
+                                  static_cast<uint32_t>(batch.getLaneCount()),
+                                  processorCount))
         producers.push_back({instance, 0});
       return producers;
     }
@@ -488,19 +493,24 @@ SmallVector<ProducerValueRef, 4> collectProducerValueRefs(Value value, const Com
 
   if (auto batch = dyn_cast<SpatComputeBatch>(op)) {
     if (batch.getNumResults() != 0) {
-      if (auto projected = collectProjectedProducerValueRefs(batch, value, consumerInstance))
+      if (auto projected = collectProjectedProducerValueRefs(
+            batch, value, consumerInstance, processorCount))
         return *projected;
-      std::optional<ProducerValueRef> producer = getProducerValueRef(value, &consumerInstance);
+      std::optional<ProducerValueRef> 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<OpResult>(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<ProducerValueRef, 4> 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<SpatComputeBatch>(producerRef.instance.op);
   if (!producerBatch || producerBatch.getNumResults() == 0)
     return transferCost;
 
   if (auto consumerBatch = dyn_cast<SpatComputeBatch>(consumerInstance.op)) {
-    if (std::optional<Cost> projectedCost = getBatchProjectedInputTransferCost(consumerBatch, input)) {
+    if (std::optional<TransferCost> projectedCost =
+          getBatchProjectedInputTransferCost(consumerBatch, input, costModel)) {
       uint32_t overlapLaneCount = getLaneOverlapCount(consumerInstance, producerRef.instance);
-      return checkedMultiply(*projectedCost, static_cast<Cost>(std::max<uint32_t>(1, overlapLaneCount)));
+      return scaleTransferCost(
+        *projectedCost,
+        static_cast<Cost>(std::max<uint32_t>(1, overlapLaneCount)));
     }
   }
 
@@ -527,8 +542,8 @@ Cost getProducerTransferCost(Value input,
     transferCost, static_cast<uint32_t>(producerBatch.getLaneCount()), producerRef.instance.laneCount);
 }
 
-static CrossbarWeight getOpaqueCrossbarWeight(Value value, std::optional<uint32_t> lane) {
-  CrossbarWeight weight;
+static ResidentWeight getOpaqueResidentWeight(Value value, std::optional<uint32_t> lane) {
+  ResidentWeight weight;
   weight.opaqueValue = value;
   weight.opaqueLane = lane.value_or(std::numeric_limits<uint32_t>::max());
   return weight;
@@ -619,7 +634,7 @@ static FailureOr<SmallVector<int64_t, 4>> evaluateIndexList(ArrayRef<OpFoldResul
   return result;
 }
 
-static Value resolveCrossbarWeightRoot(Operation* owner, Value root) {
+static Value resolveResidentWeightRoot(Operation* owner, Value root) {
   if (auto arg = dyn_cast<BlockArgument>(root)) {
     if (auto compute = dyn_cast<SpatCompute>(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<int64_t, 4> offsets,
                                              SmallVector<int64_t, 4> sizes,
                                              SmallVector<int64_t, 4> strides) {
-  CrossbarWeight weight;
+  ResidentWeight weight;
   weight.root = root;
   if (auto constant = root.getDefiningOp<arith::ConstantOp>())
     weight.rootAttr = static_cast<Attribute>(constant.getValue());
@@ -651,14 +666,14 @@ static CrossbarWeight completeCrossbarWeight(Value root,
   return weight;
 }
 
-static FailureOr<CrossbarWeight> getStaticCrossbarWeight(Operation* owner,
+static FailureOr<ResidentWeight> getStaticResidentWeight(Operation* owner,
                                                          Value value,
                                                          const DenseMap<Value, int64_t>& bindings,
                                                          std::optional<uint32_t> lane,
                                                          Value laneArg) {
   if (auto extract = value.getDefiningOp<tensor::ExtractSliceOp>()) {
-    FailureOr<CrossbarWeight> sourceWeight =
-      getStaticCrossbarWeight(owner, extract.getSource(), bindings, lane, laneArg);
+    FailureOr<ResidentWeight> 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<CrossbarWeight> getStaticCrossbarWeight(Operation* owner,
     return *sourceWeight;
   }
 
-  Value root = resolveCrossbarWeightRoot(owner, value);
+  Value root = resolveResidentWeightRoot(owner, value);
   auto type = dyn_cast<ShapedType>(root.getType());
   if (!type || !type.hasStaticShape())
     return failure();
@@ -686,18 +701,18 @@ static FailureOr<CrossbarWeight> getStaticCrossbarWeight(Operation* owner,
   SmallVector<int64_t, 4> offsets(type.getRank(), 0);
   SmallVector<int64_t, 4> sizes(type.getShape().begin(), type.getShape().end());
   SmallVector<int64_t, 4> 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<Value, int64_t>& bindings,
-                                         CrossbarUsage& usage,
+                                         ResidentWeightSet& usage,
                                          Value laneArg,
                                          std::optional<uint32_t> lane) {
   if (auto loop = dyn_cast<scf::ForOp>(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<SpatVMMOp>(op)) {
-    FailureOr<CrossbarWeight> weight = getStaticCrossbarWeight(owner, vmm.getWeight(), bindings, lane, laneArg);
+    FailureOr<ResidentWeight> 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<ComputeGraphEdge> aggregateEdges(llvm::ArrayRef<ComputeGraphEdge> edges) {
-  llvm::DenseMap<std::pair<size_t, size_t>, Cost> edgeCosts;
+  llvm::DenseMap<std::pair<size_t, size_t>, 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<ComputeGraphEdge> aggregatedEdges;
@@ -770,8 +786,8 @@ uint64_t countComputeBodyOperationInstances(Region& body) {
   return instances;
 }
 
-CrossbarUsage collectDistinctCrossbarWeights(Operation* owner, std::optional<uint32_t> lane) {
-  CrossbarUsage usage;
+ResidentWeightSet collectDistinctResidentWeights(Operation* owner, std::optional<uint32_t> lane) {
+  ResidentWeightSet usage;
   DenseMap<Value, int64_t> bindings;
   Value laneArg;
   if (auto batch = dyn_cast<SpatComputeBatch>(owner))
@@ -781,54 +797,87 @@ CrossbarUsage collectDistinctCrossbarWeights(Operation* owner, std::optional<uin
   for (Region& region : owner->getRegions())
     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<SpatCompute>(instance.op))
-    return getComputeBodyCost(spatCompute.getBody());
+    return getComputeBodyCost(spatCompute.getBody(), costModel);
   auto batch = cast<SpatComputeBatch>(instance.op);
-  return checkedMultiply(getComputeBodyCost(batch.getBody()), static_cast<Cost>(instance.laneCount));
+  return checkedMultiply(
+    getComputeBodyCost(batch.getBody(), costModel), static_cast<Cost>(instance.laneCount));
 }
 
-bool containsCrossbarWeight(ArrayRef<CrossbarWeight> usage, const CrossbarWeight& weight) {
+bool containsResidentWeight(ArrayRef<ResidentWeight> usage, const ResidentWeight& weight) {
   return llvm::is_contained(usage, weight);
 }
 
-unsigned countCrossbarOverlap(ArrayRef<CrossbarWeight> lhs, ArrayRef<CrossbarWeight> rhs) {
+unsigned countResidentWeightOverlap(ArrayRef<ResidentWeight> lhs, ArrayRef<ResidentWeight> 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<CrossbarWeight> lhs, ArrayRef<CrossbarWeight> rhs) {
+size_t getResidentWeightUnionSize(ArrayRef<ResidentWeight> lhs, ArrayRef<ResidentWeight> 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<CrossbarWeight> weights) {
-  for (const CrossbarWeight& weight : weights)
-    addCrossbarWeight(usage, weight);
+void insertResidentWeights(ResidentWeightSet& usage, ArrayRef<ResidentWeight> 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<SpatCompute>(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<int64_t>(target.matrixRows)) {
+      for (int64_t column = 0; column < weight.sizes[columnDim];
+           column += static_cast<int64_t>(target.matrixColumns)) {
+        ResidentWeight tile = weight;
+        tile.offsets[rowDim] += row * tile.strides[rowDim];
+        tile.offsets[columnDim] += column * tile.strides[columnDim];
+        tile.sizes[rowDim] =
+          std::min<int64_t>(target.matrixRows, weight.sizes[rowDim] - row);
+        tile.sizes[columnDim] =
+          std::min<int64_t>(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<SpatComputeBatch>(&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<Value, 4> 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)});
       }
     }
   }

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.hpp

@@ -11,9 +11,10 @@
 #include <vector>
 
 #include "ComputeInstance.hpp"
+#include "SchedulingTarget.hpp"
 #include "Utils.hpp"
 
-struct CrossbarWeight {
+struct ResidentWeight {
   mlir::Value root;
   mlir::Attribute rootAttr;
   llvm::SmallVector<int64_t, 4> 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<CrossbarWeight, 6>;
+using ResidentWeightSet = llvm::SmallVector<ResidentWeight, 6>;
 
 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<ComputeGraphNode> nodes;
   std::vector<ComputeGraphEdge> edges;
-  std::vector<std::vector<std::pair<size_t, Cost>>> successors;
-  std::vector<std::vector<std::pair<size_t, Cost>>> predecessors;
+  std::vector<std::vector<std::pair<size_t, TransferCost>>> successors;
+  std::vector<std::vector<std::pair<size_t, TransferCost>>> predecessors;
   llvm::DenseMap<ComputeInstance, size_t> 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<uint32_t> lane = std::nullopt);
-CrossbarUsage getComputeInstanceCrossbarUsage(const ComputeInstance& instance);
-bool containsCrossbarWeight(llvm::ArrayRef<CrossbarWeight> usage, const CrossbarWeight& weight);
-unsigned countCrossbarOverlap(llvm::ArrayRef<CrossbarWeight> lhs, llvm::ArrayRef<CrossbarWeight> rhs);
-size_t getCrossbarUnionSize(llvm::ArrayRef<CrossbarWeight> lhs, llvm::ArrayRef<CrossbarWeight> rhs);
-void insertCrossbarWeights(CrossbarUsage& usage, llvm::ArrayRef<CrossbarWeight> weights);
+Cost getComputeInstanceCost(const ComputeInstance& instance, const SchedulingTarget& target);
+ResidentWeightSet collectDistinctResidentWeights(mlir::Operation* owner,
+                                                 std::optional<uint32_t> lane = std::nullopt);
+ResidentWeightSet getComputeInstanceResidentWeights(const ComputeInstance& instance);
+ResidentWeightSet getComputeInstanceResidentWeights(const ComputeInstance& instance,
+                                                     const SchedulingTarget& target);
+bool containsResidentWeight(llvm::ArrayRef<ResidentWeight> usage, const ResidentWeight& weight);
+unsigned countResidentWeightOverlap(llvm::ArrayRef<ResidentWeight> lhs,
+                                    llvm::ArrayRef<ResidentWeight> rhs);
+size_t getResidentWeightUnionSize(llvm::ArrayRef<ResidentWeight> lhs,
+                                  llvm::ArrayRef<ResidentWeight> rhs);
+void insertResidentWeights(ResidentWeightSet& usage,
+                           llvm::ArrayRef<ResidentWeight> weights);
 
 } // namespace spatial
 } // namespace onnx_mlir

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp

@@ -2,23 +2,15 @@
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 
 #include <algorithm>
-#include <limits>
 #include <optional>
 
 #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<size_t>(coresCount.getValue());
-  return std::numeric_limits<size_t>::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<uint32_t>(start), static_cast<uint32_t>(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<size_t>(laneCount), getSchedulingCpuBudget());
+  assert(processorCount > 0 && "processorCount must be positive");
+  return std::min(static_cast<size_t>(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<uint32_t>(laneCount) && "lane out of range");
 
-  size_t chunkCount = getBatchChunkTargetCount(batch);
+  size_t chunkCount = getBatchChunkTargetCount(batch, processorCount);
   size_t laneCountSize = static_cast<size_t>(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<ComputeInstance, 4>
-getBatchChunksForRange(SpatComputeBatch batch, uint32_t laneStart, uint32_t laneCount) {
+getBatchChunksForRange(SpatComputeBatch batch, uint32_t laneStart,
+                       uint32_t laneCount, size_t processorCount) {
   llvm::SmallVector<ComputeInstance, 4> 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<uint32_t>(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<ProducerValueRef> getResultfulBatchProducerValueRef(SpatCom
   };
 }
 
-std::optional<ProducerValueRef> getProducerValueRef(Value value, const ComputeInstance* consumerInstance) {
+std::optional<ProducerValueRef> getProducerValueRef(
+    Value value, const ComputeInstance* consumerInstance,
+    size_t processorCount) {
   Operation* op = value.getDefiningOp();
   if (!op)
     return std::nullopt;
@@ -187,7 +193,8 @@ std::optional<ProducerValueRef> getProducerValueRef(Value value, const ComputeIn
     if (batch.getNumResults() != 0)
       return getResultfulBatchProducerValueRef(batch, value, consumerInstance);
     uint32_t lane = cast<OpResult>(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<ProducerValueRef> getProducerValueRef(Value value, const ComputeIn
   return std::nullopt;
 }
 
-std::optional<ComputeInstance> getComputeProducerInstance(Value value, const ComputeInstance* consumerInstance) {
-  if (std::optional<ProducerValueRef> producer = getProducerValueRef(value, consumerInstance))
+std::optional<ComputeInstance> getComputeProducerInstance(
+    Value value, const ComputeInstance* consumerInstance,
+    size_t processorCount) {
+  if (std::optional<ProducerValueRef> producer =
+        getProducerValueRef(value, consumerInstance, processorCount))
     return producer->instance;
   return std::nullopt;
 }

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<ComputeInstance, 4>
-getBatchChunksForRange(SpatComputeBatch batch, uint32_t laneStart, uint32_t laneCount);
+getBatchChunksForRange(SpatComputeBatch batch, uint32_t laneStart,
+                       uint32_t laneCount, size_t processorCount);
 
 std::optional<ProducerValueRef> getProducerValueRef(mlir::Value value,
-                                                    const ComputeInstance* consumerInstance = nullptr);
+                                                    const ComputeInstance* consumerInstance,
+                                                    size_t processorCount);
 std::optional<ComputeInstance> getComputeProducerInstance(mlir::Value value,
-                                                          const ComputeInstance* consumerInstance = nullptr);
+                                                          const ComputeInstance* consumerInstance,
+                                                          size_t processorCount);
 
 llvm::SmallVector<mlir::Value, 4> getComputeInstanceInputs(const ComputeInstance& instance);
 llvm::SmallVector<mlir::Value, 4> getComputeInstanceWeights(const ComputeInstance& instance);

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<ComputeInstance> dominanceOrderCompute;
   llvm::DenseMap<ComputeInstance, size_t> computeToCpuMap;
   llvm::DenseMap<ComputeInstance, size_t> computeToCpuSlotMap;

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<size_t, std::vector<std::pair<size_t, size_t>>> 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<Time>(result.computeToAestMap.lookup(source));
-    const Time targetStart = static_cast<Time>(result.computeToAestMap.lookup(target));
+    const Time targetStart =
+      static_cast<Time>(result.computeToAestMap.lookup(destination));
     if (sourceCpu == targetCpu && sourceSlot >= targetSlot)
       llvm::report_fatal_error("merge scheduling: same-CPU dependency order is invalid");
 
     Time earliestTargetStart = addOrMax(sourceStart, graph.nodes[edge.source].cost);
     if (sourceCpu != targetCpu)
       earliestTargetStart = addOrMax(
-        earliestTargetStart, getPeftTransferTime(edge.transferCost, sourceCpu, targetCpu, processorCount));
+        earliestTargetStart, getPeftTransferTime(edge.transferCost, sourceCpu, targetCpu, target));
     if (targetStart < earliestTargetStart) {
       std::string message = llvm::formatv("merge scheduling: dependency legality failed between tasks {0} and {1}",
                                           graph.nodes[edge.source].originalOrder,
@@ -89,30 +88,22 @@ void verifySchedule(const ComputeGraph& graph,
 
 } // namespace
 
-MergeSchedulingAnalysis::MergeSchedulingAnalysis(mlir::Operation* op)
-: entryOp(op) {
+MergeSchedulingAnalysis::MergeSchedulingAnalysis(mlir::Operation* op,
+                                                 const SchedulingTarget& schedulingTarget)
+: entryOp(op), target(schedulingTarget) {
   result = run();
 }
 
 MergeScheduleResult MergeSchedulingAnalysis::run() {
-  verifyExplicitPimCoreCount();
-  ComputeGraph graph = buildComputeGraph(entryOp);
+  ComputeGraph graph = buildComputeGraph(entryOp, target);
   if (!verifyAcyclic(graph))
     llvm::report_fatal_error("merge scheduling: compute graph is cyclic");
 
-  size_t processorCount = 0;
-  if (coresCount.getValue() > 0)
-    processorCount = static_cast<size_t>(coresCount.getValue());
-
   MergeScheduleResult schedule = runPeftScheduler(
     graph, PeftScheduleOptions {
-      processorCount,
-      static_cast<unsigned long>(crossbarCountInCore.getValue()),
+      target,
       entryOp->getContext()});
-  verifySchedule(graph,
-                 schedule,
-                 static_cast<unsigned long>(crossbarCountInCore.getValue()),
-                 processorCount);
+  verifySchedule(graph, schedule, target);
   return schedule;
 }
 

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.hpp

@@ -3,17 +3,19 @@
 #include "mlir/IR/Operation.h"
 
 #include "MergeSchedule.hpp"
+#include "SchedulingTarget.hpp"
 
 namespace onnx_mlir {
 namespace spatial {
 
 class MergeSchedulingAnalysis {
 public:
-  explicit MergeSchedulingAnalysis(mlir::Operation* op);
+  MergeSchedulingAnalysis(mlir::Operation* op, const SchedulingTarget& target);
   MergeScheduleResult& getResult() { return result; }
 
 private:
   mlir::Operation* entryOp = nullptr;
+  const SchedulingTarget& target;
   MergeScheduleResult result;
 
   MergeScheduleResult run();

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp

@@ -4,8 +4,8 @@
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/FormatVariadic.h"
 
-#include <cmath>
 #include <limits>
+#include <numeric>
 #include <optional>
 #include <queue>
 #include <tuple>
@@ -22,9 +22,7 @@ using namespace mlir;
 namespace {
 
 // Pressure means distinct weights exceed half the fleet's one-copy capacity.
-// The reserved headroom conservatively avoids greedy capacity fragmentation;
-// makespan, communication, instruction-path, and local-memory proxies remain stable.
-constexpr size_t kHighCrossbarPressureCapacityDivisor = 2;
+constexpr size_t kHighResidentWeightPressureCapacityDivisor = 2;
 
 struct ScheduledTask {
   size_t processor = std::numeric_limits<size_t>::max();
@@ -32,63 +30,22 @@ struct ScheduledTask {
   Time endTime = 0;
 };
 
-struct MeshModel {
-  size_t rows = 1;
-  size_t cols = 1;
-  long double averageDistance = 0.0L;
-
-  static MeshModel infer(size_t processorCount) {
-    MeshModel model;
-    if (processorCount == 0)
-      return model;
-
-    model.rows = static_cast<size_t>(std::sqrt(static_cast<long double>(processorCount)));
-    if (model.rows == 0)
-      model.rows = 1;
-    while (model.rows > 1 && processorCount % model.rows != 0)
-      --model.rows;
-    model.cols = (processorCount + model.rows - 1) / model.rows;
-
-    auto averageAxisDistance = [](size_t size) -> long double {
-      if (size <= 1)
-        return 0.0L;
-      return static_cast<long double>(size * size - 1) / (3.0L * static_cast<long double>(size));
-    };
-    model.averageDistance = averageAxisDistance(model.rows) + averageAxisDistance(model.cols);
-    return model;
-  }
-
-  std::pair<size_t, size_t> getCoord(size_t processor) const {
-    return {processor / cols, processor % cols};
-  }
-
-  size_t getDistance(size_t lhs, size_t rhs) const {
-    auto [lhsRow, lhsCol] = getCoord(lhs);
-    auto [rhsRow, rhsCol] = getCoord(rhs);
-    size_t rowDistance = lhsRow > rhsRow ? lhsRow - rhsRow : rhsRow - lhsRow;
-    size_t colDistance = lhsCol > rhsCol ? lhsCol - rhsCol : rhsCol - lhsCol;
-    return rowDistance + colDistance;
-  }
-
-  Time scaleTransferCost(Time transferCost, size_t sourceProcessor, size_t targetProcessor) const {
-    if (sourceProcessor == targetProcessor || transferCost == 0)
-      return 0;
-    long double distance = static_cast<long double>(getDistance(sourceProcessor, targetProcessor));
-    long double scale = averageDistance > 0.0L ? distance / averageDistance : 1.0L;
-    scale = std::max(0.25L, scale);
-    return static_cast<Time>(std::ceil(static_cast<long double>(transferCost) * scale));
-  }
+struct TopologyModel {
+  const SchedulingTarget& target;
 
   size_t getCenterDistance(size_t processor) const {
-    auto [row, col] = getCoord(processor);
-    size_t centerRow = rows / 2;
-    size_t centerCol = cols / 2;
-    size_t rowDistance = row > centerRow ? row - centerRow : centerRow - row;
-    size_t colDistance = col > centerCol ? col - centerCol : centerCol - col;
-    return rowDistance + colDistance;
+    Cost total = 0;
+    for (size_t other = 0; other < target.processorCount; ++other)
+      total = checkedAdd(total, target.getInterProcessorLatencyNs(processor, other));
+    return static_cast<size_t>(total);
   }
 };
 
+Time getAverageTransferTime(const TransferCost& transferCost, const SchedulingTarget& target) {
+  return checkedAdd(transferCost.fixed,
+                    checkedMultiply(transferCost.networkFlits, target.averageInterProcessorLatencyNs));
+}
+
 std::vector<std::vector<size_t>> buildReverseLevels(const ComputeGraph& graph) {
   std::vector<size_t> remainingSuccessors(graph.nodes.size(), 0);
   std::queue<size_t> readySinks;
@@ -143,14 +100,14 @@ void verifyOctTableSize(size_t nodeCount, size_t processorCount) {
   }
 }
 
-bool hasHighCrossbarPressure(const ComputeGraph& graph, size_t processorCount, size_t crossbarCapacity) {
-  if (crossbarCapacity > std::numeric_limits<size_t>::max() / processorCount)
+bool hasHighResidentWeightPressure(const ComputeGraph& graph, size_t processorCount, size_t residentWeightCapacity) {
+  if (residentWeightCapacity > std::numeric_limits<size_t>::max() / processorCount)
     return false;
-  const size_t threshold = processorCount * crossbarCapacity / kHighCrossbarPressureCapacityDivisor;
-  CrossbarUsage distinctWeights;
+  const size_t threshold = processorCount * residentWeightCapacity / kHighResidentWeightPressureCapacityDivisor;
+  ResidentWeightSet distinctWeights;
   for (const ComputeGraphNode& node : graph.nodes) {
-    for (const CrossbarWeight& weight : node.crossbarUsage) {
-      if (!containsCrossbarWeight(distinctWeights, weight))
+    for (const ResidentWeight& weight : node.residentWeights) {
+      if (!containsResidentWeight(distinctWeights, weight))
         distinctWeights.push_back(weight);
       if (distinctWeights.size() > threshold)
         return true;
@@ -159,36 +116,36 @@ bool hasHighCrossbarPressure(const ComputeGraph& graph, size_t processorCount, s
   return false;
 }
 
-std::vector<CrossbarUsage> planCrossbarReservations(const ComputeGraph& graph,
+std::vector<ResidentWeightSet> planResidentWeightReservations(const ComputeGraph& graph,
                                                     size_t processorCount,
-                                                    size_t crossbarCapacity,
-                                                    const MeshModel& mesh,
-                                                    bool preferCrossbarReuse) {
+                                                              size_t residentWeightCapacity,
+                                                              const TopologyModel& topology,
+                                                              bool preferWeightReuse) {
   std::vector<size_t> weightedTasks;
   for (size_t task = 0; task < graph.nodes.size(); ++task)
-    if (!graph.nodes[task].crossbarUsage.empty())
+    if (!graph.nodes[task].residentWeights.empty())
       weightedTasks.push_back(task);
   llvm::sort(weightedTasks, [&](size_t lhs, size_t rhs) {
-    if (graph.nodes[lhs].crossbarUsage.size() != graph.nodes[rhs].crossbarUsage.size())
-      return graph.nodes[lhs].crossbarUsage.size() > graph.nodes[rhs].crossbarUsage.size();
+    if (graph.nodes[lhs].residentWeights.size() != graph.nodes[rhs].residentWeights.size())
+      return graph.nodes[lhs].residentWeights.size() > graph.nodes[rhs].residentWeights.size();
     return graph.nodes[lhs].originalOrder < graph.nodes[rhs].originalOrder;
   });
 
-  std::vector<CrossbarUsage> reservations(processorCount);
+  std::vector<ResidentWeightSet> reservations(processorCount);
   std::vector<Time> reservedLoad(processorCount, 0);
   for (size_t task : weightedTasks) {
     size_t bestProcessor = std::numeric_limits<size_t>::max();
     using ReservationScore = std::tuple<Time, size_t, size_t, size_t>;
     std::optional<ReservationScore> bestScore;
     for (size_t processor = 0; processor < processorCount; ++processor) {
-      size_t crossbarUnion =
-        getCrossbarUnionSize(reservations[processor], graph.nodes[task].crossbarUsage);
-      if (crossbarUnion > crossbarCapacity)
+      size_t residentWeightUnion =
+        getResidentWeightUnionSize(reservations[processor], graph.nodes[task].residentWeights);
+      if (residentWeightUnion > residentWeightCapacity)
         continue;
-      size_t addedCrossbars = crossbarUnion - reservations[processor].size();
-      ReservationScore score {preferCrossbarReuse ? addedCrossbars : reservedLoad[processor],
-                              preferCrossbarReuse ? reservedLoad[processor] : addedCrossbars,
-                              mesh.getCenterDistance(processor),
+      size_t addedWeights = residentWeightUnion - reservations[processor].size();
+      ReservationScore score {preferWeightReuse ? addedWeights : reservedLoad[processor],
+                              preferWeightReuse ? reservedLoad[processor] : addedWeights,
+                              topology.getCenterDistance(processor),
                               processor};
       if (!bestScore || score < *bestScore) {
         bestProcessor = processor;
@@ -200,13 +157,13 @@ std::vector<CrossbarUsage> planCrossbarReservations(const ComputeGraph& graph,
         llvm::formatv("PEFT reservation planner: cannot place task {0} with {1} distinct weights in {2} "
                       "processors of capacity {3}",
                       graph.nodes[task].originalOrder,
-                      graph.nodes[task].crossbarUsage.size(),
+                      graph.nodes[task].residentWeights.size(),
                       processorCount,
-                      crossbarCapacity)
+                      residentWeightCapacity)
           .str();
       llvm::report_fatal_error(llvm::StringRef(message));
     }
-    insertCrossbarWeights(reservations[bestProcessor], graph.nodes[task].crossbarUsage);
+    insertResidentWeights(reservations[bestProcessor], graph.nodes[task].residentWeights);
     reservedLoad[bestProcessor] = addOrMax(reservedLoad[bestProcessor], graph.nodes[task].cost);
   }
   return reservations;
@@ -214,8 +171,8 @@ std::vector<CrossbarUsage> planCrossbarReservations(const ComputeGraph& graph,
 
 using LanePublicationSignatures = llvm::SmallVector<llvm::SmallVector<int64_t, 8>, 8>;
 
-FailureOr<LanePublicationSignatures>
-buildLanePublicationSignatures(SpatComputeBatch batch, GraphBatchPublicationCache& publicationCache) {
+FailureOr<LanePublicationSignatures> buildLanePublicationSignatures(SpatComputeBatch batch,
+                                                                    GraphBatchPublicationCache& publicationCache) {
   LanePublicationSignatures signatures(batch.getLaneCount());
   for (auto [resultIndex, result] : llvm::enumerate(batch.getResults())) {
     auto publicationMap = getGraphBatchPublicationMap(batch, resultIndex, publicationCache);
@@ -231,8 +188,8 @@ buildLanePublicationSignatures(SpatComputeBatch batch, GraphBatchPublicationCach
       auto sourceOffsets = blueprint.getFragmentSourceOffsets();
       auto fragmentStrides = blueprint.getFragmentStrides();
       auto outputType = dyn_cast<RankedTensorType>(blueprint.getOutput().getType());
-      if (!operandIndices || !sourceSlots || !sourceOffsets || !fragmentStrides
-          || !outputType || !outputType.hasStaticShape())
+      if (!operandIndices || !sourceSlots || !sourceOffsets || !fragmentStrides || !outputType
+          || !outputType.hasStaticShape())
         return blueprint.emitOpError("PEFT publication compatibility requires complete static fragment metadata"),
                failure();
 
@@ -241,11 +198,9 @@ buildLanePublicationSignatures(SpatComputeBatch batch, GraphBatchPublicationCach
       int64_t rank = outputType.getRank();
       if (rank <= 0 || fragmentOffsets.size() != fragmentSizes.size()
           || fragmentOffsets.size() != fragmentStrides->size()
-          || fragmentOffsets.size() != operandIndices->size() * rank
-          || sourceSlots->size() != operandIndices->size()
+          || fragmentOffsets.size() != operandIndices->size() * rank || sourceSlots->size() != operandIndices->size()
           || sourceOffsets->size() != operandIndices->size())
-        return blueprint.emitOpError("PEFT publication compatibility found inconsistent fragment metadata"),
-               failure();
+        return blueprint.emitOpError("PEFT publication compatibility found inconsistent fragment metadata"), failure();
 
       llvm::SmallVector<llvm::SmallVector<size_t, 2>, 8> fragmentsByLane(batch.getLaneCount());
       for (auto [fragmentIndex, operandIndex] : llvm::enumerate(*operandIndices)) {
@@ -287,19 +242,83 @@ buildLanePublicationSignatures(SpatComputeBatch batch, GraphBatchPublicationCach
 
 } // namespace
 
-Time getPeftTransferTime(Time transferCost, size_t sourceProcessor, size_t targetProcessor, size_t processorCount) {
-  return MeshModel::infer(processorCount).scaleTransferCost(transferCost, sourceProcessor, targetProcessor);
+std::vector<size_t> mapLogicalProcessorsToPhysicalCores(ArrayRef<Cost> logicalTrafficFlits,
+                                                        const SchedulingTarget& target) {
+  const size_t processorCount = target.processorCount;
+  assert(logicalTrafficFlits.size() == processorCount * processorCount
+         && "logical traffic matrix must cover every processor pair");
+
+  std::vector<size_t> physicalCoreForLogicalProcessor(processorCount);
+  std::iota(physicalCoreForLogicalProcessor.begin(), physicalCoreForLogicalProcessor.end(), 0);
+
+  auto transferCost = [&](size_t sourceLogicalProcessor,
+                          size_t targetLogicalProcessor,
+                          size_t sourcePhysicalCore,
+                          size_t targetPhysicalCore) {
+    Cost traffic = logicalTrafficFlits[sourceLogicalProcessor * processorCount + targetLogicalProcessor];
+    return checkedMultiply(traffic, target.getInterProcessorLatencyNs(sourcePhysicalCore, targetPhysicalCore));
+  };
+
+  for (size_t logicalProcessor = 0; logicalProcessor < processorCount; ++logicalProcessor) {
+    size_t bestPeerLogicalProcessor = logicalProcessor;
+    Cost bestSaving = 0;
+    for (size_t peerLogicalProcessor = 0; peerLogicalProcessor < processorCount; ++peerLogicalProcessor) {
+      if (peerLogicalProcessor == logicalProcessor)
+        continue;
+      size_t physicalCore = physicalCoreForLogicalProcessor[logicalProcessor];
+      size_t peerPhysicalCore = physicalCoreForLogicalProcessor[peerLogicalProcessor];
+      Cost currentCost = 0;
+      Cost swappedCost = 0;
+      for (size_t otherLogicalProcessor = 0; otherLogicalProcessor < processorCount; ++otherLogicalProcessor) {
+        if (otherLogicalProcessor == logicalProcessor || otherLogicalProcessor == peerLogicalProcessor)
+          continue;
+        size_t otherPhysicalCore = physicalCoreForLogicalProcessor[otherLogicalProcessor];
+        currentCost = checkedAdd(
+          currentCost, transferCost(logicalProcessor, otherLogicalProcessor, physicalCore, otherPhysicalCore));
+        currentCost = checkedAdd(
+          currentCost, transferCost(otherLogicalProcessor, logicalProcessor, otherPhysicalCore, physicalCore));
+        currentCost = checkedAdd(
+          currentCost, transferCost(peerLogicalProcessor, otherLogicalProcessor, peerPhysicalCore, otherPhysicalCore));
+        currentCost = checkedAdd(
+          currentCost, transferCost(otherLogicalProcessor, peerLogicalProcessor, otherPhysicalCore, peerPhysicalCore));
+        swappedCost = checkedAdd(
+          swappedCost, transferCost(logicalProcessor, otherLogicalProcessor, peerPhysicalCore, otherPhysicalCore));
+        swappedCost = checkedAdd(
+          swappedCost, transferCost(otherLogicalProcessor, logicalProcessor, otherPhysicalCore, peerPhysicalCore));
+        swappedCost = checkedAdd(
+          swappedCost, transferCost(peerLogicalProcessor, otherLogicalProcessor, physicalCore, otherPhysicalCore));
+        swappedCost = checkedAdd(
+          swappedCost, transferCost(otherLogicalProcessor, peerLogicalProcessor, otherPhysicalCore, physicalCore));
+      }
+      currentCost =
+        checkedAdd(currentCost, transferCost(logicalProcessor, peerLogicalProcessor, physicalCore, peerPhysicalCore));
+      currentCost =
+        checkedAdd(currentCost, transferCost(peerLogicalProcessor, logicalProcessor, peerPhysicalCore, physicalCore));
+      swappedCost =
+        checkedAdd(swappedCost, transferCost(logicalProcessor, peerLogicalProcessor, peerPhysicalCore, physicalCore));
+      swappedCost =
+        checkedAdd(swappedCost, transferCost(peerLogicalProcessor, logicalProcessor, physicalCore, peerPhysicalCore));
+      if (currentCost > swappedCost && currentCost - swappedCost > bestSaving) {
+        bestSaving = currentCost - swappedCost;
+        bestPeerLogicalProcessor = peerLogicalProcessor;
+      }
+    }
+    std::swap(physicalCoreForLogicalProcessor[logicalProcessor],
+              physicalCoreForLogicalProcessor[bestPeerLogicalProcessor]);
+  }
+  return physicalCoreForLogicalProcessor;
 }
 
 MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftScheduleOptions& options) {
   const size_t nodeCount = graph.nodes.size();
-  const size_t processorCount = options.processorCount;
+  const size_t processorCount = options.target.processorCount;
   if (processorCount == 0)
     llvm::report_fatal_error("PEFT scheduler: processor count must be positive");
-  MeshModel mesh = MeshModel::infer(processorCount);
-  const bool preferCrossbarReuse = hasHighCrossbarPressure(graph, processorCount, options.crossbarCapacity);
-  std::vector<CrossbarUsage> capacityReservations =
-    planCrossbarReservations(graph, processorCount, options.crossbarCapacity, mesh, preferCrossbarReuse);
+  TopologyModel topology {options.target};
+  const bool preferWeightReuse =
+    hasHighResidentWeightPressure(graph, processorCount, options.target.residentWeightCapacity);
+  std::vector<ResidentWeightSet> capacityReservations = planResidentWeightReservations(
+    graph, processorCount, options.target.residentWeightCapacity, topology, preferWeightReuse);
 
   verifyOctTableSize(nodeCount, processorCount);
   std::vector<std::vector<size_t>> reverseLevels = buildReverseLevels(graph);
@@ -317,7 +336,7 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
       std::vector<Time> maxVals(processorCount, 0);
 
       for (const auto& [succ, comm] : graph.successors[task]) {
-        Time valDifferentCpu = addOrMax(minOctPlusComp[succ], comm);
+        Time valDifferentCpu = addOrMax(minOctPlusComp[succ], getAverageTransferTime(comm, options.target));
         for (size_t processor = 0; processor < processorCount; ++processor) {
           Time valSameCpu = addOrMax(oct[succ * processorCount + processor], getComputeCost(succ, processor));
           Time bestSucc = std::min(valSameCpu, valDifferentCpu);
@@ -378,7 +397,7 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
   }
 
   std::vector<char> scheduled(nodeCount, false);
-  std::vector<CrossbarUsage> processorCrossbars(processorCount);
+  std::vector<ResidentWeightSet> processorResidentWeights(processorCount);
   std::vector<ScheduledTask> schedules(nodeCount);
   std::vector<std::vector<size_t>> tasksByProcessor(processorCount);
 
@@ -394,23 +413,25 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
     Time bestEft = 0;
     Time bestOeft = std::numeric_limits<Time>::max();
     unsigned int bestOverlapCount = 0;
+    size_t bestTaskCount = std::numeric_limits<size_t>::max();
     size_t bestCenterDistance = std::numeric_limits<size_t>::max();
-    size_t smallestCrossbarUnion = std::numeric_limits<size_t>::max();
-    bool crossbarRejected = false;
+    size_t smallestResidentWeightUnion = std::numeric_limits<size_t>::max();
+    bool residentWeightRejected = false;
 
     for (size_t processor = 0; processor < processorCount; ++processor) {
-      unsigned int overlapCount = countCrossbarOverlap(processorCrossbars[processor], graph.nodes[task].crossbarUsage);
-      size_t crossbarUnion =
-        getCrossbarUnionSize(capacityReservations[processor], graph.nodes[task].crossbarUsage);
-      smallestCrossbarUnion = std::min(smallestCrossbarUnion, crossbarUnion);
-      if (!graph.nodes[task].crossbarUsage.empty() && crossbarUnion > options.crossbarCapacity) {
-        crossbarRejected = true;
+      unsigned int overlapCount =
+        countResidentWeightOverlap(processorResidentWeights[processor], graph.nodes[task].residentWeights);
+      size_t residentWeightUnion =
+        getResidentWeightUnionSize(capacityReservations[processor], graph.nodes[task].residentWeights);
+      smallestResidentWeightUnion = std::min(smallestResidentWeightUnion, residentWeightUnion);
+      if (!graph.nodes[task].residentWeights.empty() && residentWeightUnion > options.target.residentWeightCapacity) {
+        residentWeightRejected = true;
         continue;
       }
       Time dataReady = 0;
       for (const auto& [pred, comm] : graph.predecessors[task]) {
         const ScheduledTask& predSchedule = schedules[pred];
-        Time commPenalty = getPeftTransferTime(comm, predSchedule.processor, processor, processorCount);
+        Time commPenalty = getPeftTransferTime(comm, predSchedule.processor, processor, options.target);
         dataReady = std::max(dataReady, addOrMax(predSchedule.endTime, commPenalty));
       }
 
@@ -437,9 +458,10 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
 
       Time eft = addOrMax(est, computeCost);
       Time oeft = addOrMax(eft, oct[task * processorCount + processor]);
-      size_t centerDistance = mesh.getCenterDistance(processor);
-      bool betterCrossbarChoice =
-        preferCrossbarReuse ? overlapCount > bestOverlapCount : overlapCount < bestOverlapCount;
+      size_t centerDistance = topology.getCenterDistance(processor);
+      size_t taskCount = tasksByProcessor[processor].size();
+      bool betterResidentWeightChoice =
+        preferWeightReuse ? overlapCount > bestOverlapCount : overlapCount < bestOverlapCount;
 
       if (oeft < bestOeft || (oeft == bestOeft && eft < bestEft)
           || (oeft == bestOeft && eft == bestEft && est < bestEst)) {
@@ -448,40 +470,52 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
         bestEft = eft;
         bestOeft = oeft;
         bestOverlapCount = overlapCount;
+        bestTaskCount = taskCount;
         bestCenterDistance = centerDistance;
       }
-      else if (oeft == bestOeft && eft == bestEft && est == bestEst
+      else if (oeft == bestOeft && eft == bestEft && est == bestEst && taskCount < bestTaskCount) {
+        bestProcessor = processor;
+        bestEst = est;
+        bestEft = eft;
+        bestOeft = oeft;
+        bestOverlapCount = overlapCount;
+        bestTaskCount = taskCount;
+        bestCenterDistance = centerDistance;
+      }
+      else if (oeft == bestOeft && eft == bestEft && est == bestEst && taskCount == bestTaskCount
                && centerDistance < bestCenterDistance) {
         bestProcessor = processor;
         bestEst = est;
         bestEft = eft;
         bestOeft = oeft;
         bestOverlapCount = overlapCount;
+        bestTaskCount = taskCount;
         bestCenterDistance = centerDistance;
       }
-      else if (oeft == bestOeft && eft == bestEft && est == bestEst
-               && centerDistance == bestCenterDistance && betterCrossbarChoice) {
+      else if (oeft == bestOeft && eft == bestEft && est == bestEst && taskCount == bestTaskCount
+               && centerDistance == bestCenterDistance && betterResidentWeightChoice) {
         bestProcessor = processor;
         bestEst = est;
         bestEft = eft;
         bestOeft = oeft;
         bestOverlapCount = overlapCount;
+        bestTaskCount = taskCount;
         bestCenterDistance = centerDistance;
       }
     }
 
     if (bestProcessor == std::numeric_limits<size_t>::max()) {
-      if (crossbarRejected) {
+      if (residentWeightRejected) {
         const ComputeInstance& instance = graph.nodes[task].instance;
         std::string message =
           llvm::formatv("PEFT scheduler: no valid processor for task {0} (lanes {1}..{2}, {3} distinct weights); "
-                        "smallest processor union is {4}, exceeding crossbar capacity {5}",
+                        "smallest processor union is {4}, exceeding resident-weight capacity {5}",
                         graph.nodes[task].originalOrder,
                         instance.laneStart,
                         instance.laneStart + instance.laneCount,
-                        graph.nodes[task].crossbarUsage.size(),
-                        smallestCrossbarUnion,
-                        options.crossbarCapacity)
+                        graph.nodes[task].residentWeights.size(),
+                        smallestResidentWeightUnion,
+                        options.target.residentWeightCapacity)
             .str();
         llvm::report_fatal_error(llvm::StringRef(message));
       }
@@ -495,8 +529,8 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
     schedules[task] = {bestProcessor, bestEst, bestEft};
     scheduled[task] = true;
     ++scheduledCount;
-    insertCrossbarWeights(capacityReservations[bestProcessor], graph.nodes[task].crossbarUsage);
-    insertCrossbarWeights(processorCrossbars[bestProcessor], graph.nodes[task].crossbarUsage);
+    insertResidentWeights(capacityReservations[bestProcessor], graph.nodes[task].residentWeights);
+    insertResidentWeights(processorResidentWeights[bestProcessor], graph.nodes[task].residentWeights);
 
     // 3. CRITICAL FIX: Topological Append
     // Because the readyQueue pops in strict topological order, simply pushing to the
@@ -576,6 +610,7 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
 
   // 6. Populate Final Result
   MergeScheduleResult result;
+  result.processorCount = processorCount;
   result.dominanceOrderCompute.reserve(nodeCount);
 
   for (size_t task : scheduledOrder)

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.hpp

@@ -4,18 +4,32 @@
 
 #include "ComputeGraph.hpp"
 #include "MergeSchedule.hpp"
+#include "SchedulingTarget.hpp"
 
 namespace onnx_mlir {
 namespace spatial {
 
 struct PeftScheduleOptions {
-  size_t processorCount = 0;
-  unsigned long crossbarCapacity = 0;
+  SchedulingTarget target;
   mlir::MLIRContext* context = nullptr;
 };
 
-Time getPeftTransferTime(Time transferCost, size_t sourceProcessor, size_t targetProcessor, size_t processorCount);
+inline Time getPeftTransferTime(const TransferCost& transferCost,
+                                size_t sourceProcessor,
+                                size_t targetProcessor,
+                                const SchedulingTarget& target) {
+  if (sourceProcessor == targetProcessor)
+    return 0;
+  return checkedAdd(transferCost.fixed,
+                    checkedMultiply(transferCost.networkFlits, target.averageInterProcessorLatencyNs));
+}
+
+// PEFT assigns logical processors. Physical core IDs are chosen only after
+// materialization exposes the exact transfer traffic.
 MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftScheduleOptions& options);
 
+std::vector<size_t> mapLogicalProcessorsToPhysicalCores(llvm::ArrayRef<Cost> logicalTrafficFlits,
+                                                        const SchedulingTarget& target);
+
 } // namespace spatial
 } // namespace onnx_mlir

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/SchedulingTarget.hpp

@@ -0,0 +1,54 @@
+#pragma once
+
+#include <cassert>
+#include <cstddef>
+#include <vector>
+
+#include "Utils.hpp"
+
+namespace onnx_mlir {
+namespace spatial {
+
+struct SchedulingTarget {
+  size_t processorCount = 0;
+  size_t residentWeightCapacity = 0;
+  std::vector<Cost> interProcessorLatencyNs;
+  Cost averageInterProcessorLatencyNs = 0;
+
+  Cost computeBitwidth = 8;
+  Cost processorPeriodNs = 1;
+  Cost localMemoryWidthBytes = 64;
+  Cost localMemoryReadLatencyCycles = 1;
+  Cost localMemoryWriteLatencyCycles = 1;
+  Cost transferWidthBytes = 8;
+  Cost vectorWidth = 16;
+  Cost vectorLatencyCycles = 4;
+
+  Cost matrixRows = 128;
+  Cost matrixColumns = 128;
+  Cost matrixPeriodNs = 1;
+  Cost matrixInputResolutionBits = 1;
+  Cost matrixInputLatencyCycles = 1;
+  Cost matrixInputParallelism = 128;
+  Cost matrixReadLatencyNs = 30;
+  Cost matrixSampleLatencyCycles = 1;
+  Cost matrixOutputLatencyCycles = 10;
+  Cost matrixOutputParallelism = 2;
+  Cost matrixShiftLatencyCycles = 1;
+  Cost matrixBufferLatencyCycles = 1;
+  Cost matrixInputBufferLatencyCycles = 0;
+  bool matrixPipeline = true;
+
+  Cost getInterProcessorLatencyNs(size_t source,
+                                  size_t destination) const {
+    assert(source < processorCount && destination < processorCount
+           && "processor index out of range");
+    assert(interProcessorLatencyNs.size()
+             == processorCount * processorCount
+           && "incomplete inter-processor latency matrix");
+    return interProcessorLatencyNs[source * processorCount + destination];
+  }
+};
+
+} // namespace spatial
+} // namespace onnx_mlir

src/PIM/Dialect/Spatial/Transforms/TrivialGraphComputeMergePass.cpp

@@ -36,10 +36,14 @@ static bool hasOnlyStructuralAttrs(ComputeOp op) {
   });
 }
 
-static bool hasCapacityFor(Operation* producer, Operation* consumer) {
-  CrossbarUsage producerWeights = collectDistinctCrossbarWeights(producer);
-  CrossbarUsage consumerWeights = collectDistinctCrossbarWeights(consumer);
-  return getCrossbarUnionSize(producerWeights, consumerWeights) <= static_cast<size_t>(crossbarCountInCore.getValue());
+static bool hasCapacityFor(Operation* producer, Operation* consumer,
+                           size_t residentWeightCapacity) {
+  ResidentWeightSet producerWeights =
+    collectDistinctResidentWeights(producer);
+  ResidentWeightSet consumerWeights =
+    collectDistinctResidentWeights(consumer);
+  return getResidentWeightUnionSize(producerWeights, consumerWeights)
+      <= residentWeightCapacity;
 }
 
 template <typename ConsumerOp>
@@ -156,8 +160,11 @@ static void mapExternalArguments(OldOp oldOp, NewOp newOp, IRMapping& mapper, bo
 }
 
 struct MergeTrivialScalarComputes : OpRewritePattern<SpatGraphCompute> {
-  MergeTrivialScalarComputes(MLIRContext *context, TrivialGraphMergeStats *stats)
-      : OpRewritePattern(context), stats(stats) {}
+  MergeTrivialScalarComputes(MLIRContext *context,
+                             TrivialGraphMergeStats *stats,
+                             size_t residentWeightCapacity)
+      : OpRewritePattern(context), stats(stats),
+        residentWeightCapacity(residentWeightCapacity) {}
 
   LogicalResult matchAndRewrite(SpatGraphCompute consumer, PatternRewriter& rewriter) const override {
     SpatGraphCompute producer;
@@ -166,7 +173,8 @@ struct MergeTrivialScalarComputes : OpRewritePattern<SpatGraphCompute> {
       if (candidate && candidate->getBlock() == consumer->getBlock() && hasOnlyStructuralAttrs(candidate)
           && hasOnlyStructuralAttrs(consumer) && isUniqueGraphComputePredecessor(candidate, consumer)
           && isExclusivelyConsumedBy(candidate, consumer)
-          && hasCapacityFor(candidate, consumer) && hasNoNestedArgumentCaptures(candidate)
+          && hasCapacityFor(candidate, consumer, residentWeightCapacity)
+          && hasNoNestedArgumentCaptures(candidate)
           && hasNoNestedArgumentCaptures(consumer)) {
         producer = candidate;
         break;
@@ -201,6 +209,7 @@ struct MergeTrivialScalarComputes : OpRewritePattern<SpatGraphCompute> {
 
 private:
   TrivialGraphMergeStats *stats;
+  size_t residentWeightCapacity;
 };
 
 static bool isLaneIndex(Value value, Value lane, int64_t laneCount) {
@@ -397,8 +406,11 @@ static bool hasDirectLaneConsumers(SpatGraphComputeBatch producer, SpatGraphComp
 }
 
 struct MergeTrivialBatchComputes : OpRewritePattern<SpatGraphComputeBatch> {
-  MergeTrivialBatchComputes(MLIRContext *context, TrivialGraphMergeStats *stats)
-      : OpRewritePattern(context), stats(stats) {}
+  MergeTrivialBatchComputes(MLIRContext *context,
+                            TrivialGraphMergeStats *stats,
+                            size_t residentWeightCapacity)
+      : OpRewritePattern(context), stats(stats),
+        residentWeightCapacity(residentWeightCapacity) {}
 
   LogicalResult matchAndRewrite(SpatGraphComputeBatch consumer, PatternRewriter& rewriter) const override {
     SpatGraphComputeBatch producer;
@@ -409,7 +421,8 @@ struct MergeTrivialBatchComputes : OpRewritePattern<SpatGraphComputeBatch> {
           && candidate.getLaneCount() == consumer.getLaneCount() && hasOnlyStructuralAttrs(candidate)
           && hasOnlyStructuralAttrs(consumer) && isUniqueGraphComputePredecessor(candidate, consumer)
           && isExclusivelyConsumedBy(candidate, consumer)
-          && hasCapacityFor(candidate, consumer) && hasDirectLaneConsumers(candidate, consumer)
+          && hasCapacityFor(candidate, consumer, residentWeightCapacity)
+          && hasDirectLaneConsumers(candidate, consumer)
           && succeeded(fragments = collectPublishedFragments(candidate))) {
         producer = candidate;
         break;
@@ -469,11 +482,16 @@ struct MergeTrivialBatchComputes : OpRewritePattern<SpatGraphComputeBatch> {
 
 private:
   TrivialGraphMergeStats *stats;
+  size_t residentWeightCapacity;
 };
 
 struct TrivialGraphComputeMergePass final : PassWrapper<TrivialGraphComputeMergePass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TrivialGraphComputeMergePass)
 
+  TrivialGraphComputeMergePass() = default;
+  explicit TrivialGraphComputeMergePass(size_t residentWeightCapacity)
+      : residentWeightCapacity(residentWeightCapacity) {}
+
   StringRef getArgument() const override { return "pim-trivial-graph-compute-merge"; }
   StringRef getDescription() const override {
     return "Inline linear exclusive graph compute chains while preserving fan-in boundaries.";
@@ -481,10 +499,17 @@ struct TrivialGraphComputeMergePass final : PassWrapper<TrivialGraphComputeMerge
 
   void runOnOperation() override {
     ModuleOp module = getOperation();
+    if (residentWeightCapacity == 0) {
+      module.emitError(
+        "TrivialGraphComputeMerge requires an explicit valid resident-weight capacity");
+      signalPassFailure();
+      return;
+    }
     TrivialGraphMergeStats stats;
     std::tie(stats.scalarBefore, stats.batchBefore) = countGraphComputes(module);
     RewritePatternSet patterns(&getContext());
-    patterns.add<MergeTrivialScalarComputes, MergeTrivialBatchComputes>(&getContext(), &stats);
+    patterns.add<MergeTrivialScalarComputes, MergeTrivialBatchComputes>(
+      &getContext(), &stats, residentWeightCapacity);
     if (failed(applyPatternsGreedily(module, std::move(patterns)))) {
       signalPassFailure();
       return;
@@ -510,6 +535,9 @@ struct TrivialGraphComputeMergePass final : PassWrapper<TrivialGraphComputeMerge
         signalPassFailure();
     }
   }
+
+private:
+  size_t residentWeightCapacity = 0;
 };
 
 } // namespace
@@ -519,4 +547,10 @@ std::unique_ptr<Pass> createTrivialGraphComputeMergePass() {
   return std::make_unique<spatial::TrivialGraphComputeMergePass>();
 }
 
+std::unique_ptr<Pass> createTrivialGraphComputeMergePass(
+    size_t residentWeightCapacity) {
+  return std::make_unique<spatial::TrivialGraphComputeMergePass>(
+    residentWeightCapacity);
+}
+
 } // namespace onnx_mlir

src/PIM/Pass/PIMPasses.h

@@ -2,10 +2,14 @@
 
 #include "mlir/Pass/Pass.h"
 
+#include <cstddef>
 #include <memory>
 #include <string>
 
 namespace onnx_mlir {
+namespace spatial {
+struct SchedulingTarget;
+}
 
 std::unique_ptr<mlir::Pass> createONNXToSpatialPass();
 std::unique_ptr<mlir::Pass> createSpatialLayoutPlanningPass();
@@ -16,8 +20,11 @@ std::unique_ptr<mlir::Pass> createSpatialToPimPass();
 std::unique_ptr<mlir::Pass> createPimBufferizationPass();
 
 std::unique_ptr<mlir::Pass> createMergeComputeNodesPass();
+std::unique_ptr<mlir::Pass> createMergeComputeNodesPass(const spatial::SchedulingTarget& target);
 
 std::unique_ptr<mlir::Pass> createTrivialGraphComputeMergePass();
+std::unique_ptr<mlir::Pass> createTrivialGraphComputeMergePass(
+  size_t residentWeightCapacity);
 
 std::unique_ptr<mlir::Pass> createPimHostConstantFoldingPass();
 

src/PIM/PimAccelerator.cpp

@@ -76,8 +76,8 @@ void PimAccelerator::registerPasses(int optLevel) const {
   registerPass(createLowerSpatialPlansPass);
   registerPass(createSpatialToPimPass);
   registerPass(createPimBufferizationPass);
-  registerPass(createTrivialGraphComputeMergePass);
-  registerPass(createMergeComputeNodesPass);
+  mlir::registerPass([] { return createTrivialGraphComputeMergePass(); });
+  mlir::registerPass([] { return createMergeComputeNodesPass(); });
   registerPass(createPimHostConstantFoldingPass);
   registerPass(createPimLocalMemoryPlanningPass);
   registerPass(createPimVerificationPass);

test/PIM/CMakeLists.txt

@@ -29,3 +29,10 @@ add_pim_unittest(PimMemoryLivenessPlannerTest
   LINK_LIBS PRIVATE
   OMPimCompilerUtils
 )
+
+add_pim_unittest(SpatialSchedulingTargetTest
+  SpatialSchedulingTargetTest.cpp
+
+  LINK_LIBS PRIVATE
+  OMPimCompilerUtils
+)

test/PIM/SpatialSchedulingTargetTest.cpp

@@ -0,0 +1,58 @@
+#include <cassert>
+#include <cstdlib>
+#include <vector>
+
+#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.hpp"
+
+using namespace onnx_mlir::spatial;
+
+int main() {
+  TransferCost transfer {.fixed = 50, .networkFlits = 4};
+
+  SchedulingTarget fast;
+  fast.processorCount = 2;
+  fast.interProcessorLatencyNs = {0, 3, 3, 0};
+  fast.averageInterProcessorLatencyNs = 3;
+
+  SchedulingTarget slow = fast;
+  slow.interProcessorLatencyNs = {0, 10, 10, 0};
+  slow.averageInterProcessorLatencyNs = 10;
+
+  assert(getPeftTransferTime(transfer, 0, 0, fast) == 0);
+  assert(getPeftTransferTime(transfer, 0, 1, fast) == 62);
+  assert(getPeftTransferTime(transfer, 0, 1, slow) == 90);
+  assert(fast.getInterProcessorLatencyNs(0, 1) == 3);
+  assert(slow.getInterProcessorLatencyNs(0, 1) == 10);
+
+  SchedulingTarget line;
+  line.processorCount = 3;
+  line.interProcessorLatencyNs = {
+    0,
+    1,
+    10,
+    1,
+    0,
+    1,
+    10,
+    1,
+    0,
+  };
+  std::vector<Cost> logicalTrafficFlits(9, 0);
+  logicalTrafficFlits[2] = 100;
+  assert(mapLogicalProcessorsToPhysicalCores(logicalTrafficFlits, line) == std::vector<size_t>({1, 0, 2}));
+
+  SchedulingTarget alreadyPlaced = line;
+  alreadyPlaced.interProcessorLatencyNs = {
+    0,
+    10,
+    1,
+    10,
+    0,
+    1,
+    1,
+    1,
+    0,
+  };
+  assert(mapLogicalProcessorsToPhysicalCores(logicalTrafficFlits, alreadyPlaced) == std::vector<size_t>({0, 1, 2}));
+  return EXIT_SUCCESS;
+}

validation/.gitignore

@@ -13,3 +13,4 @@ networks/**/real_image_val
 networks/**/*.png
 networks/**/*.jpg
 networks/**/*.csv
+!networks/pimcomp_models/results.csv

validation/README.md

@@ -8,7 +8,7 @@ model it can:
 3. compile PIM artifacts with Raptor;
 4. run the reference implementation and functional PIM simulator;
 5. compare their outputs;
-6. run `pimsim-nn` to report latency and power.
+6. run `pimsim-nn` to report latency, power, and energy.
 
 Run the script from the repository root with the repository Python environment.
 
@@ -65,7 +65,7 @@ The script discovers them recursively.
 The PIMCOMP paper-model suite has a one-command Arch-A comparison:
 
 ```bash
-.venv/bin/python validation/tools/run_pimcomp_paper_latency.py
+.venv/bin/python validation/tools/pimcomp/run_pimcomp_paper_latency.py
 ```
 
 The runner verifies PIMCOMP's population-200, 1000-iteration GA settings,
@@ -134,7 +134,7 @@ count with `-j` or `--jobs`:
 | `--simulator-dir PATH` | Functional `pim-simulator` crate directory. Defaults to the in-tree simulator. |
 | `--non-functional-simulator-build-dir PATH` | `pimsim-nn` build directory. Defaults to the in-tree build. |
 | `--pimcomp-config {arch-a,arch-b,arch-c}` | Non-functional hardware/timing profile. Defaults to `arch-a`. |
-| `--skip-non-functional-simulation` | Skip `pimsim-nn` latency and power measurement. |
+| `--skip-non-functional-simulation` | Skip `pimsim-nn` latency, power, and energy measurement. |
 | `--threshold FLOAT` | Absolute output-comparison tolerance. Defaults to `1e-3`. |
 | `--relative-threshold FLOAT` | Relative output-comparison tolerance. Defaults to `1e-5`. |
 | `--seed INT` | Seed for generated inputs. Defaults to `0`. |
@@ -165,17 +165,17 @@ prints the incompatible values; functional validation still runs.
 The checked-in profiles are under
 `validation/pimsim_configs/pimcomp/<profile>/latency_config.json`.
 
-Use `--skip-non-functional-simulation` when latency and power are not required.
+Use `--skip-non-functional-simulation` when latency, power, and energy are not required.
 The summary reports non-functional results as measured, failed, unsupported, or
 skipped.
 
 Overall PASS/FAIL is determined by compilation and functional output
 comparison. A non-functional simulation failure remains visible as `ERROR` in
-the latency and power columns but does not change a functional PASS.
+the latency, power, and energy columns but does not change a functional PASS.
 
 `pimsim-nn` does not currently implement the `vsoftmax` instruction. When its
 explicit unsupported-op diagnostic is encountered, Softmax validations retain
-their functional PASS and show `UNSUPPORTED` in both non-functional columns.
+their functional PASS and show `UNSUPPORTED` in the non-functional columns.
 Other `pimsim-nn` failures remain `ERROR`.
 
 ## Generated artifacts
@@ -186,7 +186,7 @@ Artifacts are written beside each model:
 |---|---|
 | `inputs/` | Generated input CSV files. |
 | `outputs/` | ONNX-MLIR reference output CSV files. |
-| `raptor/` | Exported MLIR, dialect snapshots, reports, and final `pim/` artifacts. |
+| `raptor/` | Exported MLIR, dialect snapshots, reports, final FP32 `pim/` artifacts, and the int8-equivalent `pimsim_nn/` latency view. |
 | `runner/` | Generated reference runner source, build tree, and shared library. |
 | `simulation/` | Functional simulator outputs used for numerical comparison. |
 | `pimcomp/` | PIMCOMP graph, instruction, simulator, and comparison-report artifacts. |

validation/networks/pimcomp_models/README.md

@@ -4,7 +4,8 @@ This directory contains the four networks evaluated in
 [PIMCOMP: An End-to-End DNN Compiler for Processing-In-Memory Accelerators](https://arxiv.org/pdf/2411.09159):
 VGG-8, ResNet-18, ResNet-34, and GoogLeNet.
 
-See [RESULTS.md](RESULTS.md) for the current latency-only result status.
+See the runner-generated [results.csv](results.csv) for the current latency
+and energy results.
 
 ## Models and provenance
 
@@ -15,11 +16,11 @@ See [RESULTS.md](RESULTS.md) for the current latency-only result status.
 | `googlenet/` | GoogLeNet | `1x3x224x224` | Unmodified [ONNX Model Zoo `googlenet-12`](https://huggingface.co/onnxmodelzoo/googlenet-12). |
 | `vgg8/` | VGG-8 reconstruction | `1x1x28x28` | Deterministic compiler workload with six convolution and two fully connected layers. |
 
-`googlenet/googlenet-12-no-softmax.onnx` is a derived latency model that
-exposes the original model's final FC logits (`loss3/classifier_1`) as its
-output. This matches PIMCOMP's instruction stream, which records but does not
-schedule the terminal `OP_SOFTMAX`. Keep `googlenet-12.onnx` for full-model
-functional validation.
+`googlenet/googlenet-12-latency.onnx` is the explicit common latency model.
+It exposes the original model's final FC logits (`loss3/classifier_1`) and
+removes its two LRN nodes and terminal Softmax so the comparison covers only
+operations scheduled by PIMCOMP. Keep `googlenet-12.onnx` as the unmodified
+source model.
 
 The PIMCOMP authors did not publish the ONNX checkpoints used by the paper.
 Running PIMCOMP's frontend on the three Model Zoo files above produces JSON
@@ -38,7 +39,7 @@ Current SHA-256 checksums:
 788088b908e233d924c7c26b997e89ee861290c7bc56783a306e8201d79aac8f  resnet18/resnet18-v1-7.onnx
 c3231061d081bdd47884137b02134f85142752a39e87263c529cd14ed242b096  resnet34/resnet34-v1-7.onnx
 c99c507058eaf41de8723408fdda7db8325cb57f0a89f2ee07a716d6e963e14e  googlenet/googlenet-12.onnx
-a35bad96441efbee28699cb61d1656cca7f7281f14040cf01699c3d0cfd8b202  googlenet/googlenet-12-no-softmax.onnx
+a26f9e33901c573e60c34a3f0abbb4744fff83e4e0f21b18fc66e20395e72982  googlenet/googlenet-12-latency.onnx
 396cdea21e5e7d02c3f26f14d22ef20975171702493f5c5e79b8e0d896e541ef  vgg8/vgg8-mnist-reconstructed.onnx
 ```
 
@@ -145,17 +146,17 @@ uses the fixed seed `1`, so repeated serial and parallel runs are reproducible.
 ## Compare Raptor and PIMCOMP
 
 The comparison driver uses one random input and one native ONNX-MLIR reference,
-compiles both instruction streams, runs both through `pimsim-nn`, validates
-Raptor through the Rust simulator, and writes Markdown and JSON reports.
-PIMCOMP Rust validation also runs when its optional exporter is available.
+compiles both instruction streams, runs both through `pimsim-nn`, runs
+functional validation through `pim-simulator`, and writes Markdown and JSON
+reports.
 
 To reproduce the complete Arch-A latency experiment, use the model-by-model
 runner. It verifies the paper GA settings, builds Raptor and the existing
 `third_party/PIMCOMP-NN/build` tree, then runs the `element`/batch-1 comparison
-for one model at a time:
+for one model at a time and regenerates `results.csv` from the JSON reports:
 
 ```bash
-.venv/bin/python validation/tools/run_pimcomp_paper_latency.py
+.venv/bin/python validation/tools/pimcomp/run_pimcomp_paper_latency.py
 ```
 
 Each model directory reuses regular validation's ignored `inputs/`, `outputs/`,
@@ -176,7 +177,7 @@ Arch-A low-latency example:
 ```bash
 RAPTOR_ROOT=$PWD
 
-"$RAPTOR_ROOT/.venv/bin/python" "$RAPTOR_ROOT/validation/tools/compare_raptor_pimcomp.py" \
+"$RAPTOR_ROOT/.venv/bin/python" "$RAPTOR_ROOT/validation/tools/pimcomp/compare_raptor_pimcomp.py" \
   --model "$RAPTOR_ROOT/validation/networks/pimcomp_models/resnet34/resnet34-v1-7.onnx" \
   --out-dir "$RAPTOR_ROOT/validation/networks/pimcomp_models/resnet34" \
   --pimcomp-config "$RAPTOR_ROOT/validation/pimsim_configs/pimcomp/arch-a/latency_config.json" \
@@ -202,12 +203,36 @@ The comparison runner enables `--fail-on-error`, so a failed compiler,
 simulation, or semantic validation makes the command fail while preserving the
 generated report.
 
+### Numeric precision and simulator artifacts
+
+The functional and non-functional simulators intentionally consume different
+artifacts:
+
+- Raptor and PIMCOMP are validated against the native ONNX-MLIR reference as
+  FP32 programs in the Rust simulator. Raptor's emitted program is already
+  FP32. The PIMCOMP-to-Rust export expands its element-addressed storage and
+  byte-sized transfers to FP32, emits `setbw 32, 32`, and keeps vector
+  `imm_len` fields as element counts.
+- PIMCOMP's original `SimulationInfo.gz` is copied unchanged for `pimsim-nn`.
+  PIMCOMP hardcodes `setbw 8, 8` and one byte per element without performing
+  numerical quantization; this artifact is used only for latency estimation.
+- Raptor's original FP32 artifact remains unchanged for functional validation.
+  A separate `raptor/pimsim_nn/` view uses `setbw 8, 8` and scales its
+  byte-addressed storage and transfer sizes from four bytes to one byte per
+  element. Vector `imm_len` fields remain element counts. Like PIMCOMP's
+  artifact, this view is not numerically valid and is used only for a fair
+  non-functional comparison.
+
+The ISA defines vector lengths in elements, while `ld`, `st`, `lldi`, `lmv`,
+`send`, `recv`, addresses, and non-vector offsets are byte-based. Do not use
+either latency-only artifact for semantic validation.
+
 Current Raptor status:
 
 - VGG-8, ResNet-18, fixed-batch ResNet-34, and GoogLeNet compile on Arch-A.
-- Use `googlenet-12-no-softmax.onnx` for the paper-matched latency comparison.
-  The original model's final `vsoftmax` is supported by Raptor's functional
-  simulator but not by `pimsim-nn`; PIMCOMP does not schedule that operation.
+- Use `googlenet-12-latency.onnx` for the paper-matched latency comparison.
+  It removes the two LRN nodes and terminal softmax that PIMCOMP does not
+  schedule.
 - Raptor currently accepts one square `--crossbar-size`; Arch-C's rectangular
   `512x1024` arrays can therefore be compiled by PIMCOMP but not compared
   exactly with Raptor.
@@ -227,10 +252,8 @@ rsync -azL validation/networks/pimcomp_models/ \
   "monolith:$REMOTE_REPO/validation/networks/pimcomp_models/"
 rsync -az validation/pimsim_configs/pimcomp/ \
   "monolith:$REMOTE_REPO/validation/pimsim_configs/pimcomp/"
-rsync -az validation/tools/compare_raptor_pimcomp.py \
-  "monolith:$REMOTE_REPO/validation/tools/compare_raptor_pimcomp.py"
-rsync -az validation/tools/run_pimcomp_paper_latency.py \
-  "monolith:$REMOTE_REPO/validation/tools/run_pimcomp_paper_latency.py"
+rsync -az validation/tools/pimcomp/ \
+  "monolith:$REMOTE_REPO/validation/tools/pimcomp/"
 rsync -az --exclude=.git --exclude=build --exclude=output \
   third_party/PIMCOMP-NN/ \
   "monolith:$REMOTE_REPO/third_party/PIMCOMP-NN/"
@@ -247,7 +270,7 @@ python3 -m venv .venv
 .venv/bin/python -m pip install numpy onnx onnxruntime onnxsim colorama
 
 # Run every latency comparison serially.
-.venv/bin/python validation/tools/run_pimcomp_paper_latency.py
+.venv/bin/python validation/tools/pimcomp/run_pimcomp_paper_latency.py
 ```
 
 Copy reports back without transferring large compiler artifacts:

validation/networks/pimcomp_models/RESULTS.md

@@ -1,32 +0,0 @@
-# Raptor vs PIMCOMP latency results
-
-## GoogLeNet, Arch-A, low latency
-
-Measured with `googlenet-12-no-softmax.onnx`, batch 1, Raptor's best current
-schedule, and PIMCOMP's GA/element artifacts. Both instruction streams were
-simulated by the same `pimsim-nn` build using the complete Arch-A timing and
-precision configuration.
-
-| Compiler | Latency (ms) | Instructions | Sends | Receives | MVMUL |
-| --- | ---: | ---: | ---: | ---: | ---: |
-| Raptor | 1132.458315 | 77,717,248 | 29,115 | 29,115 | 157,158 |
-| PIMCOMP | 41.790450 | 3,398,070 | 110,334 | 110,334 | 113,639 |
-
-PIMCOMP is 27.10x faster in this latency simulation.
-
-Semantic validation did not pass the comparison driver's strict default
-tolerance: the maximum logit differences from the native ONNX reference were
-`0.01995039` for Raptor and `7.768404` for PIMCOMP. Treat these as performance
-results, not as a correctness-equivalent comparison.
-
-No throughput experiment was run.
-
-## Reproduce
-
-```bash
-.venv/bin/python validation/tools/run_pimcomp_paper_latency.py \
-  --models googlenet
-```
-
-See [README.md](README.md) for model provenance, limitations, and monolith
-instructions.

validation/networks/pimcomp_models/results.csv

@@ -0,0 +1,5 @@
+model,raptor_latency_ms,pimcomp_latency_ms,raptor_energy_pj,pimcomp_energy_pj,faster_compiler,speedup
+vgg8,2.451664,7.985074,647608825.120001,1597904071.120000,raptor,3.26
+resnet18,113.483357,58.853733,23204566467.119949,13983168748.119974,pimcomp,1.93
+resnet34,120.655292,91.607980,27614808203.679939,22722922369.680016,pimcomp,1.32
+googlenet,22.936609,62.923463,7798109358.239990,14547526780.240000,raptor,2.74

validation/raptor_validation/pimsim_nn.py

@@ -0,0 +1,73 @@
+import json
+import re
+import shutil
+from pathlib import Path
+
+
+_METRIC_PATTERNS = {
+    "output_count": r"output count:\s+([0-9]+)\s+samples",
+    "throughput": r"throughput:\s+([0-9.eE+-]+)\s+samples/s",
+    "average_latency_ms": r"average latency:\s+([0-9.eE+-]+)\s+ms",
+    "latency_ms": r"latency:\s+([0-9.eE+-]+)\s+ms",
+    "average_power_mw": r"average power:\s+([0-9.eE+-]+)\s+mW",
+    "average_energy_pj": r"average energy:\s+([0-9.eE+-]+)\s+pJ(?:/it)?",
+}
+
+
+def parse_pimsim_nn_metrics(output):
+    metrics = {"raw_output": output}
+    for name, pattern in _METRIC_PATTERNS.items():
+        match = re.search(pattern, output)
+        if match:
+            value = match.group(1)
+            metrics[name] = int(value) if name == "output_count" else float(value)
+    return metrics
+
+
+def export_raptor_latency_artifact(pim_dir, output_dir):
+    pim_dir = Path(pim_dir)
+    output_dir = Path(output_dir)
+    if output_dir.exists():
+        shutil.rmtree(output_dir)
+    output_dir.mkdir(parents=True)
+
+    def int8_bytes(value, field):
+        if value % 4:
+            raise ValueError(f"Raptor {field}={value} is not aligned to its fp32 element width")
+        return value // 4
+
+    with open(pim_dir / "config.json", encoding="utf-8") as f:
+        config = json.load(f)
+    for field in ("inputs_addresses", "outputs_addresses"):
+        if field in config:
+            config[field] = [int8_bytes(value, field) for value in config[field]]
+    with open(output_dir / "config.json", "w", encoding="utf-8") as f:
+        json.dump(config, f, separators=(",", ":"))
+        f.write("\n")
+
+    byte_size_fields = {
+        "ld": "size",
+        "st": "size",
+        "lldi": "len",
+        "lmv": "len",
+        "send": "size",
+        "recv": "size",
+    }
+    for source in sorted(pim_dir.glob("core_*.json"), key=lambda path: int(path.stem.split("_")[1])):
+        with open(source, encoding="utf-8") as f:
+            instructions = json.load(f)
+        for instruction in instructions:
+            op = instruction["op"]
+            if op == "setbw":
+                instruction["ibiw"] = 8
+                instruction["obiw"] = 8
+            elif op == "sldi":
+                instruction["imm"] = int8_bytes(instruction["imm"], "address")
+            if field := byte_size_fields.get(op):
+                instruction[field] = int8_bytes(instruction[field], f"{op} {field}")
+                if offset := instruction.get("offset"):
+                    offset["offset_value"] = int8_bytes(offset["offset_value"], f"{op} offset")
+        with open(output_dir / source.name, "w", encoding="utf-8") as f:
+            json.dump(instructions, f, separators=(",", ":"))
+            f.write("\n")
+    return output_dir

validation/raptor_validation/validate_one.py

@@ -1,6 +1,5 @@
 import json
 import os
-import re
 import shutil
 import subprocess
 import sys
@@ -11,6 +10,7 @@ from colorama import Style, Fore
 from .gen_network_runner import gen_network_runner
 from .onnx_utils import gen_random_inputs, save_inputs_to_files, onnx_io, write_inputs_to_memory_bin, _ONNX_TO_NP
 from .raptor import compile_with_raptor
+from .pimsim_nn import export_raptor_latency_artifact, parse_pimsim_nn_metrics
 from .subprocess_utils import run_command_with_reporter
 
 STAGE_TITLES = (
@@ -61,9 +61,7 @@ PIMSIM_FAILED = "ERROR"
 PIMSIM_UNSUPPORTED = "UNSUPPORTED"
 PIMSIM_SKIPPED = "SKIP"
 PIMSIM_NOT_RUN = "-"
-PIMSIM_UNSUPPORTED_VSOFTMAX = "pimsim-nn does not support binary opcode vsoftmax"
-PIMSIM_LATENCY_RE = re.compile(r"\blatency:\s+([0-9.eE+-]+)\s+ms")
-PIMSIM_POWER_RE = re.compile(r"\baverage power:\s+([0-9.eE+-]+)\s+mW")
+PIMSIM_UNSUPPORTED_VSOFTMAX = "pimsim-nn does not support opcode vsoftmax"
 
 
 class PimSimUnsupportedError(RuntimeError):
@@ -80,6 +78,7 @@ class ValidationResult:
     pim_pass_timings: dict[str, float] = field(default_factory=dict)
     pimsim_latency_ms: float | None = None
     pimsim_power_mw: float | None = None
+    pimsim_energy_pj: float | None = None
     pimsim_status: str = PIMSIM_SKIPPED
 
 
@@ -262,9 +261,10 @@ def pimcomp_compatibility_errors(config_path, *, core_count, crossbar_count, cro
 
 
 def run_pimsim_nn(pimsim_nn_build_dir, pim_dir, config_path, reporter=None, timeout_sec=None):
+    latency_artifact = export_raptor_latency_artifact(pim_dir, Path(pim_dir).parent / "pimsim_nn")
     try:
         output = run_command(
-            [pimsim_nn_build_dir / "ChipTest", pim_dir, config_path, "--gui=false"],
+            [pimsim_nn_build_dir / "ChipTest", latency_artifact, config_path, "--gui=false"],
             cwd=pimsim_nn_build_dir,
             reporter=reporter,
             timeout_sec=timeout_sec,
@@ -275,11 +275,11 @@ def run_pimsim_nn(pimsim_nn_build_dir, pim_dir, config_path, reporter=None, time
         if PIMSIM_UNSUPPORTED_VSOFTMAX in error_output:
             raise PimSimUnsupportedError(PIMSIM_UNSUPPORTED_VSOFTMAX) from exc
         raise
-    latency_match = PIMSIM_LATENCY_RE.search(output)
-    power_match = PIMSIM_POWER_RE.search(output)
-    if not latency_match or not power_match:
-        raise RuntimeError("pimsim-nn output did not contain latency and average power")
-    return float(latency_match.group(1)), float(power_match.group(1))
+    metrics = parse_pimsim_nn_metrics(output)
+    required = ("latency_ms", "average_power_mw", "average_energy_pj")
+    if any(name not in metrics for name in required):
+        raise RuntimeError("pimsim-nn output did not contain latency, average power, and average energy")
+    return tuple(metrics[name] for name in required)
 
 
 def clean_workspace_artifacts(workspace_dir, model_stem):
@@ -416,6 +416,8 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
         pimsim_nn_build_dir = Path(pimsim_nn_build_dir).resolve()
         pimsim_config_path = Path(pimsim_config_path).resolve()
     compile_extra_args = list(raptor_extra_args or [])
+    if pimsim_enabled and "--pim-emit-json" not in compile_extra_args:
+        compile_extra_args.append("--pim-emit-json")
     owns_reporter = reporter is None
     reporter = reporter or ProgressReporter(model_total, stages_per_model=len(MODE_STAGE_TITLES[mode]), verbose=verbose)
 
@@ -540,10 +542,11 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
             "Run Non-functional Simulation")
         pimsim_latency_ms = None
         pimsim_power_mw = None
+        pimsim_energy_pj = None
         pimsim_status = PIMSIM_SKIPPED
         if pimsim_enabled:
             try:
-                pimsim_latency_ms, pimsim_power_mw = run_pimsim_nn(
+                pimsim_latency_ms, pimsim_power_mw, pimsim_energy_pj = run_pimsim_nn(
                     pimsim_nn_build_dir,
                     pim_dir,
                     pimsim_config_path,
@@ -554,7 +557,8 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
                 print_info(
                     reporter,
                     f"Latency: {pimsim_latency_ms:.6f} ms, "
-                    f"Power: {pimsim_power_mw:.6f} mW")
+                    f"Power: {pimsim_power_mw:.6f} mW, "
+                    f"Energy: {pimsim_energy_pj:.6f} pJ")
             except PimSimUnsupportedError as exc:
                 pimsim_status = PIMSIM_UNSUPPORTED
                 print_info(reporter, str(exc))
@@ -581,6 +585,7 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
             pim_pass_timings=pim_pass_timings,
             pimsim_latency_ms=pimsim_latency_ms,
             pimsim_power_mw=pimsim_power_mw,
+            pimsim_energy_pj=pimsim_energy_pj,
             pimsim_status=pimsim_status,
         )
     except Exception:

validation/tools/pimcomp/compare_raptor_pimcomp.py

@@ -24,9 +24,10 @@ import onnx
 from colorama import Fore, Style
 
 
-REPO = Path(__file__).resolve().parents[2]
+REPO = Path(__file__).resolve().parents[3]
 VALIDATION_DIR = REPO / "validation"
 PIMSIM_CONFIG_DIR = VALIDATION_DIR / "pimsim_configs/pimcomp"
+PIMCOMP_OUTPUT_FILES = ("SimulationInfo.gz", "VerificationInfo.json", "MappingResult.txt")
 sys.path.insert(0, str(VALIDATION_DIR))
 
 from raptor_validation.gen_network_runner import gen_network_runner  # noqa: E402
@@ -38,6 +39,10 @@ from raptor_validation.onnx_utils import (  # noqa: E402
     write_inputs_to_memory_bin,
 )
 from raptor_validation.raptor import compile_with_raptor  # noqa: E402
+from raptor_validation.pimsim_nn import (  # noqa: E402
+    export_raptor_latency_artifact,
+    parse_pimsim_nn_metrics,
+)
 from raptor_validation.validate_one import (  # noqa: E402
     STAGE_COLORS,
     build_dump_ranges,
@@ -70,7 +75,12 @@ def load_pimcomp_exporter():
     module = importlib.util.module_from_spec(spec)
     assert spec is not None and spec.loader is not None
     sys.modules.setdefault("cv2", types.ModuleType("cv2"))
-    spec.loader.exec_module(module)
+    write_bytecode = sys.dont_write_bytecode
+    sys.dont_write_bytecode = True
+    try:
+        spec.loader.exec_module(module)
+    finally:
+        sys.dont_write_bytecode = write_bytecode
     return module
 
 
@@ -183,37 +193,21 @@ def run_logged(
     return proc.stdout
 
 
-def remove_tree(path: Path) -> None:
-    if not path.exists() and not path.is_symlink():
-        return
-    if path.is_symlink() or path.is_file():
-        path.unlink()
-        return
-    while True:
-        children = list(path.iterdir())
-        if not children:
-            break
-        for child in children:
-            remove_tree(child)
-    path.rmdir()
-
-
 def load_model_inputs(model_path: Path, seed: int):
     inputs_desc, outputs_desc = onnx_io(model_path)
     arrays_in_order, _ = gen_random_inputs(inputs_desc, seed=seed)
-    return inputs_desc, outputs_desc, arrays_in_order, arrays_in_order
+    return inputs_desc, outputs_desc, arrays_in_order
 
 
 def load_saved_inputs(
-    model_path: Path,
     inputs_desc: list[tuple[int, str, int, list[int]]],
     inputs_dir: Path,
-) -> tuple[list[np.ndarray], list[np.ndarray]]:
+) -> list[np.ndarray]:
     arrays = []
     for idx, name, elem_type, shape in inputs_desc:
         array = np.loadtxt(inputs_dir / f"in{idx}.csv", delimiter=",", dtype=_ONNX_TO_NP[elem_type]).reshape(shape)
         arrays.append(array)
-    return arrays, arrays
+    return arrays
 
 
 def prepare_pimcomp_model(model_path: Path, out_dir: Path) -> Path:
@@ -328,7 +322,7 @@ def compile_reference(
     runner_base = runner_dir / stem
 
     run_logged(
-        "Reference Emit ONNX IR",
+        "Compile Reference ONNX IR",
         [str(args.raptor_path), str(model_path), "-o", str(onnx_ir_base), "--EmitONNXIR",
          "--mlir-elide-elementsattrs-if-larger=16", "--enable-conv-opt-pass=false"],
         cwd=REPO,
@@ -337,7 +331,7 @@ def compile_reference(
         stage="Compile ONNX",
     )
     run_logged(
-        "Reference Native Compile",
+        "Compile Reference Native",
         [str(args.raptor_path), "-O3", str(model_path), "-o", str(runner_base)],
         cwd=REPO,
         timeout_sec=args.timeout_seconds,
@@ -416,13 +410,18 @@ def compile_raptor_target(
         f"--crossbar-size={hardware['crossbar_size']}",
         f"--crossbar-count={hardware['crossbar_count']}",
         f"--core-count={hardware['core_count']}",
+        f"--pim-target-config={args.pimcomp_config}",
         "--pim-emit-json",
         *args.raptor_extra_arg,
     ]
     print_step("Compile Raptor PIM", cmd, REPO, "Compile PIM")
     start = time.perf_counter()
     command = shell_join(cmd)
-    raptor_extra_args = ["--pim-emit-json", *args.raptor_extra_arg]
+    raptor_extra_args = [
+        f"--pim-target-config={args.pimcomp_config}",
+        "--pim-emit-json",
+        *args.raptor_extra_arg,
+    ]
     try:
         timings = compile_with_raptor(
             model_path,
@@ -459,7 +458,7 @@ def compile_raptor_target(
     return out_dir / "pim", timings
 
 
-def run_rust_validation(
+def run_functional_validation(
     label: str,
     pim_dir: Path,
     config_path: Path,
@@ -510,7 +509,7 @@ def run_rust_validation(
 
 def copy_pimcomp_outputs(source_dir: Path, out_dir: Path):
     out_dir.mkdir(parents=True, exist_ok=True)
-    for name in ("SimulationInfo.gz", "VerificationInfo.json", "MappingResult.txt"):
+    for name in PIMCOMP_OUTPUT_FILES:
         shutil.copy2(source_dir / name, out_dir / name)
 
 
@@ -527,7 +526,7 @@ def compile_pimcomp(
         shutil.copy2(args.pimcomp_config, runtime_config)
     pimcomp_output_dir = out_dir / "output"
     pimcomp_output_dir.mkdir(parents=True, exist_ok=True)
-    for name in ("SimulationInfo.gz", "VerificationInfo.json", "MappingResult.txt"):
+    for name in PIMCOMP_OUTPUT_FILES:
         (pimcomp_output_dir / name).unlink(missing_ok=True)
     model_name = args.pimcomp_model_name or f"compare_{model_path.stem}"
     frontend_json = frontend_json_dir / f"{model_name}.json"
@@ -540,7 +539,7 @@ def compile_pimcomp(
         str(frontend_json),
     ]
     run_logged(
-        "PIMCOMP Frontend",
+        "Compile PIMCOMP Frontend",
         frontend_cmd,
         cwd=args.pimcomp_dir / "frontend",
         timeout_sec=args.timeout_seconds,
@@ -556,20 +555,20 @@ def compile_pimcomp(
         "-s=YES",
     ]
     run_logged(
-        "PIMCOMP Backend",
+        "Compile PIMCOMP Backend",
         backend_cmd,
         cwd=frontend_json_dir.parent,
         timeout_sec=args.timeout_seconds,
         steps=steps,
         stage="Compile PIM",
     )
-    remove_tree(frontend_json_dir.parent)
+    shutil.rmtree(frontend_json_dir.parent)
     return pimcomp_output_dir / "VerificationInfo.json", pimcomp_output_dir / "SimulationInfo.gz"
 
 
 def export_pimcomp_for_pimsim_nn(simulation_info: Path, output_dir: Path) -> Path:
     if output_dir.exists():
-        remove_tree(output_dir)
+        shutil.rmtree(output_dir)
     with gzip.open(simulation_info, "rt", encoding="utf-8") as f:
         sim_info = json.load(f)
 
@@ -598,7 +597,7 @@ def export_pimcomp_for_pimsim_nn(simulation_info: Path, output_dir: Path) -> Pat
 
     for core_idx in core_indices:
         core_key = f"core{core_idx}"
-        instructions = sim_info.get(core_key, []) or [{"op": "lldi", "rd": 0, "imm": 0, "len": 0}]
+        instructions = sim_info.get(core_key, [])
         with open(output_dir / f"core_{core_idx}.json", "w", encoding="utf-8") as f:
             json.dump(instructions, f, separators=(",", ":"))
             f.write("\n")
@@ -622,7 +621,7 @@ def export_pimcomp_for_rust(
     if len(runtime_inputs) != 1:
         raise ValueError("PIMCOMP export currently requires exactly one runtime input tensor")
     if output_dir.exists():
-        remove_tree(output_dir)
+        shutil.rmtree(output_dir)
     exporter = load_pimcomp_exporter()
     with open(verification_info, "r", encoding="utf-8") as f:
         final_info = json.load(f)
@@ -727,7 +726,7 @@ def export_pimcomp_for_rust(
         for sim_inst in sim_info.get(core_name, []) or []:
             op = sim_inst["op"]
             if op == "setbw":
-                instructions.append(sim_inst)
+                instructions.append({"op": "setbw", "ibiw": 32, "obiw": 32})
                 continue
             if op == "sldi":
                 translated = {"op": "sldi", "rd": sim_inst["rd"], "imm": exporter.byte_offset(sim_inst["imm"])}
@@ -782,10 +781,12 @@ def export_pimcomp_for_rust(
                         "offset": sim_inst["offset"],
                     }
                 )
-            elif op in ("lmv", "vvadd", "vvmul", "vvmax", "vrelu"):
+            elif op == "lmv":
                 translated = dict(sim_inst)
                 translated["len"] = exporter.byte_offset(sim_inst["len"])
                 instructions.append(translated)
+            elif op in ("vvadd", "vvmul", "vvmax", "vrelu"):
+                instructions.append(sim_inst)
             elif op in ("send", "recv"):
                 translated = dict(sim_inst)
                 translated["size"] = exporter.byte_offset(sim_inst["size"])
@@ -822,24 +823,6 @@ def export_pimcomp_for_rust(
     return output_dir
 
 
-def parse_pimsim_nn_report(output: str) -> dict[str, float | int | str]:
-    patterns = {
-        "output_count": r"output count:\s+([0-9]+)\s+samples",
-        "throughput": r"throughput:\s+([0-9.eE+-]+)\s+samples/s",
-        "average_latency_ms": r"average latency:\s+([0-9.eE+-]+)\s+ms",
-        "latency_ms": r"latency:\s+([0-9.eE+-]+)\s+ms",
-        "average_power_mw": r"average power:\s+([0-9.eE+-]+)\s+mW",
-        "average_energy_pj": r"average energy:\s+([0-9.eE+-]+)\s+pJ/it",
-    }
-    result: dict[str, float | int | str] = {"raw_output": output}
-    for key, pattern in patterns.items():
-        match = re.search(pattern, output)
-        if match:
-            value = match.group(1)
-            result[key] = int(value) if key == "output_count" else float(value)
-    return result
-
-
 def run_pimsim_nn(
     label: str,
     inst_path: Path,
@@ -861,7 +844,7 @@ def run_pimsim_nn(
         steps=steps,
         stage="Run Non-functional Simulation",
     )
-    return parse_pimsim_nn_report(output)
+    return parse_pimsim_nn_metrics(output)
 
 
 def parse_raptor_instructions(pim_dir: Path) -> dict[str, Any]:
@@ -1069,7 +1052,7 @@ def write_report(
 
     lines.extend(
         [
-            "## Semantic Validation",
+            "## Functional Validation",
             "",
             f"- Raptor via `pim-simulator`: `{validation_status(raptor_validation)}`",
             f"- PIMCOMP via exported `pim-simulator`: `{validation_status(pimcomp_validation)}`",
@@ -1268,6 +1251,7 @@ def main():
     runner_path: Path | None = None
     reference_dir: Path | None = None
     raptor_pim_dir: Path | None = None
+    raptor_pimsim_dir: Path | None = None
     raptor_pass_timings: dict[str, float] = {}
     verification_info: Path | None = None
     simulation_info: Path | None = None
@@ -1289,7 +1273,8 @@ def main():
 
     model_io = try_stage(failures, "Load model inputs", load_model_inputs, model_path, args.seed)
     if model_io is not None:
-        inputs_desc, outputs_desc, arrays_in_order, runtime_inputs = model_io
+        inputs_desc, outputs_desc, arrays_in_order = model_io
+        runtime_inputs = arrays_in_order
 
     if reuse_raptor and model_io is not None:
         reuse_report_path = args.reuse_raptor_report.resolve()
@@ -1300,11 +1285,11 @@ def main():
             raise ValueError(f"Reused Raptor hardware differs: {reused_hardware} != {hardware}")
         reference_dir = Path(reused["paths"]["reference_outputs"])
         raptor_pim_dir = Path(reused["paths"]["raptor_pim"])
-        arrays_in_order, runtime_inputs = load_saved_inputs(
-            model_path,
+        arrays_in_order = load_saved_inputs(
             inputs_desc,
             reference_dir.parent / "inputs",
         )
+        runtime_inputs = arrays_in_order
         raptor_validation = CompareResult(**reused["raptor_validation"])
         raptor_perf = reused["raptor_performance"]
         raptor_instr = reused["raptor_instruction_summary"]
@@ -1392,9 +1377,9 @@ def main():
         if wrote_inputs and reference_dir is not None and outputs_desc:
             validation = try_stage(
                 failures,
-                "Rust Validation Raptor",
-                run_rust_validation,
-                "Rust Validation Raptor",
+                "Functional Validation Raptor",
+                run_functional_validation,
+                "Functional Validation Raptor",
                 raptor_pim_dir,
                 raptor_pim_dir / "config.json",
                 out_dir / "simulation/out.bin",
@@ -1451,7 +1436,7 @@ def main():
     if verification_info is not None and simulation_info is not None and model_io is not None:
         exported = try_stage(
             failures,
-            "Export PIMCOMP for Rust",
+            "Export PIMCOMP for Functional Validation",
             export_pimcomp_for_rust,
             pimcomp_model_path,
             verification_info,
@@ -1464,22 +1449,22 @@ def main():
     elif verification_info is None or simulation_info is None:
         record_failure(
             failures,
-            "Export PIMCOMP for Rust",
-            "PIMCOMP Rust export failed because PIMCOMP did not produce VerificationInfo.json and SimulationInfo.gz.",
+            "Export PIMCOMP for Functional Validation",
+            "PIMCOMP functional export failed because PIMCOMP did not produce VerificationInfo.json and SimulationInfo.gz.",
         )
     else:
         record_failure(
             failures,
-            "Export PIMCOMP for Rust",
-            "PIMCOMP Rust export failed because model inputs are not available.",
+            "Export PIMCOMP for Functional Validation",
+            "PIMCOMP functional export failed because model inputs are not available.",
         )
 
     if pimcomp_export_dir is not None and reference_dir is not None and outputs_desc:
         validation = try_stage(
             failures,
-            "Rust Validation PIMCOMP",
-            run_rust_validation,
-            "Rust Validation PIMCOMP",
+            "Functional Validation PIMCOMP",
+            run_functional_validation,
+            "Functional Validation PIMCOMP",
             pimcomp_export_dir,
             pimcomp_export_dir / "config.json",
             out_dir / "simulation/pimcomp.out.bin",
@@ -1491,7 +1476,7 @@ def main():
         )
         pimcomp_validation = validation if validation is not None else failed_validation("PIMCOMP validation failed")
     elif pimcomp_export_dir is None:
-        pimcomp_validation = failed_validation("PIMCOMP Rust export is not available")
+        pimcomp_validation = failed_validation("PIMCOMP functional export is not available")
     elif reference_dir is None:
         pimcomp_validation = failed_validation("Reference outputs are not available")
     else:
@@ -1524,17 +1509,27 @@ def main():
         pimcomp_perf = skipped_perf("pimsim-nn config is not available")
     else:
         if not reuse_raptor and raptor_pim_dir is not None:
-            perf = try_stage(
+            raptor_pimsim_dir = try_stage(
                 failures,
-                "pimsim-nn Raptor",
-                run_pimsim_nn,
-                "pimsim-nn Raptor",
+                "Export Raptor for pimsim-nn",
+                export_raptor_latency_artifact,
                 raptor_pim_dir,
-                pimsim_config,
-                steps,
-                args,
+                out_dir / "raptor/pimsim_nn",
             )
-            raptor_perf = perf if perf is not None else failed_perf("pimsim-nn Raptor failed")
+            if raptor_pimsim_dir is not None:
+                perf = try_stage(
+                    failures,
+                    "Non-Functional Simulation Raptor",
+                    run_pimsim_nn,
+                    "Non-Functional Simulation Raptor",
+                    raptor_pimsim_dir,
+                    pimsim_config,
+                    steps,
+                    args,
+                )
+                raptor_perf = perf if perf is not None else failed_perf("pimsim-nn Raptor failed")
+            else:
+                raptor_perf = failed_perf("Raptor pimsim-nn export failed")
         elif not reuse_raptor:
             raptor_perf = skipped_perf("Raptor PIM directory is not available")
 
@@ -1549,9 +1544,9 @@ def main():
             if pimcomp_pimsim_dir is not None:
                 perf = try_stage(
                     failures,
-                    "pimsim-nn PIMCOMP",
+                    "Non-Functional Simulation PIMCOMP",
                     run_pimsim_nn,
-                    "pimsim-nn PIMCOMP",
+                    "Non-Functional Simulation PIMCOMP",
                     pimcomp_pimsim_dir,
                     pimsim_config,
                     steps,
@@ -1613,6 +1608,7 @@ def main():
         "paths": {
             "reference_outputs": optional_path(reference_dir),
             "raptor_pim": optional_path(raptor_pim_dir),
+            "raptor_pimsim_nn": optional_path(raptor_pimsim_dir),
             "pimcomp_simulation_info": optional_path(simulation_info),
             "pimcomp_exported_pim": optional_path(pimcomp_export_dir),
             "pimsim_config": optional_path(pimsim_config),
@@ -1625,11 +1621,11 @@ def main():
         f.write("\n")
 
     failed_steps = any(step.status != "passed" for step in steps)
-    semantic_failure = any(
+    functional_failure = any(
         result.status == "done" and not result.passed
         for result in (raptor_validation, pimcomp_validation)
     )
-    failed = bool(failures or failed_steps or semantic_failure)
+    failed = bool(failures or failed_steps or functional_failure)
     result = "FAIL" if args.fail_on_error and failed else "DONE" if failed else "PASS"
     color = Fore.RED if result == "FAIL" else Fore.YELLOW if result == "DONE" else Fore.GREEN
     print("\n" + Style.BRIGHT + f"Result: {color}{result}" + Style.RESET_ALL)

validation/tools/pimcomp/run_pimcomp_paper_latency.py

@@ -0,0 +1,219 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+import shlex
+import subprocess
+import sys
+from pathlib import Path
+
+from colorama import Fore, Style
+
+
+REPO = Path(__file__).resolve().parents[3]
+SUITE = REPO / "validation/networks/pimcomp_models"
+sys.path.insert(0, str(REPO / "validation"))
+
+from raptor_validation.pimsim_nn import parse_pimsim_nn_metrics  # noqa: E402
+from raptor_validation.validate_one import STAGE_COLORS  # noqa: E402
+
+PIMCOMP_SOURCE = REPO / "third_party/PIMCOMP-NN"
+PIMCOMP_CONFIG = REPO / "validation/pimsim_configs/pimcomp/arch-a/latency_config.json"
+COMPARE = Path(__file__).resolve().with_name("compare_raptor_pimcomp.py")
+MODELS = {
+    "vgg8": SUITE / "vgg8/vgg8-mnist-reconstructed.onnx",
+    "resnet18": SUITE / "resnet18/resnet18-v1-7.onnx",
+    "resnet34": SUITE / "resnet34/resnet34-v1-7.onnx",
+    "googlenet": SUITE / "googlenet/googlenet-12-latency.onnx",
+}
+
+
+def result_dir(root: Path | None, name: str) -> Path:
+    return root / name if root is not None else MODELS[name].parent
+
+
+def write_results_csv(root: Path | None) -> Path:
+    output = (root or SUITE) / "results.csv"
+    fields = (
+        "model",
+        "raptor_latency_ms",
+        "pimcomp_latency_ms",
+        "raptor_energy_pj",
+        "pimcomp_energy_pj",
+        "faster_compiler",
+        "speedup",
+    )
+    rows = []
+    for name in MODELS:
+        report_path = result_dir(root, name) / "pimcomp/comparison_report.json"
+        if not report_path.exists():
+            continue
+        report = json.loads(report_path.read_text(encoding="utf-8"))
+        raptor = report.get("raptor_performance") or {}
+        pimcomp = report.get("pimcomp_performance") or {}
+        raptor_latency = raptor.get("latency_ms")
+        pimcomp_latency = pimcomp.get("latency_ms")
+        if raptor_latency is None or pimcomp_latency is None:
+            continue
+        raptor_energy = (raptor.get("average_energy_pj")
+                          or parse_pimsim_nn_metrics(raptor.get("raw_output", "")).get("average_energy_pj"))
+        pimcomp_energy = (pimcomp.get("average_energy_pj")
+                           or parse_pimsim_nn_metrics(pimcomp.get("raw_output", "")).get("average_energy_pj"))
+        faster = "raptor" if raptor_latency < pimcomp_latency else "pimcomp"
+        rows.append({
+            "model": name,
+            "raptor_latency_ms": f"{raptor_latency:.6f}",
+            "pimcomp_latency_ms": f"{pimcomp_latency:.6f}",
+            "raptor_energy_pj": "" if raptor_energy is None else f"{raptor_energy:.6f}",
+            "pimcomp_energy_pj": "" if pimcomp_energy is None else f"{pimcomp_energy:.6f}",
+            "faster_compiler": faster,
+            "speedup": f"{max(raptor_latency, pimcomp_latency) / min(raptor_latency, pimcomp_latency):.2f}",
+        })
+    with open(output, "w", newline="", encoding="utf-8") as f:
+        writer = csv.DictWriter(f, fieldnames=fields, lineterminator="\n")
+        writer.writeheader()
+        writer.writerows(rows)
+    return output
+
+
+def print_stage(title: str, color: str) -> None:
+    print("\n" + Style.BRIGHT + color + f"[{title}]" + Style.RESET_ALL, flush=True)
+
+
+def run(command: list[str], *, dry_run: bool, check: bool = True) -> int:
+    print(f"  cwd: {REPO}", flush=True)
+    print(f"  $ {shlex.join(command)}", flush=True)
+    if dry_run:
+        return 0
+    return subprocess.run(command, cwd=REPO, check=check).returncode
+
+
+def validate_pimcomp_source() -> None:
+    header = PIMCOMP_SOURCE / "backend/GeneticAlgorithm.h"
+    source = header.read_text(encoding="utf-8")
+    for setting in ("int population_num = 200;", "int max_iteration = 1000;"):
+        if setting not in source:
+            raise RuntimeError(f"PIMCOMP paper setting is missing: {setting}")
+
+
+def comparison_command(model: Path, result_dir: Path, timeout: float) -> list[str]:
+    return [
+        sys.executable,
+        str(COMPARE),
+        "--model",
+        str(model),
+        "--out-dir",
+        str(result_dir),
+        "--pimcomp-dir",
+        str(PIMCOMP_SOURCE),
+        "--pimcomp-config",
+        str(PIMCOMP_CONFIG),
+        "--core-count",
+        "168",
+        "--crossbar-count",
+        "96",
+        "--crossbar-size",
+        "128",
+        "--mesh-rows",
+        "12",
+        "--mesh-cols",
+        "14",
+        "--pimsim-mode",
+        "latency",
+        "--pimcomp-pipeline",
+        "element",
+        "--pimcomp-replication",
+        "GA",
+        "--timeout-seconds",
+        str(timeout),
+        "--fail-on-error",
+    ]
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(
+        description="Reproduce the serial Arch-A latency comparison from the PIMCOMP paper."
+    )
+    parser.add_argument(
+        "--out-dir",
+        type=Path,
+        help="Result root (default: artifacts beside each model under validation/).",
+    )
+    parser.add_argument("--models", nargs="+", choices=MODELS, default=list(MODELS))
+    parser.add_argument("--timeout-seconds", type=float, default=3600.0)
+    parser.add_argument(
+        "--resume",
+        action="store_true",
+        help="Skip models with a completed JSON report.",
+    )
+    parser.add_argument("--dry-run", action="store_true", help="Print commands without modifying files.")
+    args = parser.parse_args()
+
+    out_dir = args.out_dir.resolve() if args.out_dir is not None else None
+
+    missing = [str(MODELS[name]) for name in args.models if not MODELS[name].exists()]
+    if missing:
+        parser.error(f"missing model(s): {', '.join(missing)}")
+
+    validate_pimcomp_source()
+    if out_dir is not None and not args.dry_run:
+        out_dir.mkdir(parents=True, exist_ok=True)
+
+    print(Style.BRIGHT + f"Found {len(args.models)} PIMCOMP model(s) to compare." + Style.RESET_ALL)
+    print(f"Results root: {out_dir or SUITE}")
+    print("=" * 72)
+
+    print_stage("Build Raptor", STAGE_COLORS["Build Runner"])
+    run(["cmake", "--build", str(REPO / "build_release")], dry_run=args.dry_run)
+
+    print_stage("Build PIMCOMP", STAGE_COLORS["Build Runner"])
+    run(
+        ["cmake", "--build", str(PIMCOMP_SOURCE / "build"), "--target", "PIMCOMP-NN"],
+        dry_run=args.dry_run,
+    )
+
+    failed = []
+    for index, name in enumerate(args.models, start=1):
+        model_result_dir = result_dir(out_dir, name)
+        print(
+            "\n" + Fore.CYAN + f"[{index}/{len(args.models)}]" + Style.RESET_ALL
+            + f" {Style.BRIGHT}Comparing {name}{Style.RESET_ALL}",
+            flush=True,
+        )
+        if args.resume and (model_result_dir / "pimcomp/comparison_report.json").exists():
+            print(
+                Fore.YELLOW + "  Completed report exists; skipping" + Style.RESET_ALL,
+                flush=True,
+            )
+            continue
+        returncode = run(
+            comparison_command(MODELS[name], model_result_dir, args.timeout_seconds),
+            dry_run=args.dry_run,
+            check=False,
+        )
+        if returncode:
+            failed.append(name)
+
+    if args.dry_run:
+        return 1 if failed else 0
+
+    results_path = write_results_csv(out_dir)
+    print_stage("Results", STAGE_COLORS["Compare Outputs"])
+    print(results_path.read_text(encoding="utf-8"), end="")
+    print("\n" + Style.BRIGHT + Fore.CYAN + "Summary" + Style.RESET_ALL)
+    print(Style.BRIGHT + f"Passed: {len(args.models) - len(failed)}" + Style.RESET_ALL)
+    print(Style.BRIGHT + f"Failed: {len(failed)}" + Style.RESET_ALL)
+    print(Style.BRIGHT + f"Results: {results_path}" + Style.RESET_ALL)
+    if failed:
+        print(
+            Fore.RED + f"Failed comparisons: {', '.join(failed)}" + Style.RESET_ALL,
+            file=sys.stderr,
+        )
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

validation/tools/run_pimcomp_paper_latency.py

@@ -1,148 +0,0 @@
-#!/usr/bin/env python3
-from __future__ import annotations
-
-import argparse
-import shlex
-import subprocess
-import sys
-from pathlib import Path
-
-from colorama import Fore, Style
-
-
-REPO = Path(__file__).resolve().parents[2]
-SUITE = REPO / "validation/networks/pimcomp_models"
-PIMCOMP_SOURCE = REPO / "third_party/PIMCOMP-NN"
-PIMCOMP_CONFIG = REPO / "validation/pimsim_configs/pimcomp/arch-a/latency_config.json"
-COMPARE = REPO / "validation/tools/compare_raptor_pimcomp.py"
-MODELS = {
-    "vgg8": SUITE / "vgg8/vgg8-mnist-reconstructed.onnx",
-    "resnet18": SUITE / "resnet18/resnet18-v1-7.onnx",
-    "resnet34": SUITE / "resnet34/resnet34-v1-7.onnx",
-    "googlenet": SUITE / "googlenet/googlenet-12-no-softmax.onnx",
-}
-
-
-def run(command: list[str], *, dry_run: bool, check: bool = True) -> int:
-    print(Fore.CYAN + "$ " + Style.RESET_ALL + shlex.join(command), flush=True)
-    if dry_run:
-        return 0
-    return subprocess.run(command, cwd=REPO, check=check).returncode
-
-
-def validate_pimcomp_source() -> None:
-    header = PIMCOMP_SOURCE / "backend/GeneticAlgorithm.h"
-    source = header.read_text(encoding="utf-8")
-    for setting in ("int population_num = 200;", "int max_iteration = 1000;"):
-        if setting not in source:
-            raise RuntimeError(f"PIMCOMP paper setting is missing: {setting}")
-
-
-def comparison_command(model: Path, result_dir: Path, timeout: float) -> list[str]:
-    return [
-        sys.executable,
-        str(COMPARE),
-        "--model",
-        str(model),
-        "--out-dir",
-        str(result_dir),
-        "--pimcomp-dir",
-        str(PIMCOMP_SOURCE),
-        "--pimcomp-config",
-        str(PIMCOMP_CONFIG),
-        "--core-count",
-        "168",
-        "--crossbar-count",
-        "96",
-        "--crossbar-size",
-        "128",
-        "--mesh-rows",
-        "12",
-        "--mesh-cols",
-        "14",
-        "--pimsim-mode",
-        "latency",
-        "--pimcomp-pipeline",
-        "element",
-        "--pimcomp-replication",
-        "GA",
-        "--timeout-seconds",
-        str(timeout),
-        "--fail-on-error",
-    ]
-
-
-def main() -> int:
-    parser = argparse.ArgumentParser(
-        description="Reproduce the serial Arch-A latency comparison from the PIMCOMP paper."
-    )
-    parser.add_argument(
-        "--out-dir",
-        type=Path,
-        help="Result root (default: artifacts beside each model under validation/).",
-    )
-    parser.add_argument("--models", nargs="+", choices=MODELS, default=list(MODELS))
-    parser.add_argument("--timeout-seconds", type=float, default=3600.0)
-    parser.add_argument(
-        "--resume",
-        action="store_true",
-        help="Skip models with a completed JSON report.",
-    )
-    parser.add_argument("--dry-run", action="store_true", help="Print commands without modifying files.")
-    args = parser.parse_args()
-
-    out_dir = args.out_dir.resolve() if args.out_dir is not None else None
-
-    missing = [str(MODELS[name]) for name in args.models if not MODELS[name].exists()]
-    if missing:
-        parser.error(f"missing model(s): {', '.join(missing)}")
-
-    validate_pimcomp_source()
-    if out_dir is not None and not args.dry_run:
-        out_dir.mkdir(parents=True, exist_ok=True)
-
-    run(["cmake", "--build", str(REPO / "build_release")], dry_run=args.dry_run)
-    run(
-        ["cmake", "--build", str(PIMCOMP_SOURCE / "build"), "--target", "PIMCOMP-NN"],
-        dry_run=args.dry_run,
-    )
-
-    failed = []
-    for name in args.models:
-        result_dir = out_dir / name if out_dir is not None else MODELS[name].parent
-        if args.resume and (result_dir / "pimcomp/comparison_report.json").exists():
-            print(
-                Fore.YELLOW + f"[{name}] completed report exists; skipping" + Style.RESET_ALL,
-                flush=True,
-            )
-            continue
-        print(
-            "\n" + Fore.CYAN + f"[{name}]" + Style.RESET_ALL
-            + f" {Style.BRIGHT}Arch-A latency comparison{Style.RESET_ALL}",
-            flush=True,
-        )
-        returncode = run(
-            comparison_command(MODELS[name], result_dir, args.timeout_seconds),
-            dry_run=args.dry_run,
-            check=False,
-        )
-        if returncode:
-            failed.append(name)
-
-    if failed:
-        print(
-            "\n" + Style.BRIGHT + Fore.RED + "Result: FAIL" + Style.RESET_ALL,
-            file=sys.stderr,
-        )
-        print(
-            Fore.RED + f"Failed comparisons: {', '.join(failed)}" + Style.RESET_ALL,
-            file=sys.stderr,
-        )
-        return 1
-    if not args.dry_run:
-        print("\n" + Style.BRIGHT + f"Result: {Fore.GREEN}PASS" + Style.RESET_ALL)
-    return 0
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())

validation/validate.py

@@ -255,6 +255,10 @@ def main():
     total_timing_sum = 0.0
     timed_benchmark_count = 0
     reporter = ProgressReporter(len(onnx_files), stages_per_model=1, verbose=a.verbose)
+    raptor_extra_args = list(a.raptor_extra_arg)
+    if not any(str(arg).startswith("--pim-target-config=") for arg in raptor_extra_args):
+        raptor_extra_args.append(f"--pim-target-config={pimsim_config_path}")
+
     validation_kwargs = {
         "raptor_path": a.raptor_path,
         "onnx_include_dir": a.onnx_include_dir,
@@ -262,7 +266,7 @@ def main():
         "crossbar_size": a.crossbar_size,
         "crossbar_count": a.crossbar_count,
         "core_count": a.core_count,
-        "raptor_extra_args": a.raptor_extra_arg,
+        "raptor_extra_args": raptor_extra_args,
         "pimsim_nn_build_dir": pimsim_nn_build_dir,
         "pimsim_config_path": selected_pimsim_config,
         "command_timeout_seconds": a.command_timeout_seconds,
@@ -337,28 +341,31 @@ def main():
         rel: (
             format_pimsim_metric(result, result.pimsim_latency_ms, "ms"),
             format_pimsim_metric(result, result.pimsim_power_mw, "mW"),
+            format_pimsim_metric(result, result.pimsim_energy_pj, "pJ"),
         )
         for rel, result in results.items()
     }
     latency_width = max(len("Latency"), *(len(metrics[0]) for metrics in formatted_metrics.values()))
     power_width = max(len("Power"), *(len(metrics[1]) for metrics in formatted_metrics.values()))
+    energy_width = max(len("Energy"), *(len(metrics[2]) for metrics in formatted_metrics.values()))
     separator = (
         f"+-{'-' * path_width}-+-{'-' * status_width}-+-{'-' * latency_width}"
-        f"-+-{'-' * power_width}-+")
+        f"-+-{'-' * power_width}-+-{'-' * energy_width}-+")
     print(separator)
     print(
         f"| {'Operation'.ljust(path_width)} | {'Result'.ljust(status_width)} | "
-        f"{'Latency'.rjust(latency_width)} | {'Power'.rjust(power_width)} |"
+        f"{'Latency'.rjust(latency_width)} | {'Power'.rjust(power_width)} | "
+        f"{'Energy'.rjust(energy_width)} |"
     )
     print(separator)
     for rel, result in results.items():
         plain_status = "PASS" if result.passed else "FAIL"
         status = Fore.GREEN + plain_status.ljust(status_width) + Style.RESET_ALL if result.passed else \
             Fore.RED + plain_status.ljust(status_width) + Style.RESET_ALL
-        latency, power = formatted_metrics[rel]
+        latency, power, energy = formatted_metrics[rel]
         print(
             f"| {rel.ljust(path_width)} | {status} | {latency.rjust(latency_width)} | "
-            f"{power.rjust(power_width)} |")
+            f"{power.rjust(power_width)} | {energy.rjust(energy_width)} |")
     print(separator)
     print("\n" + Style.BRIGHT + Fore.CYAN + "Summary" + Style.RESET_ALL)
     print(Style.BRIGHT + f"Passed: {n_passed}" + Style.RESET_ALL)