Merge remote changes

fast spatial node merging with batch operations

.gitignore

@@ -1,5 +1,15 @@
+.zed
 .idea
 **/.vscode
+
 .claude
+.codex
 AGENTS.md
+
+CMakeUserPresets.json
+
 build
+cmake-build-debug
+cmake-build-release
+
+**/__pycache__

README.md

@@ -1,5 +1,159 @@
 # Raptor
 
+Raptor is a domain-specific MLIR compiler for neural networks (ONNX format)
+targeting in-memory computing / processing-in-memory (PIM) architectures.
+It progressively lowers ONNX-MLIR through a set of MLIR dialects down to
+target-specific artifacts (currently JSON code for the `pimsim-nn` simulator).
+
+## Overview
+
+PIM architectures perform most of the computation directly in memory.
+Raptor's first supported target is `pimsim-nn`, which simulates a chip with:
+- a shared host memory,
+- a number of cores that do most of the computation directly in their memory
+  (vector ops, vmm/mvm on ReRAM crossbars),
+- no branching instructions (branchless architecture) and no hardware loop
+  support — any repeated work (e.g. convolutions) must be unrolled into
+  explicit per-iteration instructions.
+
+Because of this, the amount of emitted instructions explodes quickly and the
+compiler must optimize aggressively at every stage to keep compilation
+tractable.
+
+A second target, `PulPim`, is planned for an accelerator with RISC-V cores
+each carrying its own in-memory computing unit and crossbars. It will live in
+a dedicated dialect (future work).
+
+### Targets and simulators
+
+`pimsim-nn` (under `backend-simulators/pim/pimsim-nn`) is used for
+**performance** estimates (latency, energy), but does not functionally execute
+the JSON code it consumes. To validate the numerical correctness of the JSON
+code produced by Raptor (or, for comparison, by the `pimcomp` compiler), we use
+a Rust simulator we maintain in-tree at
+`backend-simulators/pim/pim-simulator`.
+
+## Compilation pipeline
+
+The PIM-related sources live under `src/PIM` and the tests under `test/PIM`.
+When working on this codebase, most changes should stay confined to those
+trees (you only need to look outside, e.g. at `onnx-mlir` or `llvm`, for
+framework-level details).
+
+High-level lowering flow:
+
+```
+ONNX-MLIR ──► Spatial ──► Pim (tensor) ──► Pim (bufferized) ──► PIM JSON
+```
+
+1. **ONNX → Spatial** (`src/PIM/Conversion/ONNXToSpatial`).
+   Lowers ONNX ops into the `spat` dialect (`src/PIM/Dialect/Spatial`).
+   Spatial models a high-level spatial in-memory accelerator: vmm/mvm
+   operations are accelerated by storing a constant RHS matrix into a
+   crossbar. Crossbars cannot be re-programmed during execution, have a
+   limited fixed size, and there is a limited number of them per core.
+   Conversion patterns are split by op family under
+   `Conversion/ONNXToSpatial/Patterns/{Math,NN,Tensor}` (Conv, Gemm, MatMul,
+   Elementwise, ReduceMean, Pool, Relu, Sigmoid, Softmax, Concat, Gather,
+   Reshape, Resize, Split).
+
+2. **Spatial → Pim** (`src/PIM/Conversion/SpatialToPim`).
+   Lowers Spatial to the `pim` dialect (`src/PIM/Dialect/Pim`), which
+   materializes PIM cores (`pim.core`), inter-core communication
+   (`pim.send` / `pim.receive`), halts, and crossbar-level operations.
+
+3. **Merge compute nodes** (`src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes`).
+   A DCP-inspired heuristic (Dynamic Critical Path — see the original
+   scheduling paper by Kwok & Ahmad,
+   [DCP-eScience2007](https://clouds.cis.unimelb.edu.au/papers/DCP-eScience2007.pdf))
+   that coarsens the virtual node graph and decides how to group compute
+   nodes onto cores. Our implementation is only DCP-*inspired*: it is a
+   heuristic with different assumptions from the paper (different cost
+   model, constraints from crossbar capacity / core resources, and a
+   windowed coarsening loop instead of full-graph reprioritization). The
+   `dcp-critical-window-size` option controls how many lowest-slack virtual
+   nodes each coarsening iteration considers (0 = legacy full-graph
+   analysis). Related sources: `DCPGraph/DCPAnalysis.cpp`, `Graph.cpp/.hpp`,
+   `MergeComputeNodesPass.cpp`.
+
+4. **Bufferization** (`src/PIM/Dialect/Pim/Transforms/Bufferization`).
+   Converts tensor-semantics PIM IR into memref-semantics PIM IR using the
+   standard MLIR `BufferizableOpInterface` machinery
+   (`OpBufferizationInterfaces.*`, `PimBufferization.td`).
+
+5. **PIM code generation** (`src/PIM/Pass/PimCodegen`):
+   - `HostConstantFolding` — folds host-side constants.
+   - `MaterializeHostConstantsPass` — materializes the remaining host
+     constants for emission.
+   - `VerificationPass` — checks invariants before emission.
+   - `EmitPimJsonPass` — emits the final PIM JSON consumed by `pimsim-nn`
+     and `pim-simulator`.
+
+Supporting pieces:
+- `src/PIM/Compiler` — PIM-specific compiler options (crossbar size/count,
+  core count, DCP window, experimental conv impl, concat error handling, …)
+  and `PimCodeGen` entry points.
+- `src/PIM/Common` — shared utilities (`PimCommon`, `LabeledList`).
+- `src/PIM/Pass` — auxiliary passes (`MessagePass`, `CountInstructionPass`)
+  and the `PIMPasses.h` registry used by `PimAccelerator`.
+- `src/PIM/PimAccelerator.{cpp,hpp}` — accelerator entry point: registers
+  dialects, passes, and plugs Raptor into the ONNX-MLIR driver.
+
+## Key compiler options
+
+Pass these on the `onnx-mlir` command line when compiling for PIM:
+
+- `--maccel=PIM` — select the PIM accelerator.
+- `--EmitSpatial` / `--EmitPim` / `--EmitPimBufferized` / `--EmitPimCodegen`
+  — stop the pipeline at the requested stage (default: `EmitPimCodegen`).
+- `--pim-only-codegen` — assume the input is already bufferized PIM IR and
+  run only the codegen tail.
+- `--crossbar-size=<N>` / `--crossbar-count=<N>` — crossbar dimensions and
+  per-core count.
+- `--core-count=<N>` — number of cores (`-1` picks the minimum).
+- `--dcp-critical-window-size=<N>` — DCP coarsening window (0 = legacy).
+- `--use-experimental-conv-impl` — alternative convolution lowering.
+- `--ignore-concat-error` — soft-fail corner case in `ConcatOp`.
+
+## Validation
+
+Functional validation lives in `validation/` and drives the Rust
+`pim-simulator` to compare Raptor's output against a reference.
+
+Per-operation validation (from `validation/`):
+
+```
+validate.py \
+    --raptor-path ../cmake-build-release/Release/bin/onnx-mlir \
+    --onnx-include-dir ../onnx-mlir/include
+```
+
+End-to-end network validation (example: first 4 layers of YOLOv11n):
+
+```
+validate.py \
+    --raptor-path ../cmake-build-release/Release/bin/onnx-mlir \
+    --onnx-include-dir ../onnx-mlir/include \
+    --operations-dir ./networks/yolo11n/depth_04 \
+    --crossbar-size 2048 --crossbar-count 256
+```
+
+Available networks under `validation/networks/`: `vgg16`, `yolo11n`.
+Available operations under `validation/operations/`: `add`, `conv`, `div`,
+`gather`, `gemm`, `gemv`, `mul`, `pool`, `reduce_mean`, `relu`, `resize`,
+`sigmoid`, `softmax`, `split`.
+
+## Rebuilding
+
+Release build (fast):
+
+```
+cmake --build /home/nico/raptor/raptor/cmake-build-release --target onnx-mlir -j 30
+```
+
+A slower debug build is also available — configure it the same way but with
+`-DCMAKE_BUILD_TYPE=Debug` (see installation instructions below).
+
 ## Build
 
 ### Protobuf

backend-simulators/pim/pim-simulator/src/bin/pim-simulator/main.rs

@@ -1,10 +1,13 @@
-use anyhow::{bail, Context, Result};
+use anyhow::{Context, Result, bail};
 use clap::Parser;
 use glob::glob;
+use pimcore::cpu::crossbar::Crossbar;
 use pimcore::json_to_instruction::json_to_executor;
+use pimcore::memory_manager::CoreMemory;
 use pimcore::tracing::TRACER;
 use serde_json::Value;
-use std::fs;
+use std::collections::HashMap;
+use std::fs::{self, read_link};
 use std::io::Write;
 use std::path::PathBuf;
 
@@ -43,8 +46,10 @@ fn main() -> Result<()> {
     let config_json = retrive_config(&args)?;
     let core_jsons = retrive_cores(&args)?;
     let memory = retrive_memory(&args)?;
-    let mut executor = json_to_executor::json_to_executor(config_json, core_jsons.iter());
-    populate_crossbar(&args, &mut executor);
+    let global_crossbars = get_crossbars(&config_json, &args).unwrap();
+    let crossbars = map_crossbars_to_cores(&config_json, &args, &global_crossbars);
+    let mut executor =
+        json_to_executor::json_to_executor(config_json, core_jsons.iter(), crossbars);
     set_memory(&mut executor, memory);
     TRACER
         .lock()
@@ -55,46 +60,100 @@ fn main() -> Result<()> {
     Ok(())
 }
 
-fn populate_crossbar(args: &Args, executor: &mut pimcore::Executable) {
-    let num_cores = executor.cpu_mut().num_core();
+fn map_crossbars_to_cores<'c>(
+    config: &Value,
+    args: &Args,
+    global_crossbars: &'c HashMap<String, Crossbar>,
+) -> Vec<Vec<&'c Crossbar>> {
+    let mut res = Vec::new();
+    let num_cores = config.get("core_cnt").unwrap().as_i64().unwrap() as i32;
 
     if let Some(folder) = args.folder.as_ref() {
         for core_idx in 0..num_cores {
             let core_folder = folder.join(format!("core_{}", core_idx));
+            res.push(Vec::new());
 
             if !core_folder.is_dir() {
                 continue;
             }
 
-            let mut bin_files: Vec<(u32, std::path::PathBuf)> = std::fs::read_dir(&core_folder)
-                .expect("Failed to read core directory")
-                .filter_map(|entry| {
-                    let path = entry.ok()?.path();
-                    let file_name = path.file_name()?.to_str()?;
+            let mut sym_link_files: Vec<(u32, std::path::PathBuf)> =
+                std::fs::read_dir(&core_folder)
+                    .expect("Failed to read core directory")
+                    .filter_map(|entry| {
+                        let entry = entry.ok()?;
+                        assert!(entry.metadata().unwrap().is_symlink());
+                        let path = entry.path();
+                        let file_name = path.file_name()?.to_str()?;
 
-                    if file_name.starts_with("crossbar_") && file_name.ends_with(".bin") {
-                        let num_str = &file_name[9..file_name.len() - 4];
-                        let num = num_str.parse::<u32>().ok()?;
-                        Some((num, path))
-                    } else {
-                        None
-                    }
-                })
-                .collect();
-            bin_files.sort_by_key(|&(num, _)| num);
-            let core = executor.cpu_mut().core(core_idx);
-            let (_memory, crossbars) = core.get_memory_crossbar();
+                        if file_name.starts_with("crossbar_") && file_name.ends_with(".bin") {
+                            let num_str = &file_name[9..file_name.len() - 4];
+                            let num = num_str.parse::<u32>().ok()?;
+                            Some((num, path))
+                        } else {
+                            None
+                        }
+                    })
+                    .collect();
+            sym_link_files.sort_by_key(|&(num, _)| num);
 
-            for (i, path) in bin_files {
-                let bytes = std::fs::read(path).expect("Failed to read binary file");
-                crossbars
-                    .get_mut(i as usize)
+            for (_, symlink) in sym_link_files {
+                let real_path = read_link(symlink).unwrap();
+                let path_as_str = real_path.to_str().unwrap();
+                assert!(
+                    global_crossbars.contains_key(path_as_str),
+                    "symlink point to {:?}\n a not stored crossbar",
+                    real_path
+                );
+
+                res.iter_mut()
+                    .next_back()
                     .unwrap()
-                    .execute_store(&bytes)
-                    .unwrap();
+                    .push(global_crossbars.get(path_as_str).unwrap());
             }
         }
     }
+    res
+}
+
+fn get_crossbars(config: &Value, args: &Args) -> anyhow::Result<HashMap<String, Crossbar>> {
+    let xbar_size = config.get("xbar_size").unwrap().as_array().unwrap();
+    let rows_crossbar = xbar_size[0].as_i64().unwrap() as usize;
+    let column_corssbar = xbar_size[1].as_i64().unwrap() as usize;
+    let mut res = HashMap::new();
+
+    if let Some(folder) = args.folder.as_ref() {
+        let weight_folder = folder.join("weights");
+        if !weight_folder.is_dir() {
+            bail!("Not a directory");
+        }
+        for weight_file in
+            std::fs::read_dir(&weight_folder).context("Weight folder not iterable")?
+        {
+            let weight_file = weight_file.context("File not iterable")?;
+            if weight_file
+                .metadata()
+                .context("Doesn't contain metadata")?
+                .is_dir()
+            {
+                continue;
+            }
+
+            let bytes = std::fs::read(weight_file.path()).expect("Failed to read binary file");
+            let mut crossbar =
+                Crossbar::new(column_corssbar * 4, rows_crossbar, CoreMemory::new());
+            crossbar.execute_store(&bytes).unwrap();
+            res.insert(
+                weight_file
+                    .path()
+                    .to_str()
+                    .context("file name not utf-8")?
+                    .to_string(),
+                crossbar,
+            );
+        }
+    }
+    Ok(res)
 }
 
 fn dump_memory(mut executor: pimcore::Executable, args: &Args) -> Result<()> {
@@ -170,7 +229,11 @@ fn retrive_cores(args: &Args) -> Result<Vec<Value>, anyhow::Error> {
         let pattern_str = pattern.to_str().context("Invalid path encoding")?;
         let mut paths: Vec<_> = glob(pattern_str)?.map(|x| x.unwrap()).collect();
         paths.sort_by_cached_key(|x| {
-            let mut x = x.file_stem().expect("Extracting the stem").to_str().expect("File not utf-8");
+            let mut x = x
+                .file_stem()
+                .expect("Extracting the stem")
+                .to_str()
+                .expect("File not utf-8");
             x = &x[5..];
             x.parse::<i32>().unwrap()
         });

backend-simulators/pim/pim-simulator/src/lib/cpu/crossbar.rs

@@ -38,14 +38,14 @@ impl Crossbar {
         self.memory.execute_store(0, element)
     }
 
-    pub fn load<T>(&mut self, size: usize) -> Result<Vec<&[T]>> where
+    pub fn load<T>(&self, size: usize) -> Result<Vec<&[T]>> where
             T: MemoryStorable, {
         if self.memory.get_len() <  size 
             //|| self.stored_bytes < size 
         {
             bail!("Loading outside crossbar boundary [{} {}] <  {}", self.stored_bytes, self.memory.get_len() ,  size);
         }
-        self.memory.load(0, size)
+        self.memory.load_const(0, size)
     }
 
 

backend-simulators/pim/pim-simulator/src/lib/cpu/mod.rs

@@ -1,5 +1,5 @@
 use std::{collections::HashMap, fmt::Debug};
-use anyhow::{Context, Result};
+use anyhow::{Context, Result, ensure};
 
 use crate::{
     cpu::crossbar::Crossbar,
@@ -10,53 +10,44 @@ use crate::{
 pub mod crossbar;
 
 #[derive(Debug, Clone)]
-pub struct CPU {
-    cores: Box<[Core]>,
+pub struct CPU<'a> {
+    cores: Box<[Core<'a>]>,
 }
 
-impl CPU {
-    pub fn new(num_cores: impl TryToUsize) -> Self {
+impl<'a> CPU<'a> {
+    pub fn new(num_cores: impl TryToUsize, crossbars: Vec<Vec<&'a Crossbar>> ) -> Self {
         let num_cores = num_cores.try_into().expect("num_cores can not be negative");
-        let mut cores: Vec<Core> = std::iter::repeat_with(Core::new)
-            .take(num_cores + 1)
-            .collect();
+        assert!(crossbars.len() == num_cores + 1);
+        let mut cores = Vec::new();
+        for crossbar in crossbars {
+            cores.push(Core::new(crossbar));
+        }
         Self {
             cores: cores.into(),
         }
     }
 
-    pub fn reserve_crossbar(
-        &mut self,
-        num_crossbar: impl TryToUsize,
-        byte_width: impl TryToUsize,
-        height: impl TryToUsize,
-    ) {
-        let num_crossbar = num_crossbar
-            .try_into()
-            .expect("num_crossbar can not be negative");
-        let byte_width = byte_width
-            .try_into()
-            .expect("byte_width can not be negative");
-        let height = height.try_into().expect("height can not be negative");
-        for core in &mut self.cores {
-            core.reserve_crossbar(num_crossbar, byte_width, height);
-        }
+    pub fn host<'b>(&'b mut self) -> &'b mut Core<'a>
+        where 'a : 'b
+    {
+        & mut self.cores[0]
     }
 
-    pub fn host(&mut self) -> &mut Core {
-        &mut self.cores[0]
-    }
-
-    pub fn core(&mut self, index: impl TryToUsize) -> &mut Core {
+    pub fn core<'b >(&'b mut self, index: impl TryToUsize) -> &'b mut Core<'a>
+        where 'a : 'b
+    {
         let index = index.try_into().expect("can not be negative");
-        &mut self.cores[index]
+        & mut self.cores[index]
     }
 
     pub fn num_core(&self) -> usize {
         self.cores.len()
     }
 
-    pub(crate) fn host_and_cores(&mut self, core: impl TryToUsize) -> (&mut Core, &mut Core) {
+    pub(crate) fn host_and_cores<'b, 'c >(&'b mut self, core: impl TryToUsize) -> (&'c mut Core<'a>, &'c mut Core<'a>) 
+    where 'a: 'b,
+          'b: 'c
+    {
         let core = core.try_into().expect("core can not be negative");
         assert_ne!(
             core, 0,
@@ -70,45 +61,29 @@ impl CPU {
         (host, core)
     }
 
-    pub fn get_multiple_cores<const N: usize>(&mut self, indices: [usize; N]) -> [&mut Core; N] {
+    pub fn get_multiple_cores<'b, const N: usize>(&'b mut self, indices: [usize; N]) -> [&'b mut Core<'a>; N] 
+    where 'a : 'b
+    {
         self.cores.get_disjoint_mut(indices).unwrap()
     }
 }
 
 #[derive(Debug, Clone)]
-pub struct Core {
-    crossbars: Vec<Crossbar>,
+pub struct Core<'a> {
+    crossbars: Vec<&'a Crossbar>,
     memory: CoreMemory,
     registers: [i32; 32],
 }
 
-impl Core {
-    fn new() -> Self {
+impl<'a> Core<'a> {
+    fn new(crossbars : Vec<&'a Crossbar>) -> Self {
         Self {
-            crossbars: Vec::new(),
+            crossbars,
             memory: CoreMemory::new(),
             registers: [0; 32],
         }
     }
 
-    pub fn reserve_crossbar(
-        &mut self,
-        num_crossbar: impl TryToUsize,
-        width: impl TryToUsize,
-        height: impl TryToUsize,
-    ) {
-        let num_crossbar = num_crossbar
-            .try_into()
-            .expect("num_crossbar can not be negative");
-        let width = width.try_into().expect("width can not be negative");
-        let height = height.try_into().expect("height can not be negative");
-        for _ in 0..num_crossbar {
-            let mut crossbar = CoreMemory::new();
-            crossbar.set_capacity(width * height);
-            self.crossbars.push(Crossbar::new(width, height, crossbar));
-        }
-    }
-
     pub fn execute_load<T>(&mut self) -> Result<Vec<&[T]>>
     where
         T: MemoryStorable,
@@ -157,7 +132,7 @@ impl Core {
         self.memory.load(address, size)
     }
 
-    pub fn get_memory_crossbar(&mut self) -> (&mut CoreMemory, &mut Vec<Crossbar>) {
+    pub fn get_memory_crossbar(&mut self) -> (&mut CoreMemory, &mut Vec<&'a Crossbar>) {
         let Self {
             crossbars,
             memory,

backend-simulators/pim/pim-simulator/src/lib/instruction_set/isa.rs

@@ -76,7 +76,8 @@ pub fn functor_to_name(functor: usize) -> &'static str {
 ///////////////////////////////////////////////////////////////
 /////////////////Scalar/register Instructions//////////////////
 ///////////////////////////////////////////////////////////////
-pub fn sldi(cores: &mut CPU, data: InstructionData) -> Result<InstructionStatus> {
+pub fn sldi(cores: &mut CPU, data: InstructionData) -> Result<InstructionStatus> 
+{
     TRACER.lock().unwrap().pre_sldi(cores, data);
     let (core_indx, rd, imm) = data.get_core_rd_imm();
     let core = cores.core(core_indx);

backend-simulators/pim/pim-simulator/src/lib/instruction_set/mod.rs

@@ -40,7 +40,9 @@ impl Instruction {
         Self { data, functor }
     }
 
-    pub fn execute(&self, cpu: &mut CPU) -> InstructionStatus {
+    pub fn execute<'a, 'b>(&'b self, cpu: &mut CPU<'a>) -> InstructionStatus
+    where 'a : 'b
+    {
         (self.functor)(cpu, self.data)
             .with_context(|| format!("Instruction: {}", functor_to_name(self.functor as usize)))
             .with_context(|| format!("Error in core: {}", self.data.core_indx() - 1))

backend-simulators/pim/pim-simulator/src/lib/json_to_instruction/json_to_executor.rs

@@ -1,10 +1,11 @@
 use core::panic;
+use std::collections::HashMap;
 
 use serde_json::{Map, Value};
 
 use crate::{
     CoreInstructionsBuilder, Executable,
-    cpu::{CPU, crossbar},
+    cpu::{CPU, crossbar::{self, Crossbar}},
     instruction_set::{
         InstructionsBuilder,
         instruction_data::{self, InstructionData, InstructionDataBuilder},
@@ -13,18 +14,20 @@ use crate::{
     memory_manager::type_traits::TryToUsize,
 };
 
+
 pub fn json_to_executor<'a>(
     config: Value,
     mut cores: impl Iterator<Item = &'a Value>,
-) -> Executable {
+    crossbars : Vec<Vec<&'a Crossbar>>
+) -> Executable<'a> {
     let cell_precision = config.get("cell_precision").unwrap().as_i64().unwrap() as i32;
     let core_cnt = config.get("core_cnt").unwrap().as_i64().unwrap() as i32 - 1;
     let xbar_count = config.get("xbar_array_count").unwrap().as_i64().unwrap() as i32;
     let xbar_size = config.get("xbar_size").unwrap().as_array().unwrap();
     let rows_crossbar = xbar_size[0].as_i64().unwrap() as i32;
     let column_corssbar = xbar_size[1].as_i64().unwrap() as i32;
-    let mut cpu = CPU::new(core_cnt);
-    cpu.reserve_crossbar(xbar_count, column_corssbar * 4, rows_crossbar);
+
+    let mut cpu = CPU::new(core_cnt, crossbars);
     let mut core_insts_builder = CoreInstructionsBuilder::new(core_cnt as usize);
     cores.next();
     for core_indx in 1..=core_cnt {

backend-simulators/pim/pim-simulator/src/lib/memory_manager/mod.rs

@@ -111,6 +111,24 @@ where {
         Ok(res)
     }
 
+    pub fn load_const<T>(&self, address: impl TryToUsize, size: impl TryToUsize) -> Result<Vec<&[T]>>
+    where
+        T: MemoryStorable,
+    {
+        let address = address.try_into().expect("address can not be negative");
+        let size = size.try_into().expect("size can not be negative");
+        let Self {
+            memory,
+            load_requests,
+        } = self;
+        let mut res = Vec::new();
+        let memory_slice = &memory[address..address + size];
+        let memory_slice = unsafe { slice_from_u8(memory_slice) }
+            .with_context(|| format!("Accessing from {} to {}",  address, address + size))?;
+        res.push(memory_slice);
+        Ok(res)
+    }
+
     pub fn load<T>(&mut self, address: impl TryToUsize, size: impl TryToUsize) -> Result<Vec<&[T]>>
     where
         T: MemoryStorable,

backend-simulators/pim/pim-simulator/src/lib/memory_manager/type_traits.rs

@@ -55,15 +55,23 @@ pub trait HasSigm {
 
 impl HasSigm for f32 {
     fn sigm(self) -> Self {
-        let ex = self.exp();
-        ex / (1.0 + ex)
+        if self >= 0.0 {
+            1.0 / (1.0 + (-self).exp())
+        } else {
+            let ex = self.exp();
+            ex / (1.0 + ex)
+        }
     }
 }
 
 impl HasSigm for f64 {
     fn sigm(self) -> Self {
-        let ex = self.exp();
-        ex / (1.0 + ex)
+        if self >= 0.0 {
+            1.0 / (1.0 + (-self).exp())
+        } else {
+            let ex = self.exp();
+            ex / (1.0 + ex)
+        }
     }
 }
 

backend-simulators/pim/pim-simulator/src/lib/pimcore.rs

@@ -1,50 +1,54 @@
 #![allow(unused)]
 
+use std::time::{Duration, SystemTime};
+
 use crate::{
-    cpu::CPU, instruction_set::{Instruction, InstructionStatus, Instructions, isa::functor_to_name}, memory_manager::type_traits::TryToUsize, send_recv::{SendRecv, handle_send_recv}, tracing::TRACER
+    cpu::CPU,
+    instruction_set::{Instruction, InstructionStatus, Instructions, isa::functor_to_name},
+    memory_manager::type_traits::TryToUsize,
+    send_recv::{SendRecv, handle_send_recv},
+    tracing::TRACER,
 };
 pub mod cpu;
 pub mod instruction_set;
+pub mod json_to_instruction;
 pub mod memory_manager;
 pub mod send_recv;
-pub mod utility;
-pub mod json_to_instruction;
 pub mod tracing;
-
-
+pub mod utility;
 
 #[derive(Debug, Clone)]
 pub struct CoreInstructionsBuilder {
-    core_instructions : Vec<CoreInstruction>
+    core_instructions: Vec<CoreInstructions>,
 }
 
 impl CoreInstructionsBuilder {
-    pub fn new(size:usize) -> Self {
+    pub fn new(size: usize) -> Self {
         let mut core_instructions = Vec::with_capacity(size);
         for _ in 0..=size {
-            core_instructions.push(CoreInstruction::empty());
+            core_instructions.push(CoreInstructions::empty());
         }
         Self { core_instructions }
     }
 
-    pub fn build(self) -> Vec<CoreInstruction> {
+    pub fn build(self) -> Vec<CoreInstructions> {
         self.core_instructions
     }
 
-    pub fn set_core(&mut self, core : impl TryToUsize, core_instruction : Instructions) -> &mut Self{
-        self.core_instructions[core.try_into().expect("Set core with not valid size")] = core_instruction.into();
-         self
+    pub fn set_core(&mut self, core: impl TryToUsize, core_instruction: Instructions) -> &mut Self {
+        self.core_instructions[core.try_into().expect("Set core with not valid size")] =
+            core_instruction.into();
+        self
     }
 }
 
-
 #[derive(Debug, Clone)]
-pub struct CoreInstruction {
+pub struct CoreInstructions {
     instructions: Instructions,
     program_counter: usize,
 }
 
-impl CoreInstruction {
+impl CoreInstructions {
     fn new(instructions: Instructions, program_counter: usize) -> Self {
         Self {
             instructions,
@@ -53,13 +57,16 @@ impl CoreInstruction {
     }
 
     fn empty() -> Self {
-        Self { instructions: Vec::new(), program_counter: 0 }
+        Self {
+            instructions: Vec::new(),
+            program_counter: 0,
+        }
     }
 }
 
-impl From<Instructions> for CoreInstruction {
+impl From<Instructions> for CoreInstructions {
     fn from(value: Instructions) -> Self {
-        CoreInstruction {
+        CoreInstructions {
             instructions: value,
             program_counter: 0,
         }
@@ -67,39 +74,64 @@ impl From<Instructions> for CoreInstruction {
 }
 
 #[derive(Debug, Clone)]
-pub struct Executable {
-    cpu: CPU,
-    core_instructions: Vec<CoreInstruction>,
-    send_recv : SendRecv,
+pub struct Executable<'a> {
+    cpu: CPU<'a>,
+    core_instructions: Vec<CoreInstructions>,
+    send_recv: SendRecv,
 }
 
-impl Executable {
-    pub fn new(cpu: CPU, core_instructions: Vec<CoreInstruction>) -> Self {
+fn print_status(core_instructions: &[CoreInstructions]) {
+    let mut tot_instructions = 0;
+    let mut progress = 0;
+    for core_instruction in core_instructions.iter() {
+        tot_instructions += core_instruction.instructions.len();
+        progress += core_instruction.program_counter;
+    }
+    println!(
+        "Progress: {}% ({}/{}) ",
+        progress as f32 / tot_instructions as f32 * 100.0,
+        progress,
+        tot_instructions
+    );
+}
+
+impl<'a> Executable<'a> {
+    pub fn new(cpu: CPU<'a>, core_instructions: Vec<CoreInstructions>) -> Executable<'a> {
         let num_core = cpu.num_core();
         let send_recv = SendRecv::new(num_core);
-        assert_eq!(num_core, core_instructions.len(), "Some core doesn't have is list of istruction (required even if empty)");
+        assert_eq!(
+            num_core,
+            core_instructions.len(),
+            "Some core doesn't have is list of istruction (required even if empty)"
+        );
         Self {
             cpu,
             core_instructions,
-            send_recv
+            send_recv,
         }
     }
 
-    pub fn execute(&mut self) {
+    pub fn execute<'b>(&'b mut self)
+    where
+        'a: 'b,
+    {
         let Self {
             cpu,
-            core_instructions,
-            send_recv
+            core_instructions: cores_instructions,
+            send_recv,
         } = self;
         let mut cpu_progressed = 0;
         let max_core = cpu.num_core();
-        let mut index_unit = 0;
+        let mut cpu_index = 0;
+        let mut now = SystemTime::now();
 
         while (cpu_progressed > -2) {
             let mut core_result = InstructionStatus::Completed;
-            while core_result.is_completed() && let Some(core_instruction) = core_instructions.get_mut(index_unit){
+            while core_result.is_completed()
+                && let Some(core_instruction) = cores_instructions.get_mut(cpu_index)
+            {
                 core_result = InstructionStatus::NotExecuted;
-                let CoreInstruction {
+                let CoreInstructions {
                     instructions,
                     program_counter,
                 } = core_instruction;
@@ -112,29 +144,44 @@ impl Executable {
                     cpu_progressed = 0;
                     *program_counter += 1;
                 }
+                if (now.elapsed().unwrap() > Duration::from_secs(1)) {
+                    print_status(&cores_instructions);
+                    now = SystemTime::now();
+                }
+            }
+            handle_wait_sync(cpu, cores_instructions, core_result);
+            match handle_send_recv(cpu, cores_instructions, send_recv, core_result) {
+                (true, other_cpu_index) => {
+                    cpu_progressed = 0;
+                    cpu_index = other_cpu_index;
+                }
+                (false, 0) => {
+                    cpu_index = if cpu_index + 1 >= cores_instructions.len() {
+                        cpu_progressed -= 1;
+                        0
+                    } else {
+                        cpu_index + 1
+                    };
+                }
+                (false, other_cpu_index) => {
+                    cpu_index = other_cpu_index;
+                }
             }
-            if handle_send_recv(cpu, core_instructions, send_recv, core_result) { cpu_progressed = 0; }
-            handle_wait_sync(cpu, core_instructions, core_result);
-            index_unit = if index_unit + 1 >= max_core {
-            cpu_progressed-=1;
-                0
-            } else {
-                index_unit + 1
-            };
         }
+        print_status(cores_instructions);
     }
 
-    pub fn cpu(&self) -> &CPU {
+    pub fn cpu(&self) -> &CPU<'a> {
         &self.cpu
     }
 
-    pub fn cpu_mut(&mut self) -> &mut CPU {
+    pub fn cpu_mut(&mut self) -> &mut CPU<'a> {
         &mut self.cpu
     }
 
-    pub fn dump(&self)  {
+    pub fn dump(&self) {
         let core_instructions = &self.core_instructions;
-        for  (i, core_instruction) in  core_instructions.iter().enumerate() {
+        for (i, core_instruction) in core_instructions.iter().enumerate() {
             eprintln!("INST OF CORE {}:", i);
             for inst in &core_instruction.instructions {
                 inst.dump();
@@ -143,64 +190,12 @@ impl Executable {
     }
 }
 
-fn handle_wait_sync(cpu: &mut CPU, core_instructions: &mut [CoreInstruction], core_result: InstructionStatus) {
-}
-
-
-#[cfg(test)]
-mod test {
-    use super::*;
-    use crate::instruction_set::instruction_data::InstructionDataBuilder;
-    use crate::instruction_set::{InstructionsBuilder, isa::*};
-
-    #[test]
-    fn test_only_host() {
-        let mut cpu = CPU::new(0);
-        cpu.host()
-            .execute_store(0, &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]);
-        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(sldi, idata_build.set_rdimm(1, 0).build());
-        inst_builder.make_inst(sld, idata_build.set_rdr1(1, 1).build());
-        inst_builder.make_inst(sldi, idata_build.set_rdimm(2, 8).build());
-        inst_builder.make_inst(sld, idata_build.set_rdr1(2, 2).build());
-        inst_builder.make_inst(sadd, idata_build.set_rdr1r2(2, 1, 2).build());
-        let mut core_instruction = vec![inst_builder.build().into()];
-        let mut executable = Executable::new(cpu, core_instruction);
-        executable.execute();
-        assert_eq!(executable.cpu_mut().host().register(2), 4, "Not sum to 4");
-    }
-
-    #[test]
-    fn test_10_core_same_code() {
-        let setup_core = |index: usize, cpu: &mut CPU| -> Instructions {
-            cpu.core(index)
-                .execute_store(0, &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]);
-            let mut inst_builder = InstructionsBuilder::new();
-            let mut idata_build = InstructionDataBuilder::new();
-            idata_build.set_core_indx(index as i32).fix_core_indx();
-            inst_builder.make_inst(sldi, idata_build.set_rdimm(1, 0).build());
-            inst_builder.make_inst(sld, idata_build.set_rdr1(1, 1).build());
-            inst_builder.make_inst(sldi, idata_build.set_rdimm(2, 8).build());
-            inst_builder.make_inst(sld, idata_build.set_rdr1(2, 2).build());
-            inst_builder.make_inst(sadd, idata_build.set_rdr1r2(2, 1, 2).build());
-            inst_builder.build()
-        };
-
-        let mut cpu = CPU::new(10);
-        let mut core_instruction = Vec::new();
-        for i in 0..cpu.num_core() {
-            core_instruction.push(setup_core(i, &mut cpu).into())
-        }
-
-        let mut executable = Executable::new(cpu, core_instruction);
-        executable.execute();
-        for i in 0.. executable.cpu.num_core() {
-            assert_eq!(executable.cpu_mut().core(i).register(2), 4, "Core {} not sum to 4", i);
-        }
-    }
-
-    
-
+fn handle_wait_sync<'a, 'b, 'c>(
+    cpu: &'b mut CPU<'a>,
+    core_instructions: &'c mut [CoreInstructions],
+    core_result: InstructionStatus,
+) where
+    'a: 'b,
+    'a: 'c,
+{
 }

backend-simulators/pim/pim-simulator/src/lib/send_recv.rs

@@ -1,7 +1,7 @@
 use anyhow::Context;
 
 use crate::{
-    CoreInstruction, cpu::CPU, instruction_set::InstructionStatus, tracing::TRACER,
+    CoreInstructions, cpu::CPU, instruction_set::InstructionStatus, tracing::TRACER,
     utility::add_offset_rd,
 };
 
@@ -41,14 +41,16 @@ impl SendRecv {
     }
 }
 
-pub fn handle_send_recv(
-    cpu: &mut CPU,
-    core_instructions: &mut [CoreInstruction],
-    send_recv: &mut SendRecv,
+pub fn handle_send_recv<'a, 'b >(
+    cpu: &'b mut CPU<'a>,
+    core_instructions: & mut [CoreInstructions],
+    send_recv: & mut SendRecv,
     core_result: InstructionStatus,
-) -> bool {
-    let transfer_memory = |cpu: &mut CPU,
-                           core_instructions: &mut [CoreInstruction],
+) -> (bool, usize)
+where 'a : 'b
+{
+    let transfer_memory = |cpu: &'b mut CPU<'a>,
+                           core_instructions: & mut [CoreInstructions],
                            sender: Option<SendRecvInfo>,
                            receiver: Option<SendRecvInfo>| {
         if let Some(sender) = sender
@@ -117,7 +119,7 @@ pub fn handle_send_recv(
                 send_recv.sending[sender] = None;
                 send_recv.receiving[receiver] = None;
             }
-            return transfered;
+            (transfered, receiver)
         }
         InstructionStatus::Reciving(instruction_data) => {
             let (core_idx, imm_core) = instruction_data.get_core_immcore();
@@ -146,8 +148,8 @@ pub fn handle_send_recv(
                 send_recv.sending[sender] = None;
                 send_recv.receiving[receiver] = None;
             }
-            return transfered;
+            (transfered, sender)
         }
-        _ => false,
+        _ => (false, 0),
     }
 }

backend-simulators/pim/pim-simulator/src/lib/tracing/mod.rs

@@ -1,11 +1,10 @@
 mod tracing_isa;
 mod disable;
 mod pretty_print;
-
-#[cfg(feature = "tracing")]
 use std::{fs::File, path::{ PathBuf}};
 use std::sync::{LazyLock, Mutex};
 
+
 use crate::Executable;
 
 #[cfg(feature = "tracing")]

src/PIM/Common/PimCommon.cpp

@@ -1,9 +1,13 @@
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/BuiltinTypeInterfaces.h"
 #include "mlir/Interfaces/DestinationStyleOpInterface.h"
 
+#include "llvm/ADT/SmallSet.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/raw_os_ostream.h"
 
 #include <filesystem>
@@ -45,7 +49,9 @@ void dumpModule(ModuleOp moduleOp, const std::string& name) {
 
   std::fstream file(dialectsDir + "/" + name + ".mlir", std::ios::out);
   llvm::raw_os_ostream os(file);
-  os << *moduleOp;
+  OpPrintingFlags flags;
+  flags.elideLargeElementsAttrs();
+  moduleOp.print(os, flags);
   os.flush();
   file.close();
 }
@@ -96,72 +102,95 @@ void markWeightAlways(Operation* op) {
   op->setAttr(PimWeightAlwaysAttrName, UnitAttr::get(op->getContext()));
 }
 
+namespace {
+
+template <typename MVMOpTy, typename VMMOpTy, typename ParentOpTy>
+bool hasMvmVmmWeightUse(ParentOpTy parentOp, unsigned weightIndex) {
+  bool found = false;
+  parentOp.walk([&](Operation* op) {
+    if (auto mvmOp = dyn_cast<MVMOpTy>(op))
+      found |= mvmOp.getWeightIndex() == weightIndex;
+    else if (auto vmmOp = dyn_cast<VMMOpTy>(op))
+      found |= vmmOp.getWeightIndex() == weightIndex;
+  });
+  return found;
+}
+
+template <typename MVMOpTy, typename VMMOpTy, typename ParentOpTy>
+void walkMvmVmmWeightUses(ParentOpTy parentOp, function_ref<void(OpOperand&)> callback) {
+  auto weights = parentOp.getWeights();
+  llvm::SmallSet<unsigned, 8> visited;
+  auto walkWeightIndex = [&](unsigned weightIndex) {
+    if (weightIndex < weights.size() && visited.insert(weightIndex).second)
+      callback(parentOp->getOpOperand(weightIndex));
+  };
+
+  parentOp.walk([&](MVMOpTy op) { walkWeightIndex(op.getWeightIndex()); });
+  parentOp.walk([&](VMMOpTy op) { walkWeightIndex(op.getWeightIndex()); });
+}
+
+} // namespace
+
+bool isSpatialMvmVmmWeightUse(OpOperand& use) {
+  Operation* user = use.getOwner();
+  unsigned operandIndex = use.getOperandNumber();
+
+  auto computeOp = dyn_cast<spatial::SpatCompute>(user);
+  if (!computeOp || operandIndex >= computeOp.getWeights().size())
+    return false;
+
+  return hasMvmVmmWeightUse<spatial::SpatWeightedMVMOp, spatial::SpatWeightedVMMOp>(computeOp, operandIndex);
+}
+
+bool hasOnlySpatialMvmVmmWeightUses(Value value) {
+  SmallPtrSet<Value, 8> visited;
+  auto walkUses = [&](Value currentValue, auto& self) -> bool {
+    if (!visited.insert(currentValue).second)
+      return true;
+    if (currentValue.use_empty())
+      return false;
+
+    return llvm::all_of(currentValue.getUses(), [&](OpOperand& use) {
+      if (isSpatialMvmVmmWeightUse(use))
+        return true;
+
+      Operation* user = use.getOwner();
+      if (auto extractSliceOp = dyn_cast<tensor::ExtractSliceOp>(user))
+        return extractSliceOp.getSource() == currentValue && self(extractSliceOp.getResult(), self);
+      if (auto expandShapeOp = dyn_cast<tensor::ExpandShapeOp>(user))
+        return expandShapeOp.getSrc() == currentValue && self(expandShapeOp.getResult(), self);
+      if (auto collapseShapeOp = dyn_cast<tensor::CollapseShapeOp>(user))
+        return collapseShapeOp.getSrc() == currentValue && self(collapseShapeOp.getResult(), self);
+      if (auto transposeOp = dyn_cast<ONNXTransposeOp>(user))
+        return transposeOp.getData() == currentValue && self(transposeOp.getResult(), self);
+
+      return false;
+    });
+  };
+
+  return walkUses(value, walkUses);
+}
+
+void walkPimMvmVmmWeightUses(Operation* root, function_ref<void(OpOperand&)> callback) {
+  assert(root && "expected valid root op");
+  root->walk([&](pim::PimCoreOp coreOp) {
+    walkMvmVmmWeightUses<pim::PimMVMOp, pim::PimVMMOp>(coreOp, callback);
+  });
+  root->walk([&](pim::PimCoreBatchOp coreBatchOp) {
+    auto weights = coreBatchOp.getWeights();
+    for (auto weight : weights)
+      for (OpOperand& use : weight.getUses())
+        if (use.getOwner() == coreBatchOp.getOperation())
+          callback(use);
+  });
+}
+
 memref::GlobalOp lookupGlobalForGetGlobal(ModuleOp moduleOp, memref::GetGlobalOp getGlobalOp) {
   if (!moduleOp || !getGlobalOp)
     return {};
   return moduleOp.lookupSymbol<memref::GlobalOp>(getGlobalOp.getName());
 }
 
-FailureOr<Operation*> getOtherEndOfChannel(Operation* op, bool opIsReceive, RewriterBase& rewriter) {
-
-  auto channelNewOp = op->getOperand(0).getDefiningOp<spatial::SpatChannelNewOp>();
-  if (!channelNewOp) {
-    op->emitError("User of Channel must have the first operand created by ChannelNewOp.");
-    return failure();
-  }
-  // channelNewOp should have two users: `op` and a
-  // `ChannelSendOp`/`ChannelReceiveOp`
-  auto channelUsers = channelNewOp->getUsers();
-  auto usersIterator = channelUsers.begin();
-  auto firstUser = *usersIterator;
-  usersIterator++;
-  if (usersIterator == channelUsers.end()) {
-    op->emitError("Operand generated by ChannelNewOp must have two users, "
-                  "only one found.");
-    channelNewOp->dump();
-    op->dump();
-    channelNewOp->getParentOp()->dump();
-    return failure();
-  }
-  auto secondUser = *usersIterator;
-  usersIterator++;
-  if (usersIterator != channelUsers.end()) {
-    op->emitError("Operand generated by ChannelNewOp must have two users, "
-                  "more than two found.");
-    return failure();
-  }
-  Operation* notOpUser;
-  if (firstUser == op) {
-    notOpUser = secondUser;
-  }
-  else if (secondUser == op) {
-    notOpUser = firstUser;
-  }
-  else {
-    op->emitError("Operand generated by ChannelNewOp must have two users, "
-                  "and one of them must be me, but"
-                  "none of them is actually me.");
-    return failure();
-  }
-
-  if (opIsReceive) {
-    if (!isa<spatial::SpatChannelSendOp>(notOpUser)) {
-      op->emitError("Operand generated by ChannelNewOp has two user, one is "
-                    "me, the other is not a ChannelSendOp.");
-      return failure();
-    }
-    return notOpUser;
-  }
-  else {
-    if (!isa<spatial::SpatChannelReceiveOp>(notOpUser)) {
-      op->emitError("Operand generated by ChannelNewOp has two user, one is "
-                    "me, the other is not a ChannelReceiveOp.");
-      return failure();
-    }
-    return notOpUser;
-  }
-}
-
 SmallVector<int64_t> computeRowMajorStrides(ArrayRef<int64_t> shape) {
   SmallVector<int64_t> strides(shape.size(), 1);
   for (int64_t dim = static_cast<int64_t>(shape.size()) - 2; dim >= 0; --dim)

src/PIM/Common/PimCommon.hpp

@@ -7,6 +7,7 @@
 #include "mlir/IR/Value.h"
 
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLFunctionalExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 
@@ -16,6 +17,8 @@ inline constexpr llvm::StringRef PimWeightAlwaysAttrName = "weightAlways";
 
 namespace onnx_mlir {
 
+inline constexpr llvm::StringLiteral kCoreIdAttrName = "core_id";
+
 struct ResolvedContiguousAddress {
   mlir::Value base;
   int64_t byteOffset = 0;
@@ -40,10 +43,12 @@ bool hasWeightAlways(mlir::Operation* op);
 
 void markWeightAlways(mlir::Operation* op);
 
-mlir::memref::GlobalOp lookupGlobalForGetGlobal(mlir::ModuleOp moduleOp, mlir::memref::GetGlobalOp getGlobalOp);
+bool isSpatialMvmVmmWeightUse(mlir::OpOperand& use);
+bool hasOnlySpatialMvmVmmWeightUses(mlir::Value value);
 
-llvm::FailureOr<mlir::Operation*>
-getOtherEndOfChannel(mlir::Operation* op, bool opIsReceive, mlir::RewriterBase& rewriter);
+void walkPimMvmVmmWeightUses(mlir::Operation* root, llvm::function_ref<void(mlir::OpOperand&)> callback);
+
+mlir::memref::GlobalOp lookupGlobalForGetGlobal(mlir::ModuleOp moduleOp, mlir::memref::GetGlobalOp getGlobalOp);
 
 llvm::SmallVector<int64_t> computeRowMajorStrides(llvm::ArrayRef<int64_t> shape);
 

src/PIM/Compiler/PimCodeGen.cpp

@@ -1,10 +1,15 @@
+#include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/IRMapping.h"
+#include "mlir/IR/Value.h"
 
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/JSON.h"
 #include "llvm/Support/raw_ostream.h"
@@ -33,6 +38,12 @@ MemEntry* PimMemory::gatherMemEntry(mlir::Value value) {
   return &memEntries.emplace_back(memEntry, value).first;
 }
 
+void PimMemory::allocateGatheredMemory() {
+  llvm::sort(memEntries, [](auto a, auto b) -> bool { return a.first.size > b.first.size; });
+  for (auto& [memEntry, value] : memEntries)
+    allocateMemoryForValue(value, memEntry);
+}
+
 void PimMemory::allocateMemoryForValue(mlir::Value value, MemEntry& memEntry) {
   memEntry.address = firstAvailableAddress;
   firstAvailableAddress += memEntry.size;
@@ -44,35 +55,49 @@ void PimMemory::allocateMemoryForValue(mlir::Value value, MemEntry& memEntry) {
 }
 
 void PimMemory::allocateHost(ModuleOp moduleOp, func::FuncOp funcOp) {
-  // More than one SSA value per single global constant:
-  // Cannot call gatherMemEntry for each of them, otherwise memory will be allocated multiple times
-  // Thus, call gatherMemEntry only for the first SSA value and assign the same memEntry to all others
-  SmallDenseMap<memref::GlobalOp, MemEntry*, 8> globalConstants;
+  SmallDenseMap<memref::GlobalOp, mlir::Value, 8> globalConstants;
+  SmallVector<std::pair<mlir::Value, mlir::Value>, 16> globalAliases;
+  SmallVector<mlir::Value> args;
+
+
+  for (mlir::Value arg : funcOp.getArguments()){
+    gatherMemEntry(arg);
+    args.push_back(arg);
+  }
+
   funcOp.walk([&](memref::GetGlobalOp getGlobalOp) {
     if (!hasWeightAlways(getGlobalOp)) {
       auto globalMemrefOp = lookupGlobalForGetGlobal(moduleOp, getGlobalOp);
-      auto iter = globalConstants.find(globalMemrefOp);
-      if (iter == globalConstants.end())
-        globalConstants[globalMemrefOp] = gatherMemEntry(getGlobalOp);
-      else {
-        MemEntry memEntry = *iter->second;
-        globalMemEntriesMap[getGlobalOp] = memEntry;
+      if (globalMemrefOp.getName().starts_with("arg")){
+        StringRef indexStr = globalMemrefOp.getName().substr(4);
+        int  index = 0;
+        llvm::to_integer(indexStr,index, 10);
+        globalAliases.push_back({getGlobalOp.getResult(),  args[index]});
       }
+      auto [iter, inserted] = globalConstants.try_emplace(globalMemrefOp, getGlobalOp.getResult());
+      if (inserted)
+        gatherMemEntry(getGlobalOp.getResult());
+      else
+        globalAliases.push_back({getGlobalOp.getResult(), iter->second});
     }
   });
 
-  for (mlir::Value arg : funcOp.getArguments())
-    gatherMemEntry(arg);
 
-  allocateCore(funcOp);
+  funcOp.walk([&](memref::AllocOp allocOp) {
+    if (!allocOp->getParentOfType<pim::PimCoreOp>())
+      gatherMemEntry(allocOp.getResult());
+  });
+
+  allocateGatheredMemory();
+
+  for (auto [alias, original] : globalAliases)
+    globalMemEntriesMap[alias] = getMemEntry(original);
 }
 
 void PimMemory::allocateCore(Operation* op) {
   op->walk([&](memref::AllocOp allocOp) { gatherMemEntry(allocOp); });
 
-  llvm::sort(memEntries, [](auto a, auto b) -> bool { return a.first.size > b.first.size; });
-  for (auto& [memEntry, value] : memEntries)
-    allocateMemoryForValue(value, memEntry);
+  allocateGatheredMemory();
 }
 
 MemEntry PimMemory::getMemEntry(mlir::Value value) const {
@@ -122,6 +147,12 @@ json::Object PimCodeGen::createEmptyOffset() {
   return offset;
 }
 
+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;
+}
+
 static json::Object createRs1OnlyOffset() {
   json::Object offset;
   offset["offset_select"] = 1;
@@ -181,7 +212,7 @@ void PimCodeGen::emitCommunicationOp(StringRef opName, size_t bufferAddr, size_t
   json::Object json;
   json["op"] = opName;
   json["rd"] = 0;
-  json["core"] = coreId;
+  json["core"] = remapCoreId(coreId);
   json["size"] = size;
   json["offset"] = createEmptyOffset();
   emitInstruction(std::move(json));
@@ -403,6 +434,9 @@ void PimCodeGen::codeGenVSoftmaxOp(pim::PimVSoftmaxOp vsoftmaxOp, const StaticVa
   emitInstruction(std::move(json));
 }
 
+void PimCodeGen::codeGetGlobalOp(memref::GetGlobalOp getGlobalOp, const StaticValueKnowledge& knowledge) const {
+}
+
 void PimCodeGen::codeGenTransposeOp(pim::PimTransposeOp transposeOp, const StaticValueKnowledge& knowledge) const {
   auto srcAddr = addressOf(transposeOp.getInput(), knowledge);
   auto dstAddr = addressOf(transposeOp.getOutputBuffer(), knowledge);
@@ -465,6 +499,136 @@ std::string getMemorySizeAsString(size_t size) {
   return std::to_string(size) + " Bytes";
 }
 
+static SmallVector<unsigned, 8> getUsedWeightIndices(Block& block) {
+  SmallVector<unsigned, 8> indices;
+  auto addIndex = [&](unsigned weightIndex) {
+    if (!llvm::is_contained(indices, weightIndex))
+      indices.push_back(weightIndex);
+  };
+
+  block.walk([&](pim::PimMVMOp mvmOp) { addIndex(mvmOp.getWeightIndex()); });
+  block.walk([&](pim::PimVMMOp vmmOp) { addIndex(vmmOp.getWeightIndex()); });
+  llvm::sort(indices);
+  return indices;
+}
+
+static SmallVector<unsigned, 8> getUsedWeightIndices(pim::PimCoreOp coreOp) {
+  return getUsedWeightIndices(coreOp.getBody().front());
+}
+
+static SmallVector<int32_t> getBatchCoreIds(pim::PimCoreBatchOp coreBatchOp) {
+  auto coreIdsAttr = coreBatchOp->getAttrOfType<DenseI32ArrayAttr>(onnx_mlir::kCoreIdAttrName);
+  assert(coreIdsAttr && "pim.core_batch requires core_id array attribute");
+  return SmallVector<int32_t>(coreIdsAttr.asArrayRef().begin(), coreIdsAttr.asArrayRef().end());
+}
+
+static SmallVector<Operation*> collectTopLevelCoreLikeOps(func::FuncOp funcOp) {
+  SmallVector<Operation*> coreLikeOps;
+  for (Operation& op : funcOp.getBody().front()) {
+    if (dyn_cast<pim::PimCoreOp>(&op) || dyn_cast<pim::PimCoreBatchOp>(&op))
+      coreLikeOps.push_back(&op);
+  }
+  return coreLikeOps;
+}
+
+static pim::PimCoreOp materializeScalarCoreFromBatchLane(pim::PimCoreBatchOp coreBatchOp, unsigned lane) {
+  OpBuilder builder(coreBatchOp);
+  builder.setInsertionPointAfter(coreBatchOp);
+
+  size_t laneCount = static_cast<size_t>(coreBatchOp.getLaneCount());
+  size_t weightsPerLane = coreBatchOp.getWeights().size() / laneCount;
+  SmallVector<mlir::Value> laneWeights;
+  laneWeights.reserve(weightsPerLane);
+  for (size_t weightIndex = 0; weightIndex < weightsPerLane; ++weightIndex)
+    laneWeights.push_back(coreBatchOp.getWeights()[lane * weightsPerLane + weightIndex]);
+
+  auto coreIds = getBatchCoreIds(coreBatchOp);
+  auto scalarCore = pim::PimCoreOp::create(builder,
+                                           coreBatchOp.getLoc(),
+                                           ValueRange(laneWeights),
+                                           builder.getI32IntegerAttr(coreIds[lane]));
+  Block* block = builder.createBlock(&scalarCore.getBody(), scalarCore.getBody().end());
+  IRMapping mapper;
+  if (coreBatchOp.getBody().front().getNumArguments() == 1)
+    mapper.map(coreBatchOp.getBody().front().getArgument(0), coreBatchOp.getInputs()[lane]);
+
+  builder.setInsertionPointToEnd(block);
+  for (Operation& op : coreBatchOp.getBody().front()) {
+    if (isa<pim::PimHaltOp>(op)) {
+      pim::PimHaltOp::create(builder, op.getLoc());
+      continue;
+    }
+
+    if (auto sendBatchOp = dyn_cast<pim::PimSendBatchOp>(op)) {
+      pim::PimSendOp::create(builder,
+                             sendBatchOp.getLoc(),
+                             mapper.lookup(sendBatchOp.getInput()),
+                             sendBatchOp.getSizeAttr(),
+                             builder.getI32IntegerAttr(sendBatchOp.getTargetCoreIds()[lane]));
+      continue;
+    }
+
+    if (auto receiveBatchOp = dyn_cast<pim::PimReceiveBatchOp>(op)) {
+      auto scalarReceive = pim::PimReceiveOp::create(builder,
+                                                     receiveBatchOp.getLoc(),
+                                                     receiveBatchOp.getOutput().getType(),
+                                                     mapper.lookup(receiveBatchOp.getOutputBuffer()),
+                                                     receiveBatchOp.getSizeAttr(),
+                                                     builder.getI32IntegerAttr(receiveBatchOp.getSourceCoreIds()[lane]));
+      mapper.map(receiveBatchOp.getOutput(), scalarReceive.getOutput());
+      continue;
+    }
+
+    if (auto memcpBatchOp = dyn_cast<pim::PimMemCopyHostToDevBatchOp>(op)) {
+      mlir::Value hostSource = mapper.lookupOrNull(memcpBatchOp.getHostSource());
+      if (!hostSource)
+        hostSource = memcpBatchOp.getHostSource();
+
+      auto scalarCopy = pim::PimMemCopyHostToDevOp::create(builder,
+                                                           memcpBatchOp.getLoc(),
+                                                           memcpBatchOp.getOutput().getType(),
+                                                           mapper.lookup(memcpBatchOp.getDeviceTarget()),
+                                                           hostSource,
+                                                           memcpBatchOp.getDeviceTargetOffsetAttr(),
+                                                           memcpBatchOp.getHostSourceOffsetAttr(),
+                                                           memcpBatchOp.getSizeAttr());
+      mapper.map(memcpBatchOp.getOutput(), scalarCopy.getOutput());
+      continue;
+    }
+
+    Operation* cloned = builder.clone(op, mapper);
+    for (auto [originalResult, clonedResult] : llvm::zip(op.getResults(), cloned->getResults()))
+      mapper.map(originalResult, clonedResult);
+  }
+
+  if (block->empty() || !isa<pim::PimHaltOp>(block->back()))
+    pim::PimHaltOp::create(builder, coreBatchOp.getLoc());
+  return scalarCore;
+}
+
+static void aliasMaterializedHostGlobals(
+  ModuleOp moduleOp, func::FuncOp funcOp, pim::PimCoreOp coreOp, PimAcceleratorMemory& memory) {
+  coreOp.walk([&](memref::GetGlobalOp getGlobalOp) {
+    if (hasWeightAlways(getGlobalOp) || memory.memEntriesMap.contains(getGlobalOp.getResult()))
+      return;
+
+    auto targetGlobal = lookupGlobalForGetGlobal(moduleOp, getGlobalOp);
+    if (!targetGlobal)
+      return;
+
+    mlir::Value aliasedValue;
+    funcOp.walk([&](memref::GetGlobalOp candidate) {
+      if (aliasedValue || candidate == getGlobalOp || !memory.memEntriesMap.contains(candidate.getResult()))
+        return;
+      if (lookupGlobalForGetGlobal(moduleOp, candidate) == targetGlobal)
+        aliasedValue = candidate.getResult();
+    });
+
+    if (aliasedValue)
+      memory.memEntriesMap[getGlobalOp.getResult()] = memory.memEntriesMap[aliasedValue];
+  });
+}
+
 /// Write global constant data into a binary memory image at their allocated addresses.
 static OnnxMlirCompilerErrorCodes
 writeMemoryBinary(ModuleOp moduleOp, func::FuncOp funcOp, PimAcceleratorMemory& memory, StringRef outputDirPath) {
@@ -478,12 +642,15 @@ writeMemoryBinary(ModuleOp moduleOp, func::FuncOp funcOp, PimAcceleratorMemory&
 
   std::vector<char> memoryBuffer(memory.hostMem.getFirstAvailableAddress(), 0);
 
+  SmallPtrSet<Operation*, 16> writtenGlobals;
   funcOp.walk([&](memref::GetGlobalOp getGlobalOp) {
     if (hasWeightAlways(getGlobalOp))
       return;
     auto globalOp = lookupGlobalForGetGlobal(moduleOp, getGlobalOp);
     if (!globalOp)
       return;
+    if (!writtenGlobals.insert(globalOp.getOperation()).second)
+      return;
     auto initialValue = globalOp.getInitialValue();
     if (!initialValue)
       return;
@@ -556,6 +723,8 @@ static int64_t codeGenCoreOps(Block& block, PimCodeGen& coreCodeGen) {
         coreCodeGen.codeGenVSigmOp(vsigmOp, knowledge);
       else if (auto vsoftmaxOp = dyn_cast<pim::PimVSoftmaxOp>(op))
         coreCodeGen.codeGenVSoftmaxOp(vsoftmaxOp, knowledge);
+      else if (auto getGlobalOp = dyn_cast<memref::GetGlobalOp>(op))
+        coreCodeGen.codeGetGlobalOp(getGlobalOp, knowledge);
       else {
         op.emitError("Unsupported codegen for this operation");
         op.dump();
@@ -645,7 +814,7 @@ static OnnxMlirCompilerErrorCodes writeCrossbarWeights(ModuleOp moduleOp,
   return CompilerSuccess;
 }
 
-llvm::DenseMap<pim::PimCoreOp, llvm::DenseMap<mlir::Value, std::string>>
+llvm::DenseMap<size_t, llvm::DenseMap<mlir::Value, std::string>>
 createAndPopulateWeightFolder(func::FuncOp funcOp, StringRef outputDirPath) {
   ModuleOp moduleOp = funcOp->getParentOfType<ModuleOp>();
   auto coreWeightsDirPath = outputDirPath + "/weights";
@@ -654,80 +823,104 @@ createAndPopulateWeightFolder(func::FuncOp funcOp, StringRef outputDirPath) {
   size_t indexFileName = 0;
 
   int64_t xbarSize = crossbarSize.getValue();
-  llvm::DenseMap<pim::PimCoreOp, llvm::DenseMap<mlir::Value, std::string>> mapCoreWeightToFileName;
+  llvm::DenseMap<size_t, llvm::DenseMap<mlir::Value, std::string>> mapCoreWeightToFileName;
   llvm::DenseMap<memref::GlobalOp, std::string> mapGlobalOpToFileName;
 
-  for (pim::PimCoreOp coreOp : funcOp.getOps<pim::PimCoreOp>()) {
-    for (auto [index, weight] : llvm::enumerate(coreOp.getWeights())) {
+  SmallVector<Operation*> coreLikeOps = collectTopLevelCoreLikeOps(funcOp);
 
-      auto getGlobalOp = weight.getDefiningOp<memref::GetGlobalOp>();
-      if (!getGlobalOp) {
-        coreOp.emitWarning("Weight is not from a memref.get_global at index " + std::to_string(index));
-        assert(!getGlobalOp && "Weight is not from a memref.get_global");
-      }
+  for (Operation* op : coreLikeOps) {
+    SmallVector<pim::PimCoreOp> scalarCores;
+    if (auto coreOp = dyn_cast<pim::PimCoreOp>(op)) {
+      scalarCores.push_back(coreOp);
+    }
+    else {
+      auto coreBatchOp = cast<pim::PimCoreBatchOp>(op);
+      for (unsigned lane = 0; lane < static_cast<unsigned>(coreBatchOp.getLaneCount()); ++lane)
+        scalarCores.push_back(materializeScalarCoreFromBatchLane(coreBatchOp, lane));
+    }
 
-      auto globalOp = lookupGlobalForGetGlobal(moduleOp, getGlobalOp);
-      if (!globalOp) {
-        coreOp.emitWarning("Could not find memref.global for weight at index " + std::to_string(index));
-        assert(!globalOp && "Could not find memref.global");
-      }
+    for (pim::PimCoreOp coreOp : scalarCores) {
+      size_t coreId = static_cast<size_t>(coreOp.getCoreId());
+      for (unsigned index : getUsedWeightIndices(coreOp)) {
+        if (index >= coreOp.getWeights().size()) {
+          coreOp.emitWarning("Weight index " + std::to_string(index) + " is out of range");
+          assert(index < coreOp.getWeights().size() && "Weight index is out of range");
+        }
+        mlir::Value weight = coreOp.getWeights()[index];
 
-      auto initialValue = globalOp.getInitialValue();
-      if (!initialValue) {
-        coreOp.emitWarning("memref.global has no initial value at index " + std::to_string(index));
-        assert(!initialValue && "memref.global has no initial value");
-      }
+        auto getGlobalOp = weight.getDefiningOp<memref::GetGlobalOp>();
+        if (!getGlobalOp) {
+          coreOp.emitWarning("Weight is not from a memref.get_global at index " + std::to_string(index));
+          assert(!getGlobalOp && "Weight is not from a memref.get_global");
+        }
 
-      auto denseAttr = dyn_cast<DenseElementsAttr>(*initialValue);
-      if (!denseAttr) {
-        coreOp.emitWarning("memref.global initial value is not dense at index " + std::to_string(index));
-        assert(!denseAttr && "memref.global initial value is not dense");
-      }
+        auto globalOp = lookupGlobalForGetGlobal(moduleOp, getGlobalOp);
+        if (!globalOp) {
+          coreOp.emitWarning("Could not find memref.global for weight at index " + std::to_string(index));
+          assert(!globalOp && "Could not find memref.global");
+        }
 
-      if (mapGlobalOpToFileName.contains(globalOp)) {
-        auto& fileName = mapGlobalOpToFileName[globalOp];
-        std::pair<mlir::Value, std::string> weightToFile = {weight, fileName};
-        mapCoreWeightToFileName[coreOp].insert(weightToFile);
-        continue;
-      }
+        auto initialValue = globalOp.getInitialValue();
+        if (!initialValue) {
+          coreOp.emitWarning("memref.global has no initial value at index " + std::to_string(index));
+          assert(!initialValue && "memref.global has no initial value");
+        }
 
-      auto type = denseAttr.getType();
-      auto shape = type.getShape();
-      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");
+        auto denseAttr = dyn_cast<DenseElementsAttr>(*initialValue);
+        if (!denseAttr) {
+          coreOp.emitWarning("memref.global initial value is not dense at index " + std::to_string(index));
+          assert(!denseAttr && "memref.global initial value is not dense");
+        }
 
-      size_t elementByteWidth = type.getElementType().getIntOrFloatBitWidth() / 8;
+        if (mapGlobalOpToFileName.contains(globalOp)) {
+          auto& fileName = mapGlobalOpToFileName[globalOp];
+          std::pair<mlir::Value, std::string> weightToFile = {weight, fileName};
+          mapCoreWeightToFileName[coreId].insert(weightToFile);
+          continue;
+        }
 
-      std::string newFileName = "crossbar_" + std::to_string(indexFileName++) + ".bin";
-      auto weightFilePath = (coreWeightsDirPath + "/" + newFileName).str();
-      std::error_code errorCode;
-      raw_fd_ostream weightFileStream(weightFilePath, errorCode, sys::fs::OF_None);
-      if (errorCode) {
-        errs() << "Error while opening weight file `" << weightFilePath << "`: " << errorCode.message() << '\n';
-        assert(errorCode);
-      }
+        auto type = denseAttr.getType();
+        auto shape = type.getShape();
+        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");
 
-      uint64_t zero = 0;
-      for (int64_t row = 0; row < xbarSize; row++) {
-        for (int64_t col = 0; col < xbarSize; col++) {
-          if (row < numRows && col < numCols) {
-            int64_t index = row * numCols + col;
-            APInt bits = denseAttr.getValues<APFloat>()[index].bitcastToAPInt();
-            uint64_t word = bits.getZExtValue();
-            weightFileStream.write(reinterpret_cast<const char*>(&word), elementByteWidth);
-          }
-          else {
-            weightFileStream.write(reinterpret_cast<const char*>(&zero), elementByteWidth);
+        size_t elementByteWidth = type.getElementType().getIntOrFloatBitWidth() / 8;
+
+        std::string newFileName = "crossbar_" + std::to_string(indexFileName++) + ".bin";
+        auto weightFilePath = (coreWeightsDirPath + "/" + newFileName).str();
+        std::error_code errorCode;
+        raw_fd_ostream weightFileStream(weightFilePath, errorCode, sys::fs::OF_None);
+        if (errorCode) {
+          errs() << "Error while opening weight file `" << weightFilePath << "`: " << errorCode.message() << '\n';
+          assert(errorCode);
+        }
+
+        uint64_t zero = 0;
+        for (int64_t row = 0; row < xbarSize; row++) {
+          for (int64_t col = 0; col < xbarSize; col++) {
+            if (row < numRows && col < numCols) {
+              int64_t index = row * numCols + col;
+              APInt bits = denseAttr.getValues<APFloat>()[index].bitcastToAPInt();
+              uint64_t word = bits.getZExtValue();
+              weightFileStream.write(reinterpret_cast<const char*>(&word), elementByteWidth);
+            }
+            else {
+              weightFileStream.write(reinterpret_cast<const char*>(&zero), elementByteWidth);
+            }
           }
         }
-      }
 
-      weightFileStream.close();
-      mapGlobalOpToFileName.insert({globalOp, newFileName});
-      mapCoreWeightToFileName[coreOp].insert({weight, newFileName});
+        weightFileStream.close();
+        mapGlobalOpToFileName.insert({globalOp, newFileName});
+        mapCoreWeightToFileName[coreId].insert({weight, newFileName});
+      }
     }
+
+    for (pim::PimCoreOp coreOp : scalarCores)
+      if (coreOp.getOperation() != op)
+        coreOp.erase();
   }
   return mapCoreWeightToFileName;
 }
@@ -735,13 +928,14 @@ createAndPopulateWeightFolder(func::FuncOp funcOp, StringRef outputDirPath) {
 /// Write the top-level PIM configuration JSON (core count, crossbar config, I/O addresses).
 static OnnxMlirCompilerErrorCodes writeConfigJson(func::FuncOp funcOp,
                                                   PimAcceleratorMemory& memory,
-                                                  size_t coreCount,
+                                                  size_t maxCoreId,
                                                   json::Object xbarsPerArrayGroup,
                                                   StringRef outputDirPath) {
   json::Object configJson;
 
-  // +1 because pimsim-nn also considers the host as a core
-  configJson["core_cnt"] = coreCount + 1;
+  // pimsim-nn indexes cores directly by their numeric core ID, with the host
+  // occupying core 0.
+  configJson["core_cnt"] = maxCoreId + 1;
 
   // TODO: Should this be based on the floating point type used in the model?
   // The 2 following values determine the bitwidth of the vectors' elements: bitwidth = adc_count * cell_precision
@@ -815,61 +1009,103 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimJson(ModuleOp& moduleOp, std::
   // For each core, specify the number of crossbar per array group.
   // This implementation always assigns one crossbar per group.
   json::Object xbarsPerArrayGroup;
-  size_t coreCount = 0;
+  size_t maxCoreId = 0;
 
   // Create Weight Folder
   auto mapCoreWeightToFileName = createAndPopulateWeightFolder(funcOp, outputDirPath);
 
-  for (auto coreOp : funcOp.getOps<pim::PimCoreOp>()) {
-    auto coreId = coreOp.getCoreId();
-    coreCount++;
+  SmallVector<Operation*> coreLikeOps = collectTopLevelCoreLikeOps(funcOp);
+  llvm::DenseMap<size_t, size_t> emittedCoreIds;
+  size_t nextEmittedCoreId = 1;
 
-    std::error_code errorCode;
-    auto outputCorePath = outputDirPath + "/core_" + std::to_string(coreId) + ".json";
-    raw_fd_ostream coreFileStream(outputCorePath, errorCode);
-    if (errorCode) {
-      errs() << "Error while opening core file `" << outputCorePath << "`: " << errorCode.message() << '\n';
-      return InvalidOutputFileAccess;
-    }
-    coreFileStream << '[';
-
-    PimCodeGen coreCodeGen(memory, coreFileStream);
-    memory.getOrCreateDeviceMem(coreId).allocateCore(coreOp);
-
-    int64_t processedOperations = codeGenCoreOps(coreOp.getBody().front(), coreCodeGen);
-    if (processedOperations < 0)
-      return CompilerFailure;
-    assert(processedOperations > 0);
-
-    // Remove trailing comma, close JSON array
-    coreFileStream.seek(coreFileStream.tell() - 1);
-    coreFileStream << ']';
-    coreFileStream.close();
-
-    // Write crossbar weights for this core
-    auto coreWeightsDirPath = outputDirPath + "/core_" + std::to_string(coreId);
-    if (auto error = sys::fs::create_directory(coreWeightsDirPath)) {
-      errs() << "Error creating core directory: " << coreWeightsDirPath << ": " << error.message() << '\n';
-      return InvalidOutputFileAccess;
+  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& mapWeightToFile = mapCoreWeightToFileName[coreOp];
-    json::Array xbarsPerGroup;
-    for (auto [index, weight] : llvm::enumerate(coreOp.getWeights())) {
-      xbarsPerGroup.push_back(index);
-      assert(mapWeightToFile.contains(weight) && "Weight was not materialized into a file!!");
-      auto& fileName = mapWeightToFile[weight];
-      if (auto error = sys::fs::create_link(outputDirPath + "/weights/" + fileName,
-                                            coreWeightsDirPath + "/crossbar_" + std::to_string(index) + ".bin")) {
-        errs() << "Error creating link file: " << (outputDirPath + "/weights/" + fileName) << " to "
-               << (coreWeightsDirPath + "/crossbar_" + std::to_string(index) + ".bin") << "\nError:" << error.message()
-               << '\n';
-        return InvalidOutputFileAccess;
-      }
+    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++;
     }
-
-    xbarsPerArrayGroup["core" + std::to_string(coreId)] = std::move(xbarsPerGroup);
   }
 
-  return writeConfigJson(funcOp, memory, coreCount, std::move(xbarsPerArrayGroup), outputDirPath);
+  for (Operation* op : coreLikeOps) {
+    SmallVector<pim::PimCoreOp> scalarCores;
+    if (auto coreOp = dyn_cast<pim::PimCoreOp>(op)) {
+      scalarCores.push_back(coreOp);
+    }
+    else {
+      auto coreBatchOp = cast<pim::PimCoreBatchOp>(op);
+      for (unsigned lane = 0; lane < static_cast<unsigned>(coreBatchOp.getLaneCount()); ++lane)
+        scalarCores.push_back(materializeScalarCoreFromBatchLane(coreBatchOp, lane));
+    }
+
+    for (pim::PimCoreOp coreOp : scalarCores) {
+      size_t originalCoreId = static_cast<size_t>(coreOp.getCoreId());
+      size_t coreId = emittedCoreIds.lookup(originalCoreId);
+      maxCoreId = std::max(maxCoreId, coreId);
+
+      std::error_code errorCode;
+      auto outputCorePath = outputDirPath + "/core_" + std::to_string(coreId) + ".json";
+      raw_fd_ostream coreFileStream(outputCorePath, errorCode);
+      if (errorCode) {
+        errs() << "Error while opening core file `" << outputCorePath << "`: " << errorCode.message() << '\n';
+        return InvalidOutputFileAccess;
+      }
+      coreFileStream << '[';
+
+      PimCodeGen coreCodeGen(memory, coreFileStream, emittedCoreIds);
+      aliasMaterializedHostGlobals(moduleOp, funcOp, coreOp, memory);
+      memory.getOrCreateDeviceMem(coreId).allocateCore(coreOp);
+
+      int64_t processedOperations = codeGenCoreOps(coreOp.getBody().front(), coreCodeGen);
+      if (processedOperations < 0)
+        return CompilerFailure;
+      assert(processedOperations > 0);
+
+      coreFileStream.seek(coreFileStream.tell() - 1);
+      coreFileStream << ']';
+      coreFileStream.close();
+
+      auto coreWeightsDirPath = outputDirPath + "/core_" + std::to_string(coreId);
+      if (auto error = sys::fs::create_directory(coreWeightsDirPath)) {
+        errs() << "Error creating core directory: " << coreWeightsDirPath << ": " << error.message() << '\n';
+        return InvalidOutputFileAccess;
+      }
+
+      auto& mapWeightToFile = mapCoreWeightToFileName[originalCoreId];
+      json::Array xbarsPerGroup;
+      for (unsigned index : getUsedWeightIndices(coreOp)) {
+        if (index >= coreOp.getWeights().size()) {
+          coreOp.emitWarning("Weight index " + std::to_string(index) + " is out of range");
+          assert(index < coreOp.getWeights().size() && "Weight index is out of range");
+        }
+        mlir::Value weight = coreOp.getWeights()[index];
+        xbarsPerGroup.push_back(index);
+        assert(mapWeightToFile.contains(weight) && "Weight was not materialized into a file!!");
+        auto& fileName = mapWeightToFile[weight];
+        if (auto error = sys::fs::create_link(outputDirPath + "/weights/" + fileName,
+                                              coreWeightsDirPath + "/crossbar_" + std::to_string(index) + ".bin")) {
+          errs() << "Error creating link file: " << (outputDirPath + "/weights/" + fileName) << " to "
+                 << (coreWeightsDirPath + "/crossbar_" + std::to_string(index) + ".bin") << "\nError:"
+                 << error.message() << '\n';
+          return InvalidOutputFileAccess;
+        }
+      }
+
+      xbarsPerArrayGroup["core" + std::to_string(coreId)] = std::move(xbarsPerGroup);
+    }
+
+    for (pim::PimCoreOp coreOp : scalarCores)
+      if (coreOp.getOperation() != op)
+        coreOp.erase();
+  }
+
+  return writeConfigJson(funcOp, memory, maxCoreId, std::move(xbarsPerArrayGroup), outputDirPath);
 }

src/PIM/Compiler/PimCodeGen.hpp

@@ -1,5 +1,6 @@
 #pragma once
 
+#include "llvm/ADT/DenseMap.h"
 #include "llvm-project/clang/include/clang/Basic/LLVM.h"
 #include "llvm/Support/JSON.h"
 
@@ -24,6 +25,7 @@ class PimMemory {
   size_t firstAvailableAddress = 0;
 
   MemEntry* gatherMemEntry(mlir::Value value);
+  void allocateGatheredMemory();
   void allocateMemoryForValue(mlir::Value value, MemEntry& memEntry);
 
 public:
@@ -57,10 +59,12 @@ public:
 class PimCodeGen {
   PimAcceleratorMemory& memory;
   llvm::raw_fd_ostream& coreFileStream;
+  const llvm::DenseMap<size_t, size_t>& emittedCoreIds;
 
   size_t addressOf(mlir::Value value, const StaticValueKnowledge& knowledge) const {
     return memory.getValueAddress(value, knowledge);
   }
+  size_t remapCoreId(size_t coreId) const;
 
   static llvm::json::Object createEmptyOffset();
   void emitInstruction(llvm::json::Object instruction) const;
@@ -82,8 +86,10 @@ class PimCodeGen {
   void emitMvmOp(size_t groupId, size_t rdAddr, size_t rdOffset, size_t rs1Addr, size_t rs1Offset) const;
 
 public:
-  PimCodeGen(PimAcceleratorMemory& memory, llvm::raw_fd_ostream& coreJson)
-  : memory(memory), coreFileStream(coreJson) {}
+  PimCodeGen(PimAcceleratorMemory& memory,
+             llvm::raw_fd_ostream& coreJson,
+             const llvm::DenseMap<size_t, size_t>& emittedCoreIds)
+  : memory(memory), coreFileStream(coreJson), emittedCoreIds(emittedCoreIds) {}
 
   void codeGenLoadOp(pim::PimMemCopyHostToDevOp loadOp, const StaticValueKnowledge& knowledge) const;
   void codeGenStoreOp(pim::PimMemCopyDevToHostOp storeOp, const StaticValueKnowledge& knowledge) const;
@@ -105,6 +111,7 @@ public:
   void codeGenVTanhOp(pim::PimVTanhOp vtanhOp, const StaticValueKnowledge& knowledge) const;
   void codeGenVSigmOp(pim::PimVSigmOp vsigmOp, const StaticValueKnowledge& knowledge) const;
   void codeGenVSoftmaxOp(pim::PimVSoftmaxOp vsoftmaxOp, const StaticValueKnowledge& knowledge) const;
+  void codeGetGlobalOp(mlir::memref::GetGlobalOp getGlobalOp, const StaticValueKnowledge& knowledge) const;
   void codeGenTransposeOp(pim::PimTransposeOp transposeOp, const StaticValueKnowledge& knowledge) const;
 };
 

src/PIM/Compiler/PimCompilerOptions.cpp

@@ -41,12 +41,18 @@ llvm::cl::opt<size_t>
   crossbarSize("crossbar-size", llvm::cl::desc("Width and heigth of a single crossbar"), llvm::cl::init(2));
 
 llvm::cl::opt<size_t>
-  crossbarCountInCore("crossbar-count", llvm::cl::desc("Number of crossbars in each core"), llvm::cl::init(2));
+  crossbarCountInCore("crossbar-count", llvm::cl::desc("Number of crossbars in each core"), llvm::cl::init(256));
 
 llvm::cl::opt<long> coresCount("core-count",
                                llvm::cl::desc("Number of cores in the chip. `-1` to use the minimum amount of cores."),
                                llvm::cl::init(-1));
 
+llvm::cl::opt<size_t> dcpCriticalWindowSize(
+  "dcp-critical-window-size",
+  llvm::cl::desc("Number of lowest-slack virtual nodes considered by each DCP coarsening iteration. "
+                 "Use 0 to run the legacy full-graph DCP analysis."),
+  llvm::cl::init(4000));
+
 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

@@ -29,6 +29,7 @@ extern llvm::cl::opt<bool> useExperimentalConvImpl;
 extern llvm::cl::opt<size_t> crossbarSize;
 extern llvm::cl::opt<size_t> crossbarCountInCore;
 extern llvm::cl::opt<long> coresCount;
+extern llvm::cl::opt<size_t> dcpCriticalWindowSize;
 
 // This option, by default set to false, will ignore an error when resolving a
 // specific tiles of the operands of a concat. This specific case is when the

src/PIM/Conversion/ONNXToSpatial/Common.hpp

@@ -12,6 +12,7 @@
 #include <utility>
 
 #include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/STLExtras.h"
 
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
@@ -174,6 +175,31 @@ using InvokeWithValueRangeResultT = std::invoke_result_t<Fn, mlir::ValueRange>;
 
 } // namespace detail
 
+template <typename RewriterT>
+inline mlir::Value createSpatConcat(RewriterT& rewriter, mlir::Location loc, int64_t axis, mlir::ValueRange inputs) {
+  assert(!inputs.empty() && "spat.concat requires at least one input");
+  if (inputs.size() == 1)
+    return inputs.front();
+
+  auto firstType = mlir::cast<mlir::RankedTensorType>(inputs.front().getType());
+  auto outputShape = llvm::to_vector(firstType.getShape());
+  int64_t concatDimSize = 0;
+  bool concatDimDynamic = false;
+
+  for (mlir::Value input : inputs) {
+    auto inputType = mlir::cast<mlir::RankedTensorType>(input.getType());
+    assert(inputType.getRank() == firstType.getRank() && "spat.concat expects same-rank inputs");
+    if (mlir::ShapedType::isDynamic(inputType.getDimSize(axis)))
+      concatDimDynamic = true;
+    else
+      concatDimSize += inputType.getDimSize(axis);
+  }
+
+  outputShape[axis] = concatDimDynamic ? mlir::ShapedType::kDynamic : concatDimSize;
+  auto outputType = mlir::RankedTensorType::get(outputShape, firstType.getElementType(), firstType.getEncoding());
+  return spatial::SpatConcatOp::create(rewriter, loc, outputType, rewriter.getI64IntegerAttr(axis), inputs).getOutput();
+}
+
 template <size_t NumInputs, typename RewriterT, typename BodyFn>
 auto createSpatCompute(RewriterT& rewriter,
                        mlir::Location loc,
@@ -182,7 +208,7 @@ auto createSpatCompute(RewriterT& rewriter,
                        mlir::ValueRange inputs,
                        BodyFn&& body) {
   assert(inputs.size() == NumInputs && "NumInputs must match the number of input values");
-  auto computeOp = spatial::SpatWeightedCompute::create(rewriter, loc, resultTypes, weights, inputs);
+  auto computeOp = spatial::SpatCompute::create(rewriter, loc, resultTypes, weights, inputs);
 
   auto* block = new mlir::Block();
   for (mlir::Value input : inputs)
@@ -198,10 +224,10 @@ auto createSpatCompute(RewriterT& rewriter,
     if (mlir::failed(bodyResult)) {
       rewriter.setInsertionPointAfter(computeOp);
       rewriter.eraseOp(computeOp);
-      return mlir::FailureOr<spatial::SpatWeightedCompute>(mlir::failure());
+      return mlir::FailureOr<spatial::SpatCompute>(mlir::failure());
     }
     rewriter.setInsertionPointAfter(computeOp);
-    return mlir::FailureOr<spatial::SpatWeightedCompute>(computeOp);
+    return mlir::FailureOr<spatial::SpatCompute>(computeOp);
   }
   else {
     static_assert(std::is_same_v<BodyResult, void>, "createSpatCompute body must return void or mlir::LogicalResult");
@@ -219,7 +245,7 @@ auto createSpatCompute(RewriterT& rewriter,
                        mlir::ValueRange weights,
                        mlir::ValueRange inputs,
                        BodyFn&& body) {
-  auto computeOp = spatial::SpatWeightedCompute::create(rewriter, loc, resultTypes, weights, inputs);
+  auto computeOp = spatial::SpatCompute::create(rewriter, loc, resultTypes, weights, inputs);
 
   auto* block = new mlir::Block();
   for (mlir::Value input : inputs)
@@ -234,10 +260,10 @@ auto createSpatCompute(RewriterT& rewriter,
     if (mlir::failed(bodyResult)) {
       rewriter.setInsertionPointAfter(computeOp);
       rewriter.eraseOp(computeOp);
-      return mlir::FailureOr<spatial::SpatWeightedCompute>(mlir::failure());
+      return mlir::FailureOr<spatial::SpatCompute>(mlir::failure());
     }
     rewriter.setInsertionPointAfter(computeOp);
-    return mlir::FailureOr<spatial::SpatWeightedCompute>(computeOp);
+    return mlir::FailureOr<spatial::SpatCompute>(computeOp);
   }
   else {
     static_assert(std::is_same_v<BodyResult, void>, "createSpatCompute body must return void or mlir::LogicalResult");

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp

@@ -1,3 +1,4 @@
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
@@ -8,10 +9,10 @@
 #include "mlir/Transforms/Passes.h"
 
 #include "llvm/ADT/STLExtras.h"
-#include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/raw_os_ostream.h"
 
 #include <fstream>
@@ -24,8 +25,6 @@
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/DCPAnalysis.hpp"
-#include "src/Accelerators/PIM/Pass/PIMPasses.h"
 #include "src/Compiler/CompilerOptions.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
@@ -52,12 +51,48 @@ struct ONNXToSpatialPass : PassWrapper<ONNXToSpatialPass, OperationPass<ModuleOp
 private:
   void annotateWeightsConstants(func::FuncOp funcOp) const;
   void encapsulateGlobalInstruction(func::FuncOp funcOp);
-  void mergeTriviallyConnectedComputes(func::FuncOp funcOp);
   LogicalResult promoteConstantInputsToWeights(func::FuncOp funcOp);
 };
 
 } // namespace
 
+static void foldSingleLaneComputeBatches(func::FuncOp funcOp) {
+  IRRewriter rewriter(funcOp.getContext());
+  SmallVector<spatial::SpatComputeBatch> batchOps;
+  funcOp.walk([&](spatial::SpatComputeBatch batchOp) { batchOps.push_back(batchOp); });
+
+  for (auto batchOp : batchOps) {
+    if (batchOp.getLaneCount() != 1)
+      continue;
+
+    auto loc = batchOp.getLoc();
+    rewriter.setInsertionPoint(batchOp);
+    auto computeOp = spatial::SpatCompute::create(rewriter, loc, batchOp.getResultTypes(), batchOp.getWeights(), batchOp.getInputs());
+    computeOp.getProperties().setOperandSegmentSizes(
+      {static_cast<int>(batchOp.getWeights().size()), static_cast<int>(batchOp.getInputs().size())});
+
+    Block& templateBlock = batchOp.getBody().front();
+    SmallVector<Type> blockArgTypes;
+    SmallVector<Location> blockArgLocs;
+    for (BlockArgument arg : templateBlock.getArguments()) {
+      blockArgTypes.push_back(arg.getType());
+      blockArgLocs.push_back(loc);
+    }
+    auto* newBlock = rewriter.createBlock(
+      &computeOp.getBody(), computeOp.getBody().end(), TypeRange(blockArgTypes), blockArgLocs);
+
+    IRMapping mapper;
+    for (auto [oldArg, newArg] : llvm::zip(templateBlock.getArguments(), newBlock->getArguments()))
+      mapper.map(oldArg, newArg);
+    rewriter.setInsertionPointToEnd(newBlock);
+    for (Operation& op : templateBlock)
+      rewriter.clone(op, mapper);
+
+    batchOp.replaceAllUsesWith(computeOp.getResults());
+    rewriter.eraseOp(batchOp);
+  }
+}
+
 void ONNXToSpatialPass::runOnOperation() {
   ModuleOp moduleOp = getOperation();
   MLIRContext* ctx = &getContext();
@@ -87,8 +122,7 @@ void ONNXToSpatialPass::runOnOperation() {
                          tensor::TensorDialect,
                          arith::ArithDialect,
                          scf::SCFDialect>();
-  target.addDynamicallyLegalOp<ONNXMatMulOp>(
-    [](ONNXMatMulOp op) { return cast<ShapedType>(op.getY().getType()).getRank() != 2; });
+  target.addIllegalOp<ONNXMatMulOp>();
   target.addIllegalOp<ONNXAddOp>();
   target.addIllegalOp<ONNXDivOp>();
   target.addIllegalOp<ONNXMulOp>();
@@ -129,11 +163,13 @@ void ONNXToSpatialPass::runOnOperation() {
     return;
   }
 
+  foldSingleLaneComputeBatches(*entryFunc);
+
   // Count the number of compute ops and check they do not exceed the core count
   if (coresCount != -1) {
     int computeOpsCount = 0;
     for (auto& op : entryFunc->getFunctionBody().front().getOperations())
-      if (isa<spatial::SpatWeightedCompute>(op))
+      if (isa<spatial::SpatCompute>(op))
         computeOpsCount++;
 
     if (computeOpsCount > coresCount) {
@@ -149,6 +185,7 @@ void ONNXToSpatialPass::runOnOperation() {
     llvm::dbgs() << "Failed to run canonicalization cleanup, continuing...\n";
 
   annotateWeightsConstants(*entryFunc);
+
   encapsulateGlobalInstruction(*entryFunc);
 
   if (failed(promoteConstantInputsToWeights(*entryFunc))) {
@@ -156,8 +193,6 @@ void ONNXToSpatialPass::runOnOperation() {
     return;
   }
 
-  mergeTriviallyConnectedComputes(*entryFunc);
-
   // Dump to file for debug
   dumpModule(moduleOp, "spatial0");
 }
@@ -167,19 +202,36 @@ bool encapsulator(IRRewriter& rewriter, Location loc, Operation* inst, std::func
   if (T toRemoveOp = llvm::dyn_cast_if_present<T>(inst)) {
     Value source = funcSource(toRemoveOp);
     rewriter.setInsertionPointAfter(toRemoveOp);
-    if (isa_and_present<spatial::SpatWeightedCompute>(source.getDefiningOp())) {
-      auto newCompute = spatial::SpatWeightedCompute::create(rewriter, loc, inst->getResultTypes().front(), source);
-      auto BB = rewriter.createBlock(&newCompute.getBody(), newCompute.getBody().end(), {source.getType()}, {loc});
-      newCompute.getProperties().setOperandSegmentSizes({(int) 0, (int) 1});
-      rewriter.setInsertionPointToEnd(BB);
-      IRMapping mapper;
-      mapper.map(source, BB->getArgument(0));
-      auto newInst = rewriter.clone(*inst, mapper);
-      spatial::SpatYieldOp::create(rewriter, loc, newInst->getResult(0));
-      inst->replaceAllUsesWith(newCompute);
-      inst->erase();
-      return true;
-    }
+    auto newCompute = spatial::SpatCompute::create(rewriter, loc, inst->getResultTypes(), source);
+    auto BB = rewriter.createBlock(&newCompute.getBody(), newCompute.getBody().end(), {source.getType()}, {loc});
+    newCompute.getProperties().setOperandSegmentSizes({(int) 0, (int) 1});
+    rewriter.setInsertionPointToEnd(BB);
+    IRMapping mapper;
+    mapper.map(source, BB->getArgument(0));
+    auto newInst = rewriter.clone(*inst, mapper);
+    spatial::SpatYieldOp::create(rewriter, loc, newInst->getResults());
+    inst->replaceAllUsesWith(newCompute->getResults());
+    inst->erase();
+    return true;
+  }
+  return false;
+}
+
+bool encapsulateSlice(IRRewriter& rewriter, Location loc, Operation* inst) {
+  if (tensor::ExtractSliceOp toRemoveOp = llvm::dyn_cast_if_present<tensor::ExtractSliceOp>(inst)) {
+    auto source = toRemoveOp.getSource();
+    rewriter.setInsertionPointAfter(toRemoveOp);
+    auto newCompute = spatial::SpatCompute::create(rewriter, loc, inst->getResultTypes(), source);
+    auto BB = rewriter.createBlock(&newCompute.getBody(), newCompute.getBody().end(), {source.getType()}, {loc});
+    newCompute.getProperties().setOperandSegmentSizes({(int) 0, (int) 1});
+    rewriter.setInsertionPointToEnd(BB);
+    IRMapping mapper;
+    mapper.map(source, BB->getArgument(0));
+    auto newInst = rewriter.clone(*inst, mapper);
+    spatial::SpatYieldOp::create(rewriter, loc, newInst->getResults());
+    inst->replaceAllUsesWith(newCompute->getResults());
+    inst->erase();
+    return true;
   }
   return false;
 }
@@ -188,9 +240,9 @@ bool encapsulateConcat(IRRewriter& rewriter, Location loc, Operation* inst) {
   if (auto toRemoveOp = llvm::dyn_cast_if_present<tensor::ConcatOp>(inst)) {
     auto sources = toRemoveOp.getInputs();
     rewriter.setInsertionPointAfter(toRemoveOp);
-    if (llvm::any_of(
-          sources, [](auto source) { return isa_and_present<spatial::SpatWeightedCompute>(source.getDefiningOp()); })) {
-      auto newCompute = spatial::SpatWeightedCompute::create(rewriter, loc, inst->getResultTypes().front(), sources);
+    if (llvm::any_of(sources,
+                     [](auto source) { return isa_and_present<spatial::SpatCompute>(source.getDefiningOp()); })) {
+      auto newCompute = spatial::SpatCompute::create(rewriter, loc, inst->getResultTypes(), sources);
       SmallVector<Type> sourceTypes;
       SmallVector<Location> sourceLoc;
       for (auto source : sources) {
@@ -203,12 +255,34 @@ bool encapsulateConcat(IRRewriter& rewriter, Location loc, Operation* inst) {
       IRMapping mapper;
       for (auto [source, bbArg] : llvm::zip(sources, BB->getArguments()))
         mapper.map(source, bbArg);
-      auto newConcat = rewriter.clone(*inst, mapper);
-      spatial::SpatYieldOp::create(rewriter, loc, newConcat->getResult(0));
-      inst->replaceAllUsesWith(newCompute);
+      auto newConcat = spatial::SpatConcatOp::create(rewriter,
+                                                     loc,
+                                                     toRemoveOp.getType(),
+                                                     rewriter.getI64IntegerAttr(toRemoveOp.getDim()),
+                                                     ValueRange(BB->getArguments()));
+      spatial::SpatYieldOp::create(rewriter, loc, newConcat.getOutput());
+      inst->replaceAllUsesWith(newCompute->getResults());
       inst->erase();
       return true;
     }
+    auto newCompute = spatial::SpatCompute::create(rewriter, loc, inst->getResultTypes(), sources);
+    SmallVector<Type> sourceTypes;
+    SmallVector<Location> sourceLoc;
+    for (auto source : sources) {
+      sourceTypes.push_back(source.getType());
+      sourceLoc.push_back(loc);
+    }
+    auto BB = rewriter.createBlock(&newCompute.getBody(), newCompute.getBody().end(), sourceTypes, sourceLoc);
+    newCompute.getProperties().setOperandSegmentSizes({(int) 0, (int) sources.size()});
+    rewriter.setInsertionPointToEnd(BB);
+    IRMapping mapper;
+    for (auto [source, bbArg] : llvm::zip(sources, BB->getArguments()))
+      mapper.map(source, bbArg);
+    auto newConcat = rewriter.clone(*inst, mapper);
+    spatial::SpatYieldOp::create(rewriter, loc, newConcat->getResults());
+    inst->replaceAllUsesWith(newCompute->getResults());
+    inst->erase();
+    return true;
   }
   return false;
 }
@@ -270,6 +344,89 @@ static FailureOr<Value> materializeWeightLikeValueInBlock(Value value, IRRewrite
   return cast<Value>(mapped);
 }
 
+bool sourceOpernadHasWeightAlways(Operation* op) {
+  if (op == nullptr)
+    return false;
+
+  Operation* source = nullptr;
+  do {
+
+    if (isa<spatial::SpatCompute, spatial::SpatComputeBatch>(*op)) {
+      return false;
+    }
+    else if (auto extractSliceOp = dyn_cast<tensor::ExtractSliceOp>(*op)) {
+      auto tmpSource = extractSliceOp.getSource();
+      auto definingOp = tmpSource.getDefiningOp();
+      if (definingOp)
+        op = definingOp;
+      else
+        return false;
+    }
+    else if (auto extractRowsOp = dyn_cast<spatial::SpatExtractRowsOp>(*op)) {
+      auto tmpSource = extractRowsOp.getInput();
+      auto definingOp = tmpSource.getDefiningOp();
+      if (definingOp)
+        op = definingOp;
+      else
+        return false;
+    }
+    else if (auto expandShapeOp = dyn_cast<tensor::ExpandShapeOp>(*op)) {
+      auto tmpSource = expandShapeOp.getSrc();
+      auto definingOp = tmpSource.getDefiningOp();
+      if (definingOp)
+        op = definingOp;
+      else
+        return false;
+    }
+    else if (auto transposeOp = dyn_cast<ONNXTransposeOp>(*op)) {
+      auto tmpSource = transposeOp.getData();
+      auto definingOp = tmpSource.getDefiningOp();
+      if (definingOp)
+        op = definingOp;
+      else
+        return false;
+    }
+    else if (auto collapseShapeOp = dyn_cast<tensor::CollapseShapeOp>(*op)) {
+      auto tmpSource = collapseShapeOp.getSrc();
+      auto definingOp = tmpSource.getDefiningOp();
+      if (definingOp)
+        op = definingOp;
+      else
+        return false;
+    }
+    else if (auto constantOp = dyn_cast<arith::ConstantOp>(*op)) {
+      source = constantOp;
+    }
+    else if (auto concatOp = dyn_cast<tensor::ConcatOp>(*op)) {
+      bool res = false;
+      for (auto operand : concatOp.getOperands()) {
+        res |= hasWeightAlways(operand.getDefiningOp());
+        if (res)
+          return res;
+      }
+      return res;
+    }
+    else if (auto concatOp = dyn_cast<spatial::SpatConcatOp>(*op)) {
+      bool res = false;
+      for (auto operand : concatOp.getOperands()) {
+        res |= hasWeightAlways(operand.getDefiningOp());
+        if (res)
+          return res;
+      }
+      return res;
+    }
+    else {
+      op->dump();
+      llvm_unreachable("Global instruction not handle in func");
+    }
+  }
+  while (source == nullptr);
+
+  if (hasWeightAlways(source))
+    return true;
+  return false;
+}
+
 // TODO what we want to keep in global?
 void ONNXToSpatialPass::encapsulateGlobalInstruction(func::FuncOp funcOp) {
   Location loc = funcOp.getLoc();
@@ -278,8 +435,14 @@ void ONNXToSpatialPass::encapsulateGlobalInstruction(func::FuncOp funcOp) {
   while (keep) {
     keep = false;
     for (auto& instruction : llvm::make_early_inc_range(funcOp.getOps())) {
-      keep |= encapsulator<tensor::ExtractSliceOp>(
-        rewriter, loc, &instruction, [](tensor::ExtractSliceOp extract) { return extract.getSource(); });
+
+      if (isa<spatial::SpatCompute, spatial::SpatComputeBatch, spatial::SpatConcatOp, spatial::SpatExtractRowsOp>(
+            instruction)
+          || isa<func::ReturnOp>(instruction)
+          || sourceOpernadHasWeightAlways(&instruction))
+        continue;
+
+      keep |= encapsulateSlice(rewriter, loc, &instruction);
 
       keep |= encapsulator<tensor::ExpandShapeOp>(
         rewriter, loc, &instruction, [](tensor::ExpandShapeOp expand) { return expand.getSrc(); });
@@ -295,106 +458,16 @@ void ONNXToSpatialPass::encapsulateGlobalInstruction(func::FuncOp funcOp) {
   }
 }
 
-void ONNXToSpatialPass::mergeTriviallyConnectedComputes(func::FuncOp funcOp) {
-  Location loc = funcOp.getLoc();
-  IRRewriter rewriter(&getContext());
-  SmallVector<spatial::SpatWeightedCompute> trivialComputes;
-  llvm::SmallSet<spatial::SpatWeightedCompute, 8> toErase;
-
-  for (auto compute : funcOp.getOps<spatial::SpatWeightedCompute>())
-    if (compute->hasOneUse()) {
-      auto user = dyn_cast<spatial::SpatWeightedCompute>(*compute->getUsers().begin());
-
-      if (user && user.getInputs().size() == 1)
-        trivialComputes.push_back(compute);
-    }
-
-  while (!trivialComputes.empty()) {
-    auto compute = trivialComputes.front();
-
-    if (compute.use_empty()) {
-      std::swap(trivialComputes.front(), trivialComputes.back());
-      trivialComputes.pop_back();
-      continue;
-    }
-    auto child = cast<spatial::SpatWeightedCompute>(*compute->getUsers().begin());
-
-    rewriter.setInsertionPointAfter(compute.getOperation());
-
-    auto newCompute =
-      spatial::SpatWeightedCompute::create(rewriter, loc, child.getResultTypes(), compute.getOperands());
-    newCompute.getProperties().setOperandSegmentSizes(
-      {static_cast<int>(compute.getWeights().size()), static_cast<int>(compute.getInputs().size())});
-
-    IRMapping mapper;
-    auto weightMutableIter = newCompute.getWeightsMutable();
-    for (auto weight : child.getWeights()) {
-      auto founded = llvm::find(newCompute.getWeights(), weight);
-      if (founded == newCompute.getWeights().end()) {
-        weightMutableIter.append(weight);
-        auto last = weightMutableIter.end();
-        last = std::prev(last, 1);
-        mapper.map(weight, last->get());
-      }
-      else {
-        mapper.map(weight, *founded);
-      }
-    }
-
-    compute.getBodyRegion().cloneInto(&newCompute.getBodyRegion(), mapper);
-    auto newTerminator = newCompute.getBody().front().getTerminator();
-    mapper.map(*child.getBody().front().getArguments().begin(), newTerminator->getOperand(0));
-    newTerminator->erase();
-    rewriter.setInsertionPoint(&newCompute.getBody().front(), newCompute.getBody().front().end());
-    for (auto& op : child.getBody().front()) {
-      auto newInst = rewriter.clone(op, mapper);
-
-      if (auto vmOp = llvm::dyn_cast<spatial::SpatWeightedMVMOp>(newInst)) {
-        auto oldIndex = vmOp.getWeightIndex();
-        auto newWeight = mapper.lookup(*std::next(child.getWeights().begin(), oldIndex));
-        auto newIndex = std::distance(newCompute.getWeights().begin(), llvm::find(newCompute.getWeights(), newWeight));
-        vmOp.setWeightIndex(newIndex);
-      }
-      if (auto vmOp = llvm::dyn_cast<spatial::SpatWeightedVMMOp>(newInst)) {
-        auto oldIndex = vmOp.getWeightIndex();
-        auto newWeight = mapper.lookup(*std::next(child.getWeights().begin(), oldIndex));
-        auto newIndex = std::distance(newCompute.getWeights().begin(), llvm::find(newCompute.getWeights(), newWeight));
-        vmOp.setWeightIndex(newIndex);
-      }
-    }
-
-    child.replaceAllUsesWith(newCompute);
-    toErase.insert(child);
-
-    std::swap(trivialComputes.front(), trivialComputes.back());
-    trivialComputes.pop_back();
-    toErase.insert(compute);
-
-    if (newCompute->hasOneUse()) {
-      auto user = dyn_cast<spatial::SpatWeightedCompute>(*newCompute->getUsers().begin());
-      if (user && user.getInputs().size() == 1)
-        trivialComputes.push_back(newCompute);
-    }
-  }
-
-  for (auto compute : toErase) {
-    compute.getResult(0).dropAllUses();
-    compute.erase();
-  }
-}
-
 void ONNXToSpatialPass::annotateWeightsConstants(func::FuncOp funcOp) const {
   funcOp.walk([&](arith::ConstantOp constantOp) {
-    bool isAlwaysWeight =
-      llvm::all_of(constantOp->getUsers(), [](auto user) -> bool { return isa<spatial::SpatWeightedCompute>(user); });
-    if (isAlwaysWeight)
+    if (hasOnlySpatialMvmVmmWeightUses(constantOp.getResult()))
       markWeightAlways(constantOp);
   });
 }
 
 LogicalResult ONNXToSpatialPass::promoteConstantInputsToWeights(func::FuncOp funcOp) {
   IRRewriter rewriter(&getContext());
-  SmallVector<spatial::SpatWeightedCompute> computes(funcOp.getOps<spatial::SpatWeightedCompute>());
+  SmallVector<spatial::SpatCompute> computes(funcOp.getOps<spatial::SpatCompute>());
 
   for (auto compute : computes) {
     SmallVector<bool> promoteInput(compute.getInputs().size(), false);
@@ -430,7 +503,7 @@ LogicalResult ONNXToSpatialPass::promoteConstantInputsToWeights(func::FuncOp fun
     }
 
     auto newCompute =
-      spatial::SpatWeightedCompute::create(rewriter, compute.getLoc(), compute.getResultTypes(), newWeights, newInputs);
+      spatial::SpatCompute::create(rewriter, compute.getLoc(), compute.getResultTypes(), newWeights, newInputs);
     auto* newBlock =
       rewriter.createBlock(&newCompute.getBody(), newCompute.getBody().end(), newInputTypes, newInputLocs);
     newCompute.getProperties().setOperandSegmentSizes(

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

@@ -147,230 +147,212 @@ static Value buildPackedBias(bool hasBias,
   return arith::ConstantOp::create(rewriter, loc, packedBiasType, packedBiasAttr).getResult();
 }
 
-static Value createIm2colCompute(Value x,
-                                 RankedTensorType xType,
-                                 RankedTensorType im2colType,
-                                 RankedTensorType rowType,
-                                 int64_t batchSize,
-                                 int64_t numChannelsIn,
-                                 int64_t xHeight,
-                                 int64_t xWidth,
-                                 int64_t wHeight,
-                                 int64_t wWidth,
-                                 int64_t padHeightBegin,
-                                 int64_t padHeightEnd,
-                                 int64_t padWidthBegin,
-                                 int64_t padWidthEnd,
-                                 int64_t strideHeight,
-                                 int64_t strideWidth,
-                                 int64_t dilationHeight,
-                                 int64_t dilationWidth,
-                                 int64_t outWidth,
-                                 int64_t patchSize,
-                                 int64_t numPatches,
-                                 int64_t numPatchesPerBatch,
-                                 ConversionPatternRewriter& rewriter,
-                                 Location loc) {
+static Value createIm2colRowComputes(Value x,
+                                     RankedTensorType xType,
+                                     RankedTensorType im2colType,
+                                     RankedTensorType im2colRowType,
+                                     RankedTensorType gemmInputRowsType,
+                                     int64_t batchSize,
+                                     int64_t numChannelsIn,
+                                     int64_t xHeight,
+                                     int64_t xWidth,
+                                     int64_t wHeight,
+                                     int64_t wWidth,
+                                     int64_t padHeightBegin,
+                                     int64_t padHeightEnd,
+                                     int64_t padWidthBegin,
+                                     int64_t padWidthEnd,
+                                     int64_t strideHeight,
+                                     int64_t strideWidth,
+                                     int64_t dilationHeight,
+                                     int64_t dilationWidth,
+                                     int64_t outWidth,
+                                     int64_t patchSize,
+                                     int64_t numPatches,
+                                     int64_t numPatchesPerBatch,
+                                     int64_t packFactor,
+                                     ConversionPatternRewriter& rewriter,
+                                     Location loc) {
   auto elemType = xType.getElementType();
   constexpr size_t numInputs = 1;
-  auto im2colComputeOp = createSpatCompute<numInputs>(rewriter, loc, im2colType, {}, x, [&](Value xArg) {
-    Value paddedInput = xArg;
+  const int64_t packedNumRows = ceilIntegerDivide(numPatches, packFactor);
+  auto im2colComputeOp =
+    createSpatCompute<numInputs>(rewriter, loc, TypeRange {gemmInputRowsType}, {}, x, [&](Value xArg) {
+      Value paddedInput = xArg;
 
-    // Pad input with zeros if needed:
-    // [1, numChannelsIn, xHeight, xWidth] -> [1, numChannelsIn, xHeight+padHeight, xWidth+padWidth]
-    if (padHeightBegin || padHeightEnd || padWidthBegin || padWidthEnd) {
-      const int64_t paddedHeight = xHeight + padHeightBegin + padHeightEnd;
-      const int64_t paddedWidth = xWidth + padWidthBegin + padWidthEnd;
-      auto paddedType = RankedTensorType::get({batchSize, numChannelsIn, paddedHeight, paddedWidth}, elemType);
-      SmallVector<OpFoldResult> lowPads = {rewriter.getIndexAttr(0),
-                                           rewriter.getIndexAttr(0),
-                                           rewriter.getIndexAttr(padHeightBegin),
-                                           rewriter.getIndexAttr(padWidthBegin)};
-      SmallVector<OpFoldResult> highPads = {rewriter.getIndexAttr(0),
-                                            rewriter.getIndexAttr(0),
-                                            rewriter.getIndexAttr(padHeightEnd),
-                                            rewriter.getIndexAttr(padWidthEnd)};
-      auto padOp = tensor::PadOp::create(rewriter, loc, paddedType, paddedInput, lowPads, highPads);
-      auto* padBlock = new Block();
-      for (int i = 0; i < 4; i++)
-        padBlock->addArgument(rewriter.getIndexType(), loc);
-      padOp.getRegion().push_back(padBlock);
-      rewriter.setInsertionPointToStart(padBlock);
-      auto zero = arith::ConstantOp::create(rewriter, loc, elemType, rewriter.getFloatAttr(elemType, 0.0));
-      tensor::YieldOp::create(rewriter, loc, zero.getResult());
-      rewriter.setInsertionPointAfter(padOp);
-      paddedInput = padOp.getResult();
-    }
+      // Pad input with zeros if needed:
+      // [1, numChannelsIn, xHeight, xWidth] -> [1, numChannelsIn, xHeight+padHeight, xWidth+padWidth]
+      if (padHeightBegin || padHeightEnd || padWidthBegin || padWidthEnd) {
+        const int64_t paddedHeight = xHeight + padHeightBegin + padHeightEnd;
+        const int64_t paddedWidth = xWidth + padWidthBegin + padWidthEnd;
+        auto paddedType = RankedTensorType::get({batchSize, numChannelsIn, paddedHeight, paddedWidth}, elemType);
+        SmallVector<OpFoldResult> lowPads = {rewriter.getIndexAttr(0),
+                                             rewriter.getIndexAttr(0),
+                                             rewriter.getIndexAttr(padHeightBegin),
+                                             rewriter.getIndexAttr(padWidthBegin)};
+        SmallVector<OpFoldResult> highPads = {rewriter.getIndexAttr(0),
+                                              rewriter.getIndexAttr(0),
+                                              rewriter.getIndexAttr(padHeightEnd),
+                                              rewriter.getIndexAttr(padWidthEnd)};
+        auto padOp = tensor::PadOp::create(rewriter, loc, paddedType, paddedInput, lowPads, highPads);
+        auto* padBlock = new Block();
+        for (int i = 0; i < 4; i++)
+          padBlock->addArgument(rewriter.getIndexType(), loc);
+        padOp.getRegion().push_back(padBlock);
+        rewriter.setInsertionPointToStart(padBlock);
+        auto zero = arith::ConstantOp::create(rewriter, loc, elemType, rewriter.getFloatAttr(elemType, 0.0));
+        tensor::YieldOp::create(rewriter, loc, zero.getResult());
+        rewriter.setInsertionPointAfter(padOp);
+        paddedInput = padOp.getResult();
+      }
 
-    // Build im2col [numPatches, patchSize] incrementally to keep the IR small
-    // until the late PIM unrolling step.
-    Value im2colInit = tensor::EmptyOp::create(rewriter, loc, im2colType.getShape(), elemType);
-    auto c0 = arith::ConstantIndexOp::create(rewriter, loc, 0);
-    auto c1 = arith::ConstantIndexOp::create(rewriter, loc, 1);
-    auto cNumPatches = arith::ConstantIndexOp::create(rewriter, loc, numPatches);
-    auto cNumPatchesPerBatch = arith::ConstantIndexOp::create(rewriter, loc, numPatchesPerBatch);
-    auto cOutWidth = arith::ConstantIndexOp::create(rewriter, loc, outWidth);
-    auto cStrideHeight = arith::ConstantIndexOp::create(rewriter, loc, strideHeight);
-    auto cStrideWidth = arith::ConstantIndexOp::create(rewriter, loc, strideWidth);
+      // Build im2col [numPatches, patchSize] incrementally to keep the IR small
+      // until the late PIM unrolling step.
+      Value im2colInit = tensor::EmptyOp::create(rewriter, loc, im2colType.getShape(), elemType);
+      auto c0 = arith::ConstantIndexOp::create(rewriter, loc, 0);
+      auto c1 = arith::ConstantIndexOp::create(rewriter, loc, 1);
+      auto cNumPatches = arith::ConstantIndexOp::create(rewriter, loc, numPatches);
+      auto cNumPatchesPerBatch = arith::ConstantIndexOp::create(rewriter, loc, numPatchesPerBatch);
+      auto cOutWidth = arith::ConstantIndexOp::create(rewriter, loc, outWidth);
+      auto cStrideHeight = arith::ConstantIndexOp::create(rewriter, loc, strideHeight);
+      auto cStrideWidth = arith::ConstantIndexOp::create(rewriter, loc, strideWidth);
 
-    auto im2colLoop = scf::ForOp::create(rewriter, loc, c0, cNumPatches, c1, ValueRange {im2colInit});
-    rewriter.setInsertionPointToStart(im2colLoop.getBody());
+      auto im2colLoop = scf::ForOp::create(rewriter, loc, c0, cNumPatches, c1, ValueRange {im2colInit});
+      rewriter.setInsertionPointToStart(im2colLoop.getBody());
 
-    Value patchIndex = im2colLoop.getInductionVar();
-    Value im2colAcc = im2colLoop.getRegionIterArgs().front();
+      Value patchIndex = im2colLoop.getInductionVar();
+      Value im2colAcc = im2colLoop.getRegionIterArgs().front();
 
-    Value batchIndex = arith::DivUIOp::create(rewriter, loc, patchIndex, cNumPatchesPerBatch);
-    Value batchPatchIndex = arith::RemUIOp::create(rewriter, loc, patchIndex, cNumPatchesPerBatch);
-    Value outHeightIndex = arith::DivUIOp::create(rewriter, loc, batchPatchIndex, cOutWidth);
-    Value outWidthIndex = arith::RemUIOp::create(rewriter, loc, batchPatchIndex, cOutWidth);
-    Value inputHeightOffset = arith::MulIOp::create(rewriter, loc, outHeightIndex, cStrideHeight);
-    Value inputWidthOffset = arith::MulIOp::create(rewriter, loc, outWidthIndex, cStrideWidth);
+      Value batchIndex = arith::DivUIOp::create(rewriter, loc, patchIndex, cNumPatchesPerBatch);
+      Value batchPatchIndex = arith::RemUIOp::create(rewriter, loc, patchIndex, cNumPatchesPerBatch);
+      Value outHeightIndex = arith::DivUIOp::create(rewriter, loc, batchPatchIndex, cOutWidth);
+      Value outWidthIndex = arith::RemUIOp::create(rewriter, loc, batchPatchIndex, cOutWidth);
+      Value inputHeightOffset = arith::MulIOp::create(rewriter, loc, outHeightIndex, cStrideHeight);
+      Value inputWidthOffset = arith::MulIOp::create(rewriter, loc, outWidthIndex, cStrideWidth);
 
-    SmallVector<OpFoldResult> offsets = {batchIndex, rewriter.getIndexAttr(0), inputHeightOffset, inputWidthOffset};
-    SmallVector<OpFoldResult> sizes = {rewriter.getIndexAttr(1),
-                                       rewriter.getIndexAttr(numChannelsIn),
-                                       rewriter.getIndexAttr(wHeight),
-                                       rewriter.getIndexAttr(wWidth)};
-    SmallVector<OpFoldResult> strides = {rewriter.getIndexAttr(1),
-                                         rewriter.getIndexAttr(1),
-                                         rewriter.getIndexAttr(dilationHeight),
-                                         rewriter.getIndexAttr(dilationWidth)};
-    auto patchType = RankedTensorType::get({1, numChannelsIn, wHeight, wWidth}, elemType);
-    Value patch = tensor::ExtractSliceOp::create(rewriter, loc, patchType, paddedInput, offsets, sizes, strides);
+      SmallVector<OpFoldResult> offsets = {batchIndex, rewriter.getIndexAttr(0), inputHeightOffset, inputWidthOffset};
+      SmallVector<OpFoldResult> sizes = {rewriter.getIndexAttr(1),
+                                         rewriter.getIndexAttr(numChannelsIn),
+                                         rewriter.getIndexAttr(wHeight),
+                                         rewriter.getIndexAttr(wWidth)};
+      SmallVector<OpFoldResult> strides = {rewriter.getIndexAttr(1),
+                                           rewriter.getIndexAttr(1),
+                                           rewriter.getIndexAttr(dilationHeight),
+                                           rewriter.getIndexAttr(dilationWidth)};
+      auto patchType = RankedTensorType::get({1, numChannelsIn, wHeight, wWidth}, elemType);
+      Value patch = tensor::ExtractSliceOp::create(rewriter, loc, patchType, paddedInput, offsets, sizes, strides);
 
-    Value row = tensor::CollapseShapeOp::create(rewriter,
-                                                loc,
-                                                rowType,
-                                                patch,
-                                                SmallVector<ReassociationIndices> {
-                                                  {0},
-                                                  {1, 2, 3}
+      Value row = tensor::CollapseShapeOp::create(rewriter,
+                                                  loc,
+                                                  im2colRowType,
+                                                  patch,
+                                                  SmallVector<ReassociationIndices> {
+                                                    {0},
+                                                    {1, 2, 3}
+      });
+
+      SmallVector<OpFoldResult> rowOffsets = {patchIndex, rewriter.getIndexAttr(0)};
+      SmallVector<OpFoldResult> rowSizes = {rewriter.getIndexAttr(1), rewriter.getIndexAttr(patchSize)};
+      SmallVector<OpFoldResult> rowStrides = {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
+      Value updatedIm2col =
+        tensor::InsertSliceOp::create(rewriter, loc, row, im2colAcc, rowOffsets, rowSizes, rowStrides);
+      scf::YieldOp::create(rewriter, loc, updatedIm2col);
+
+      rewriter.setInsertionPointAfter(im2colLoop);
+      Value im2col = im2colLoop.getResult(0);
+
+      Value gemmInputRows = im2col;
+      if (packFactor != 1) {
+        const int64_t paddedNumPatches = packedNumRows * packFactor;
+        auto groupedType = RankedTensorType::get({packedNumRows, packFactor, patchSize}, elemType);
+        auto packedType = RankedTensorType::get({packedNumRows, packFactor * patchSize}, elemType);
+        Value paddedIm2col = createPaddedRows(im2col, im2colType, paddedNumPatches, rewriter, loc);
+        Value groupedIm2col = tensor::ExpandShapeOp::create(rewriter,
+                                                            loc,
+                                                            groupedType,
+                                                            paddedIm2col,
+                                                            SmallVector<ReassociationIndices> {
+                                                              {0, 1},
+                                                              {2}
+        });
+        gemmInputRows = tensor::CollapseShapeOp::create(rewriter,
+                                                        loc,
+                                                        packedType,
+                                                        groupedIm2col,
+                                                        SmallVector<ReassociationIndices> {
+                                                          {0},
+                                                          {1, 2}
+        });
+      }
+
+      spatial::SpatYieldOp::create(rewriter, loc, gemmInputRows);
     });
 
-    SmallVector<OpFoldResult> rowOffsets = {patchIndex, rewriter.getIndexAttr(0)};
-    SmallVector<OpFoldResult> rowSizes = {rewriter.getIndexAttr(1), rewriter.getIndexAttr(patchSize)};
-    SmallVector<OpFoldResult> rowStrides = {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-    Value updatedIm2col =
-      tensor::InsertSliceOp::create(rewriter, loc, row, im2colAcc, rowOffsets, rowSizes, rowStrides);
-    scf::YieldOp::create(rewriter, loc, updatedIm2col);
-
-    rewriter.setInsertionPointAfter(im2colLoop);
-    Value im2col = im2colLoop.getResult(0);
-    spatial::SpatYieldOp::create(rewriter, loc, im2col);
-  });
   return im2colComputeOp.getResult(0);
 }
 
-static Value createPackedIm2colRows(Value im2col,
-                                    RankedTensorType im2colType,
-                                    Type elemType,
-                                    int64_t numPatches,
-                                    int64_t patchSize,
-                                    int64_t packFactor,
-                                    ConversionPatternRewriter& rewriter,
-                                    Location loc) {
-  if (packFactor == 1)
-    return im2col;
-
-  const int64_t packedNumRows = ceilIntegerDivide(numPatches, packFactor);
-  const int64_t paddedNumPatches = packedNumRows * packFactor;
-  auto groupedType = RankedTensorType::get({packedNumRows, packFactor, patchSize}, elemType);
-  auto packedType = RankedTensorType::get({packedNumRows, packFactor * patchSize}, elemType);
-  auto packedComputeOp = createSpatCompute<1>(rewriter, loc, packedType, {}, im2col, [&](Value im2colArg) {
-    Value paddedIm2col = createPaddedRows(im2colArg, im2colType, paddedNumPatches, rewriter, loc);
-    Value groupedIm2col = tensor::ExpandShapeOp::create(rewriter,
-                                                        loc,
-                                                        groupedType,
-                                                        paddedIm2col,
-                                                        SmallVector<ReassociationIndices> {
-                                                          {0, 1},
-                                                          {2}
-    });
-    Value packedIm2col = tensor::CollapseShapeOp::create(rewriter,
-                                                         loc,
-                                                         packedType,
-                                                         groupedIm2col,
-                                                         SmallVector<ReassociationIndices> {
-                                                           {0},
-                                                           {1, 2}
-    });
-    spatial::SpatYieldOp::create(rewriter, loc, packedIm2col);
-  });
-  return packedComputeOp.getResult(0);
-}
-
-static Value createUnpackedOutput(Value packedOutput,
-                                  RankedTensorType gemmOutType,
-                                  RankedTensorType outType,
-                                  int64_t numPatches,
-                                  int64_t numChannelsOut,
-                                  int64_t packFactor,
-                                  ConversionPatternRewriter& rewriter,
-                                  Location loc) {
-  if (packFactor == 1)
-    return packedOutput;
-
-  const int64_t packedNumRows = ceilIntegerDivide(numPatches, packFactor);
-  const int64_t paddedNumPatches = packedNumRows * packFactor;
-  auto expandedType = RankedTensorType::get({packedNumRows, packFactor, numChannelsOut}, outType.getElementType());
-  auto paddedType = RankedTensorType::get({paddedNumPatches, numChannelsOut}, outType.getElementType());
-  auto unpackComputeOp = createSpatCompute<1>(rewriter, loc, gemmOutType, {}, packedOutput, [&](Value packedOutputArg) {
-    Value expandedOutput = tensor::ExpandShapeOp::create(rewriter,
-                                                         loc,
-                                                         expandedType,
-                                                         packedOutputArg,
-                                                         SmallVector<ReassociationIndices> {
-                                                           {0},
-                                                           {1, 2}
-    });
-    Value paddedOutput = tensor::CollapseShapeOp::create(rewriter,
-                                                         loc,
-                                                         paddedType,
-                                                         expandedOutput,
-                                                         SmallVector<ReassociationIndices> {
-                                                           {0, 1},
-                                                           {2}
-    });
-
-    Value unpackedOutput = paddedOutput;
-    if (paddedNumPatches != numPatches) {
-      SmallVector<OpFoldResult> offsets = {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
-      SmallVector<OpFoldResult> sizes = {rewriter.getIndexAttr(numPatches), rewriter.getIndexAttr(numChannelsOut)};
-      SmallVector<OpFoldResult> strides = {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-      unpackedOutput =
-        tensor::ExtractSliceOp::create(rewriter, loc, gemmOutType, paddedOutput, offsets, sizes, strides);
-    }
-
-    spatial::SpatYieldOp::create(rewriter, loc, unpackedOutput);
-  });
-  return unpackComputeOp.getResult(0);
-}
-
-static Value createCollectedConvOutput(Value gemmOut,
+static Value createCollectedConvOutput(ValueRange gemmRows,
                                        Type convType,
+                                       RankedTensorType gemmOutType,
                                        RankedTensorType nhwcType,
                                        RankedTensorType outType,
+                                       int64_t numPatches,
+                                       int64_t numChannelsOut,
+                                       int64_t packFactor,
                                        ConversionPatternRewriter& rewriter,
                                        Location loc) {
-  auto collectComputeOp =
-    createSpatCompute(rewriter, loc, convType, {}, ValueRange {gemmOut}, [&](ValueRange gemmOutArgs) {
-      Value gemmOutArg = gemmOutArgs.front();
-
-      // Restore to NCHW layout:
-      // [numPatches, numChannelsOut]
-      // -> [1, outHeight, outWidth, numChannelsOut]
-      // -> [1, numChannelsOut, outHeight, outWidth]
-      Value nhwcOut = tensor::ExpandShapeOp::create(rewriter,
-                                                    loc,
-                                                    nhwcType,
-                                                    gemmOutArg,
-                                                    SmallVector<ReassociationIndices> {
-                                                      {0, 1, 2},
-                                                      {3}
+  const int64_t packedNumRows = ceilIntegerDivide(numPatches, packFactor);
+  const int64_t paddedNumPatches = packedNumRows * packFactor;
+  auto collectComputeOp = createSpatCompute(rewriter, loc, convType, {}, gemmRows, [&](ValueRange gemmRowArgs) {
+    Value gemmOut;
+    if (packFactor == 1) {
+      gemmOut = createSpatConcat(rewriter, loc, /*axis=*/0, gemmRowArgs);
+    }
+    else {
+      auto expandedType = RankedTensorType::get({packedNumRows, packFactor, numChannelsOut}, outType.getElementType());
+      auto paddedType = RankedTensorType::get({paddedNumPatches, numChannelsOut}, outType.getElementType());
+      Value packedOutput = createSpatConcat(rewriter, loc, /*axis=*/0, gemmRowArgs);
+      Value expandedOutput = tensor::ExpandShapeOp::create(rewriter,
+                                                           loc,
+                                                           expandedType,
+                                                           packedOutput,
+                                                           SmallVector<ReassociationIndices> {
+                                                             {0},
+                                                             {1, 2}
       });
-      Value nchwOut = ONNXTransposeOp::create(rewriter, loc, outType, nhwcOut, rewriter.getI64ArrayAttr({0, 3, 1, 2}));
-      spatial::SpatYieldOp::create(rewriter, loc, nchwOut);
+      Value paddedOutput = tensor::CollapseShapeOp::create(rewriter,
+                                                           loc,
+                                                           paddedType,
+                                                           expandedOutput,
+                                                           SmallVector<ReassociationIndices> {
+                                                             {0, 1},
+                                                             {2}
+      });
+
+      gemmOut = paddedOutput;
+      if (paddedNumPatches != numPatches) {
+        SmallVector<OpFoldResult> offsets = {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
+        SmallVector<OpFoldResult> sizes = {rewriter.getIndexAttr(numPatches), rewriter.getIndexAttr(numChannelsOut)};
+        SmallVector<OpFoldResult> strides = {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
+        gemmOut = tensor::ExtractSliceOp::create(rewriter, loc, gemmOutType, paddedOutput, offsets, sizes, strides);
+      }
+    }
+
+    // Restore to NCHW layout:
+    // [numPatches, numChannelsOut]
+    // -> [1, outHeight, outWidth, numChannelsOut]
+    // -> [1, numChannelsOut, outHeight, outWidth]
+    Value nhwcOut = tensor::ExpandShapeOp::create(rewriter,
+                                                  loc,
+                                                  nhwcType,
+                                                  gemmOut,
+                                                  SmallVector<ReassociationIndices> {
+                                                    {0, 1, 2},
+                                                    {3}
     });
+    Value nchwOut = ONNXTransposeOp::create(rewriter, loc, outType, nhwcOut, rewriter.getI64ArrayAttr({0, 3, 1, 2}));
+    spatial::SpatYieldOp::create(rewriter, loc, nchwOut);
+  });
   return collectComputeOp.getResult(0);
 }
 
@@ -487,11 +469,11 @@ LogicalResult ConvToGemm::matchAndRewrite(ONNXConvOp convOp,
 
   // Pass bias through directly; Gemm handles rank-1 C canonicalization.
   bool hasB = !isa<ONNXNoneOp>(b.getDefiningOp());
-  Value gemmC = ONNXNoneOp::create(rewriter, loc, rewriter.getNoneType());
+  Value gemmBias = ONNXNoneOp::create(rewriter, loc, rewriter.getNoneType());
   Value biasMatrix;
   DenseElementsAttr biasDenseAttr;
   if (hasB) {
-    gemmC = b;
+    gemmBias = b;
     biasDenseAttr = getDenseConstantAttr(b);
     biasMatrix = expandBiasIfNeeded(b, rewriter, loc);
   }
@@ -500,94 +482,85 @@ LogicalResult ConvToGemm::matchAndRewrite(ONNXConvOp convOp,
   const int64_t effectiveMaxParallelPixels =
     (canPackWeightsAsConstants && canPackBiasAsConstants) ? maxParallelPixels : 1;
 
-  Value im2col = createIm2colCompute(x,
-                                     xType,
-                                     im2colType,
-                                     rowType,
-                                     batchSize,
-                                     numChannelsIn,
-                                     xHeight,
-                                     xWidth,
-                                     wHeight,
-                                     wWidth,
-                                     padHeightBegin,
-                                     padHeightEnd,
-                                     padWidthBegin,
-                                     padWidthEnd,
-                                     strideHeight,
-                                     strideWidth,
-                                     dilationHeight,
-                                     dilationWidth,
-                                     outWidth,
-                                     patchSize,
-                                     numPatches,
-                                     numPatchesPerBatch,
-                                     rewriter,
-                                     loc);
+  // Keep the standard im2col view of convolution:
+  //   A (im2col):  [numPatches, patchSize]  -- one row per output spatial position
+  //   B (weights): [patchSize, cOut]        -- W^T, stored in crossbar columns
+  // and optionally repack several old rows into one GEMM row to use the available crossbar size better.
+  //
+  // We want to process N pixels at the same time. Instead of doing N separate operations
+  // of (1 x patchSize) x (patchSize x cOut), we construct a block-diagonal weight matrix
+  // containing N copies of W^T and concatenate N im2col rows into one longer row:
+  //   A_packed: [ceil(numPatches / N), N * patchSize]
+  //   B_packed: [N * patchSize, N * cOut]
+  //   Y_packed: [ceil(numPatches / N), N * cOut]
+  const int64_t packedNumRows = ceilIntegerDivide(numPatches, effectiveMaxParallelPixels);
+  auto gemmInputRowsType = RankedTensorType::get({packedNumRows, effectiveMaxParallelPixels * patchSize}, elemType);
+  auto gemmOutputRowsType =
+    RankedTensorType::get({packedNumRows, effectiveMaxParallelPixels * numChannelsOut}, outType.getElementType());
+  Value gemmInputRows = createIm2colRowComputes(x,
+                                                xType,
+                                                im2colType,
+                                                rowType,
+                                                gemmInputRowsType,
+                                                batchSize,
+                                                numChannelsIn,
+                                                xHeight,
+                                                xWidth,
+                                                wHeight,
+                                                wWidth,
+                                                padHeightBegin,
+                                                padHeightEnd,
+                                                padWidthBegin,
+                                                padWidthEnd,
+                                                strideHeight,
+                                                strideWidth,
+                                                dilationHeight,
+                                                dilationWidth,
+                                                outWidth,
+                                                patchSize,
+                                                numPatches,
+                                                numPatchesPerBatch,
+                                                effectiveMaxParallelPixels,
+                                                rewriter,
+                                                loc);
 
-  Value gemmOut;
-  if (effectiveMaxParallelPixels == 1) {
-    // Fallback to the plain im2col GEMM when a single crossbar cannot fit multiple pixels.
-    gemmOut = ONNXGemmOp::create(rewriter,
-                                 loc,
-                                 gemmOutType,
-                                 im2col,
-                                 wTrans,
-                                 gemmC,
-                                 rewriter.getF32FloatAttr(1.0f),
-                                 rewriter.getF32FloatAttr(1.0f),
-                                 rewriter.getBoolAttr(false),
-                                 rewriter.getBoolAttr(false))
-                .getY();
-  }
-  else {
-    // Keep the standard im2col view of convolution:
-    //   A (im2col):  [numPatches, patchSize]  -- one row per output spatial position
-    //   B (weights): [patchSize, cOut]        -- W^T, stored in crossbar columns
-    // but repack several old rows into one new row so we use the available crossbar size better.
-    //
-    // We want to process N spatial pixels at the exact same time. Instead of doing N separate
-    // operations of (1 x patchSize) x (patchSize x cOut), we construct a block-diagonal weight matrix
-    // containing N copies of W^T and concatenate N im2col rows into one longer row:
-    //   A_packed: [ceil(numPatches / N), N * patchSize]
-    //   B_packed: [N * patchSize, N * cOut]
-    //   Y_packed: [ceil(numPatches / N), N * cOut]
-    // The downstream GemmToManyGemv pass still splits by row, but now there are fewer, longer rows.
-    const int64_t packedNumRows = ceilIntegerDivide(numPatches, effectiveMaxParallelPixels);
-    auto packedOutType =
-      RankedTensorType::get({packedNumRows, effectiveMaxParallelPixels * numChannelsOut}, outType.getElementType());
+  Value gemmB = buildPackedWeight(wDenseAttr,
+                                  wTrans,
+                                  wType,
+                                  numChannelsIn,
+                                  numChannelsOut,
+                                  wHeight,
+                                  wWidth,
+                                  patchSize,
+                                  effectiveMaxParallelPixels,
+                                  rewriter,
+                                  loc);
+  Value gemmC = buildPackedBias(
+    hasB, gemmBias, biasMatrix, biasDenseAttr, outType, numChannelsOut, effectiveMaxParallelPixels, rewriter, loc);
 
-    Value packedA = createPackedIm2colRows(
-      im2col, im2colType, elemType, numPatches, patchSize, effectiveMaxParallelPixels, rewriter, loc);
-    Value packedB = buildPackedWeight(wDenseAttr,
-                                      wTrans,
-                                      wType,
-                                      numChannelsIn,
-                                      numChannelsOut,
-                                      wHeight,
-                                      wWidth,
-                                      patchSize,
-                                      effectiveMaxParallelPixels,
-                                      rewriter,
-                                      loc);
-    Value packedC = buildPackedBias(
-      hasB, gemmC, biasMatrix, biasDenseAttr, outType, numChannelsOut, effectiveMaxParallelPixels, rewriter, loc);
-    Value packedOut = ONNXGemmOp::create(rewriter,
-                                         loc,
-                                         packedOutType,
-                                         packedA,
-                                         packedB,
-                                         packedC,
-                                         rewriter.getF32FloatAttr(1.0f),
-                                         rewriter.getF32FloatAttr(1.0f),
-                                         rewriter.getBoolAttr(false),
-                                         rewriter.getBoolAttr(false))
-                        .getY();
-    gemmOut = createUnpackedOutput(
-      packedOut, gemmOutType, outType, numPatches, numChannelsOut, effectiveMaxParallelPixels, rewriter, loc);
-  }
+  Value gemmRows = ONNXGemmOp::create(rewriter,
+                                      loc,
+                                      gemmOutputRowsType,
+                                      gemmInputRows,
+                                      gemmB,
+                                      gemmC,
+                                      rewriter.getF32FloatAttr(1.0f),
+                                      rewriter.getF32FloatAttr(1.0f),
+                                      rewriter.getBoolAttr(false),
+                                      rewriter.getBoolAttr(false))
+                     .getY();
 
-  rewriter.replaceOp(convOp, createCollectedConvOutput(gemmOut, convOp.getType(), nhwcType, outType, rewriter, loc));
+  rewriter.replaceOp(convOp,
+                     createCollectedConvOutput(ValueRange {gemmRows},
+                                               convOp.getType(),
+                                               gemmOutType,
+                                               nhwcType,
+                                               outType,
+                                               numPatches,
+                                               numChannelsOut,
+                                               effectiveMaxParallelPixels,
+                                               rewriter,
+                                               loc));
   return success();
 }
 

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

@@ -1,6 +1,7 @@
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/IRMapping.h"
 #include "mlir/IR/Location.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Transforms/DialectConversion.h"
@@ -65,6 +66,66 @@ struct GemvToSpatialCompute : OpConversionPattern<ONNXGemmOp> {
                                 ConversionPatternRewriter& rewriter) const override;
 };
 
+struct GemmToSpatialComputeBatch : OpConversionPattern<ONNXGemmOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(ONNXGemmOp gemmOp,
+                                ONNXGemmOpAdaptor gemmOpAdaptor,
+                                ConversionPatternRewriter& rewriter) const override;
+};
+
+static SmallVector<Value> materializeBatchRowSlices(Value matrix,
+                                                    RankedTensorType matrixType,
+                                                    ConversionPatternRewriter& rewriter,
+                                                    Location loc) {
+  const int64_t numRows = matrixType.getDimSize(0);
+  auto rowType = RankedTensorType::get({1, matrixType.getDimSize(1)}, matrixType.getElementType());
+  SmallVector<Type> resultTypes(static_cast<size_t>(numRows), rowType);
+
+  auto buildRowSlices = [&](Value matrixArg) {
+    auto extractRowsOp = spatial::SpatExtractRowsOp::create(rewriter, loc, TypeRange(resultTypes), matrixArg);
+    return SmallVector<Value>(extractRowsOp->result_begin(), extractRowsOp->result_end());
+  };
+
+  auto cloneBatchInputChainIntoSliceCompute =
+    [&](Value rootInput, SmallVector<Operation*> chainOps, Value rootValue) -> SmallVector<Value> {
+    auto sliceCompute =
+      createSpatCompute<1>(rewriter, loc, TypeRange(resultTypes), {}, ValueRange {rootInput}, [&](Value input) {
+        Value transformedMatrix = input;
+        if (!chainOps.empty()) {
+          IRMapping mapper;
+          mapper.map(rootValue, input);
+          for (Operation* chainOp : chainOps)
+            rewriter.clone(*chainOp, mapper);
+          transformedMatrix = cast<Value>(mapper.lookup(matrix));
+        }
+        spatial::SpatYieldOp::create(rewriter, loc, buildRowSlices(transformedMatrix));
+      });
+    SmallVector<Value> rowSlices(sliceCompute->result_begin(), sliceCompute->result_end());
+    return rowSlices;
+  };
+
+  SmallVector<Operation*> chainOps;
+  Value rootValue = matrix;
+  while (Operation* definingOp = rootValue.getDefiningOp()) {
+    if (auto rootCompute = dyn_cast<spatial::SpatCompute>(definingOp)) {
+      SmallVector<Operation*> reversedChainOps(chainOps.rbegin(), chainOps.rend());
+      return cloneBatchInputChainIntoSliceCompute(
+        rootCompute.getResult(cast<OpResult>(rootValue).getResultNumber()), reversedChainOps, rootValue);
+    }
+
+    if (definingOp->getNumOperands() != 1)
+      break;
+    if (!isa<tensor::ExtractSliceOp, tensor::ExpandShapeOp, tensor::CollapseShapeOp, ONNXTransposeOp>(definingOp))
+      break;
+
+    chainOps.push_back(definingOp);
+    rootValue = definingOp->getOperand(0);
+  }
+
+  return buildRowSlices(matrix);
+}
+
 } // namespace
 
 LogicalResult GemmToManyGemv::matchAndRewrite(ONNXGemmOp gemmOp,
@@ -156,8 +217,7 @@ LogicalResult GemmToManyGemv::matchAndRewrite(ONNXGemmOp gemmOp,
   }
 
   auto concatComputeOp = createSpatCompute(rewriter, loc, gemmOp.getType(), {}, gemvOps, [&](ValueRange gemvOpsArgs) {
-    auto concatOp = tensor::ConcatOp::create(rewriter, loc, /*axis=*/0, gemvOpsArgs);
-    spatial::SpatYieldOp::create(rewriter, loc, concatOp.getResult());
+    spatial::SpatYieldOp::create(rewriter, loc, createSpatConcat(rewriter, loc, /*axis=*/0, gemvOpsArgs));
   });
 
   rewriter.replaceOp(gemmOp, concatComputeOp);
@@ -313,15 +373,116 @@ LogicalResult GemvToSpatialCompute::matchAndRewrite(ONNXGemmOp gemmOp,
 
   auto concatComputeOp =
     createSpatCompute(rewriter, gemmLoc, gemmOp.getType(), {}, outHSlices, [&](ValueRange blockArgs) {
-      auto concatOp = tensor::ConcatOp::create(rewriter, gemmLoc, /*axis=*/1, blockArgs);
-      spatial::SpatYieldOp::create(rewriter, gemmLoc, concatOp.getResult());
+      spatial::SpatYieldOp::create(rewriter, gemmLoc, createSpatConcat(rewriter, gemmLoc, /*axis=*/1, blockArgs));
     });
 
   rewriter.replaceOp(gemmOp, concatComputeOp);
   return success();
 }
 
+LogicalResult GemmToSpatialComputeBatch::matchAndRewrite(ONNXGemmOp gemmOp,
+                                                         ONNXGemmOpAdaptor gemmOpAdaptor,
+                                                         ConversionPatternRewriter& rewriter) const {
+  Location loc = gemmOp.getLoc();
+  Value a = gemmOpAdaptor.getA();
+  Value b = gemmOpAdaptor.getB();
+  Value c = gemmOpAdaptor.getC();
+
+  assert("A should have been transposed already" && !gemmOpAdaptor.getTransA());
+
+  bool hasC = !isa<ONNXNoneOp>(c.getDefiningOp());
+
+  auto aType = cast<RankedTensorType>(a.getType());
+  auto bType = cast<RankedTensorType>(b.getType());
+  auto outType = cast<RankedTensorType>(gemmOp.getY().getType());
+  assert("Only support static shapes" && aType.hasStaticShape() && bType.hasStaticShape() && outType.hasStaticShape());
+
+  const int64_t numOutRows = aType.getDimSize(0);
+  if (numOutRows <= 1)
+    return failure();
+
+  // Only handle the single-tile case: K <= crossbarSize and N <= crossbarSize
+  if (aType.getDimSize(1) > static_cast<int64_t>(crossbarSize.getValue())
+      || outType.getDimSize(1) > static_cast<int64_t>(crossbarSize.getValue()))
+    return failure();
+
+  auto scaledB = materializeScaledConstantTensor(b, gemmOpAdaptor.getAlpha().convertToFloat(), rewriter, loc);
+  if (failed(scaledB))
+    return failure();
+  b = *scaledB;
+  bType = cast<RankedTensorType>(b.getType());
+
+  if (gemmOpAdaptor.getTransB()) {
+    auto bShape = bType.getShape();
+    auto transposedType = bType.cloneWith(ArrayRef({bShape[1], bShape[0]}), bType.getElementType());
+    b = ONNXTransposeOp::create(rewriter, loc, transposedType, b, rewriter.getI64ArrayAttr({1, 0}));
+    bType = cast<RankedTensorType>(b.getType());
+  }
+  (void) bType;
+
+  Value sharedBias;
+  if (hasC) {
+    auto scaledC = materializeScaledConstantTensor(c, gemmOpAdaptor.getBeta().convertToFloat(), rewriter, loc);
+    if (failed(scaledC))
+      return failure();
+    c = *scaledC;
+    auto cType = cast<RankedTensorType>(c.getType());
+    if (cType.getRank() == 1) {
+      auto expandedType = RankedTensorType::get({1, cType.getDimSize(0)}, cType.getElementType());
+      c = tensor::ExpandShapeOp::create(rewriter,
+                                        loc,
+                                        expandedType,
+                                        c,
+                                        SmallVector<ReassociationIndices> {
+                                          {0, 1}
+      });
+      cType = cast<RankedTensorType>(c.getType());
+    }
+    assert("Only support rank 2 tensor for C" && cType.getRank() == 2);
+    // Row-specific bias can't share a single template body; fall through to GemmToManyGemv
+    if (cType.getDimSize(0) == numOutRows && numOutRows > 1)
+      return failure();
+    if (cType.getDimSize(0) == 1 && cType.getDimSize(1) == 1)
+      c = broadcastToVector(c, outType.getDimSize(1), rewriter, loc);
+    sharedBias = c;
+  }
+
+  SmallVector<Value> aSlices = materializeBatchRowSlices(a, aType, rewriter, loc);
+  auto aSliceType = cast<RankedTensorType>(aSlices.front().getType());
+
+  auto outRowType = RankedTensorType::get({1, outType.getDimSize(1)}, outType.getElementType());
+  SmallVector<Type> resultTypes(static_cast<size_t>(numOutRows), outRowType);
+  SmallVector<Value> weights(static_cast<size_t>(numOutRows), b);
+
+  auto batchOp = spatial::SpatComputeBatch::create(rewriter,
+                                                   loc,
+                                                   TypeRange(resultTypes),
+                                                   rewriter.getI32IntegerAttr(static_cast<int32_t>(numOutRows)),
+                                                   ValueRange(weights),
+                                                   ValueRange(aSlices));
+
+  Block* body = rewriter.createBlock(
+    &batchOp.getBody(), batchOp.getBody().end(), TypeRange {aSliceType}, SmallVector<Location>(1, loc));
+  rewriter.setInsertionPointToEnd(body);
+
+  Value vmmResult = spatial::SpatWeightedVMMOp::create(rewriter, loc, outRowType, 0, body->getArgument(0)).getResult();
+  Value laneResult = vmmResult;
+  if (sharedBias)
+    laneResult = spatial::SpatVAddOp::create(rewriter, loc, outRowType, vmmResult, sharedBias).getResult();
+  spatial::SpatYieldOp::create(rewriter, loc, laneResult);
+
+  rewriter.setInsertionPointAfter(batchOp);
+  SmallVector<Value> laneResults(batchOp->result_begin(), batchOp->result_end());
+  auto concatComputeOp = createSpatCompute(rewriter, loc, gemmOp.getType(), {}, laneResults, [&](ValueRange args) {
+    spatial::SpatYieldOp::create(rewriter, loc, createSpatConcat(rewriter, loc, /*axis=*/0, args));
+  });
+
+  rewriter.replaceOp(gemmOp, concatComputeOp);
+  return success();
+}
+
 void populateGemmPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
+  patterns.insert<GemmToSpatialComputeBatch>(ctx, PatternBenefit(2));
   patterns.insert<GemmToManyGemv>(ctx);
   patterns.insert<GemvToSpatialCompute>(ctx);
 }

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

@@ -2,6 +2,7 @@
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/PatternMatch.h"
 
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common.hpp"
@@ -14,7 +15,108 @@ using namespace mlir;
 namespace onnx_mlir {
 namespace {
 
-struct MatMulRank3ToGemm : OpRewritePattern<ONNXMatMulOp> {
+static bool haveStaticPositiveShape(ArrayRef<int64_t> shape) {
+  return llvm::all_of(shape, [](int64_t dim) { return dim > 0; });
+}
+
+static Value extractBatchMatrix(Value value,
+                                int64_t batchIndex,
+                                int64_t batchSize,
+                                int64_t rows,
+                                int64_t cols,
+                                PatternRewriter& rewriter,
+                                Location loc) {
+  auto type = cast<RankedTensorType>(value.getType());
+  if (type.getRank() == 2)
+    return value;
+
+  auto sliceType = RankedTensorType::get({1, rows, cols}, type.getElementType());
+  SmallVector<OpFoldResult> offsets = {
+    rewriter.getIndexAttr(batchSize == 1 ? 0 : batchIndex), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
+  SmallVector<OpFoldResult> sizes = {
+    rewriter.getIndexAttr(1), rewriter.getIndexAttr(rows), rewriter.getIndexAttr(cols)};
+  SmallVector<OpFoldResult> strides = {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
+  Value slice = tensor::ExtractSliceOp::create(rewriter, loc, sliceType, value, offsets, sizes, strides);
+
+  auto matrixType = RankedTensorType::get({rows, cols}, type.getElementType());
+  return tensor::CollapseShapeOp::create(rewriter,
+                                         loc,
+                                         matrixType,
+                                         slice,
+                                         SmallVector<ReassociationIndices> {
+                                           {0, 1},
+                                           {2}
+  });
+}
+
+static bool isConstantLikeOperand(Value value) {
+  llvm::SmallPtrSet<Operation*, 8> visited;
+
+  while (auto* definingOp = value.getDefiningOp()) {
+    if (!visited.insert(definingOp).second)
+      return false;
+    if (definingOp->hasTrait<OpTrait::ConstantLike>())
+      return true;
+
+    if (auto extractSliceOp = dyn_cast<tensor::ExtractSliceOp>(definingOp)) {
+      value = extractSliceOp.getSource();
+      continue;
+    }
+    if (auto expandShapeOp = dyn_cast<tensor::ExpandShapeOp>(definingOp)) {
+      value = expandShapeOp.getSrc();
+      continue;
+    }
+    if (auto collapseShapeOp = dyn_cast<tensor::CollapseShapeOp>(definingOp)) {
+      value = collapseShapeOp.getSrc();
+      continue;
+    }
+    if (auto transposeOp = dyn_cast<ONNXTransposeOp>(definingOp)) {
+      value = transposeOp.getData();
+      continue;
+    }
+
+    return false;
+  }
+
+  return false;
+}
+
+static Value transposeLastTwoDims(Value value, PatternRewriter& rewriter, Location loc) {
+  auto type = cast<RankedTensorType>(value.getType());
+  auto shape = type.getShape();
+  if (type.getRank() == 2) {
+    auto transposedType = RankedTensorType::get({shape[1], shape[0]}, type.getElementType());
+    return ONNXTransposeOp::create(rewriter, loc, transposedType, value, rewriter.getI64ArrayAttr({1, 0}));
+  }
+
+  auto transposedType = RankedTensorType::get({shape[0], shape[2], shape[1]}, type.getElementType());
+  return ONNXTransposeOp::create(rewriter, loc, transposedType, value, rewriter.getI64ArrayAttr({0, 2, 1}));
+}
+
+static Value transposeLastTwoDimsInCompute(Value value, PatternRewriter& rewriter, Location loc) {
+  auto type = cast<RankedTensorType>(value.getType());
+  auto shape = type.getShape();
+  RankedTensorType transposedType;
+  SmallVector<int64_t> perm;
+  if (type.getRank() == 2) {
+    transposedType = RankedTensorType::get({shape[1], shape[0]}, type.getElementType());
+    perm = {1, 0};
+  }
+  else {
+    transposedType = RankedTensorType::get({shape[0], shape[2], shape[1]}, type.getElementType());
+    perm = {0, 2, 1};
+  }
+
+  auto transposeCompute =
+    createSpatCompute<1>(rewriter, loc, transposedType, {}, ValueRange {value}, [&](Value input) {
+      Value transposed =
+        ONNXTransposeOp::create(rewriter, loc, transposedType, input, rewriter.getI64ArrayAttr(perm));
+      spatial::SpatYieldOp::create(rewriter, loc, transposed);
+    });
+  return transposeCompute.getResult(0);
+}
+
+struct MatMulToGemm : OpRewritePattern<ONNXMatMulOp> {
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(ONNXMatMulOp matmulOp, PatternRewriter& rewriter) const override {
@@ -24,80 +126,113 @@ struct MatMulRank3ToGemm : OpRewritePattern<ONNXMatMulOp> {
     if (!lhsType || !rhsType || !outType || !lhsType.hasStaticShape() || !rhsType.hasStaticShape()
         || !outType.hasStaticShape())
       return failure();
-    if (lhsType.getRank() != 2 || rhsType.getRank() != 3 || outType.getRank() != 3)
+    if ((lhsType.getRank() != 2 && lhsType.getRank() != 3) || (rhsType.getRank() != 2 && rhsType.getRank() != 3)
+        || (outType.getRank() != 2 && outType.getRank() != 3))
+      return failure();
+    if (!haveStaticPositiveShape(lhsType.getShape()) || !haveStaticPositiveShape(rhsType.getShape())
+        || !haveStaticPositiveShape(outType.getShape()))
       return failure();
 
-    const int64_t batch = rhsType.getDimSize(0);
-    const int64_t k = rhsType.getDimSize(1);
-    const int64_t n = rhsType.getDimSize(2);
-    const int64_t m = lhsType.getDimSize(0);
-    if (lhsType.getDimSize(1) != k || outType.getDimSize(0) != batch || outType.getDimSize(1) != m
-        || outType.getDimSize(2) != n)
+    const int64_t lhsBatch = lhsType.getRank() == 3 ? lhsType.getDimSize(0) : 1;
+    const int64_t rhsBatch = rhsType.getRank() == 3 ? rhsType.getDimSize(0) : 1;
+    const int64_t batch = std::max(lhsBatch, rhsBatch);
+
+    if ((lhsBatch != 1 && lhsBatch != batch) || (rhsBatch != 1 && rhsBatch != batch))
       return failure();
 
-    Location loc = matmulOp.getLoc();
-    auto lhsTransposedType = RankedTensorType::get({k, m}, lhsType.getElementType());
-    auto rhsSliceType = RankedTensorType::get({1, k, 1}, rhsType.getElementType());
-    auto rhsRowType = RankedTensorType::get({1, k}, rhsType.getElementType());
-    auto gemmRowType = RankedTensorType::get({1, m}, outType.getElementType());
-    auto gemmOutType = RankedTensorType::get({batch * n, m}, outType.getElementType());
-    auto gemmExpandedType = RankedTensorType::get({batch, n, m}, outType.getElementType());
+    const int64_t m = lhsType.getRank() == 3 ? lhsType.getDimSize(1) : lhsType.getDimSize(0);
+    const int64_t k = lhsType.getRank() == 3 ? lhsType.getDimSize(2) : lhsType.getDimSize(1);
+    const int64_t rhsK = rhsType.getRank() == 3 ? rhsType.getDimSize(1) : rhsType.getDimSize(0);
+    const int64_t n = rhsType.getRank() == 3 ? rhsType.getDimSize(2) : rhsType.getDimSize(1);
+    if (k != rhsK)
+      return failure();
 
-    Value lhsTransposed =
-      ONNXTransposeOp::create(rewriter, loc, lhsTransposedType, matmulOp.getA(), rewriter.getI64ArrayAttr({1, 0}));
-    Value none = ONNXNoneOp::create(rewriter, loc, rewriter.getNoneType());
-
-    SmallVector<Value> gemmRows;
-    gemmRows.reserve(batch * n);
-    for (int64_t batchIdx = 0; batchIdx < batch; batchIdx++) {
-      for (int64_t colIdx = 0; colIdx < n; colIdx++) {
-        SmallVector<OpFoldResult> offsets = {
-          rewriter.getIndexAttr(batchIdx), rewriter.getIndexAttr(0), rewriter.getIndexAttr(colIdx)};
-        SmallVector<OpFoldResult> sizes = {
-          rewriter.getIndexAttr(1), rewriter.getIndexAttr(k), rewriter.getIndexAttr(1)};
-        SmallVector<OpFoldResult> strides = {
-          rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-        Value rhsSlice =
-          tensor::ExtractSliceOp::create(rewriter, loc, rhsSliceType, matmulOp.getB(), offsets, sizes, strides);
-        Value rhsRow = tensor::CollapseShapeOp::create(rewriter,
-                                                       loc,
-                                                       rhsRowType,
-                                                       rhsSlice,
-                                                       SmallVector<ReassociationIndices> {
-                                                         {0},
-                                                         {1, 2}
-        });
-
-        auto gemmOp = ONNXGemmOp::create(rewriter,
-                                         loc,
-                                         gemmRowType,
-                                         rhsRow,
-                                         lhsTransposed,
-                                         none,
-                                         rewriter.getF32FloatAttr(1.0f),
-                                         rewriter.getF32FloatAttr(1.0f),
-                                         rewriter.getBoolAttr(false),
-                                         rewriter.getBoolAttr(false));
-        gemmRows.push_back(gemmOp.getY());
-      }
+    if (outType.getRank() == 2) {
+      if (batch != 1 || outType.getDimSize(0) != m || outType.getDimSize(1) != n)
+        return failure();
+    }
+    else {
+      if (outType.getDimSize(0) != batch || outType.getDimSize(1) != m || outType.getDimSize(2) != n)
+        return failure();
     }
 
-    auto concatComputeOp = createSpatCompute(rewriter, loc, gemmOutType, {}, gemmRows, [&](ValueRange gemmRowsArgs) {
-      auto concatOp = tensor::ConcatOp::create(rewriter, loc, /*axis=*/0, gemmRowsArgs);
-      spatial::SpatYieldOp::create(rewriter, loc, concatOp.getResult());
-    });
+    Location loc = matmulOp.getLoc();
+    bool useTransposedForm = isConstantLikeOperand(matmulOp.getA()) && !isConstantLikeOperand(matmulOp.getB());
 
-    Value gemmOut = concatComputeOp.getResult(0);
-    Value gemmExpanded = tensor::ExpandShapeOp::create(rewriter,
-                                                       loc,
-                                                       gemmExpandedType,
-                                                       gemmOut,
-                                                       SmallVector<ReassociationIndices> {
-                                                         {0, 1},
-                                                         {2}
-    });
-    Value result = ONNXTransposeOp::create(rewriter, loc, outType, gemmExpanded, rewriter.getI64ArrayAttr({0, 2, 1}));
+    Value lhs = matmulOp.getA();
+    Value rhs = matmulOp.getB();
+    int64_t lhsBatchForGemm = lhsBatch;
+    int64_t rhsBatchForGemm = rhsBatch;
+    int64_t gemmM = m;
+    int64_t gemmK = k;
+    int64_t gemmN = n;
+    if (useTransposedForm) {
+      lhs = transposeLastTwoDimsInCompute(matmulOp.getB(), rewriter, loc);
+      lhsBatchForGemm = rhsBatch;
+      rhs = transposeLastTwoDims(matmulOp.getA(), rewriter, loc);
+      rhsBatchForGemm = lhsBatch;
+      gemmM = n;
+      gemmN = m;
+    }
 
+    auto gemmType = RankedTensorType::get({gemmM, gemmN}, outType.getElementType());
+    auto batchedOutType = RankedTensorType::get({1, m, n}, outType.getElementType());
+    Value none = ONNXNoneOp::create(rewriter, loc, rewriter.getNoneType());
+
+    if (outType.getRank() == 2) {
+      Value lhsMatrix = extractBatchMatrix(lhs, /*batchIndex=*/0, lhsBatchForGemm, gemmM, gemmK, rewriter, loc);
+      Value rhsMatrix = extractBatchMatrix(rhs, /*batchIndex=*/0, rhsBatchForGemm, gemmK, gemmN, rewriter, loc);
+      Value gemmResult = ONNXGemmOp::create(rewriter,
+                                            loc,
+                                            gemmType,
+                                            lhsMatrix,
+                                            rhsMatrix,
+                                            none,
+                                            rewriter.getF32FloatAttr(1.0f),
+                                            rewriter.getF32FloatAttr(1.0f),
+                                            rewriter.getBoolAttr(false),
+                                            rewriter.getBoolAttr(false))
+                           .getY();
+      if (useTransposedForm)
+        gemmResult = ONNXTransposeOp::create(rewriter, loc, outType, gemmResult, rewriter.getI64ArrayAttr({1, 0}));
+      rewriter.replaceOp(matmulOp, gemmResult);
+      return success();
+    }
+
+    SmallVector<Value> batchResults;
+    batchResults.reserve(batch);
+    for (int64_t batchIdx = 0; batchIdx < batch; batchIdx++) {
+      Value lhsMatrix = extractBatchMatrix(lhs, batchIdx, lhsBatchForGemm, gemmM, gemmK, rewriter, loc);
+      Value rhsMatrix = extractBatchMatrix(rhs, batchIdx, rhsBatchForGemm, gemmK, gemmN, rewriter, loc);
+      Value gemmResult = ONNXGemmOp::create(rewriter,
+                                            loc,
+                                            gemmType,
+                                            lhsMatrix,
+                                            rhsMatrix,
+                                            none,
+                                            rewriter.getF32FloatAttr(1.0f),
+                                            rewriter.getF32FloatAttr(1.0f),
+                                            rewriter.getBoolAttr(false),
+                                            rewriter.getBoolAttr(false))
+                           .getY();
+      if (useTransposedForm)
+        gemmResult = ONNXTransposeOp::create(
+          rewriter,
+          loc,
+          RankedTensorType::get({m, n}, outType.getElementType()),
+          gemmResult,
+          rewriter.getI64ArrayAttr({1, 0}));
+      batchResults.push_back(tensor::ExpandShapeOp::create(rewriter,
+                                                           loc,
+                                                           batchedOutType,
+                                                           gemmResult,
+                                                           SmallVector<ReassociationIndices> {
+                                                             {0, 1},
+                                                             {2}
+      }));
+    }
+
+    Value result = createSpatConcat(rewriter, loc, /*axis=*/0, batchResults);
     rewriter.replaceOp(matmulOp, result);
     return success();
   }
@@ -106,7 +241,7 @@ struct MatMulRank3ToGemm : OpRewritePattern<ONNXMatMulOp> {
 } // namespace
 
 void populateMatMulRewritePatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
-  patterns.insert<MatMulRank3ToGemm>(ctx);
+  patterns.insert<MatMulToGemm>(ctx);
 }
 
 } // namespace onnx_mlir

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

@@ -100,8 +100,7 @@ static Value buildReduceMeanKeepdims(Value input,
   for (Value slice : slices)
     reducedSlices.push_back(buildReduceMeanKeepdims(slice, reducedAxes, axis + 1, leafType, rewriter, loc));
 
-  return reducedSlices.size() == 1 ? reducedSlices.front()
-                                   : tensor::ConcatOp::create(rewriter, loc, axis, reducedSlices).getResult();
+  return createSpatConcat(rewriter, loc, axis, reducedSlices);
 }
 
 static Value squeezeReducedAxes(Value keepdimsValue,

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

@@ -33,9 +33,7 @@ static int64_t getOptionalI64(std::optional<ArrayAttrT> arrayAttr, size_t index,
 
 static Value concatAlongAxis(ConversionPatternRewriter& rewriter, Location loc, int64_t axis, ArrayRef<Value> values) {
   assert(!values.empty() && "Expected at least one value to concatenate.");
-  if (values.size() == 1)
-    return values.front();
-  return tensor::ConcatOp::create(rewriter, loc, axis, values);
+  return createSpatConcat(rewriter, loc, axis, values);
 }
 
 static Value materializeContiguousTile(ConversionPatternRewriter& rewriter, Location loc, Value tile) {

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

@@ -47,8 +47,7 @@ buildSoftmax(Value input, int64_t softmaxAxis, int64_t axis, ConversionPatternRe
   for (Value slice : slices)
     rebuiltSlices.push_back(buildSoftmax(slice, softmaxAxis, axis + 1, rewriter, loc));
 
-  return rebuiltSlices.size() == 1 ? rebuiltSlices.front()
-                                   : tensor::ConcatOp::create(rewriter, loc, axis, rebuiltSlices).getResult();
+  return createSpatConcat(rewriter, loc, axis, rebuiltSlices);
 }
 
 struct SoftmaxToSpatialCompute : OpConversionPattern<ONNXSoftmaxOp> {

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

@@ -2,6 +2,7 @@
 #include "mlir/IR/PatternMatch.h"
 
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
 using namespace mlir;
@@ -17,7 +18,7 @@ struct Concat : public OpConversionPattern<ONNXConcatOp> {
     auto inputs = adaptor.getInputs();
     int64_t axis = adaptor.getAxis();
 
-    rewriter.replaceOpWithNewOp<tensor::ConcatOp>(maxpoolOp, axis, inputs);
+    rewriter.replaceOp(maxpoolOp, createSpatConcat(rewriter, maxpoolOp.getLoc(), axis, inputs));
 
     return success();
   }

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

@@ -49,7 +49,7 @@ static Value concatGatherSlices(Value data,
   }
   if (slices.empty())
     return {};
-  return slices.size() == 1 ? slices.front() : tensor::ConcatOp::create(rewriter, loc, axis, slices).getResult();
+  return createSpatConcat(rewriter, loc, axis, slices);
 }
 
 static Value addLeadingGatherDim(Value value, int64_t axis, ConversionPatternRewriter& rewriter, Location loc) {
@@ -130,9 +130,7 @@ struct Gather : OpConversionPattern<ONNXGatherOp> {
                                    return failure();
                                  rows.push_back(addLeadingGatherDim(gatheredRow, axis, rewriter, loc));
                                }
-                               result = rows.size() == 1
-                                        ? rows.front()
-                                        : tensor::ConcatOp::create(rewriter, loc, /*axis=*/axis, rows).getResult();
+                               result = createSpatConcat(rewriter, loc, /*axis=*/axis, rows);
                              }
                              else {
                                return failure();

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

@@ -50,7 +50,7 @@ static Value buildNearestResize(Value input,
     slices.push_back(buildNearestResize(slice, inputShape, outputShape, axis + 1, rewriter, loc));
   }
 
-  return slices.size() == 1 ? slices.front() : tensor::ConcatOp::create(rewriter, loc, axis, slices).getResult();
+  return createSpatConcat(rewriter, loc, axis, slices);
 }
 
 struct Resize : OpConversionPattern<ONNXResizeOp> {

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

@@ -23,7 +23,10 @@ static Value extractSliceAt(
     sizes.push_back(rewriter.getIndexAttr(dim));
   offsets[axis] = rewriter.getIndexAttr(offset);
   sizes[axis] = rewriter.getIndexAttr(size);
-  return tensor::ExtractSliceOp::create(rewriter, loc, input, offsets, sizes, strides);
+  SmallVector<int64_t> resultShape(inputType.getShape());
+  resultShape[axis] = size;
+  auto resultType = RankedTensorType::get(resultShape, inputType.getElementType());
+  return tensor::ExtractSliceOp::create(rewriter, loc, resultType, input, offsets, sizes, strides);
 }
 
 struct Split : OpConversionPattern<ONNXSplitOp> {
@@ -49,12 +52,7 @@ struct Split : OpConversionPattern<ONNXSplitOp> {
       if (!resultType || !resultType.hasStaticShape())
         return failure();
       int64_t sliceSize = resultType.getShape()[axis];
-      auto computeOp =
-        createSpatCompute<1>(rewriter, splitOp.getLoc(), TypeRange {resultType}, {}, adaptor.getInput(), [&](Value x) {
-          Value output = extractSliceAt(x, axis, offset, sliceSize, rewriter, splitOp.getLoc());
-          spatial::SpatYieldOp::create(rewriter, splitOp.getLoc(), output);
-        });
-      outputs.push_back(computeOp.getResult(0));
+      outputs.push_back(extractSliceAt(adaptor.getInput(), axis, offset, sliceSize, rewriter, splitOp.getLoc()));
       offset += sliceSize;
     }
 

src/PIM/Conversion/SpatialToGraphviz/SpatialToGraphviz.cpp

@@ -42,15 +42,15 @@ private:
   raw_ostream& os;
 
   /**
-   * Draws the subgraph for a given spatial::SpatWeightedCompute, including:
+   * Draws the subgraph for a given spatial::SpatCompute, including:
    * 1. Input nodes (block arguments)
    * 2. Operations
    * 3. Edges between yield (output) and its users
    *
-   * @param op The spatial::SpatWeightedCompute to draw the subgraph for.
+   * @param op The spatial::SpatCompute to draw the subgraph for.
    * @param computeNum The number of the compute operation.
    */
-  void drawComputeOpSubgraph(spatial::SpatWeightedCompute op, size_t computeNum) {
+  void drawComputeOpSubgraph(spatial::SpatCompute op, size_t computeNum) {
     os << "\tsubgraph cluster" << computeNum << " {\n\t\tlabel=\"Compute" << computeNum << "\";\n"
        << "\t\tstyle=filled;\n"
        << "\t\tcolor=lightblue;\n";
@@ -217,7 +217,7 @@ void SpatialToGraphvizPass::runOnOperation() {
   // 1. Print their subgraph
   // 2. Print the edges from its inputs to its outputs
   for (Operation& op : func.getOps()) {
-    if (auto computeOp = dyn_cast<spatial::SpatWeightedCompute>(op)) {
+    if (auto computeOp = dyn_cast<spatial::SpatCompute>(op)) {
       drawComputeOpSubgraph(computeOp, computeNum++);
     }
     else if (auto concatOp = dyn_cast<tensor::ConcatOp>(op)) {

src/PIM/Conversion/SpatialToPim/CMakeLists.txt

@@ -5,6 +5,7 @@ add_public_tablegen_target(SpatialToPimIncGen)
 add_pim_library(OMSpatialToPim
   SpatialToPimPass.cpp
   Common.cpp
+  Patterns.cpp
 
   EXCLUDE_FROM_OM_LIBS
 

src/PIM/Conversion/SpatialToPim/Common.cpp

@@ -7,23 +7,12 @@
 #include <cstddef>
 
 #include "Common.hpp"
-#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 
 using namespace llvm;
 using namespace mlir;
 
 namespace onnx_mlir {
 
-namespace {
-
-IntegerAttr getRequiredI32Attr(Builder& builder, Operation* op, llvm::StringRef attrName) {
-  auto attr = op->getAttrOfType<IntegerAttr>(attrName);
-  assert(attr && "required precomputed channel attr is missing");
-  return IntegerAttr::get(builder.getI32Type(), attr.getInt());
-}
-
-} // namespace
-
 size_t getSliceActualOffset(tensor::ExtractSliceOp& sliceOp, ShapedType& inputShape) {
   /*
     EXAMPLE RUN:
@@ -74,37 +63,6 @@ IntegerAttr getTensorSizeInBytesAttr(Builder& builder, mlir::Value value) {
   return builder.getI32IntegerAttr(static_cast<int32_t>(getShapedTypeSizeInBytes(cast<ShapedType>(value.getType()))));
 }
 
-IntegerAttr getSpatialChannelSourceCoreIdAttr(Builder& builder, mlir::Value channel) {
-  auto channelNewOp = channel.getDefiningOp<spatial::SpatChannelNewOp>();
-  assert(channelNewOp && "spatial channel value must come from spat.channel_new");
-  return getRequiredI32Attr(builder, channelNewOp, kChannelSourceCoreIdAttrName);
-}
-
-IntegerAttr getSpatialChannelTargetCoreIdAttr(Builder& builder, mlir::Value channel) {
-  auto channelNewOp = channel.getDefiningOp<spatial::SpatChannelNewOp>();
-  assert(channelNewOp && "spatial channel value must come from spat.channel_new");
-  return getRequiredI32Attr(builder, channelNewOp, kChannelTargetCoreIdAttrName);
-}
-
-bool hasSpatialChannelSourceCoreIdAttr(mlir::Value channel) {
-  auto channelNewOp = channel.getDefiningOp<spatial::SpatChannelNewOp>();
-  return channelNewOp && channelNewOp->hasAttr(kChannelSourceCoreIdAttrName);
-}
-
-bool hasSpatialChannelTargetCoreIdAttr(mlir::Value channel) {
-  auto channelNewOp = channel.getDefiningOp<spatial::SpatChannelNewOp>();
-  return channelNewOp && channelNewOp->hasAttr(kChannelTargetCoreIdAttrName);
-}
-
-mlir::Value createPimReceiveFromSpatialChannel(
-  PatternRewriter& rewriter, Location loc, mlir::Value output, mlir::Value channel) {
-  mlir::Value outputBuffer = getBestOutputTensorFromOperandsOrAllocate(rewriter, output.getDefiningOp());
-  auto sizeAttr = getTensorSizeInBytesAttr(rewriter, output);
-  auto sourceCoreIdAttr = getSpatialChannelSourceCoreIdAttr(rewriter, channel);
-  return pim::PimReceiveOp::create(rewriter, loc, outputBuffer.getType(), outputBuffer, sizeAttr, sourceCoreIdAttr)
-    .getOutput();
-}
-
 Operation* getEarliestUserWithinBlock(mlir::Value value) {
   auto users = value.getUsers();
 
@@ -127,6 +85,16 @@ SmallVector<mlir::Value> getOpOperandsSortedByUses(Operation* operation) {
   return map_to_vector(operandsAndUses, [](auto operandAndUse) { return operandAndUse.first; });
 }
 
+bool hasLaterUserInBlock(mlir::Value value, Operation* operation) {
+  for (Operation* user : value.getUsers()) {
+    if (user->getBlock() != operation->getBlock())
+      return true;
+    if (operation->isBeforeInBlock(user))
+      return true;
+  }
+  return false;
+}
+
 mlir::Value getBestOutputTensorFromOperandsOrAllocate(PatternRewriter& rewriter, Operation* operation) {
   assert("Only support operations with a single result" && operation->getNumResults() == 1);
   mlir::Value result = operation->getResult(0);
@@ -134,8 +102,9 @@ mlir::Value getBestOutputTensorFromOperandsOrAllocate(PatternRewriter& rewriter,
   assert("Only support result ShapedType as result type" && isa<ShapedType>(resultType));
 
   SmallVector<mlir::Value> operands = getOpOperandsSortedByUses(operation);
-  auto validOperands =
-    make_filter_range(operands, [resultType](mlir::Value operand) { return operand.getType() == resultType; });
+  auto validOperands = make_filter_range(operands, [operation, resultType](mlir::Value operand) {
+    return operand.getType() == resultType && !hasLaterUserInBlock(operand, operation);
+  });
   auto bestOperand = validOperands.begin();
 
   if (bestOperand != validOperands.end())

src/PIM/Conversion/SpatialToPim/Common.hpp

@@ -2,16 +2,10 @@
 
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 
-#include "llvm/ADT/StringRef.h"
-
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 namespace onnx_mlir {
 
-inline constexpr llvm::StringLiteral kChannelSourceCoreIdAttrName = "precomp_source_core_id";
-inline constexpr llvm::StringLiteral kChannelTargetCoreIdAttrName = "precomp_target_core_id";
-
 /**
  * \brief Get the offset of the ExtractSliceOp based on its static offsets and
  * its static tensor input.
@@ -30,17 +24,6 @@ size_t getShapedTypeSizeInBytes(mlir::ShapedType shapedType);
 
 mlir::IntegerAttr getTensorSizeInBytesAttr(mlir::Builder& builder, mlir::Value value);
 
-mlir::IntegerAttr getSpatialChannelSourceCoreIdAttr(mlir::Builder& builder, mlir::Value channel);
-
-mlir::IntegerAttr getSpatialChannelTargetCoreIdAttr(mlir::Builder& builder, mlir::Value channel);
-
-bool hasSpatialChannelSourceCoreIdAttr(mlir::Value channel);
-
-bool hasSpatialChannelTargetCoreIdAttr(mlir::Value channel);
-
-mlir::Value createPimReceiveFromSpatialChannel(
-  mlir::PatternRewriter& rewriter, mlir::Location loc, mlir::Value output, mlir::Value channel);
-
 template <class T>
 size_t rangeLength(const mlir::iterator_range<T> range) {
   return std::distance(range.begin(), range.end());

src/PIM/Conversion/SpatialToPim/Patterns.cpp

@@ -0,0 +1,385 @@
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/Value.h"
+#include "mlir/Support/LLVM.h"
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/LogicalResult.h"
+
+#include "Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace {
+
+struct MoveExtractSliceIntoCompute final : OpRewritePattern<mlir::tensor::ExtractSliceOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(mlir::tensor::ExtractSliceOp extractSliceOp, PatternRewriter& rewriter) const override {
+    Location loc = extractSliceOp.getLoc();
+
+    if (!isa<func::FuncOp>(extractSliceOp->getParentOp()))
+      return failure();
+
+    for (auto& uses : extractSliceOp->getUses()) {
+      if (isa<spatial::SpatCompute>(uses.getOwner())) {
+        auto spatCompute = cast<spatial::SpatCompute>(uses.getOwner());
+        if (spatCompute.getInputs().empty())
+          return failure();
+        if (uses.getOperandNumber() < spatCompute.getInputs().getBeginOperandIndex())
+          return failure();
+      }
+      else if (isa_and_present<func::FuncOp>(uses.getOwner()->getParentOp())) {
+        return failure();
+      }
+    }
+
+    llvm::DenseMap<Operation*, Value> mapSpatToExtract;
+
+    for (auto& uses : llvm::make_early_inc_range(extractSliceOp->getUses())) {
+
+      if (auto spatCompute = dyn_cast<spatial::SpatCompute>(uses.getOwner())) {
+        auto BBArgIndex = uses.getOperandNumber() - spatCompute.getInputs().getBeginOperandIndex();
+        auto BBArgValue = spatCompute.getBody().front().getArgument(BBArgIndex);
+
+        if (BBArgValue.use_empty())
+          continue;
+
+        rewriter.setInsertionPoint(&spatCompute.getBody().front().front());
+        if (!mapSpatToExtract.contains(spatCompute.getOperation())) {
+          auto newExtractSlice = rewriter.clone(*extractSliceOp.getOperation());
+          mapSpatToExtract.insert({spatCompute.getOperation(), newExtractSlice->getResult(0)});
+        }
+
+        rewriter.startOpModification(spatCompute.getOperation());
+        BBArgValue.replaceAllUsesWith(mapSpatToExtract[spatCompute.getOperation()]);
+        spatCompute.getInputsMutable().erase(BBArgIndex);
+        spatCompute.getBody().front().eraseArgument(BBArgIndex);
+        rewriter.finalizeOpModification(spatCompute.getOperation());
+      }
+      else if (auto spatComputeBatch = dyn_cast<spatial::SpatComputeBatch>(uses.getOwner())) {
+        auto BBArgIndex = uses.getOperandNumber() - spatComputeBatch.getInputs().getBeginOperandIndex();
+        auto BBArgValue = spatComputeBatch.getBody().front().getArgument(BBArgIndex);
+
+        if (BBArgValue.use_empty())
+          continue;
+
+        rewriter.setInsertionPoint(&spatComputeBatch.getBody().front().front());
+        if (!mapSpatToExtract.contains(spatComputeBatch.getOperation())) {
+          auto newExtractSlice = rewriter.clone(*extractSliceOp.getOperation());
+          mapSpatToExtract.insert({spatComputeBatch.getOperation(), newExtractSlice->getResult(0)});
+        }
+
+        rewriter.startOpModification(spatComputeBatch.getOperation());
+        BBArgValue.replaceAllUsesWith(mapSpatToExtract[spatComputeBatch.getOperation()]);
+        spatComputeBatch.getInputsMutable().erase(BBArgIndex);
+        spatComputeBatch.getBody().front().eraseArgument(BBArgIndex);
+        rewriter.finalizeOpModification(spatComputeBatch.getOperation());
+      }
+      else {
+        {
+          if (auto spatCompute = uses.getOwner()->getParentOfType<spatial::SpatCompute>()) {
+            rewriter.setInsertionPoint(&spatCompute.getBody().front().front());
+            if (!mapSpatToExtract.contains(spatCompute.getOperation())) {
+              auto newExtractSlice = rewriter.clone(*extractSliceOp.getOperation());
+              mapSpatToExtract.insert({spatCompute.getOperation(), newExtractSlice->getResult(0)});
+            }
+
+            rewriter.startOpModification(spatCompute.getOperation());
+            uses.set(mapSpatToExtract[spatCompute.getOperation()]);
+            rewriter.finalizeOpModification(spatCompute.getOperation());
+          }
+          else if (auto spatComputeBatch = uses.getOwner()->getParentOfType<spatial::SpatComputeBatch>()) {
+            rewriter.setInsertionPoint(&spatComputeBatch.getBody().front().front());
+            if (!mapSpatToExtract.contains(spatComputeBatch.getOperation())) {
+              auto newExtractSlice = rewriter.clone(*extractSliceOp.getOperation());
+              mapSpatToExtract.insert({spatComputeBatch.getOperation(), newExtractSlice->getResult(0)});
+            }
+
+            rewriter.startOpModification(spatComputeBatch.getOperation());
+            uses.set(mapSpatToExtract[spatComputeBatch.getOperation()]);
+            rewriter.finalizeOpModification(spatComputeBatch.getOperation());
+          }
+        }
+      }
+    }
+
+    rewriter.eraseOp(extractSliceOp);
+    return success();
+  }
+};
+
+struct ArithConstToGlobalMemoryPattern final : OpRewritePattern<mlir::arith::ConstantOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(mlir::arith::ConstantOp constantOp, PatternRewriter& rewriter) const override {
+    static int i = 0;
+    Location loc = constantOp.getLoc();
+
+    if (hasWeightAlways(constantOp))
+      return failure();
+
+    if (!isa<func::FuncOp>(constantOp->getParentOp()))
+      return failure();
+
+    if (llvm::all_of(constantOp->getUsers(), [](Operation* op) {
+          if (isa<spatial::SpatCompute>(op))
+            return false;
+          if (isa<func::FuncOp>(op->getParentOp()))
+            return true;
+          return false;
+        }))
+      return failure();
+
+    rewriter.setInsertionPoint(constantOp->getParentOfType<func::FuncOp>());
+
+    auto constRankedTensorType = llvm::dyn_cast<mlir::RankedTensorType>(constantOp.getType());
+
+    if (constRankedTensorType) {
+      mlir::MemRefType memRefType =
+        mlir::MemRefType::get(constRankedTensorType.getShape(), constRankedTensorType.getElementType());
+      std::string argName = "const_" + std::to_string(i++);
+      memref::GlobalOp::create(rewriter,
+                               loc,
+                               rewriter.getStringAttr(argName),
+                               rewriter.getStringAttr("private"),
+                               TypeAttr::get(memRefType),
+                               constantOp.getValueAttr(),
+                               rewriter.getUnitAttr(),
+                               {});
+
+      llvm::DenseMap<Operation*, Value> mapSpatComputeToConst;
+
+      for (auto& constUses : llvm::make_early_inc_range(constantOp->getUses())) {
+        auto constUsers = constUses.getOwner();
+
+        if (auto spatCompute = llvm::dyn_cast<spatial::SpatCompute>(constUsers)) {
+
+          auto BBArgIndex = constUses.getOperandNumber() - spatCompute.getInputs().getBeginOperandIndex();
+          auto BBArgValue = spatCompute.getBody().front().getArgument(BBArgIndex);
+          rewriter.setInsertionPoint(&spatCompute.getBody().front().front());
+          if (!mapSpatComputeToConst.contains(spatCompute.getOperation())) {
+            auto getGlobalOp = memref::GetGlobalOp::create(rewriter, loc, memRefType, argName);
+            auto toTensor = bufferization::ToTensorOp::create(
+              rewriter, loc, constRankedTensorType, getGlobalOp, rewriter.getUnitAttr(), rewriter.getUnitAttr());
+            mapSpatComputeToConst.insert({spatCompute.getOperation(), toTensor.getResult()});
+          }
+
+          rewriter.startOpModification(spatCompute.getOperation());
+          BBArgValue.replaceAllUsesWith(mapSpatComputeToConst[spatCompute.getOperation()]);
+          spatCompute.getInputsMutable().erase(BBArgIndex);
+          spatCompute.getBody().front().eraseArgument(BBArgIndex);
+          rewriter.finalizeOpModification(spatCompute.getOperation());
+        }
+        else if (auto spatComputeBatch = llvm::dyn_cast<spatial::SpatComputeBatch>(constUsers)) {
+
+          auto BBArgIndex = constUses.getOperandNumber() - spatComputeBatch.getInputs().getBeginOperandIndex();
+          auto BBArgValue = spatComputeBatch.getBody().front().getArgument(BBArgIndex);
+          rewriter.setInsertionPoint(&spatComputeBatch.getBody().front().front());
+          if (!mapSpatComputeToConst.contains(spatComputeBatch.getOperation())) {
+            auto getGlobalOp = memref::GetGlobalOp::create(rewriter, loc, memRefType, argName);
+            auto toTensor = bufferization::ToTensorOp::create(
+              rewriter, loc, constRankedTensorType, getGlobalOp, rewriter.getUnitAttr(), rewriter.getUnitAttr());
+            mapSpatComputeToConst.insert({spatComputeBatch.getOperation(), toTensor.getResult()});
+          }
+
+          rewriter.startOpModification(spatComputeBatch.getOperation());
+          BBArgValue.replaceAllUsesWith(mapSpatComputeToConst[spatComputeBatch.getOperation()]);
+          spatComputeBatch.getInputsMutable().erase(BBArgIndex);
+          spatComputeBatch.getBody().front().eraseArgument(BBArgIndex);
+          rewriter.finalizeOpModification(spatComputeBatch.getOperation());
+        }
+        else {
+        {
+
+            if (auto spatCompute = constUses.getOwner()->getParentOfType<spatial::SpatCompute>()) {
+              rewriter.setInsertionPoint(&spatCompute.getBody().front().front());
+              if (!mapSpatComputeToConst.contains(spatCompute.getOperation())) {
+                auto getGlobalOp = memref::GetGlobalOp::create(rewriter, loc, memRefType, argName);
+                auto toTensor = bufferization::ToTensorOp::create(
+                  rewriter, loc, constRankedTensorType, getGlobalOp, rewriter.getUnitAttr(), rewriter.getUnitAttr());
+                mapSpatComputeToConst.insert({spatCompute.getOperation(), toTensor.getResult()});
+              }
+
+              rewriter.startOpModification(spatCompute.getOperation());
+              constUses.set(mapSpatComputeToConst[spatCompute.getOperation()]);
+              rewriter.finalizeOpModification(spatCompute.getOperation());
+            }
+            else if (auto spatComputeBatch = constUses.getOwner()->getParentOfType<spatial::SpatComputeBatch>()) {
+              rewriter.setInsertionPoint(&spatComputeBatch.getBody().front().front());
+              if (!mapSpatComputeToConst.contains(spatComputeBatch.getOperation())) {
+                auto getGlobalOp = memref::GetGlobalOp::create(rewriter, loc, memRefType, argName);
+                auto toTensor = bufferization::ToTensorOp::create(
+                  rewriter, loc, constRankedTensorType, getGlobalOp, rewriter.getUnitAttr(), rewriter.getUnitAttr());
+                mapSpatComputeToConst.insert({spatComputeBatch.getOperation(), toTensor.getResult()});
+              }
+
+              rewriter.startOpModification(spatComputeBatch.getOperation());
+              constUses.set(mapSpatComputeToConst[spatComputeBatch.getOperation()]);
+              rewriter.finalizeOpModification(spatComputeBatch.getOperation());
+            }
+          }
+        }
+      }
+    }
+    else if (constantOp.getType().isIntOrIndexOrFloat()) {
+      llvm::DenseMap<Operation*, Value> mapSpatComputeToConst;
+
+      for (auto& constUses : llvm::make_early_inc_range(constantOp->getUses())) {
+        auto constUsers = constUses.getOwner();
+
+        if (auto spatCompute = llvm::dyn_cast<spatial::SpatCompute>(constUsers)) {
+
+          auto BBArgIndex = constUses.getOperandNumber() - spatCompute.getInputs().getBeginOperandIndex();
+          auto BBArgValue = spatCompute.getBody().front().getArgument(BBArgIndex);
+          rewriter.setInsertionPoint(&spatCompute.getBody().front().front());
+          auto newConst = rewriter.clone(*constantOp);
+
+          rewriter.startOpModification(spatCompute.getOperation());
+          BBArgValue.replaceAllUsesWith(newConst->getResult(0));
+          spatCompute.getInputsMutable().erase(BBArgIndex);
+          spatCompute.getBody().front().eraseArgument(BBArgIndex);
+          rewriter.finalizeOpModification(spatCompute.getOperation());
+        }
+        else if (auto spatComputeBatch = llvm::dyn_cast<spatial::SpatComputeBatch>(constUsers)) {
+
+          auto BBArgIndex = constUses.getOperandNumber() - spatComputeBatch.getInputs().getBeginOperandIndex();
+          auto BBArgValue = spatComputeBatch.getBody().front().getArgument(BBArgIndex);
+          rewriter.setInsertionPoint(&spatComputeBatch.getBody().front().front());
+          auto newConst = rewriter.clone(*constantOp);
+
+          rewriter.startOpModification(spatComputeBatch.getOperation());
+          BBArgValue.replaceAllUsesWith(newConst->getResult(0));
+          spatComputeBatch.getInputsMutable().erase(BBArgIndex);
+          spatComputeBatch.getBody().front().eraseArgument(BBArgIndex);
+          rewriter.finalizeOpModification(spatComputeBatch.getOperation());
+        }
+        else {
+          if (auto parent = constUsers->getParentOfType<spatial::SpatCompute>()) {
+            if (!mapSpatComputeToConst.contains(parent)) {
+              rewriter.setInsertionPoint(&parent.getBody().front().front());
+              auto newConst = rewriter.clone(*constantOp);
+            mapSpatComputeToConst.insert({parent.getOperation(), newConst->getResult(0)});
+          }
+          constUses.set(mapSpatComputeToConst[parent.getOperation()]);
+        }
+        else {
+          auto batchParent = constUsers->getParentOfType<spatial::SpatComputeBatch>();
+          assert(batchParent && "Global Constant used direcly not within a compute");
+          if (!mapSpatComputeToConst.contains(batchParent.getOperation())) {
+            rewriter.setInsertionPoint(&batchParent.getBody().front().front());
+            auto newConst = rewriter.clone(*constantOp);
+            mapSpatComputeToConst.insert({batchParent.getOperation(), newConst->getResult(0)});
+          }
+          constUses.set(mapSpatComputeToConst[batchParent.getOperation()]);
+        }
+      }
+      }
+    }
+    auto parent = constantOp->getParentOp();
+    rewriter.eraseOp(constantOp);
+    return success();
+  }
+};
+
+struct FuncOpArgToGlobalMemoryPattern final : OpRewritePattern<mlir::func::FuncOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(mlir::func::FuncOp funcOp, PatternRewriter& rewriter) const override {
+
+    if (funcOp.getArguments().empty())
+      return failure();
+
+    if (llvm::all_of(funcOp.getArguments(),
+                     [](mlir::BlockArgument blockArgument) { return blockArgument.use_empty(); }))
+      return failure();
+
+    Location loc = funcOp.getLoc();
+
+    for (auto [index, arg] : llvm::enumerate(funcOp.getArguments())) {
+      if (arg.getUses().empty())
+        continue;
+
+      rewriter.setInsertionPoint(funcOp.getOperation());
+
+      assert(isa<mlir::RankedTensorType>(arg.getType()));
+
+      auto argRankedTensorType = llvm::dyn_cast<mlir::RankedTensorType>(arg.getType());
+      mlir::MemRefType memRefType =
+        mlir::MemRefType::get(argRankedTensorType.getShape(), argRankedTensorType.getElementType());
+
+      std::string argName = "arg_" + std::to_string(index);
+
+      memref::GlobalOp::create(rewriter,
+                               loc,
+                               rewriter.getStringAttr(argName),
+                               rewriter.getStringAttr("private"),
+                               TypeAttr::get(memRefType),
+                               {},
+                               {},
+                               {});
+
+      for (auto& argUses : llvm::make_early_inc_range(arg.getUses())) {
+        auto argUser = argUses.getOwner();
+        if (auto spatCompute = dyn_cast<spatial::SpatCompute>(argUser)) {
+          auto BBArgIndex = argUses.getOperandNumber() - spatCompute.getInputs().getBeginOperandIndex();
+          auto BBArgValue = spatCompute.getBody().front().getArgument(BBArgIndex);
+          rewriter.setInsertionPoint(&spatCompute.getBody().front().front());
+          auto getGlobalOp = memref::GetGlobalOp::create(rewriter, loc, memRefType, argName);
+          auto toTensor = bufferization::ToTensorOp::create(
+            rewriter, loc, argRankedTensorType, getGlobalOp, rewriter.getUnitAttr(), rewriter.getUnitAttr());
+
+          rewriter.startOpModification(spatCompute.getOperation());
+          BBArgValue.replaceAllUsesWith(toTensor);
+          spatCompute.getInputsMutable().erase(BBArgIndex);
+          spatCompute.getBody().front().eraseArgument(BBArgIndex);
+          rewriter.finalizeOpModification(spatCompute.getOperation());
+        }
+        else if (auto spatComputeBatch = dyn_cast<spatial::SpatComputeBatch>(argUser)) {
+          auto BBArgIndex = argUses.getOperandNumber() - spatComputeBatch.getInputs().getBeginOperandIndex();
+          auto BBArgValue = spatComputeBatch.getBody().front().getArgument(BBArgIndex);
+          rewriter.setInsertionPoint(&spatComputeBatch.getBody().front().front());
+          auto getGlobalOp = memref::GetGlobalOp::create(rewriter, loc, memRefType, argName);
+          auto toTensor = bufferization::ToTensorOp::create(
+            rewriter, loc, argRankedTensorType, getGlobalOp, rewriter.getUnitAttr(), rewriter.getUnitAttr());
+
+          rewriter.startOpModification(spatComputeBatch.getOperation());
+          BBArgValue.replaceAllUsesWith(toTensor);
+          spatComputeBatch.getInputsMutable().erase(BBArgIndex);
+          spatComputeBatch.getBody().front().eraseArgument(BBArgIndex);
+          rewriter.finalizeOpModification(spatComputeBatch.getOperation());
+        }
+        else {
+          rewriter.setInsertionPoint(argUser);
+          auto getGlobalOp = memref::GetGlobalOp::create(rewriter, loc, memRefType, argName);
+          auto toTensor = bufferization::ToTensorOp::create(
+            rewriter, loc, argRankedTensorType, getGlobalOp, rewriter.getUnitAttr(), rewriter.getUnitAttr());
+          rewriter.startOpModification(argUser);
+          argUses.set(toTensor);
+          rewriter.finalizeOpModification(argUser);
+        }
+      }
+    }
+
+    return success();
+  }
+};
+
+} // namespace
+void populateGlobalTensorToMemrefPatterns(RewritePatternSet& patterns) {
+  patterns.add<MoveExtractSliceIntoCompute, FuncOpArgToGlobalMemoryPattern, ArithConstToGlobalMemoryPattern>(
+    patterns.getContext());
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/Patterns.hpp

@@ -0,0 +1,10 @@
+#pragma once
+
+#include "mlir/IR/PatternMatch.h"
+
+
+namespace onnx_mlir {
+
+void populateGlobalTensorToMemrefPatterns(mlir::RewritePatternSet& patterns); 
+
+}

src/PIM/Conversion/SpatialToPim/SpatialToPim.td

@@ -9,17 +9,6 @@ include "src/Accelerators/PIM/Dialect/Spatial/Spatial.td"
 include "src/Accelerators/PIM/Dialect/Pim/Pim.td"
 #endif // OP_BASE
 
-def HasSpatialChannelSourceCoreIdAttr: Constraint<
-  CPred<"onnx_mlir::hasSpatialChannelSourceCoreIdAttr($0)">,
-  "spatial channel has precomputed source core id">;
-
-def HasSpatialChannelTargetCoreIdAttr: Constraint<
-  CPred<"onnx_mlir::hasSpatialChannelTargetCoreIdAttr($0)">,
-  "spatial channel has precomputed target core id">;
-
-def createPimReceiveFromSpatialChannelValue: NativeCodeCall<
-  "onnx_mlir::createPimReceiveFromSpatialChannel($_builder, $_loc, $0, $1)">;
-
 def onnxToPimTranspose : Pat<
   (ONNXTransposeOp:$srcOpRes $data, $perms),
   (PimTransposeOp $data, $perms,
@@ -80,18 +69,4 @@ def spatToPimVSoftmax : Pat<
     (NativeCodeCall<"onnx_mlir::getBestOutputTensorFromOperandsOrAllocate($_builder, $0.getDefiningOp())"> $srcOpRes))
 >;
 
-def spatChannelSendToPimSend : Pat<
-  (SpatChannelSendOp $channel, $input),
-  (PimSendOp $input,
-    (NativeCodeCall<"onnx_mlir::getTensorSizeInBytesAttr($_builder, $0)"> $input),
-    (NativeCodeCall<"onnx_mlir::getSpatialChannelTargetCoreIdAttr($_builder, $0)"> $channel)),
-  [(HasSpatialChannelTargetCoreIdAttr $channel)]
->;
-
-def spatChannelReceiveToPimReceive : Pat<
-  (SpatChannelReceiveOp:$srcOpRes $channel),
-  (createPimReceiveFromSpatialChannelValue $srcOpRes, $channel),
-  [(HasSpatialChannelSourceCoreIdAttr $channel)]
->;
-
 #endif // SPATIAL_TO_PIM

src/PIM/Dialect/Pim/Pim.td

@@ -24,7 +24,7 @@ def PimTensor :
 // Execution
 //===----------------------------------------------------------------------===//
 
-def PimCoreOp : PimOp<"core", [SingleBlock]> {
+def PimCoreOp : PimOp<"core", [SingleBlock, IsolatedFromAbove]> {
   let summary = "Execute a block on a PIM core";
 
   let regions = (region SizedRegion<1>:$body);
@@ -39,6 +39,22 @@ def PimCoreOp : PimOp<"core", [SingleBlock]> {
   }];
 }
 
+def PimCoreBatchOp : PimOp<"core_batch", [SingleBlock, AttrSizedOperandSegments]> {
+  let summary = "Execute equivalent batched core bodies";
+
+  let regions = (region SizedRegion<1>:$body);
+
+  let arguments = (ins
+    I32Attr:$laneCount,
+    Variadic<PimTensor>:$weights,
+    Variadic<PimTensor>:$inputs
+  );
+
+  let assemblyFormat = [{
+    `lanes` $laneCount `(` $weights `)` `[` $inputs `]` attr-dict regions `:` type($weights) `[` type($inputs) `]` `->` `(` `)`
+  }];
+}
+
 def PimHaltOp : PimOp<"halt", [Terminator]> {
   let summary = "Halt execution of the core";
 
@@ -65,6 +81,20 @@ def PimSendOp : PimOp<"send", []> {
   }];
 }
 
+def PimSendBatchOp : PimOp<"send_batch", []> {
+  let summary = "Send a per-lane tensor to target cores from a batched core";
+
+  let arguments = (ins
+    PimTensor:$input,
+    I32Attr:$size,
+    DenseI32ArrayAttr:$targetCoreIds
+  );
+
+  let assemblyFormat = [{
+    `(` $input `)` attr-dict `:` type($input) `->` `(` `)`
+  }];
+}
+
 def PimReceiveOp : PimOp<"receive", [DestinationStyleOpInterface]> {
   let summary = "Receive a tensor from another core";
 
@@ -89,6 +119,30 @@ def PimReceiveOp : PimOp<"receive", [DestinationStyleOpInterface]> {
   }];
 }
 
+def PimReceiveBatchOp : PimOp<"receive_batch", [DestinationStyleOpInterface]> {
+  let summary = "Receive per-lane tensors from source cores into a batched core";
+
+  let arguments = (ins
+    PimTensor:$outputBuffer,
+    I32Attr:$size,
+    DenseI32ArrayAttr:$sourceCoreIds
+  );
+
+  let results = (outs
+    PimTensor:$output
+  );
+
+  let extraClassDeclaration = [{
+    mlir::MutableOperandRange getDpsInitsMutable() {
+      return getOutputBufferMutable();
+    }
+  }];
+
+  let assemblyFormat = [{
+    `(` $outputBuffer `)` attr-dict `:` type($outputBuffer) `->` type($output)
+  }];
+}
+
 def PimMemCopyHostToDevOp : PimOp<"memcp_hd", [DestinationStyleOpInterface]> {
   let summary = "Copy a memory region from host memory into device memory";
 
@@ -115,6 +169,32 @@ def PimMemCopyHostToDevOp : PimOp<"memcp_hd", [DestinationStyleOpInterface]> {
   }];
 }
 
+def PimMemCopyHostToDevBatchOp : PimOp<"memcp_hd_batch", [DestinationStyleOpInterface]> {
+  let summary = "Copy a per-lane tensor from host memory into device memory inside a batched core";
+
+  let arguments = (ins
+    PimTensor:$deviceTarget,
+    PimTensor:$hostSource,
+    I32Attr:$deviceTargetOffset,
+    I32Attr:$hostSourceOffset,
+    I32Attr:$size
+  );
+
+  let results = (outs
+    PimTensor:$output
+  );
+
+  let extraClassDeclaration = [{
+    mlir::MutableOperandRange getDpsInitsMutable() {
+      return getDeviceTargetMutable();
+    }
+  }];
+
+  let assemblyFormat = [{
+    `(` $deviceTarget `,` $hostSource `)` attr-dict `:` `(` type($deviceTarget) `,` type($hostSource) `)` `->` type($output)
+  }];
+}
+
 def PimMemCopyDevToHostOp : PimOp<"memcp_dh", [DestinationStyleOpInterface]> {
   let summary = "Copy a memory region from device memory into host memory";
 

src/PIM/Dialect/Pim/Transforms/Bufferization/OpBufferizationInterfaces.cpp

@@ -1,6 +1,7 @@
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Bufferization/IR/DstBufferizableOpInterfaceImpl.h"
+#include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 
 #include "OpBufferizationInterfaces.hpp"
@@ -65,6 +66,32 @@ struct MemCopyHostToDevOpInterface
   }
 };
 
+struct MemCopyHostToDevBatchOpInterface
+: DstBufferizableOpInterfaceExternalModel<MemCopyHostToDevBatchOpInterface, PimMemCopyHostToDevBatchOp> {
+  LogicalResult bufferize(Operation* op,
+                          RewriterBase& rewriter,
+                          const BufferizationOptions& options,
+                          BufferizationState& state) const {
+    auto memCopyHostToDevOp = cast<PimMemCopyHostToDevBatchOp>(op);
+    auto deviceTargetOpt = getBuffer(rewriter, memCopyHostToDevOp.getDeviceTarget(), options, state);
+    if (failed(deviceTargetOpt))
+      return failure();
+    auto hostSourceOpt = getBuffer(rewriter, memCopyHostToDevOp.getHostSource(), options, state);
+    if (failed(hostSourceOpt))
+      return failure();
+
+    replaceOpWithNewBufferizedOp<PimMemCopyHostToDevBatchOp>(rewriter,
+                                                             memCopyHostToDevOp,
+                                                             deviceTargetOpt->getType(),
+                                                             *deviceTargetOpt,
+                                                             *hostSourceOpt,
+                                                             memCopyHostToDevOp.getDeviceTargetOffsetAttr(),
+                                                             memCopyHostToDevOp.getHostSourceOffsetAttr(),
+                                                             memCopyHostToDevOp.getSizeAttr());
+    return success();
+  }
+};
+
 struct MemCopyDevToHostOpInterface
 : DstBufferizableOpInterfaceExternalModel<MemCopyDevToHostOpInterface, PimMemCopyDevToHostOp> {
   LogicalResult bufferize(Operation* op,
@@ -122,6 +149,127 @@ struct ReceiveOpInterface : DstBufferizableOpInterfaceExternalModel<ReceiveOpInt
   }
 };
 
+struct ReceiveBatchOpInterface : DstBufferizableOpInterfaceExternalModel<ReceiveBatchOpInterface, PimReceiveBatchOp> {
+  bool bufferizesToMemoryRead(Operation* op, OpOperand& opOperand, const AnalysisState& state) const {
+    return !cast<DestinationStyleOpInterface>(op).isDpsInit(&opOperand);
+  }
+
+  LogicalResult bufferize(Operation* op,
+                          RewriterBase& rewriter,
+                          const BufferizationOptions& options,
+                          BufferizationState& state) const {
+    auto receiveOp = cast<PimReceiveBatchOp>(op);
+    auto outputBufferOpt = getBuffer(rewriter, receiveOp.getOutputBuffer(), options, state);
+    if (failed(outputBufferOpt))
+      return failure();
+
+    replaceOpWithNewBufferizedOp<PimReceiveBatchOp>(rewriter,
+                                                    op,
+                                                    outputBufferOpt->getType(),
+                                                    *outputBufferOpt,
+                                                    receiveOp.getSizeAttr(),
+                                                    receiveOp.getSourceCoreIdsAttr());
+    return success();
+  }
+};
+
+struct CoreBatchOpInterface : BufferizableOpInterface::ExternalModel<CoreBatchOpInterface, PimCoreBatchOp> {
+  bool bufferizesToMemoryRead(Operation* op, OpOperand& opOperand, const AnalysisState& state) const {
+    return true;
+  }
+
+  bool bufferizesToMemoryWrite(Operation* op, OpOperand& opOperand, const AnalysisState& state) const {
+    return false;
+  }
+
+  AliasingValueList getAliasingValues(Operation* op, OpOperand& opOperand, const AnalysisState& state) const {
+    return {};
+  }
+
+  AliasingOpOperandList getAliasingOpOperands(Operation* op, Value value, const AnalysisState& state) const {
+    auto coreBatchOp = cast<PimCoreBatchOp>(op);
+    auto bbArg = dyn_cast<BlockArgument>(value);
+    if (!bbArg || bbArg.getOwner() != &coreBatchOp.getBody().front())
+      return {};
+
+    unsigned inputOperandIndex = coreBatchOp.getWeights().size() + bbArg.getArgNumber();
+    return {{&coreBatchOp->getOpOperand(inputOperandIndex), BufferRelation::Equivalent}};
+  }
+
+  bool isWritable(Operation* op, Value value, const AnalysisState& state) const {
+    return false;
+  }
+
+  FailureOr<BufferLikeType>
+  getBufferType(Operation* op,
+                Value value,
+                const BufferizationOptions& options,
+                const BufferizationState& state,
+                SmallVector<Value>& invocationStack) const {
+    auto coreBatchOp = cast<PimCoreBatchOp>(op);
+    auto bbArg = dyn_cast<BlockArgument>(value);
+    if (!bbArg || bbArg.getOwner() != &coreBatchOp.getBody().front())
+      return failure();
+
+    Value tiedInput = coreBatchOp.getInputs()[bbArg.getArgNumber()];
+    if (auto memRefType = dyn_cast<BufferLikeType>(tiedInput.getType()))
+      return memRefType;
+
+    return bufferization::getBufferType(tiedInput, options, state, invocationStack);
+  }
+
+  LogicalResult bufferize(Operation* op,
+                          RewriterBase& rewriter,
+                          const BufferizationOptions& options,
+                          BufferizationState& state) const {
+    auto coreBatchOp = cast<PimCoreBatchOp>(op);
+
+    SmallVector<Value> weights;
+    SmallVector<Value> inputs;
+    weights.reserve(coreBatchOp.getWeights().size());
+    inputs.reserve(coreBatchOp.getInputs().size());
+
+    for (Value weight : coreBatchOp.getWeights()) {
+      if (isa<TensorType>(weight.getType())) {
+        auto weightOpt = getBuffer(rewriter, weight, options, state);
+        if (failed(weightOpt))
+          return failure();
+        weights.push_back(*weightOpt);
+      }
+      else {
+        weights.push_back(weight);
+      }
+    }
+
+    for (Value input : coreBatchOp.getInputs()) {
+      if (isa<TensorType>(input.getType())) {
+        auto inputOpt = getBuffer(rewriter, input, options, state);
+        if (failed(inputOpt))
+          return failure();
+        inputs.push_back(*inputOpt);
+      }
+      else {
+        inputs.push_back(input);
+      }
+    }
+
+    rewriter.setInsertionPoint(coreBatchOp);
+    auto newOp = PimCoreBatchOp::create(
+      rewriter, coreBatchOp.getLoc(), coreBatchOp.getLaneCountAttr(), ValueRange(weights), ValueRange(inputs));
+    newOp.getProperties().setOperandSegmentSizes({static_cast<int>(weights.size()), static_cast<int>(inputs.size())});
+    if (auto coreIdsAttr = coreBatchOp->getAttr(onnx_mlir::kCoreIdAttrName))
+      newOp->setAttr(onnx_mlir::kCoreIdAttrName, coreIdsAttr);
+
+    rewriter.inlineRegionBefore(coreBatchOp.getBody(), newOp.getBody(), newOp.getBody().begin());
+    for (Block& block : newOp.getBody())
+      if (failed(bufferization::bufferizeBlockSignature(&block, rewriter, options, state)))
+        return failure();
+
+    rewriter.eraseOp(coreBatchOp);
+    return success();
+  }
+};
+
 struct TransposeOpInterface : DstBufferizableOpInterfaceExternalModel<TransposeOpInterface, PimTransposeOp> {
   bool bufferizesToMemoryRead(Operation* op, OpOperand& opOperand, const AnalysisState& state) const {
     return !cast<DestinationStyleOpInterface>(op).isDpsInit(&opOperand);
@@ -178,8 +326,10 @@ struct VMMOpInterface : DstBufferizableOpInterfaceExternalModel<VMMOpInterface,
     if (failed(outputBufferOpt))
       return failure();
 
+    Value contiguousInput = materializeContiguousMemRef(*inputOpt, op->getLoc(), rewriter);
+
     replaceOpWithNewBufferizedOp<PimVMMOp>(
-      rewriter, op, outputBufferOpt->getType(), vmmOp.getWeightIndexAttr(), *inputOpt, *outputBufferOpt);
+      rewriter, op, outputBufferOpt->getType(), vmmOp.getWeightIndexAttr(), contiguousInput, *outputBufferOpt);
     return success();
   }
 };
@@ -203,8 +353,10 @@ struct MVMOpInterface : DstBufferizableOpInterfaceExternalModel<MVMOpInterface,
     if (failed(outputBufferOpt))
       return failure();
 
+    Value contiguousInput = materializeContiguousMemRef(*inputOpt, op->getLoc(), rewriter);
+
     replaceOpWithNewBufferizedOp<PimMVMOp>(
-      rewriter, op, outputBufferOpt->getType(), mvmOp.getWeightIndexAttr(), *inputOpt, *outputBufferOpt);
+      rewriter, op, outputBufferOpt->getType(), mvmOp.getWeightIndexAttr(), contiguousInput, *outputBufferOpt);
     return success();
   }
 };
@@ -283,8 +435,11 @@ struct UnaryDstOpInterface : DstBufferizableOpInterfaceExternalModel<UnaryDstOpI
 
 void registerOpBufferizationInterfaces(DialectRegistry& registry) {
   registry.addExtension(+[](MLIRContext* ctx, PimDialect* dialect) {
+    PimCoreBatchOp::attachInterface<CoreBatchOpInterface>(*ctx);
     PimReceiveOp::attachInterface<ReceiveOpInterface>(*ctx);
+    PimReceiveBatchOp::attachInterface<ReceiveBatchOpInterface>(*ctx);
     PimMemCopyHostToDevOp::attachInterface<MemCopyHostToDevOpInterface>(*ctx);
+    PimMemCopyHostToDevBatchOp::attachInterface<MemCopyHostToDevBatchOpInterface>(*ctx);
     PimMemCopyDevToHostOp::attachInterface<MemCopyDevToHostOpInterface>(*ctx);
     PimTransposeOp::attachInterface<TransposeOpInterface>(*ctx);
     PimVMMOp::attachInterface<VMMOpInterface>(*ctx);

src/PIM/Dialect/Pim/Transforms/Bufferization/PimBufferizationPass.cpp

@@ -3,12 +3,17 @@
 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/Threading.h"
 #include "mlir/Pass/Pass.h"
 
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/Debug.h"
+
 #include "Common/PimCommon.hpp"
 #include "Compiler/PimCodeGen.hpp"
 #include "Dialect/Pim/PimOps.hpp"
 #include "Dialect/Pim/Transforms/Bufferization/Common.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Accelerators/PIM/Pass/PIMPasses.h"
 #include "src/Compiler/CompilerOptions.hpp"
 
@@ -40,14 +45,44 @@ private:
 
 void PimBufferizationPass::runOnOperation() {
   auto moduleOp = getOperation();
+  // Refactor this into a function
+  {
+    auto funcOp = getPimEntryFunc(moduleOp);
 
-  // One-Shot-Bufferization
-  bufferization::OneShotBufferizationOptions options;
-  options.allowUnknownOps = true;
-  bufferization::BufferizationState state;
-  if (failed(bufferization::runOneShotBufferize(moduleOp, options, state))) {
-    moduleOp.emitError("Failed to bufferize PIM and Spatial ops");
-    signalPassFailure();
+    auto coreOps = llvm::to_vector(funcOp->getOps<pim::PimCoreOp>());
+    MLIRContext* ctx = moduleOp.getContext();
+    // failableParallelForEach will run the lambda in parallel and stop if any thread fails
+    LogicalResult result = mlir::failableParallelForEach(ctx, coreOps, [&](pim::PimCoreOp coreOp) {
+      // Again, allocate state LOCALLY per thread/function
+      bufferization::OneShotBufferizationOptions options;
+      options.allowUnknownOps = true;
+      bufferization::BufferizationState state;
+      if (failed(bufferization::runOneShotBufferize(coreOp, options, state))) {
+        coreOp.emitError("Failed to bufferize PIM and Spatial ops");
+        return failure();
+      }
+      return success();
+    });
+
+    if (failed(result)) {
+      moduleOp.emitError("Failed to bufferize-parallel PIM and Spatial ops");
+      signalPassFailure();
+    }
+
+    funcOp->walk([&](bufferization::ToTensorOp toTensorOp) {
+      if (llvm::isa_and_present<pim::PimCoreOp>(toTensorOp->getParentOp()))
+        toTensorOp->setAttr("restrict", UnitAttr::get(ctx));
+    });
+
+    // One-Shot-Bufferization
+    bufferization::OneShotBufferizationOptions options;
+    options.allowUnknownOps = true;
+    bufferization::BufferizationState state;
+
+    if (failed(bufferization::runOneShotBufferize(moduleOp, options, state))) {
+      moduleOp.emitError("Failed to bufferize PIM and Spatial ops");
+      signalPassFailure();
+    }
   }
 
   MLIRContext* ctx = moduleOp.getContext();
@@ -57,7 +92,18 @@ void PimBufferizationPass::runOnOperation() {
   RewritePatternSet patterns(ctx);
   populateWithGenerated(patterns);
 
-  if (failed(applyPartialConversion(moduleOp, target, std::move(patterns)))) {
+  // Only convert memref.copy → pim.memcp inside pim.core / pim.core_batch bodies.
+  // Host-level copies (e.g. from split/slice ops) must remain as memref.copy for CPU lowering.
+  FrozenRewritePatternSet frozenPatterns(std::move(patterns));
+  bool hasFailed = false;
+  moduleOp.walk<WalkOrder::PreOrder>([&](Operation* op) {
+    if (!isa<pim::PimCoreOp, pim::PimCoreBatchOp>(op))
+      return WalkResult::advance();
+    if (failed(applyPartialConversion(op, target, frozenPatterns)))
+      hasFailed = true;
+    return WalkResult::skip();
+  });
+  if (hasFailed) {
     signalPassFailure();
     return;
   }
@@ -93,16 +139,21 @@ void PimBufferizationPass::runOnOperation() {
 }
 
 void PimBufferizationPass::annotateWeightsMemrefs(ModuleOp moduleOp, func::FuncOp funcOp) const {
-  funcOp.walk([&](memref::GetGlobalOp getGlobalOp) {
-    bool isAlwaysWeight = !getGlobalOp->getUsers().empty()
-                       && all_of(getGlobalOp->getUsers(), [](auto user) -> bool { return isa<PimCoreOp>(user); });
-    if (isAlwaysWeight) {
+  auto markWeights = [&](Operation* op) {
+    walkPimMvmVmmWeightUses(op, [&](OpOperand& weightUse) {
+      Value weight = weightUse.get();
+      auto getGlobalOp = weight.getDefiningOp<memref::GetGlobalOp>();
+      if (!getGlobalOp)
+        return;
       auto globalMemrefOp = lookupGlobalForGetGlobal(moduleOp, getGlobalOp);
       assert("Weights must be constants" && globalMemrefOp.getConstant());
       markWeightAlways(getGlobalOp);
       markWeightAlways(globalMemrefOp);
-    }
-  });
+    });
+  };
+
+  funcOp.walk([&](PimCoreOp coreOp) { markWeights(coreOp); });
+  funcOp.walk([&](PimCoreBatchOp coreBatchOp) { markWeights(coreBatchOp); });
 }
 
 std::unique_ptr<Pass> createPimBufferizationPass() { return std::make_unique<PimBufferizationPass>(); }

src/PIM/Dialect/Spatial/CMakeLists.txt

@@ -2,6 +2,7 @@ add_onnx_mlir_dialect(Spatial spat)
 add_onnx_mlir_dialect_doc(spat Spatial.td)
 
 add_pim_library(SpatialOps
+  Channels.cpp
   SpatialOps.cpp
   Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp
   Transforms/MergeComputeNodes/DCPGraph/Graph.cpp

src/PIM/Dialect/Spatial/Channels.cpp

@@ -0,0 +1,120 @@
+#include "src/Accelerators/PIM/Dialect/Spatial/Channels.hpp"
+
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/Diagnostics.h"
+
+using namespace mlir;
+
+namespace onnx_mlir::spatial {
+
+namespace {
+
+static Channels::ChannelId getChannelId(SpatChannelSendOp sendOp) { return sendOp.getChannelId(); }
+
+static Channels::ChannelId getChannelId(SpatChannelReceiveOp receiveOp) { return receiveOp.getChannelId(); }
+
+static LogicalResult verifyEndpointPair(ChannelEndpoints endpoints) {
+  if (!endpoints.send || !endpoints.receive)
+    return failure();
+
+  if (endpoints.send.getSourceCoreId() != endpoints.receive.getSourceCoreId()) {
+    endpoints.send.emitOpError("sourceCoreId does not match paired spat.channel_receive");
+    return failure();
+  }
+  if (endpoints.send.getTargetCoreId() != endpoints.receive.getTargetCoreId()) {
+    endpoints.send.emitOpError("targetCoreId does not match paired spat.channel_receive");
+    return failure();
+  }
+  if (endpoints.send.getInput().getType() != endpoints.receive.getOutput().getType()) {
+    endpoints.send.emitOpError("input type does not match paired spat.channel_receive result type");
+    return failure();
+  }
+
+  return success();
+}
+
+} // namespace
+
+Channels::Channels(func::FuncOp funcOp) {
+  if (!funcOp)
+    return;
+
+  funcOp.walk([&](SpatChannelSendOp sendOp) { insertSend(sendOp); });
+  funcOp.walk([&](SpatChannelReceiveOp receiveOp) { insertReceive(receiveOp); });
+}
+
+Channels::ChannelId Channels::allocate() { return nextChannelId++; }
+
+void Channels::insertSend(SpatChannelSendOp sendOp) {
+  ChannelId channelId = getChannelId(sendOp);
+  nextChannelId = std::max(nextChannelId, channelId + 1);
+  endpoints[channelId].send = sendOp;
+}
+
+void Channels::insertReceive(SpatChannelReceiveOp receiveOp) {
+  ChannelId channelId = getChannelId(receiveOp);
+  nextChannelId = std::max(nextChannelId, channelId + 1);
+  endpoints[channelId].receive = receiveOp;
+}
+
+void Channels::eraseSend(SpatChannelSendOp sendOp) {
+  ChannelId channelId = getChannelId(sendOp);
+  auto it = endpoints.find(channelId);
+  if (it == endpoints.end())
+    return;
+  it->second.send = {};
+  if (!it->second.receive)
+    endpoints.erase(it);
+}
+
+void Channels::eraseReceive(SpatChannelReceiveOp receiveOp) {
+  ChannelId channelId = getChannelId(receiveOp);
+  auto it = endpoints.find(channelId);
+  if (it == endpoints.end())
+    return;
+  it->second.receive = {};
+  if (!it->second.send)
+    endpoints.erase(it);
+}
+
+FailureOr<ChannelEndpoints> Channels::lookup(ChannelId id) const {
+  auto it = endpoints.find(id);
+  if (it == endpoints.end())
+    return failure();
+  return it->second;
+}
+
+FailureOr<SpatChannelReceiveOp> Channels::getReceiveFor(SpatChannelSendOp sendOp) const {
+  auto endpointsOr = lookup(getChannelId(sendOp));
+  if (failed(endpointsOr) || !endpointsOr->receive)
+    return failure();
+  return endpointsOr->receive;
+}
+
+FailureOr<SpatChannelSendOp> Channels::getSendFor(SpatChannelReceiveOp receiveOp) const {
+  auto endpointsOr = lookup(getChannelId(receiveOp));
+  if (failed(endpointsOr) || !endpointsOr->send)
+    return failure();
+  return endpointsOr->send;
+}
+
+LogicalResult Channels::verify() const {
+  for (const auto& [channelId, pair] : endpoints) {
+    if (!pair.send || !pair.receive) {
+      if (pair.send) {
+        auto sendOp = pair.send;
+        sendOp.emitOpError("channel_id ") << channelId << " is missing a paired spat.channel_receive";
+      }
+      else if (pair.receive) {
+        auto receiveOp = pair.receive;
+        receiveOp.emitOpError("channel_id ") << channelId << " is missing a paired spat.channel_send";
+      }
+      return failure();
+    }
+    if (failed(verifyEndpointPair(pair)))
+      return failure();
+  }
+  return success();
+}
+
+} // namespace onnx_mlir::spatial

src/PIM/Dialect/Spatial/Channels.hpp

@@ -0,0 +1,43 @@
+#pragma once
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Support/LogicalResult.h"
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/StringRef.h"
+
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+namespace onnx_mlir::spatial {
+
+struct ChannelEndpoints {
+  SpatChannelSendOp send;
+  SpatChannelReceiveOp receive;
+};
+
+class Channels {
+public:
+  using ChannelId = int64_t;
+
+  explicit Channels(mlir::func::FuncOp funcOp);
+
+  ChannelId allocate();
+
+  void insertSend(SpatChannelSendOp sendOp);
+  void insertReceive(SpatChannelReceiveOp receiveOp);
+  void eraseSend(SpatChannelSendOp sendOp);
+  void eraseReceive(SpatChannelReceiveOp receiveOp);
+
+  llvm::FailureOr<ChannelEndpoints> lookup(ChannelId id) const;
+  llvm::FailureOr<SpatChannelReceiveOp> getReceiveFor(SpatChannelSendOp sendOp) const;
+  llvm::FailureOr<SpatChannelSendOp> getSendFor(SpatChannelReceiveOp receiveOp) const;
+
+  mlir::LogicalResult verify() const;
+
+private:
+  ChannelId nextChannelId = 0;
+  llvm::DenseMap<ChannelId, ChannelEndpoints> endpoints;
+};
+
+} // namespace onnx_mlir::spatial

src/PIM/Dialect/Spatial/Spatial.td

@@ -9,7 +9,6 @@ def SpatialDialect : Dialect {
   let name = "spat";
   let summary = "Dialect designed for deep learning computation in a spatial architecture";
   let cppNamespace = "::onnx_mlir::spatial";
-  let useDefaultTypePrinterParser = 1;
 }
 
 class SpatOp<string mnemonic, list<Trait> traits = []> :
@@ -19,20 +18,11 @@ class SpatOp<string mnemonic, list<Trait> traits = []> :
 def SpatTensor :
   AnyTypeOf<[AnyMemRef, AnyRankedTensor], "", "::mlir::ShapedType">;
 
-class SpatType<string name, string typeMnemonic, list<Trait> traits = []>
-    : TypeDef<SpatialDialect, name, traits> {
-  let mnemonic = typeMnemonic;
-}
-
-def SpatChannelType : SpatType<"SpatChannel", "ch"> {
-  let summary = "Virtual channel type";
-}
-
 //===----------------------------------------------------------------------===//
 // Execution
 //===----------------------------------------------------------------------===//
 
-def SpatWeightedCompute : SpatOp<"compute", [SingleBlock, AttrSizedOperandSegments]> {
+def SpatCompute : SpatOp<"compute", [SingleBlock, AttrSizedOperandSegments]> {
   let summary = "Compute region with attached constant weights";
 
   let arguments = (ins
@@ -48,10 +38,27 @@ def SpatWeightedCompute : SpatOp<"compute", [SingleBlock, AttrSizedOperandSegmen
 
   let hasVerifier = 1;
   let hasFolder = 1;
+  let hasCustomAssemblyFormat = 1;
+}
 
-  let assemblyFormat = [{
-    `[` $weights `]` `(` $inputs `)` attr-dict `:` `[` type($weights) `]` `(` type($inputs) `)` `->` type($outputs) $body
-  }];
+def SpatComputeBatch : SpatOp<"compute_batch",
+    [SingleBlock, AttrSizedOperandSegments]> {
+  let summary = "Compressed batch of independent equivalent compute lanes";
+
+  let arguments = (ins
+    I32Attr:$laneCount,
+    Variadic<SpatTensor>:$weights,
+    Variadic<SpatTensor>:$inputs
+  );
+
+  let results = (outs
+    Variadic<SpatTensor>:$outputs
+  );
+
+  let regions = (region SizedRegion<1>:$body);
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
 }
 
 def SpatYieldOp : SpatOp<"yield", [Terminator]> {
@@ -61,51 +68,66 @@ def SpatYieldOp : SpatOp<"yield", [Terminator]> {
     Variadic<SpatTensor>:$outputs
   );
 
-  let assemblyFormat = [{
-    $outputs attr-dict `:` type($outputs)
-  }];
+  let hasCustomAssemblyFormat = 1;
+}
+
+def SpatExtractRowsOp : SpatOp<"extract_rows", []> {
+  let summary = "Extract every row of a rank-2 tensor as separate rank-2 row tensors";
+
+  let arguments = (ins
+    SpatTensor:$input
+  );
+
+  let results = (outs
+    Variadic<SpatTensor>:$outputs
+  );
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def SpatConcatOp : SpatOp<"concat", []> {
+  let summary = "Concatenate tensors with compact Spatial operand syntax";
+
+  let arguments = (ins
+    I64Attr:$axis,
+    Variadic<SpatTensor>:$inputs
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Communication
 //===----------------------------------------------------------------------===//
 
-def SpatChannelNewOp : SpatOp<"channel_new", []> {
-  let summary = "Create a new virtual channel";
-
-  let results = (outs
-    SpatChannelType:$channel
-  );
-
-  let builders = [
-    OpBuilder<(ins ), [{
-      $_state.addTypes(SpatChannelType());
-    }]>
-  ];
-
-  let assemblyFormat = [{
-    attr-dict
-  }];
-}
-
 def SpatChannelSendOp : SpatOp<"channel_send", []> {
-  let summary = "Send a tensor through a channel";
+  let summary = "Send a tensor through a logical channel";
 
   let arguments = (ins
-    SpatChannelType:$channel,
+    I64Attr:$channelId,
+    I32Attr:$sourceCoreId,
+    I32Attr:$targetCoreId,
     SpatTensor:$input
   );
 
   let assemblyFormat = [{
-    $input `to` $channel attr-dict `:` `(` type($input) `->` type($channel) `)`
+    $input attr-dict `:` type($input)
   }];
 }
 
 def SpatChannelReceiveOp : SpatOp<"channel_receive", []> {
-  let summary = "Receive a tensor from a channel";
+  let summary = "Receive a tensor from a logical channel";
 
   let arguments = (ins
-    SpatChannelType:$channel
+    I64Attr:$channelId,
+    I32Attr:$sourceCoreId,
+    I32Attr:$targetCoreId
   );
 
   let results = (outs
@@ -113,37 +135,70 @@ def SpatChannelReceiveOp : SpatOp<"channel_receive", []> {
   );
 
   let assemblyFormat = [{
-    $channel attr-dict `:` `(` type($channel) `->` type($output) `)`
+    attr-dict `:` type($output)
   }];
 }
 
-def SpatChannelBroadcastSendOp : SpatOp<"channel_broadcast_send", []> {
-  let summary = "Broadcast a tensor through a shared channel buffer";
+def SpatChannelSendManyOp : SpatOp<"channel_send_many", []> {
+  let summary = "Send multiple tensors through logical channels";
 
   let arguments = (ins
-    SpatChannelType:$channel,
+    DenseI64ArrayAttr:$channelIds,
+    DenseI32ArrayAttr:$sourceCoreIds,
+    DenseI32ArrayAttr:$targetCoreIds,
+    Variadic<SpatTensor>:$inputs
+  );
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def SpatChannelReceiveManyOp : SpatOp<"channel_receive_many", []> {
+  let summary = "Receive multiple tensors from logical channels";
+
+  let arguments = (ins
+    DenseI64ArrayAttr:$channelIds,
+    DenseI32ArrayAttr:$sourceCoreIds,
+    DenseI32ArrayAttr:$targetCoreIds
+  );
+
+  let results = (outs
+    Variadic<SpatTensor>:$outputs
+  );
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def SpatChannelSendBatchOp : SpatOp<"channel_send_batch", []> {
+  let summary = "Send per-lane tensors through logical channels in a batch body";
+
+  let arguments = (ins
+    DenseI64ArrayAttr:$channelIds,
+    DenseI32ArrayAttr:$sourceCoreIds,
+    DenseI32ArrayAttr:$targetCoreIds,
     SpatTensor:$input
   );
 
-  let assemblyFormat = [{
-    $input `to` $channel attr-dict `:` `(` type($input) `->` type($channel) `)`
-  }];
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
 }
 
-def SpatChannelBroadcastReceiveOp : SpatOp<"channel_broadcast_receive", []> {
-  let summary = "Receive a tensor from a shared channel buffer";
+def SpatChannelReceiveBatchOp : SpatOp<"channel_receive_batch", []> {
+  let summary = "Receive a per-lane tensor through logical channels in a batch body";
 
   let arguments = (ins
-    SpatChannelType:$channel
+    DenseI64ArrayAttr:$channelIds,
+    DenseI32ArrayAttr:$sourceCoreIds,
+    DenseI32ArrayAttr:$targetCoreIds
   );
 
   let results = (outs
     SpatTensor:$output
   );
 
-  let assemblyFormat = [{
-    $channel attr-dict `:` `(` type($channel) `->` type($output) `)`
-  }];
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
 }
 
 //===----------------------------------------------------------------------===//

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/DCPAnalysis.cpp

@@ -3,13 +3,23 @@
 #include "mlir/IR/Value.h"
 #include "mlir/IR/ValueRange.h"
 
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/Casting.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/raw_ostream.h"
 
+#include <algorithm>
 #include <iterator>
+#include <numeric>
+#include <optional>
+#include <queue>
+#include <utility>
+#include <vector>
 
 #include "DCPAnalysis.hpp"
 #include "Graph.hpp"
+#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Support/TypeUtilities.hpp"
 
 namespace onnx_mlir {
@@ -17,46 +27,731 @@ namespace spatial {
 
 using namespace mlir;
 
-SpatWeightedCompute getOriginalSpatWeightedCompute(Operation* op) {
+namespace {
+using SpatCompute = onnx_mlir::spatial::SpatCompute;
+using SpatComputeBatch = onnx_mlir::spatial::SpatComputeBatch;
+
+struct VirtualNode {
+  SmallVector<size_t, 4> originalComputeIndices;
+  Weight weight = 0;
+  CrossbarUsage crossbarUsage = 0;
+};
+
+struct VirtualGraph {
+  std::vector<VirtualNode> nodes;
+  std::vector<IndexedEdge> edges;
+};
+
+struct TimingInfo {
+  std::vector<Time> aest;
+  std::vector<Time> alst;
+  std::vector<size_t> topologicalOrder;
+  bool valid = false;
+};
+
+struct WindowScheduleResult {
+  std::vector<std::vector<size_t>> mergeGroups;
+  CPU cpuCount = 0;
+  size_t mergedNodeCount = 0;
+  size_t maxMergeGroupSize = 0;
+};
+
+constexpr CPU kDefaultMaxCpuCount = 1000;
+
+size_t getSchedulingCpuBudget() {
+  if (coresCount.getValue() > 0)
+    return static_cast<size_t>(coresCount.getValue());
+  return static_cast<size_t>(kDefaultMaxCpuCount);
+}
+
+size_t getBatchChunkTargetCount(int32_t laneCount) {
+  assert(laneCount > 0 && "laneCount must be positive");
+  return std::min(static_cast<size_t>(laneCount), std::max<size_t>(1, getSchedulingCpuBudget()));
+}
+
+ComputeInstance getBatchChunkForIndex(SpatComputeBatch batch, size_t chunkIndex) {
+  size_t totalLanes = static_cast<size_t>(batch.getLaneCount());
+  size_t chunkCount = getBatchChunkTargetCount(batch.getLaneCount());
+  size_t baseChunkSize = totalLanes / chunkCount;
+  size_t largeChunkCount = totalLanes % chunkCount;
+
+  size_t laneStart = chunkIndex * baseChunkSize + std::min(chunkIndex, largeChunkCount);
+  size_t laneCount = baseChunkSize + (chunkIndex < largeChunkCount ? 1 : 0);
+  return {batch.getOperation(), static_cast<uint32_t>(laneStart), static_cast<uint32_t>(laneCount)};
+}
+
+ComputeInstance getBatchChunkForLane(SpatComputeBatch batch, uint32_t lane) {
+  size_t totalLanes = static_cast<size_t>(batch.getLaneCount());
+  size_t chunkCount = getBatchChunkTargetCount(batch.getLaneCount());
+  size_t baseChunkSize = totalLanes / chunkCount;
+  size_t largeChunkCount = totalLanes % chunkCount;
+  size_t largeChunkSpan = largeChunkCount * (baseChunkSize + 1);
+
+  size_t chunkIndex = 0;
+  if (static_cast<size_t>(lane) < largeChunkSpan)
+    chunkIndex = static_cast<size_t>(lane) / (baseChunkSize + 1);
+  else
+    chunkIndex = largeChunkCount + (static_cast<size_t>(lane) - largeChunkSpan) / baseChunkSize;
+  return getBatchChunkForIndex(batch, chunkIndex);
+}
+
+std::vector<IndexedEdge> aggregateEdges(ArrayRef<IndexedEdge> edges) {
+  llvm::DenseMap<std::pair<size_t, size_t>, Weight> edgeWeights;
+  for (auto [start, end, weight] : edges) {
+    size_t startIndex = static_cast<size_t>(start);
+    size_t endIndex = static_cast<size_t>(end);
+    if (startIndex == endIndex)
+      continue;
+    auto key = std::make_pair(startIndex, endIndex);
+    Weight edgeWeight = static_cast<Weight>(weight);
+    auto inserted = edgeWeights.try_emplace(key, edgeWeight);
+    if (!inserted.second)
+      inserted.first->second = std::max(inserted.first->second, edgeWeight);
+  }
+
+  std::vector<IndexedEdge> aggregatedEdges;
+  aggregatedEdges.reserve(edgeWeights.size());
+  for (auto [key, weight] : edgeWeights)
+    aggregatedEdges.push_back(
+      {static_cast<int64_t>(key.first), static_cast<int64_t>(key.second), static_cast<int64_t>(weight)});
+  llvm::sort(aggregatedEdges, [](const IndexedEdge& lhs, const IndexedEdge& rhs) {
+    if (std::get<0>(lhs) != std::get<0>(rhs))
+      return std::get<0>(lhs) < std::get<0>(rhs);
+    return std::get<1>(lhs) < std::get<1>(rhs);
+  });
+  return aggregatedEdges;
+}
+
+Weight getComputeBodyWeight(Region& body) {
+  constexpr Weight kOperationWeight = 100;
+  Weight numOperations = 0;
+  for (auto& block : body)
+    for ([[maybe_unused]] auto& op : block)
+      numOperations = checkedAdd(numOperations, static_cast<Weight>(1));
+  return checkedMultiply(numOperations, kOperationWeight);
+}
+
+CrossbarUsage getComputeBodyCrossbarUsage(Region& body) {
+  CrossbarUsage crossbarUsage = 0;
+  for (auto& block : body)
+    for (auto& op : block)
+      if (isa<SpatWeightedVMMOp>(op))
+        crossbarUsage = checkedAdd(crossbarUsage, static_cast<CrossbarUsage>(1));
+  return crossbarUsage;
+}
+
+Weight getComputeInstanceWeight(const ComputeInstance& instance) {
+  if (auto spatCompute = dyn_cast<SpatCompute>(instance.op))
+    return getSpatComputeWeight(spatCompute);
+  auto batch = cast<SpatComputeBatch>(instance.op);
+  return checkedMultiply(getComputeBodyWeight(batch.getBody()), static_cast<Weight>(instance.laneCount));
+}
+
+CrossbarUsage getComputeInstanceCrossbarUsage(const ComputeInstance& instance) {
+  if (auto spatCompute = dyn_cast<SpatCompute>(instance.op))
+    return getSpatComputeCrossbarUsage(spatCompute);
+  auto batch = cast<SpatComputeBatch>(instance.op);
+  return checkedMultiply(getComputeBodyCrossbarUsage(batch.getBody()), static_cast<CrossbarUsage>(instance.laneCount));
+}
+
+SmallVector<Value, 4> getComputeInstanceInputs(const ComputeInstance& instance) {
+  if (auto spatCompute = dyn_cast<SpatCompute>(instance.op))
+    return SmallVector<Value, 4>(spatCompute.getInputs().begin(), spatCompute.getInputs().end());
+  auto batch = cast<SpatComputeBatch>(instance.op);
+  SmallVector<Value, 4> inputs;
+  inputs.reserve(instance.laneCount);
+  for (uint32_t lane = instance.laneStart; lane < instance.laneStart + instance.laneCount; ++lane)
+    inputs.push_back(batch.getInputs()[lane]);
+  return inputs;
+}
+
+std::optional<ComputeInstance> getOriginalComputeInstance(Value value) {
+  Operation* op = value.getDefiningOp();
+  if (!op)
+    return std::nullopt;
+
+  while (auto extract = dyn_cast<tensor::ExtractSliceOp>(op)) {
+    value = extract.getSource();
+    op = value.getDefiningOp();
+    if (!op)
+      return std::nullopt;
+  }
+
+  if (auto spatCompute = dyn_cast<SpatCompute>(op))
+    return ComputeInstance {spatCompute.getOperation(), 0, 1};
+  if (auto batch = dyn_cast<SpatComputeBatch>(op))
+    return getBatchChunkForLane(batch, static_cast<uint32_t>(cast<OpResult>(value).getResultNumber()));
+  return std::nullopt;
+}
+
+SmallVector<ComputeInstance> collectComputeInstances(Operation* entryOp) {
+  SmallVector<ComputeInstance> instances;
+  for (Region& region : entryOp->getRegions()) {
+    for (Block& block : region) {
+      for (Operation& op : block) {
+        if (auto spatCompute = dyn_cast<SpatCompute>(&op)) {
+          instances.push_back({spatCompute.getOperation(), 0, 1});
+          continue;
+        }
+        if (auto batch = dyn_cast<SpatComputeBatch>(&op)) {
+          size_t chunkCount = getBatchChunkTargetCount(batch.getLaneCount());
+          for (size_t chunkIndex = 0; chunkIndex < chunkCount; ++chunkIndex)
+            instances.push_back(getBatchChunkForIndex(batch, chunkIndex));
+        }
+      }
+    }
+  }
+  return instances;
+}
+
+VirtualGraph buildInitialVirtualGraph(ArrayRef<ComputeInstance> computeInstances, ArrayRef<IndexedEdge> edges) {
+  VirtualGraph graph;
+  graph.nodes.reserve(computeInstances.size());
+  for (auto [index, computeInstance] : llvm::enumerate(computeInstances)) {
+    VirtualNode node;
+    node.originalComputeIndices.push_back(index);
+    node.weight = getComputeInstanceWeight(computeInstance);
+    node.crossbarUsage = getComputeInstanceCrossbarUsage(computeInstance);
+    graph.nodes.push_back(std::move(node));
+  }
+  graph.edges = aggregateEdges(edges);
+  return graph;
+}
+
+TimingInfo computeTiming(const VirtualGraph& graph) {
+  TimingInfo timing;
+  size_t nodeCount = graph.nodes.size();
+  timing.aest.assign(nodeCount, 0);
+  timing.alst.assign(nodeCount, 0);
+  timing.topologicalOrder.reserve(nodeCount);
+
+  std::vector<std::vector<std::pair<size_t, Weight>>> parents(nodeCount);
+  std::vector<std::vector<std::pair<size_t, Weight>>> children(nodeCount);
+  std::vector<size_t> incomingEdgeCount(nodeCount, 0);
+
+  for (auto [start, end, weight] : graph.edges) {
+    size_t startIndex = static_cast<size_t>(start);
+    size_t endIndex = static_cast<size_t>(end);
+    Weight edgeWeight = static_cast<Weight>(weight);
+    assert(startIndex < nodeCount && endIndex < nodeCount && "virtual edge endpoint out of range");
+    children[startIndex].push_back({endIndex, edgeWeight});
+    parents[endIndex].push_back({startIndex, edgeWeight});
+    incomingEdgeCount[endIndex]++;
+  }
+
+  auto getVirtualNodeOrderKey = [&](size_t nodeIndex) {
+    const VirtualNode& node = graph.nodes[nodeIndex];
+    if (!node.originalComputeIndices.empty())
+      return node.originalComputeIndices.front();
+    return nodeIndex;
+  };
+  auto readyNodeGreater = [&](size_t lhs, size_t rhs) {
+    size_t lhsKey = getVirtualNodeOrderKey(lhs);
+    size_t rhsKey = getVirtualNodeOrderKey(rhs);
+    if (lhsKey != rhsKey)
+      return lhsKey > rhsKey;
+    return lhs > rhs;
+  };
+  std::priority_queue<size_t, std::vector<size_t>, decltype(readyNodeGreater)> readyNodes(readyNodeGreater);
+  for (size_t i = 0; i < nodeCount; ++i)
+    if (incomingEdgeCount[i] == 0)
+      readyNodes.push(i);
+
+  while (!readyNodes.empty()) {
+    size_t current = readyNodes.top();
+    readyNodes.pop();
+    timing.topologicalOrder.push_back(current);
+    for (auto [child, weight] : children[current]) {
+      (void) weight;
+      assert(incomingEdgeCount[child] > 0 && "incoming edge count underflow");
+      incomingEdgeCount[child]--;
+      if (incomingEdgeCount[child] == 0)
+        readyNodes.push(child);
+    }
+  }
+
+  if (timing.topologicalOrder.size() != nodeCount)
+    return timing;
+
+  Time dcpl = 0;
+  for (size_t nodeIndex : timing.topologicalOrder) {
+    Time maxParentAest = 0;
+    for (auto [parent, transferCost] : parents[nodeIndex]) {
+      maxParentAest =
+        std::max(maxParentAest, addOrMax(addOrMax(timing.aest[parent], graph.nodes[parent].weight), transferCost));
+    }
+    timing.aest[nodeIndex] = maxParentAest;
+    dcpl = std::max(dcpl, addOrMax(maxParentAest, graph.nodes[nodeIndex].weight));
+  }
+
+  for (size_t nodeIndex : llvm::reverse(timing.topologicalOrder)) {
+    Time minAlst = std::numeric_limits<Time>::max();
+    if (children[nodeIndex].empty())
+      minAlst = subtractOrZero(dcpl, graph.nodes[nodeIndex].weight);
+    for (auto [child, transferCost] : children[nodeIndex]) {
+      minAlst =
+        std::min(minAlst, subtractOrZero(timing.alst[child], addOrMax(graph.nodes[nodeIndex].weight, transferCost)));
+    }
+    timing.alst[nodeIndex] = minAlst;
+  }
+
+  timing.valid = true;
+  return timing;
+}
+
+std::vector<std::vector<size_t>> buildUndirectedAdjacency(const VirtualGraph& graph) {
+  std::vector<std::vector<size_t>> adjacency(graph.nodes.size());
+  for (auto [start, end, weight] : graph.edges) {
+    (void) weight;
+    size_t startIndex = static_cast<size_t>(start);
+    size_t endIndex = static_cast<size_t>(end);
+    assert(startIndex < graph.nodes.size() && endIndex < graph.nodes.size() && "virtual edge endpoint out of range");
+    adjacency[startIndex].push_back(endIndex);
+    adjacency[endIndex].push_back(startIndex);
+  }
+  for (auto& neighbours : adjacency) {
+    llvm::sort(neighbours);
+    neighbours.erase(std::unique(neighbours.begin(), neighbours.end()), neighbours.end());
+  }
+  return adjacency;
+}
+
+std::vector<size_t> selectCriticalWindow(const VirtualGraph& graph, const TimingInfo& timing, size_t windowSize) {
+  std::vector<size_t> ranked(timing.aest.size());
+  std::iota(ranked.begin(), ranked.end(), 0);
+  auto isHigherPriority = [&](size_t lhs, size_t rhs) {
+    Time lhsSlack = slackOrZero(timing.aest[lhs], timing.alst[lhs]);
+    Time rhsSlack = slackOrZero(timing.aest[rhs], timing.alst[rhs]);
+    if (lhsSlack != rhsSlack)
+      return lhsSlack < rhsSlack;
+    if (timing.aest[lhs] != timing.aest[rhs])
+      return timing.aest[lhs] < timing.aest[rhs];
+    return lhs < rhs;
+  };
+
+  windowSize = std::min(windowSize, ranked.size());
+  if (windowSize == 0)
+    return {};
+  if (windowSize == ranked.size()) {
+    llvm::sort(ranked, isHigherPriority);
+    return ranked;
+  }
+
+  size_t criticalPoolSize = std::min(ranked.size(), std::max(windowSize, windowSize * 2));
+  if (criticalPoolSize < ranked.size())
+    std::nth_element(
+      ranked.begin(), ranked.begin() + static_cast<std::ptrdiff_t>(criticalPoolSize), ranked.end(), isHigherPriority);
+
+  std::vector<char> inCriticalPool(ranked.size(), false);
+  for (size_t i = 0; i < criticalPoolSize; ++i)
+    inCriticalPool[ranked[i]] = true;
+
+  size_t seed = *std::min_element(ranked.begin(), ranked.end(), isHigherPriority);
+  std::vector<std::vector<size_t>> adjacency = buildUndirectedAdjacency(graph);
+  std::vector<size_t> selected;
+  std::vector<char> inWindow(ranked.size(), false);
+  selected.reserve(windowSize);
+
+  struct FrontierEntry {
+    size_t node;
+  };
+  auto frontierCompare = [&](FrontierEntry lhs, FrontierEntry rhs) { return isHigherPriority(rhs.node, lhs.node); };
+  std::priority_queue<FrontierEntry, std::vector<FrontierEntry>, decltype(frontierCompare)> frontier(frontierCompare);
+
+  auto addToWindow = [&](size_t node, const std::vector<char>& eligible) {
+    if (inWindow[node])
+      return;
+    inWindow[node] = true;
+    selected.push_back(node);
+    for (size_t neighbour : adjacency[node])
+      if (!inWindow[neighbour] && eligible[neighbour])
+        frontier.push({neighbour});
+  };
+
+  addToWindow(seed, inCriticalPool);
+  while (!frontier.empty() && selected.size() < windowSize) {
+    size_t node = frontier.top().node;
+    frontier.pop();
+    if (!inWindow[node])
+      addToWindow(node, inCriticalPool);
+  }
+
+  if (selected.size() < windowSize) {
+    std::vector<char> anyNode(ranked.size(), true);
+    for (size_t node : selected)
+      for (size_t neighbour : adjacency[node])
+        if (!inWindow[neighbour])
+          frontier.push({neighbour});
+    while (!frontier.empty() && selected.size() < windowSize) {
+      size_t node = frontier.top().node;
+      frontier.pop();
+      if (!inWindow[node])
+        addToWindow(node, anyNode);
+    }
+  }
+
+  if (selected.size() < windowSize) {
+    llvm::sort(ranked, isHigherPriority);
+    for (size_t node : ranked) {
+      if (selected.size() == windowSize)
+        break;
+      if (!inWindow[node]) {
+        inWindow[node] = true;
+        selected.push_back(node);
+      }
+    }
+  }
+
+  llvm::sort(selected, isHigherPriority);
+  return selected;
+}
+
+std::vector<IndexedEdge> buildWindowEdges(const VirtualGraph& graph, const std::vector<int64_t>& nodeToWindowIndex) {
+  std::vector<IndexedEdge> windowEdges;
+  windowEdges.reserve(graph.edges.size());
+  for (auto [start, end, weight] : graph.edges) {
+    int64_t mappedStart = nodeToWindowIndex[static_cast<size_t>(start)];
+    int64_t mappedEnd = nodeToWindowIndex[static_cast<size_t>(end)];
+    if (mappedStart == -1 || mappedEnd == -1)
+      continue;
+    windowEdges.push_back({mappedStart, mappedEnd, weight});
+  }
+  return aggregateEdges(windowEdges);
+}
+
+WindowScheduleResult scheduleWindow(const VirtualGraph& graph, ArrayRef<size_t> selectedNodes, MLIRContext* context) {
+  std::vector<Weight> windowWeights;
+  std::vector<CrossbarUsage> windowCrossbarUsage;
+  std::vector<int64_t> windowNodeOrderKeys;
+  std::vector<int64_t> nodeToWindowIndex(graph.nodes.size(), -1);
+  windowWeights.reserve(selectedNodes.size());
+  windowCrossbarUsage.reserve(selectedNodes.size());
+  windowNodeOrderKeys.reserve(selectedNodes.size());
+
+  for (auto [windowIndex, nodeIndex] : llvm::enumerate(selectedNodes)) {
+    nodeToWindowIndex[nodeIndex] = static_cast<int64_t>(windowIndex);
+    windowWeights.push_back(graph.nodes[nodeIndex].weight);
+    windowCrossbarUsage.push_back(graph.nodes[nodeIndex].crossbarUsage);
+    windowNodeOrderKeys.push_back(static_cast<int64_t>(nodeIndex));
+  }
+
+  GraphDCP windowGraph(
+    windowWeights, buildWindowEdges(graph, nodeToWindowIndex), windowNodeOrderKeys, windowCrossbarUsage);
+  if (coresCount.getValue() > 0)
+    windowGraph.setMaxCpuCount(static_cast<int>(coresCount.getValue()));
+  windowGraph.setContext(context);
+  windowGraph.runDcp();
+
+  WindowScheduleResult result;
+  result.cpuCount = windowGraph.cpuCount();
+  for (CPU cpu = 0; cpu < windowGraph.cpuCount(); ++cpu) {
+    auto scheduledTasks = windowGraph.getScheduledTasks(cpu);
+    if (scheduledTasks.size() < 2)
+      continue;
+
+    result.mergedNodeCount += scheduledTasks.size();
+    result.maxMergeGroupSize = std::max(result.maxMergeGroupSize, scheduledTasks.size());
+    std::vector<size_t> mergeGroup;
+    mergeGroup.reserve(scheduledTasks.size());
+    for (const auto& task : scheduledTasks)
+      mergeGroup.push_back(selectedNodes[task.nodeIndex]);
+    result.mergeGroups.push_back(std::move(mergeGroup));
+  }
+  return result;
+}
+
+bool coarsenGraph(const VirtualGraph& graph,
+                  ArrayRef<std::vector<size_t>> mergeGroups,
+                  VirtualGraph& coarsenedGraph,
+                  std::vector<size_t>& oldToNewNode) {
+  TimingInfo timing = computeTiming(graph);
+  std::vector<size_t> topologicalRank(graph.nodes.size());
+  std::iota(topologicalRank.begin(), topologicalRank.end(), 0);
+  if (timing.valid)
+    for (auto [rank, nodeIndex] : llvm::enumerate(timing.topologicalOrder))
+      topologicalRank[nodeIndex] = rank;
+
+  std::vector<std::vector<size_t>> orderedMergeGroups;
+  orderedMergeGroups.reserve(mergeGroups.size());
+  for (const auto& mergeGroup : mergeGroups) {
+    orderedMergeGroups.emplace_back(mergeGroup.begin(), mergeGroup.end());
+    std::stable_sort(orderedMergeGroups.back().begin(), orderedMergeGroups.back().end(), [&](size_t lhs, size_t rhs) {
+      if (topologicalRank[lhs] != topologicalRank[rhs])
+        return topologicalRank[lhs] < topologicalRank[rhs];
+      return lhs < rhs;
+    });
+  }
+
+  std::vector<int64_t> nodeToMergeGroup(graph.nodes.size(), -1);
+  for (auto [groupIndex, mergeGroup] : llvm::enumerate(orderedMergeGroups)) {
+    if (mergeGroup.size() < 2)
+      continue;
+    for (size_t nodeIndex : mergeGroup) {
+      assert(nodeIndex < graph.nodes.size() && "merge group node out of range");
+      nodeToMergeGroup[nodeIndex] = static_cast<int64_t>(groupIndex);
+    }
+  }
+
+  std::vector<std::optional<size_t>> mergeGroupToNewNode(orderedMergeGroups.size());
+  std::vector<size_t> newNodeRank;
+  oldToNewNode.assign(graph.nodes.size(), 0);
+  bool mergedAny = false;
+  coarsenedGraph.nodes.clear();
+  coarsenedGraph.edges.clear();
+  coarsenedGraph.nodes.reserve(graph.nodes.size());
+  newNodeRank.reserve(graph.nodes.size());
+
+  for (size_t nodeIndex = 0; nodeIndex < graph.nodes.size(); ++nodeIndex) {
+    int64_t mergeGroupIndex = nodeToMergeGroup[nodeIndex];
+    if (mergeGroupIndex == -1) {
+      oldToNewNode[nodeIndex] = coarsenedGraph.nodes.size();
+      coarsenedGraph.nodes.push_back(graph.nodes[nodeIndex]);
+      newNodeRank.push_back(topologicalRank[nodeIndex]);
+      continue;
+    }
+
+    auto& newNodeIndex = mergeGroupToNewNode[static_cast<size_t>(mergeGroupIndex)];
+    if (newNodeIndex.has_value()) {
+      oldToNewNode[nodeIndex] = *newNodeIndex;
+      continue;
+    }
+
+    VirtualNode mergedNode;
+    for (size_t memberIndex : orderedMergeGroups[static_cast<size_t>(mergeGroupIndex)]) {
+      const VirtualNode& memberNode = graph.nodes[memberIndex];
+      mergedNode.originalComputeIndices.append(memberNode.originalComputeIndices.begin(),
+                                               memberNode.originalComputeIndices.end());
+      mergedNode.weight = addOrMax(mergedNode.weight, memberNode.weight);
+      mergedNode.crossbarUsage = addOrMax(mergedNode.crossbarUsage, memberNode.crossbarUsage);
+    }
+    std::sort(mergedNode.originalComputeIndices.begin(), mergedNode.originalComputeIndices.end());
+
+    mergedAny = true;
+    newNodeIndex = coarsenedGraph.nodes.size();
+    for (size_t memberIndex : orderedMergeGroups[static_cast<size_t>(mergeGroupIndex)])
+      oldToNewNode[memberIndex] = *newNodeIndex;
+    newNodeRank.push_back(topologicalRank[orderedMergeGroups[static_cast<size_t>(mergeGroupIndex)].front()]);
+    coarsenedGraph.nodes.push_back(std::move(mergedNode));
+  }
+
+  if (!mergedAny)
+    return false;
+
+  std::vector<IndexedEdge> remappedEdges;
+  remappedEdges.reserve(graph.edges.size());
+  for (auto [start, end, weight] : graph.edges) {
+    size_t newStart = oldToNewNode[static_cast<size_t>(start)];
+    size_t newEnd = oldToNewNode[static_cast<size_t>(end)];
+    if (newStart == newEnd)
+      continue;
+    if (newNodeRank[newStart] >= newNodeRank[newEnd])
+      continue;
+    remappedEdges.push_back({static_cast<int64_t>(newStart), static_cast<int64_t>(newEnd), weight});
+  }
+  coarsenedGraph.edges = aggregateEdges(remappedEdges);
+
+  return true;
+}
+
+CPU getVirtualGraphMaxCpuCount() { return static_cast<CPU>(getSchedulingCpuBudget()); }
+
+size_t getDcpCoarseningWindowSize(size_t nodeCount) {
+  size_t windowSize = std::min(dcpCriticalWindowSize.getValue(), nodeCount);
+  CPU maxCpuCount = std::max<CPU>(1, getVirtualGraphMaxCpuCount());
+  if (nodeCount > static_cast<size_t>(maxCpuCount))
+    windowSize = std::max(windowSize, std::min(nodeCount, static_cast<size_t>(maxCpuCount) + 1));
+  return windowSize;
+}
+
+DCPAnalysisResult buildResultFromVirtualGraph(const VirtualGraph& graph, ArrayRef<ComputeInstance> computeInstances) {
+  DCPAnalysisResult result;
+
+  TimingInfo timing = computeTiming(graph);
+  std::vector<size_t> virtualNodeOrder;
+  if (timing.valid) {
+    virtualNodeOrder = std::move(timing.topologicalOrder);
+  }
+  else {
+    virtualNodeOrder.resize(graph.nodes.size());
+    std::iota(virtualNodeOrder.begin(), virtualNodeOrder.end(), 0);
+  }
+
+  std::vector<size_t> originalComputeToCpu(computeInstances.size(), 0);
+  for (auto [cpu, virtualNodeIndex] : llvm::enumerate(virtualNodeOrder)) {
+    const VirtualNode& virtualNode = graph.nodes[virtualNodeIndex];
+    for (size_t originalIndex : virtualNode.originalComputeIndices)
+      originalComputeToCpu[originalIndex] = cpu;
+  }
+
+  result.dominanceOrderCompute.reserve(computeInstances.size());
+  for (auto [originalIndex, computeInstance] : llvm::enumerate(computeInstances)) {
+    size_t cpu = originalComputeToCpu[originalIndex];
+    result.dominanceOrderCompute.push_back(computeInstance);
+    result.computeToCpuMap[computeInstance] = cpu;
+    result.cpuToLastComputeMap[cpu] = computeInstance;
+  }
+  for (const auto& [cpu, lastCompute] : result.cpuToLastComputeMap)
+    result.isLastComputeOfCpu.insert(lastCompute);
+
+  return result;
+}
+
+DCPAnalysisResult buildResultFromScheduledGraph(GraphDCP& graphDCP, ArrayRef<ComputeInstance> computeInstances) {
+  DCPAnalysisResult result;
+  result.dominanceOrderCompute.assign(computeInstances.begin(), computeInstances.end());
+
+  for (CPU cpu = 0; cpu < graphDCP.cpuCount(); ++cpu) {
+    auto scheduledTasks = graphDCP.getScheduledTasks(cpu);
+    if (scheduledTasks.empty())
+      continue;
+
+    for (const auto& task : scheduledTasks)
+      result.computeToCpuMap[computeInstances[task.nodeIndex]] = cpu;
+    result.cpuToLastComputeMap[cpu] = computeInstances[scheduledTasks.back().nodeIndex];
+    result.isLastComputeOfCpu.insert(computeInstances[scheduledTasks.back().nodeIndex]);
+  }
+
+  return result;
+}
+
+DCPAnalysisResult
+runLegacyDcp(ArrayRef<ComputeInstance> computeInstances, ArrayRef<IndexedEdge> edges, MLIRContext* context) {
+  SmallVector<Weight> nodeWeights;
+  SmallVector<CrossbarUsage> nodeCrossbarUsage;
+  SmallVector<int64_t> nodeOrderKeys;
+  nodeWeights.reserve(computeInstances.size());
+  nodeCrossbarUsage.reserve(computeInstances.size());
+  nodeOrderKeys.reserve(computeInstances.size());
+  for (auto [index, instance] : llvm::enumerate(computeInstances)) {
+    nodeWeights.push_back(getComputeInstanceWeight(instance));
+    nodeCrossbarUsage.push_back(getComputeInstanceCrossbarUsage(instance));
+    nodeOrderKeys.push_back(static_cast<int64_t>(index));
+  }
+
+  GraphDCP graphDCP(nodeWeights, edges, nodeOrderKeys, nodeCrossbarUsage);
+  if (coresCount.getValue() > 0)
+    graphDCP.setMaxCpuCount(static_cast<int>(coresCount.getValue()));
+  graphDCP.setContext(context);
+  graphDCP.runDcp();
+  return buildResultFromScheduledGraph(graphDCP, computeInstances);
+}
+
+} // namespace
+
+SpatCompute getOriginalSpatCompute(Operation* op) {
   if (!op)
     return {};
-  while (auto extract = llvm::dyn_cast<tensor::ExtractSliceOp>(op)) {
+  while (auto extract = dyn_cast<tensor::ExtractSliceOp>(op)) {
     op = extract.getSource().getDefiningOp();
     if (!op)
       return {};
   }
-  if (auto res = llvm::dyn_cast<SpatWeightedCompute>(op))
+  if (auto res = dyn_cast<SpatCompute>(op))
     return res;
   return {};
 }
 
 DCPAnalysisResult DCPAnalysis::run() {
-  llvm::SmallVector<SpatWeightedCompute, 10> spatWeightedComputes;
-  llvm::SmallVector<IndexedEdge, 10> edges;
-  for (auto& region : entryOp->getRegions())
-    for (SpatWeightedCompute spatWeightedCompute : region.getOps<SpatWeightedCompute>())
-      spatWeightedComputes.push_back(spatWeightedCompute);
+  SmallVector<ComputeInstance> computeInstances = collectComputeInstances(entryOp);
+  SmallVector<IndexedEdge, 10> edges;
 
-  for (auto [indexEndEdge, spatWeightedCompute] : llvm::enumerate(spatWeightedComputes)) {
-    for (Value input : spatWeightedCompute.getInputs()) {
-      if (auto producerCompute = getOriginalSpatWeightedCompute(input.getDefiningOp())) {
-        auto producerIt = llvm::find(spatWeightedComputes, producerCompute);
-        assert(producerIt != spatWeightedComputes.end());
-        auto indexStartEdge = std::distance(spatWeightedComputes.begin(), producerIt);
-        ResultRange outputs = producerCompute.getResults();
-        int64_t totalSize = 0;
-        for (auto output : outputs) {
-          ShapedType resultType = cast<ShapedType>(output.getType());
-          totalSize += getSizeInBytes(resultType);
-        }
-        edges.push_back({indexStartEdge, indexEndEdge, totalSize});
+  llvm::DenseMap<ComputeInstance, size_t> instanceToIndex;
+  instanceToIndex.reserve(computeInstances.size());
+  for (auto [index, instance] : llvm::enumerate(computeInstances))
+    instanceToIndex[instance] = index;
+
+  for (auto [indexEndEdge, computeInstance] : llvm::enumerate(computeInstances)) {
+    for (Value input : getComputeInstanceInputs(computeInstance)) {
+      if (auto producerInstance = getOriginalComputeInstance(input)) {
+        auto producerIt = instanceToIndex.find(*producerInstance);
+        assert(producerIt != instanceToIndex.end());
+        auto indexStartEdge = producerIt->second;
+        edges.push_back({static_cast<int64_t>(indexStartEdge),
+                         static_cast<int64_t>(indexEndEdge),
+                         static_cast<int64_t>(getSizeInBytes(cast<ShapedType>(input.getType())))});
       }
     }
   }
-  GraphDCP graphDCP(spatWeightedComputes, edges);
-  graphDCP.setContext(entryOp->getContext());
-  graphDCP.runDcp();
-  return graphDCP.getResult();
+
+  if (dcpCriticalWindowSize.getValue() == 0)
+    return runLegacyDcp(computeInstances, edges, entryOp->getContext());
+
+  VirtualGraph virtualGraph = buildInitialVirtualGraph(computeInstances, edges);
+  size_t iteration = 0;
+  auto tryCoarsenSelectedNodes = [&](ArrayRef<size_t> selectedNodes) {
+    size_t oldNodeCount = virtualGraph.nodes.size();
+    WindowScheduleResult windowSchedule = scheduleWindow(virtualGraph, selectedNodes, entryOp->getContext());
+    if (windowSchedule.mergeGroups.empty()) {
+      if (oldNodeCount >= 200)
+        llvm::errs() << llvm::formatv("[DCP-COARSEN] iter={0} old={1} selected={2} windowCpus={3} "
+                                      "groups=0 mergedNodes=0 maxGroup=0 new={1} changed=0\n",
+                                      iteration,
+                                      oldNodeCount,
+                                      selectedNodes.size(),
+                                      windowSchedule.cpuCount);
+      return false;
+    }
+
+    VirtualGraph coarsenedGraph;
+    std::vector<size_t> oldToNewNode;
+    if (!coarsenGraph(virtualGraph, windowSchedule.mergeGroups, coarsenedGraph, oldToNewNode))
+      return false;
+    if (oldNodeCount >= 200 || coarsenedGraph.nodes.size() >= 200)
+      llvm::errs() << llvm::formatv("[DCP-COARSEN] iter={0} old={1} selected={2} windowCpus={3} "
+                                    "groups={4} mergedNodes={5} maxGroup={6} new={7} changed={8}\n",
+                                    iteration,
+                                    oldNodeCount,
+                                    selectedNodes.size(),
+                                    windowSchedule.cpuCount,
+                                    windowSchedule.mergeGroups.size(),
+                                    windowSchedule.mergedNodeCount,
+                                    windowSchedule.maxMergeGroupSize,
+                                    coarsenedGraph.nodes.size(),
+                                    oldNodeCount - coarsenedGraph.nodes.size());
+    virtualGraph = std::move(coarsenedGraph);
+    return true;
+  };
+
+  while (virtualGraph.nodes.size() > 1) {
+    if (virtualGraph.nodes.size() <= getSchedulingCpuBudget()) {
+      if (virtualGraph.nodes.size() >= 200)
+        llvm::errs() << llvm::formatv(
+          "[DCP-COARSEN] iter={0} old={1} stop=cpu-budget\n", iteration, virtualGraph.nodes.size());
+      break;
+    }
+
+    iteration++;
+    TimingInfo timing = computeTiming(virtualGraph);
+    if (!timing.valid) {
+      if (virtualGraph.nodes.size() >= 200)
+        llvm::errs() << llvm::formatv(
+          "[DCP-COARSEN] iter={0} old={1} invalid-timing\n", iteration, virtualGraph.nodes.size());
+      break;
+    }
+
+    SmallVector<size_t> selectedNodes;
+    auto criticalWindow =
+      selectCriticalWindow(virtualGraph, timing, getDcpCoarseningWindowSize(virtualGraph.nodes.size()));
+    selectedNodes.append(criticalWindow.begin(), criticalWindow.end());
+
+    if (selectedNodes.size() < 2) {
+      if (virtualGraph.nodes.size() >= 200)
+        llvm::errs() << llvm::formatv("[DCP-COARSEN] iter={0} old={1} selected={2} stop=small-window\n",
+                                      iteration,
+                                      virtualGraph.nodes.size(),
+                                      selectedNodes.size());
+      break;
+    }
+
+    if (tryCoarsenSelectedNodes(selectedNodes))
+      continue;
+    if (virtualGraph.nodes.size() >= 200)
+      llvm::errs() << llvm::formatv(
+        "[DCP-COARSEN] iter={0} old={1} stop=no-merge\n", iteration, virtualGraph.nodes.size());
+    break;
+  }
+
+  return buildResultFromVirtualGraph(virtualGraph, computeInstances);
 }
 
 } // namespace spatial

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/DCPAnalysis.hpp

@@ -5,15 +5,28 @@
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 
+#include <cstdint>
 #include <vector>
 
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
+// A scheduling identity that covers both spat.compute and scheduled shards of
+// spat.compute_batch.
+struct ComputeInstance {
+  mlir::Operation* op = nullptr;
+  uint32_t laneStart = 0;
+  uint32_t laneCount = 1;
+
+  bool operator==(const ComputeInstance& other) const {
+    return op == other.op && laneStart == other.laneStart && laneCount == other.laneCount;
+  }
+};
+
 struct DCPAnalysisResult {
-  std::vector<onnx_mlir::spatial::SpatWeightedCompute> dominanceOrderCompute;
-  llvm::DenseMap<onnx_mlir::spatial::SpatWeightedCompute, size_t> computeToCpuMap;
-  llvm::DenseSet<onnx_mlir::spatial::SpatWeightedCompute> isLastComputeOfCpu;
-  llvm::DenseMap<size_t, onnx_mlir::spatial::SpatWeightedCompute> cpuToLastComputeMap;
+  std::vector<ComputeInstance> dominanceOrderCompute;
+  llvm::DenseMap<ComputeInstance, size_t> computeToCpuMap;
+  llvm::DenseSet<ComputeInstance> isLastComputeOfCpu;
+  llvm::DenseMap<size_t, ComputeInstance> cpuToLastComputeMap;
 };
 
 namespace onnx_mlir {
@@ -34,3 +47,21 @@ public:
 
 } // namespace spatial
 } // namespace onnx_mlir
+
+namespace llvm {
+template <>
+struct DenseMapInfo<ComputeInstance> {
+  static ComputeInstance getEmptyKey() {
+    return {DenseMapInfo<mlir::Operation*>::getEmptyKey(), UINT32_MAX, UINT32_MAX};
+  }
+  static ComputeInstance getTombstoneKey() {
+    return {DenseMapInfo<mlir::Operation*>::getTombstoneKey(), UINT32_MAX, UINT32_MAX};
+  }
+  static unsigned getHashValue(const ComputeInstance& v) {
+    return llvm::hash_combine(v.op, v.laneStart, v.laneCount);
+  }
+  static bool isEqual(const ComputeInstance& a, const ComputeInstance& b) {
+    return a == b;
+  }
+};
+} // namespace llvm

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/Graph.cpp

@@ -6,7 +6,7 @@
 // consumer land on different CPUs.
 //
 // Output: an assignment of every task to a CPU and an order within that CPU,
-// aiming to minimise the overall critical-path length (DCPL).
+// aiming to minimize the overall critical-path length (DCPL).
 //
 // Every task keeps two timing estimates:
 //   AEST - earliest start time, driven by parent completions + transfers.
@@ -16,9 +16,9 @@
 // Main loop (runDcp):
 //   1. Build a topological order and seed AEST/ALST from the unscheduled DAG.
 //   2. While there are ready tasks (all dependency parents scheduled):
-//        a. Pick the candidate with tightest slack (earliest AEST breaks ties).
+//        a. Pick the candidate with the tightest slack (earliest AEST breaks ties).
 //        b. selectProcessor() tries every candidate CPU and picks the one that
-//           minimises a composite cost (own slot + smallest unscheduled child).
+//           minimizes a composite cost (own slot + the smallest unscheduled child).
 //        c. Commit the placement and refresh AEST/ALST.
 //        d. Release any child whose dependency parents are now all scheduled.
 //
@@ -38,12 +38,14 @@
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/ErrorHandling.h"
 
 #include <algorithm>
 #include <cassert>
 #include <chrono>
+#include <cstdint>
 #include <cstdio>
-#include <cstdlib>
+#include <queue>
 #include <vector>
 
 #include "DCPAnalysis.hpp"
@@ -61,6 +63,7 @@ namespace {
 // Coarse-grained phase timers printed when DCP_SELECT_PROFILE is set.
 struct SelectTimers {
   double findSlot = 0.0;
+  double dedup = 0.0;
   double precheck = 0.0;
   double snapshotInsertUpdate = 0.0;
   double childSlot = 0.0;
@@ -71,9 +74,19 @@ struct SelectTimers {
   long tasksProcessed = 0;
   void dump(const char* label) const {
     std::fprintf(stderr,
-                 "[selectProfile:%s] tasks=%ld findSlot=%.2fs precheck=%.2fs snapUpd=%.2fs childSlot=%.2fs rollback=%.2fs iter=%ld precheckPass=%ld dcplPass=%ld\n",
-                 label, tasksProcessed, findSlot, precheck, snapshotInsertUpdate, childSlot,
-                 rollbackRestore, iterations, passedPrecheck, passedDcpl);
+                 "[selectProfile:%s] tasks=%ld dedup=%.2fs findSlot=%.2fs precheck=%.2fs snapUpd=%.2fs "
+                 "childSlot=%.2fs rollback=%.2fs iter=%ld precheckPass=%ld dcplPass=%ld\n",
+                 label,
+                 tasksProcessed,
+                 dedup,
+                 findSlot,
+                 precheck,
+                 snapshotInsertUpdate,
+                 childSlot,
+                 rollbackRestore,
+                 iterations,
+                 passedPrecheck,
+                 passedDcpl);
   }
   ~SelectTimers() {
     if (std::getenv("DCP_SELECT_PROFILE"))
@@ -84,6 +97,101 @@ static SelectTimers gSelectTimers;
 } // namespace
 #endif
 
+namespace {
+
+uint64_t mixHash(uint64_t seed, uint64_t value) {
+  seed ^= value + 0x9e3779b97f4a7c15ULL + (seed << 6) + (seed >> 2);
+  return seed;
+}
+
+uint64_t finishHash(uint64_t seed) {
+  seed ^= seed >> 33;
+  seed *= 0xff51afd7ed558ccdULL;
+  seed ^= seed >> 33;
+  seed *= 0xc4ceb9fe1a85ec53ULL;
+  seed ^= seed >> 33;
+  return seed;
+}
+
+uint64_t hashEdgeSignature(uint64_t neighborHash, Weight weight, uint64_t direction) {
+  uint64_t hash = mixHash(0x84222325cbf29ce4ULL, direction);
+  hash = mixHash(hash, neighborHash);
+  hash = mixHash(hash, static_cast<uint64_t>(weight));
+  return finishHash(hash);
+}
+
+struct CpuAestCache {
+  Time defaultAest = 0;
+  llvm::SmallDenseMap<CPU, Time, 8> colocatedParentAests;
+
+  Time get(CPU cpu) const {
+    auto it = colocatedParentAests.find(cpu);
+    if (it == colocatedParentAests.end())
+      return defaultAest;
+    return it->second;
+  }
+};
+
+struct CpuTimeMax {
+  CPU cpu = -1;
+  Time time = 0;
+};
+
+void updateCpuTimeMax(CpuTimeMax& first, CpuTimeMax& second, CPU cpu, Time time) {
+  if (first.cpu == cpu) {
+    first.time = std::max(first.time, time);
+    return;
+  }
+  if (second.cpu == cpu) {
+    second.time = std::max(second.time, time);
+    if (second.time > first.time)
+      std::swap(first, second);
+    return;
+  }
+  if (time >= first.time) {
+    second = first;
+    first = {cpu, time};
+    return;
+  }
+  if (time > second.time)
+    second = {cpu, time};
+}
+
+CpuAestCache computeCpuAestCache(TaskDCP* task) {
+  CpuAestCache cache;
+  llvm::SmallDenseMap<CPU, Time, 8> transferAestByCpu;
+  llvm::SmallDenseMap<CPU, Time, 8> localAestByCpu;
+  Time unscheduledTransferAest = 0;
+
+  for (const Edge& parentEdge : task->parents) {
+    Time parentFinish = addOrMax(parentEdge.first->getAest(), parentEdge.first->getWeight());
+    Time transferAest = addOrMax(parentFinish, getTransferCost(parentEdge.first, task));
+    if (std::optional<CPU> parentCpu = parentEdge.first->getCpu()) {
+      Time& cpuTransferAest = transferAestByCpu[*parentCpu];
+      cpuTransferAest = std::max(cpuTransferAest, transferAest);
+      Time& cpuLocalAest = localAestByCpu[*parentCpu];
+      cpuLocalAest = std::max(cpuLocalAest, parentFinish);
+      continue;
+    }
+    unscheduledTransferAest = std::max(unscheduledTransferAest, transferAest);
+  }
+
+  CpuTimeMax firstOther {-1, unscheduledTransferAest};
+  CpuTimeMax secondOther {-1, 0};
+  for (const auto& entry : transferAestByCpu)
+    updateCpuTimeMax(firstOther, secondOther, entry.first, entry.second);
+
+  cache.defaultAest = firstOther.time;
+  for (const auto& entry : localAestByCpu) {
+    CPU cpu = entry.first;
+    Time bestNonLocalParentAest = firstOther.cpu == cpu ? secondOther.time : firstOther.time;
+    cache.colocatedParentAests[cpu] = std::max(bestNonLocalParentAest, entry.second);
+  }
+  return cache;
+}
+
+} // namespace
+
 //===----------------------------------------------------------------------===//
 // Edge manipulation
 //===----------------------------------------------------------------------===//
@@ -157,6 +265,49 @@ std::vector<TaskDCP*> GraphDCP::getRoots() {
   return tmp;
 }
 
+void GraphDCP::initTaskStructureHashes() {
+  taskStructureHashes.resize(nodes.size());
+  for (auto [index, task] : llvm::enumerate(nodes)) {
+    uint64_t hash = mixHash(0x7442b1129fd01363ULL, static_cast<uint64_t>(task.getWeight()));
+    hash = mixHash(hash, static_cast<uint64_t>(task.getCrossbarUsage()));
+    taskStructureHashes[index] = finishHash(hash);
+  }
+
+  std::vector<uint64_t> nextHashes(nodes.size());
+  std::vector<uint64_t> edgeHashes;
+  for (int iteration = 0; iteration < 4; ++iteration) {
+    for (auto [index, task] : llvm::enumerate(nodes)) {
+      uint64_t hash = mixHash(0x464dcab27ac82291ULL, taskStructureHashes[index]);
+      edgeHashes.clear();
+      edgeHashes.reserve(task.parents.size() + task.children.size());
+      for (const Edge& parent : task.parents)
+        if (!parent.isScheduling)
+          edgeHashes.push_back(
+            hashEdgeSignature(taskStructureHashes[getNodeIndex(parent.first)], parent.second, /*direction=*/0));
+      for (const Edge& child : task.children)
+        if (!child.isScheduling)
+          edgeHashes.push_back(
+            hashEdgeSignature(taskStructureHashes[getNodeIndex(child.first)], child.second, /*direction=*/1));
+      llvm::sort(edgeHashes);
+      hash = mixHash(hash, static_cast<uint64_t>(edgeHashes.size()));
+      for (uint64_t edgeHash : edgeHashes)
+        hash = mixHash(hash, edgeHash);
+      nextHashes[index] = finishHash(hash);
+    }
+    taskStructureHashes.swap(nextHashes);
+  }
+}
+
+// Compact dedup key for CPU `c` vs `candidate`: mixes candidateAest, crossbar
+// usage, and the incremental cpu structure hash. No heap allocation.
+uint64_t GraphDCP::computeCpuCandidateKey(Time candidateAest, CPU cpu) {
+  uint64_t hash = mixHash(0xd6e8feb86659fd93ULL, static_cast<uint64_t>(candidateAest));
+  hash = mixHash(hash, static_cast<uint64_t>(getCpuCrossbarUsage(cpu)));
+  auto it = cpuStructureHashes.find(cpu);
+  hash = mixHash(hash, it != cpuStructureHashes.end() ? it->second : 0ULL);
+  return finishHash(hash);
+}
+
 // Inserts `task` at `position` on `cpu`, wiring up scheduling edges with the
 // neighbouring tasks and keeping the global topological order consistent.
 TaskInsertion GraphDCP::insertTaskInCPU(CPU cpu, TaskDCP* task, size_t position) {
@@ -165,6 +316,7 @@ TaskInsertion GraphDCP::insertTaskInCPU(CPU cpu, TaskDCP* task, size_t position)
   task->setCpu(cpu);
   task->setWeight(scheduledWeight);
   reserveTaskCrossbars(cpu, task);
+  cpuStructureHashes[cpu] ^= taskStructureHashes[getNodeIndex(task)];
   auto& tasksInCpu = getOrCreateCpuTasks(cpu);
   unsigned int numCpuTasks = tasksInCpu.size();
   assert(position <= numCpuTasks && "Inserting in a not valid position");
@@ -202,6 +354,7 @@ TaskInsertion GraphDCP::insertTaskInCPU(CPU cpu, TaskDCP* task, size_t position)
 
 void GraphDCP::removeTaskFromCPU(CPU cpu, TaskDCP* task) {
   releaseTaskCrossbars(cpu, task);
+  cpuStructureHashes[cpu] ^= taskStructureHashes[getNodeIndex(task)];
   task->resetCpu();
   task->resetWeight();
   auto& scheduledTasks = getOrCreateCpuTasks(cpu);
@@ -272,6 +425,21 @@ bool GraphDCP::wouldExhaustCrossbarCapacity(CPU cpu, const TaskDCP* task) const
   return nextUsage >= getCpuCrossbarCapacity();
 }
 
+size_t GraphDCP::crossbarsUsed() const {
+  CrossbarUsage crossbarEdge = static_cast<CrossbarUsage>(onnx_mlir::crossbarSize.getValue());
+  CrossbarUsage crossbarArea = crossbarEdge * crossbarEdge;
+  if (crossbarArea == 0)
+    return 0;
+  CrossbarUsage totalArea = 0;
+  for (const auto& [cpu, usage] : cpuCrossbarUsage)
+    totalArea = checkedAdd(totalArea, usage);
+  return static_cast<size_t>(totalArea / crossbarArea);
+}
+
+size_t GraphDCP::crossbarsAvailable() const {
+  return static_cast<size_t>(lastCpu) * onnx_mlir::crossbarCountInCore.getValue();
+}
+
 //===----------------------------------------------------------------------===//
 // AEST / ALST computation
 //===----------------------------------------------------------------------===//
@@ -457,9 +625,9 @@ void GraphDCP::updateAestFromTaskWithDescendants(TaskDCP* task, llvm::ArrayRef<T
   for (TaskDCP* descendant : descendantsTopoOrder)
     recomputeAest(descendant);
 
-  const bool oldMaxInvalidated = maxCompletionTask != nullptr
-                              && (maxCompletionTask == task
-                                  || llvm::is_contained(descendantsTopoOrder, maxCompletionTask));
+  const bool oldMaxInvalidated =
+    maxCompletionTask != nullptr
+    && (maxCompletionTask == task || llvm::is_contained(descendantsTopoOrder, maxCompletionTask));
   if (oldMaxInvalidated) {
     // The pre-update max came from a modified task; its completion has moved
     // upward, so modifiedMaxCompletion is an upper bound covering it. The
@@ -524,9 +692,9 @@ bool GraphDCP::tryUpdateAestWithinBudget(TaskDCP* task,
     if (!process(descendant))
       return false;
 
-  const bool oldMaxInvalidated = maxCompletionTask != nullptr
-                              && (maxCompletionTask == task
-                                  || llvm::is_contained(descendantsTopoOrder, maxCompletionTask));
+  const bool oldMaxInvalidated =
+    maxCompletionTask != nullptr
+    && (maxCompletionTask == task || llvm::is_contained(descendantsTopoOrder, maxCompletionTask));
   if (oldMaxInvalidated) {
     dcpl = modifiedMaxCompletion;
     maxCompletion = modifiedMaxCompletion;
@@ -547,6 +715,109 @@ bool GraphDCP::tryUpdateAestWithinBudget(TaskDCP* task,
   return true;
 }
 
+// Incrementally refreshes ALST after `task` was placed. The set of nodes whose
+// ALST is structurally affected by the insertion is exactly
+// `relations.ancestors ∪ {task}`: the task's outgoing transfer costs to
+// same-CPU real children become 0, and new scheduling edges create parent
+// relationships between `task` and its same-CPU neighbors. Every other node
+// keeps its relative distance to the sink boundary and only absorbs the
+// signed DCPL delta captured between `oldDcpl` and the now-updated `dcpl`.
+void GraphDCP::updateAlstFromScheduledTask(TaskDCP* task, const CandidateRelations& relations, Time oldDcpl) {
+  Time newDcpl = getDcpl();
+  // If the AEST update saturated dcpl (e.g. rescue placement on a
+  // crossbar-exhausted CPU sets task weight to UINT64_MAX), the shift delta
+  // would be meaningless. Fall back to a full recompute for this step only.
+  if (newDcpl == std::numeric_limits<Time>::max()) {
+    initAlst();
+    return;
+  }
+
+  if (newDcpl != oldDcpl) {
+    const bool increased = newDcpl > oldDcpl;
+    const Time delta = increased ? (newDcpl - oldDcpl) : (oldDcpl - newDcpl);
+    for (TaskDCP& node : topologicalOrder) {
+      if (&node == task || relations.ancestors.contains(&node))
+        continue;
+      Time alst = node.getAlst();
+      node.setAlst(increased ? addOrMax(alst, delta) : subtractOrZero(alst, delta));
+    }
+  }
+
+  auto recomputeAlst = [&](TaskDCP* node) {
+    Time minAlst = std::numeric_limits<Time>::max();
+    if (!node->hasChildren())
+      minAlst = subtractOrZero(newDcpl, node->getWeight());
+    for (const Edge& childEdge : node->children)
+      minAlst = std::min(minAlst,
+                         subtractOrZero(childEdge.first->getAlst(),
+                                        addOrMax(node->getWeight(), getTransferCost(node, childEdge.first))));
+    node->setAlst(minAlst);
+  };
+
+  // Walk the backward cone with a pending-children counter so that every
+  // ancestor is recomputed only after all of its affected children have
+  // been refreshed. This is resilient to staleness in the global
+  // `topologicalOrder` relative to freshly added scheduling edges.
+  llvm::DenseSet<TaskDCP*> affected = relations.ancestors;
+  affected.insert(task);
+  llvm::DenseMap<TaskDCP*, int> pendingAffectedChildren;
+  pendingAffectedChildren.reserve(affected.size());
+  std::vector<TaskDCP*> worklist;
+  worklist.reserve(affected.size());
+  for (TaskDCP* node : affected) {
+    int count = 0;
+    for (const Edge& childEdge : node->children)
+      if (affected.contains(childEdge.first))
+        count++;
+    pendingAffectedChildren[node] = count;
+    if (count == 0)
+      worklist.push_back(node);
+  }
+  while (!worklist.empty()) {
+    TaskDCP* node = worklist.back();
+    worklist.pop_back();
+    recomputeAlst(node);
+    for (const Edge& parentEdge : node->parents) {
+      if (!affected.contains(parentEdge.first))
+        continue;
+      auto it = pendingAffectedChildren.find(parentEdge.first);
+      assert(it != pendingAffectedChildren.end());
+      if (--it->second == 0)
+        worklist.push_back(parentEdge.first);
+    }
+  }
+
+// Opt-in consistency check: verifies the incremental ALST result against a
+// full initAlst() recomputation. Very expensive (O(V+E) per placement) - only
+// enable when investigating suspected drift.
+#ifdef DCP_DEBUG_CHECK_ALST
+  std::vector<Time> afterIncremental(nodes.size());
+  for (size_t i = 0; i < nodes.size(); ++i)
+    afterIncremental[i] = nodes[i].getAlst();
+  initAlst();
+  bool mismatched = false;
+  for (size_t i = 0; i < nodes.size(); ++i) {
+    if (afterIncremental[i] != nodes[i].getAlst()) {
+      if (!mismatched) {
+        llvm::errs() << "[alst-mismatch] placed=" << getNodeIndex(task) << " oldDcpl=" << oldDcpl
+                     << " newDcpl=" << newDcpl << " ancestors={";
+        for (TaskDCP* a : relations.ancestors)
+          llvm::errs() << getNodeIndex(a) << ",";
+        llvm::errs() << "}\n";
+        mismatched = true;
+      }
+      llvm::errs() << "  node=" << i << " incremental=" << afterIncremental[i] << " full=" << nodes[i].getAlst()
+                   << " weight=" << nodes[i].getWeight()
+                   << " cpu=" << (nodes[i].isScheduled() ? (int) *nodes[i].getCpu() : -1) << " children=[";
+      for (const Edge& e : nodes[i].children)
+        llvm::errs() << getNodeIndex(e.first) << (e.isScheduling ? "s" : "")
+                     << "(tc=" << getTransferCost(&nodes[i], e.first) << ",alst=" << e.first->getAlst() << "),";
+      llvm::errs() << "]\n";
+    }
+  }
+#endif
+}
+
 // Computes a localised ALST: only ancestors of the candidate (plus the
 // candidate itself) get recomputed, every other task keeps its current ALST.
 // Processes nodes in reverse dependency order using a pending-children
@@ -906,32 +1177,6 @@ GraphDCP::FindSlot GraphDCP::findSlotWithFixedFinalTime(
 // Candidate selection and processor assignment
 //===----------------------------------------------------------------------===//
 
-// Lowest slack wins; earliest AEST breaks ties. Critical-path tasks (zero
-// slack) naturally float to the front.
-TaskDCP* GraphDCP::findCandidate(const std::vector<TaskDCP*>& readyNodes) {
-  auto findBestNode = [](auto lft, auto rgt) {
-    Time leftSlack = slackOrZero((*lft)->getAest(), (*lft)->getAlst());
-    Time rightSlack = slackOrZero((*rgt)->getAest(), (*rgt)->getAlst());
-    if (leftSlack < rightSlack)
-      return lft;
-    if (rightSlack < leftSlack)
-      return rgt;
-    if ((*lft)->getAest() < (*rgt)->getAest())
-      return lft;
-    return rgt;
-  };
-
-  assert(!readyNodes.empty() && "expected at least one ready node");
-  auto validNode = readyNodes.begin();
-  auto bestNode = validNode;
-
-  while (validNode != readyNodes.end()) {
-    bestNode = findBestNode(validNode, bestNode);
-    std::advance(validNode, 1);
-  }
-  return *bestNode;
-}
-
 // Picks the best CPU + slot for `candidate`:
 //   * Phase 1 (parallel, read-only): call findSlot on every candidate CPU.
 //   * Phase 2 (sequential): process CPUs in ascending slot.aest order. For
@@ -940,7 +1185,7 @@ TaskDCP* GraphDCP::findCandidate(const std::vector<TaskDCP*>& readyNodes) {
 //     evaluate a slot for the smallest-slack child, then roll back.
 //   * Rescue (sequential): if nothing fit, grow the CPU count if allowed,
 //     otherwise pick the CPU that leads to the smallest DCPL increase.
-void GraphDCP::selectProcessor(TaskDCP* candidate, bool push) {
+GraphDCP::CandidateRelations GraphDCP::selectProcessor(TaskDCP* candidate, bool push) {
   CandidateRelations relations = dcp_graph::computeCandidateRelations(candidate);
   relations.descendantsTopoOrder.reserve(relations.descendants.size());
   for (auto it = candidate->getTopologicalIterator(); it != topologicalOrder.end(); ++it) {
@@ -960,22 +1205,43 @@ void GraphDCP::selectProcessor(TaskDCP* candidate, bool push) {
   const CrossbarUsage candidateFootprint = getTaskCrossbarFootprint(candidate);
   const bool candidateHasCrossbar = candidateFootprint != 0;
   const CrossbarUsage cpuCapacity = candidateHasCrossbar ? getCpuCrossbarCapacity() : 0;
+  DCP_DEBUG_IF(auto dedupStart = std::chrono::steady_clock::now();)
+  CpuAestCache cpuAests = computeCpuAestCache(candidate);
+  DCP_DEBUG_IF(const bool checkCpuAestCache = std::getenv("DCP_CHECK_CPU_AEST_CACHE") != nullptr;)
+  llvm::SmallDenseSet<uint64_t, 32> seenProcessorKeys;
+  seenProcessorKeys.reserve(static_cast<size_t>(topCpu + 1));
   for (CPU c = 0; c <= topCpu; c++) {
     if (candidateHasCrossbar && c != getLastCpu()) {
       CrossbarUsage nextUsage = checkedAdd(getCpuCrossbarUsage(c), candidateFootprint);
       if (nextUsage >= cpuCapacity)
         continue;
     }
+    Time candidateAest = cpuAests.get(c);
+    DCP_DEBUG_IF(if (checkCpuAestCache) {
+      Time recomputedAest = computeAestOnCpu(candidate, c);
+      if (candidateAest != recomputedAest) {
+        std::fprintf(stderr,
+                     "[DCP_CHECK_CPU_AEST_CACHE] mismatch candidate=%zu cpu=%d cached=%llu recomputed=%llu\n",
+                     getNodeIndex(candidate),
+                     c,
+                     static_cast<unsigned long long>(candidateAest),
+                     static_cast<unsigned long long>(recomputedAest));
+        llvm::report_fatal_error("DCP CPU AEST cache mismatch");
+      }
+    })
+    if (!seenProcessorKeys.insert(computeCpuCandidateKey(candidateAest, c)).second)
+      continue;
     processors.push_back(c);
   }
-
+  DCP_DEBUG_IF(gSelectTimers.dedup +=
+               std::chrono::duration<double>(std::chrono::steady_clock::now() - dedupStart).count();)
   if (processors.empty()) {
-    CPU bestCpu = canCreateNewCpu ? getLastCpu() : 0;
-    FindSlot bestSlot = {computeAestOnCpu(candidate, bestCpu), static_cast<int>(getOrCreateCpuTasks(bestCpu).size())};
-    if (canCreateNewCpu)
-      incrementLastCpu();
-    insertTaskInCPU(bestCpu, candidate, bestSlot.index);
-    return;
+    // processors.empty() implies !canCreateNewCpu: a fresh CPU always passes
+    // the crossbar filter and would have been added. Reaching here means every
+    // existing CPU is crossbar-exhausted and the task requires crossbar
+    // capacity — the placement is impossible.
+    llvm::report_fatal_error("DCP scheduler: crossbar capacity exhausted on all CPUs; "
+                             "cannot schedule task that requires crossbar allocation");
   }
 
   // Phase 1: parallel findSlot sweep (read-only over graph state).
@@ -1001,21 +1267,20 @@ void GraphDCP::selectProcessor(TaskDCP* candidate, bool push) {
     for (size_t i = 0; i < processors.size(); ++i)
       sweep(i);
   DCP_DEBUG_IF(gSelectTimers.findSlot +=
-                 std::chrono::duration<double>(std::chrono::steady_clock::now() - sweepStart).count();)
+               std::chrono::duration<double>(std::chrono::steady_clock::now() - sweepStart).count();)
 
 #ifdef DCP_DEBUG_ENABLED
   {
     static bool reported = false;
     if (!reported) {
       reported = true;
-      std::fprintf(stderr,
-                   "[dcp] selectProcessor parallel sweep: context=%p mt=%d procs=%zu pool=%u\n",
-                   (void*) context,
-                   context != nullptr ? (int) context->isMultithreadingEnabled() : -1,
-                   processors.size(),
-                   context != nullptr && context->isMultithreadingEnabled()
-                     ? context->getThreadPool().getMaxConcurrency()
-                     : 0u);
+      std::fprintf(
+        stderr,
+        "[dcp] selectProcessor parallel sweep: context=%p mt=%d procs=%zu pool=%u\n",
+        (void*) context,
+        context != nullptr ? (int) context->isMultithreadingEnabled() : -1,
+        processors.size(),
+        context != nullptr && context->isMultithreadingEnabled() ? context->getThreadPool().getMaxConcurrency() : 0u);
     }
   }
 #endif
@@ -1056,9 +1321,10 @@ void GraphDCP::selectProcessor(TaskDCP* candidate, bool push) {
       DCP_DEBUG_IF(auto t2 = std::chrono::steady_clock::now();)
       Weight candidateWeight = candidate->computeWeightOnCpu(this, currentCpu);
       Time candidateCompletion = addOrMax(slot.aest, candidateWeight);
-      bool skip = (!emptyCpu && candidateCompletion > currentDcpl)
-               || addOrMax(slot.aest, candidateCompletion) >= bestComposite;
-      DCP_DEBUG_IF(gSelectTimers.precheck += std::chrono::duration<double>(std::chrono::steady_clock::now() - t2).count();)
+      bool skip =
+        (!emptyCpu && candidateCompletion > currentDcpl) || addOrMax(slot.aest, candidateCompletion) >= bestComposite;
+      DCP_DEBUG_IF(gSelectTimers.precheck +=
+                   std::chrono::duration<double>(std::chrono::steady_clock::now() - t2).count();)
       if (skip)
         continue;
       DCP_DEBUG_IF(++gSelectTimers.passedPrecheck;)
@@ -1074,8 +1340,8 @@ void GraphDCP::selectProcessor(TaskDCP* candidate, bool push) {
         scheduleSnapshot = dcp_graph::captureLocalScheduleState(
           candidate, relations.descendants, dcpl, maxCompletion, secondMaxCompletion, maxCompletionTask);
         taskInsertion = insertTaskInCPU(currentCpu, candidate, slot.index);
-        bool withinBudget = tryUpdateAestWithinBudget(
-          candidate, llvm::ArrayRef<TaskDCP*>(relations.descendantsTopoOrder), currentDcpl);
+        bool withinBudget =
+          tryUpdateAestWithinBudget(candidate, llvm::ArrayRef<TaskDCP*>(relations.descendantsTopoOrder), currentDcpl);
         if (!withinBudget) {
           DCP_DEBUG_IF(auto t4 = std::chrono::steady_clock::now();)
           taskInsertion.rollBack();
@@ -1088,7 +1354,7 @@ void GraphDCP::selectProcessor(TaskDCP* candidate, bool push) {
         }
       }
       DCP_DEBUG_IF(gSelectTimers.snapshotInsertUpdate +=
-                     std::chrono::duration<double>(std::chrono::steady_clock::now() - t3).count();)
+                   std::chrono::duration<double>(std::chrono::steady_clock::now() - t3).count();)
       DCP_DEBUG_IF(++gSelectTimers.passedDcpl;)
 
       // Pick the tightest unscheduled child (smallest slack) and measure what
@@ -1136,7 +1402,7 @@ void GraphDCP::selectProcessor(TaskDCP* candidate, bool push) {
         dcp_graph::restoreLocalScheduleState(
           scheduleSnapshot, dcpl, maxCompletion, secondMaxCompletion, maxCompletionTask);
       DCP_DEBUG_IF(gSelectTimers.rollbackRestore +=
-                     std::chrono::duration<double>(std::chrono::steady_clock::now() - t6).count();)
+                   std::chrono::duration<double>(std::chrono::steady_clock::now() - t6).count();)
     }
   }
 
@@ -1151,7 +1417,9 @@ void GraphDCP::selectProcessor(TaskDCP* candidate, bool push) {
     else {
       Time bestDcpl = std::numeric_limits<Time>::max();
       Time currentDcpl = getDcpl();
-      for (CPU c = 0; c < getLastCpu(); c++) {
+      for (CPU c : processors) {
+        if (c == getLastCpu())
+          continue;
         auto slot = findSlot(candidate, c, false, relations);
         if (slot.aest == std::numeric_limits<Time>::max())
           slot = findSlot(candidate, c, true, relations);
@@ -1160,8 +1428,7 @@ void GraphDCP::selectProcessor(TaskDCP* candidate, bool push) {
         // Cheap lower bound: post-insertion DCPL is at least max(currentDcpl,
         // candidate completion on this slot). Skip CPUs already worse than
         // the best seen.
-        Time lowerBound =
-          std::max(currentDcpl, addOrMax(slot.aest, candidate->computeWeightOnCpu(this, c)));
+        Time lowerBound = std::max(currentDcpl, addOrMax(slot.aest, candidate->computeWeightOnCpu(this, c)));
         if (lowerBound >= bestDcpl)
           continue;
         auto snapshot = dcp_graph::captureLocalScheduleState(
@@ -1170,23 +1437,37 @@ void GraphDCP::selectProcessor(TaskDCP* candidate, bool push) {
         updateAestFromTaskWithDescendants(candidate, llvm::ArrayRef<TaskDCP*>(relations.descendantsTopoOrder));
         Time candidateDcpl = getDcpl();
         taskInsertion.rollBack();
-        dcp_graph::restoreLocalScheduleState(
-          snapshot, dcpl, maxCompletion, secondMaxCompletion, maxCompletionTask);
+        dcp_graph::restoreLocalScheduleState(snapshot, dcpl, maxCompletion, secondMaxCompletion, maxCompletionTask);
         if (candidateDcpl < bestDcpl) {
           bestDcpl = candidateDcpl;
           bestCpu = c;
           bestSlot = slot;
         }
       }
-      if (bestCpu == -1) {
-        bestCpu = 0;
-        bestSlot = {computeAestOnCpu(candidate, bestCpu), static_cast<int>(getOrCreateCpuTasks(bestCpu).size())};
-      }
+      if (bestCpu == -1)
+        llvm::report_fatal_error("DCP scheduler: no valid slot found for task on any eligible CPU — "
+                                 "all slots are blocked by already-placed descendants");
     }
   }
   if (bestCpu == getLastCpu() && getLastCpu() < maxCpuCount)
     incrementLastCpu();
   insertTaskInCPU(bestCpu, candidate, bestSlot.index);
+
+  // Incremental AEST/ALST refresh replacing the full initAest/initAlst that
+  // used to run after every placement. Post-insertion relations pick up any
+  // new scheduling-edge ancestors/descendants introduced by the insertion.
+  Time oldDcpl = getDcpl();
+  CandidateRelations postRelations = dcp_graph::computeCandidateRelations(candidate);
+  llvm::SmallVector<TaskDCP*, 32> postDescendantsTopoOrder;
+  postDescendantsTopoOrder.reserve(postRelations.descendants.size());
+  for (auto it = candidate->getTopologicalIterator(); it != topologicalOrder.end(); ++it) {
+    TaskDCP* current = &*it;
+    if (current != candidate && postRelations.descendants.contains(current))
+      postDescendantsTopoOrder.push_back(current);
+  }
+  updateAestFromTaskWithDescendants(candidate, llvm::ArrayRef<TaskDCP*>(postDescendantsTopoOrder));
+  updateAlstFromScheduledTask(candidate, postRelations, oldDcpl);
+  return postRelations;
 }
 
 //===----------------------------------------------------------------------===//
@@ -1195,61 +1476,102 @@ void GraphDCP::selectProcessor(TaskDCP* candidate, bool push) {
 
 void GraphDCP::runDcp() {
   initTopological();
+  initTaskStructureHashes();
   initAest();
   initAlst();
   dumpDot();
 
   dcp_graph::DcpProgressLogger progressLogger(nodes.size());
   llvm::DenseMap<TaskDCP*, int> unscheduledParents;
-  std::vector<TaskDCP*> readyNodes;
-  readyNodes.reserve(nodes.size());
+
+  // Min-heap over ready tasks: tightest slack first, earliest AEST as tiebreak.
+  // Lazy deletion: when AEST/ALST change after a placement, fresh entries are
+  // pushed for the affected tasks. Stale ones are detected on pop by comparing
+  // stored vs current (slack, aest) and re-pushed with the current values.
+  struct ReadyEntry {
+    Time slack;
+    Time aest;
+    int64_t orderKey;
+    TaskDCP* task;
+    bool operator>(const ReadyEntry& other) const {
+      if (slack != other.slack)
+        return slack > other.slack;
+      if (aest != other.aest)
+        return aest > other.aest;
+      return orderKey > other.orderKey;
+    }
+  };
+  std::priority_queue<ReadyEntry, std::vector<ReadyEntry>, std::greater<ReadyEntry>> readyQueue;
+  size_t readyCount = 0;
+
+  auto pushReady = [&](TaskDCP* node) {
+    readyQueue.push({slackOrZero(node->getAest(), node->getAlst()), node->getAest(), node->Id(), node});
+  };
+
   for (auto& node : nodes) {
     int dependencyParents = dcp_graph::countDependencyParents(&node);
     unscheduledParents[&node] = dependencyParents;
-    if (dependencyParents == 0)
-      readyNodes.push_back(&node);
+    if (dependencyParents == 0) {
+      pushReady(&node);
+      ++readyCount;
+    }
   }
-  progressLogger.printStart(readyNodes.size());
+  size_t xbarsCapacity = static_cast<size_t>(maxCpuCount) * onnx_mlir::crossbarCountInCore.getValue();
+  progressLogger.printStart(readyCount, maxCpuCount, xbarsCapacity);
 
-  while (!readyNodes.empty()) {
-    DCP_DEBUG_IF(auto findStart = std::chrono::steady_clock::now();)
-    TaskDCP* candidate = findCandidate(readyNodes);
-    DCP_DEBUG_IF(progressLogger.recordFindDuration(
-      std::chrono::duration<double>(std::chrono::steady_clock::now() - findStart).count());)
-    fastRemove(readyNodes, candidate);
+  while (readyCount > 0) {
+    // Pop with lazy deletion: skip stale entries and re-push with current values.
+    TaskDCP* candidate = nullptr;
+    while (!readyQueue.empty()) {
+      auto entry = readyQueue.top();
+      readyQueue.pop();
+      Time curSlack = slackOrZero(entry.task->getAest(), entry.task->getAlst());
+      Time curAest = entry.task->getAest();
+      if (entry.slack == curSlack && entry.aest == curAest) {
+        candidate = entry.task;
+        break;
+      }
+      readyQueue.push({curSlack, curAest, entry.orderKey, entry.task});
+    }
+    assert(candidate != nullptr && "readyCount > 0 but heap exhausted");
+    --readyCount;
 
     DCP_DEBUG_IF(auto selectStart = std::chrono::steady_clock::now();)
-    selectProcessor(candidate, candidate->isCriticalPath());
+    CandidateRelations postRelations = selectProcessor(candidate, candidate->isCriticalPath());
     DCP_DEBUG_IF(
       double selectSeconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - selectStart).count();
       progressLogger.recordSelectDuration(selectSeconds);
-      progressLogger.maybePrintSlowCandidate(getNodeIndex(candidate), selectSeconds, readyNodes.size(), getLastCpu());
-    )
+      progressLogger.maybePrintSlowCandidate(getNodeIndex(candidate), selectSeconds, readyCount, getLastCpu());)
+
+    // Proactively refresh the heap priority for ready nodes whose AEST or ALST
+    // changed: ancestors had their ALST individually recomputed; descendants had
+    // their AEST bumped. Both may now sort differently than their stale entries.
+    for (TaskDCP* node : postRelations.ancestors)
+      if (!node->isScheduled() && unscheduledParents[node] == 0)
+        pushReady(node);
+    for (TaskDCP* node : postRelations.descendants)
+      if (!node->isScheduled() && unscheduledParents[node] == 0)
+        pushReady(node);
 
-    DCP_DEBUG_IF(auto updateStart = std::chrono::steady_clock::now();)
-    initAest();
-    initAlst();
-    DCP_DEBUG_IF(progressLogger.recordUpdateDuration(
-      std::chrono::duration<double>(std::chrono::steady_clock::now() - updateStart).count());)
     progressLogger.advanceCompleted();
-    progressLogger.printProgress(readyNodes.size(), getLastCpu(), "recompute", false);
+    progressLogger.printProgress(readyCount, getLastCpu(), maxCpuCount, crossbarsUsed(), crossbarsAvailable(), false);
 
     for (const auto& childEdge : candidate->children) {
       if (childEdge.isScheduling || childEdge.first->isScheduled())
         continue;
       int& dependencyParents = unscheduledParents[childEdge.first];
       assert(dependencyParents > 0 && "dependency parent count must stay positive");
-      dependencyParents--;
-      if (dependencyParents == 0)
-        readyNodes.push_back(childEdge.first);
+      --dependencyParents;
+      if (dependencyParents == 0) {
+        pushReady(childEdge.first);
+        ++readyCount;
+      }
     }
-    DCP_DEBUG_IF(
-      ++gSelectTimers.tasksProcessed;
-      if (std::getenv("DCP_SELECT_PROFILE") && (gSelectTimers.tasksProcessed % 100 == 0))
-        gSelectTimers.dump("tick");
-    )
+    DCP_DEBUG_IF(++gSelectTimers.tasksProcessed;
+                 if (std::getenv("DCP_SELECT_PROFILE") && (gSelectTimers.tasksProcessed % 100 == 0))
+                   gSelectTimers.dump("tick");)
   }
-  progressLogger.printProgress(readyNodes.size(), getLastCpu(), "done", true);
+  progressLogger.printProgress(0, getLastCpu(), maxCpuCount, crossbarsUsed(), crossbarsAvailable(), true);
   dumpDot();
 }
 
@@ -1260,8 +1582,11 @@ DCPAnalysisResult GraphDCP::getResult() {
 
   auto dominanceOrder = dcp_graph::collectDominanceOrder(getRoots(), nodes.size());
   ret.dominanceOrderCompute.reserve(dominanceOrder.size());
-  for (auto elem : dominanceOrder)
-    ret.dominanceOrderCompute.push_back(elem->getSpatWeightedCompute());
+  for (auto elem : dominanceOrder) {
+    auto spatCompute = elem->getSpatCompute();
+    if (spatCompute)
+      ret.dominanceOrderCompute.push_back({spatCompute.getOperation(), 0});
+  }
 
   for (CPU cpu = 0; cpu < getLastCpu(); ++cpu) {
     const CpuTaskList* tasks = findCpuTasks(cpu);
@@ -1269,10 +1594,14 @@ DCPAnalysisResult GraphDCP::getResult() {
       continue;
     size_t i = 0;
     for (auto node : *tasks) {
-      ret.computeToCpuMap[node->getSpatWeightedCompute()] = cpu;
+      auto spatCompute = node->getSpatCompute();
+      if (!spatCompute)
+        continue;
+      ComputeInstance instance {spatCompute.getOperation(), 0};
+      ret.computeToCpuMap[instance] = cpu;
       if (i++ == tasks->size() - 1) {
-        ret.isLastComputeOfCpu.insert(node->getSpatWeightedCompute());
-        ret.cpuToLastComputeMap[cpu] = node->getSpatWeightedCompute();
+        ret.isLastComputeOfCpu.insert(instance);
+        ret.cpuToLastComputeMap[cpu] = instance;
       }
     }
   }

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/Graph.hpp

@@ -4,6 +4,7 @@
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 
+#include <cstdint>
 #include <list>
 #include <optional>
 #include <unordered_map>
@@ -48,8 +49,10 @@ private:
 
   std::vector<TaskDCP> nodes;
   onnx_mlir::LabeledList<TaskDCP> topologicalOrder;
+  std::vector<uint64_t> taskStructureHashes;
   std::vector<CpuTaskList> cpuTasks;
   std::unordered_map<CPU, CrossbarUsage> cpuCrossbarUsage;
+  llvm::DenseMap<CPU, uint64_t> cpuStructureHashes;
   CPU lastCpu = 0;
   long long flag = 1;
   Time dcpl = 0;
@@ -70,6 +73,7 @@ private:
 
   void initAest();
   void initAlst();
+  void initTaskStructureHashes();
 
   Time computeAestOnCpu(TaskDCP* task, CPU cpu);
   Time computeDcplOnCpu(TaskDCP* task, CPU cpu);
@@ -83,9 +87,15 @@ private:
   // `dcplBudget`, signalling that the new DCPL would exceed the budget.
   // Returns true iff the full propagation completed without exceeding the
   // budget. Uses the caller's snapshot to restore AEST on the aborted tail.
-  bool tryUpdateAestWithinBudget(TaskDCP* task,
-                                 llvm::ArrayRef<TaskDCP*> descendantsTopoOrder,
-                                 Time dcplBudget);
+  bool tryUpdateAestWithinBudget(TaskDCP* task, llvm::ArrayRef<TaskDCP*> descendantsTopoOrder, Time dcplBudget);
+
+  // Incrementally refreshes ALST after `task` has been scheduled. Nodes
+  // outside the backward cone (`relations.ancestors` plus `task`) retain
+  // their relative distance to the sink boundary and only absorb the signed
+  // DCPL delta (`newDcpl - oldDcpl`). `task` itself and its ancestors are
+  // recomputed in reverse topological order so that new same-CPU transfer
+  // costs (now zero) and scheduling-edge children are reflected.
+  void updateAlstFromScheduledTask(TaskDCP* task, const CandidateRelations& relations, Time oldDcpl);
 
   void initTopological();
   void topologicalMoveAfter(TaskDCP* task, TaskDCP* pivotPoint, TaskInsertion* insertion = nullptr);
@@ -94,8 +104,11 @@ private:
   llvm::DenseMap<TaskDCP*, Time> computeAlst(TaskDCP* task, CPU cpu, const CandidateRelations& relations);
   size_t getNodeIndex(const TaskDCP* task) const;
 
-  TaskDCP* findCandidate(const std::vector<TaskDCP*>& readyNodes);
-  void selectProcessor(TaskDCP* candidate, bool push);
+  // Returns a compact dedup key for CPU `c` when evaluating `candidate`:
+  // mixes candidateAest, crossbar usage, and the incremental cpu structure
+  // hash into a single uint64_t. Zero heap allocation.
+  uint64_t computeCpuCandidateKey(Time candidateAest, CPU cpu);
+  CandidateRelations selectProcessor(TaskDCP* candidate, bool push);
   CPU getLastCpu() const { return lastCpu; }
   void incrementLastCpu() { lastCpu++; }
   FindSlot findSlot(TaskDCP* candidate, CPU cpu, bool push, const CandidateRelations& relations);
@@ -115,24 +128,28 @@ private:
 
 public:
   void runDcp();
-  GraphDCP(llvm::ArrayRef<onnx_mlir::spatial::SpatWeightedCompute> spatWeightedComputes,
-           llvm::ArrayRef<IndexedEdge> edges)
+  GraphDCP(llvm::ArrayRef<onnx_mlir::spatial::SpatCompute> spatComputes, llvm::ArrayRef<IndexedEdge> edges)
   : nodes(), cpuTasks(), cpuCrossbarUsage() {
-    for (auto spatWeightedCompute : spatWeightedComputes)
-      nodes.emplace_back(spatWeightedCompute);
+    for (auto spatCompute : spatComputes)
+      nodes.emplace_back(spatCompute);
     for (auto [start, end, weight] : edges)
       makeEdge(start, end, weight);
   }
 
   GraphDCP(llvm::ArrayRef<Weight> nodeWeights,
            llvm::ArrayRef<IndexedEdge> edges,
+           llvm::ArrayRef<int64_t> nodeOrderKeys = {},
            llvm::ArrayRef<CrossbarUsage> nodeCrossbarUsage = {})
   : nodes(), cpuTasks(), cpuCrossbarUsage() {
     assert((nodeCrossbarUsage.empty() || nodeCrossbarUsage.size() == nodeWeights.size())
            && "synthetic crossbar usage must match synthetic node weights");
+    assert((nodeOrderKeys.empty() || nodeOrderKeys.size() == nodeWeights.size())
+           && "synthetic node order keys must match synthetic node weights");
     nodes.reserve(nodeWeights.size());
     for (auto [index, weight] : llvm::enumerate(nodeWeights))
-      nodes.emplace_back(index, weight, nodeCrossbarUsage.empty() ? 0 : nodeCrossbarUsage[index]);
+      nodes.emplace_back(nodeOrderKeys.empty() ? static_cast<int64_t>(index) : nodeOrderKeys[index],
+                         weight,
+                         nodeCrossbarUsage.empty() ? 0 : nodeCrossbarUsage[index]);
     for (auto [start, end, weight] : edges)
       makeEdge(start, end, weight);
   }
@@ -150,6 +167,11 @@ public:
   void setMaxCpuCount(int value) { maxCpuCount = value; }
   int getMaxCpuCount() const { return maxCpuCount; }
 
+  // Total crossbar units allocated across all active CPUs.
+  size_t crossbarsUsed() const;
+  // Maximum crossbar units available across all active CPUs (lastCpu * per-CPU capacity).
+  size_t crossbarsAvailable() const;
+
   // Optional MLIR context used to drive mlir::parallelFor inside runDcp. If
   // null the scheduler runs single-threaded (tests use this path).
   void setContext(mlir::MLIRContext* ctx) { context = ctx; }

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/GraphDebug.cpp

@@ -35,10 +35,12 @@ void DcpProgressLogger::recordSelectDuration(double seconds) { selectProcessorSe
 void DcpProgressLogger::recordUpdateDuration(double seconds) { updateTimingSeconds += seconds; }
 void DcpProgressLogger::advanceCompleted(size_t taskCount) { completedTasks += taskCount; }
 
-void DcpProgressLogger::printStart(size_t readyCount) const {
+void DcpProgressLogger::printStart(size_t readyCount, int maxCpuCount, size_t xbarsCapacity) const {
   if (!logProgress)
     return;
-  llvm::errs() << llvm::formatv("[DCP] start: tasks={0} ready={1}\n", totalTasks, readyCount);
+  llvm::errs() << llvm::formatv(
+    "[DCP] start   tasks={0}   ready={1}   cpus=0/{2}   crossbars=0/{3}\n",
+    totalTasks, readyCount, maxCpuCount, xbarsCapacity);
 }
 
 void DcpProgressLogger::maybePrintSlowCandidate(size_t nodeIndex,
@@ -48,14 +50,15 @@ void DcpProgressLogger::maybePrintSlowCandidate(size_t nodeIndex,
   if (!logProgress || elapsedSeconds < 1.0)
     return;
 
-  llvm::errs() << llvm::formatv("[DCP] slow candidate node={0} elapsed={1} ready={2} cpus={3}\n",
+  llvm::errs() << llvm::formatv("[DCP] slow node={0}   elapsed={1}   ready={2}   cpus={3}\n",
                                 nodeIndex,
                                 formatDuration(elapsedSeconds),
                                 readyCount,
                                 cpuCount);
 }
 
-void DcpProgressLogger::printProgress(size_t readyCount, CPU cpuCount, llvm::StringRef stage, bool force) {
+void DcpProgressLogger::printProgress(
+  size_t readyCount, CPU cpuCount, int maxCpuCount, size_t xbarsUsed, size_t xbarsAvailable, bool force) {
   if (!logProgress)
     return;
 
@@ -68,19 +71,19 @@ void DcpProgressLogger::printProgress(size_t readyCount, CPU cpuCount, llvm::Str
   double etaSeconds = rate > 0.0 ? static_cast<double>(totalTasks - completedTasks) / rate : 0.0;
   double percent = totalTasks == 0 ? 100.0 : (100.0 * static_cast<double>(completedTasks) / totalTasks);
 
-  llvm::errs() << llvm::formatv("[DCP] {0}/{1} ({2:F1}%) ready={3} cpus={4} stage={5} elapsed={6} eta={7}\n",
-                                completedTasks,
-                                totalTasks,
-                                percent,
-                                readyCount,
-                                cpuCount,
-                                stage,
-                                formatDuration(elapsedSeconds),
-                                completedTasks == totalTasks ? "0:00" : formatDuration(etaSeconds));
-  llvm::errs() << llvm::formatv("        time(find={0}, select={1}, update={2})\n",
-                                formatDuration(findCandidateSeconds),
-                                formatDuration(selectProcessorSeconds),
-                                formatDuration(updateTimingSeconds));
+  bool done = completedTasks == totalTasks;
+  llvm::errs() << llvm::formatv(
+    "[DCP] {0}/{1} ({2:F0}%)   ready={3}   cpus={4}/{5}   crossbars={6}/{7}   {8}{9}\n",
+    completedTasks,
+    totalTasks,
+    percent,
+    readyCount,
+    cpuCount,
+    maxCpuCount,
+    xbarsUsed,
+    xbarsAvailable,
+    llvm::formatv("elapsed={0}", formatDuration(elapsedSeconds)).str(),
+    done ? "" : llvm::formatv("   eta={0}", formatDuration(etaSeconds)).str());
   lastProgressPrint = now;
 }
 
@@ -91,9 +94,9 @@ void DcpProgressLogger::recordFindDuration(double) {}
 void DcpProgressLogger::recordSelectDuration(double) {}
 void DcpProgressLogger::recordUpdateDuration(double) {}
 void DcpProgressLogger::advanceCompleted(size_t) {}
-void DcpProgressLogger::printStart(size_t) const {}
+void DcpProgressLogger::printStart(size_t, int, size_t) const {}
 void DcpProgressLogger::maybePrintSlowCandidate(size_t, double, size_t, CPU) const {}
-void DcpProgressLogger::printProgress(size_t, CPU, llvm::StringRef, bool) {}
+void DcpProgressLogger::printProgress(size_t, CPU, int, size_t, size_t, bool) {}
 
 #endif
 

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/GraphDebug.hpp

@@ -31,9 +31,10 @@ public:
   void recordUpdateDuration(double seconds);
   void advanceCompleted(size_t taskCount = 1);
 
-  void printStart(size_t readyCount) const;
+  void printStart(size_t readyCount, int maxCpuCount, size_t xbarsCapacity) const;
   void maybePrintSlowCandidate(size_t nodeIndex, double elapsedSeconds, size_t readyCount, CPU cpuCount) const;
-  void printProgress(size_t readyCount, CPU cpuCount, llvm::StringRef stage, bool force);
+  void printProgress(size_t readyCount, CPU cpuCount, int maxCpuCount,
+                     size_t xbarsUsed, size_t xbarsAvailable, bool force);
 
 #ifdef DCP_DEBUG_ENABLED
 private:

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/Task.hpp

@@ -8,7 +8,7 @@
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 class TaskDCP : public onnx_mlir::LabeledListNode<TaskDCP> {
-  onnx_mlir::spatial::SpatWeightedCompute spatWeightedCompute;
+  onnx_mlir::spatial::SpatCompute spatCompute;
   Time aest;
   Time alst;
   std::optional<CPU> scheduledCpu;
@@ -38,22 +38,22 @@ public:
   std::vector<Edge> parents;
   std::vector<Edge> children;
   TaskDCP() = default;
-  TaskDCP(onnx_mlir::spatial::SpatWeightedCompute spatWeightedCompute)
+  TaskDCP(onnx_mlir::spatial::SpatCompute spatCompute)
   : onnx_mlir::LabeledListNode<TaskDCP>(),
-    spatWeightedCompute(spatWeightedCompute),
+    spatCompute(spatCompute),
     aest(0),
     alst(0),
     scheduledCpu(),
-    weight(getSpatComputeWeight(spatWeightedCompute)),
+    weight(getSpatComputeWeight(spatCompute)),
     baseWeight(weight),
-    crossbarUsage(getSpatComputeCrossbarUsage(spatWeightedCompute)),
+    crossbarUsage(getSpatComputeCrossbarUsage(spatCompute)),
     syntheticId(-1),
     parents(),
     children() {}
 
   TaskDCP(int64_t id, Weight weight, CrossbarUsage crossbarUsage = 0)
   : onnx_mlir::LabeledListNode<TaskDCP>(),
-    spatWeightedCompute(),
+    spatCompute(),
     aest(0),
     alst(0),
     scheduledCpu(),
@@ -90,14 +90,14 @@ public:
   void setAlst(Time value) { alst = value; }
   bool hasDescendant(TaskDCP* child);
   int64_t Id() const {
-    if (spatWeightedCompute)
-      return reinterpret_cast<int64_t>(spatWeightedCompute.getAsOpaquePointer());
+    if (spatCompute)
+      return reinterpret_cast<int64_t>(spatCompute.getAsOpaquePointer());
     return syntheticId;
   }
 
   bool isCriticalPath() const { return alst == aest; }
   bool isScheduled() const { return scheduledCpu.has_value(); }
-  onnx_mlir::spatial::SpatWeightedCompute getSpatWeightedCompute() const { return spatWeightedCompute; }
+  onnx_mlir::spatial::SpatCompute getSpatCompute() const { return spatCompute; }
 
   void setFlag(long long val) { flag = val; }
   long long getFlag() const { return flag; }

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/Utils.hpp

@@ -92,18 +92,18 @@ inline T subtractOrZero(T lhs, T rhs) {
 
 inline Time slackOrZero(Time earliestStart, Time latestStart) { return subtractOrZero(latestStart, earliestStart); }
 
-inline Weight getSpatComputeWeight(onnx_mlir::spatial::SpatWeightedCompute spatWeightedCompute) {
+inline Weight getSpatComputeWeight(onnx_mlir::spatial::SpatCompute spatCompute) {
   constexpr Weight kOperationWeight = 100;
   Weight numOperations = 0;
-  for (auto& block : spatWeightedCompute.getBody())
+  for (auto& block : spatCompute.getBody())
     for ([[maybe_unused]] auto& op : block)
       numOperations = checkedAdd(numOperations, static_cast<Weight>(1));
   return checkedMultiply(numOperations, kOperationWeight);
 }
 
-inline CrossbarUsage getSpatComputeCrossbarUsage(onnx_mlir::spatial::SpatWeightedCompute spatWeightedCompute) {
+inline CrossbarUsage getSpatComputeCrossbarUsage(onnx_mlir::spatial::SpatCompute spatCompute) {
   CrossbarUsage crossbarUsage = 0;
-  for (auto& region : spatWeightedCompute.getBody())
+  for (auto& region : spatCompute.getBody())
     for (auto& inst : region)
       if (llvm::isa<onnx_mlir::spatial::SpatWeightedVMMOp>(inst))
         crossbarUsage = checkedAdd(crossbarUsage, static_cast<CrossbarUsage>(1));

src/PIM/Pass/CountInstructionPass.cpp

@@ -31,7 +31,7 @@ struct CountInstructionPass : public PassWrapper<CountInstructionPass, Operation
     unsigned totalInstructionCount = 0;
 
     unsigned computeId = 0;
-    for (auto computeOp : func.getOps<spatial::SpatWeightedCompute>()) {
+    for (auto computeOp : func.getOps<spatial::SpatCompute>()) {
       unsigned instructionCount = 0;
       instructionCount += computeOp.getBody().front().getOperations().size();
       llvm::outs() << "Compute " << computeId << ": " << instructionCount << " instructions\n";

src/PIM/Pass/PimCodegen/HostConstantFolding/Patterns/Constant.cpp

@@ -116,10 +116,9 @@ struct FoldConstantCoreMapPattern final : OpRewritePattern<linalg::MapOp> {
     auto globalOp = createFoldedGlobal(moduleOp, mapOp.getLoc(), initType, splatAttr, "pim_core_fill");
 
     OpBuilder::InsertionGuard guard(rewriter);
-    rewriter.setInsertionPoint(coreOp);
-    auto getGlobalOp = memref::GetGlobalOp::create(rewriter, mapOp.getLoc(), initType, globalOp.getName());
 
     rewriter.setInsertionPoint(mapOp);
+    auto getGlobalOp = memref::GetGlobalOp::create(rewriter, mapOp.getLoc(), initType, globalOp.getName());
     auto sizeInBytes = initType.getNumElements() * initType.getElementTypeBitWidth() / 8;
     pim::PimMemCopyOp::create(rewriter,
                               mapOp.getLoc(),
@@ -258,9 +257,18 @@ struct FoldConstantTransposePattern final : OpRewritePattern<pim::PimTransposeOp
     if (!resultType || !resultType.hasStaticShape())
       return failure();
 
+    // Look through an optional pim.memcp_hd to find the source get_global.
+    // This occurs when the constant was staged into device memory before transposing.
+    pim::PimMemCopyHostToDevOp memcpHd;
     auto sourceGetGlobal = transposeOp.getInput().getDefiningOp<memref::GetGlobalOp>();
-    if (!sourceGetGlobal)
-      return failure();
+    if (!sourceGetGlobal) {
+      memcpHd = transposeOp.getInput().getDefiningOp<pim::PimMemCopyHostToDevOp>();
+      if (!memcpHd)
+        return failure();
+      sourceGetGlobal = memcpHd.getHostSource().getDefiningOp<memref::GetGlobalOp>();
+      if (!sourceGetGlobal)
+        return failure();
+    }
 
     auto moduleOp = transposeOp->getParentOfType<ModuleOp>();
     if (!moduleOp)
@@ -298,13 +306,26 @@ struct FoldConstantTransposePattern final : OpRewritePattern<pim::PimTransposeOp
 
     bool isAlwaysWeight =
       !transposeOp->getUsers().empty()
-      && llvm::all_of(transposeOp->getUsers(), [](Operation* user) { return isa<pim::PimCoreOp>(user); });
+      && llvm::all_of(transposeOp->getUsers(), [](Operation* user) {
+           return isa<pim::PimCoreOp, pim::PimCoreBatchOp>(user);
+         });
     if (isAlwaysWeight) {
       markWeightAlways(newGlobal);
       markWeightAlways(newGetGlobal);
     }
 
+    auto outputAllocOp = transposeOp.getOutputBuffer().getDefiningOp<memref::AllocOp>();
     rewriter.replaceOp(transposeOp, newGetGlobal.getResult());
+
+    if (memcpHd && memcpHd.use_empty()) {
+      auto deviceAllocOp = memcpHd.getDeviceTarget().getDefiningOp<memref::AllocOp>();
+      rewriter.eraseOp(memcpHd);
+      if (deviceAllocOp && deviceAllocOp->use_empty())
+        rewriter.eraseOp(deviceAllocOp);
+    }
+    if (outputAllocOp && outputAllocOp->use_empty())
+      rewriter.eraseOp(outputAllocOp);
+
     return success();
   }
 };
@@ -341,18 +362,25 @@ struct FoldConstantAllocPattern final : OpRewritePattern<memref::AllocOp> {
         continue;
       }
 
-      if (!isa<pim::PimCoreOp>(user))
+      if (!isa<pim::PimCoreOp, pim::PimCoreBatchOp>(user))
         return failure();
     }
 
     if (!llvm::all_of(castsToReplace, [](memref::CastOp castOp) {
-          return llvm::all_of(castOp->getUsers(), [](Operation* user) { return isa<pim::PimCoreOp>(user); });
+          return llvm::all_of(castOp->getUsers(), [](Operation* user) {
+            return isa<pim::PimCoreOp, pim::PimCoreBatchOp>(user);
+          });
         })) {
       allLiveUsersAreCoreOps = false;
     }
 
     if (!llvm::all_of(allocOp->getUsers(), [](Operation* user) {
-          return isa<linalg::MapOp, memref::SubViewOp, memref::DeallocOp, memref::CastOp, pim::PimCoreOp>(user);
+          return isa<linalg::MapOp,
+                     memref::SubViewOp,
+                     memref::DeallocOp,
+                     memref::CastOp,
+                     pim::PimCoreOp,
+                     pim::PimCoreBatchOp>(user);
         })) {
       return failure();
     }
@@ -389,6 +417,83 @@ struct FoldConstantAllocPattern final : OpRewritePattern<memref::AllocOp> {
   }
 };
 
+struct FoldConstantHostCopyPattern final : OpRewritePattern<memref::CopyOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(memref::CopyOp copyOp, PatternRewriter& rewriter) const override {
+    if (copyOp->getParentOfType<pim::PimCoreOp>())
+      return failure();
+
+    auto allocOp = copyOp.getTarget().getDefiningOp<memref::AllocOp>();
+    if (!allocOp)
+      return failure();
+    auto allocType = dyn_cast<MemRefType>(allocOp.getType());
+    if (!allocType || !allocType.hasStaticShape())
+      return failure();
+
+    auto srcSubview = getStaticSubviewInfo(copyOp.getSource());
+    Value globalSource = succeeded(srcSubview) ? srcSubview->source : stripMemRefCasts(copyOp.getSource());
+
+    auto moduleOp = copyOp->getParentOfType<ModuleOp>();
+    if (!moduleOp)
+      return failure();
+
+    auto denseAttr = getDenseGlobalValue(moduleOp, globalSource);
+    if (failed(denseAttr))
+      return failure();
+
+    DenseElementsAttr foldedAttr;
+    if (succeeded(srcSubview)) {
+      if (llvm::any_of(srcSubview->strides, [](int64_t stride) { return stride != 1; }))
+        return failure();
+      auto staticOffsets = getStaticSubviewOffsets(*srcSubview);
+      if (failed(staticOffsets))
+        return failure();
+
+      auto maybeFoldedAttr = foldDenseSubview(*denseAttr, *staticOffsets, allocType.getShape());
+      if (failed(maybeFoldedAttr))
+        return failure();
+      foldedAttr = *maybeFoldedAttr;
+    }
+    else {
+      auto resultTensorType = RankedTensorType::get(allocType.getShape(), allocType.getElementType());
+      if (resultTensorType != denseAttr->getType())
+        return failure();
+      foldedAttr = *denseAttr;
+    }
+
+    bool allLiveUsersAreCores = true;
+    for (Operation* user : allocOp->getUsers()) {
+      if (user == copyOp)
+        continue;
+      if (isa<memref::DeallocOp>(user))
+        continue;
+      if (isa<pim::PimCoreOp, pim::PimCoreBatchOp>(user))
+        continue;
+      if (isa<memref::SubViewOp>(user)) {
+        allLiveUsersAreCores = false;
+        continue;
+      }
+      return failure();
+    }
+
+    auto newGlobal = createFoldedGlobal(moduleOp, allocOp.getLoc(), allocType, foldedAttr, "pim_folded_host_copy");
+    if (allLiveUsersAreCores)
+      markWeightAlways(newGlobal);
+
+    rewriter.setInsertionPoint(allocOp);
+    auto newGetGlobal = memref::GetGlobalOp::create(rewriter, allocOp.getLoc(), allocType, newGlobal.getName());
+    if (allLiveUsersAreCores)
+      markWeightAlways(newGetGlobal);
+
+    rewriter.replaceAllUsesWith(allocOp.getResult(), newGetGlobal.getResult());
+    rewriter.eraseOp(copyOp);
+    if (allocOp.use_empty())
+      rewriter.eraseOp(allocOp);
+    return success();
+  }
+};
+
 struct FoldConstantMemCpPattern final : OpRewritePattern<pim::PimMemCopyOp> {
   using OpRewritePattern::OpRewritePattern;
 
@@ -443,7 +548,7 @@ struct FoldConstantMemCpPattern final : OpRewritePattern<pim::PimMemCopyOp> {
         continue;
       if (isa<memref::DeallocOp>(user))
         continue;
-      if (isa<pim::PimCoreOp>(user))
+      if (isa<pim::PimCoreOp, pim::PimCoreBatchOp>(user))
         continue;
       if (isa<memref::SubViewOp>(user)) {
         allLiveUsersAreCores = false;
@@ -473,7 +578,11 @@ struct FoldConstantMemCpPattern final : OpRewritePattern<pim::PimMemCopyOp> {
 
 void populateConstantFoldingConstantPatterns(RewritePatternSet& patterns) {
   patterns
-    .add<FoldConstantTransposePattern, FoldConstantAllocPattern, FoldConstantCoreMapPattern, FoldConstantMemCpPattern>(
+    .add<FoldConstantTransposePattern,
+         FoldConstantAllocPattern,
+         FoldConstantCoreMapPattern,
+         FoldConstantHostCopyPattern,
+         FoldConstantMemCpPattern>(
       patterns.getContext());
 }
 

src/PIM/Pass/PimCodegen/VerificationPass.cpp

@@ -24,7 +24,26 @@ static bool isAddressOnlyHostOp(Operation* op) {
              memref::CastOp,
              memref::CollapseShapeOp,
              memref::ExpandShapeOp,
-             spatial::SpatChannelNewOp>(op);
+             memref::CopyOp>(op);
+}
+
+// Looser than isCodegenAddressableValue: follows view ops without requiring contiguity.
+// Used for memref.copy operands which may be non-contiguous subviews.
+static bool isBaseAddressableValue(Value value) {
+  while (true) {
+    if (isa<BlockArgument>(value))
+      return true;
+    Operation* defOp = value.getDefiningOp();
+    if (!defOp)
+      return false;
+    if (isa<memref::AllocOp, memref::GetGlobalOp>(defOp))
+      return true;
+    if (auto subview = dyn_cast<memref::SubViewOp>(defOp)) { value = subview.getSource(); continue; }
+    if (auto cast = dyn_cast<memref::CastOp>(defOp)) { value = cast.getSource(); continue; }
+    if (auto collapse = dyn_cast<memref::CollapseShapeOp>(defOp)) { value = collapse.getSrc(); continue; }
+    if (auto expand = dyn_cast<memref::ExpandShapeOp>(defOp)) { value = expand.getSrc(); continue; }
+    return false;
+  }
 }
 
 static bool isCodegenAddressableValue(Value value) {
@@ -38,6 +57,8 @@ static bool isCodegenAddressableValue(Value value) {
 static bool isExplicitHostOperand(Operation* op, unsigned operandIndex) {
   if (isa<pim::PimMemCopyHostToDevOp>(op))
     return operandIndex == 1;
+  if (isa<pim::PimMemCopyHostToDevBatchOp>(op))
+    return operandIndex == 1;
   if (isa<pim::PimMemCopyDevToHostOp>(op))
     return operandIndex == 0;
   return false;
@@ -69,6 +90,12 @@ struct VerificationPass : PassWrapper<VerificationPass, OperationPass<ModuleOp>>
           continue;
         }
 
+        if (auto coreBatchOp = dyn_cast<pim::PimCoreBatchOp>(&op)) {
+          if (failed(verifyCoreWeights(moduleOp, coreBatchOp)) || failed(verifyCoreOperands(coreBatchOp)))
+            hasFailure = true;
+          continue;
+        }
+
         if (auto returnOp = dyn_cast<func::ReturnOp>(&op)) {
           if (failed(verifyReturnOp(returnOp)))
             hasFailure = true;
@@ -92,10 +119,11 @@ struct VerificationPass : PassWrapper<VerificationPass, OperationPass<ModuleOp>>
   }
 
 private:
-  static LogicalResult verifyCoreWeights(ModuleOp moduleOp, pim::PimCoreOp coreOp) {
+  template <typename CoreOpTy>
+  static LogicalResult verifyCoreWeights(ModuleOp moduleOp, CoreOpTy coreOp) {
     bool hasFailure = false;
     for (auto [weightIndex, weight] : llvm::enumerate(coreOp.getWeights())) {
-      auto getGlobalOp = weight.getDefiningOp<memref::GetGlobalOp>();
+      auto getGlobalOp = weight.template getDefiningOp<memref::GetGlobalOp>();
       if (!getGlobalOp) {
         coreOp.emitOpError() << "weight #" << weightIndex
                              << " must be materialized as memref.get_global before JSON codegen";
@@ -131,7 +159,8 @@ private:
     return success(!hasFailure);
   }
 
-  static LogicalResult verifyCoreOperands(pim::PimCoreOp coreOp) {
+  template <typename CoreOpTy>
+  static LogicalResult verifyCoreOperands(CoreOpTy coreOp) {
     return walkPimCoreBlock(
       coreOp.getBody().front(), StaticValueKnowledge {}, [](Operation& op, const StaticValueKnowledge& knowledge) {
         bool hasFailure = false;
@@ -174,6 +203,13 @@ private:
       return verifyAddressOnlySource(op, collapseOp.getSrc());
     if (auto expandOp = dyn_cast<memref::ExpandShapeOp>(op))
       return verifyAddressOnlySource(op, expandOp.getSrc());
+    if (auto copyOp = dyn_cast<memref::CopyOp>(op)) {
+      if (!isBaseAddressableValue(copyOp.getSource()) || !isBaseAddressableValue(copyOp.getTarget())) {
+        op->emitOpError("depends on a value that is not backed by addressable storage");
+        return failure();
+      }
+      return success();
+    }
     return success();
   }
 

test/PIM/CMakeLists.txt

@@ -1,5 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
-
 add_custom_target(pim-unittest)
 set_target_properties(pim-unittest PROPERTIES FOLDER "Tests")
 

test/PIM/DCPTest.cpp

@@ -457,6 +457,10 @@ int testDCPGraphDiamondDependencies() {
   return 0;
 }
 
+// crossbarSize=4, crossbarCount=2 => capacity = 4*4*2 = 32.
+// Each task with crossbarUsage=1 needs footprint = 4*4 = 16, so at most 1 task
+// can fit per CPU (16+16 = 32 >= capacity). The scheduler must open a fresh CPU
+// for each task; all three end up on separate CPUs with their base weight.
 int testDCPGraphCrossbarExhaustion() {
   std::cout << "testDCPGraphCrossbarExhaustion:" << std::endl;
   configureDcpDotOutput();
@@ -474,36 +478,35 @@ int testDCPGraphCrossbarExhaustion() {
   const std::vector<Weight> nodeWeights = {10, 10, 10};
   const std::vector<CrossbarUsage> nodeCrossbarUsage = {1, 1, 1};
   GraphDCP graph(nodeWeights, {}, nodeCrossbarUsage);
-  graph.setMaxCpuCount(1);
+  graph.setMaxCpuCount(3);
   graph.runDcp();
 
-  if (graph.cpuCount() != 1) {
+  if (graph.cpuCount() != 3) {
     restoreCrossbarOptions();
-    std::cerr << "Expected exactly 1 CPU with maxCpuCount=1, got " << graph.cpuCount() << "\n";
+    std::cerr << "Expected 3 CPUs (one per task due to crossbar limit), got " << graph.cpuCount() << "\n";
     dumpDcpFailureArtifacts();
     return 1;
   }
 
-  auto scheduledTasks = graph.getScheduledTasks(0);
-  if (scheduledTasks.size() != 3) {
-    restoreCrossbarOptions();
-    std::cerr << "Expected all three tasks to be scheduled on CPU 0\n";
-    printCpuSchedule(graph, 0);
-    dumpDcpFailureArtifacts();
-    return 1;
-  }
-
-  if (scheduledTasks[0].weight != 10 || scheduledTasks[1].weight != std::numeric_limits<Weight>::max()
-      || scheduledTasks[2].weight != std::numeric_limits<Weight>::max()) {
-    restoreCrossbarOptions();
-    std::cerr << "Unexpected effective weights under crossbar exhaustion\n";
-    printCpuSchedule(graph, 0);
-    dumpDcpFailureArtifacts();
-    return 1;
+  int failures = 0;
+  for (CPU c = 0; c < 3; c++) {
+    auto scheduledTasks = graph.getScheduledTasks(c);
+    if (scheduledTasks.size() != 1) {
+      std::cerr << "Expected exactly 1 task on CPU " << c << ", got " << scheduledTasks.size() << "\n";
+      printCpuSchedule(graph, c);
+      failures++;
+      continue;
+    }
+    if (scheduledTasks[0].weight != 10) {
+      std::cerr << "Expected weight=10 on CPU " << c << ", got " << scheduledTasks[0].weight << "\n";
+      printCpuSchedule(graph, c);
+      failures++;
+    }
   }
 
   restoreCrossbarOptions();
-  return 0;
+  if (failures) dumpDcpFailureArtifacts();
+  return failures;
 }
 
 } // namespace

validation/validate_one.py

@@ -26,6 +26,10 @@ STAGE_COUNT = len(STAGE_TITLES)
 GENERATED_DIR_NAMES = ("inputs", "outputs", "raptor", "runner", "simulation")
 
 
+def sanitize_output_name(name):
+    return "".join(ch if ch.isalnum() or ch in "_.-" else "_" for ch in name[:255])
+
+
 @dataclass
 class ValidationResult:
     passed: bool
@@ -33,7 +37,7 @@ class ValidationResult:
 
 
 class ProgressReporter:
-    def __init__(self, total_models, stages_per_model=STAGE_COUNT):
+    def __init__(self, total_models, stages_per_model=STAGE_COUNT, enabled=None):
         self.total_models = total_models
         self.stages_per_model = stages_per_model
         self.total_steps = max(1, total_models * stages_per_model)
@@ -41,7 +45,7 @@ class ProgressReporter:
         self.passed_models = 0
         self.failed_models = 0
         self.current_label = ""
-        self.enabled = True
+        self.enabled = sys.stdout.isatty() if enabled is None else enabled
         self.columns = shutil.get_terminal_size((100, 20)).columns
         self.suspended = False
 
@@ -205,7 +209,7 @@ def build_dump_ranges(config_path, outputs_descriptor):
 
 def run_pim_simulator(simulator_dir, pim_dir, output_bin_path, dump_ranges, reporter=None):
     run_command(
-        ["cargo", "run", "--release", "--package", "pim-simulator", "--bin", "pim-simulator", "--",
+        ["cargo", "run", "--no-default-features", "--release", "--package", "pim-simulator", "--bin", "pim-simulator", "--",
          "-f", str(pim_dir), "-o", str(output_bin_path), "-d", dump_ranges],
         cwd=simulator_dir,
         reporter=reporter,
@@ -229,7 +233,7 @@ def validate_outputs(sim_arrays, runner_out_dir, outputs_descriptor, threshold=1
     all_passed = True
     rows = []
     for sim_array, (oi, name, _, shape) in zip(sim_arrays, outputs_descriptor):
-        csv_name = f"output{oi}_{name}.csv"
+        csv_name = f"output{oi}_{sanitize_output_name(name)}.csv"
         runner_array = np.loadtxt(runner_out_dir / csv_name, delimiter=',', dtype=np.float32).reshape(shape)
         max_diff = float(np.max(np.abs(sim_array.astype(np.float64) - runner_array.astype(np.float64))))
         passed = max_diff <= threshold