test pimcomp adversarial memory scheduling

normalize names and artifact paths

.agents/invariants/PERFORMANCE_OPTIMIZATION_INVARIANT.md

@@ -44,6 +44,19 @@ input size and target linear or sublinear time and space. Avoid repeated full-IR
 walks, nested scans, and per-operation recomputation when indexing, caching, or
 a single traversal can express the same behavior.
 
+Bufferization cost scales with both the number of operations and the number of
+MLIR values it receives. Upstream lowering and scheduling must therefore keep
+repeated work in compact structured operations, such as statically evaluable
+loops and batches, until after bufferization. Do not compensate for avoidable
+pre-bufferization expansion by weakening, partitioning, or special-casing the
+bufferization analysis; measure both operation and value counts at its input.
+Compactness must also be preserved after bufferization so that liveness,
+verification, memory planning, and other downstream passes do not repeat work
+per logical lane or iteration. Keep structured loops and batches intact until
+PIM ISA code generation is forced to scalarize them into concrete per-core
+instructions; earlier expansion requires an explicit semantic necessity and
+before/after operation and value counts.
+
 When linear-or-better complexity is not possible, use the lowest justified
 complexity and report the actual time and space Big-O, the input variable, and
 why a lower bound is not practical. Include that cost in the final report; do

.agents/invariants/PIMSIM_NN_ORACLE_INVARIANT.md

@@ -0,0 +1,22 @@
+# pimsim-nn Oracle Invariant
+
+`backend-simulators/pim/pimsim-nn` is the performance oracle. Its simulation
+behavior defines the hardware model used for Raptor/PIMCOMP comparisons.
+
+## Required invariant
+
+- Do not add instruction or operator support to pimsim-nn.
+- Changes to pimsim-nn must preserve simulation behavior exactly. Acceptable
+  changes are limited to behavior-neutral maintenance proven not to alter
+  simulated timing, scheduling, power, energy, or supported input programs.
+- Unsupported pimsim-nn operations must remain unsupported; do not approximate
+  their timing or map them onto another operation.
+- Adapt compiler inputs to the oracle instead. For YOLO, use
+  `validation/networks/pimcomp_models/yolo11n/yolo11n-pimsim-nn.onnx`, the
+  dedicated pimsim-ready performance artifact with Softmax operations removed.
+  Use `validation/networks/yolo11n/depth_51/yolo11n_depth_51.onnx` for YOLO
+  functional validation; removing Softmax changes the model's numerical
+  behavior, so the pimsim-ready artifact is not a correctness reference.
+
+Any proposed pimsim-nn behavior change requires explicit user authorization and
+must not be introduced as part of a compiler optimization.

.agents/invariants/PIPELINE_SCHEDULING_INVARIANT.md

@@ -0,0 +1,53 @@
+# Pipeline Scheduling Invariant
+
+## Scope
+
+This invariant applies to pipeline stage partitioning, physical-core
+assignment, scheduled materialization, deferred transfers, and pipeline
+synchronization.
+
+## Invariant
+
+A scheduled compute operation and all of its lanes belong to exactly one
+pipeline stage. An operation may consume results produced in its own stage or
+the immediately preceding stage only. Therefore every compute-graph edge from
+stage `S` targets stage `S` or `S + 1`; backward edges and dependencies that
+skip a stage are invalid.
+
+Dynamic function inputs are stage-zero sources. Any operation that directly
+consumes one must belong to stage 0. A later stage may consume that data only
+through an explicit result forwarded by the preceding stage.
+
+Each logical core belongs to exactly one stage capacity range before physical
+placement. Those ranges cover every core but may have different sizes when the
+initial partitioner predicts a lower maximum stage interval. Physical placement
+may map a stage to arbitrary core IDs using the injected target topology.
+Synchronization and deferred transfers consume the explicit stage identity;
+they must not infer it from a physical core number after placement.
+
+## Ownership
+
+Logical PEFT remains pipeline-agnostic. Stage partitioning is the first phase
+of pipeline scheduling and owns this invariant. It must construct a valid
+operation-level partition before physical-core packing. Operations split for
+physical capacity retain one shared stage identity. Repacking may move work
+only within its assigned stage. Deferred-transfer planning and
+synchronization consume the verified stage assignment; they must not repair
+or reinterpret it.
+
+## Verification
+
+Before scheduled materialization, verify that:
+
+- every compute instance has one valid physical core and stage;
+- all lanes of one compute operation have the same stage;
+- every direct dynamic-function-input consumer belongs to stage 0;
+- every compute-graph edge stays within a stage or advances exactly one stage;
+- every stage-local resident-weight set fits its assigned physical core; and
+- stage capacities cover all logical cores exactly once; and
+- physical placement is a permutation of all target cores.
+
+Pipeline scheduling tests must include an uneven physical-core layout and a
+graph with a long-lived dependency that would cross multiple naive stage
+cuts. End-to-end validation must preserve functional results and exercise the
+existing synchronization lowering without simulator changes.

AGENTS.md

@@ -6,6 +6,8 @@ Before modifying the relevant subsystem, read:
 
 * `.agents/invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md`
 * `.agents/invariants/PERFORMANCE_OPTIMIZATION_INVARIANT.md`
+* `.agents/invariants/PIMSIM_NN_ORACLE_INVARIANT.md`
+* `.agents/invariants/PIPELINE_SCHEDULING_INVARIANT.md`
 * `.agents/invariants/SPATIAL_TARGET_GENERALITY_INVARIANT.md`
 * Build commands:
   * `cmake --build ./build_release`

README.md

@@ -5,7 +5,7 @@ targeting in-memory computing / processing-in-memory (PIM) architectures. It
 extends ONNX-MLIR with a PIM accelerator and progressively lowers ONNX-MLIR
 through custom MLIR dialects to simulator artifacts.
 
-The current target is the PIM simulator stack under `backend-simulators/pim`.
+The current target is the Pim simulator stack under `backend-simulators/pim`.
 Raptor emits binary per-core `.pim` instruction files by default, plus
 `memory.bin`, `config.json`, and weight binaries. It can also emit per-core JSON
 instruction files with `--pim-emit-json`.
@@ -29,9 +29,9 @@ lowering, scheduling, memory layout, and code-generation optimizations.
 - `backend-simulators/pim/pim-simulator` is the in-tree Rust functional
   simulator used by validation. It reads Raptor's `pim/` artifact directory and
   compares simulator output against native ONNX-MLIR execution.
-- `backend-simulators/pim/pimsim-nn` is the non-functional simulator submodule
-  used internally by validation for latency, power, and energy.
-  The helper scripts in `pimcomp_utils/` are for comparison with PIMCOMP-NN and
+- `backend-simulators/pim/pimsim-nn` contains the non-functional Pimsim
+  simulator used internally by validation for latency, power, and energy.
+  The helper scripts in `pimcomp_utils/` are for comparison with Pimcomp and
   contain local paths; treat them as local utilities, not portable workflows.
 
 ## Compilation pipeline
@@ -43,7 +43,7 @@ them to ONNX-MLIR through generated shim directories under
 High-level lowering flow:
 
 ```
-ONNX-MLIR -> Spatial -> Pim (tensor) -> Pim (bufferized) -> PIM artifacts
+ONNX-MLIR -> Spatial -> Pim (tensor) -> Pim (bufferized) -> Pim artifacts
 ```
 
 1. **ONNX -> Spatial** (`src/PIM/Conversion/ONNXToSpatial`).
@@ -52,84 +52,147 @@ ONNX-MLIR -> Spatial -> Pim (tensor) -> Pim (bufferized) -> PIM artifacts
    `Patterns/{Math,NN,Tensor}` and currently cover Conv, Gemm, MatMul,
    elementwise Add/Mul/Div, ReduceMean, pooling, Relu, Sigmoid, Softmax,
    Concat, Gather, Reshape, Resize, and Split.
+   The compiler-layer target adapter supplies the target-neutral
+   `SpatialTargetResources`. Layout-aware plan ops advertise typed alternatives
+   through the Spatial layout interface; the layout planner records the
+   selected layout and explicit materialization edges. `LowerSpatialPlans`
+   then pattern-lowers those selected plans. Contraction and Conv lowering
+   keep semantic problems, target-dependent plans, and IR materializers in
+   separate layers. Passes and their invariant/layout analyses live under
+   `Passes/Transforms` and `Passes/Analyses`.
 
-2. **Merge compute nodes**
-   (`src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes`).
-   Builds a compute graph, schedules it with the PEFT scheduler, and materializes
-   the merge schedule into Spatial IR. Supporting scheduling code lives under
-   `MergeComputeNodes/Scheduling`.
+2. **Merge, schedule, and realize Spatial communication**
+   (`src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes`).
+   `TrivialGraphComputeMerge` performs local graph merging. One
+   `ScheduleAndRealizeSpatial` pass then owns scheduling, intermediate
+   verification, communication realization, and final verification. Supporting
+   scheduling code lives under `MergeComputeNodes/Scheduling`.
 
 3. **Spatial -> Pim** (`src/PIM/Conversion/SpatialToPim`).
    Lowers Spatial operations to the `pim` dialect (`src/PIM/Dialect/Pim`),
    including `pim.core`, `pim.core_batch`, communication, tensor packing, global
    tensor materialization, and return-path normalization.
 
-4. **Bufferization** (`src/PIM/Dialect/Pim/Transforms/Bufferization`).
-   Converts tensor-semantics PIM IR into memref-semantics PIM IR using MLIR's
-   bufferization interfaces.
+4. **Bufferization** (`src/PIM/Dialect/Pim/Passes/Transforms/Bufferization`).
+   `PimBufferizationPreparation` establishes writable destinations without
+   duplicating the one-shot copy analysis, `PimOneShotBufferization` runs
+   MLIR's one-shot analysis,
+   `PimMemoryNormalization` forwards/removes redundant copies and normalizes
+   addressable accesses, and `PimBufferizationVerification` checks tensor
+   absence, contiguity, and copy address spaces.
 
-5. **PIM local-memory planning**
-   (`src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning`).
+5. **Pim local-memory planning**
+   (`src/PIM/Dialect/Pim/Passes/Transforms/LocalMemoryPlanning`).
    Computes whole-core lifetimes, reuses addresses for non-overlapping
-   allocations, and records the explicit plan in PIM IR.
-6. **PIM verification and code generation** (`src/PIM/Pass/PimCodegen` and
+   allocations, and records the explicit plan in Pim IR. Reusable lifetime
+   analysis lives under `src/PIM/Dialect/Pim/Passes/Analyses`.
+6. **Pim verification and code generation** (`src/PIM/Passes/PimCodegen` and
    `src/PIM/Compiler`).
-   Verifies the memory plan and other PIM invariants, then emits `.pim` core
+   Verifies the memory plan and other Pim invariants, then emits `.pim` core
    files, weights, and `memory.bin` / `config.json` without rerunning liveness.
 
 Supporting pieces:
 - `src/PIM/Common` - shared IR, filesystem, diagnostics, reports, and utility
   helpers.
-- `src/PIM/Compiler` - PIM compiler options, planned-address materialization, binary
+- `src/PIM/Compiler` - Pim compiler options, planned-address materialization, binary
   instruction format, artifact writing, weight emission, and codegen entry
   points.
 - `src/PIM/Conversion/SpatialToGraphviz` - optional Spatial graphviz conversion
   pass.
-- `src/PIM/Pass` - pass registration and auxiliary passes.
+- `src/PIM/Passes` - pass registration and auxiliary passes.
 - `src/PIM/PimAccelerator.{cpp,hpp}` - ONNX-MLIR accelerator entry point.
 
-## PIM compiler options
+## Pim compiler options
 
-Pass these to `onnx-mlir` when compiling for PIM. These are all Raptor/PIM-specific
+Pass these to `onnx-mlir` when compiling for Pim. These are all Raptor/Pim-specific
 options; `onnx-mlir --help` lists the inherited ONNX-MLIR options.
 
-- `--maccel=PIM` - select the PIM accelerator.
+- `--maccel=PIM` - select the Pim accelerator. Default: no Pim accelerator.
 - `--EmitSpatial`, `--EmitPim`, `--EmitPimBufferized`,
-  `--EmitPimCodegen` - stop the PIM pipeline at the requested stage. The PIM
-  default is `--EmitPimCodegen`.
-- `--core-count=<N>` - required positive core count for PIM compilation.
-- `--crossbar-size=<N>` - crossbar width/height. Default in code is `128`.
-- `--crossbar-count=<N>` - crossbars per core. Default in code is `64`.
-- `--pim-target-config=<PATH>` - optional PIM target configuration used by the
+  `--EmitPimCodegen` - stop the Pim pipeline at the requested stage. Default:
+  `--EmitPimCodegen` for Pim compilation.
+- `--core-count=<N>` - required positive core count for Pim compilation.
+  Default: none; this option is required.
+- `--crossbar-size=<N>` - required positive crossbar width/height for Pim
+  compilation. Default: none; this option is required.
+- `--crossbar-count=<N>` - required positive crossbar count per core for Pim
+  compilation. Default: none; this option is required.
+- `--pipeline=<N>` - number of throughput pipeline stages; `1` preserves
+  latency scheduling. Default: `1`.
+- `--pim-target-config=<PATH>` - optional Pim target configuration used by the
   target adapter to construct the target-neutral Spatial scheduling cost and
   topology model. Resource values must match the explicit core/crossbar flags.
+  Default: empty; use the built-in target model.
 - `--pim-memory-report=<summary|none>` - emit the concise combined memory report
-  under `reports/memory_report.txt`, or disable it. Default is `summary`.
-- `--pim-only-codegen` - assume input is already bufferized PIM IR and only run
-  the codegen tail.
+  under `reports/memory_report.txt`, or disable it. Default: `summary`.
+- `--pim-only-codegen` - assume input is already bufferized Pim IR and only run
+  the codegen tail. Default: off.
+- `--pim-disable-synchronization` - omit generated `wait` and `sync`
+  instructions for performance ablation. Default: off.
+- `--pim-disable-spatial-planning` - select the first, trivial DenseNCHW layout
+  alternative for every Spatial plan operation, disabling cost-based layout
+  planning while leaving ONNX rewrites and graph-compute merging enabled.
+  Default: off.
+
+### Spatial layout plan variants
+
+Spatial plan operations advertise alternatives as an exact combination of
+operand physical layouts and one result physical layout. Every plan operation
+has the default `DenseNCHW -> DenseNCHW` alternative. The planner can select
+the following additional variants when the operation, tensor shapes, and
+target resources make them legal:
+
+| Physical layout or plan | Meaning and current use |
+|---|---|
+| `DenseNCHW` | Ordinary dense NCHW storage. This is the first alternative and the one selected by `--pim-disable-spatial-planning`. |
+| `NHWCRowStrip` | Row-strip storage for NCHW logical tensors: spatial rows are processed as channel vectors. This enables row-strip lowering through compatible chains. |
+| `Fragmented` | Fragmented physical input accepted by `Flatten`, which reassembles it to dense NCHW. It is not currently selected as a plan result. |
+| `NCHWRowStrip` | A Spatial IR layout enum value reserved for NCHW-oriented row strips; current layout-capability implementations do not advertise it as a plan alternative. |
+
+The operation-specific non-trivial alternatives are:
+
+| Plan operation | Additional alternatives beyond dense NCHW |
+|---|---|
+| `Conv2D` | Dense input to row-strip output, or row-strip input to row-strip output when the target-dependent Conv lowering supports it. |
+| `Flatten` | Fragmented input to dense output, or row-strip input to dense output when legal. |
+| `Relu` | Row-strip input to row-strip output. |
+| `SiLU` | Row-strip input to row-strip output, with a stronger intrinsic cost preference than the generic row-strip variant. |
+| `ResizeNearest` | Row-strip input to row-strip output when its lowering is legal. |
+| `MaxPool2D` | Dense input to row-strip output, or row-strip input to row-strip output. |
+| `GlobalAveragePool` | Dense input to row-strip output, or row-strip input to row-strip output. |
+| `BiasAdd` | Row-strip data input plus a dense bias input to row-strip output when the bias shape is supported. |
+| `Add` | All data inputs row-strip to row-strip output. |
+| `Concat` | All inputs row-strip to row-strip output. |
+
+Cost-based planning scores intrinsic alternative cost, operand layout
+mismatches, and downstream incompatibility, then iterates in alternating
+forward and reverse operation order until the bounded analysis converges.
+Function results are required to remain `DenseNCHW`; explicit materialization
+operations reconcile layout mismatches at boundaries. With
+`--pim-disable-spatial-planning`, the pass still runs and records a valid plan,
+but chooses the first dense alternative for every plan operation. Later graph
+compute merging is unchanged, so elementwise operations such as `Relu` remain
+separate from neighboring parallel operations and can create fan-out/fan-in
+diamonds.
 - `--pim-emit-json` - also emit `core_*.json` instruction files alongside
-  `core_*.pim`.
+  `core_*.pim`. Default: off.
 - `--pim-export-spatial-dataflow=<none|spatial1|spatial2|spatial3|spatial4|all>` -
   control Spatial dataflow CSV reports for the graph, trivially merged graph,
-  scheduled, and realized snapshots under `reports/`. Default is `none`.
+  scheduled, and realized snapshots under `reports/`. Default: `none`.
 - `--pim-conv-lowering=<auto|legacy|depthwise|packed-im2col|streamed-patch|streamed-packed|output-channel-tiled|input-k-tiled|tiled-2d>` -
-  select the convolution lowering strategy. Default is `auto`.
+  select the convolution lowering strategy. Default: `auto`.
 - `--pim-conv-im2col-max-elements=<N>` - maximum globally materialized im2col
-  elements per convolution before streaming. Default is `1048576`.
+  elements per convolution before streaming. Default: `1048576`.
 - `--pim-conv-stream-chunk-positions=<N>` - maximum output positions per
-  streamed convolution chunk. Default is `1024`.
-- `--pim-report-conv-lowering=<true|false>` - emit the bounded convolution
-  lowering report. Default is `true`.
-- `--use-experimental-conv-impl` - use the alternate convolution lowering.
+  streamed convolution chunk. Default: `1024`.
+- `--pim-report-conv-lowering=<true|false>` - emit a bounded convolution
+  lowering report. Default: `true`.
 - `--pim-detect-communication-deadlock` - statically simulate expanded
-  send/receive ordering and reject blocking deadlocks. Default is off.
-- `--pim-materialize-scalar-fanout-global-order` - use the experimental,
-  expensive globally ordered scalar-fanout materializer. Default is off.
-- `--pim-trace-communication-materialization` - emit verbose communication
-  materialization diagnostics and provenance attributes. Default is off.
-- `--ignore-concat-error` - soft-fail a ConcatOp corner case.
+  send/receive ordering and reject blocking deadlocks. Default: off.
+- `--pim-verify-bufferization-copy-freedom` - run the expensive official Pim
+  tensor-copy freedom proof before bufferization. Default: off.
 
-## Standard PIM hardware profile
+## Standard Pim hardware profile
 
 Raptor's standard development and YOLO validation profile is:
 
@@ -143,7 +206,8 @@ Canonical compiler flags:
 
 `--crossbar-count=64 --crossbar-size=128 --core-count=144`
 
-`--core-count` remains mandatory and must be passed explicitly to the compiler.
+`--crossbar-size`, `--crossbar-count`, and `--core-count` remain mandatory and
+must be passed explicitly to the compiler.
 
 Example:
 
@@ -153,11 +217,11 @@ Example:
   --crossbar-count=64 --crossbar-size=128 --core-count=144
 ```
 
-This writes PIM artifacts under `/tmp/raptor/pim/`.
+This writes Pim artifacts under `/tmp/raptor/pim/`.
 
 ## Validation
 
-Functional validation compiles ONNX models, compares native ONNX-MLIR and PIM
+Functional validation compiles ONNX models, compares native ONNX-MLIR and Pim
 simulator outputs, and optionally reports latency, power, and energy. See
 [`validation/README.md`](validation/README.md) for prerequisites, usage,
 options, artifacts, and results.
@@ -270,7 +334,7 @@ cd backend-simulators/pim/pim-simulator
 cargo test
 ```
 
-## Repository Layout
+## Repository layout
 
 - `src/PIM/` - PIM accelerator implementation.
 - `test/PIM/` - PIM C++ unit tests.
@@ -278,6 +342,6 @@ cargo test
   slices, and pimsim config generation.
 - `backend-simulators/pim/pim-simulator/` - in-tree Rust functional simulator.
 - `backend-simulators/pim/pimsim-nn/` - non-functional simulator submodule.
-- `pimcomp_utils/` - local comparison helpers for PIMCOMP-NN.
+- `pimcomp_utils/` - local comparison helpers for Pimcomp.
 - `.github/actions/` and `.github/workflows/validate_operations.yml` - CI setup
   for MLIR/Protobuf caching, building Raptor, and validation.

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

@@ -4,17 +4,18 @@ use mimalloc::MiMalloc;
 static GLOBAL: MiMalloc = MiMalloc;
 
 use anyhow::{Context, Result, bail};
-use clap::Parser;
+use clap::{Parser, ValueEnum};
 use glob::glob;
 use pimcore::binary_to_instruction::binary_to_executor;
 use pimcore::cpu::crossbar::Crossbar;
 use pimcore::json_to_instruction::json_to_executor;
 use pimcore::memory_manager::CoreMemory;
 use pimcore::tracing::TRACER;
+use pimcore::{DiagnosticSchedulePolicy, DiagnosticScheduleTarget};
 use serde_json::Value;
 use std::collections::HashMap;
-use std::fs::{self, File, read_link};
-use std::io::{BufReader, Write};
+use std::fs::{self, File};
+use std::io::BufReader;
 use std::path::PathBuf;
 
 /// Program to simulate core execution configuration
@@ -44,14 +45,79 @@ struct Args {
     /// Comma separated list of (address,size) for memory output dump
     #[arg(short, long, value_delimiter = ',', num_args = 1.., value_name = "ADDR,SIZE")]
     dump: Vec<usize>,
+
+    /// Simulator execution mode
+    #[arg(long, value_enum, default_value_t = ExecutionMode::Latency)]
+    mode: ExecutionMode,
+
+    /// Number of inputs to execute (required in throughput mode)
+    #[arg(long)]
+    batch_size: Option<u32>,
+
+    /// Directory containing input_*.bin files, one per batch entry
+    #[arg(long = "input-dir")]
+    input_dir: PathBuf,
+
+    /// Optional directory for per-iteration output dumps
+    #[arg(long)]
+    batch_output_dir: Option<PathBuf>,
+
+    /// Optional JSONL shadow provenance trace
+    #[arg(long)]
+    provenance_trace: Option<PathBuf>,
+
+    /// Diagnostic-only barrier between global throughput iterations
+    #[arg(long)]
+    provenance_global_barrier: bool,
+
+    /// Diagnostic-only ready-core delay, formatted as CORE:CYCLES
+    #[arg(long, value_name = "CORE:CYCLES")]
+    provenance_core_stall: Option<String>,
+
+    /// Diagnostic scheduler policy; greedy is the unchanged default
+    #[arg(long, value_enum, default_value_t = DiagnosticSchedulePolicyArg::Greedy)]
+    diagnostic_schedule_policy: DiagnosticSchedulePolicyArg,
+
+    /// Deterministic seed for the randomized diagnostic scheduler
+    #[arg(long, default_value_t = 0)]
+    diagnostic_schedule_seed: u64,
+
+    /// Adversarial target: WCORE:WPC:RCORE:RPC:BEGIN:END[:READER_ITER:WRITER_MIN_ITER]
+    #[arg(long, value_name = "WCORE:WPC:RCORE:RPC:BEGIN:END[:RITER:WMIN]")]
+    diagnostic_schedule_target: Option<String>,
+
+    /// Maximum number of target-consumer deferrals
+    #[arg(long, default_value_t = 10_000)]
+    diagnostic_schedule_deferral_budget: u64,
+
+    /// Diagnostic delay applied when a target reader reaches its PC, formatted as CORE:PC:CYCLES or CORE:PC:ITERATION:CYCLES
+    #[arg(long, value_name = "CORE:PC[:ITERATION]:CYCLES")]
+    diagnostic_target_stall: Option<String>,
+}
+
+#[derive(Clone, Debug, ValueEnum)]
+enum ExecutionMode {
+    Latency,
+    Throughput,
+}
+
+#[derive(Clone, Copy, Debug, ValueEnum)]
+enum DiagnosticSchedulePolicyArg {
+    Greedy,
+    Randomized,
+    Adversarial,
 }
 
 fn main() -> Result<()> {
     let args = Args::parse();
 
-    let config_json = retrive_config(&args)?;
-    let mut core_inputs = retrive_cores(&args)?;
-    let memory = retrive_memory(&args)?;
+    let config_json = retrieve_config(&args)?;
+    let batch_size = batch_size(&args)?;
+    let input_regions = input_regions(&config_json)?;
+    let input_data = retrieve_inputs(&args, batch_size)?;
+    let inputs: Vec<&[u8]> = input_data.iter().map(Vec::as_slice).collect();
+    let mut core_inputs = retrieve_cores(&args)?;
+    let memory = retrieve_memory(&args)?;
     let global_crossbars = get_crossbars(&config_json, &args).unwrap();
     let crossbars = map_crossbars_to_cores(&config_json, &args, &global_crossbars);
     let mut executor = match &mut core_inputs {
@@ -63,15 +129,161 @@ fn main() -> Result<()> {
         }
     };
     set_memory(&mut executor, memory);
+    if let Some(path) = &args.provenance_trace {
+        executor.enable_provenance(path)?;
+    }
+    executor.set_provenance_global_barrier(args.provenance_global_barrier);
+    if let Some(spec) = args.provenance_core_stall.as_deref() {
+        let (core, cycles) = parse_core_stall(spec)?;
+        executor.set_provenance_core_stall(core, cycles);
+    }
+    executor.set_diagnostic_schedule_policy(match args.diagnostic_schedule_policy {
+        DiagnosticSchedulePolicyArg::Greedy => DiagnosticSchedulePolicy::Greedy,
+        DiagnosticSchedulePolicyArg::Randomized => DiagnosticSchedulePolicy::Randomized,
+        DiagnosticSchedulePolicyArg::Adversarial => DiagnosticSchedulePolicy::Adversarial,
+    });
+    executor.set_diagnostic_schedule_seed(args.diagnostic_schedule_seed);
+    executor.set_diagnostic_schedule_deferral_budget(args.diagnostic_schedule_deferral_budget);
+    if let Some(spec) = args.diagnostic_target_stall.as_deref() {
+        let (core, pc, iteration, cycles) = parse_target_stall(spec)?;
+        executor.set_diagnostic_target_stall(core, pc, iteration, cycles);
+    }
+    if let Some(spec) = args.diagnostic_schedule_target.as_deref() {
+        executor.set_diagnostic_schedule_target(parse_schedule_target(spec)?);
+    } else if matches!(
+        args.diagnostic_schedule_policy,
+        DiagnosticSchedulePolicyArg::Adversarial
+    ) {
+        bail!("adversarial scheduling requires --diagnostic-schedule-target");
+    }
     TRACER
         .lock()
         .unwrap()
         .init(executor.cpu().num_core(), args.output.clone());
-    executor.execute()?;
-    dump_memory(executor, &args)?;
+    let dumps = dump_ranges(&args.dump)?;
+    let batch_outputs = executor.execute_batch(&inputs, &input_regions, &dumps)?;
+    fs::write(
+        &args.output,
+        batch_outputs
+            .last()
+            .context("simulation produced no output")?,
+    )?;
+    if let Some(batch_output_dir) = args.batch_output_dir {
+        write_batch_outputs(batch_output_dir, batch_outputs)?;
+    }
     Ok(())
 }
 
+fn parse_core_stall(spec: &str) -> Result<(usize, u64)> {
+    let (core, cycles) = spec
+        .split_once(':')
+        .context("--provenance-core-stall must be CORE:CYCLES")?;
+    let core = core.parse().context("invalid stalled core")?;
+    let cycles = cycles.parse().context("invalid stall cycle count")?;
+    if cycles == 0 {
+        bail!("--provenance-core-stall cycles must be positive");
+    }
+    Ok((core, cycles))
+}
+
+fn parse_schedule_target(spec: &str) -> Result<DiagnosticScheduleTarget> {
+    let values: Vec<usize> = spec
+        .split(':')
+        .map(|value| {
+            value
+                .parse()
+                .with_context(|| format!("invalid schedule target field: {value}"))
+        })
+        .collect::<Result<_>>()?;
+    if values.len() != 6 && values.len() != 8 {
+        bail!(
+            "--diagnostic-schedule-target requires WCORE:WPC:RCORE:RPC:BEGIN:END with optional RITER:WMIN"
+        );
+    }
+    if values[5] <= values[4] {
+        bail!("schedule target address end must be greater than begin");
+    }
+    Ok(DiagnosticScheduleTarget {
+        writer_core: values[0],
+        writer_pc: values[1],
+        reader_core: values[2],
+        reader_pc: values[3],
+        address_begin: values[4],
+        address_end: values[5],
+        reader_iteration: (values.len() == 8).then_some(values[6] as u32),
+        writer_min_iteration: (values.len() == 8).then_some(values[7] as u32),
+    })
+}
+
+fn parse_target_stall(spec: &str) -> Result<(usize, usize, Option<u32>, u64)> {
+    let values: Vec<&str> = spec.split(':').collect();
+    if values.len() != 3 && values.len() != 4 {
+        bail!("--diagnostic-target-stall must be CORE:PC:CYCLES or CORE:PC:ITERATION:CYCLES");
+    }
+    let core = values[0].parse().context("invalid target-stall core")?;
+    let pc = values[1].parse().context("invalid target-stall PC")?;
+    let (iteration, cycle_field) = if values.len() == 4 {
+        (
+            Some(
+                values[2]
+                    .parse()
+                    .context("invalid target-stall iteration")?,
+            ),
+            values[3],
+        )
+    } else {
+        (None, values[2])
+    };
+    let cycles = cycle_field
+        .parse()
+        .context("invalid target-stall cycle count")?;
+    if cycles == 0 {
+        bail!("--diagnostic-target-stall cycles must be positive");
+    }
+    Ok((core, pc, iteration, cycles))
+}
+
+fn batch_size(args: &Args) -> Result<u32> {
+    match (&args.mode, args.batch_size) {
+        (ExecutionMode::Latency, None | Some(1)) => Ok(1),
+        (ExecutionMode::Latency, Some(_)) => bail!("latency mode requires batch size 1"),
+        (ExecutionMode::Throughput, Some(0)) => bail!("batch size must be positive"),
+        (ExecutionMode::Throughput, Some(batch_size)) => Ok(batch_size),
+        (ExecutionMode::Throughput, None) => bail!("throughput mode requires --batch-size"),
+    }
+}
+
+fn input_regions(config: &Value) -> Result<Vec<(usize, usize)>> {
+    let addresses = config
+        .get("inputs_addresses")
+        .and_then(Value::as_array)
+        .context("config.json has no inputs_addresses array")?;
+    let sizes = config
+        .get("inputs_sizes")
+        .and_then(Value::as_array)
+        .context("config.json has no inputs_sizes array")?;
+    if addresses.len() != sizes.len() {
+        bail!("config.json input address/size count mismatch");
+    }
+    addresses
+        .iter()
+        .zip(sizes)
+        .map(|(address, size)| {
+            Ok((
+                usize::try_from(address.as_u64().context("invalid input address")?)?,
+                usize::try_from(size.as_u64().context("invalid input size")?)?,
+            ))
+        })
+        .collect()
+}
+
+fn retrieve_inputs(args: &Args, batch_size: u32) -> Result<Vec<Vec<u8>>> {
+    (0..batch_size)
+        .map(|index| args.input_dir.join(format!("input_{index}.bin")))
+        .map(|path| fs::read(&path).with_context(|| format!("Failed to read input file: {path:?}")))
+        .collect()
+}
+
 fn map_crossbars_to_cores<'c>(
     config: &Value,
     args: &Args,
@@ -110,11 +322,11 @@ fn map_crossbars_to_cores<'c>(
             sym_link_files.sort_by_key(|&(num, _)| num);
 
             for (_, symlink) in sym_link_files {
-                let real_path = read_link(symlink).unwrap();
+                let real_path = symlink.canonicalize().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",
+                    "symlink points to {:?}\n a crossbar that was not stored",
                     real_path
                 );
 
@@ -131,7 +343,7 @@ fn map_crossbars_to_cores<'c>(
 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 column_crossbar = xbar_size[1].as_i64().unwrap() as usize;
     let mut res = HashMap::new();
 
     if let Some(folder) = args.folder.as_ref() {
@@ -154,7 +366,7 @@ fn get_crossbars(config: &Value, args: &Args) -> anyhow::Result<HashMap<String,
             let bytes = std::fs::read(weight_file.path()).expect("Failed to read binary file");
             let stored_row_bytes = bytes.len() / rows_crossbar;
             let mut crossbar = Crossbar::new(
-                std::cmp::max(column_corssbar * 4, stored_row_bytes),
+                std::cmp::max(column_crossbar * 4, stored_row_bytes),
                 rows_crossbar,
                 CoreMemory::new(),
             );
@@ -162,6 +374,8 @@ fn get_crossbars(config: &Value, args: &Args) -> anyhow::Result<HashMap<String,
             res.insert(
                 weight_file
                     .path()
+                    .canonicalize()
+                    .context("Failed to resolve crossbar path")?
                     .to_str()
                     .context("file name not utf-8")?
                     .to_string(),
@@ -172,21 +386,22 @@ fn get_crossbars(config: &Value, args: &Args) -> anyhow::Result<HashMap<String,
     Ok(res)
 }
 
-fn dump_memory(mut executor: pimcore::Executable, args: &Args) -> Result<()> {
-    let dumps: Vec<(usize, usize)> = args
-        .dump
+fn dump_ranges(values: &[usize]) -> Result<Vec<(usize, usize)>> {
+    if !values.len().is_multiple_of(2) {
+        bail!("memory dump requires address,size pairs");
+    }
+    Ok(values
         .chunks_exact(2)
         .map(|chunk| (chunk[0], chunk[1]))
-        .collect();
-    let mut out_file = fs::OpenOptions::new()
-        .create(true)
-        .write(true)
-        .truncate(true)
-        .open(&args.output)
-        .with_context(|| format!("cannot open file {:?} for writing", args.output))?;
+        .collect())
+}
 
-    for (address, size) in dumps {
-        out_file.write_all(executor.cpu_mut().host().load::<u8>(address, size).unwrap()[0])?;
+fn write_batch_outputs(output_dir: PathBuf, outputs: Vec<Vec<u8>>) -> Result<()> {
+    fs::create_dir_all(&output_dir)
+        .with_context(|| format!("cannot create batch output directory {output_dir:?}"))?;
+    for (iteration, output) in outputs.into_iter().enumerate() {
+        let path = output_dir.join(format!("output_{iteration:06}.bin"));
+        fs::write(&path, output).with_context(|| format!("cannot write batch output {path:?}"))?;
     }
     Ok(())
 }
@@ -195,7 +410,7 @@ fn set_memory(executor: &mut pimcore::Executable, memory: Vec<u8>) {
     executor.cpu_mut().host().execute_store(0, &memory).unwrap();
 }
 
-fn retrive_memory(args: &Args) -> Result<Vec<u8>> {
+fn retrieve_memory(args: &Args) -> Result<Vec<u8>> {
     let memory_path = if let Some(mem_override) = &args.memory {
         mem_override.clone()
     } else if let Some(folder) = &args.folder.as_ref() {
@@ -235,7 +450,7 @@ enum CoreInputs {
     Binary(Vec<Vec<u8>>),
 }
 
-fn retrive_cores(args: &Args) -> Result<CoreInputs, anyhow::Error> {
+fn retrieve_cores(args: &Args) -> Result<CoreInputs, anyhow::Error> {
     if let Some(cores_override) = &args.cores {
         let first_extension = cores_override
             .first()
@@ -308,7 +523,7 @@ fn core_sort_key(path: &PathBuf) -> i32 {
     stem.parse::<i32>().unwrap()
 }
 
-fn retrive_config(args: &Args) -> Result<Value, anyhow::Error> {
+fn retrieve_config(args: &Args) -> Result<Value, anyhow::Error> {
     let config_path: PathBuf = {
         let override_path = args.config.as_ref();
         let folder = args.folder.as_ref();

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

@@ -80,19 +80,19 @@ fn read_i32_le(bytes: &[u8], offset: usize) -> i32 {
 
 fn parse_binary_records(bytes: &[u8]) -> Result<Vec<InstructionRecord>> {
     ensure!(bytes.len() >= HEADER_SIZE, "binary core file too small");
-    ensure!(&bytes[0..4] == MAGIC, "invalid PIM binary magic");
+    ensure!(&bytes[0..4] == MAGIC, "invalid Pim binary magic");
 
     let version = read_u32_le(bytes, 4);
     ensure!(
         version == VERSION,
-        "unsupported PIM binary version {version}"
+        "unsupported Pim binary version {version}"
     );
 
     let instruction_count = read_u32_le(bytes, 8) as usize;
     let expected_len = HEADER_SIZE + instruction_count * RECORD_SIZE;
     ensure!(
         bytes.len() == expected_len,
-        "PIM binary size mismatch: expected {expected_len} bytes, got {}",
+        "Pim binary size mismatch: expected {expected_len} bytes, got {}",
         bytes.len()
     );
 
@@ -326,12 +326,16 @@ fn append_record(
             inst_builder.make_inst(recv, inst_data_builder.build());
         }
         31 => {
+            inst_data_builder.set_offset_select_value(generic1, generic2);
             inst_builder.make_inst(wait, inst_data_builder.build());
         }
         32 => {
+            inst_data_builder
+                .set_imm_core(r2_or_imm + 1)
+                .set_offset_select_value(generic1, 0);
             inst_builder.make_inst(sync, inst_data_builder.build());
         }
-        _ => bail!("unsupported PIM binary opcode {opcode}"),
+        _ => bail!("unsupported Pim binary opcode {opcode}"),
     }
     Ok(())
 }

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

@@ -1,7 +1,6 @@
 use crate::memory_manager::{CoreMemory, MemoryStorable};
 use anyhow::{Result, bail, ensure};
 
-
 #[derive(Debug, Clone)]
 pub struct Crossbar {
     max_width: usize,
@@ -12,7 +11,12 @@ pub struct Crossbar {
 
 impl Crossbar {
     pub fn new(width: usize, height: usize, memory: CoreMemory) -> Self {
-        Self { max_width: width, max_height: height, memory, stored_bytes:0 }
+        Self {
+            max_width: width,
+            max_height: height,
+            memory,
+            stored_bytes: 0,
+        }
     }
 
     pub fn width(&self) -> usize {
@@ -27,27 +31,36 @@ impl Crossbar {
         self.stored_bytes
     }
 
-    pub fn execute_store<T>(&mut self, element: &[T]) -> Result<()> where
-            T: MemoryStorable, {
+    pub fn execute_store<T>(&mut self, element: &[T]) -> Result<()>
+    where
+        T: MemoryStorable,
+    {
         self.memory.clear();
         let total_size = self.max_width * self.max_height;
         self.memory.set_capacity(total_size);
         let stored_size = std::mem::size_of_val(element);
-        ensure!(stored_size <= total_size, "Storing more than crossbar can handle");
-        self.stored_bytes=stored_size;
+        ensure!(
+            stored_size <= total_size,
+            "Storing more than crossbar can handle"
+        );
+        self.stored_bytes = stored_size;
         self.memory.execute_store(0, element)
     }
 
-    pub fn load<T>(&self, size: usize) -> Result<Vec<&[T]>> where
-            T: MemoryStorable, {
-        if self.memory.get_len() <  size 
-            //|| self.stored_bytes < size 
+    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);
+            bail!(
+                "Loading outside crossbar boundary [{} {}] <  {}",
+                self.stored_bytes,
+                self.memory.get_len(),
+                size
+            );
         }
         self.memory.load_const(0, size)
     }
-
-
-
 }

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

@@ -1,23 +1,53 @@
 use crate::utility::AddressArg;
-use std::{collections::HashMap, fmt::Debug};
 use anyhow::{Context, Result, ensure};
+use std::{collections::HashMap, fmt::Debug};
 
+use super::{DiagnosticSchedulePolicy, DiagnosticScheduleTarget};
 use crate::{
     cpu::crossbar::Crossbar,
     instruction_set::Instructions,
     memory_manager::{CoreMemory, MemoryStorable, type_traits::TryToUsize},
+    provenance::ProvenanceTracker,
 };
+use serde_json::json;
 
 pub mod crossbar;
 
-
 #[derive(Debug, Clone)]
 pub struct CPU<'a> {
     cores: Box<[Core<'a>]>,
+    batch_outputs: Option<BatchOutputs>,
+    provenance: Option<ProvenanceTracker>,
+}
+
+#[derive(Debug, Clone)]
+struct BatchOutputs {
+    iteration: usize,
+    ranges: Vec<(usize, usize)>,
+    outputs: Vec<Vec<u8>>,
+}
+
+impl BatchOutputs {
+    fn record(&mut self, address: usize, bytes: &[u8]) {
+        let output = &mut self.outputs[self.iteration];
+        let store_end = address + bytes.len();
+        let mut output_offset = 0;
+        for &(range_address, range_size) in &self.ranges {
+            let start = address.max(range_address);
+            let end = store_end.min(range_address + range_size);
+            if start < end {
+                let size = end - start;
+                output[output_offset + start - range_address
+                    ..output_offset + start - range_address + size]
+                    .copy_from_slice(&bytes[start - address..start - address + size]);
+            }
+            output_offset += range_size;
+        }
+    }
 }
 
 impl<'a> CPU<'a> {
-    pub fn new(num_cores: impl TryToUsize, crossbars: Vec<Vec<&'a Crossbar>> ) -> Self {
+    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");
         assert!(crossbars.len() == num_cores + 1);
         let mut cores = Vec::new();
@@ -26,29 +56,346 @@ impl<'a> CPU<'a> {
         }
         Self {
             cores: cores.into(),
+            batch_outputs: None,
+            provenance: None,
         }
     }
 
-    pub fn host<'b>(&'b mut self) -> &'b mut Core<'a>
-        where 'a : 'b
-    {
-        & mut self.cores[0]
+    pub(crate) fn enable_provenance(
+        &mut self,
+        path: impl AsRef<std::path::Path>,
+    ) -> std::io::Result<()> {
+        self.provenance = Some(ProvenanceTracker::new(self.cores.len(), path)?);
+        Ok(())
     }
 
-    pub fn core<'b >(&'b mut self, index: impl TryToUsize) -> &'b mut Core<'a>
-        where 'a : 'b
+    pub(crate) fn begin_provenance_batch(&mut self, batch_size: usize) {
+        if let Some(provenance) = &mut self.provenance {
+            provenance.begin_batch(batch_size);
+        }
+    }
+
+    pub(crate) fn set_execution_context(
+        &mut self,
+        cycle: u64,
+        core: usize,
+        pc: usize,
+        iteration: u32,
+    ) {
+        if let Some(provenance) = &mut self.provenance {
+            provenance.set_context(cycle, core, pc, iteration);
+        }
+    }
+
+    pub(crate) fn provenance_input_store(&mut self, address: usize, size: usize, sample: u32) {
+        if let Some(provenance) = &mut self.provenance {
+            provenance.input_store(address, size, sample);
+        }
+    }
+
+    pub(crate) fn provenance_global_store_from_local(
+        &mut self,
+        core: usize,
+        global_address: usize,
+        local_address: usize,
+        size: usize,
+    ) {
+        if let Some(provenance) = &mut self.provenance {
+            provenance.global_store_from_local(core, global_address, local_address, size);
+        }
+    }
+
+    pub(crate) fn provenance_global_load_to_local(
+        &mut self,
+        core: usize,
+        global_address: usize,
+        local_address: usize,
+        size: usize,
+    ) {
+        if let Some(provenance) = &mut self.provenance {
+            provenance.global_load_to_local(core, global_address, local_address, size);
+        }
+    }
+
+    pub(crate) fn provenance_local_copy(
+        &mut self,
+        core: usize,
+        destination: usize,
+        source: usize,
+        size: usize,
+        operation: &'static str,
+    ) {
+        if let Some(provenance) = &mut self.provenance {
+            provenance.local_copy(core, destination, source, size, operation);
+        }
+    }
+
+    pub(crate) fn provenance_local_strided_copy(
+        &mut self,
+        core: usize,
+        destination: usize,
+        source: usize,
+        element_size: usize,
+        stride: usize,
+        element_count: usize,
+        operation: &'static str,
+    ) {
+        if let Some(provenance) = &mut self.provenance {
+            provenance.local_strided_copy(
+                core,
+                destination,
+                source,
+                element_size,
+                stride,
+                element_count,
+                operation,
+            );
+        }
+    }
+
+    pub(crate) fn provenance_local_transform(
+        &mut self,
+        core: usize,
+        destination: usize,
+        sources: &[(usize, usize)],
+        output_size: usize,
+        operation: &'static str,
+    ) {
+        if let Some(provenance) = &mut self.provenance {
+            provenance.local_transform(core, destination, sources, output_size, operation);
+        }
+    }
+
+    pub(crate) fn provenance_local_broadcast_transform(
+        &mut self,
+        core: usize,
+        destination: usize,
+        source: usize,
+        source_size: usize,
+        output_size: usize,
+        operation: &'static str,
+    ) {
+        if let Some(provenance) = &mut self.provenance {
+            provenance.local_broadcast_transform(
+                core,
+                destination,
+                source,
+                source_size,
+                output_size,
+                operation,
+            );
+        }
+    }
+
+    pub(crate) fn provenance_local_mvm_transform(
+        &mut self,
+        core: usize,
+        destination: usize,
+        source: usize,
+        element_size: usize,
+        output_size: usize,
+        used_rows: &[bool],
+        operation: &'static str,
+    ) {
+        if let Some(provenance) = &mut self.provenance {
+            provenance.local_mvm_transform(
+                core,
+                destination,
+                source,
+                element_size,
+                output_size,
+                used_rows,
+                operation,
+            );
+        }
+    }
+
+    pub(crate) fn provenance_send_transfer(
+        &mut self,
+        sender: usize,
+        receiver: usize,
+        source: usize,
+        destination: usize,
+        size: usize,
+    ) {
+        if let Some(provenance) = &mut self.provenance {
+            provenance.send_transfer(sender, receiver, source, destination, size);
+        }
+    }
+
+    pub(crate) fn finish_provenance(&mut self) {
+        if let Some(provenance) = &mut self.provenance {
+            provenance.flush();
+        }
+    }
+
+    pub(crate) fn provenance_schedule_stall(&mut self, core: usize, remaining_cycles: u64) {
+        if let Some(provenance) = &self.provenance {
+            provenance.schedule_stall(core, remaining_cycles);
+        }
+    }
+
+    pub(crate) fn provenance_schedule_target_stall(
+        &mut self,
+        core: usize,
+        pc: usize,
+        remaining_cycles: u64,
+    ) {
+        if let Some(provenance) = &self.provenance {
+            provenance.schedule_target_stall(core, pc, remaining_cycles);
+        }
+    }
+
+    pub(crate) fn provenance_schedule_config(&mut self, config: super::DiagnosticScheduleConfig) {
+        if let Some(provenance) = &self.provenance {
+            provenance.schedule_config(
+                match config.policy {
+                    DiagnosticSchedulePolicy::Greedy => "greedy",
+                    DiagnosticSchedulePolicy::Randomized => "randomized",
+                    DiagnosticSchedulePolicy::Adversarial => "adversarial",
+                },
+                config.seed,
+                config.target.map(|target| {
+                    json!({
+                        "writer_core": target.writer_core,
+                        "writer_pc": target.writer_pc,
+                        "reader_core": target.reader_core,
+                        "reader_pc": target.reader_pc,
+                        "address_begin": target.address_begin,
+                        "address_end": target.address_end,
+                        "reader_iteration": target.reader_iteration,
+                        "writer_min_iteration": target.writer_min_iteration,
+                    })
+                }),
+                config.deferral_budget,
+                config.fixed_stall,
+                config.fixed_target_stall,
+            );
+        }
+    }
+
+    pub(crate) fn provenance_scheduler_event(
+        &mut self,
+        event: &'static str,
+        core: usize,
+        pc: usize,
+        iteration: u32,
+        reason: &'static str,
+        selected_core: Option<usize>,
+        target: Option<DiagnosticScheduleTarget>,
+        deferrals: u64,
+    ) {
+        if let Some(provenance) = &self.provenance {
+            let target = target.map(|target| {
+                json!({
+                    "writer_core": target.writer_core,
+                    "writer_pc": target.writer_pc,
+                    "reader_core": target.reader_core,
+                    "reader_pc": target.reader_pc,
+                    "address_begin": target.address_begin,
+                    "address_end": target.address_end,
+                    "reader_iteration": target.reader_iteration,
+                    "writer_min_iteration": target.writer_min_iteration,
+                })
+            });
+            provenance.scheduler_event(json!({
+                "event": event,
+                "cycle": self.provenance_cycle(),
+                "core": core,
+                "pc": pc,
+                "core_iteration": iteration,
+                "reason": reason,
+                "selected_core": selected_core,
+                "target": target,
+                "deferrals": deferrals,
+            }));
+        }
+    }
+
+    fn provenance_cycle(&self) -> u64 {
+        self.provenance.as_ref().map_or(0, ProvenanceTracker::cycle)
+    }
+
+    pub(crate) fn set_current_iteration(&mut self, iteration: u32) {
+        if let Some(batch_outputs) = &mut self.batch_outputs {
+            batch_outputs.iteration = iteration as usize;
+        }
+    }
+
+    pub(crate) fn begin_host_store_recording(
+        &mut self,
+        batch_size: usize,
+        dump_ranges: &[(usize, usize)],
+    ) -> Result<()> {
+        let mut initial = Vec::new();
+        for &(address, size) in dump_ranges {
+            initial.extend_from_slice(self.host().load::<u8>(address, size)?[0]);
+        }
+        self.batch_outputs = Some(BatchOutputs {
+            iteration: 0,
+            ranges: dump_ranges.to_vec(),
+            outputs: vec![initial; batch_size],
+        });
+        Ok(())
+    }
+
+    pub(crate) fn store_to_host(
+        &mut self,
+        core: impl TryToUsize,
+        host_address: impl AddressArg,
+        core_address: impl AddressArg,
+        size: impl TryToUsize,
+    ) -> Result<()> {
+        let core = core.try_into().expect("core can not be negative");
+        let host_address = host_address.to_address_usize()?;
+        let core_address = core_address.to_address_usize()?;
+        let size = size.try_into().context("size can not be negative")?;
+        let Self {
+            cores,
+            batch_outputs,
+            ..
+        } = self;
+        let (host, cores) = cores.split_at_mut(1);
+        let bytes = cores[core - 1].load::<u8>(core_address, size)?[0];
+        host[0].execute_store(host_address, bytes)?;
+        if let Some(batch_outputs) = batch_outputs {
+            batch_outputs.record(host_address, bytes);
+        }
+        Ok(())
+    }
+
+    pub(crate) fn finish_host_store_recording(&mut self) -> Vec<Vec<u8>> {
+        self.batch_outputs
+            .take()
+            .map_or_else(Vec::new, |batch_outputs| batch_outputs.outputs)
+    }
+
+    pub fn host<'b>(&'b mut self) -> &'b mut Core<'a>
+    where
+        'a: 'b,
+    {
+        &mut self.cores[0]
+    }
+
+    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<'b, 'c >(&'b mut self, core: impl TryToUsize) -> (&'c mut Core<'a>, &'c mut Core<'a>) 
-    where 'a: 'b,
-          'b: 'c
+    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!(
@@ -63,8 +410,12 @@ impl<'a> CPU<'a> {
         (host, core)
     }
 
-    pub fn get_multiple_cores<'b, const N: usize>(&'b mut self, indices: [usize; N]) -> [&'b mut Core<'a>; N] 
-    where 'a : 'b
+    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()
     }
@@ -78,7 +429,7 @@ pub struct Core<'a> {
 }
 
 impl<'a> Core<'a> {
-    fn new(crossbars : Vec<&'a Crossbar>) -> Self {
+    fn new(crossbars: Vec<&'a Crossbar>) -> Self {
         Self {
             crossbars,
             memory: CoreMemory::new(),
@@ -139,7 +490,12 @@ impl<'a> Core<'a> {
         (memory, crossbars)
     }
 
-    pub fn memset(&mut self, address: impl AddressArg, size: impl TryToUsize, val: u8) -> Result<()> {
+    pub fn memset(
+        &mut self,
+        address: impl AddressArg,
+        size: impl TryToUsize,
+        val: u8,
+    ) -> Result<()> {
         let address = address.to_address_usize()?;
         let size = size.try_into().context("size can not be negative")?;
         self.memory.memset(address, size, val)

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

@@ -278,9 +278,18 @@ impl InstructionDataBuilder {
     }
 
     fn check_sanity(&self) {
-        assert!(!(self.get_r2() != 0 && self.get_imm() != 0 && self.get_mbiw() != 0 && self.get_imm_core() != 0));
-        assert!(!(self.get_ibiw() != 0 && self.get_offset_select() != 0 && self.get_imm_relu() != 0));
-        assert!(!(self.get_obiw() != 0 && self.get_offset_value() != 0 && self.get_imm_group() != 0));
+        assert!(
+            !(self.get_r2() != 0
+                && self.get_imm() != 0
+                && self.get_mbiw() != 0
+                && self.get_imm_core() != 0)
+        );
+        assert!(
+            !(self.get_ibiw() != 0 && self.get_offset_select() != 0 && self.get_imm_relu() != 0)
+        );
+        assert!(
+            !(self.get_obiw() != 0 && self.get_offset_value() != 0 && self.get_imm_group() != 0)
+        );
     }
 
     pub fn build(&mut self) -> InstructionData {

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

@@ -1,7 +1,7 @@
 use crate::{
     cpu::{CPU, crossbar},
     instruction_set::{
-        Instruction, InstructionData, InstructionStatus, InstructionType, VectorBitWith,
+        Instruction, InstructionData, InstructionStatus, InstructionType, VectorBitWidth,
         helper::add_all,
     },
     memory_manager::{
@@ -9,7 +9,7 @@ use crate::{
         type_traits::{FromFloat, UpcastDestTraits, UpcastSlice},
     },
     tracing::TRACER,
-    utility::{add_offset_r1, add_offset_r2, add_offset_rd},
+    utility::{AddressArg, add_offset_r1, add_offset_r2, add_offset_rd},
 };
 use aligned_vec::{AVec, ConstAlign};
 use anyhow::{Context, Result, ensure};
@@ -58,7 +58,7 @@ pub static NAMES: LazyLock<HashMap<usize, &'static str>> = LazyLock::new(|| {
     add_name_simd!(hash, vtanh);
     add_name_simd!(hash, vsigm);
     add_name_simd!(hash, vsoftmax);
-    add_name!(hash, vmv);
+    add_name_simd!(hash, vmv);
     add_name!(hash, vrsu);
     add_name!(hash, vrsl);
     add_name!(hash, ld);
@@ -189,6 +189,7 @@ static SIMD: LazyLock<HashMap<usize, HashMap<(usize, usize), InstructionType>>>
         add_simd_to_map!(storage, vtanh);
         add_simd_to_map!(storage, vsigm);
         add_simd_to_map!(storage, vsoftmax);
+        add_simd_to_map!(storage, vmv);
         add_simd_to_map!(storage, mvmul);
         storage
     });
@@ -199,20 +200,20 @@ pub fn isa_simd(functor: InstructionType) -> bool {
 
 pub fn dispatch_simd(
     functor: InstructionType,
-    vector_bit_with: VectorBitWith,
+    vector_bit_width: VectorBitWidth,
 ) -> Result<InstructionType> {
-    let VectorBitWith {
-        vector_input_bitwith,
-        vector_output_bitwith,
-    } = vector_bit_with;
+    let VectorBitWidth {
+        vector_input_bitwidth,
+        vector_output_bitwidth,
+    } = vector_bit_width;
     let res = SIMD
         .get(&(functor as usize))
         .context("Request a non present simd")?
-        .get(&(vector_input_bitwith, vector_output_bitwith))
+        .get(&(vector_input_bitwidth, vector_output_bitwidth))
         .with_context(|| {
             format!(
                 "Function not found for the requested size input:{} output:{}",
-                vector_input_bitwith, vector_output_bitwith
+                vector_input_bitwidth, vector_output_bitwidth
             )
         })?;
     Ok(*res)
@@ -273,7 +274,10 @@ where
         "Stored crossbar bytes do not describe an integral number of columns"
     );
     let crossbar_elem_width = crossbar_stored_bytes / bytes_per_column;
-    ensure!(crossbar_elem_width != 0, "Crossbar contains no stored columns");
+    ensure!(
+        crossbar_elem_width != 0,
+        "Crossbar contains no stored columns"
+    );
 
     let loads = memory
         .reserve_load(r1_val, crossbar_height * size_of::<F>())?
@@ -281,6 +285,10 @@ where
     let load = loads[0];
     let vec: Cow<[M]> = load.up();
     let matrix = crossbar.load::<M>(crossbar_stored_bytes)?[0];
+    let used_rows: Vec<bool> = matrix
+        .chunks_exact(crossbar_elem_width)
+        .map(|row| row.iter().any(|value| *value != M::from_f32(0.0)))
+        .collect();
 
     // --- FAER IMPLEMENTATION ---
 
@@ -319,6 +327,16 @@ where
 
     let res_up: Cow<[T]> = res.as_slice().up();
     core.execute_store(rd_val, res_up.as_ref());
+    let _ = core;
+    cores.provenance_local_mvm_transform(
+        core_indx as usize,
+        rd_val as usize,
+        r1_val as usize,
+        size_of::<F>(),
+        res_up.len() * size_of::<T>(),
+        &used_rows,
+        "mvmul",
+    );
 
     TRACER.lock().unwrap().post_mvm::<F, M, T>(cores, data);
     Ok(InstructionStatus::Completed)
@@ -385,6 +403,14 @@ where
     );
     let res_up: Cow<[T]> = res.as_slice().up();
     core.execute_store(rd_val, res_up.as_ref());
+    let _ = core;
+    cores.provenance_local_transform(
+        core_indx as usize,
+        rd_val,
+        &[(r1_val, byte_len), (r2_val, byte_len)],
+        byte_len,
+        "vvadd",
+    );
     TRACER.lock().unwrap().post_vvadd::<F, T>(cores, data);
     Ok(InstructionStatus::Completed)
 }
@@ -470,6 +496,13 @@ where
     );
     let res_up: Cow<[T]> = res.as_slice().up();
     core.execute_store(rd_val, res_up.as_ref());
+    cores.provenance_local_transform(
+        core_indx as usize,
+        rd_val,
+        &[(r1_val, byte_len), (r2_val, byte_len)],
+        byte_len,
+        "vvmul",
+    );
     Ok(InstructionStatus::Completed)
 }
 
@@ -742,7 +775,50 @@ where
 
 #[inline(never)]
 pub fn vmv(cores: &mut CPU, data: InstructionData) -> Result<InstructionStatus> {
-    todo!()
+    panic!("You are calling a placeholder, the real call is the generic version");
+}
+
+#[inline(never)]
+pub(super) fn vmv_impl<F, T>(cores: &mut CPU, data: InstructionData) -> Result<InstructionStatus>
+where
+    F: Copy + MemoryStorable,
+{
+    let (core_indx, rd, r1, r2, imm_len, _, _) = data.get_core_rd_r1_r2_immlen_offset();
+    let core = cores.core(core_indx);
+    let source = core.register(r1).to_address_usize()?;
+    let destination = core.register(rd);
+    let stride: usize = core
+        .register(r2)
+        .try_into()
+        .context("vmv stride can not be negative")?;
+    let (element_count, _) = vector_lengths::<F>(imm_len)?;
+    let stride_bytes = stride
+        .checked_mul(size_of::<F>())
+        .context("vmv byte stride overflow")?;
+    let mut result = Vec::with_capacity(element_count);
+    for index in 0..element_count {
+        let offset = index
+            .checked_mul(stride_bytes)
+            .context("vmv source offset overflow")?;
+        let address = source
+            .checked_add(offset)
+            .context("vmv source address overflow")?;
+        result.push(
+            core.reserve_load(address, size_of::<F>())?
+                .execute_load::<F>()?[0][0],
+        );
+    }
+    core.execute_store(destination, &result)?;
+    cores.provenance_local_strided_copy(
+        core_indx as usize,
+        destination as usize,
+        source,
+        size_of::<F>(),
+        stride,
+        element_count,
+        "vmv",
+    );
+    Ok(InstructionStatus::Completed)
 }
 
 #[inline(never)]
@@ -761,16 +837,23 @@ pub fn vrsl(cores: &mut CPU, data: InstructionData) -> Result<InstructionStatus>
 #[inline(never)]
 pub fn ld(cores: &mut CPU, data: InstructionData) -> Result<InstructionStatus> {
     TRACER.lock().unwrap().pre_ld(cores, data);
-    let (core, rd, r1, _, imm_len, offset_select, offset_value) =
+    let (core_index, rd, r1, _, imm_len, offset_select, offset_value) =
         data.get_core_rd_r1_r2_immlen_offset();
-    ensure!(core != 0, "LD cannot be used to move from host to host");
-    let (host, core) = cores.host_and_cores(core);
-    let r1_val = core.register(r1);
-    let rd_val = core.register(rd);
-    let r1_val = add_offset_r1(r1_val, offset_select, offset_value);
-    let rd_val = add_offset_rd(rd_val, offset_select, offset_value);
-    let global_memory = host.load::<u8>(r1_val, imm_len)?;
-    core.execute_store(rd_val, global_memory[0])?;
+    ensure!(
+        core_index != 0,
+        "LD cannot be used to move from host to host"
+    );
+    let (r1_val, rd_val) = {
+        let (host, core) = cores.host_and_cores(core_index);
+        let r1_val = core.register(r1);
+        let rd_val = core.register(rd);
+        let r1_val = add_offset_r1(r1_val, offset_select, offset_value);
+        let rd_val = add_offset_rd(rd_val, offset_select, offset_value);
+        let global_memory = host.load::<u8>(r1_val, imm_len)?;
+        core.execute_store(rd_val, global_memory[0])?;
+        (r1_val, rd_val)
+    };
+    cores.provenance_global_load_to_local(core_index as usize, r1_val, rd_val, imm_len as usize);
     TRACER.lock().unwrap().post_ld(cores, data);
     Ok(InstructionStatus::Completed)
 }
@@ -781,13 +864,16 @@ pub fn st(cores: &mut CPU, data: InstructionData) -> Result<InstructionStatus> {
     let (core, rd, r1, _, imm_len, offset_select, offset_value) =
         data.get_core_rd_r1_r2_immlen_offset();
     ensure!(core != 0, "ST cannot be used to move from host to host");
-    let (host, core) = cores.host_and_cores(core);
-    let r1_val = core.register(r1);
-    let rd_val = core.register(rd);
-    let r1_val = add_offset_r1(r1_val, offset_select, offset_value);
-    let rd_val = add_offset_rd(rd_val, offset_select, offset_value);
-    let local_memory = core.load::<u8>(r1_val, imm_len)?;
-    host.execute_store(rd_val, local_memory[0]);
+    let (rd_val, r1_val) = {
+        let core = cores.core(core);
+        let r1_val = core.register(r1);
+        let rd_val = core.register(rd);
+        let r1_val = add_offset_r1(r1_val, offset_select, offset_value);
+        let rd_val = add_offset_rd(rd_val, offset_select, offset_value);
+        (rd_val, r1_val)
+    };
+    cores.store_to_host(core, rd_val, r1_val, imm_len)?;
+    cores.provenance_global_store_from_local(core as usize, rd_val, r1_val, imm_len as usize);
     TRACER.lock().unwrap().post_st(cores, data);
     Ok(InstructionStatus::Completed)
 }
@@ -812,9 +898,9 @@ pub fn lldi(cores: &mut CPU, data: InstructionData) -> Result<InstructionStatus>
 #[inline(never)]
 pub fn lmv(cores: &mut CPU, data: InstructionData) -> Result<InstructionStatus> {
     TRACER.lock().unwrap().pre_lmv(cores, data);
-    let (core, rd, r1, _, imm_len, offset_select, offset_value) =
+    let (core_index, rd, r1, _, imm_len, offset_select, offset_value) =
         data.get_core_rd_r1_r2_immlen_offset();
-    let core = cores.core(core);
+    let core = cores.core(core_index);
     let r1_val = core.register(r1);
     let rd_val = core.register(rd);
     let r1_val = add_offset_r1(r1_val, offset_select, offset_value);
@@ -822,12 +908,14 @@ pub fn lmv(cores: &mut CPU, data: InstructionData) -> Result<InstructionStatus>
     let local_memory = core.load::<u8>(r1_val, imm_len)?;
     let tmp = local_memory[0].to_vec();
     core.execute_store(rd_val, tmp.as_slice());
+    let _ = core;
+    cores.provenance_local_copy(core_index as usize, rd_val, r1_val, imm_len as usize, "lmv");
     TRACER.lock().unwrap().post_lmv(cores, data);
     Ok(InstructionStatus::Completed)
 }
 
 #[inline(never)]
-pub fn isa_send(functor : usize) -> bool{
+pub fn isa_send(functor: usize) -> bool {
     (send as *const () as usize) == functor
 }
 
@@ -837,13 +925,13 @@ pub fn send(cores: &mut CPU, data: InstructionData) -> Result<InstructionStatus>
 }
 
 #[inline(never)]
-pub fn isa_recv(functor : usize) -> bool{
+pub fn isa_recv(functor: usize) -> bool {
     (recv as *const () as usize) == functor
 }
 
 #[inline(never)]
 pub fn recv(cores: &mut CPU, data: InstructionData) -> Result<InstructionStatus> {
-    Ok(InstructionStatus::Reciving(data))
+    Ok(InstructionStatus::Receiving(data))
 }
 
 #[inline(never)]

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

@@ -7,9 +7,9 @@ use crate::{
 };
 use anyhow::{Context, Result};
 use std::mem::swap;
+pub mod helper;
 pub mod instruction_data;
 pub mod isa;
-pub mod helper;
 
 #[derive(Clone, Copy, Debug)]
 pub struct Instruction {
@@ -22,7 +22,7 @@ pub enum InstructionStatus {
     Completed,
     Waiting(InstructionData),
     Sending(InstructionData),
-    Reciving(InstructionData),
+    Receiving(InstructionData),
     Sync(InstructionData),
     #[default]
     NotExecuted,
@@ -41,7 +41,8 @@ impl Instruction {
     }
 
     pub fn execute<'a, 'b>(&'b self, cpu: &mut CPU<'a>) -> InstructionStatus
-    where 'a : 'b
+    where
+        'a: 'b,
     {
         (self.functor)(cpu, self.data)
             .with_context(|| format!("Instruction: {}", functor_to_name(self.functor as usize)))
@@ -58,21 +59,21 @@ pub type Instructions = Vec<Instruction>;
 pub type InstructionType = fn(&mut CPU, InstructionData) -> Result<InstructionStatus>;
 
 #[derive(Debug, Clone, Copy, Default)]
-pub struct VectorBitWith {
-    pub vector_input_bitwith: usize,
-    pub vector_output_bitwith: usize,
+pub struct VectorBitWidth {
+    pub vector_input_bitwidth: usize,
+    pub vector_output_bitwidth: usize,
 }
 
 /// Support for the
 /// setbw ibiw, obiw
 /// Set the bit-widths of each element for input vectors and output vectors. Related vector instructions
-/// use the configured bit-widths. Once setbw is caled, all subsequent related vector instructions will
+/// use the configured bit-widths. Once setbw is called, all subsequent related vector instructions will
 /// use the configured bit-widths, until a new setbw is called. Once ibiw and obiw are set, ibyw and
 /// obyw are also set accordingly by the hardware.
 /// If the hardware does not support variable bit-width, this instruction is invalid and the matrix/vector
 /// instructions use the fixed bit-width of the hardware.
 pub struct InstructionsBuilder {
-    vector_bit_with: VectorBitWith,
+    vector_bit_width: VectorBitWidth,
     instructions: Instructions,
 }
 
@@ -85,9 +86,9 @@ impl Default for InstructionsBuilder {
 impl InstructionsBuilder {
     pub fn new() -> Self {
         Self {
-            vector_bit_with: VectorBitWith {
-                vector_input_bitwith: 32,
-                vector_output_bitwith: 32,
+            vector_bit_width: VectorBitWidth {
+                vector_input_bitwidth: 32,
+                vector_output_bitwidth: 32,
             },
             instructions: Instructions::new(),
         }
@@ -96,9 +97,9 @@ impl InstructionsBuilder {
     pub fn make_inst(&mut self, functor: InstructionType, data: InstructionData) {
         if is_setbw(functor) {
             let (ibiw, obiw) = data.get_ibiw_obiw();
-            self.vector_bit_with.vector_input_bitwith =
+            self.vector_bit_width.vector_input_bitwidth =
                 ibiw.try_into().expect("ibiw can not be negative");
-            self.vector_bit_with.vector_output_bitwith =
+            self.vector_bit_width.vector_output_bitwidth =
                 obiw.try_into().expect("obiw can not be negative");
             return;
         }
@@ -106,7 +107,7 @@ impl InstructionsBuilder {
         if (isa_simd(functor)) {
             self.instructions.push(Instruction::new(
                 data,
-                dispatch_simd(functor, self.vector_bit_with).unwrap(),
+                dispatch_simd(functor, self.vector_bit_width).unwrap(),
             ))
         } else {
             self.instructions.push(Instruction::new(data, functor))

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

@@ -601,7 +601,11 @@ fn json_to_wait(
     inst_data_builder: &mut InstructionDataBuilder,
     json: &Value,
 ) -> Result<()> {
-    todo!("Not present in the compiler");
+    inst_data_builder.set_offset_select_value(
+        json_i64!(json, "event_register") as i32,
+        json_i64!(json, "wait_value") as i32,
+    );
+    inst_builder.make_inst(wait, inst_data_builder.build());
     Ok(())
 }
 
@@ -610,7 +614,10 @@ fn json_to_sync(
     inst_data_builder: &mut InstructionDataBuilder,
     json: &Value,
 ) -> Result<()> {
-    todo!("Not present in the compiler");
+    inst_data_builder
+        .set_imm_core(json_i64!(json, "core") as i32 + 1)
+        .set_offset_select_value(json_i64!(json, "event_register") as i32, 0);
+    inst_builder.make_inst(sync, inst_data_builder.build());
     Ok(())
 }
 

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

@@ -78,7 +78,11 @@ impl CoreMemory {
         }
     }
 
-    pub fn reserve_load(&mut self, address: impl TryToUsize, size: impl TryToUsize) -> Result<&mut Self>
+    pub fn reserve_load(
+        &mut self,
+        address: impl TryToUsize,
+        size: impl TryToUsize,
+    ) -> Result<&mut Self>
 where {
         let address = address.try_into().context("address can not be negative")?;
         let size = size.try_into().context("size can not be negative")?;
@@ -87,7 +91,8 @@ where {
             size,
         };
         if self.memory.len() < address + size {
-            self.memory.resize(min((address + size) * 2, u32::MAX as usize), 0);
+            self.memory
+                .resize(min((address + size) * 2, u32::MAX as usize), 0);
         }
         self.load_requests.push(load_request);
         Ok(self)
@@ -105,14 +110,24 @@ where {
         for (load_index, load_request) in load_requests.drain(..).enumerate() {
             let LoadRequest { index, size } = load_request;
             let memory_slice = &memory[index..index + size];
-            let memory_slice = unsafe { slice_from_u8(memory_slice) }
-                .with_context(|| format!("Load number: {} Accessing from {} to {}", load_index, index, index + size))?;
+            let memory_slice = unsafe { slice_from_u8(memory_slice) }.with_context(|| {
+                format!(
+                    "Load number: {} Accessing from {} to {}",
+                    load_index,
+                    index,
+                    index + size
+                )
+            })?;
             res.push(memory_slice);
         }
         Ok(res)
     }
 
-    pub fn load_const<T>(&self, address: impl TryToUsize, size: impl TryToUsize) -> Result<Vec<&[T]>>
+    pub fn load_const<T>(
+        &self,
+        address: impl TryToUsize,
+        size: impl TryToUsize,
+    ) -> Result<Vec<&[T]>>
     where
         T: MemoryStorable,
     {
@@ -125,7 +140,7 @@ where {
         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))?;
+            .with_context(|| format!("Accessing from {} to {}", address, address + size))?;
         res.push(memory_slice);
         Ok(res)
     }
@@ -155,7 +170,12 @@ where {
         Ok(())
     }
 
-    pub fn memset(&mut self, address: impl TryToUsize, size: impl TryToUsize, val: u8) -> Result<()> {
+    pub fn memset(
+        &mut self,
+        address: impl TryToUsize,
+        size: impl TryToUsize,
+        val: u8,
+    ) -> Result<()> {
         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, .. } = self;
@@ -174,11 +194,11 @@ where {
         }
     }
 
-    pub fn get_len(&self) ->usize {
+    pub fn get_len(&self) -> usize {
         self.memory.len()
     }
 
-    pub(crate) fn clear(&mut self)  {
+    pub(crate) fn clear(&mut self) {
         self.memory.clear();
     }
 }
@@ -257,7 +277,11 @@ mod test {
         let mut data = [0_f32; 2];
         data[0] = loads[0][0] + loads[1][0];
         core_memory.execute_store(4, &data[0..1]).unwrap();
-        let loads: &[f32] = core_memory.reserve_load(0, 16).unwrap().execute_load().unwrap()[0];
+        let loads: &[f32] = core_memory
+            .reserve_load(0, 16)
+            .unwrap()
+            .execute_load()
+            .unwrap()[0];
         println!("{:?}", loads);
         assert!(loads[0] == 5_f32 && loads[1] == 12_f32 && loads[2] == 7_f32 && loads[3] == 15_f32)
     }

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

@@ -1,5 +1,3 @@
-
-
 use std::{
     borrow::Cow,
     fmt::Debug,
@@ -9,32 +7,32 @@ use std::{
 use anyhow::Context;
 
 pub trait FromFloat {
-    fn from_f32(val :f32) -> Self;
-    fn from_f64(val :f64) -> Self;
+    fn from_f32(val: f32) -> Self;
+    fn from_f64(val: f64) -> Self;
 }
 
 impl FromFloat for f32 {
-    fn from_f32(val :f32) -> Self {
+    fn from_f32(val: f32) -> Self {
         val
     }
 
-    fn from_f64(val :f64) -> Self {
+    fn from_f64(val: f64) -> Self {
         val as f32
     }
 }
 
 impl FromFloat for f64 {
-    fn from_f32(val :f32) -> Self {
+    fn from_f32(val: f32) -> Self {
         val as f64
     }
 
-    fn from_f64(val :f64) -> Self {
-        val 
+    fn from_f64(val: f64) -> Self {
+        val
     }
 }
 
 pub trait HasTanh {
-    fn tanh(self) -> Self ;
+    fn tanh(self) -> Self;
 }
 
 impl HasTanh for f32 {
@@ -50,7 +48,7 @@ impl HasTanh for f64 {
 }
 
 pub trait HasSigm {
-    fn sigm(self) -> Self ;
+    fn sigm(self) -> Self;
 }
 
 impl HasSigm for f32 {
@@ -91,34 +89,36 @@ impl HasExp for f64 {
     }
 }
 
-
-
-pub trait TryToUsize: TryInto<usize, Error = Self::TryError> 
-where std::result::Result<usize, Self::TryError> : Context<usize, Self::TryError>
+pub trait TryToUsize: TryInto<usize, Error = Self::TryError>
+where
+    std::result::Result<usize, Self::TryError>: Context<usize, Self::TryError>,
 {
     type TryError: Debug + Send + Sync + 'static + std::error::Error;
 }
 
-impl<T, E> TryToUsize for T 
-where 
+impl<T, E> TryToUsize for T
+where
     T: TryInto<usize, Error = E>,
     E: Debug + Send + Sync + 'static + std::error::Error,
-    std::result::Result<usize, E> : Context<usize, E>
+    std::result::Result<usize, E>: Context<usize, E>,
 {
     type TryError = E;
 }
 
-
 pub trait FromUsize {
     fn from_usize(v: usize) -> Self;
 }
 
 impl FromUsize for f32 {
-    fn from_usize(v: usize) -> Self { v as f32 }
+    fn from_usize(v: usize) -> Self {
+        v as f32
+    }
 }
 
 impl FromUsize for f64 {
-    fn from_usize(v: usize) -> Self { v as f64 }
+    fn from_usize(v: usize) -> Self {
+        v as f64
+    }
 }
 
 pub trait UpcastDestTraits<T>:

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

@@ -1,8 +1,14 @@
 #![allow(unused)]
 
-use anyhow::{Result, bail};
+use anyhow::{Context, Result, bail};
+use serde_json::json;
 use std::{
     collections::{HashMap, HashSet},
+    path::Path,
+    sync::{
+        Mutex,
+        atomic::{AtomicU32, Ordering},
+    },
     time::{Duration, SystemTime},
 };
 
@@ -21,10 +27,14 @@ pub mod cpu;
 pub mod instruction_set;
 pub mod json_to_instruction;
 pub mod memory_manager;
+pub mod provenance;
 pub mod send_recv;
 pub mod tracing;
 pub mod utility;
 
+static GLOBAL_ITERATION: AtomicU32 = AtomicU32::new(0);
+static EXECUTION_LOCK: Mutex<()> = Mutex::new(());
+
 #[derive(Debug, Clone)]
 pub struct CoreInstructionsBuilder {
     core_instructions: Vec<CoreInstructions>,
@@ -54,6 +64,7 @@ impl CoreInstructionsBuilder {
 pub struct CoreInstructions {
     instructions: Instructions,
     program_counter: usize,
+    current_iteration: u32,
 }
 
 impl CoreInstructions {
@@ -61,6 +72,7 @@ impl CoreInstructions {
         Self {
             instructions,
             program_counter,
+            current_iteration: 0,
         }
     }
 
@@ -68,6 +80,7 @@ impl CoreInstructions {
         Self {
             instructions: Vec::new(),
             program_counter: 0,
+            current_iteration: 0,
         }
     }
 }
@@ -77,15 +90,320 @@ impl From<Instructions> for CoreInstructions {
         CoreInstructions {
             instructions: value,
             program_counter: 0,
+            current_iteration: 0,
         }
     }
 }
 
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
+pub enum DiagnosticSchedulePolicy {
+    #[default]
+    Greedy,
+    Randomized,
+    Adversarial,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct DiagnosticScheduleTarget {
+    pub writer_core: usize,
+    pub writer_pc: usize,
+    pub reader_core: usize,
+    pub reader_pc: usize,
+    pub address_begin: usize,
+    pub address_end: usize,
+    pub reader_iteration: Option<u32>,
+    pub writer_min_iteration: Option<u32>,
+}
+
+#[derive(Clone, Copy, Debug)]
+struct DiagnosticScheduleConfig {
+    policy: DiagnosticSchedulePolicy,
+    seed: u64,
+    target: Option<DiagnosticScheduleTarget>,
+    deferral_budget: u64,
+    fixed_stall: Option<(usize, u64)>,
+    fixed_target_stall: Option<(usize, usize, Option<u32>, u64)>,
+}
+
+impl Default for DiagnosticScheduleConfig {
+    fn default() -> Self {
+        Self {
+            policy: DiagnosticSchedulePolicy::Greedy,
+            seed: 0,
+            target: None,
+            deferral_budget: 10_000,
+            fixed_stall: None,
+            fixed_target_stall: None,
+        }
+    }
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+struct ScheduleCoreState {
+    program_counter: usize,
+    instruction_count: usize,
+    current_iteration: u32,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+enum ScheduleAction {
+    Execute,
+    Defer {
+        next_core: usize,
+        reason: &'static str,
+    },
+    Stall {
+        remaining: u64,
+        target: bool,
+    },
+    Force {
+        reason: &'static str,
+    },
+}
+
+#[derive(Clone, Copy, Debug)]
+struct ScheduleChoice {
+    core: usize,
+    reason: &'static str,
+}
+
+#[derive(Debug)]
+struct DiagnosticScheduler {
+    config: DiagnosticScheduleConfig,
+    random_state: u64,
+    deferrals: u64,
+    writer_store_iteration: Option<u32>,
+    last_writer_state: Option<ScheduleCoreState>,
+    writer_stagnation: u64,
+}
+
+impl DiagnosticScheduler {
+    fn new(config: DiagnosticScheduleConfig) -> Self {
+        let random_state = if config.seed == 0 {
+            0x9e37_79b9_7f4a_7c15
+        } else {
+            config.seed
+        };
+        Self {
+            config,
+            random_state,
+            deferrals: 0,
+            writer_store_iteration: None,
+            last_writer_state: None,
+            writer_stagnation: 0,
+        }
+    }
+
+    fn state(cores: &[CoreInstructions]) -> Vec<ScheduleCoreState> {
+        cores
+            .iter()
+            .map(|core| ScheduleCoreState {
+                program_counter: core.program_counter,
+                instruction_count: core.instructions.len(),
+                current_iteration: core.current_iteration,
+            })
+            .collect()
+    }
+
+    fn has_work(core: ScheduleCoreState, batch_size: u32) -> bool {
+        !((core.instruction_count == 0)
+            || (core.program_counter == core.instruction_count
+                && core.current_iteration + 1 >= batch_size))
+    }
+
+    fn candidates(states: &[ScheduleCoreState], current: usize, batch_size: u32) -> Vec<usize> {
+        states
+            .iter()
+            .enumerate()
+            .filter_map(|(index, &core)| {
+                (index != current && Self::has_work(core, batch_size)).then_some(index)
+            })
+            .collect()
+    }
+
+    fn next_random(&mut self) -> u64 {
+        let mut value = self.random_state;
+        value ^= value << 13;
+        value ^= value >> 7;
+        value ^= value << 17;
+        self.random_state = value;
+        value
+    }
+
+    fn random_choice(&mut self, candidates: &[usize]) -> Option<ScheduleChoice> {
+        (!candidates.is_empty()).then(|| ScheduleChoice {
+            core: candidates[(self.next_random() as usize) % candidates.len()],
+            reason: "randomized_ready_core",
+        })
+    }
+
+    fn target_reader(&self, core: usize, state: ScheduleCoreState) -> bool {
+        self.config.target.is_some_and(|target| {
+            target.reader_core == core
+                && target.reader_pc == state.program_counter
+                && target
+                    .reader_iteration
+                    .is_none_or(|iteration| iteration == state.current_iteration)
+        })
+    }
+
+    fn target_writer_has_work(&self, states: &[ScheduleCoreState], batch_size: u32) -> bool {
+        self.config.target.is_some_and(|target| {
+            states
+                .get(target.writer_core)
+                .is_some_and(|&state| Self::has_work(state, batch_size))
+        })
+    }
+
+    fn reader_may_execute(&self, reader_iteration: u32) -> bool {
+        let minimum = self
+            .config
+            .target
+            .and_then(|target| target.writer_min_iteration)
+            .unwrap_or(reader_iteration + 1);
+        self.writer_store_iteration
+            .is_some_and(|iteration| iteration >= minimum)
+    }
+
+    fn before_instruction(
+        &mut self,
+        current: usize,
+        states: &[ScheduleCoreState],
+        batch_size: u32,
+    ) -> ScheduleAction {
+        if let Some((core, pc, iteration, remaining)) = self.config.fixed_target_stall
+            && core == current
+            && states.get(current).is_some_and(|state| {
+                state.program_counter == pc
+                    && iteration.is_none_or(|iteration| state.current_iteration == iteration)
+            })
+            && remaining > 0
+        {
+            self.config.fixed_target_stall = Some((core, pc, iteration, remaining - 1));
+            return ScheduleAction::Stall {
+                remaining: remaining - 1,
+                target: true,
+            };
+        }
+        if let Some((core, remaining)) = self.config.fixed_stall
+            && core == current
+            && remaining > 0
+        {
+            self.config.fixed_stall = Some((core, remaining - 1));
+            return ScheduleAction::Stall {
+                remaining: remaining - 1,
+                target: false,
+            };
+        }
+
+        let Some(state) = states.get(current).copied() else {
+            return ScheduleAction::Execute;
+        };
+        let candidates = Self::candidates(states, current, batch_size);
+        match self.config.policy {
+            DiagnosticSchedulePolicy::Greedy => ScheduleAction::Execute,
+            // Random choices are made after a blocking/ready boundary in
+            // `after_block`.  Deferring here would allow two ready cores to
+            // defer each other forever without executing an instruction.
+            DiagnosticSchedulePolicy::Randomized => ScheduleAction::Execute,
+            DiagnosticSchedulePolicy::Adversarial => {
+                let Some(target) = self.config.target else {
+                    return ScheduleAction::Execute;
+                };
+                if !self.target_reader(current, state)
+                    || self.reader_may_execute(state.current_iteration)
+                {
+                    return ScheduleAction::Execute;
+                }
+                if self.deferrals >= self.config.deferral_budget {
+                    return ScheduleAction::Force {
+                        reason: "DEFERRAL_LIMIT_REACHED",
+                    };
+                }
+                let writer_state = states.get(target.writer_core).copied();
+                let next_core = if Self::has_work(writer_state.unwrap_or(state), batch_size) {
+                    if self.last_writer_state == writer_state {
+                        self.writer_stagnation += 1;
+                    } else {
+                        self.last_writer_state = writer_state;
+                        self.writer_stagnation = 0;
+                    }
+                    if self.writer_stagnation >= 256 {
+                        return ScheduleAction::Force {
+                            reason: "PRODUCER_BLOCKED_BY_REAL_DEPENDENCY",
+                        };
+                    }
+                    target.writer_core
+                } else {
+                    candidates.first().copied().unwrap_or(current)
+                };
+                if next_core == current {
+                    ScheduleAction::Force {
+                        reason: "NO_ALTERNATIVE_READY_EVENT",
+                    }
+                } else {
+                    self.deferrals += 1;
+                    ScheduleAction::Defer {
+                        next_core,
+                        reason: "target_consumer",
+                    }
+                }
+            }
+        }
+    }
+
+    fn after_block(
+        &mut self,
+        current: usize,
+        states: &[ScheduleCoreState],
+        batch_size: u32,
+    ) -> Option<ScheduleChoice> {
+        let candidates = Self::candidates(states, current, batch_size);
+        match self.config.policy {
+            DiagnosticSchedulePolicy::Greedy => None,
+            DiagnosticSchedulePolicy::Randomized => self.random_choice(&candidates),
+            DiagnosticSchedulePolicy::Adversarial => {
+                let target = self.config.target?;
+                if self.target_writer_has_work(states, batch_size)
+                    && target.writer_core != current
+                    && candidates.contains(&target.writer_core)
+                {
+                    Some(ScheduleChoice {
+                        core: target.writer_core,
+                        reason: "target_producer",
+                    })
+                } else {
+                    candidates.first().copied().map(|core| ScheduleChoice {
+                        core,
+                        reason: "adversarial_ready_core",
+                    })
+                }
+            }
+        }
+    }
+
+    fn note_completed(&mut self, core: usize, pc: usize, iteration: u32) {
+        if self
+            .config
+            .target
+            .is_some_and(|target| target.writer_core == core && target.writer_pc == pc)
+        {
+            self.writer_store_iteration = Some(iteration);
+        }
+    }
+
+    fn config(&self) -> DiagnosticScheduleConfig {
+        self.config
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct Executable<'a> {
     cpu: CPU<'a>,
     core_instructions: Vec<CoreInstructions>,
     send_recv: SendRecv,
+    provenance_global_barrier: bool,
+    diagnostic_schedule: DiagnosticScheduleConfig,
 }
 
 struct DeadlockInfo {
@@ -93,6 +411,8 @@ struct DeadlockInfo {
     states: String,
 }
 
+type SyncEvents = Vec<[i32; 32]>;
+
 fn print_status(core_instructions: &[CoreInstructions]) {
     let mut tot_instructions = 0;
     let mut progress = 0;
@@ -121,60 +441,321 @@ impl<'a> Executable<'a> {
             cpu,
             core_instructions,
             send_recv,
+            provenance_global_barrier: false,
+            diagnostic_schedule: DiagnosticScheduleConfig::default(),
         }
     }
 
+    pub fn enable_provenance(&mut self, path: impl AsRef<Path>) -> Result<()> {
+        self.cpu
+            .enable_provenance(path)
+            .context("cannot enable provenance tracing")?;
+        Ok(())
+    }
+
+    pub fn set_provenance_global_barrier(&mut self, enabled: bool) {
+        self.provenance_global_barrier = enabled;
+    }
+
+    pub fn set_provenance_core_stall(&mut self, core: usize, cycles: u64) {
+        self.diagnostic_schedule.fixed_stall = Some((core, cycles));
+    }
+
+    pub fn set_diagnostic_target_stall(
+        &mut self,
+        core: usize,
+        pc: usize,
+        iteration: Option<u32>,
+        cycles: u64,
+    ) {
+        self.diagnostic_schedule.fixed_target_stall = Some((core, pc, iteration, cycles));
+    }
+
+    pub fn set_diagnostic_schedule_policy(&mut self, policy: DiagnosticSchedulePolicy) {
+        self.diagnostic_schedule.policy = policy;
+    }
+
+    pub fn set_diagnostic_schedule_seed(&mut self, seed: u64) {
+        self.diagnostic_schedule.seed = seed;
+    }
+
+    pub fn set_diagnostic_schedule_target(&mut self, target: DiagnosticScheduleTarget) {
+        self.diagnostic_schedule.target = Some(target);
+    }
+
+    pub fn set_diagnostic_schedule_deferral_budget(&mut self, budget: u64) {
+        self.diagnostic_schedule.deferral_budget = budget;
+    }
+
     pub fn execute<'b>(&'b mut self) -> Result<()>
     where
         'a: 'b,
     {
+        self.execute_batch(&[&[]], &[], &[]).map(|_| ())
+    }
+
+    pub fn execute_batch<'b>(
+        &'b mut self,
+        inputs: &[&[u8]],
+        input_regions: &[(usize, usize)],
+        dump_ranges: &[(usize, usize)],
+    ) -> Result<Vec<Vec<u8>>>
+    where
+        'a: 'b,
+    {
+        validate_inputs(inputs, input_regions)?;
+        self.execute_iterations(inputs, input_regions, dump_ranges)
+    }
+
+    fn execute_iterations<'b>(
+        &'b mut self,
+        inputs: &[&[u8]],
+        input_regions: &[(usize, usize)],
+        dump_ranges: &[(usize, usize)],
+    ) -> Result<Vec<Vec<u8>>>
+    where
+        'a: 'b,
+    {
+        let _execution_lock = EXECUTION_LOCK.lock().unwrap();
+        let batch_size = u32::try_from(inputs.len().max(1)).context("batch size exceeds u32")?;
+        GLOBAL_ITERATION.store(0, Ordering::SeqCst);
+        self.cpu.begin_provenance_batch(batch_size as usize);
+        if let Some(input) = inputs.first() {
+            store_input(&mut self.cpu, input, input_regions, 0)?;
+        }
+        self.cpu
+            .begin_host_store_recording(batch_size as usize, dump_ranges)?;
+
+        let provenance_global_barrier = self.provenance_global_barrier;
+        let mut scheduler = DiagnosticScheduler::new(self.diagnostic_schedule);
+        self.cpu.provenance_schedule_config(scheduler.config());
         let Self {
             cpu,
             core_instructions: cores_instructions,
             send_recv,
+            ..
         } = self;
+        let active_cores: Vec<usize> = cores_instructions
+            .iter()
+            .enumerate()
+            .filter_map(|(index, core)| (!core.instructions.is_empty()).then_some(index))
+            .collect();
+        let mut barrier_iteration = None;
         let mut cpu_progressed = 0;
         let max_core = cpu.num_core();
+        let mut sync_events: SyncEvents = vec![[0; 32]; max_core];
         let mut cpu_index = 0;
+        let mut cycle = 0;
+        let mut scheduler_no_progress = 0usize;
+        let scheduler_no_progress_limit = max_core.saturating_mul(4).max(8);
         let mut now = SystemTime::now();
 
         while (cpu_progressed > -2) {
             let mut core_result = InstructionStatus::Completed;
-            while core_result.is_completed()
-                && let Some(core_instruction) = cores_instructions.get_mut(cpu_index)
+            let mut scheduler_next = None;
+            if provenance_global_barrier
+                && barrier_iteration.is_some()
+                && active_cores.iter().all(|&index| {
+                    cores_instructions[index].current_iteration >= barrier_iteration.unwrap()
+                })
             {
-                core_result = InstructionStatus::NotExecuted;
-                let CoreInstructions {
-                    instructions,
-                    program_counter,
-                } = core_instruction;
-                core_result = instructions
-                    .get(*program_counter)
-                    .map_or(InstructionStatus::default(), |inst: &Instruction| {
-                        inst.execute(cpu)
-                    });
-                if core_result.is_completed() {
-                    cpu_progressed = 0;
-                    *program_counter += 1;
+                barrier_iteration = None;
+            }
+            while core_result.is_completed() {
+                let barrier_ready = if provenance_global_barrier {
+                    let current_iteration = cores_instructions[cpu_index].current_iteration;
+                    active_cores.iter().all(|&index| {
+                        let core = &cores_instructions[index];
+                        core.program_counter == core.instructions.len()
+                            && core.current_iteration == current_iteration
+                    })
+                } else {
+                    false
+                };
+                let current_pc = cores_instructions[cpu_index].program_counter;
+                let current_iteration = cores_instructions[cpu_index].current_iteration;
+                let states = if scheduler.config().policy == DiagnosticSchedulePolicy::Greedy
+                    && scheduler.config().fixed_stall.is_none()
+                    && scheduler.config().fixed_target_stall.is_none()
+                {
+                    None
+                } else {
+                    Some(DiagnosticScheduler::state(cores_instructions))
+                };
+                let schedule_action = states.as_deref().map_or(ScheduleAction::Execute, |states| {
+                    scheduler.before_instruction(cpu_index, states, batch_size)
+                });
+                if let ScheduleAction::Defer { next_core, reason } = schedule_action {
+                    cpu.set_execution_context(cycle, cpu_index, current_pc, current_iteration);
+                    cpu.provenance_scheduler_event(
+                        "scheduler_defer",
+                        cpu_index,
+                        current_pc,
+                        current_iteration,
+                        reason,
+                        Some(next_core),
+                        scheduler.config().target,
+                        scheduler.deferrals,
+                    );
+                    scheduler_next = Some(next_core);
+                    break;
                 }
-                if (now.elapsed().unwrap() > Duration::from_secs(5)) {
-                    print_status(cores_instructions);
-                    if let Some(deadlock) = detect_deadlock(cores_instructions) {
-                        bail!(
-                            "Deadlock cycle detected: {} [{}]",
-                            deadlock.cycle,
-                            deadlock.states
-                        );
+                if let ScheduleAction::Stall { remaining, target } = schedule_action {
+                    cpu_progressed = 0;
+                    cpu.set_execution_context(cycle, cpu_index, current_pc, current_iteration);
+                    if target {
+                        cpu.provenance_schedule_target_stall(cpu_index, current_pc, remaining);
+                    } else {
+                        cpu.provenance_schedule_stall(cpu_index, remaining);
+                    }
+                    break;
+                }
+                if let ScheduleAction::Force { reason } = schedule_action {
+                    cpu.set_execution_context(cycle, cpu_index, current_pc, current_iteration);
+                    cpu.provenance_scheduler_event(
+                        "scheduler_force",
+                        cpu_index,
+                        current_pc,
+                        current_iteration,
+                        reason,
+                        None,
+                        scheduler.config().target,
+                        scheduler.deferrals,
+                    );
+                }
+                let Some(core_instruction) = cores_instructions.get_mut(cpu_index) else {
+                    break;
+                };
+                core_result = InstructionStatus::NotExecuted;
+                if core_instruction.program_counter == core_instruction.instructions.len() {
+                    if core_instruction.instructions.is_empty()
+                        || core_instruction.current_iteration + 1 >= batch_size
+                    {
+                        break;
+                    }
+                    let next_iteration = core_instruction.current_iteration + 1;
+                    if provenance_global_barrier {
+                        if barrier_iteration != Some(next_iteration) {
+                            if !barrier_ready {
+                                break;
+                            }
+                            barrier_iteration = Some(next_iteration);
+                        }
+                    }
+                    core_instruction.current_iteration += 1;
+                    core_instruction.program_counter = 0;
+                    let iteration = core_instruction.current_iteration;
+                    if iteration > GLOBAL_ITERATION.fetch_max(iteration, Ordering::SeqCst) {
+                        cpu.set_execution_context(cycle, cpu_index, 0, iteration);
+                        store_input(cpu, inputs[iteration as usize], input_regions, iteration)?;
+                    }
+                }
+                if !matches!(
+                    schedule_action,
+                    ScheduleAction::Stall { .. } | ScheduleAction::Defer { .. }
+                ) {
+                    cpu.set_current_iteration(core_instruction.current_iteration);
+                    let CoreInstructions {
+                        instructions,
+                        program_counter,
+                        ..
+                    } = core_instruction;
+                    cpu.set_execution_context(
+                        cycle,
+                        cpu_index,
+                        *program_counter,
+                        core_instruction.current_iteration,
+                    );
+                    cycle += 1;
+                    core_result = instructions
+                        .get(*program_counter)
+                        .map_or(InstructionStatus::default(), |inst: &Instruction| {
+                            inst.execute(cpu)
+                        });
+                    if core_result.is_completed() {
+                        scheduler.note_completed(
+                            cpu_index,
+                            *program_counter,
+                            core_instruction.current_iteration,
+                        );
+                        cpu_progressed = 0;
+                        scheduler_no_progress = 0;
+                        *program_counter += 1;
+                    }
+                    if (now.elapsed().unwrap() > Duration::from_secs(5)) {
+                        print_status(cores_instructions);
+                        if let Some(deadlock) = detect_deadlock(cores_instructions) {
+                            bail!(
+                                "Deadlock cycle detected: {} [{}]",
+                                deadlock.cycle,
+                                deadlock.states
+                            );
+                        }
+                        now = SystemTime::now();
                     }
-                    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;
-                }
+            if handle_wait_sync(cores_instructions, &mut sync_events, core_result) {
+                cpu_progressed = 0;
+                scheduler_no_progress = 0;
+            }
+            if let Some(next_core) = scheduler_next {
+                cpu_index = next_core;
+                continue;
+            }
+            let send_recv_result =
+                handle_send_recv(cpu, cores_instructions, send_recv, core_result);
+            if let (true, other_cpu_index) = send_recv_result {
+                cpu_progressed = 0;
+                scheduler_no_progress = 0;
+                cpu_index = other_cpu_index;
+                continue;
+            }
+            if !core_result.is_completed() {
+                scheduler_no_progress += 1;
+            }
+            let states = if scheduler.config().policy == DiagnosticSchedulePolicy::Greedy {
+                None
+            } else {
+                Some(DiagnosticScheduler::state(cores_instructions))
+            };
+            let scheduler_choice = (scheduler_no_progress <= scheduler_no_progress_limit)
+                .then(|| {
+                    states
+                        .as_deref()
+                        .and_then(|states| scheduler.after_block(cpu_index, states, batch_size))
+                })
+                .flatten();
+            if let Some(choice) = scheduler_choice {
+                cpu.provenance_scheduler_event(
+                    "scheduler_prefer",
+                    cpu_index,
+                    cores_instructions[cpu_index].program_counter,
+                    cores_instructions[cpu_index].current_iteration,
+                    choice.reason,
+                    Some(choice.core),
+                    scheduler.config().target,
+                    scheduler.deferrals,
+                );
+                cpu_index = choice.core;
+                continue;
+            }
+            if scheduler_no_progress == scheduler_no_progress_limit + 1
+                && scheduler.config().policy != DiagnosticSchedulePolicy::Greedy
+            {
+                cpu.provenance_scheduler_event(
+                    "scheduler_force",
+                    cpu_index,
+                    cores_instructions[cpu_index].program_counter,
+                    cores_instructions[cpu_index].current_iteration,
+                    "NO_ALTERNATIVE_READY_EVENT",
+                    None,
+                    scheduler.config().target,
+                    scheduler.deferrals,
+                );
+            }
+            match send_recv_result {
+                (true, _) => unreachable!("completed SEND/RECV was handled above"),
                 (false, 0) => {
                     cpu_index = if cpu_index + 1 >= cores_instructions.len() {
                         cpu_progressed -= 1;
@@ -206,7 +787,8 @@ impl<'a> Executable<'a> {
 
         #[cfg(feature = "profile_time")]
         TRACER.lock().unwrap().report();
-        Ok(())
+        cpu.finish_provenance();
+        Ok(cpu.finish_host_store_recording())
     }
 
     pub fn cpu(&self) -> &CPU<'a> {
@@ -220,7 +802,7 @@ impl<'a> Executable<'a> {
     pub fn dump(&self) {
         let core_instructions = &self.core_instructions;
         for (i, core_instruction) in core_instructions.iter().enumerate() {
-            eprintln!("INST OF CORE {}:", i);
+            eprintln!("Instructions for core {}:", i);
             for inst in &core_instruction.instructions {
                 inst.dump();
             }
@@ -228,6 +810,35 @@ impl<'a> Executable<'a> {
     }
 }
 
+fn validate_inputs(inputs: &[&[u8]], input_regions: &[(usize, usize)]) -> Result<()> {
+    let input_size = input_regions.iter().try_fold(0usize, |total, (_, size)| {
+        total.checked_add(*size).context("input size overflow")
+    })?;
+    if inputs.is_empty() {
+        bail!("at least one input is required");
+    }
+    if inputs.iter().any(|input| input.len() != input_size) {
+        bail!("each input must contain exactly {input_size} bytes");
+    }
+    Ok(())
+}
+
+fn store_input(
+    cpu: &mut CPU,
+    input: &[u8],
+    input_regions: &[(usize, usize)],
+    sample: u32,
+) -> Result<()> {
+    let mut offset = 0;
+    for &(address, size) in input_regions {
+        cpu.host()
+            .execute_store(address, &input[offset..offset + size])?;
+        cpu.provenance_input_store(address, size, sample);
+        offset += size;
+    }
+    Ok(())
+}
+
 fn detect_deadlock(cores_instructions: &[CoreInstructions]) -> Option<DeadlockInfo> {
     #[derive(Debug, PartialEq, Eq)]
     enum CoreState {
@@ -250,7 +861,10 @@ fn detect_deadlock(cores_instructions: &[CoreInstructions]) -> Option<DeadlockIn
         let (this_core, target_core) = data.get_core_immcore();
 
         if isa_recv(functor_address) {
-            states.insert(this_core, CoreState::ReceivingFrom(target_core, data.imm_len()));
+            states.insert(
+                this_core,
+                CoreState::ReceivingFrom(target_core, data.imm_len()),
+            );
         } else if isa_send(functor_address) {
             states.insert(this_core, CoreState::SendingTo(target_core, data.imm_len()));
         } else {
@@ -274,8 +888,7 @@ fn detect_deadlock(cores_instructions: &[CoreInstructions]) -> Option<DeadlockIn
                     wait_for.insert(core_id, *target_core);
                 }
             }
-            CoreState::Working | CoreState::Halted => {
-            }
+            CoreState::Working | CoreState::Halted => {}
         }
     }
 
@@ -347,12 +960,171 @@ fn detect_deadlock(cores_instructions: &[CoreInstructions]) -> Option<DeadlockIn
     None
 }
 
-fn handle_wait_sync<'a, 'b, 'c>(
-    cpu: &'b mut CPU<'a>,
-    core_instructions: &'c mut [CoreInstructions],
+fn handle_wait_sync(
+    core_instructions: &mut [CoreInstructions],
+    events: &mut SyncEvents,
     core_result: InstructionStatus,
-) where
-    'a: 'b,
-    'a: 'c,
-{
+) -> bool {
+    match core_result {
+        InstructionStatus::Sync(data) => {
+            let (source, target) = data.get_core_immcore();
+            let register = data.offset_select() as usize;
+            events[target as usize][register] += 1;
+            core_instructions[source as usize].program_counter += 1;
+            true
+        }
+        InstructionStatus::Waiting(data) => {
+            let core = data.core_indx() as usize;
+            let register = data.offset_select() as usize;
+            let value = data.offset_value();
+            if events[core][register] >= value {
+                events[core][register] -= value;
+                core_instructions[core].program_counter += 1;
+                true
+            } else {
+                false
+            }
+        }
+        _ => false,
+    }
+}
+
+#[cfg(test)]
+mod scheduler_tests {
+    use super::*;
+
+    fn target() -> DiagnosticScheduleTarget {
+        DiagnosticScheduleTarget {
+            writer_core: 0,
+            writer_pc: 3,
+            reader_core: 1,
+            reader_pc: 2,
+            address_begin: 100,
+            address_end: 200,
+            reader_iteration: Some(0),
+            writer_min_iteration: Some(1),
+        }
+    }
+
+    fn states() -> Vec<ScheduleCoreState> {
+        vec![
+            ScheduleCoreState {
+                program_counter: 3,
+                instruction_count: 8,
+                current_iteration: 1,
+            },
+            ScheduleCoreState {
+                program_counter: 2,
+                instruction_count: 8,
+                current_iteration: 0,
+            },
+        ]
+    }
+
+    #[test]
+    fn adversarial_policy_defers_reader_until_writer_store() {
+        let mut scheduler = DiagnosticScheduler::new(DiagnosticScheduleConfig {
+            policy: DiagnosticSchedulePolicy::Adversarial,
+            target: Some(target()),
+            deferral_budget: 4,
+            ..Default::default()
+        });
+        assert_eq!(
+            scheduler.before_instruction(1, &states(), 4),
+            ScheduleAction::Defer {
+                next_core: 0,
+                reason: "target_consumer"
+            }
+        );
+        scheduler.note_completed(0, 3, 1);
+        assert_eq!(
+            scheduler.before_instruction(1, &states(), 4),
+            ScheduleAction::Execute
+        );
+    }
+
+    #[test]
+    fn randomized_policy_is_deterministic_for_fixed_seed() {
+        let config = DiagnosticScheduleConfig {
+            policy: DiagnosticSchedulePolicy::Randomized,
+            seed: 17,
+            ..Default::default()
+        };
+        let mut left = DiagnosticScheduler::new(config);
+        let mut right = DiagnosticScheduler::new(config);
+        let states = vec![
+            ScheduleCoreState {
+                program_counter: 0,
+                instruction_count: 4,
+                current_iteration: 0,
+            },
+            ScheduleCoreState {
+                program_counter: 1,
+                instruction_count: 4,
+                current_iteration: 0,
+            },
+            ScheduleCoreState {
+                program_counter: 2,
+                instruction_count: 4,
+                current_iteration: 0,
+            },
+        ];
+        for current in [0, 1, 2, 0, 1] {
+            assert_eq!(
+                left.before_instruction(current, &states, 2),
+                right.before_instruction(current, &states, 2)
+            );
+        }
+    }
+
+    #[test]
+    fn target_stall_is_consumed_without_changing_program_order() {
+        let mut scheduler = DiagnosticScheduler::new(DiagnosticScheduleConfig {
+            fixed_target_stall: Some((1, 2, None, 2)),
+            ..Default::default()
+        });
+        let states = states();
+        assert_eq!(
+            scheduler.before_instruction(1, &states, 4),
+            ScheduleAction::Stall {
+                remaining: 1,
+                target: true
+            }
+        );
+        assert_eq!(
+            scheduler.before_instruction(1, &states, 4),
+            ScheduleAction::Stall {
+                remaining: 0,
+                target: true
+            }
+        );
+        assert_eq!(
+            scheduler.before_instruction(1, &states, 4),
+            ScheduleAction::Execute
+        );
+        assert_eq!(states[1].program_counter, 2);
+    }
+
+    #[test]
+    fn adversarial_policy_releases_reader_when_writer_state_stagnates() {
+        let mut scheduler = DiagnosticScheduler::new(DiagnosticScheduleConfig {
+            policy: DiagnosticSchedulePolicy::Adversarial,
+            target: Some(target()),
+            deferral_budget: 10_000,
+            ..Default::default()
+        });
+        let states = states();
+        for _ in 0..256 {
+            assert!(matches!(
+                scheduler.before_instruction(1, &states, 4),
+                ScheduleAction::Defer { .. }
+            ));
+        }
+        assert_eq!(
+            scheduler.before_instruction(1, &states, 4),
+            ScheduleAction::Force {
+                reason: "PRODUCER_BLOCKED_BY_REAL_DEPENDENCY"
+            }
+        );
+    }
 }

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

@@ -0,0 +1,639 @@
+use serde_json::{Value, json};
+use std::{
+    collections::BTreeSet,
+    fs::File,
+    io::{BufWriter, Write},
+    path::Path,
+    sync::{Arc, Mutex},
+};
+
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord)]
+pub enum Provenance {
+    #[default]
+    Uninitialized,
+    Unknown,
+    Samples(u64),
+}
+
+impl Provenance {
+    pub fn sample(sample: u32) -> Self {
+        if sample < 64 {
+            Self::Samples(1_u64 << sample)
+        } else {
+            Self::Unknown
+        }
+    }
+
+    fn merge(self, other: Self) -> Self {
+        match (self, other) {
+            (Self::Unknown, _) | (_, Self::Unknown) => Self::Unknown,
+            (Self::Uninitialized, value) | (value, Self::Uninitialized) => value,
+            (Self::Samples(left), Self::Samples(right)) => Self::Samples(left | right),
+        }
+    }
+
+    fn merge_all(values: impl IntoIterator<Item = Self>) -> Self {
+        let mut values = values.into_iter();
+        values
+            .next()
+            .map_or(Self::Uninitialized, |first| values.fold(first, Self::merge))
+    }
+
+    fn samples(self) -> Vec<u32> {
+        match self {
+            Self::Samples(mask) => (0..64)
+                .filter(|sample| mask & (1_u64 << sample) != 0)
+                .collect(),
+            Self::Unknown | Self::Uninitialized => Vec::new(),
+        }
+    }
+
+    fn state(self) -> &'static str {
+        match self {
+            Self::Samples(_) => "known",
+            Self::Unknown => "unknown",
+            Self::Uninitialized => "uninitialized",
+        }
+    }
+
+    fn is_mixed(self) -> bool {
+        matches!(self, Self::Samples(mask) if mask.count_ones() > 1)
+    }
+
+    fn json(self) -> Value {
+        json!({
+            "provenance": self.samples(),
+            "provenance_state": self.state(),
+            "mixed": self.is_mixed(),
+        })
+    }
+}
+
+#[derive(Clone, Copy, Debug)]
+struct ExecutionContext {
+    cycle: u64,
+    core: usize,
+    pc: usize,
+    iteration: u32,
+}
+
+impl Default for ExecutionContext {
+    fn default() -> Self {
+        Self {
+            cycle: 0,
+            core: 0,
+            pc: 0,
+            iteration: 0,
+        }
+    }
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
+struct Writer {
+    version: u64,
+    cycle: u64,
+    core: usize,
+    pc: usize,
+    iteration: u32,
+    provenance: Provenance,
+}
+
+#[derive(Clone, Copy, Debug, Default)]
+struct GlobalCell {
+    provenance: Provenance,
+    writer: Option<Writer>,
+}
+
+#[derive(Debug)]
+struct TraceSink {
+    output: BufWriter<File>,
+}
+
+impl TraceSink {
+    fn new(path: impl AsRef<Path>) -> std::io::Result<Self> {
+        let file = File::create(path)?;
+        Ok(Self {
+            output: BufWriter::new(file),
+        })
+    }
+
+    fn event(&mut self, value: Value) {
+        serde_json::to_writer(&mut self.output, &value).expect("write provenance event");
+        self.output
+            .write_all(b"\n")
+            .expect("write provenance newline");
+    }
+
+    fn flush(&mut self) {
+        self.output.flush().expect("flush provenance trace");
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct ProvenanceTracker {
+    global: Vec<GlobalCell>,
+    local: Vec<Vec<Provenance>>,
+    next_version: u64,
+    context: ExecutionContext,
+    sink: Arc<Mutex<TraceSink>>,
+}
+
+impl ProvenanceTracker {
+    pub fn new(core_count: usize, path: impl AsRef<Path>) -> std::io::Result<Self> {
+        if let Some(parent) = path.as_ref().parent() {
+            std::fs::create_dir_all(parent)?;
+        }
+        let sink = Arc::new(Mutex::new(TraceSink::new(path)?));
+        let tracker = Self {
+            global: Vec::new(),
+            local: vec![Vec::new(); core_count],
+            next_version: 0,
+            context: ExecutionContext::default(),
+            sink,
+        };
+        tracker.write(json!({
+            "event": "provenance_trace_start",
+            "schema": 1,
+            "core_count": core_count,
+        }));
+        Ok(tracker)
+    }
+
+    pub fn begin_batch(&mut self, batch_size: usize) {
+        self.global.fill(GlobalCell::default());
+        for memory in &mut self.local {
+            memory.fill(Provenance::Uninitialized);
+        }
+        self.next_version = 0;
+        self.write(json!({
+            "event": "batch_start",
+            "batch_size": batch_size,
+        }));
+    }
+
+    pub fn set_context(&mut self, cycle: u64, core: usize, pc: usize, iteration: u32) {
+        self.context = ExecutionContext {
+            cycle,
+            core,
+            pc,
+            iteration,
+        };
+    }
+
+    pub fn cycle(&self) -> u64 {
+        self.context.cycle
+    }
+
+    pub fn flush(&mut self) {
+        self.sink.lock().unwrap().flush();
+    }
+
+    pub fn schedule_stall(&self, core: usize, remaining_cycles: u64) {
+        self.write(json!({
+            "event": "diagnostic_core_stall",
+            "cycle": self.context.cycle,
+            "core": core,
+            "pc": self.context.pc,
+            "core_iteration": self.context.iteration,
+            "remaining_cycles": remaining_cycles,
+        }));
+    }
+
+    pub fn schedule_target_stall(&self, core: usize, pc: usize, remaining_cycles: u64) {
+        self.write(json!({
+            "event": "diagnostic_target_stall",
+            "cycle": self.context.cycle,
+            "core": core,
+            "pc": pc,
+            "core_iteration": self.context.iteration,
+            "remaining_cycles": remaining_cycles,
+        }));
+    }
+
+    pub fn schedule_config(
+        &self,
+        policy: &'static str,
+        seed: u64,
+        target: Option<Value>,
+        deferral_budget: u64,
+        fixed_stall: Option<(usize, u64)>,
+        fixed_target_stall: Option<(usize, usize, Option<u32>, u64)>,
+    ) {
+        self.write(json!({
+            "event": "scheduler_config",
+            "schedule_policy": policy,
+            "schedule_seed": seed,
+            "target_dependency": target,
+            "deferral_budget": deferral_budget,
+            "fixed_stall": fixed_stall.map(|(core, cycles)| json!({
+                "core": core,
+                "cycles": cycles,
+            })),
+            "fixed_target_stall": fixed_target_stall.map(|(core, pc, iteration, cycles)| json!({
+                "core": core,
+                "pc": pc,
+                "iteration": iteration,
+                "cycles": cycles,
+            })),
+        }));
+    }
+
+    pub fn scheduler_event(&self, event: Value) {
+        self.write(event);
+    }
+
+    fn write(&self, value: Value) {
+        self.sink.lock().unwrap().event(value);
+    }
+
+    fn ensure_global(&mut self, end: usize) {
+        if self.global.len() < end {
+            self.global.resize(end, GlobalCell::default());
+        }
+    }
+
+    fn ensure_local(&mut self, core: usize, end: usize) {
+        if let Some(memory) = self.local.get_mut(core)
+            && memory.len() < end
+        {
+            memory.resize(end, Provenance::Uninitialized);
+        }
+    }
+
+    fn local_tags(&mut self, core: usize, address: usize, size: usize) -> Vec<Provenance> {
+        let Some(end) = address.checked_add(size) else {
+            return vec![Provenance::Unknown; size];
+        };
+        self.ensure_local(core, end);
+        self.local[core][address..end].to_vec()
+    }
+
+    fn store_local_tags(&mut self, core: usize, address: usize, tags: &[Provenance]) {
+        let Some(end) = address.checked_add(tags.len()) else {
+            return;
+        };
+        self.ensure_local(core, end);
+        self.local[core][address..end].copy_from_slice(tags);
+    }
+
+    fn unique_versions(cells: &[GlobalCell]) -> Vec<u64> {
+        cells
+            .iter()
+            .filter_map(|cell| cell.writer.map(|writer| writer.version))
+            .collect::<BTreeSet<_>>()
+            .into_iter()
+            .collect()
+    }
+
+    fn unique_writers(cells: &[GlobalCell]) -> Vec<Writer> {
+        cells
+            .iter()
+            .filter_map(|cell| cell.writer)
+            .collect::<BTreeSet<_>>()
+            .into_iter()
+            .collect()
+    }
+
+    fn writer_json(writer: Writer) -> Value {
+        json!({
+            "version": writer.version,
+            "cycle": writer.cycle,
+            "core": writer.core,
+            "pc": writer.pc,
+            "core_iteration": writer.iteration,
+            "provenance": writer.provenance.samples(),
+            "provenance_state": writer.provenance.state(),
+        })
+    }
+
+    pub fn input_store(&mut self, address: usize, size: usize, sample: u32) {
+        let Some(end) = address.checked_add(size) else {
+            return;
+        };
+        self.ensure_global(end);
+        let provenance = Provenance::sample(sample);
+        self.next_version += 1;
+        let writer = Writer {
+            version: self.next_version,
+            cycle: self.context.cycle,
+            core: 0,
+            pc: 0,
+            iteration: sample,
+            provenance,
+        };
+        let overwritten_versions = Self::unique_versions(&self.global[address..end]);
+        for cell in &mut self.global[address..end] {
+            *cell = GlobalCell {
+                provenance,
+                writer: Some(writer),
+            };
+        }
+        let mut event = json!({
+            "event": "external_input_store",
+            "cycle": self.context.cycle,
+            "core": 0,
+            "pc": 0,
+            "core_iteration": sample,
+            "address": address,
+            "size": size,
+            "sample": sample,
+            "version": writer.version,
+            "overwritten_versions": overwritten_versions,
+        });
+        if let Some(object) = event.as_object_mut() {
+            object.extend(provenance.json().as_object().unwrap().clone());
+        }
+        self.write(event);
+    }
+
+    pub fn global_store_from_local(
+        &mut self,
+        core: usize,
+        global_address: usize,
+        local_address: usize,
+        size: usize,
+    ) {
+        let Some(end) = global_address.checked_add(size) else {
+            return;
+        };
+        let tags = self.local_tags(core, local_address, size);
+        self.ensure_global(end);
+        self.next_version += 1;
+        let provenance = Provenance::merge_all(tags.iter().copied());
+        let writer = Writer {
+            version: self.next_version,
+            cycle: self.context.cycle,
+            core,
+            pc: self.context.pc,
+            iteration: self.context.iteration,
+            provenance,
+        };
+        let overwritten_versions = Self::unique_versions(&self.global[global_address..end]);
+        let overwritten_writers = Self::unique_writers(&self.global[global_address..end]);
+        for (cell, tag) in self.global[global_address..end].iter_mut().zip(tags) {
+            *cell = GlobalCell {
+                provenance: tag,
+                writer: Some(writer),
+            };
+        }
+        let mut event = json!({
+            "event": "global_store",
+            "cycle": self.context.cycle,
+            "core": core,
+            "pc": self.context.pc,
+            "core_iteration": self.context.iteration,
+            "address": global_address,
+            "local_address": local_address,
+            "size": size,
+            "version": writer.version,
+            "overwritten_versions": overwritten_versions,
+            "overwritten_writers": overwritten_writers.into_iter().map(Self::writer_json).collect::<Vec<_>>(),
+        });
+        if let Some(object) = event.as_object_mut() {
+            object.extend(provenance.json().as_object().unwrap().clone());
+        }
+        self.write(event);
+    }
+
+    pub fn global_load_to_local(
+        &mut self,
+        core: usize,
+        global_address: usize,
+        local_address: usize,
+        size: usize,
+    ) {
+        let Some(end) = global_address.checked_add(size) else {
+            return;
+        };
+        self.ensure_global(end);
+        let cells = self.global[global_address..end].to_vec();
+        let tags: Vec<_> = cells.iter().map(|cell| cell.provenance).collect();
+        let provenance = Provenance::merge_all(tags.iter().copied());
+        let writers = Self::unique_writers(&cells);
+        let versions = Self::unique_versions(&cells);
+        self.store_local_tags(core, local_address, &tags);
+        let mut event = json!({
+            "event": "global_load",
+            "cycle": self.context.cycle,
+            "core": core,
+            "pc": self.context.pc,
+            "core_iteration": self.context.iteration,
+            "address": global_address,
+            "local_address": local_address,
+            "size": size,
+            "versions": versions,
+            "last_writers": writers.into_iter().map(Self::writer_json).collect::<Vec<_>>(),
+        });
+        if let Some(object) = event.as_object_mut() {
+            object.extend(provenance.json().as_object().unwrap().clone());
+        }
+        self.write(event);
+    }
+
+    pub fn local_copy(
+        &mut self,
+        core: usize,
+        destination: usize,
+        source: usize,
+        size: usize,
+        operation: &'static str,
+    ) {
+        let tags = self.local_tags(core, source, size);
+        let provenance = Provenance::merge_all(tags.iter().copied());
+        self.store_local_tags(core, destination, &tags);
+        self.local_event(
+            operation,
+            core,
+            destination,
+            size,
+            provenance,
+            &[provenance],
+        );
+    }
+
+    pub fn local_strided_copy(
+        &mut self,
+        core: usize,
+        destination: usize,
+        source: usize,
+        element_size: usize,
+        stride: usize,
+        element_count: usize,
+        operation: &'static str,
+    ) {
+        let mut tags = Vec::with_capacity(element_size.saturating_mul(element_count));
+        for index in 0..element_count {
+            let address =
+                source.saturating_add(index.saturating_mul(stride).saturating_mul(element_size));
+            tags.extend(self.local_tags(core, address, element_size));
+        }
+        let provenance = Provenance::merge_all(tags.iter().copied());
+        self.store_local_tags(core, destination, &tags);
+        self.local_event(
+            operation,
+            core,
+            destination,
+            tags.len(),
+            provenance,
+            &[provenance],
+        );
+    }
+
+    pub fn local_transform(
+        &mut self,
+        core: usize,
+        destination: usize,
+        sources: &[(usize, usize)],
+        output_size: usize,
+        operation: &'static str,
+    ) {
+        let source_tags: Vec<Vec<_>> = sources
+            .iter()
+            .map(|&(address, size)| self.local_tags(core, address, size))
+            .collect();
+        let mut output = Vec::with_capacity(output_size);
+        for index in 0..output_size {
+            let provenance = Provenance::merge_all(
+                source_tags
+                    .iter()
+                    .filter_map(|tags| tags.get(index).copied()),
+            );
+            output.push(provenance);
+        }
+        let provenance = Provenance::merge_all(output.iter().copied());
+        self.store_local_tags(core, destination, &output);
+        let operands: Vec<_> = source_tags
+            .iter()
+            .map(|tags| Provenance::merge_all(tags.iter().copied()))
+            .collect();
+        self.local_event(
+            operation,
+            core,
+            destination,
+            output_size,
+            provenance,
+            &operands,
+        );
+    }
+
+    pub fn local_broadcast_transform(
+        &mut self,
+        core: usize,
+        destination: usize,
+        source: usize,
+        source_size: usize,
+        output_size: usize,
+        operation: &'static str,
+    ) {
+        let tags = self.local_tags(core, source, source_size);
+        let provenance = Provenance::merge_all(tags.iter().copied());
+        self.store_local_tags(core, destination, &vec![provenance; output_size]);
+        self.local_event(
+            operation,
+            core,
+            destination,
+            output_size,
+            provenance,
+            &[provenance],
+        );
+    }
+
+    pub fn local_mvm_transform(
+        &mut self,
+        core: usize,
+        destination: usize,
+        source: usize,
+        element_size: usize,
+        output_size: usize,
+        used_rows: &[bool],
+        operation: &'static str,
+    ) {
+        let mut provenance = Provenance::Uninitialized;
+        for (row, used) in used_rows.iter().copied().enumerate() {
+            if used {
+                provenance = provenance.merge(Provenance::merge_all(self.local_tags(
+                    core,
+                    source + row * element_size,
+                    element_size,
+                )));
+            }
+        }
+        self.store_local_tags(core, destination, &vec![provenance; output_size]);
+        self.local_event(
+            operation,
+            core,
+            destination,
+            output_size,
+            provenance,
+            &[provenance],
+        );
+    }
+
+    pub fn send_transfer(
+        &mut self,
+        sender: usize,
+        receiver: usize,
+        source: usize,
+        destination: usize,
+        size: usize,
+    ) {
+        let tags = self.local_tags(sender, source, size);
+        let provenance = Provenance::merge_all(tags.iter().copied());
+        self.store_local_tags(receiver, destination, &tags);
+        let mut event = json!({
+            "event": "send_recv_transfer",
+            "cycle": self.context.cycle,
+            "pc": self.context.pc,
+            "core_iteration": self.context.iteration,
+            "sender_core": sender,
+            "receiver_core": receiver,
+            "source_address": source,
+            "destination_address": destination,
+            "size": size,
+        });
+        if let Some(object) = event.as_object_mut() {
+            object.extend(provenance.json().as_object().unwrap().clone());
+        }
+        self.write(event);
+    }
+
+    fn local_event(
+        &self,
+        operation: &'static str,
+        core: usize,
+        destination: usize,
+        size: usize,
+        provenance: Provenance,
+        operands: &[Provenance],
+    ) {
+        let mut event = json!({
+            "event": "local_compute",
+            "operation": operation,
+            "cycle": self.context.cycle,
+            "core": core,
+            "pc": self.context.pc,
+            "core_iteration": self.context.iteration,
+            "destination_address": destination,
+            "size": size,
+            "operand_provenance": operands.iter().map(|tag| tag.json()).collect::<Vec<_>>(),
+        });
+        if let Some(object) = event.as_object_mut() {
+            object.extend(provenance.json().as_object().unwrap().clone());
+        }
+        self.write(event);
+        if provenance.is_mixed() {
+            self.write(json!({
+                "event": "cross_sample_data_mix",
+                "cycle": self.context.cycle,
+                "core": core,
+                "pc": self.context.pc,
+                "core_iteration": self.context.iteration,
+                "operation": operation,
+                "destination_address": destination,
+                "size": size,
+                "operand_provenance": operands.iter().map(|tag| tag.json()).collect::<Vec<_>>(),
+                "provenance": provenance.samples(),
+            }));
+        }
+    }
+}

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

@@ -33,24 +33,25 @@ pub struct SendRecv {
 impl SendRecv {
     pub fn new(num_core: usize) -> Self {
         let sending = [Option::None].repeat(num_core);
-        let reciving = [Option::None].repeat(num_core);
+        let receiving = [Option::None].repeat(num_core);
         Self {
             sending: sending.into(),
-            receiving: reciving.into(),
+            receiving: receiving.into(),
         }
     }
 }
 
-pub fn handle_send_recv<'a, 'b >(
+pub fn handle_send_recv<'a, 'b>(
     cpu: &'b mut CPU<'a>,
-    core_instructions: & mut [CoreInstructions],
-    send_recv: & mut SendRecv,
+    core_instructions: &mut [CoreInstructions],
+    send_recv: &mut SendRecv,
     core_result: InstructionStatus,
 ) -> (bool, usize)
-where 'a : 'b
+where
+    'a: 'b,
 {
     let transfer_memory = |cpu: &'b mut CPU<'a>,
-                           core_instructions: & mut [CoreInstructions],
+                           core_instructions: &mut [CoreInstructions],
                            sender: Option<SendRecvInfo>,
                            receiver: Option<SendRecvInfo>| {
         if let Some(sender) = sender
@@ -72,18 +73,27 @@ where 'a : 'b
                 let data = inst.data;
                 TRACER.lock().unwrap().pre_recv(cpu, data);
             }
-            let [sender_core, reciver_core] =
-                cpu.get_multiple_cores([sender.internal_core, receiver.internal_core]);
-            let memory = sender_core
-                .load::<u8>(sender.address, sender.size)
-                .with_context(|| {
-                    format!(
-                        "Sender crash tranfering memroy from {} with size {}",
-                        sender.address, sender.size
-                    )
-                })
-                .unwrap();
-            reciver_core.execute_store(receiver.address, memory[0]);
+            {
+                let [sender_core, receiver_core] =
+                    cpu.get_multiple_cores([sender.internal_core, receiver.internal_core]);
+                let memory = sender_core
+                    .load::<u8>(sender.address, sender.size)
+                    .with_context(|| {
+                        format!(
+                            "Sender crashed while transferring memory from {} with size {}",
+                            sender.address, sender.size
+                        )
+                    })
+                    .unwrap();
+                receiver_core.execute_store(receiver.address, memory[0]);
+            }
+            cpu.provenance_send_transfer(
+                sender.internal_core,
+                receiver.internal_core,
+                sender.address,
+                receiver.address,
+                sender.size,
+            );
             {
                 let sender = &mut core_instructions[sender.internal_core];
                 let pc = sender.program_counter;
@@ -123,19 +133,19 @@ where 'a : 'b
             let receiver: usize = imm_core.try_into().expect("imm_core can not be negative");
             assert_ne!(receiver, 0, "Host can not use receive");
             send_recv.sending[sender] = Some(SendRecvInfo::new(sender, receiver, address, imm_len));
-            let transfered = transfer_memory(
+            let transferred = transfer_memory(
                 cpu,
                 core_instructions,
                 send_recv.sending[sender],
                 send_recv.receiving[receiver],
             );
-            if transfered {
+            if transferred {
                 send_recv.sending[sender] = None;
                 send_recv.receiving[receiver] = None;
             }
-            (transfered, receiver)
+            (transferred, if transferred { receiver } else { 0 })
         }
-        InstructionStatus::Reciving(instruction_data) => {
+        InstructionStatus::Receiving(instruction_data) => {
             let (core_idx, imm_core) = instruction_data.get_core_immcore();
             let rd = instruction_data.rd();
             let imm_len = instruction_data
@@ -152,17 +162,17 @@ where 'a : 'b
             assert_ne!(sender, 0, "Host can not use send");
             send_recv.receiving[receiver] =
                 Some(SendRecvInfo::new(receiver, sender, address, imm_len));
-            let transfered = transfer_memory(
+            let transferred = transfer_memory(
                 cpu,
                 core_instructions,
                 send_recv.sending[sender],
                 send_recv.receiving[receiver],
             );
-            if transfered {
+            if transferred {
                 send_recv.sending[sender] = None;
                 send_recv.receiving[receiver] = None;
             }
-            (transfered, sender)
+            (transferred, if transferred { sender } else { 0 })
         }
         _ => (false, 0),
     }

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

@@ -1,4 +1,3 @@
-
 use std::fs::File;
 
 use crate::{
@@ -19,58 +18,41 @@ impl Trace {
     /////////////////Scalar/register Instructions//////////////////
     ///////////////////////////////////////////////////////////////
 
-    pub fn pre_sldi(&mut self, cores: &mut CPU, data: InstructionData) {
+    pub fn pre_sldi(&mut self, cores: &mut CPU, data: InstructionData) {}
 
-    }
+    pub fn post_sldi(&mut self, cores: &mut CPU, data: InstructionData) {}
 
-    pub fn post_sldi(&mut self, cores: &mut CPU, data: InstructionData) {
-    }
+    pub fn pre_sld(&mut self, cores: &mut CPU, data: InstructionData) {}
 
-    pub fn pre_sld(&mut self, cores: &mut CPU, data: InstructionData) {
-    }
+    pub fn post_sld(&mut self, cores: &mut CPU, data: InstructionData) {}
 
-    pub fn post_sld(&mut self, cores: &mut CPU, data: InstructionData) {
-    }
+    pub fn pre_sadd(&mut self, cores: &mut CPU, data: InstructionData) {}
 
-    pub fn pre_sadd(&mut self, cores: &mut CPU, data: InstructionData) {
-    }
+    pub fn post_sadd(&mut self, cores: &mut CPU, data: InstructionData) {}
 
-    pub fn post_sadd(&mut self, cores: &mut CPU, data: InstructionData) {
-    }
+    pub fn pre_ssub(&mut self, cores: &mut CPU, data: InstructionData) {}
 
-    pub fn pre_ssub(&mut self, cores: &mut CPU, data: InstructionData) {
-    }
+    pub fn post_ssub(&mut self, cores: &mut CPU, data: InstructionData) {}
 
-    pub fn post_ssub(&mut self, cores: &mut CPU, data: InstructionData) {
-    }
+    pub fn pre_smul(&mut self, cores: &mut CPU, data: InstructionData) {}
 
-    pub fn pre_smul(&mut self, cores: &mut CPU, data: InstructionData) {
-    }
+    pub fn post_smul(&mut self, cores: &mut CPU, data: InstructionData) {}
 
-    pub fn post_smul(&mut self, cores: &mut CPU, data: InstructionData) {
-    }
+    pub fn pre_saddi(&mut self, cores: &mut CPU, data: InstructionData) {}
 
-    pub fn pre_saddi(&mut self, cores: &mut CPU, data: InstructionData) {
-    }
+    pub fn post_saddi(&mut self, cores: &mut CPU, data: InstructionData) {}
 
-    pub fn post_saddi(&mut self, cores: &mut CPU, data: InstructionData) {
-    }
+    pub fn pre_smuli(&mut self, cores: &mut CPU, data: InstructionData) {}
 
-    pub fn pre_smuli(&mut self, cores: &mut CPU, data: InstructionData) {
-    }
-
-    pub fn post_smuli(&mut self, cores: &mut CPU, data: InstructionData) {
-    }
+    pub fn post_smuli(&mut self, cores: &mut CPU, data: InstructionData) {}
 
     /////////////////////////////////////////////////////////////////
     ///////////////////Matrix/vector Instructions////////////////////
     /////////////////////////////////////////////////////////////////
 
-    pub fn pre_setbw(&mut self, cores: &mut CPU, data: InstructionData) {
-    }
+    pub fn pre_setbw(&mut self, cores: &mut CPU, data: InstructionData) {}
 
-    pub fn post_setbw(&mut self, cores: &mut CPU, data: InstructionData) {
-    }
+    pub fn post_setbw(&mut self, cores: &mut CPU, data: InstructionData) {}
 
     pub fn pre_mvm<F, M, T>(&mut self, cores: &mut CPU, data: InstructionData)
     where
@@ -80,7 +62,7 @@ impl Trace {
         M: UpcastDestTraits<M> + MemoryStorable + FromFloat,
         F: UpcastDestTraits<F> + MemoryStorable,
     {
-        self.mvm_impl::<F,M,T>(cores, data, "Pre");
+        self.mvm_impl::<F, M, T>(cores, data, "Pre");
     }
 
     pub fn post_mvm<F, M, T>(&mut self, cores: &mut CPU, data: InstructionData)
@@ -91,11 +73,15 @@ impl Trace {
         M: UpcastDestTraits<M> + MemoryStorable + FromFloat,
         F: UpcastDestTraits<F> + MemoryStorable,
     {
-        self.mvm_impl::<F,M,T>(cores, data, "Post");
+        self.mvm_impl::<F, M, T>(cores, data, "Post");
     }
 
-    pub fn mvm_impl<F, M, T>(&mut self, cores: &mut CPU, data: InstructionData, prefix : &'static str)
-    where
+    pub fn mvm_impl<F, M, T>(
+        &mut self,
+        cores: &mut CPU,
+        data: InstructionData,
+        prefix: &'static str,
+    ) where
         [F]: UpcastSlice<T> + UpcastSlice<M>,
         [M]: UpcastSlice<T>,
         T: UpcastDestTraits<T> + MemoryStorable,
@@ -267,39 +253,27 @@ impl Trace {
     /////////////////////////////////////////////////////////////////
     /////Communication/synchronization Instructions/////////////////
     /////////////////////////////////////////////////////////////////
-    pub fn pre_ld(&mut self, cores: &mut CPU, data: InstructionData) {
-    }
+    pub fn pre_ld(&mut self, cores: &mut CPU, data: InstructionData) {}
 
-    pub fn post_ld(&mut self, cores: &mut CPU, data: InstructionData) {
-    }
+    pub fn post_ld(&mut self, cores: &mut CPU, data: InstructionData) {}
 
-    pub fn pre_st(&mut self, cores: &mut CPU, data: InstructionData) {
-    }
+    pub fn pre_st(&mut self, cores: &mut CPU, data: InstructionData) {}
 
-    pub fn post_st(&mut self, cores: &mut CPU, data: InstructionData) {
-    }
+    pub fn post_st(&mut self, cores: &mut CPU, data: InstructionData) {}
 
-    pub fn pre_lldi(&mut self, cores: &mut CPU, data: InstructionData) {
-    }
+    pub fn pre_lldi(&mut self, cores: &mut CPU, data: InstructionData) {}
 
-    pub fn post_lldi(&mut self, cores: &mut CPU, data: InstructionData) {
-    }
+    pub fn post_lldi(&mut self, cores: &mut CPU, data: InstructionData) {}
 
-    pub fn pre_lmv(&mut self, cores: &mut CPU, data: InstructionData) {
-    }
+    pub fn pre_lmv(&mut self, cores: &mut CPU, data: InstructionData) {}
 
-    pub fn post_lmv(&mut self, cores: &mut CPU, data: InstructionData) {
-    }
+    pub fn post_lmv(&mut self, cores: &mut CPU, data: InstructionData) {}
 
-    pub fn pre_send(&mut self, cores: &mut CPU, data: InstructionData) {
-    }
+    pub fn pre_send(&mut self, cores: &mut CPU, data: InstructionData) {}
 
-    pub fn post_send(&mut self, cores: &mut CPU, data: InstructionData) {
-    }
+    pub fn post_send(&mut self, cores: &mut CPU, data: InstructionData) {}
 
-    pub fn pre_recv(&mut self, cores: &mut CPU, data: InstructionData) {
-    }
+    pub fn pre_recv(&mut self, cores: &mut CPU, data: InstructionData) {}
 
-    pub fn post_recv(&mut self, cores: &mut CPU, data: InstructionData) {
-    }
+    pub fn post_recv(&mut self, cores: &mut CPU, data: InstructionData) {}
 }

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

@@ -15,7 +15,6 @@ use crate::Executable;
 use std::path::PathBuf;
 use std::sync::{LazyLock, Mutex};
 
-
 #[cfg(not(any(feature = "tracing", feature = "profile_time")))]
 pub struct Trace {}
 
@@ -28,5 +27,4 @@ impl Trace {
     pub fn init(&mut self, num_core: usize, path: PathBuf) {}
 }
 
-
 pub static TRACER: LazyLock<Mutex<Trace>> = LazyLock::new(|| Trace::new().into());

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

@@ -8,7 +8,7 @@ pub mod profile_analysis;
 pub mod profile_isa;
 
 pub struct Trace {
-    instruction_times: HashMap<String, Vec<(u128,u128)>>,
+    instruction_times: HashMap<String, Vec<(u128, u128)>>,
     core_start_time: HashMap<usize, Option<Instant>>,
     start_time: Instant,
 }
@@ -51,7 +51,7 @@ impl Trace {
         Self {
             instruction_times,
             core_start_time: HashMap::new(),
-            start_time: Instant::now()
+            start_time: Instant::now(),
         }
     }
 

backend-simulators/pim/pim-simulator/src/lib/tracing/profile/profile_analysis.rs

@@ -145,17 +145,15 @@ pub fn print_textual_report(stats: &[InstructionStats]) {
     println!("{table}");
 }
 
-
 pub fn generate_interactive_report(
     timings: &HashMap<String, Vec<(u128, u128)>>,
     instructions_to_plot: &[&str], // <-- NEW: Only plot these
     file_path: &str,
 ) {
-
-use plotly::common::{Mode, Marker, Line};
-use plotly::layout::{Axis, Layout};
-use plotly::{Plot, Scatter};
-use std::collections::HashMap;
+    use plotly::common::{Line, Marker, Mode};
+    use plotly::layout::{Axis, Layout};
+    use plotly::{Plot, Scatter};
+    use std::collections::HashMap;
     let mut plot = Plot::new();
 
     for &instruction_name in instructions_to_plot {
@@ -163,8 +161,9 @@ use std::collections::HashMap;
         if let Some(times) = timings.get(instruction_name) {
             let x_axis: Vec<f64> = times.iter().map(|&(ts, _)| ts as f64).collect();
             let y_axis: Vec<f64> = times.iter().map(|&(_, dur)| dur as f64).collect();
-            
-            let text_array: Vec<String> = times.iter()
+
+            let text_array: Vec<String> = times
+                .iter()
                 .map(|&(_, dur)| format_time(dur as f64))
                 .collect();
 
@@ -181,7 +180,9 @@ use std::collections::HashMap;
     }
 
     let layout = Layout::new()
-        .title(plotly::common::Title::new("Simulator Timeline: Top Offenders"))
+        .title(plotly::common::Title::new(
+            "Simulator Timeline: Top Offenders",
+        ))
         .x_axis(Axis::new().title(plotly::common::Title::new("Absolute Time (ns)")))
         .y_axis(Axis::new().title(plotly::common::Title::new("Execution Duration")));
 
@@ -189,4 +190,3 @@ use std::collections::HashMap;
     plot.write_html(file_path);
     println!("🌐 Interactive timeline saved to {}", file_path);
 }
-

backend-simulators/pim/pim-simulator/src/lib/tracing/profile/profile_isa.rs

@@ -34,10 +34,11 @@ impl Trace {
             start_time,
         } = self;
         let now = Instant::now();
-        instruction_times
-            .get_mut(name)
-            .unwrap()
-            .push((now.duration_since(*start_time).as_nanos(), now.duration_since(core_start_time[&core_indx].unwrap()).as_nanos()));
+        instruction_times.get_mut(name).unwrap().push((
+            now.duration_since(*start_time).as_nanos(),
+            now.duration_since(core_start_time[&core_indx].unwrap())
+                .as_nanos(),
+        ));
         self.core_start_time.insert(core_indx, None);
     }
 

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

@@ -7,7 +7,6 @@ pub struct Trace {
     out_files: Vec<File>,
 }
 
-
 impl Trace {
     pub fn new() -> Self {
         Self {
@@ -25,4 +24,3 @@ impl Trace {
         }
     }
 }
-

backend-simulators/pim/pim-simulator/src/lib/tracing/trace/tracing_isa.rs

@@ -283,7 +283,6 @@ impl Trace {
         M: UpcastDestTraits<M> + MemoryStorable + FromFloat,
         F: UpcastDestTraits<F> + MemoryStorable,
     {
-
         let (core_indx, rd, r1, mbiw, relu, group) = data.get_core_rd_r1_mbiw_immrelu_immgroup();
         let file: &mut File = self
             .out_files
@@ -337,7 +336,7 @@ impl Trace {
         );
         pretty_print::print_slice::<_, F>(file, loads[0], 30);
         writeln!(file, "\tCrossbar[{}:{}](B): ", 0, crossbar_byte_size);
-        pretty_print::print_slice::<_,F>(file, matrix, 30);
+        pretty_print::print_slice::<_, F>(file, matrix, 30);
         writeln!(
             file,
             "\tLocal[{}:{}](out): ",
@@ -409,11 +408,11 @@ impl Trace {
             .unwrap();
         writeln!(file, "{} Memory:", prefix);
         write!(file, "\tLocal[{}:{}](A): ", r1_final, r1_final + byte_len);
-        pretty_print::print_slice::<_,f32>(file, loads[0], 30);
-        write!(file, "\tLocal[{}:{}](B): ", r2_final , r2_final + byte_len);
-        pretty_print::print_slice::<_,f32>(file, loads[1], 30);
-        write!(file, "\tLocal[{}:{}](out): ", rd_final, rd_final+ byte_len);
-        pretty_print::print_slice::<_,f32>(file, loads[2], 30);
+        pretty_print::print_slice::<_, f32>(file, loads[0], 30);
+        write!(file, "\tLocal[{}:{}](B): ", r2_final, r2_final + byte_len);
+        pretty_print::print_slice::<_, f32>(file, loads[1], 30);
+        write!(file, "\tLocal[{}:{}](out): ", rd_final, rd_final + byte_len);
+        pretty_print::print_slice::<_, f32>(file, loads[2], 30);
         if prefix == "Post" {
             writeln!(file, "\n###############################################\n");
         }
@@ -1027,9 +1026,9 @@ impl Trace {
         let core_memory = core.load::<u8>(rd_val, imm_len).unwrap();
         writeln!(file, "{} Memory:", prefix);
         writeln!(file, "\tHost[{}:{}]: ", r1_val, r1_val + imm_len as usize,);
-        pretty_print::print_slice::<_,f32>(file, global_memory[0], 30);
+        pretty_print::print_slice::<_, f32>(file, global_memory[0], 30);
         writeln!(file, "\tLocal[{}:{}]: ", rd_val, rd_val + imm_len as usize,);
-        pretty_print::print_slice::<_,f32>(file, core_memory[0], 30);
+        pretty_print::print_slice::<_, f32>(file, core_memory[0], 30);
 
         if prefix == "Post" {
             writeln!(file, "\n###############################################\n");
@@ -1079,9 +1078,9 @@ impl Trace {
         let global_memory = host.load::<u8>(rd_val, imm_len).unwrap();
         writeln!(file, "{} Memory:", prefix);
         writeln!(file, "\tLocal[{}:{}]: ", r1_val, r1_val + imm_len as usize,);
-        pretty_print::print_slice::<_,f32>(file, core_memory[0], 30);
+        pretty_print::print_slice::<_, f32>(file, core_memory[0], 30);
         writeln!(file, "\tHost[{}:{}]: ", rd_val, rd_val + imm_len as usize,);
-        pretty_print::print_slice::<_,f32>(file, global_memory[0], 30);
+        pretty_print::print_slice::<_, f32>(file, global_memory[0], 30);
 
         if prefix == "Post" {
             writeln!(file, "\n###############################################\n");
@@ -1096,7 +1095,6 @@ impl Trace {
         self.st_impl(cores, data, "Pre");
     }
 
-
     pub fn pre_lldi(&mut self, cores: &mut CPU, data: InstructionData) {
         let (core, rd, imm) = data.get_core_rd_imm();
         let file: &mut File = self
@@ -1136,8 +1134,7 @@ impl Trace {
         //    Ok(InstructionStatus::Completed)
     }
 
-    fn lmv_impl (&mut self, cores: &mut CPU, data: InstructionData, prefix: &'static str) {
-
+    fn lmv_impl(&mut self, cores: &mut CPU, data: InstructionData, prefix: &'static str) {
         let (core, rd, r1, _, imm_len, offset_select, offset_value) =
             data.get_core_rd_r1_r2_immlen_offset();
         let file: &mut File = self
@@ -1169,14 +1166,17 @@ impl Trace {
         let r1_val = add_offset_r1(r1_val, offset_select, offset_value);
         let rd_val = add_offset_rd(rd_val, offset_select, offset_value);
         let core_memory = core
-            .reserve_load(r1_val, imm_len).unwrap()
-            .reserve_load(rd_val, imm_len).unwrap()
-            .execute_load::<u8>().unwrap();
+            .reserve_load(r1_val, imm_len)
+            .unwrap()
+            .reserve_load(rd_val, imm_len)
+            .unwrap()
+            .execute_load::<u8>()
+            .unwrap();
         writeln!(file, "{} Memory:", prefix);
         writeln!(file, "\tLocal[{}:{}]: ", r1_val, r1_val + imm_len as usize,);
-        pretty_print::print_slice::<_,f32>(file, core_memory[0], 30);
+        pretty_print::print_slice::<_, f32>(file, core_memory[0], 30);
         writeln!(file, "\tLocal[{}:{}]: ", rd_val, rd_val + imm_len as usize,);
-        pretty_print::print_slice::<_,f32>(file, core_memory[1], 30);
+        pretty_print::print_slice::<_, f32>(file, core_memory[1], 30);
 
         if prefix == "Post" {
             writeln!(file, "\n###############################################\n");

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

@@ -1,7 +1,6 @@
-use anyhow::{Result,Context};
+use anyhow::{Context, Result};
 use std::{fmt::Debug, mem::transmute};
 
-
 pub trait AddressArg {
     fn to_address_usize(self) -> Result<usize>;
 }
@@ -36,42 +35,56 @@ impl AddressArg for i64 {
     }
 }
 
-
 fn address_to_usize(address: i32) -> usize {
     address as u32 as usize
 }
 
-fn add_offset_impl(address: usize, offset_select : i32, offset_value : i32, id:i32) -> usize{
-     assert!(offset_select == 1 || offset_select == 2 || offset_select == 4 || offset_value == 0, "offset_select not a bit field");
-     let offset_value = (offset_select & id) * offset_value;
-     if offset_value > 0 {
-        address + offset_value as usize
-     } else {
-         address - offset_value as usize
-     }
+fn add_offset_impl(address: usize, offset_select: i32, offset_value: i32, id: i32) -> usize {
+    assert!(
+        (0..=7).contains(&offset_select),
+        "offset_select is not a 3-bit field"
+    );
+    if offset_select & id == 0 {
+        address
+    } else {
+        address
+            .checked_add_signed(offset_value as isize)
+            .expect("address offset overflow")
+    }
 }
 
-
-pub fn add_offset_rd(address: i32, offset_select : i32, offset_value : i32) -> usize
-{
-    let address = address_to_usize(address);
-    add_offset_impl(address, offset_select, offset_value, 4)
-}
-
-pub fn add_offset_r1(address: i32, offset_select : i32, offset_value : i32) -> usize
-{
+pub fn add_offset_rd(address: i32, offset_select: i32, offset_value: i32) -> usize {
     let address = address_to_usize(address);
     add_offset_impl(address, offset_select, offset_value, 1)
 }
 
-pub fn add_offset_r2(address: i32, offset_select : i32, offset_value : i32) -> usize
-{
+pub fn add_offset_r1(address: i32, offset_select: i32, offset_value: i32) -> usize {
     let address = address_to_usize(address);
     add_offset_impl(address, offset_select, offset_value, 2)
 }
 
-
-pub fn pack_float_in_i32(val : impl TryInto<f32>) -> i32 {
-    let val = val.try_into().unwrap_or_else( |x| panic!("Cannot parse into f32"));
-     f32::to_bits(val).cast_signed()
+pub fn add_offset_r2(address: i32, offset_select: i32, offset_value: i32) -> usize {
+    let address = address_to_usize(address);
+    add_offset_impl(address, offset_select, offset_value, 4)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::{add_offset_r1, add_offset_r2, add_offset_rd};
+
+    #[test]
+    fn offset_select_uses_rd_rs1_rs2_bit_order() {
+        assert_eq!(add_offset_rd(100, 1, 7), 107);
+        assert_eq!(add_offset_r1(100, 2, 7), 107);
+        assert_eq!(add_offset_r2(100, 4, 7), 107);
+        assert_eq!(add_offset_rd(100, 6, 7), 100);
+        assert_eq!(add_offset_r1(100, 7, -7), 93);
+    }
+}
+
+pub fn pack_float_in_i32(val: impl TryInto<f32>) -> i32 {
+    let val = val
+        .try_into()
+        .unwrap_or_else(|x| panic!("Cannot parse into f32"));
+    f32::to_bits(val).cast_signed()
 }

backend-simulators/pim/pim-simulator/tests/batch.rs

@@ -0,0 +1,70 @@
+mod common;
+
+use pimcore::{
+    CoreInstructionsBuilder, Executable,
+    instruction_set::{InstructionsBuilder, instruction_data::InstructionDataBuilder, isa::*},
+};
+
+#[test]
+fn restarts_cores_and_loads_each_input() {
+    let cpu = common::empty_cpu(1);
+    let mut cores = CoreInstructionsBuilder::new(1);
+    let mut instructions = InstructionsBuilder::new();
+    let mut data = InstructionDataBuilder::new();
+    data.set_core_indx(1).fix_core_indx();
+    instructions.make_inst(sldi, data.set_rdimm(1, 0).build());
+    instructions.make_inst(sldi, data.set_rdimm(2, 0).build());
+    instructions.make_inst(ld, data.set_rdr1(2, 1).set_imm_len(4).build());
+    instructions.make_inst(sldi, data.set_rdimm(3, 4).build());
+    instructions.make_inst(st, data.set_rdr1(3, 2).set_imm_len(4).build());
+    cores.set_core(1, instructions.build());
+
+    let mut executable = Executable::new(cpu, cores.build());
+    let first = 1.0f32.to_ne_bytes();
+    let second = 2.0f32.to_ne_bytes();
+    assert!(
+        executable
+            .execute_batch(&[&first[..3]], &[(0, 4)], &[])
+            .is_err()
+    );
+    executable
+        .execute_batch(&[&first, &second], &[(0, 4)], &[])
+        .unwrap();
+
+    assert_eq!(
+        executable.cpu_mut().host().load::<f32>(4, 4).unwrap()[0],
+        [2.0]
+    );
+}
+
+#[test]
+fn records_each_iteration_output() {
+    let cpu = common::empty_cpu(1);
+    let mut cores = CoreInstructionsBuilder::new(1);
+    let mut instructions = InstructionsBuilder::new();
+    let mut data = InstructionDataBuilder::new();
+    data.set_core_indx(1).fix_core_indx();
+    instructions.make_inst(sldi, data.set_rdimm(1, 0).build());
+    instructions.make_inst(sldi, data.set_rdimm(2, 0).build());
+    instructions.make_inst(ld, data.set_rdr1(2, 1).set_imm_len(4).build());
+    instructions.make_inst(sldi, data.set_rdimm(3, 4).build());
+    instructions.make_inst(st, data.set_rdr1(3, 2).set_imm_len(4).build());
+    cores.set_core(1, instructions.build());
+
+    let mut executable = Executable::new(cpu, cores.build());
+    let first = 1.0f32.to_ne_bytes();
+    let second = 2.0f32.to_ne_bytes();
+    let outputs = executable
+        .execute_batch(&[&first, &second], &[(0, 4)], &[(4, 2), (6, 2)])
+        .unwrap();
+
+    assert_eq!(outputs.len(), 2);
+    assert_eq!(
+        f32::from_ne_bytes(outputs[0].as_slice().try_into().unwrap()),
+        1.0
+    );
+    assert_eq!(
+        f32::from_ne_bytes(outputs[1].as_slice().try_into().unwrap()),
+        2.0
+    );
+}

backend-simulators/pim/pim-simulator/tests/big_mul.rs

@@ -7,7 +7,7 @@ use pimcore::{
     memory_manager::CoreMemory,
 };
 
-fn simple_read(path:  &Path) -> Vec<f32> {
+fn simple_read(path: &Path) -> Vec<f32> {
     if !path.exists() {
         panic!("{:?} not exists", path)
     }
@@ -19,9 +19,7 @@ fn simple_read(path:  &Path) -> Vec<f32> {
 }
 
 /// mvmul Test
-fn mvmul_f32(err: &str)
-where
-{
+fn mvmul_f32(err: &str) {
     let matrix = simple_read(Path::new("tests/B.txt"));
     let mut crossbar = Crossbar::new(1024 * size_of::<f32>(), 1024, CoreMemory::new());
     crossbar.execute_store(&matrix).unwrap();
@@ -36,7 +34,9 @@ where
     inst_builder.make_inst(sldi, idata_build.set_rdimm(1, 0).build());
     inst_builder.make_inst(
         sldi,
-        idata_build.set_rdimm(3, 1024 * size_of::<f32>() as i32).build(),
+        idata_build
+            .set_rdimm(3, 1024 * size_of::<f32>() as i32)
+            .build(),
     );
     inst_builder.make_inst(
         setbw,
@@ -48,7 +48,7 @@ where
         mvmul,
         idata_build
             .set_rdr1(3, 1)
-            .set_mbiw_immrelu_immgroup(8*size_of::<f32>() as i32, 0, 0)
+            .set_mbiw_immrelu_immgroup(8 * size_of::<f32>() as i32, 0, 0)
             .build(),
     );
     let core_instruction = vec![inst_builder.build().into()];
@@ -59,8 +59,11 @@ where
         executable
             .cpu_mut()
             .host()
-            .load::<f32>(1024 * size_of::<f32>(), 1024*size_of::<f32>()).unwrap()[0].iter().zip(
-        simple_read(Path::new("tests/X.txt")) ).all(|(&a,b) : (&f32, f32)| {a-b < 0.001}),
+            .load::<f32>(1024 * size_of::<f32>(), 1024 * size_of::<f32>())
+            .unwrap()[0]
+            .iter()
+            .zip(simple_read(Path::new("tests/X.txt")))
+            .all(|(&a, b): (&f32, f32)| { a - b < 0.001 }),
         "Wrong result for {}",
         err
     );
@@ -69,6 +72,4 @@ where
 #[test]
 fn mvmul_big_test() {
     mvmul_f32("mvmul_f32");
-   
-
 }

backend-simulators/pim/pim-simulator/tests/es_runner.rs

@@ -6,9 +6,7 @@ use std::{
 
 use anyhow::{Context, Result};
 use pimcore::{
-    cpu::crossbar::Crossbar,
-    json_to_instruction::json_to_executor,
-    memory_manager::CoreMemory,
+    cpu::crossbar::Crossbar, json_to_instruction::json_to_executor, memory_manager::CoreMemory,
 };
 use serde_json::Value;
 
@@ -95,9 +93,14 @@ fn json_folder_tester() {
             .map(|core_crossbars| core_crossbars.iter().collect())
             .collect();
 
-        let mut executable = json_to_executor::json_to_executor(config, &mut core_readers, crossbars);
+        let mut executable =
+            json_to_executor::json_to_executor(config, &mut core_readers, crossbars);
         let memory = fs::read(folder.join("memory.bin")).unwrap();
-        executable.cpu_mut().host().execute_store(0, &memory).unwrap();
+        executable
+            .cpu_mut()
+            .host()
+            .execute_store(0, &memory)
+            .unwrap();
         executable.execute();
     }
 }

backend-simulators/pim/pim-simulator/tests/placeholder.rs

@@ -7,26 +7,17 @@ use pimcore::{
     },
 };
 
-
 #[test]
-#[should_panic(expected = "Function not found for the requested size") ]
+#[should_panic(expected = "Function not found for the requested size")]
 fn wrong_size_place_holder() {
     let cpu = common::empty_cpu(0);
     let mut inst_builder = InstructionsBuilder::new();
     let mut idata_build = InstructionDataBuilder::new();
     idata_build.set_core_indx(0).fix_core_indx();
-    inst_builder.make_inst(
-        setbw,
-        idata_build
-            .set_ibiw_obiw(55, 55)
-            .build(),
-    );
+    inst_builder.make_inst(setbw, idata_build.set_ibiw_obiw(55, 55).build());
     inst_builder.make_inst(
         vvadd,
-        idata_build
-            .set_rdr1r2(3, 1, 2)
-            .set_imm_len(8)
-            .build(),
+        idata_build.set_rdr1r2(3, 1, 2).set_imm_len(8).build(),
     );
     let core_instruction = vec![inst_builder.build().into()];
     let mut executable = Executable::new(cpu, core_instruction);
@@ -39,97 +30,88 @@ fn unsupported_8_bit_vectors_do_not_alias_f32() {
     let mut inst_builder = InstructionsBuilder::new();
     let mut idata_build = InstructionDataBuilder::new();
     idata_build.set_core_indx(0).fix_core_indx();
-    inst_builder.make_inst(
-        setbw,
-        idata_build.set_ibiw_obiw(8, 8).build(),
-    );
+    inst_builder.make_inst(setbw, idata_build.set_ibiw_obiw(8, 8).build());
     inst_builder.make_inst(
         vvadd,
-        idata_build
-            .set_rdr1r2(3, 1, 2)
-            .set_imm_len(8)
-            .build(),
+        idata_build.set_rdr1r2(3, 1, 2).set_imm_len(8).build(),
     );
 }
 
-
-
-fn  place_holder(inst :  InstructionType) {
+fn place_holder(inst: InstructionType) {
     let mut cpu = common::empty_cpu(0);
     let mut idata_build = InstructionDataBuilder::new();
     idata_build.set_core_indx(0).fix_core_indx();
     inst(&mut cpu, idata_build.build()).unwrap();
 }
 
-
 #[test]
-#[should_panic(expected = "You are calling a placeholder, the real call is the generic version") ]
+#[should_panic(expected = "You are calling a placeholder, the real call is the generic version")]
 fn vvadd_placeholder() {
     place_holder(vvadd);
 }
 
 #[test]
-#[should_panic(expected = "You are calling a placeholder, the real call is the generic version") ]
+#[should_panic(expected = "You are calling a placeholder, the real call is the generic version")]
 fn vvsub_placeholder() {
     place_holder(vvsub);
 }
 
 #[test]
-#[should_panic(expected = "You are calling a placeholder, the real call is the generic version") ]
+#[should_panic(expected = "You are calling a placeholder, the real call is the generic version")]
 fn vvmul_placeholder() {
     place_holder(vvmul);
 }
 
 #[test]
-#[should_panic(expected = "You are calling a placeholder, the real call is the generic version") ]
+#[should_panic(expected = "You are calling a placeholder, the real call is the generic version")]
 fn vvdmul_placeholder() {
     place_holder(vvdmul);
 }
 
 #[test]
-#[should_panic(expected = "You are calling a placeholder, the real call is the generic version") ]
+#[should_panic(expected = "You are calling a placeholder, the real call is the generic version")]
 fn vvmax_placeholder() {
     place_holder(vvmax);
 }
 
 #[test]
-#[should_panic(expected = "You are calling a placeholder, the real call is the generic version") ]
+#[should_panic(expected = "You are calling a placeholder, the real call is the generic version")]
 fn vavg_placeholder() {
     place_holder(vavg);
 }
 
 #[test]
-#[should_panic(expected = "You are calling a placeholder, the real call is the generic version") ]
+#[should_panic(expected = "You are calling a placeholder, the real call is the generic version")]
 fn vrelu_placeholder() {
     place_holder(vrelu);
 }
 
 #[test]
-#[should_panic(expected = "You are calling a placeholder, the real call is the generic version") ]
+#[should_panic(expected = "You are calling a placeholder, the real call is the generic version")]
 fn vtanh_placeholder() {
     place_holder(vtanh);
 }
 
 #[test]
-#[should_panic(expected = "You are calling a placeholder, the real call is the generic version") ]
+#[should_panic(expected = "You are calling a placeholder, the real call is the generic version")]
 fn vsigm_placeholder() {
     place_holder(vsigm);
 }
 
 #[test]
-#[should_panic(expected = "You are calling a placeholder, the real call is the generic version") ]
+#[should_panic(expected = "You are calling a placeholder, the real call is the generic version")]
 fn mvmul_placeholder() {
     place_holder(mvmul);
 }
 
 #[test]
-#[should_panic ]
+#[should_panic]
 fn vvsll_why_inst() {
     place_holder(vvsll);
 }
 
 #[test]
-#[should_panic ]
+#[should_panic]
 fn vvsra_why_inst() {
     place_holder(vvsra);
 }

backend-simulators/pim/pim-simulator/tests/simd.rs

@@ -53,10 +53,7 @@ where
     );
     inst_builder.make_inst(
         vvadd,
-        idata_build
-            .set_rdr1r2(3, 1, 2)
-            .set_imm_len(8)
-            .build(),
+        idata_build.set_rdr1r2(3, 1, 2).set_imm_len(8).build(),
     );
     let core_instruction = vec![inst_builder.build().into()];
     let mut executable = Executable::new(cpu, core_instruction);
@@ -67,7 +64,8 @@ where
         executable
             .cpu_mut()
             .host()
-            .load::<T>(16 * size_of::<F>(), 8 * size_of::<T>()).unwrap()[0],
+            .load::<T>(16 * size_of::<F>(), 8 * size_of::<T>())
+            .unwrap()[0],
         vec![
             10.0.into(),
             12.0.into(),
@@ -86,17 +84,22 @@ where
         executable
             .cpu_mut()
             .host()
-            .load::<F>(0, 16 * size_of::<F>()).unwrap()[0],
+            .load::<F>(0, 16 * size_of::<F>())
+            .unwrap()[0],
         &buff,
         "Altered first part for {}",
         err
     );
     //Check that later is 0
     assert_eq!(
-        executable.cpu_mut().host().load::<i32>(
-            16 * size_of::<F>() + 8 * size_of::<T>(),
-            4 * size_of::<i32>()
-        ).unwrap()[0],
+        executable
+            .cpu_mut()
+            .host()
+            .load::<i32>(
+                16 * size_of::<F>() + 8 * size_of::<T>(),
+                4 * size_of::<i32>()
+            )
+            .unwrap()[0],
         [0, 0, 0, 0],
         "Altered first part for {}",
         err
@@ -157,10 +160,7 @@ where
     );
     inst_builder.make_inst(
         vvsub,
-        idata_build
-            .set_rdr1r2(3, 1, 2)
-            .set_imm_len(8)
-            .build(),
+        idata_build.set_rdr1r2(3, 1, 2).set_imm_len(8).build(),
     );
     let core_instruction = vec![inst_builder.build().into()];
     let mut executable = Executable::new(cpu, core_instruction);
@@ -171,7 +171,8 @@ where
         executable
             .cpu_mut()
             .host()
-            .load::<T>(16 * size_of::<F>(), 8 * size_of::<T>()).unwrap()[0],
+            .load::<T>(16 * size_of::<F>(), 8 * size_of::<T>())
+            .unwrap()[0],
         vec![
             (-8.0).into(),
             (-8.0).into(),
@@ -190,17 +191,22 @@ where
         executable
             .cpu_mut()
             .host()
-            .load::<F>(0, 16 * size_of::<F>()).unwrap()[0],
+            .load::<F>(0, 16 * size_of::<F>())
+            .unwrap()[0],
         &buff,
         "Altered first part for {}",
         err
     );
     //Check that later is 0
     assert_eq!(
-        executable.cpu_mut().host().load::<i32>(
-            16 * size_of::<F>() + 8 * size_of::<T>(),
-            4 * size_of::<i32>()
-        ).unwrap()[0],
+        executable
+            .cpu_mut()
+            .host()
+            .load::<i32>(
+                16 * size_of::<F>() + 8 * size_of::<T>(),
+                4 * size_of::<i32>()
+            )
+            .unwrap()[0],
         [0, 0, 0, 0],
         "Altered first part for {}",
         err
@@ -261,10 +267,7 @@ where
     );
     inst_builder.make_inst(
         vvmul,
-        idata_build
-            .set_rdr1r2(3, 1, 2)
-            .set_imm_len(8)
-            .build(),
+        idata_build.set_rdr1r2(3, 1, 2).set_imm_len(8).build(),
     );
     let core_instruction = vec![inst_builder.build().into()];
     let mut executable = Executable::new(cpu, core_instruction);
@@ -275,7 +278,8 @@ where
         executable
             .cpu_mut()
             .host()
-            .load::<T>(16 * size_of::<F>(), 8 * size_of::<T>()).unwrap()[0],
+            .load::<T>(16 * size_of::<F>(), 8 * size_of::<T>())
+            .unwrap()[0],
         vec![
             (9.0).into(),
             (20.0).into(),
@@ -294,17 +298,22 @@ where
         executable
             .cpu_mut()
             .host()
-            .load::<F>(0, 16 * size_of::<F>()).unwrap()[0],
+            .load::<F>(0, 16 * size_of::<F>())
+            .unwrap()[0],
         &buff,
         "Altered first part for {}",
         err
     );
     //Check that later is 0
     assert_eq!(
-        executable.cpu_mut().host().load::<i32>(
-            16 * size_of::<F>() + 8 * size_of::<T>(),
-            4 * size_of::<i32>()
-        ).unwrap()[0],
+        executable
+            .cpu_mut()
+            .host()
+            .load::<i32>(
+                16 * size_of::<F>() + 8 * size_of::<T>(),
+                4 * size_of::<i32>()
+            )
+            .unwrap()[0],
         [0, 0, 0, 0],
         "Altered first part for {}",
         err
@@ -365,10 +374,7 @@ where
     );
     inst_builder.make_inst(
         vvdmul,
-        idata_build
-            .set_rdr1r2(3, 1, 2)
-            .set_imm_len(8)
-            .build(),
+        idata_build.set_rdr1r2(3, 1, 2).set_imm_len(8).build(),
     );
     let core_instruction = vec![inst_builder.build().into()];
     let mut executable = Executable::new(cpu, core_instruction);
@@ -379,10 +385,9 @@ where
         executable
             .cpu_mut()
             .host()
-            .load::<T>(16 * size_of::<F>(), size_of::<T>()).unwrap()[0],
-        vec![
-            (492.0).into(),
-        ],
+            .load::<T>(16 * size_of::<F>(), size_of::<T>())
+            .unwrap()[0],
+        vec![(492.0).into(),],
         "Wrong result for {}",
         err
     );
@@ -391,17 +396,19 @@ where
         executable
             .cpu_mut()
             .host()
-            .load::<F>(0, 16 * size_of::<F>()).unwrap()[0],
+            .load::<F>(0, 16 * size_of::<F>())
+            .unwrap()[0],
         &buff,
         "Altered first part for {}",
         err
     );
     //Check that later is 0
     assert_eq!(
-        executable.cpu_mut().host().load::<i32>(
-            16 * size_of::<F>() + size_of::<T>(),
-            4 * size_of::<i32>()
-        ).unwrap()[0],
+        executable
+            .cpu_mut()
+            .host()
+            .load::<i32>(16 * size_of::<F>() + size_of::<T>(), 4 * size_of::<i32>())
+            .unwrap()[0],
         [0, 0, 0, 0],
         "Altered first part for {}",
         err
@@ -462,10 +469,7 @@ where
     );
     inst_builder.make_inst(
         vvmax,
-        idata_build
-            .set_rdr1r2(3, 1, 2)
-            .set_imm_len(8)
-            .build(),
+        idata_build.set_rdr1r2(3, 1, 2).set_imm_len(8).build(),
     );
     let core_instruction = vec![inst_builder.build().into()];
     let mut executable = Executable::new(cpu, core_instruction);
@@ -476,16 +480,17 @@ where
         executable
             .cpu_mut()
             .host()
-            .load::<T>(16 * size_of::<F>(), 8 * size_of::<T>()).unwrap()[0],
+            .load::<T>(16 * size_of::<F>(), 8 * size_of::<T>())
+            .unwrap()[0],
         vec![
-        9.0.into(),
-        10.0.into(),
-        11.0.into(),
-        12.0.into(),
-        13.0.into(),
-        14.0.into(),
-        15.0.into(),
-        16.0.into(),
+            9.0.into(),
+            10.0.into(),
+            11.0.into(),
+            12.0.into(),
+            13.0.into(),
+            14.0.into(),
+            15.0.into(),
+            16.0.into(),
         ],
         "Wrong result for {}",
         err
@@ -495,17 +500,22 @@ where
         executable
             .cpu_mut()
             .host()
-            .load::<F>(0, 16 * size_of::<F>()).unwrap()[0],
+            .load::<F>(0, 16 * size_of::<F>())
+            .unwrap()[0],
         &buff,
         "Altered first part for {}",
         err
     );
     //Check that later is 0
     assert_eq!(
-        executable.cpu_mut().host().load::<i32>(
-            16 * size_of::<F>() + 8 * size_of::<T>(),
-            4 * size_of::<i32>()
-        ).unwrap()[0],
+        executable
+            .cpu_mut()
+            .host()
+            .load::<i32>(
+                16 * size_of::<F>() + 8 * size_of::<T>(),
+                4 * size_of::<i32>()
+            )
+            .unwrap()[0],
         [0, 0, 0, 0],
         "Altered first part for {}",
         err
@@ -577,10 +587,9 @@ where
         executable
             .cpu_mut()
             .host()
-            .load::<T>(16 * size_of::<F>(), size_of::<T>()).unwrap()[0],
-        vec![
-        7.5.into(),
-        ],
+            .load::<T>(16 * size_of::<F>(), size_of::<T>())
+            .unwrap()[0],
+        vec![7.5.into(),],
         "Wrong result for {}",
         err
     );
@@ -589,17 +598,19 @@ where
         executable
             .cpu_mut()
             .host()
-            .load::<F>(0, 16 * size_of::<F>()).unwrap()[0],
+            .load::<F>(0, 16 * size_of::<F>())
+            .unwrap()[0],
         &buff,
         "Altered first part for {}",
         err
     );
     //Check that later is 0
     assert_eq!(
-        executable.cpu_mut().host().load::<i32>(
-            16 * size_of::<F>() +  size_of::<T>(),
-            4 * size_of::<i32>()
-        ).unwrap()[0],
+        executable
+            .cpu_mut()
+            .host()
+            .load::<i32>(16 * size_of::<F>() + size_of::<T>(), 4 * size_of::<i32>())
+            .unwrap()[0],
         [0, 0, 0, 0],
         "Altered first part for {}",
         err
@@ -656,10 +667,7 @@ where
     );
     inst_builder.make_inst(
         vrelu,
-        idata_build
-            .set_rdr1r2(3, 1, 1)
-            .set_imm_len(8)
-            .build(),
+        idata_build.set_rdr1r2(3, 1, 1).set_imm_len(8).build(),
     );
     let core_instruction = vec![inst_builder.build().into()];
     let mut executable = Executable::new(cpu, core_instruction);
@@ -670,16 +678,17 @@ where
         executable
             .cpu_mut()
             .host()
-            .load::<T>(16 * size_of::<F>(), 8*size_of::<T>()).unwrap()[0],
+            .load::<T>(16 * size_of::<F>(), 8 * size_of::<T>())
+            .unwrap()[0],
         vec![
-        0.0.into(),
-        2.0.into(),
-        11.0.into(),
-        0.0.into(),
-        13.0.into(),
-        0.0.into(),
-        7.0.into(),
-        0.0.into(),
+            0.0.into(),
+            2.0.into(),
+            11.0.into(),
+            0.0.into(),
+            13.0.into(),
+            0.0.into(),
+            7.0.into(),
+            0.0.into(),
         ],
         "Wrong result for {}",
         err
@@ -689,17 +698,22 @@ where
         executable
             .cpu_mut()
             .host()
-            .load::<F>(0, 16 * size_of::<F>()).unwrap()[0],
+            .load::<F>(0, 16 * size_of::<F>())
+            .unwrap()[0],
         &buff,
         "Altered first part for {}",
         err
     );
     //Check that later is 0
     assert_eq!(
-        executable.cpu_mut().host().load::<i32>(
-            16 * size_of::<F>() +  8*size_of::<T>(),
-            4 * size_of::<i32>()
-        ).unwrap()[0],
+        executable
+            .cpu_mut()
+            .host()
+            .load::<i32>(
+                16 * size_of::<F>() + 8 * size_of::<T>(),
+                4 * size_of::<i32>()
+            )
+            .unwrap()[0],
         [0, 0, 0, 0],
         "Altered first part for {}",
         err
@@ -756,32 +770,32 @@ where
     );
     inst_builder.make_inst(
         vtanh,
-        idata_build
-            .set_rdr1r2(3, 1, 1)
-            .set_imm_len(8)
-            .build(),
+        idata_build.set_rdr1r2(3, 1, 1).set_imm_len(8).build(),
     );
     let core_instruction = vec![inst_builder.build().into()];
     let mut executable = Executable::new(cpu, core_instruction);
     executable.execute();
 
     // Check result correct
-    
+
     assert!(
         executable
             .cpu_mut()
             .host()
-            .load::<T>(16 * size_of::<F>(), 8*size_of::<T>()).unwrap()[0].iter().zip(
-        vec![
-        T::from(0.1).tanh(),
-        T::from(0.2).tanh(),
-        T::from(0.3).tanh(),
-        T::from(0.4).tanh(),
-        T::from(0.5).tanh(),
-        T::from(0.6).tanh(),
-        T::from(0.7).tanh(),
-        T::from(0.8).tanh(),
-        ]).all(|(&a,b) : (&T, T)| {a-b < 0.001.into()}),
+            .load::<T>(16 * size_of::<F>(), 8 * size_of::<T>())
+            .unwrap()[0]
+            .iter()
+            .zip(vec![
+                T::from(0.1).tanh(),
+                T::from(0.2).tanh(),
+                T::from(0.3).tanh(),
+                T::from(0.4).tanh(),
+                T::from(0.5).tanh(),
+                T::from(0.6).tanh(),
+                T::from(0.7).tanh(),
+                T::from(0.8).tanh(),
+            ])
+            .all(|(&a, b): (&T, T)| { a - b < 0.001.into() }),
         "Wrong result for {}",
         err
     );
@@ -790,17 +804,22 @@ where
         executable
             .cpu_mut()
             .host()
-            .load::<F>(0, 16 * size_of::<F>()).unwrap()[0],
+            .load::<F>(0, 16 * size_of::<F>())
+            .unwrap()[0],
         &buff,
         "Altered first part for {}",
         err
     );
     //Check that later is 0
     assert_eq!(
-        executable.cpu_mut().host().load::<i32>(
-            16 * size_of::<F>() +  8*size_of::<T>(),
-            4 * size_of::<i32>()
-        ).unwrap()[0],
+        executable
+            .cpu_mut()
+            .host()
+            .load::<i32>(
+                16 * size_of::<F>() + 8 * size_of::<T>(),
+                4 * size_of::<i32>()
+            )
+            .unwrap()[0],
         [0, 0, 0, 0],
         "Altered first part for {}",
         err
@@ -815,7 +834,6 @@ fn vtanh_test() {
     vtanh_test_generic::<f64, f64>("vtanh<f64,f64>");
 }
 
-
 /// vsigm Test
 fn vsigm_test_generic<F, T>(err: &str)
 where
@@ -858,32 +876,32 @@ where
     );
     inst_builder.make_inst(
         vsigm,
-        idata_build
-            .set_rdr1r2(3, 1, 1)
-            .set_imm_len(8)
-            .build(),
+        idata_build.set_rdr1r2(3, 1, 1).set_imm_len(8).build(),
     );
     let core_instruction = vec![inst_builder.build().into()];
     let mut executable = Executable::new(cpu, core_instruction);
     executable.execute();
 
     // Check result correct
-    
+
     assert!(
         executable
             .cpu_mut()
             .host()
-            .load::<T>(16 * size_of::<F>(), 8*size_of::<T>()).unwrap()[0].iter().zip(
-        vec![
-        T::from(0.1).sigm(),
-        T::from(0.2).sigm(),
-        T::from(0.3).sigm(),
-        T::from(0.4).sigm(),
-        T::from(0.5).sigm(),
-        T::from(0.6).sigm(),
-        T::from(0.7).sigm(),
-        T::from(0.8).sigm(),
-        ]).all(|(&a,b) : (&T, T)| {a-b < 0.001.into()}),
+            .load::<T>(16 * size_of::<F>(), 8 * size_of::<T>())
+            .unwrap()[0]
+            .iter()
+            .zip(vec![
+                T::from(0.1).sigm(),
+                T::from(0.2).sigm(),
+                T::from(0.3).sigm(),
+                T::from(0.4).sigm(),
+                T::from(0.5).sigm(),
+                T::from(0.6).sigm(),
+                T::from(0.7).sigm(),
+                T::from(0.8).sigm(),
+            ])
+            .all(|(&a, b): (&T, T)| { a - b < 0.001.into() }),
         "Wrong result for {}",
         err
     );
@@ -892,17 +910,22 @@ where
         executable
             .cpu_mut()
             .host()
-            .load::<F>(0, 16 * size_of::<F>()).unwrap()[0],
+            .load::<F>(0, 16 * size_of::<F>())
+            .unwrap()[0],
         &buff,
         "Altered first part for {}",
         err
     );
     //Check that later is 0
     assert_eq!(
-        executable.cpu_mut().host().load::<i32>(
-            16 * size_of::<F>() +  8*size_of::<T>(),
-            4 * size_of::<i32>()
-        ).unwrap()[0],
+        executable
+            .cpu_mut()
+            .host()
+            .load::<i32>(
+                16 * size_of::<F>() + 8 * size_of::<T>(),
+                4 * size_of::<i32>()
+            )
+            .unwrap()[0],
         [0, 0, 0, 0],
         "Altered first part for {}",
         err
@@ -917,10 +940,8 @@ fn vsigm_test() {
     vsigm_test_generic::<f64, f64>("vsigm<f64,f64>");
 }
 
-
-
 /// mvmul Test
-fn mvmul_test_generic<F,M, T>(err: &str, relu:i32)
+fn mvmul_test_generic<F, M, T>(err: &str, relu: i32)
 where
     F: From<f32> + std::fmt::Debug + PartialEq<F> + MemoryStorable,
     M: From<f32> + std::fmt::Debug + PartialEq<M> + MemoryStorable,
@@ -948,12 +969,7 @@ where
     crossbar.execute_store(&matrix).unwrap();
     let mut cpu = pimcore::cpu::CPU::new(0, vec![vec![&crossbar]]);
     let (memory, _) = cpu.host().get_memory_crossbar();
-    let vector: [F; _] = [
-        1.0.into(),
-        2.0.into(),
-        3.0.into(),
-        4.0.into(),
-    ];
+    let vector: [F; _] = [1.0.into(), 2.0.into(), 3.0.into(), 4.0.into()];
     memory.execute_store(0, &vector).unwrap();
 
     let mut inst_builder = InstructionsBuilder::new();
@@ -974,7 +990,7 @@ where
         mvmul,
         idata_build
             .set_rdr1(3, 1)
-            .set_mbiw_immrelu_immgroup(8*size_of::<M>() as i32, relu, 0)
+            .set_mbiw_immrelu_immgroup(8 * size_of::<M>() as i32, relu, 0)
             .build(),
     );
     let core_instruction = vec![inst_builder.build().into()];
@@ -982,54 +998,50 @@ where
     executable.execute();
 
     // Check result correct
-   if relu == 0 {
-    assert_eq!(
-        executable
-            .cpu_mut()
-            .host()
-            .load::<T>(4 * size_of::<F>(), 4*size_of::<T>()).unwrap()[0],
-        vec![
-        90.0.into(),
-        (-24.0).into(),
-        110.0.into(),
-        120.0.into(),
-        ],
-        "Wrong result for {}",
-        err
-    );
-   }
-    else {
-    assert_eq!(
-        executable
-            .cpu_mut()
-            .host()
-            .load::<T>(4 * size_of::<F>(), 4*size_of::<T>()).unwrap()[0],
-        vec![
-        90.0.into(),
-        0.0.into(),
-        110.0.into(),
-        120.0.into(),
-        ],
-        "Wrong result for {}",
-        err
-    );
-   }
+    if relu == 0 {
+        assert_eq!(
+            executable
+                .cpu_mut()
+                .host()
+                .load::<T>(4 * size_of::<F>(), 4 * size_of::<T>())
+                .unwrap()[0],
+            vec![90.0.into(), (-24.0).into(), 110.0.into(), 120.0.into(),],
+            "Wrong result for {}",
+            err
+        );
+    } else {
+        assert_eq!(
+            executable
+                .cpu_mut()
+                .host()
+                .load::<T>(4 * size_of::<F>(), 4 * size_of::<T>())
+                .unwrap()[0],
+            vec![90.0.into(), 0.0.into(), 110.0.into(), 120.0.into(),],
+            "Wrong result for {}",
+            err
+        );
+    }
     // Check first part equal
     assert_eq!(
         executable
             .cpu_mut()
             .host()
-            .load::<F>(0, 4 * size_of::<F>()).unwrap()[0],
+            .load::<F>(0, 4 * size_of::<F>())
+            .unwrap()[0],
         &vector,
         "Altered first part for {}",
         err
     );
     //Check that later is 0
     assert_eq!(
-        executable.cpu_mut().host().load::<i32>(
-            4 * size_of::<F>() +  4*size_of::<T>(),
-            4 * size_of::<i32>()
-        ).unwrap()[0],
+        executable
+            .cpu_mut()
+            .host()
+            .load::<i32>(
+                4 * size_of::<F>() + 4 * size_of::<T>(),
+                4 * size_of::<i32>()
+            )
+            .unwrap()[0],
         [0, 0, 0, 0],
         "Altered first part for {}",
         err
@@ -1038,22 +1050,21 @@ where
 
 #[test]
 fn mvmul_test() {
-    mvmul_test_generic::<f32,f32,f32>("mvmul<f32,f32,f32>",0);
-    mvmul_test_generic::<f32,f32,f64>("mvmul<f32,f32,f64>",0);
-    mvmul_test_generic::<f32,f64,f32>("mvmul<f32,f64,f32>",0);
-    mvmul_test_generic::<f32,f64,f64>("mvmul<f32,f64,f64>",0);
-    mvmul_test_generic::<f64,f32,f32>("mvmul<f64,f32,f32>",0);
-    mvmul_test_generic::<f64,f32,f64>("mvmul<f64,f32,f64>",0);
-    mvmul_test_generic::<f64,f64,f32>("mvmul<f64,f64,f32>",0);
-    mvmul_test_generic::<f64,f64,f64>("mvmul<f64,f64,f64>",0);
-
-    mvmul_test_generic::<f32,f32,f32>("mvmul<f32,f32,f32>",1);
-    mvmul_test_generic::<f32,f32,f64>("mvmul<f32,f32,f64>",1);
-    mvmul_test_generic::<f32,f64,f32>("mvmul<f32,f64,f32>",1);
-    mvmul_test_generic::<f32,f64,f64>("mvmul<f32,f64,f64>",1);
-    mvmul_test_generic::<f64,f32,f32>("mvmul<f64,f32,f32>",1);
-    mvmul_test_generic::<f64,f32,f64>("mvmul<f64,f32,f64>",1);
-    mvmul_test_generic::<f64,f64,f32>("mvmul<f64,f64,f32>",1);
-    mvmul_test_generic::<f64,f64,f64>("mvmul<f64,f64,f64>",1);
+    mvmul_test_generic::<f32, f32, f32>("mvmul<f32,f32,f32>", 0);
+    mvmul_test_generic::<f32, f32, f64>("mvmul<f32,f32,f64>", 0);
+    mvmul_test_generic::<f32, f64, f32>("mvmul<f32,f64,f32>", 0);
+    mvmul_test_generic::<f32, f64, f64>("mvmul<f32,f64,f64>", 0);
+    mvmul_test_generic::<f64, f32, f32>("mvmul<f64,f32,f32>", 0);
+    mvmul_test_generic::<f64, f32, f64>("mvmul<f64,f32,f64>", 0);
+    mvmul_test_generic::<f64, f64, f32>("mvmul<f64,f64,f32>", 0);
+    mvmul_test_generic::<f64, f64, f64>("mvmul<f64,f64,f64>", 0);
 
+    mvmul_test_generic::<f32, f32, f32>("mvmul<f32,f32,f32>", 1);
+    mvmul_test_generic::<f32, f32, f64>("mvmul<f32,f32,f64>", 1);
+    mvmul_test_generic::<f32, f64, f32>("mvmul<f32,f64,f32>", 1);
+    mvmul_test_generic::<f32, f64, f64>("mvmul<f32,f64,f64>", 1);
+    mvmul_test_generic::<f64, f32, f32>("mvmul<f64,f32,f32>", 1);
+    mvmul_test_generic::<f64, f32, f64>("mvmul<f64,f32,f64>", 1);
+    mvmul_test_generic::<f64, f64, f32>("mvmul<f64,f64,f32>", 1);
+    mvmul_test_generic::<f64, f64, f64>("mvmul<f64,f64,f64>", 1);
 }

backend-simulators/pim/pim-simulator/tests/sync.rs

@@ -1,7 +1,7 @@
 mod common;
 
 use pimcore::{
-    Executable, CoreInstructionsBuilder,
+    CoreInstructionsBuilder, Executable,
     instruction_set::{InstructionsBuilder, instruction_data::InstructionDataBuilder, isa::*},
 };
 
@@ -158,7 +158,12 @@ fn simple_send_recv_test() {
     let mut inst_builder = InstructionsBuilder::new();
     let mut idata_build = InstructionDataBuilder::new();
     idata_build.set_core_indx(1).fix_core_indx();
-    inst_builder.make_inst(sldi, idata_build.set_rdimm(1, 3*size_of::<f32>() as i32).build());
+    inst_builder.make_inst(
+        sldi,
+        idata_build
+            .set_rdimm(1, 3 * size_of::<f32>() as i32)
+            .build(),
+    );
     inst_builder.make_inst(
         send,
         idata_build
@@ -188,15 +193,11 @@ fn simple_send_recv_test() {
 
     assert_eq!(
         res.unwrap()[0],
-        [
-             4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0
-        ],
+        [4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0],
         "send_recv failed to store"
     );
 }
 
-
-
 // 1 -> 3
 // 2 -> 3
 // 3 <- 2
@@ -210,53 +211,54 @@ fn simple_send_recv_test() {
 fn multiple_send_recv_test() {
     let mut cpu = common::empty_cpu(4);
     let mut core_instruction_builder = CoreInstructionsBuilder::new(4);
-    let buff: [f32; _] = [
-       1.0, 1.0, 1.0, 1.0, 1.0 
-    ];
+    let buff: [f32; _] = [1.0, 1.0, 1.0, 1.0, 1.0];
     cpu.core(1).execute_store(0, &buff).unwrap();
-    let buff: [f32; _] = [
-       2.0, 2.0, 2.0, 2.0, 2.0
-    ];
+    let buff: [f32; _] = [2.0, 2.0, 2.0, 2.0, 2.0];
     cpu.core(2).execute_store(0, &buff).unwrap();
-    let buff: [f32; _] = [
-       3.0, 3.0, 3.0, 3.0, 3.0
-    ];
+    let buff: [f32; _] = [3.0, 3.0, 3.0, 3.0, 3.0];
     cpu.core(3).execute_store(0, &buff).unwrap();
-    let buff: [f32; _] = [
-       4.0, 4.0, 4.0, 4.0, 4.0
-    ];
+    let buff: [f32; _] = [4.0, 4.0, 4.0, 4.0, 4.0];
     cpu.core(4).execute_store(0, &buff).unwrap();
 
     let send_inst = |inst_builder: &mut InstructionsBuilder, from: i32, to: i32| {
-    let mut idata_build = InstructionDataBuilder::new();
-    idata_build.set_core_indx(from).fix_core_indx();
-    inst_builder.make_inst(sldi, idata_build.set_rdimm(1, from*size_of::<f32>() as i32).build());
-    inst_builder.make_inst(
-        send,
-        idata_build
-            .set_r1(1)
-            .set_imm_core(to)
-            .set_imm_len(size_of::<f32>() as i32)
-            .build(),
-    );
+        let mut idata_build = InstructionDataBuilder::new();
+        idata_build.set_core_indx(from).fix_core_indx();
+        inst_builder.make_inst(
+            sldi,
+            idata_build
+                .set_rdimm(1, from * size_of::<f32>() as i32)
+                .build(),
+        );
+        inst_builder.make_inst(
+            send,
+            idata_build
+                .set_r1(1)
+                .set_imm_core(to)
+                .set_imm_len(size_of::<f32>() as i32)
+                .build(),
+        );
     };
 
     let recv_inst = |inst_builder: &mut InstructionsBuilder, to: i32, from: i32| {
-    let mut idata_build = InstructionDataBuilder::new();
-    idata_build.set_core_indx(to).fix_core_indx();
-    inst_builder.make_inst(sldi, idata_build.set_rdimm(1, from*size_of::<f32>() as i32).build());
-    inst_builder.make_inst(
-        recv,
-        idata_build
-            .set_rd(1)
-            .set_imm_core(from)
-            .set_imm_len(size_of::<f32>() as i32)
-            .build(),
-    );
+        let mut idata_build = InstructionDataBuilder::new();
+        idata_build.set_core_indx(to).fix_core_indx();
+        inst_builder.make_inst(
+            sldi,
+            idata_build
+                .set_rdimm(1, from * size_of::<f32>() as i32)
+                .build(),
+        );
+        inst_builder.make_inst(
+            recv,
+            idata_build
+                .set_rd(1)
+                .set_imm_core(from)
+                .set_imm_len(size_of::<f32>() as i32)
+                .build(),
+        );
     };
     let mut inst_builder = InstructionsBuilder::new();
 
-
     // 1 -> 3
     send_inst(&mut inst_builder, 1, 3);
     core_instruction_builder.set_core(1, inst_builder.build());
@@ -289,7 +291,72 @@ fn multiple_send_recv_test() {
 
     assert_eq!(
         res.unwrap()[0],
-        [ 1.0, 2.0, 3.0, 4.0    ],
+        [1.0, 2.0, 3.0, 4.0],
         "send_recv failed to store"
     );
 }
+
+#[test]
+fn sync_wait_tokens_test() {
+    let cpu = common::empty_cpu(2);
+    let mut cores = CoreInstructionsBuilder::new(2);
+    let mut instructions = InstructionsBuilder::new();
+    let mut data = InstructionDataBuilder::new();
+
+    data.set_core_indx(1).fix_core_indx();
+    for _ in 0..2 {
+        instructions.make_inst(
+            sync,
+            data.set_imm_core(2).set_offset_select_value(0, 0).build(),
+        );
+    }
+    cores.set_core(1, instructions.build());
+
+    data.set_core_indx(2).fix_core_indx();
+    for _ in 0..2 {
+        instructions.make_inst(wait, data.set_offset_select_value(0, 1).build());
+    }
+    cores.set_core(2, instructions.build());
+
+    Executable::new(cpu, cores.build()).execute().unwrap();
+}
+
+#[test]
+fn blocked_transfers_do_not_starve_sync_producer() {
+    let cpu = common::empty_cpu(4);
+    let mut cores = CoreInstructionsBuilder::new(4);
+
+    let mut instructions = InstructionsBuilder::new();
+    let mut data = InstructionDataBuilder::new();
+    data.set_core_indx(1).fix_core_indx();
+    instructions.make_inst(sldi, data.set_rdimm(1, 0).build());
+    instructions.make_inst(recv, data.set_rd(1).set_imm_core(2).set_imm_len(1).build());
+    instructions.make_inst(send, data.set_r1(1).set_imm_core(3).set_imm_len(1).build());
+    cores.set_core(1, instructions.build());
+
+    let mut instructions = InstructionsBuilder::new();
+    let mut data = InstructionDataBuilder::new();
+    data.set_core_indx(2).fix_core_indx();
+    instructions.make_inst(sldi, data.set_rdimm(1, 0).build());
+    instructions.make_inst(wait, data.set_offset_select_value(0, 1).build());
+    instructions.make_inst(send, data.set_r1(1).set_imm_core(1).set_imm_len(1).build());
+    cores.set_core(2, instructions.build());
+
+    let mut instructions = InstructionsBuilder::new();
+    let mut data = InstructionDataBuilder::new();
+    data.set_core_indx(3).fix_core_indx();
+    instructions.make_inst(sldi, data.set_rdimm(1, 0).build());
+    instructions.make_inst(recv, data.set_rd(1).set_imm_core(1).set_imm_len(1).build());
+    cores.set_core(3, instructions.build());
+
+    let mut instructions = InstructionsBuilder::new();
+    let mut data = InstructionDataBuilder::new();
+    data.set_core_indx(4).fix_core_indx();
+    instructions.make_inst(
+        sync,
+        data.set_imm_core(2).set_offset_select_value(0, 0).build(),
+    );
+    cores.set_core(4, instructions.build());
+
+    Executable::new(cpu, cores.build()).execute().unwrap();
+}

src/PIM/CMakeLists.txt

@@ -94,7 +94,7 @@ endfunction()
 
 add_subdirectory(Dialect)
 add_subdirectory(Common)
-add_subdirectory(Pass)
+add_subdirectory(Passes)
 add_subdirectory(Compiler)
 add_subdirectory(Conversion)
 
@@ -121,6 +121,7 @@ add_pim_library(OMPIMAccel
   OMPimCommon
   OMPimBufferization
   OMPimHostConstantFolding
+  OMPimInstructionSelection
   OMPimVerification
   MLIRTensorInferTypeOpInterfaceImpl
 )

src/PIM/Common/IR/AddressAnalysis.cpp

@@ -1,10 +1,13 @@
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/Interfaces/DestinationStyleOpInterface.h"
 
+#include "llvm/ADT/SmallPtrSet.h"
+
 #include <limits>
 
 #include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp"
@@ -36,6 +39,10 @@ mlir::Value resolveAlias(mlir::Value value, const StaticValueKnowledge* knowledg
 
 llvm::FailureOr<CompiledIndexExpr> compileIndexValueImpl(mlir::Value value);
 llvm::FailureOr<CompiledAddressExpr> compileContiguousAddressExprImpl(mlir::Value value);
+using AliasResolutionSet = llvm::SmallPtrSet<mlir::Value, 8>;
+mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value,
+                                        const StaticValueKnowledge* knowledge,
+                                        AliasResolutionSet& visited);
 mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnowledge* knowledge);
 
 template <typename... Args>
@@ -45,18 +52,23 @@ CompiledIndexExpr makeCompiledIndexExpr(Args&&... args) {
 
 static mlir::Value resolveForYieldedAliasToInit(mlir::scf::ForOp forOp,
                                                 mlir::Value yieldedValue,
-                                                const StaticValueKnowledge* knowledge) {
-  yieldedValue = resolveLoopCarriedAliasImpl(yieldedValue, knowledge);
+                                                const StaticValueKnowledge* knowledge,
+                                                AliasResolutionSet& visited) {
+  yieldedValue = resolveLoopCarriedAliasImpl(yieldedValue, knowledge, visited);
   if (auto blockArgument = mlir::dyn_cast<mlir::BlockArgument>(yieldedValue)) {
     if (blockArgument.getOwner() == forOp.getBody() && blockArgument.getArgNumber() > 0
         && static_cast<unsigned>(blockArgument.getArgNumber() - 1) < forOp.getInitArgs().size())
-      return resolveLoopCarriedAliasImpl(forOp.getInitArgs()[blockArgument.getArgNumber() - 1], knowledge);
+      return resolveLoopCarriedAliasImpl(forOp.getInitArgs()[blockArgument.getArgNumber() - 1], knowledge, visited);
   }
   return yieldedValue;
 }
 
-mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnowledge* knowledge) {
+mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value,
+                                        const StaticValueKnowledge* knowledge,
+                                        AliasResolutionSet& visited) {
   value = resolveAlias(value, knowledge);
+  if (!value || !visited.insert(value).second)
+    return value;
 
   if (auto blockArgument = mlir::dyn_cast<mlir::BlockArgument>(value)) {
     auto forOp = mlir::dyn_cast_or_null<mlir::scf::ForOp>(blockArgument.getOwner()->getParentOp());
@@ -64,9 +76,12 @@ mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnow
       const unsigned iterArgIndex = blockArgument.getArgNumber() - 1;
       auto yieldOp = mlir::dyn_cast<mlir::scf::YieldOp>(forOp.getBody()->getTerminator());
       if (iterArgIndex < forOp.getInitArgs().size() && yieldOp
-          && iterArgIndex < yieldOp.getNumOperands()
-          && resolveAlias(yieldOp.getOperand(iterArgIndex), knowledge) == blockArgument)
-        return resolveLoopCarriedAliasImpl(forOp.getInitArgs()[iterArgIndex], knowledge);
+          && iterArgIndex < yieldOp.getNumOperands()) {
+        mlir::Value yieldedValue = resolveAlias(yieldOp.getOperand(iterArgIndex), knowledge);
+        if (yieldedValue == blockArgument
+            || (yieldedValue && resolveLoopCarriedAliasImpl(yieldedValue, knowledge, visited) == blockArgument))
+          return resolveLoopCarriedAliasImpl(forOp.getInitArgs()[iterArgIndex], knowledge, visited);
+      }
     }
     return value;
   }
@@ -75,10 +90,15 @@ mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnow
   if (!definingOp)
     return value;
 
+  if (auto toBufferOp = mlir::dyn_cast<mlir::bufferization::ToBufferOp>(definingOp))
+    return resolveLoopCarriedAliasImpl(toBufferOp.getTensor(), knowledge, visited);
+  if (auto toTensorOp = mlir::dyn_cast<mlir::bufferization::ToTensorOp>(definingOp))
+    return resolveLoopCarriedAliasImpl(toTensorOp.getBuffer(), knowledge, visited);
+
   if (auto dpsDefiningOp = mlir::dyn_cast<mlir::DestinationStyleOpInterface>(definingOp)) {
     if (auto result = mlir::dyn_cast<mlir::OpResult>(value))
       if (mlir::OpOperand* tiedOperand = dpsDefiningOp.getTiedOpOperand(result))
-        return resolveLoopCarriedAliasImpl(tiedOperand->get(), knowledge);
+        return resolveLoopCarriedAliasImpl(tiedOperand->get(), knowledge, visited);
   }
 
   if (auto forOp = mlir::dyn_cast<mlir::scf::ForOp>(definingOp)) {
@@ -86,20 +106,26 @@ mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnow
     if (result) {
       auto yieldOp = mlir::dyn_cast<mlir::scf::YieldOp>(forOp.getBody()->getTerminator());
       if (yieldOp && result.getResultNumber() < yieldOp.getNumOperands())
-        return resolveForYieldedAliasToInit(forOp, yieldOp.getOperand(result.getResultNumber()), knowledge);
+        return resolveForYieldedAliasToInit(
+          forOp, yieldOp.getOperand(result.getResultNumber()), knowledge, visited);
     }
   }
 
   if (auto castOp = mlir::dyn_cast<mlir::memref::CastOp>(definingOp))
-    return resolveLoopCarriedAliasImpl(castOp.getSource(), knowledge);
+    return resolveLoopCarriedAliasImpl(castOp.getSource(), knowledge, visited);
   if (auto collapseOp = mlir::dyn_cast<mlir::memref::CollapseShapeOp>(definingOp))
-    return resolveLoopCarriedAliasImpl(collapseOp.getSrc(), knowledge);
+    return resolveLoopCarriedAliasImpl(collapseOp.getSrc(), knowledge, visited);
   if (auto expandOp = mlir::dyn_cast<mlir::memref::ExpandShapeOp>(definingOp))
-    return resolveLoopCarriedAliasImpl(expandOp.getSrc(), knowledge);
+    return resolveLoopCarriedAliasImpl(expandOp.getSrc(), knowledge, visited);
 
   return value;
 }
 
+mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnowledge* knowledge) {
+  AliasResolutionSet visited;
+  return resolveLoopCarriedAliasImpl(value, knowledge, visited);
+}
+
 llvm::FailureOr<int64_t> resolveOpFoldResult(mlir::OpFoldResult ofr, const StaticValueKnowledge* knowledge);
 llvm::FailureOr<int64_t> resolveIndexValueImpl(mlir::Value value, const StaticValueKnowledge* knowledge);
 
@@ -524,6 +550,15 @@ llvm::FailureOr<ResolvedContiguousAddress> resolveContiguousAddressImpl(mlir::Va
     if (!definingOp)
       return mlir::failure();
 
+    if (auto toBufferOp = mlir::dyn_cast<mlir::bufferization::ToBufferOp>(definingOp)) {
+      value = resolveAlias(toBufferOp.getTensor(), knowledge);
+      continue;
+    }
+    if (auto toTensorOp = mlir::dyn_cast<mlir::bufferization::ToTensorOp>(definingOp)) {
+      value = resolveAlias(toTensorOp.getBuffer(), knowledge);
+      continue;
+    }
+
     if (auto dpsDefiningOp = mlir::dyn_cast<mlir::DestinationStyleOpInterface>(definingOp)) {
       mlir::OpOperand* tiedOperand = dpsDefiningOp.getTiedOpOperand(mlir::dyn_cast<mlir::OpResult>(value));
       if (!tiedOperand)
@@ -538,7 +573,9 @@ llvm::FailureOr<ResolvedContiguousAddress> resolveContiguousAddressImpl(mlir::Va
         return mlir::failure();
 
       auto yieldOp = mlir::cast<mlir::scf::YieldOp>(forOp.getBody()->getTerminator());
-      value = resolveForYieldedAliasToInit(forOp, yieldOp.getOperand(result.getResultNumber()), knowledge);
+      AliasResolutionSet visited;
+      value = resolveForYieldedAliasToInit(
+        forOp, yieldOp.getOperand(result.getResultNumber()), knowledge, visited);
       continue;
     }
 
@@ -643,6 +680,15 @@ llvm::FailureOr<CompiledAddressExpr> compileContiguousAddressExprImpl(mlir::Valu
     if (!definingOp)
       return mlir::failure();
 
+    if (auto toBufferOp = mlir::dyn_cast<mlir::bufferization::ToBufferOp>(definingOp)) {
+      value = toBufferOp.getTensor();
+      continue;
+    }
+    if (auto toTensorOp = mlir::dyn_cast<mlir::bufferization::ToTensorOp>(definingOp)) {
+      value = toTensorOp.getBuffer();
+      continue;
+    }
+
     if (auto dpsDefiningOp = mlir::dyn_cast<mlir::DestinationStyleOpInterface>(definingOp)) {
       mlir::OpOperand* tiedOperand = dpsDefiningOp.getTiedOpOperand(mlir::dyn_cast<mlir::OpResult>(value));
       if (!tiedOperand)
@@ -657,7 +703,9 @@ llvm::FailureOr<CompiledAddressExpr> compileContiguousAddressExprImpl(mlir::Valu
         return mlir::failure();
 
       auto yieldOp = mlir::cast<mlir::scf::YieldOp>(forOp.getBody()->getTerminator());
-      value = resolveForYieldedAliasToInit(forOp, yieldOp.getOperand(result.getResultNumber()), nullptr);
+      AliasResolutionSet visited;
+      value = resolveForYieldedAliasToInit(
+        forOp, yieldOp.getOperand(result.getResultNumber()), nullptr, visited);
       continue;
     }
 

src/PIM/Common/IR/AddressAnalysis.hpp

@@ -19,7 +19,7 @@ struct ResolvedContiguousAddress {
 };
 
 /// Records compile-time facts used when interpreting address arithmetic and
-/// loop-carried aliases inside PIM regions.
+/// loop-carried aliases inside Pim regions.
 struct StaticValueKnowledge {
   llvm::DenseMap<mlir::Value, int64_t> indexValues;
   llvm::DenseMap<mlir::Value, mlir::Value> aliases;

src/PIM/Common/IR/CoreBlockUtils.cpp

@@ -85,12 +85,12 @@ static mlir::LogicalResult walkPimCoreBlockImpl(mlir::Block& block,
       auto step = resolveIndexValue(forOp.getStep(), knowledge);
       if (failed(lower) || failed(upper) || failed(step)
           || (mode == CoreWalkMode::ExecuteCommunication && *step <= 0)) {
-        forOp.emitOpError() << "requires statically evaluable scf.for bounds for PIM " << purpose;
+        forOp.emitOpError() << "requires statically evaluable scf.for bounds for Pim " << purpose;
         hasFailure = true;
         continue;
       }
       if (*step <= 0) {
-        forOp.emitOpError("requires positive scf.for step for PIM verification");
+        forOp.emitOpError("requires positive scf.for step for Pim verification");
         hasFailure = true;
         continue;
       }
@@ -126,7 +126,7 @@ static mlir::LogicalResult walkPimCoreBlockImpl(mlir::Block& block,
     if (auto ifOp = mlir::dyn_cast<mlir::scf::IfOp>(op)) {
       auto condition = resolveIndexValue(ifOp.getCondition(), knowledge);
       if (failed(condition)) {
-        ifOp.emitOpError() << "requires statically evaluable scf.if condition for PIM " << purpose;
+        ifOp.emitOpError() << "requires statically evaluable scf.if condition for Pim " << purpose;
         hasFailure = true;
         continue;
       }
@@ -147,7 +147,7 @@ static mlir::LogicalResult walkPimCoreBlockImpl(mlir::Block& block,
     if (auto switchOp = mlir::dyn_cast<mlir::scf::IndexSwitchOp>(op)) {
       auto selector = resolveIndexValue(switchOp.getArg(), knowledge);
       if (failed(selector)) {
-        switchOp.emitOpError() << "requires a statically evaluable scf.index_switch selector for PIM " << purpose;
+        switchOp.emitOpError() << "requires a statically evaluable scf.index_switch selector for Pim " << purpose;
         hasFailure = true;
         continue;
       }

src/PIM/Common/IR/CoreBlockUtils.hpp

@@ -14,7 +14,7 @@ namespace onnx_mlir {
 using PimCoreCommunicationPlan = llvm::DenseMap<mlir::Block*, llvm::SmallVector<mlir::Operation*, 8>>;
 
 /// Returns true for ops in a `pim.core` body that only participate in static
-/// address or index computation and therefore do not emit PIM instructions.
+/// address or index computation and therefore do not emit Pim instructions.
 bool isCoreStaticAddressOp(mlir::Operation* op);
 
 /// Walks a `pim.core` body's communication stream, statically unrolling

src/PIM/Common/IR/EntryPointUtils.cpp

@@ -9,7 +9,7 @@ llvm::FailureOr<mlir::func::FuncOp> getPimEntryFunc(mlir::ModuleOp moduleOp) {
 
   llvm::SmallVector<mlir::ONNXEntryPointOp> entryPoints(moduleOp.getOps<mlir::ONNXEntryPointOp>());
   if (entryPoints.size() > 1) {
-    moduleOp.emitError("PIM pipeline requires a single ONNX entry point, but found ") << entryPoints.size();
+    moduleOp.emitError("Pim pipeline requires a single ONNX entry point, but found ") << entryPoints.size();
     return mlir::failure();
   }
   if (!entryPoints.empty()) {
@@ -38,7 +38,7 @@ llvm::FailureOr<mlir::func::FuncOp> getPimEntryFunc(mlir::ModuleOp moduleOp) {
   if (nonExternalFuncs.size() == 1)
     return nonExternalFuncs.front();
 
-  moduleOp.emitError("could not resolve a unique PIM entry function");
+  moduleOp.emitError("could not resolve a unique Pim entry function");
   return mlir::failure();
 }
 

src/PIM/Common/IR/EntryPointUtils.hpp

@@ -5,7 +5,7 @@
 
 namespace onnx_mlir {
 
-/// Resolves the function the PIM pipeline should treat as its entry point.
+/// Resolves the function the Pim pipeline should treat as its entry point.
 /// Prefers ONNX entry-point metadata, then `main_graph`, then the only
 /// non-external function if the module is otherwise unambiguous.
 llvm::FailureOr<mlir::func::FuncOp> getPimEntryFunc(mlir::ModuleOp moduleOp);

src/PIM/Common/IR/WeightUtils.hpp

@@ -32,7 +32,7 @@ struct ResolvedWeightView {
 bool hasWeightAlways(mlir::Operation* op);
 
 /// Tags an op as producing a value that should stay materialized as a reusable
-/// weight across later PIM lowering/codegen stages.
+/// weight across later Pim lowering/codegen stages.
 void markWeightAlways(mlir::Operation* op);
 
 bool isSpatialMvmVmmWeightUse(mlir::OpOperand& use);

src/PIM/Common/PimCommon.hpp

@@ -32,6 +32,9 @@ inline constexpr llvm::StringLiteral kCoreIdAttrName = "coreId";
 inline constexpr llvm::StringLiteral kCoreIdsAttrName = "coreIds";
 inline constexpr llvm::StringLiteral kLocalMemoryAddressAttrName = "pim.local_memory_address";
 inline constexpr llvm::StringLiteral kLocalMemorySizeAttrName = "pim.local_memory_size";
+inline constexpr llvm::StringLiteral kPipelineHostBufferBytesAttrName = "pim.pipeline_host_buffer_bytes";
+inline constexpr llvm::StringLiteral kPipelineHostBufferName = "pim_pipeline_channels";
+inline constexpr size_t kPimEventRegisterCount = 32;
 inline constexpr std::array<llvm::StringLiteral, 4> kRemovedLocalMemoryPlanAttrNames = {
   "pim.local_memory_slot",
   "pim.local_memory_slot_size",

src/PIM/Common/Support/CheckedArithmetic.cpp

@@ -11,7 +11,7 @@ namespace onnx_mlir::pim {
 namespace {
 
 static void emitCrashMessage(llvm::StringRef fieldName, llvm::StringRef message) {
-  llvm::errs() << "PIM " << fieldName << " " << message << "\n";
+  llvm::errs() << "Pim " << fieldName << " " << message << "\n";
 }
 
 template <typename To, typename From>
@@ -65,7 +65,7 @@ InFlightDiagnostic emitCheckedArithmeticError(Operation* anchor, llvm::StringRef
 }
 
 InFlightDiagnostic emitCheckedArithmeticError(Location loc, llvm::StringRef fieldName, llvm::StringRef message) {
-  return emitError(loc) << "PIM " << fieldName << " " << message;
+  return emitError(loc) << "Pim " << fieldName << " " << message;
 }
 
 FailureOr<int32_t> checkedI32(int64_t value, Operation* anchor, llvm::StringRef fieldName) {
@@ -174,7 +174,7 @@ FailureOr<uint64_t> getCheckedShapedTypeSizeInBytes(ShapedType type, Location lo
 int32_t checkedI32OrCrash(int64_t value, llvm::StringRef fieldName) {
   if (value < std::numeric_limits<int32_t>::min() || value > std::numeric_limits<int32_t>::max()) {
     emitCrashMessage(fieldName, "is outside representable range");
-    llvm_unreachable("PIM checked arithmetic failure");
+    llvm_unreachable("Pim checked arithmetic failure");
   }
   return static_cast<int32_t>(value);
 }
@@ -182,7 +182,7 @@ int32_t checkedI32OrCrash(int64_t value, llvm::StringRef fieldName) {
 int32_t checkedI32OrCrash(uint64_t value, llvm::StringRef fieldName) {
   if (value > static_cast<uint64_t>(std::numeric_limits<int32_t>::max())) {
     emitCrashMessage(fieldName, "is outside representable range");
-    llvm_unreachable("PIM checked arithmetic failure");
+    llvm_unreachable("Pim checked arithmetic failure");
   }
   return static_cast<int32_t>(value);
 }
@@ -190,7 +190,7 @@ int32_t checkedI32OrCrash(uint64_t value, llvm::StringRef fieldName) {
 uint8_t checkedU8OrCrash(uint64_t value, llvm::StringRef fieldName) {
   if (value > static_cast<uint64_t>(std::numeric_limits<uint8_t>::max())) {
     emitCrashMessage(fieldName, "is outside representable range");
-    llvm_unreachable("PIM checked arithmetic failure");
+    llvm_unreachable("Pim checked arithmetic failure");
   }
   return static_cast<uint8_t>(value);
 }
@@ -198,7 +198,7 @@ uint8_t checkedU8OrCrash(uint64_t value, llvm::StringRef fieldName) {
 size_t checkedSizeOrCrash(int64_t value, llvm::StringRef fieldName) {
   if (value < 0) {
     emitCrashMessage(fieldName, "is outside representable range");
-    llvm_unreachable("PIM checked arithmetic failure");
+    llvm_unreachable("Pim checked arithmetic failure");
   }
   return static_cast<size_t>(value);
 }
@@ -206,7 +206,7 @@ size_t checkedSizeOrCrash(int64_t value, llvm::StringRef fieldName) {
 size_t checkedAddOrCrash(size_t lhs, size_t rhs, llvm::StringRef fieldName) {
   if (rhs > std::numeric_limits<size_t>::max() - lhs) {
     emitCrashMessage(fieldName, "addition overflow");
-    llvm_unreachable("PIM checked arithmetic failure");
+    llvm_unreachable("Pim checked arithmetic failure");
   }
   return lhs + rhs;
 }
@@ -214,7 +214,7 @@ size_t checkedAddOrCrash(size_t lhs, size_t rhs, llvm::StringRef fieldName) {
 size_t checkedMulOrCrash(size_t lhs, size_t rhs, llvm::StringRef fieldName) {
   if (lhs != 0 && rhs > std::numeric_limits<size_t>::max() / lhs) {
     emitCrashMessage(fieldName, "multiplication overflow");
-    llvm_unreachable("PIM checked arithmetic failure");
+    llvm_unreachable("Pim checked arithmetic failure");
   }
   return lhs * rhs;
 }

src/PIM/Common/Support/FileSystemUtils.hpp

@@ -4,7 +4,7 @@
 
 namespace onnx_mlir {
 
-/// Returns the directory that should hold PIM artifacts/debug dumps for the
+/// Returns the directory that should hold Pim artifacts/debug dumps for the
 /// current compiler invocation.
 std::string getOutputDir();
 

src/PIM/Compiler/CMakeLists.txt

@@ -31,6 +31,7 @@ add_pim_library(OMPimCompilerUtils
   OMPimCommon
   OMPimBufferization
   OMPimHostConstantFolding
+  OMPimInstructionSelection
   OMPimLocalMemoryPlanning
   OMPimVerification
   OMPimPasses

src/PIM/Compiler/PimArtifactWriter.cpp

@@ -140,9 +140,13 @@ OnnxMlirCompilerErrorCodes writeConfigJson(func::FuncOp funcOp,
   configJson["array_group_map"] = std::move(xbarsPerArrayGroup);
 
   json::Array inputsAddresses;
-  for (BlockArgument input : funcOp.getArguments())
+  json::Array inputsSizes;
+  for (BlockArgument input : funcOp.getArguments()) {
     inputsAddresses.push_back(memory.getValueAddress(input));
+    inputsSizes.push_back(memory.hostMem.getMemEntry({input, std::nullopt}).size);
+  }
   configJson["inputs_addresses"] = std::move(inputsAddresses);
+  configJson["inputs_sizes"] = std::move(inputsSizes);
 
   json::Array outputsAddresses;
   for (func::ReturnOp returnOp : funcOp.getOps<func::ReturnOp>())

src/PIM/Compiler/PimBinaryFormat.hpp

@@ -162,8 +162,8 @@ inline constexpr std::array<InstructionJsonFormat, kOpcodeCount> kInstructionJso
   {true,  true,  true,  "",     "",     "",      "len" }, // lmv
   {true,  false, true,  "core", "",     "",      "size"}, // send
   {true,  false, true,  "core", "",     "",      "size"}, // recv
-  {false, false, false, "",     "",     "",      ""    }, // wait
-  {false, false, false, "",     "",     "",      ""    }, // sync
+  {false, false, false, "",     "event_register", "wait_value", ""}, // wait
+  {false, false, false, "core", "event_register", "", ""}, // sync
 }};
 static_assert(kInstructionJsonFormats.size() == kOpcodeCount);
 
@@ -171,19 +171,19 @@ inline Opcode opcodeFromString(llvm::StringRef opName) {
   for (auto [index, name] : llvm::enumerate(kOpcodeNames))
     if (opName == name)
       return static_cast<Opcode>(index);
-  llvm_unreachable("Unsupported PIM binary opcode");
+  llvm_unreachable("Unsupported Pim binary opcode");
 }
 
 inline llvm::StringRef opcodeToString(Opcode opcode) {
   size_t index = static_cast<size_t>(opcode);
-  assert(index < kOpcodeNames.size() && "Unsupported PIM binary opcode");
+  assert(index < kOpcodeNames.size() && "Unsupported Pim binary opcode");
   return kOpcodeNames[index];
 }
 
 inline InstructionRecord makeInstructionRecord(const llvm::json::Object& instruction) {
   InstructionRecord record;
   std::optional<llvm::StringRef> opName = instruction.getString("op");
-  assert(opName && "Missing op field in PIM instruction");
+  assert(opName && "Missing op field in Pim instruction");
   record.opcode = opcodeFromString(*opName);
   const auto& format = kInstructionJsonFormats[static_cast<size_t>(record.opcode)];
   if (format.rd)

src/PIM/Compiler/PimCodeGen.cpp

@@ -108,7 +108,24 @@ static Operation* getDiagnosticAnchor(mlir::Value value) {
   return nullptr;
 }
 
-// PIM instruction immediates are serialized as signed int32_t fields today
+static bool isZeroSplatGlobal(mlir::Value value) {
+  auto getGlobalOp = value.getDefiningOp<memref::GetGlobalOp>();
+  auto moduleOp = getGlobalOp ? getGlobalOp->getParentOfType<ModuleOp>() : ModuleOp();
+  auto globalOp = moduleOp ? lookupGlobalForGetGlobal(moduleOp, getGlobalOp) : memref::GlobalOp();
+  if (!globalOp || !globalOp.getConstant() || !globalOp.getInitialValue())
+    return false;
+  auto denseAttr = dyn_cast<DenseElementsAttr>(*globalOp.getInitialValue());
+  if (!denseAttr || !denseAttr.isSplat())
+    return false;
+  Attribute valueAttr = denseAttr.getSplatValue<Attribute>();
+  if (auto floatAttr = dyn_cast<FloatAttr>(valueAttr))
+    return floatAttr.getValue().isZero() && !floatAttr.getValue().isNegative();
+  if (auto integerAttr = dyn_cast<IntegerAttr>(valueAttr))
+    return integerAttr.getValue().isZero();
+  return false;
+}
+
+// Pim instruction immediates are serialized as signed int32_t fields today
 // (`sldi` goes through checkedI32OrCrash), so local addresses must stay within
 // the non-negative int32_t range.
 static FailureOr<size_t> checkedAlignTo(size_t value, size_t alignment, Operation* anchor, StringRef fieldName) {
@@ -124,7 +141,7 @@ static void printMemoryOverflowDiagnostic(const MemoryValueKey& key,
                                           size_t requestedSize,
                                           size_t currentFirstAvailableAddress,
                                           size_t alignedEndAddress) {
-  llvm::errs() << "PIM local memory allocation overflow\n";
+  llvm::errs() << "Pim local memory allocation overflow\n";
   llvm::errs() << "Requested allocation size: " << requestedSize << " bytes\n";
   llvm::errs() << "Current firstAvailableAddress: " << currentFirstAvailableAddress << "\n";
   llvm::errs() << "Aligned end address: " << alignedEndAddress << "\n";
@@ -170,7 +187,7 @@ size_t PimMemory::allocateAddress(size_t size, const MemoryValueKey& key) {
       size,
       firstAvailableAddress,
       succeeded(checkedAlignedEnd) ? *checkedAlignedEnd : kPimLocalMemoryAddressLimit);
-    llvm_unreachable("PIM local memory allocation overflow");
+    llvm_unreachable("Pim local memory allocation overflow");
   }
   firstAvailableAddress = *checkedAlignedEnd;
   return address;
@@ -259,7 +276,7 @@ void PimMemory::allocateCore(const CompiledCoreMemoryPlan& plan, std::optional<u
   }
   else if (*localArenaSize != plan.arenaSize || reportRow.logicalLocalAllocationCount != plan.logicalAllocationCount
            || reportRow.logicalLocalBytes != plan.logicalBytes)
-    llvm_unreachable("inconsistent PIM local-memory plan across core-batch lanes");
+    llvm_unreachable("inconsistent Pim local-memory plan across core-batch lanes");
   for (const CompiledLocalMemoryEntry& entry : plan.entries) {
     MemoryValueKey key = getMemoryValueKey(entry.value, lane);
     ownedMemEntriesMap[key] = entry.memory;
@@ -335,8 +352,8 @@ size_t PimAcceleratorMemory::getValueAddress(mlir::Value value,
     llvm_unreachable("Missing mem entry");
   }
 
-  size_t byteOffset = pim::checkedSizeOrCrash(resolvedAddress->byteOffset, "resolved PIM byte offset");
-  return pim::checkedAddOrCrash(iter->second.address, byteOffset, "resolved PIM address");
+  size_t byteOffset = pim::checkedSizeOrCrash(resolvedAddress->byteOffset, "resolved Pim byte offset");
+  return pim::checkedAddOrCrash(iter->second.address, byteOffset, "resolved Pim address");
 }
 
 llvm::FailureOr<int64_t> PimAcceleratorMemory::getIndexValue(mlir::Value value,
@@ -527,11 +544,23 @@ void PimCodeGen::setupRdRs1(size_t rdAddress, size_t rdOffset, size_t rs1Address
   genSetRegisterImmediateUnsigned(1, pim::checkedAddOrCrash(rs1Address, rs1Offset, "rs1 address"));
 }
 
-void PimCodeGen::setupRdRs1Rs2(
+std::array<uint8_t, 3> PimCodeGen::setupRdRs1Rs2(
   size_t rdAddress, size_t rdOffset, size_t rs1Address, size_t rs1Offset, size_t rs2Address, size_t rs2Offset) const {
-  genSetRegisterImmediateUnsigned(0, pim::checkedAddOrCrash(rdAddress, rdOffset, "rd address"));
-  genSetRegisterImmediateUnsigned(1, pim::checkedAddOrCrash(rs1Address, rs1Offset, "rs1 address"));
-  genSetRegisterImmediateUnsigned(2, pim::checkedAddOrCrash(rs2Address, rs2Offset, "rs2 address"));
+  size_t rd = pim::checkedAddOrCrash(rdAddress, rdOffset, "rd address");
+  size_t rs1 = pim::checkedAddOrCrash(rs1Address, rs1Offset, "rs1 address");
+  size_t rs2 = pim::checkedAddOrCrash(rs2Address, rs2Offset, "rs2 address");
+  genSetRegisterImmediateUnsigned(0, rd);
+  uint8_t rs1Register = 0;
+  if (rd != rs1) {
+    genSetRegisterImmediateUnsigned(1, rs1);
+    rs1Register = 1;
+  }
+  if (rd == rs2)
+    return {0, rs1Register, 0};
+  if (rs1 == rs2)
+    return {0, rs1Register, rs1Register};
+  genSetRegisterImmediateUnsigned(2, rs2);
+  return {0, rs1Register, 2};
 }
 
 void PimCodeGen::emitMemCopyOp(pim_binary::Opcode opcode,
@@ -541,14 +570,23 @@ void PimCodeGen::emitMemCopyOp(pim_binary::Opcode opcode,
                                size_t rs1Offset,
                                size_t size,
                                StringRef sizeFieldName) const {
-  setupRdRs1(rdAddr, rdOffset, rs1Addr, rs1Offset);
-
   pim_binary::InstructionRecord instruction;
   instruction.opcode = opcode;
   instruction.rd = 0;
   instruction.r1 = 1;
-  instruction.generic1 = 0;
-  instruction.generic2 = 0;
+  if (rdOffset == rs1Offset) {
+    setupRdRs1(rdAddr, 0, rs1Addr, 0);
+    instruction.generic1 = rdOffset == 0 ? 0 : 3;
+    instruction.generic2 = pim::checkedI32OrCrash(rdOffset, "shared address offset");
+  } else if (rdOffset != 0) {
+    setupRdRs1(rdAddr, 0, rs1Addr, rs1Offset);
+    instruction.generic1 = 1;
+    instruction.generic2 = pim::checkedI32OrCrash(rdOffset, "rd address offset");
+  } else {
+    setupRdRs1(rdAddr, 0, rs1Addr, 0);
+    instruction.generic1 = rs1Offset == 0 ? 0 : 2;
+    instruction.generic2 = pim::checkedI32OrCrash(rs1Offset, "rs1 address offset");
+  }
   instruction.generic3 = pim::checkedI32OrCrash(size, sizeFieldName);
   emitInstruction(instruction);
 }
@@ -584,6 +622,16 @@ void PimCodeGen::codeGenLoadOp(pim::PimMemCopyHostToDevOp loadOp, const StaticVa
   auto hostSourceOffset = indexOf(loadOp.getHostSourceOffset(), knowledge);
   assert(succeeded(deviceTargetOffset) && succeeded(hostSourceOffset)
          && "pim.memcp_hd offsets must be statically resolvable during codegen");
+  if (isZeroSplatGlobal(loadOp.getHostSource())) {
+    setupRd(addressOf(loadOp.getDeviceTarget(), knowledge), *deviceTargetOffset);
+    pim_binary::InstructionRecord instruction;
+    instruction.opcode = pim_binary::Opcode::lldi;
+    instruction.rd = 0;
+    instruction.r2OrImm = 0;
+    instruction.generic3 = loadOp.getSize();
+    emitInstruction(instruction);
+    return;
+  }
   emitMemCopyOp(pim_binary::Opcode::ld,
                 addressOf(loadOp.getDeviceTarget(), knowledge),
                 *deviceTargetOffset,
@@ -619,11 +667,34 @@ void PimCodeGen::codeGenLmvOp(pim::PimMemCopyOp lmvOp, const StaticValueKnowledg
                 "len");
 }
 
+void PimCodeGen::codeGenVMVOp(pim::PimVMVOp vmvOp, const StaticValueKnowledge& knowledge) const {
+  auto targetOffset = indexOf(vmvOp.getTargetOffset(), knowledge);
+  auto sourceOffset = indexOf(vmvOp.getSourceOffset(), knowledge);
+  auto sourceStride = indexOf(vmvOp.getSourceStride(), knowledge);
+  assert(succeeded(targetOffset) && succeeded(sourceOffset) && succeeded(sourceStride)
+         && "pim.vmv operands must be statically resolvable during codegen");
+  auto sourceType = cast<ShapedType>(vmvOp.getSource().getType());
+  int32_t bitwidth = getVectorElementBitwidthOrCrash(sourceType);
+  ensureVectorBitwidth(bitwidth, bitwidth);
+  auto registers = setupRdRs1Rs2(addressOf(vmvOp.getTarget(), knowledge), *targetOffset,
+                                 addressOf(vmvOp.getSource(), knowledge), *sourceOffset, 0, *sourceStride);
+  pim_binary::InstructionRecord instruction;
+  instruction.opcode = pim_binary::Opcode::vmv;
+  instruction.rd = registers[0];
+  instruction.r1 = registers[1];
+  instruction.r2OrImm = registers[2];
+  instruction.generic3 = vmvOp.getLength();
+  emitInstruction(instruction);
+}
+
 void PimCodeGen::codeGenReceiveOp(pim::PimReceiveOp receiveOp, const StaticValueKnowledge& knowledge) const {
+  auto outputOffset = indexOf(receiveOp.getOutputOffset(), knowledge);
   auto sourceCoreId = indexOf(receiveOp.getSourceCoreId(), knowledge);
-  assert(succeeded(sourceCoreId) && "pim.receive source core id must be statically resolvable during codegen");
+  assert(succeeded(outputOffset) && succeeded(sourceCoreId)
+         && "pim.receive offset and source core id must be statically resolvable during codegen");
   emitCommunicationOp(
-    pim_binary::Opcode::recv, addressOf(receiveOp.getOutputBuffer(), knowledge), *sourceCoreId, receiveOp.getSize());
+    pim_binary::Opcode::recv, addressOf(receiveOp.getOutputBuffer(), knowledge) + *outputOffset,
+    *sourceCoreId, receiveOp.getSize());
 }
 
 void PimCodeGen::codeGenSendOp(pim::PimSendOp sendOp, const StaticValueKnowledge& knowledge) const {
@@ -633,6 +704,41 @@ void PimCodeGen::codeGenSendOp(pim::PimSendOp sendOp, const StaticValueKnowledge
     pim_binary::Opcode::send, addressOf(sendOp.getInput(), knowledge), *targetCoreId, sendOp.getSize());
 }
 
+void PimCodeGen::codeGenWaitOp(
+    pim::PimWaitOp waitOp, const StaticValueKnowledge& knowledge) const {
+  if (pimDisableSynchronization)
+    return;
+  auto eventRegister = indexOf(waitOp.getEventRegister(), knowledge);
+  auto waitValue = indexOf(waitOp.getWaitValue(), knowledge);
+  assert(succeeded(eventRegister) && succeeded(waitValue)
+         && "pim.wait operands must be statically resolvable during codegen");
+  if (*waitValue == 0)
+    return;
+  pim_binary::InstructionRecord instruction;
+  instruction.opcode = pim_binary::Opcode::wait;
+  instruction.generic1 = pim::checkedI32OrCrash(
+    *eventRegister, "wait event register");
+  instruction.generic2 = pim::checkedI32OrCrash(*waitValue, "wait value");
+  emitInstruction(instruction);
+}
+
+void PimCodeGen::codeGenSyncOp(
+    pim::PimSyncOp syncOp, const StaticValueKnowledge& knowledge) const {
+  if (pimDisableSynchronization)
+    return;
+  auto targetCoreId = indexOf(syncOp.getTargetCoreId(), knowledge);
+  auto eventRegister = indexOf(syncOp.getEventRegister(), knowledge);
+  assert(succeeded(targetCoreId) && succeeded(eventRegister)
+         && "pim.sync operands must be statically resolvable during codegen");
+  pim_binary::InstructionRecord instruction;
+  instruction.opcode = pim_binary::Opcode::sync;
+  instruction.r2OrImm = pim::checkedI32OrCrash(
+    *targetCoreId, "sync target core id");
+  instruction.generic1 = pim::checkedI32OrCrash(
+    *eventRegister, "sync event register");
+  emitInstruction(instruction);
+}
+
 void PimCodeGen::codeGenConcatOp(pim::PimConcatOp concatOp, const StaticValueKnowledge& knowledge) const {
   auto outputType = cast<ShapedType>(concatOp.getOutputBuffer().getType());
   assert(outputType.hasStaticShape() && "concat codegen requires static output shape");
@@ -690,12 +796,13 @@ void PimCodeGen::emitBinaryVectorOp(pim_binary::Opcode opcode,
   auto inputType = cast<ShapedType>(lhs.getType());
   ensureVectorBitwidth(getVectorElementBitwidthOrCrash(inputType),
                        getVectorElementBitwidthOrCrash(cast<ShapedType>(output.getType())));
-  setupRdRs1Rs2(addressOf(output, knowledge), 0, addressOf(lhs, knowledge), 0, addressOf(rhs, knowledge), 0);
+  auto registers = setupRdRs1Rs2(
+    addressOf(output, knowledge), 0, addressOf(lhs, knowledge), 0, addressOf(rhs, knowledge), 0);
   pim_binary::InstructionRecord instruction;
   instruction.opcode = opcode;
-  instruction.rd = 0;
-  instruction.r1 = 1;
-  instruction.r2OrImm = 2;
+  instruction.rd = registers[0];
+  instruction.r1 = registers[1];
+  instruction.r2OrImm = registers[2];
   instruction.generic3 = getVectorElementCountOrCrash(inputType);
   emitInstruction(instruction);
 }
@@ -720,89 +827,6 @@ void PimCodeGen::emitUnaryVectorOp(pim_binary::Opcode opcode,
   emitInstruction(instruction);
 }
 
-void PimCodeGen::codeGenTransposeOp(const CompiledTransposePlan& plan, const StaticValueKnowledge& knowledge) const {
-  auto srcAddr = addressOf(plan.source, knowledge);
-  auto dstAddr = addressOf(plan.destination, knowledge);
-
-  size_t maxElementOffset = plan.totalBytes == 0 ? 0 : plan.totalBytes - plan.elementBytes;
-  int32_t maxSourceAddress = pim::checkedI32OrCrash(
-    pim::checkedAddOrCrash(srcAddr, maxElementOffset, "transpose source address"), "transpose source address");
-  int32_t maxDestinationAddress =
-    pim::checkedI32OrCrash(pim::checkedAddOrCrash(dstAddr, maxElementOffset, "transpose destination address"),
-                           "transpose destination address");
-  (void) maxSourceAddress;
-  (void) maxDestinationAddress;
-
-  pim_binary::InstructionRecord copyInstruction;
-  copyInstruction.opcode = pim_binary::Opcode::lmv;
-  copyInstruction.rd = 0;
-  copyInstruction.r1 = 1;
-
-  size_t maxRunElements = static_cast<size_t>(std::numeric_limits<int32_t>::max()) / plan.elementBytes;
-  auto emitRun = [&](size_t sourceStart, size_t destinationStart, size_t runLength) {
-    while (runLength != 0) {
-      size_t chunkElements = std::min(runLength, maxRunElements);
-      // totalBytes was checked when the plan was compiled, so these bounded
-      // products cannot overflow.
-      size_t sourceOffset = sourceStart * plan.elementBytes;
-      size_t destinationOffset = destinationStart * plan.elementBytes;
-      size_t byteSize = chunkElements * plan.elementBytes;
-      assert(sourceOffset <= plan.totalBytes - byteSize && destinationOffset <= plan.totalBytes - byteSize);
-      genSetRegisterImmediate(0, static_cast<int32_t>(dstAddr + destinationOffset));
-      genSetRegisterImmediate(1, static_cast<int32_t>(srcAddr + sourceOffset));
-      copyInstruction.generic3 = static_cast<int32_t>(byteSize);
-      emitInstruction(copyInstruction);
-      sourceStart += chunkElements;
-      destinationStart += chunkElements;
-      runLength -= chunkElements;
-    }
-  };
-
-  if (plan.storagePreserving) {
-    emitRun(0, 0, plan.totalElements);
-    return;
-  }
-
-  size_t rank = plan.sourceShape.size();
-  SmallVector<size_t> sourceIndices(rank, 0);
-  size_t destinationFlat = 0;
-  size_t runSourceStart = 0;
-  size_t runDestinationStart = 0;
-  size_t runLength = 0;
-
-  for (size_t sourceFlat = 0; sourceFlat < plan.totalElements; ++sourceFlat) {
-    if (runLength != 0 && destinationFlat != runDestinationStart + runLength) {
-      emitRun(runSourceStart, runDestinationStart, runLength);
-      runSourceStart = sourceFlat;
-      runDestinationStart = destinationFlat;
-      runLength = 0;
-    }
-    if (runLength == 0) {
-      runSourceStart = sourceFlat;
-      runDestinationStart = destinationFlat;
-    }
-    ++runLength;
-
-    if (runLength == maxRunElements) {
-      emitRun(runSourceStart, runDestinationStart, runLength);
-      runLength = 0;
-    }
-
-    if (sourceFlat + 1 == plan.totalElements)
-      break;
-    for (size_t sourceDim = rank; sourceDim-- > 0;) {
-      destinationFlat += plan.destinationStrides[plan.destinationDimensionForSource[sourceDim]];
-      if (++sourceIndices[sourceDim] < plan.sourceShape[sourceDim])
-        break;
-      sourceIndices[sourceDim] = 0;
-      destinationFlat -= plan.destinationRewinds[sourceDim];
-    }
-  }
-
-  if (runLength != 0)
-    emitRun(runSourceStart, runDestinationStart, runLength);
-}
-
 static SmallVector<Operation*> collectTopLevelCoreLikeOps(func::FuncOp funcOp) {
   SmallVector<Operation*> coreLikeOps;
   for (Operation& op : funcOp.getBody().front())
@@ -938,7 +962,7 @@ static LogicalResult executeCompiledCorePlan(
       auto step = node.step.evaluate(knowledge);
       auto forOp = cast<mlir::scf::ForOp>(node.op);
       if (failed(lowerBound) || failed(upperBound) || failed(step) || *step <= 0) {
-        forOp.emitOpError("requires statically evaluable scf.for bounds for PIM codegen");
+        forOp.emitOpError("requires statically evaluable scf.for bounds for Pim codegen");
         return failure();
       }
 
@@ -964,7 +988,7 @@ static LogicalResult executeCompiledCorePlan(
       auto condition = node.condition.evaluate(knowledge);
       auto ifOp = cast<mlir::scf::IfOp>(node.op);
       if (failed(condition)) {
-        ifOp.emitOpError("requires statically evaluable scf.if condition for PIM codegen");
+        ifOp.emitOpError("requires statically evaluable scf.if condition for Pim codegen");
         return failure();
       }
 
@@ -978,7 +1002,7 @@ static LogicalResult executeCompiledCorePlan(
       auto selector = node.condition.evaluate(knowledge);
       auto switchOp = cast<mlir::scf::IndexSwitchOp>(node.op);
       if (failed(selector)) {
-        switchOp.emitOpError("requires a statically evaluable scf.index_switch selector for PIM codegen");
+        switchOp.emitOpError("requires a statically evaluable scf.index_switch selector for Pim codegen");
         return failure();
       }
       const llvm::SmallVectorImpl<CompiledCoreNode>* selectedBody = node.defaultBody.get();
@@ -1012,8 +1036,11 @@ static LogicalResult executeCompiledCorePlan(
       coreCodeGen.codeGenStoreOp(cast<pim::PimMemCopyDevToHostOp>(node.op), knowledge);
       break;
     case CompiledCoreOpKind::Lmv:     coreCodeGen.codeGenLmvOp(cast<pim::PimMemCopyOp>(node.op), knowledge); break;
+    case CompiledCoreOpKind::VMV:     coreCodeGen.codeGenVMVOp(cast<pim::PimVMVOp>(node.op), knowledge); break;
     case CompiledCoreOpKind::Receive: coreCodeGen.codeGenReceiveOp(cast<pim::PimReceiveOp>(node.op), knowledge); break;
     case CompiledCoreOpKind::Send:    coreCodeGen.codeGenSendOp(cast<pim::PimSendOp>(node.op), knowledge); break;
+    case CompiledCoreOpKind::Wait:    coreCodeGen.codeGenWaitOp(cast<pim::PimWaitOp>(node.op), knowledge); break;
+    case CompiledCoreOpKind::Sync:    coreCodeGen.codeGenSyncOp(cast<pim::PimSyncOp>(node.op), knowledge); break;
     case CompiledCoreOpKind::Concat:  coreCodeGen.codeGenConcatOp(cast<pim::PimConcatOp>(node.op), knowledge); break;
     case CompiledCoreOpKind::Vmm:
       if (auto weightSlot = resolveWeightSlot(cast<pim::PimVMMOp>(node.op), knowledge); succeeded(weightSlot))
@@ -1021,7 +1048,6 @@ static LogicalResult executeCompiledCorePlan(
       else
         return failure();
       break;
-    case CompiledCoreOpKind::Transpose: coreCodeGen.codeGenTransposeOp(*node.transposePlan, knowledge); break;
     case CompiledCoreOpKind::VVAdd:     emitBinary(cast<pim::PimVVAddOp>(node.op), pim_binary::Opcode::vvadd); break;
     case CompiledCoreOpKind::VVSub:     emitBinary(cast<pim::PimVVSubOp>(node.op), pim_binary::Opcode::vvsub); break;
     case CompiledCoreOpKind::VVMul:     emitBinary(cast<pim::PimVVMulOp>(node.op), pim_binary::Opcode::vvmul); break;
@@ -1162,12 +1188,12 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
   }
   auto getCompiledProgram = [&](Operation* op) {
     auto it = compiledPrograms.find(op);
-    assert(it != compiledPrograms.end() && "missing compiled PIM core program");
+    assert(it != compiledPrograms.end() && "missing compiled Pim core program");
     return it->second.get();
   };
   auto getMemoryPlan = [&](Operation* op) {
     auto it = memoryPlans.find(op);
-    assert(it != memoryPlans.end() && "missing PIM core memory plan");
+    assert(it != memoryPlans.end() && "missing Pim core memory plan");
     return it->second.get();
   };
 
@@ -1221,7 +1247,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
 
     for (auto [slot, fileName] : llvm::enumerate(weightFiles)) {
       xbarsPerGroup.push_back(weights[slot].shape[1] / static_cast<int64_t>(crossbarSize));
-      std::string sourcePath = outputDirPath + "/weights/" + fileName;
+      std::string sourcePath = "../weights/" + fileName;
       std::string targetPath = coreWeightsDirPath + "/crossbar_" + std::to_string(slot) + ".bin";
       sys::fs::remove(targetPath);
       if (auto error = sys::fs::create_link(sourcePath, targetPath)) {
@@ -1245,7 +1271,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
       if (failed(weightView)) {
         std::string message;
         llvm::raw_string_ostream os(message);
-        os << "requires a statically resolvable dense global weight view during PIM codegen; weight="
+        os << "requires a statically resolvable dense global weight view during Pim codegen; weight="
            << vmmOp.getWeight() << " type=" << vmmOp.getWeight().getType();
         result.recordDiagnostic(vmmOp, os.str());
         return failure();
@@ -1253,7 +1279,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
       if (weightView->shape.size() != 2) {
         std::string message;
         llvm::raw_string_ostream os(message);
-        os << "requires a rank-2 matrix weight view during PIM codegen; resolved shape=[";
+        os << "requires a rank-2 matrix weight view during Pim codegen; resolved shape=[";
         llvm::interleaveComma(weightView->shape, os);
         os << "] weight=" << vmmOp.getWeight() << " type=" << vmmOp.getWeight().getType();
         result.recordDiagnostic(vmmOp, os.str());
@@ -1365,7 +1391,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
   }
   if (diagnostics.hasFailure())
     diagnostics.emitSuppressedSummary(summaryAnchor ? summaryAnchor : moduleOp.getOperation(),
-                                      "PIM codegen diagnostic(s)");
+                                      "Pim codegen diagnostic(s)");
 
   for (size_t jobIndex = 0; jobIndex < jobs.size(); ++jobIndex)
     if (jobResults[jobIndex].status != CompilerSuccess)
@@ -1431,7 +1457,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
       if (!batchPerCoreRow)
         batchPerCoreRow = result.reportRow;
       else if (!(*batchPerCoreRow == result.reportRow))
-        llvm_unreachable("one PIM core batch produced inconsistent per-core memory reports");
+        llvm_unreachable("one Pim core batch produced inconsistent per-core memory reports");
     }
 
     uint64_t batchReportId = jobs[group.front()].batchReportId.value_or(0);

src/PIM/Compiler/PimCodeGen.hpp

@@ -25,7 +25,6 @@
 namespace onnx_mlir {
 
 struct CompiledCoreProgram;
-struct CompiledTransposePlan;
 class PimInstructionWriter;
 
 struct MemEntry {
@@ -177,7 +176,7 @@ class PimCodeGen {
   void genSetRegisterImmediateUnsigned(size_t registerNumber, size_t immediate) const;
   void setupRd(size_t rdAddress, size_t rdOffset) const;
   void setupRdRs1(size_t rdAddress, size_t rdOffset, size_t rs1Address, size_t rs1Offset) const;
-  void setupRdRs1Rs2(
+  std::array<uint8_t, 3> setupRdRs1Rs2(
     size_t rdAddress, size_t rdOffset, size_t rs1Address, size_t rs1Offset, size_t rs2Address, size_t rs2Offset) const;
 
   void emitMemCopyOp(pim_binary::Opcode opcode,
@@ -214,15 +213,17 @@ public:
   void codeGenLoadOp(pim::PimMemCopyHostToDevOp loadOp, const StaticValueKnowledge& knowledge) const;
   void codeGenStoreOp(pim::PimMemCopyDevToHostOp storeOp, const StaticValueKnowledge& knowledge) const;
   void codeGenLmvOp(pim::PimMemCopyOp lmvOp, const StaticValueKnowledge& knowledge) const;
+  void codeGenVMVOp(pim::PimVMVOp vmvOp, const StaticValueKnowledge& knowledge) const;
 
   void codeGenReceiveOp(pim::PimReceiveOp receiveOp, const StaticValueKnowledge& knowledge) const;
   void codeGenSendOp(pim::PimSendOp sendOp, const StaticValueKnowledge& knowledge) const;
+  void codeGenWaitOp(pim::PimWaitOp waitOp, const StaticValueKnowledge& knowledge) const;
+  void codeGenSyncOp(pim::PimSyncOp syncOp, const StaticValueKnowledge& knowledge) const;
   void codeGenConcatOp(pim::PimConcatOp concatOp, const StaticValueKnowledge& knowledge) const;
 
   template <typename MVMTy>
   void codeGenMVMLikeOp(size_t mvmId, MVMTy mvmLikeOp, bool transposeMatrix, const StaticValueKnowledge& knowledge);
 
-  void codeGenTransposeOp(const CompiledTransposePlan& plan, const StaticValueKnowledge& knowledge) const;
 };
 
 OnnxMlirCompilerErrorCodes compileToPimCode(mlir::ModuleOp& moduleOpRef, std::string& outputDirName);

src/PIM/Compiler/PimCompilerOptions.cpp

@@ -2,30 +2,32 @@
 
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 
+#include <limits>
+
 #define DEBUG_TYPE "PimCompilerOptions"
 
 namespace onnx_mlir {
 
 llvm::cl::opt<PimEmissionTargetType> pimEmissionTarget(
-  llvm::cl::desc("[Optional] Choose PIM-related target to emit (once selected it will cancel the other targets):"),
-  llvm::cl::values(clEnumVal(EmitSpatial, "Lower model to spatial IR")),
-  llvm::cl::values(clEnumVal(EmitPim, "Lower model to PIM IR")),
-  llvm::cl::values(clEnumVal(EmitPimBufferized, "Lower model to PIM IR and bufferize it")),
-  llvm::cl::values(clEnumVal(EmitPimCodegen, "Lower model to PIM IR and generate code for PIM")),
+  llvm::cl::desc("[Optional] Choose Pim-related target to emit (once selected it will cancel the other targets):"),
+  llvm::cl::values(clEnumVal(EmitSpatial, "Lower model to Spatial IR")),
+  llvm::cl::values(clEnumVal(EmitPim, "Lower model to Pim IR")),
+  llvm::cl::values(clEnumVal(EmitPimBufferized, "Lower model to Pim IR and bufferize it")),
+  llvm::cl::values(clEnumVal(EmitPimCodegen, "Lower model to Pim IR and generate code for Pim")),
   llvm::cl::init(EmitPimCodegen),
   llvm::cl::cat(OnnxMlirOptions));
 
 llvm::cl::opt<PimMemoryReportLevel> pimMemoryReport(
   "pim-memory-report",
-  llvm::cl::desc("Emit a human-readable PIM memory planning report"),
-  llvm::cl::values(clEnumValN(PimMemoryReportNone, "none", "Do not emit any PIM memory planning report")),
-  llvm::cl::values(clEnumValN(PimMemoryReportSummary, "summary", "Emit a concise PIM memory summary")),
+  llvm::cl::desc("Emit a human-readable Pim memory planning report"),
+  llvm::cl::values(clEnumValN(PimMemoryReportNone, "none", "Do not emit any Pim memory planning report")),
+  llvm::cl::values(clEnumValN(PimMemoryReportSummary, "summary", "Emit a concise Pim memory summary")),
   llvm::cl::init(PimMemoryReportSummary),
   llvm::cl::cat(OnnxMlirOptions));
 
 llvm::cl::opt<PimConvLoweringType> pimConvLowering(
   "pim-conv-lowering",
-  llvm::cl::desc("Convolution lowering strategy for PIM"),
+  llvm::cl::desc("Convolution lowering strategy for Pim"),
   llvm::cl::values(clEnumValN(PimConvLoweringAuto, "auto", "Select the Conv lowering strategy automatically")),
   llvm::cl::values(clEnumValN(PimConvLoweringLegacy, "legacy", "Use the legacy explicit-im2col Conv lowering")),
   llvm::cl::values(clEnumValN(PimConvLoweringDepthwise, "depthwise", "Force the depthwise-specialized Conv lowering")),
@@ -53,28 +55,23 @@ llvm::cl::opt<PimSpatialDataflowExportType> pimExportSpatialDataflow(
   llvm::cl::desc("Emit Gephi-importable CSV dataflow reports for Spatial pipeline snapshots"),
   llvm::cl::values(clEnumValN(SpatialDataflowExportNone, "none", "Do not emit Spatial dataflow CSV reports")),
   llvm::cl::values(
-    clEnumValN(SpatialDataflowExportSpatial1, "spatial1", "Emit spatial1 graph dataflow CSV reports")),
+    clEnumValN(SpatialDataflowExportSpatial1, "spatial1", "Emit Spatial1 graph dataflow CSV reports")),
   llvm::cl::values(
-    clEnumValN(SpatialDataflowExportSpatial2, "spatial2", "Emit spatial2 trivially merged graph dataflow CSV reports")),
+    clEnumValN(SpatialDataflowExportSpatial2, "spatial2", "Emit Spatial2 trivially merged graph dataflow CSV reports")),
   llvm::cl::values(
-    clEnumValN(SpatialDataflowExportSpatial3, "spatial3", "Emit spatial3 scheduled dataflow CSV reports")),
+    clEnumValN(SpatialDataflowExportSpatial3, "spatial3", "Emit Spatial3 scheduled dataflow CSV reports")),
   llvm::cl::values(
-    clEnumValN(SpatialDataflowExportSpatial4, "spatial4", "Emit spatial4 realized dataflow CSV reports")),
+    clEnumValN(SpatialDataflowExportSpatial4, "spatial4", "Emit Spatial4 realized dataflow CSV reports")),
   llvm::cl::values(clEnumValN(SpatialDataflowExportAll, "all", "Emit all Spatial dataflow CSV reports")),
   llvm::cl::init(SpatialDataflowExportNone),
   llvm::cl::cat(OnnxMlirOptions));
 
 llvm::cl::opt<bool>
   pimOnlyCodegen("pim-only-codegen",
-                 llvm::cl::desc("Only generate code for PIM (assume input is already in bufferized PIM IR)"),
+                 llvm::cl::desc("Only generate code for Pim (assume input is already in bufferized Pim IR)"),
                  llvm::cl::init(false),
                  llvm::cl::cat(OnnxMlirOptions));
 
-llvm::cl::opt<bool> useExperimentalConvImpl("use-experimental-conv-impl",
-                                            llvm::cl::desc("Use experimental implementation for convolution"),
-                                            llvm::cl::init(false),
-                                            llvm::cl::cat(OnnxMlirOptions));
-
 llvm::cl::opt<uint64_t> pimConvIm2colMaxElements(
   "pim-conv-im2col-max-elements",
   llvm::cl::desc("Maximum number of im2col elements to materialize globally for one Conv before streaming/chunking"),
@@ -99,50 +96,74 @@ llvm::cl::opt<bool> pimEmitJson("pim-emit-json",
 
 llvm::cl::opt<bool> pimDetectCommunicationDeadlock(
   "pim-detect-communication-deadlock",
-  llvm::cl::desc("Expensively simulate the statically expanded PIM send/receive order at verification time and fail if a blocking communication deadlock is found"),
+  llvm::cl::desc("Expensively simulate the statically expanded Pim send/receive order at verification time and fail if a blocking communication deadlock is found"),
   llvm::cl::init(false),
   llvm::cl::cat(OnnxMlirOptions));
 
-llvm::cl::opt<bool> pimMaterializeScalarFanoutGlobalOrder(
-  "pim-materialize-scalar-fanout-global-order",
-  llvm::cl::desc("Experimental expensive materializer mode: emit scalar-source fanout as globally ordered communication events instead of all-send fanout loops"),
+llvm::cl::opt<bool> pimVerifyBufferizationCopyFreedom(
+  "pim-verify-bufferization-copy-freedom",
+  llvm::cl::desc("Run the expensive official Pim tensor-copy freedom proof before bufferization"),
   llvm::cl::init(false),
   llvm::cl::cat(OnnxMlirOptions));
 
-llvm::cl::opt<bool> pimTraceCommunicationMaterialization(
-  "pim-trace-communication-materialization",
-  llvm::cl::desc("Emit verbose materializer-time diagnostics and provenance attributes for every Spatial communication op"),
+llvm::cl::opt<bool> pimDisableSynchronization(
+  "pim-disable-synchronization",
+  llvm::cl::desc("Omit Pim wait/sync instructions from generated code for performance ablation"),
+  llvm::cl::init(false),
+  llvm::cl::cat(OnnxMlirOptions));
+
+llvm::cl::opt<bool> pimDisableSpatialPlanning(
+  "pim-disable-spatial-planning",
+  llvm::cl::desc("Select the trivial Spatial layout plan for performance ablation"),
   llvm::cl::init(false),
   llvm::cl::cat(OnnxMlirOptions));
 
 llvm::cl::opt<size_t>
-  crossbarSize("crossbar-size", llvm::cl::desc("Width and height of a single crossbar"), llvm::cl::init(128));
+  crossbarSize("crossbar-size",
+               llvm::cl::desc("Width and height of a single crossbar (required for Pim compilation)"),
+               llvm::cl::init(0));
 
 llvm::cl::opt<size_t>
-  crossbarCountInCore("crossbar-count", llvm::cl::desc("Number of crossbars in each core"), llvm::cl::init(64));
+  crossbarCountInCore("crossbar-count",
+                      llvm::cl::desc("Number of crossbars in each core (required for Pim compilation)"),
+                      llvm::cl::init(0));
+
+llvm::cl::opt<size_t> pipelineStages(
+  "pipeline",
+  llvm::cl::desc("Number of throughput pipeline stages (1 preserves latency scheduling)"),
+  llvm::cl::init(1),
+  llvm::cl::cat(OnnxMlirOptions));
 
 llvm::cl::opt<long> coresCount("core-count",
-                               llvm::cl::desc("Number of cores in the chip. Required for PIM compilation."),
+                               llvm::cl::desc("Number of cores in the chip. Required for Pim compilation."),
                                llvm::cl::init(-1));
 
 llvm::cl::opt<std::string> pimTargetConfig(
   "pim-target-config",
-  llvm::cl::desc("PIM target configuration used to construct the Spatial scheduling cost model"),
+  llvm::cl::desc("Pim target configuration used to construct the Spatial scheduling cost model"),
   llvm::cl::init(""),
   llvm::cl::cat(OnnxMlirOptions));
 
-llvm::cl::opt<bool>
-  ignoreConcatError("ignore-concat-error",
-                    llvm::cl::desc("Ignore ConcatOp corner case: do not assert and do a simplification"),
-                    llvm::cl::init(false));
-
-bool hasExplicitPimCoreCount() { return coresCount.getNumOccurrences() != 0; }
-
-void verifyExplicitPimCoreCount() {
-  if (!hasExplicitPimCoreCount())
-    llvm::report_fatal_error("PIM compilation requires an explicit --core-count=<positive integer>");
+void verifyPimCompilerOptions() {
+  if (coresCount.getNumOccurrences() == 0)
+    llvm::report_fatal_error("Pim compilation requires an explicit --core-count=<positive integer>");
   if (coresCount.getValue() <= 0)
-    llvm::report_fatal_error("PIM compilation requires --core-count to be a positive integer");
+    llvm::report_fatal_error("Pim compilation requires --core-count to be a positive integer");
+  if (crossbarSize.getNumOccurrences() == 0)
+    llvm::report_fatal_error("Pim compilation requires an explicit --crossbar-size=<positive integer>");
+  if (crossbarSize.getValue() == 0)
+    llvm::report_fatal_error("Pim compilation requires --crossbar-size to be a positive integer");
+  if (crossbarCountInCore.getNumOccurrences() == 0)
+    llvm::report_fatal_error("Pim compilation requires an explicit --crossbar-count=<positive integer>");
+  if (crossbarCountInCore.getValue() == 0)
+    llvm::report_fatal_error("Pim compilation requires --crossbar-count to be a positive integer");
+  if (pipelineStages.getValue() == 0)
+    llvm::report_fatal_error("Pim compilation requires --pipeline to be positive");
+  if (static_cast<size_t>(coresCount.getValue()) < pipelineStages.getValue())
+    llvm::report_fatal_error("Pim compilation requires --pipeline not to exceed --core-count");
+  if (crossbarCountInCore.getValue()
+      > std::numeric_limits<size_t>::max() / pipelineStages.getValue())
+    llvm::report_fatal_error("Pim compilation --crossbar-count * --pipeline overflows");
 }
 
 } // namespace onnx_mlir

src/PIM/Compiler/PimCompilerOptions.hpp

@@ -55,29 +55,21 @@ extern llvm::cl::opt<PimConvLoweringType> pimConvLowering;
 extern llvm::cl::opt<PimSpatialDataflowExportType> pimExportSpatialDataflow;
 
 extern llvm::cl::opt<bool> pimOnlyCodegen;
-extern llvm::cl::opt<bool> useExperimentalConvImpl;
 extern llvm::cl::opt<bool> pimEmitJson;
 extern llvm::cl::opt<bool> pimReportConvLowering;
 extern llvm::cl::opt<bool> pimDetectCommunicationDeadlock;
-extern llvm::cl::opt<bool> pimMaterializeScalarFanoutGlobalOrder;
-extern llvm::cl::opt<bool> pimTraceCommunicationMaterialization;
+extern llvm::cl::opt<bool> pimVerifyBufferizationCopyFreedom;
+extern llvm::cl::opt<bool> pimDisableSynchronization;
+extern llvm::cl::opt<bool> pimDisableSpatialPlanning;
 
 extern llvm::cl::opt<size_t> crossbarSize;
 extern llvm::cl::opt<size_t> crossbarCountInCore;
+extern llvm::cl::opt<size_t> pipelineStages;
 extern llvm::cl::opt<long> coresCount;
 extern llvm::cl::opt<std::string> pimTargetConfig;
 extern llvm::cl::opt<uint64_t> pimConvIm2colMaxElements;
 extern llvm::cl::opt<uint64_t> pimConvStreamChunkPositions;
 
-bool hasExplicitPimCoreCount();
-void verifyExplicitPimCoreCount();
-
-// 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
-// wanted tile is generated by two separate operands of the concat. If this is
-// set to false, this corner case will assert an error. If this is set to true,
-// a simplification is performed and only the tile from the first operand is
-// taken.
-extern llvm::cl::opt<bool> ignoreConcatError;
+void verifyPimCompilerOptions();
 
 } // namespace onnx_mlir

src/PIM/Compiler/PimCompilerUtils.cpp

@@ -12,11 +12,15 @@
 #include <limits>
 #include <tuple>
 
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/SchedulingTarget.hpp"
-#include "src/Accelerators/PIM/Pass/PIMPasses.h"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/SchedulingTarget.hpp"
+#include "src/Accelerators/PIM/Passes/PIMPasses.h"
 #include "src/Compiler/CompilerPasses.hpp"
 
 #define DEBUG_TYPE "PimCompilerUtils"
@@ -75,17 +79,66 @@ spatial::SchedulingTarget getDefaultPimSchedulingTarget() {
   target.residentWeightCapacity = crossbarCountInCore.getValue();
   target.matrixRows = crossbarSize.getValue();
   target.matrixColumns = crossbarSize.getValue();
+  target.synchronizationRegisterCount = kPimEventRegisterCount;
 
   setDefaultPimInterProcessorLatencies(target);
   return target;
 }
 
+spatial::ConvLoweringStrategy getSpatialConvLoweringStrategy(PimConvLoweringType strategy) {
+  switch (strategy) {
+  case PimConvLoweringAuto: return spatial::ConvLoweringStrategy::Auto;
+  case PimConvLoweringLegacy: return spatial::ConvLoweringStrategy::Legacy;
+  case PimConvLoweringDepthwise: return spatial::ConvLoweringStrategy::Depthwise;
+  case PimConvLoweringPackedIm2Col: return spatial::ConvLoweringStrategy::PackedIm2Col;
+  case PimConvLoweringStreamedPatch: return spatial::ConvLoweringStrategy::StreamedPatch;
+  case PimConvLoweringStreamedPacked: return spatial::ConvLoweringStrategy::StreamedPacked;
+  case PimConvLoweringOutputChannelTiled: return spatial::ConvLoweringStrategy::OutputChannelTiled;
+  case PimConvLoweringInputKTiled: return spatial::ConvLoweringStrategy::InputKTiled;
+  case PimConvLoweringTiled2D: return spatial::ConvLoweringStrategy::Tiled2D;
+  }
+  llvm_unreachable("unknown Pim Conv lowering strategy");
+}
+
+spatial::SpatialDataflowExportStage getPimSpatialDataflowExportStage(
+    PimSpatialDataflowExportType stage) {
+  switch (stage) {
+  case SpatialDataflowExportNone:     return spatial::SpatialDataflowExportStage::None;
+  case SpatialDataflowExportSpatial1: return spatial::SpatialDataflowExportStage::Spatial1;
+  case SpatialDataflowExportSpatial2: return spatial::SpatialDataflowExportStage::Spatial2;
+  case SpatialDataflowExportSpatial3: return spatial::SpatialDataflowExportStage::Spatial3;
+  case SpatialDataflowExportSpatial4: return spatial::SpatialDataflowExportStage::Spatial4;
+  case SpatialDataflowExportAll:      return spatial::SpatialDataflowExportStage::All;
+  }
+  llvm_unreachable("unknown Pim Spatial dataflow export stage");
+}
+
+spatial::SpatialTargetResources getPimSpatialTargetResources(const spatial::SchedulingTarget& target) {
+  spatial::SpatialTargetResources resources;
+  resources.matrixShape = {target.matrixRows, target.matrixColumns};
+  resources.matrixUnitsPerProcessor = target.residentWeightCapacity;
+  resources.processorCount = target.processorCount;
+  resources.vectorWidth = target.vectorWidth;
+  if (failed(resources.verify()))
+    llvm::report_fatal_error("Pim target resources are incomplete");
+  return resources;
+}
+
+ONNXToSpatialPlanningOptions getPimONNXToSpatialPlanningOptions() {
+  ONNXToSpatialPlanningOptions options;
+  options.convIm2colMaxElements = pimConvIm2colMaxElements.getValue();
+  options.convStreamChunkPositions = pimConvStreamChunkPositions.getValue();
+  options.forcedConvStrategy = getSpatialConvLoweringStrategy(pimConvLowering.getValue());
+  options.reportConvLowering = pimReportConvLowering.getValue();
+  return options;
+}
+
 const llvm::json::Object& requireObject(const llvm::json::Object& object,
                                         llvm::StringRef key,
                                         llvm::StringRef path) {
   const llvm::json::Object* nested = object.getObject(key);
   if (!nested)
-    llvm::report_fatal_error("PIM target config is missing object '" + path + "." + key + "'");
+    llvm::report_fatal_error("Pim target config is missing object '" + path + "." + key + "'");
   return *nested;
 }
 
@@ -98,7 +151,7 @@ Cost getConfigCost(const llvm::json::Object& object,
     return fallback;
   if (!std::isfinite(*number) || *number < 0.0 || (!allowZero && *number == 0.0)
       || *number > static_cast<double>(std::numeric_limits<Cost>::max()))
-    llvm::report_fatal_error("PIM target config field '" + key + "' must be a valid positive number");
+    llvm::report_fatal_error("Pim target config field '" + key + "' must be a valid positive number");
   return static_cast<Cost>(std::ceil(*number));
 }
 
@@ -106,11 +159,11 @@ std::pair<size_t, size_t> getConfigPair(const llvm::json::Object& object,
                                         llvm::StringRef key) {
   const llvm::json::Array* values = object.getArray(key);
   if (!values || values->size() != 2)
-    llvm::report_fatal_error("PIM target config field '" + key + "' must contain two integers");
+    llvm::report_fatal_error("Pim target config field '" + key + "' must contain two integers");
   std::optional<int64_t> first = (*values)[0].getAsInteger();
   std::optional<int64_t> second = (*values)[1].getAsInteger();
   if (!first || !second || *first <= 0 || *second <= 0)
-    llvm::report_fatal_error("PIM target config field '" + key + "' must contain two positive integers");
+    llvm::report_fatal_error("Pim target config field '" + key + "' must contain two positive integers");
   return {static_cast<size_t>(*first), static_cast<size_t>(*second)};
 }
 
@@ -121,7 +174,7 @@ void loadPimInterProcessorLatencies(
     network.getString("net_config_file_path");
   if (!filename)
     llvm::report_fatal_error(
-      "PIM target config is missing network latency file path");
+      "Pim target config is missing network latency file path");
 
   llvm::SmallString<256> networkPath(*filename);
   if (!llvm::sys::path::is_absolute(networkPath)) {
@@ -134,19 +187,19 @@ void loadPimInterProcessorLatencies(
   auto buffer = llvm::MemoryBuffer::getFile(networkPath);
   if (!buffer)
     llvm::report_fatal_error(
-      llvm::Twine("failed to read PIM network config '")
+      llvm::Twine("failed to read Pim network config '")
       + networkPath + "': " + buffer.getError().message());
   auto parsed = llvm::json::parse(buffer.get()->getBuffer());
   if (!parsed)
     llvm::report_fatal_error(
-      llvm::Twine("failed to parse PIM network config '")
+      llvm::Twine("failed to parse Pim network config '")
       + networkPath + "': " + llvm::toString(parsed.takeError()));
   const llvm::json::Object* root = parsed->getAsObject();
   const llvm::json::Object* latencies =
     root ? root->getObject("latency") : nullptr;
   if (!latencies)
     llvm::report_fatal_error(
-      "PIM network config is missing its latency matrix");
+      "Pim network config is missing its latency matrix");
 
   target.interProcessorLatencyNs.assign(
     target.processorCount * target.processorCount, 0);
@@ -157,7 +210,7 @@ void loadPimInterProcessorLatencies(
     const llvm::json::Object* row = latencies->getObject(sourceKey);
     if (!row)
       llvm::report_fatal_error(
-        llvm::Twine("PIM network config is missing latency row ")
+        llvm::Twine("Pim network config is missing latency row ")
         + sourceKey);
     for (size_t destination = 0;
          destination < target.processorCount; ++destination) {
@@ -167,7 +220,7 @@ void loadPimInterProcessorLatencies(
       std::optional<double> latency = row->getNumber(destinationKey);
       if (!latency || !std::isfinite(*latency) || *latency <= 0.0)
         llvm::report_fatal_error(
-          llvm::Twine("PIM network config is missing latency ")
+          llvm::Twine("Pim network config is missing latency ")
           + sourceKey + " -> " + destinationKey);
       Cost roundedLatency = static_cast<Cost>(std::ceil(*latency));
       target.interProcessorLatencyNs[
@@ -191,17 +244,17 @@ spatial::SchedulingTarget getPimSchedulingTarget() {
   auto buffer = llvm::MemoryBuffer::getFile(pimTargetConfig);
   if (!buffer)
     llvm::report_fatal_error(
-      llvm::Twine("failed to read PIM target config '")
+      llvm::Twine("failed to read Pim target config '")
       + pimTargetConfig.getValue() + "': " + buffer.getError().message());
   auto parsed = llvm::json::parse(buffer.get()->getBuffer());
   if (!parsed)
     llvm::report_fatal_error(
-      llvm::Twine("failed to parse PIM target config '")
+      llvm::Twine("failed to parse Pim target config '")
       + pimTargetConfig.getValue() + "': "
       + llvm::toString(parsed.takeError()));
   const llvm::json::Object* root = parsed->getAsObject();
   if (!root)
-    llvm::report_fatal_error("PIM target config must contain a JSON object");
+    llvm::report_fatal_error("Pim target config must contain a JSON object");
 
   const llvm::json::Object& chip = requireObject(*root, "chip_config", "root");
   const llvm::json::Object& core = requireObject(chip, "core_config", "chip_config");
@@ -214,7 +267,7 @@ spatial::SchedulingTarget getPimSchedulingTarget() {
 
   std::optional<int64_t> coreCount = chip.getInteger("core_cnt");
   if (!coreCount || *coreCount <= 0)
-    llvm::report_fatal_error("PIM target config field 'core_cnt' must be a positive integer");
+    llvm::report_fatal_error("Pim target config field 'core_cnt' must be a positive integer");
   target.processorCount = static_cast<size_t>(*coreCount);
   target.residentWeightCapacity =
     getConfigCost(matrix, "xbar_array_count", target.residentWeightCapacity);
@@ -225,7 +278,7 @@ spatial::SchedulingTarget getPimSchedulingTarget() {
       || target.residentWeightCapacity != crossbarCountInCore.getValue()
       || target.matrixRows != crossbarSize.getValue()
       || target.matrixColumns != crossbarSize.getValue())
-    llvm::report_fatal_error("PIM target config resources do not match --core-count, "
+    llvm::report_fatal_error("Pim target config resources do not match --core-count, "
                              "--crossbar-count, and --crossbar-size");
   loadPimInterProcessorLatencies(target, network);
 
@@ -278,11 +331,14 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
                   PassManager& pm,
                   EmissionTargetType& emissionTarget,
                   std::string outputNameNoExt) {
-  verifyExplicitPimCoreCount();
+  verifyPimCompilerOptions();
+  spatial::SchedulingTarget schedulingTarget = getPimSchedulingTarget();
+  spatial::SpatialTargetResources targetResources = getPimSpatialTargetResources(schedulingTarget);
 
   if (pimOnlyCodegen) {
+    pm.addPass(createPimInstructionSelectionPass());
     pm.addPass(createPimLocalMemoryPlanningPass());
-    pm.addPass(createPimVerificationPass());
+    pm.addPass(createPimVerificationPass(targetResources, pimDetectCommunicationDeadlock.getValue()));
     pm.addPass(createEmitPimCodePass());
     return;
   }
@@ -291,23 +347,30 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
     addONNXToMLIRPasses(pm, /*target CPU*/ false);
 
   if (pimEmissionTarget >= EmitSpatial) {
-    spatial::SchedulingTarget schedulingTarget = getPimSchedulingTarget();
-    pm.addPass(createONNXToSpatialPass());
-    pm.addPass(createSpatialLayoutPlanningPass());
-    pm.addPass(createLowerSpatialPlansPass());
+    ONNXToSpatialPlanningOptions planningOptions = getPimONNXToSpatialPlanningOptions();
+    spatial::SpatialDataflowExportStage exportStage =
+      getPimSpatialDataflowExportStage(pimExportSpatialDataflow.getValue());
+    pm.addPass(createONNXToSpatialPass(targetResources, planningOptions));
+    pm.addPass(createSpatialLayoutPlanningPass(
+      targetResources, pimDisableSpatialPlanning.getValue()));
+    pm.addPass(createLowerSpatialPlansPass(targetResources, planningOptions, exportStage));
     pm.addPass(createTrivialGraphComputeMergePass(
-      schedulingTarget.residentWeightCapacity));
-    pm.addPass(createMergeComputeNodesPass(schedulingTarget));
+      schedulingTarget.residentWeightCapacity, exportStage));
+    pm.addPass(spatial::createScheduleAndRealizeSpatialPass(
+      schedulingTarget, exportStage, pipelineStages.getValue()));
     pm.addPass(createMessagePass("Onnx lowered to Spatial"));
   }
 
   if (pimEmissionTarget >= EmitPim) {
-    pm.addPass(createSpatialToPimPass());
+    pm.addPass(createSpatialToPimPass(targetResources));
     pm.addPass(createMessagePass("Spatial lowered to Pim"));
   }
 
   if (pimEmissionTarget >= EmitPimBufferized) {
-    pm.addPass(createPimBufferizationPass());
+    pm.addPass(createPimBufferizationPreparationPass(pimVerifyBufferizationCopyFreedom.getValue()));
+    pm.addPass(createPimOneShotBufferizationPass());
+    pm.addPass(createPimMemoryNormalizationPass());
+    pm.addPass(createPimBufferizationVerificationPass());
     pm.addPass(createMessagePass("Pim bufferized"));
   }
 
@@ -315,9 +378,11 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
     pm.addPass(mlir::createLowerAffinePass());
     pm.addPass(createPimHostConstantFoldingPass());
     pm.addPass(createMessagePass("Pim host constants folded"));
+    pm.addPass(createPimInstructionSelectionPass());
+    pm.addPass(createMessagePass("Pim instructions selected"));
     pm.addPass(createPimLocalMemoryPlanningPass());
     pm.addPass(createMessagePass("Pim local memory planned"));
-    pm.addPass(createPimVerificationPass());
+    pm.addPass(createPimVerificationPass(targetResources, pimDetectCommunicationDeadlock.getValue()));
     pm.addPass(createMessagePass("Pim verified"));
     pm.addPass(createEmitPimCodePass());
     pm.addPass(createMessagePass("Pim code emitted"));

src/PIM/Compiler/PimCoreProgram.cpp

@@ -1,12 +1,6 @@
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
-#include "mlir/IR/BuiltinTypes.h"
-
-#include "llvm/ADT/STLExtras.h"
-
 #include "src/Accelerators/PIM/Common/IR/CoreBlockUtils.hpp"
-#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
-#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCoreProgram.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 
@@ -20,11 +14,13 @@ static FailureOr<CompiledCoreOpKind> classifyCompiledCoreOpKind(Operation& op) {
   if (isa<pim::PimMemCopyHostToDevOp>(op)) return CompiledCoreOpKind::Load;
   if (isa<pim::PimMemCopyDevToHostOp>(op)) return CompiledCoreOpKind::Store;
   if (isa<pim::PimMemCopyOp>(op)) return CompiledCoreOpKind::Lmv;
+  if (isa<pim::PimVMVOp>(op)) return CompiledCoreOpKind::VMV;
   if (isa<pim::PimReceiveOp>(op)) return CompiledCoreOpKind::Receive;
   if (isa<pim::PimSendOp>(op)) return CompiledCoreOpKind::Send;
+  if (isa<pim::PimWaitOp>(op)) return CompiledCoreOpKind::Wait;
+  if (isa<pim::PimSyncOp>(op)) return CompiledCoreOpKind::Sync;
   if (isa<pim::PimConcatOp>(op)) return CompiledCoreOpKind::Concat;
   if (isa<pim::PimVMMOp>(op)) return CompiledCoreOpKind::Vmm;
-  if (isa<pim::PimTransposeOp>(op)) return CompiledCoreOpKind::Transpose;
   if (isa<pim::PimVVAddOp>(op)) return CompiledCoreOpKind::VVAdd;
   if (isa<pim::PimVVSubOp>(op)) return CompiledCoreOpKind::VVSub;
   if (isa<pim::PimVVMulOp>(op)) return CompiledCoreOpKind::VVMul;
@@ -38,77 +34,6 @@ static FailureOr<CompiledCoreOpKind> classifyCompiledCoreOpKind(Operation& op) {
   return failure();
 }
 
-static bool isStoragePreservingTranspose(ArrayRef<size_t> sourceShape, ArrayRef<int64_t> permutation) {
-  SmallVector<unsigned> sourceNonUnitDims;
-  SmallVector<unsigned> destinationSourceNonUnitDims;
-  for (auto [dim, size] : llvm::enumerate(sourceShape))
-    if (size != 1)
-      sourceNonUnitDims.push_back(dim);
-  for (int64_t sourceDim : permutation)
-    if (sourceShape[sourceDim] != 1)
-      destinationSourceNonUnitDims.push_back(static_cast<unsigned>(sourceDim));
-  return sourceNonUnitDims == destinationSourceNonUnitDims;
-}
-
-static FailureOr<CompiledTransposePlan> compileTransposePlan(pim::PimTransposeOp transposeOp) {
-  auto sourceType = cast<ShapedType>(transposeOp.getInput().getType());
-  ArrayRef<int64_t> sourceShape = sourceType.getShape();
-  size_t rank = sourceShape.size();
-  CompiledTransposePlan plan;
-  plan.source = transposeOp.getInput();
-  plan.destination = transposeOp.getOutputBuffer();
-  plan.elementBytes = getElementTypeSizeInBytes(sourceType.getElementType());
-  auto totalElements = pim::checkedSize(sourceType.getNumElements(), transposeOp, "transpose elements");
-  if (failed(totalElements)) return failure();
-  plan.totalElements = *totalElements;
-  auto totalBytes = pim::checkedMul(plan.totalElements, plan.elementBytes, transposeOp, "transpose byte size");
-  if (failed(totalBytes)) return failure();
-  plan.totalBytes = *totalBytes;
-
-  SmallVector<int64_t> permutation = map_to_vector(transposeOp.getPermutation().getAsRange<IntegerAttr>(),
-                                                   [](IntegerAttr attr) { return attr.getInt(); });
-  if (permutation.size() != rank) {
-    transposeOp.emitOpError("requires permutation rank to match source rank for PIM codegen");
-    return failure();
-  }
-
-  SmallVector<size_t> destinationShape(rank);
-  plan.destinationStrides.assign(rank, 1);
-  plan.destinationDimensionForSource.assign(rank, 0);
-  plan.destinationRewinds.assign(rank, 0);
-  SmallVector<bool> seenSourceDimensions(rank, false);
-  for (size_t dim = 0; dim < rank; ++dim) {
-    auto size = pim::checkedSize(sourceShape[dim], transposeOp, "transpose source dimension");
-    if (failed(size)) return failure();
-    plan.sourceShape.push_back(*size);
-  }
-  for (auto [destinationDim, sourceDim] : llvm::enumerate(permutation)) {
-    if (sourceDim < 0 || static_cast<size_t>(sourceDim) >= rank || seenSourceDimensions[sourceDim]) {
-      transposeOp.emitOpError("requires a valid permutation containing each source dimension exactly once");
-      return failure();
-    }
-    seenSourceDimensions[sourceDim] = true;
-    destinationShape[destinationDim] = plan.sourceShape[sourceDim];
-    plan.destinationDimensionForSource[sourceDim] = destinationDim;
-  }
-  for (size_t dim = rank; dim > 1; --dim) {
-    auto stride = pim::checkedMul(
-      plan.destinationStrides[dim - 1], destinationShape[dim - 1], transposeOp, "transpose destination stride");
-    if (failed(stride)) return failure();
-    plan.destinationStrides[dim - 2] = *stride;
-  }
-  for (size_t sourceDim = 0; sourceDim < rank; ++sourceDim) {
-    auto rewind = pim::checkedMul(plan.sourceShape[sourceDim],
-                                  plan.destinationStrides[plan.destinationDimensionForSource[sourceDim]],
-                                  transposeOp,
-                                  "transpose destination rewind");
-    if (failed(rewind)) return failure();
-    plan.destinationRewinds[sourceDim] = *rewind;
-  }
-  plan.storagePreserving = isStoragePreservingTranspose(plan.sourceShape, permutation);
-  return plan;
-}
-
 static LogicalResult compileCoreEmissionPlan(Block& block, SmallVectorImpl<CompiledCoreNode>& plan) {
   for (Operation& op : block) {
     if (isa<pim::PimHaltOp, scf::YieldOp, memref::GetGlobalOp>(op) || isCoreStaticAddressOp(&op))
@@ -121,7 +46,7 @@ static LogicalResult compileCoreEmissionPlan(Block& block, SmallVectorImpl<Compi
       auto upper = compileIndexExpr(forOp.getUpperBound());
       auto step = compileIndexExpr(forOp.getStep());
       if (failed(lower) || failed(upper) || failed(step)) {
-        forOp.emitOpError("requires statically evaluable scf.for bounds for PIM codegen");
+        forOp.emitOpError("requires statically evaluable scf.for bounds for Pim codegen");
         return failure();
       }
       CompiledCoreNode node;
@@ -138,7 +63,7 @@ static LogicalResult compileCoreEmissionPlan(Block& block, SmallVectorImpl<Compi
     if (auto ifOp = dyn_cast<scf::IfOp>(op)) {
       auto condition = compileIndexExpr(ifOp.getCondition());
       if (failed(condition)) {
-        ifOp.emitOpError("requires statically evaluable scf.if condition for PIM codegen");
+        ifOp.emitOpError("requires statically evaluable scf.if condition for Pim codegen");
         return failure();
       }
       CompiledCoreNode node;
@@ -157,7 +82,7 @@ static LogicalResult compileCoreEmissionPlan(Block& block, SmallVectorImpl<Compi
     if (auto switchOp = dyn_cast<scf::IndexSwitchOp>(op)) {
       auto selector = compileIndexExpr(switchOp.getArg());
       if (failed(selector)) {
-        switchOp.emitOpError("requires a statically evaluable scf.index_switch selector for PIM codegen");
+        switchOp.emitOpError("requires a statically evaluable scf.index_switch selector for Pim codegen");
         return failure();
       }
       CompiledCoreNode node;
@@ -188,11 +113,6 @@ static LogicalResult compileCoreEmissionPlan(Block& block, SmallVectorImpl<Compi
     CompiledCoreNode node;
     node.op = &op;
     node.opKind = *opKind;
-    if (*opKind == CompiledCoreOpKind::Transpose) {
-      auto transposePlan = compileTransposePlan(cast<pim::PimTransposeOp>(op));
-      if (failed(transposePlan)) return failure();
-      node.transposePlan = *transposePlan;
-    }
     plan.push_back(std::move(node));
   }
   return success();

src/PIM/Compiler/PimCoreProgram.hpp

@@ -6,34 +6,21 @@
 #include "llvm/ADT/SmallVector.h"
 
 #include <memory>
-#include <optional>
-
 #include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp"
 
 namespace onnx_mlir {
 
-struct CompiledTransposePlan {
-  mlir::Value source;
-  mlir::Value destination;
-  size_t elementBytes = 0;
-  size_t totalElements = 0;
-  size_t totalBytes = 0;
-  llvm::SmallVector<size_t> sourceShape;
-  llvm::SmallVector<size_t> destinationStrides;
-  llvm::SmallVector<unsigned> destinationDimensionForSource;
-  llvm::SmallVector<size_t> destinationRewinds;
-  bool storagePreserving = false;
-};
-
 enum class CompiledCoreOpKind : uint8_t {
   Load,
   Store,
   Lmv,
+  VMV,
   Receive,
   Send,
+  Wait,
+  Sync,
   Concat,
   Vmm,
-  Transpose,
   VVAdd,
   VVSub,
   VVMul,
@@ -62,7 +49,6 @@ struct CompiledCoreNode {
   llvm::SmallVector<int64_t> caseValues;
   llvm::SmallVector<std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>>> caseBodies;
   std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> defaultBody;
-  std::optional<CompiledTransposePlan> transposePlan;
 };
 
 struct CompiledCoreProgram {

src/PIM/Conversion/ONNXToSpatial/CMakeLists.txt

@@ -5,7 +5,7 @@ add_public_tablegen_target(ONNXToSpatialIncGen)
 add_pim_library(OMONNXToSpatial
   Patterns.cpp
   CompileTime.cpp
-  ONNXToSpatialVerifier.cpp
+  Passes/Analyses/ONNXToSpatialVerifier.cpp
   Patterns/Pre.cpp
   Patterns/Post.cpp
   Patterns/Math/Conv.cpp
@@ -26,12 +26,15 @@ add_pim_library(OMONNXToSpatial
   Patterns/Tensor/Slice.cpp
   Patterns/Tensor/Split.cpp
   Patterns/Tensor/Transpose.cpp
-  ONNXToSpatialPass.cpp
-  SpatialLayoutPlanningPass.cpp
-  LowerSpatialPlansPass.cpp
+  Passes/Transforms/ONNXToSpatialPass.cpp
+  Passes/Analyses/SpatialLayoutCapabilities.cpp
+  Passes/Transforms/SpatialLayoutPlanningPass.cpp
+  Passes/Transforms/SpatialPlanLoweringPatterns.cpp
+  Passes/Transforms/LowerSpatialPlansPass.cpp
   Common/AttributeUtils.cpp
   Common/BiasAddUtils.cpp
   Common/ComputeRegionBuilder.cpp
+  Common/ContractionPlanning.cpp
   Common/MatrixProductLowering.cpp
   Common/RowStripLayoutUtils.cpp
   Common/ShapeTilingUtils.cpp
@@ -46,8 +49,6 @@ add_pim_library(OMONNXToSpatial
   MLIRLinalgDialect
   MLIRSCFDialect
   MLIRTosaDialect
-  OMCompilerOptions
-  OMPimCompilerOptions
   OMONNXOps
   SpatialOps
   OMPimCommon

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

@@ -25,6 +25,9 @@ FailureOr<Value> createFragmentAssemblyBlueprint(Value physicalBatch,
   const int64_t laneCount = physicalType.getDimSize(0);
   if (laneCount <= 0)
     return emitError(loc, "fragment assembly requires at least one physical source slot"), failure();
+  auto physicalLayoutValue = spatial::symbolizePhysicalLayout(physicalLayout);
+  if (!physicalLayoutValue)
+    return emitError(loc, "unknown physical layout for fragment assembly"), failure();
   const int64_t fragmentElements = physicalType.getNumElements() / laneCount;
   SmallVector<int64_t> operandIndices(entries.size(), 0), sourceSlots, sourceOffsets, offsets, sizes,
     strides(entries.size() * rank, 1);
@@ -47,13 +50,18 @@ FailureOr<Value> createFragmentAssemblyBlueprint(Value physicalBatch,
     llvm::append_range(offsets, entry.destinationOffsets);
     llvm::append_range(sizes, entry.sizes);
   }
-  return spatial::SpatBlueprintOp::create(rewriter, loc, logicalType, physicalBatch, ValueRange {},
-    rewriter.getStringAttr("nchw"), rewriter.getStringAttr(physicalLayout),
+  auto blueprint = spatial::SpatBlueprintOp::create(rewriter, loc, logicalType, physicalBatch, ValueRange {},
+    spatial::getNCHWLayout(rewriter.getContext()),
+    spatial::PhysicalLayoutAttr::get(rewriter.getContext(), *physicalLayoutValue),
     rewriter.getDenseI64ArrayAttr(offsets), rewriter.getDenseI64ArrayAttr(sizes),
-    rewriter.getStringAttr(indexMap), rewriter.getStringAttr("fragment_assembly"),
+    rewriter.getStringAttr(indexMap), spatial::getFragmentAssemblyMode(rewriter.getContext()),
     rewriter.getDenseI64ArrayAttr(operandIndices), rewriter.getDenseI64ArrayAttr(sourceSlots),
     rewriter.getDenseI64ArrayAttr(sourceOffsets), rewriter.getDenseI64ArrayAttr(strides),
-    rewriter.getStringAttr("disjoint"), rewriter.getStringAttr("complete")).getOutput();
+    rewriter.getStringAttr("disjoint"), rewriter.getStringAttr("complete"));
+  if (indexMap == spatial::kContiguousRowMajorFragments
+      && !spatial::isCanonicalContiguousRowMajorFragmentAssembly(blueprint))
+    blueprint.setIndexMapAttr(rewriter.getStringAttr("fragment_assembly"));
+  return blueprint.getOutput();
 }
 
 Value sumTensors(ArrayRef<Value> tensors, PatternRewriter& rewriter) {

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

@@ -249,7 +249,7 @@ auto createEmptySpatGraphComputeBatch(RewriterT& rewriter,
   if (laneCount <= 0 || laneCount > std::numeric_limits<int32_t>::max())
     return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());
 
-  auto laneCountAttr = pim::getCheckedI32Attr(rewriter, loc, laneCount, "spatial compute_batch lane count");
+  auto laneCountAttr = pim::getCheckedI32Attr(rewriter, loc, laneCount, "Spatial compute_batch lane count");
   if (mlir::failed(laneCountAttr))
     return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());
 
@@ -394,6 +394,39 @@ extractGraphBatchPhysicalFragment(mlir::PatternRewriter& rewriter,
     rewriter, loc, physicalBatch, fragmentType, {offsets, sizes, strides});
 }
 
+template <typename BodyFn>
+mlir::FailureOr<mlir::Value> mapGraphBatchFragments(mlir::Value input,
+                                                    mlir::RankedTensorType outputType,
+                                                    mlir::PatternRewriter& rewriter,
+                                                    mlir::Location loc,
+                                                    BodyFn&& build) {
+  auto inputType = mlir::dyn_cast<mlir::RankedTensorType>(input.getType());
+  if (!inputType || !inputType.hasStaticShape() || !outputType.hasStaticShape()
+      || inputType.getRank() != outputType.getRank() || inputType.getRank() < 2
+      || inputType.getDimSize(0) != outputType.getDimSize(0))
+    return mlir::failure();
+  auto inputFragmentType = mlir::RankedTensorType::get(
+    inputType.getShape().drop_front(), inputType.getElementType(), inputType.getEncoding());
+  auto outputFragmentType = mlir::RankedTensorType::get(
+    outputType.getShape().drop_front(), outputType.getElementType(), outputType.getEncoding());
+  auto batch = createSpatComputeBatch(
+    rewriter, loc, mlir::TypeRange {outputType}, inputType.getDimSize(0), {}, mlir::ValueRange {input},
+    [&](detail::SpatComputeBatchBodyArgs args) -> mlir::LogicalResult {
+      auto fragment = extractGraphBatchPhysicalFragment(
+        rewriter, loc, args.inputs.front(), args.lane, inputFragmentType);
+      if (mlir::failed(fragment))
+        return mlir::failure();
+      mlir::FailureOr<mlir::Value> result = build(*fragment, outputFragmentType);
+      if (mlir::failed(result) || result->getType() != outputFragmentType)
+        return mlir::failure();
+      publishGraphBatchPhysicalFragment(rewriter, loc, *result, args.outputs.front(), args.lane);
+      return mlir::success();
+    });
+  if (mlir::failed(batch))
+    return mlir::failure();
+  return batch->getResult(0);
+}
+
 template <typename BodyFn>
 mlir::Value materializeOrComputeUnary(mlir::Value input,
                                       mlir::RankedTensorType resultType,

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

@@ -0,0 +1,42 @@
+#include "ContractionPlanning.hpp"
+
+#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+
+#include <algorithm>
+
+namespace onnx_mlir {
+
+namespace {
+
+static int64_t ceilDivide(int64_t value, int64_t divisor) {
+  return divisor == 0 ? 0 : (value + divisor - 1) / divisor;
+}
+
+} // namespace
+
+ContractionPlan makeContractionPlan(
+    const ContractionProblem& problem,
+    const spatial::SpatialTargetResources& target,
+    ContractionPlanKind kind,
+    int64_t laneCount,
+    int64_t fragmentRows) {
+  ContractionPlan plan;
+  plan.tileK = std::max<int64_t>(1, target.matrixShape.rows);
+  plan.tileN = std::max<int64_t>(1, target.matrixShape.columns);
+  plan.reductionSlices = std::max<int64_t>(1, ceilDivide(problem.k, plan.tileK));
+  plan.outputTiles = std::max<int64_t>(1, ceilDivide(problem.n, plan.tileN));
+  const int64_t rowsPerLane = std::max<int64_t>(
+      1, fragmentRows != 0 ? fragmentRows : target.matrixShape.rows);
+
+  if (laneCount != 0)
+    plan.laneCount = laneCount;
+  else if (kind == ContractionPlanKind::StaticTiled)
+    plan.laneCount = problem.batch * problem.m * plan.reductionSlices * plan.outputTiles;
+  else if (kind == ContractionPlanKind::GroupedRowDynamicVVD)
+    plan.laneCount = problem.batch * ceilDivide(problem.m, rowsPerLane);
+  else
+    plan.laneCount = problem.batch * problem.m * problem.n;
+  return plan;
+}
+
+} // namespace onnx_mlir

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

@@ -0,0 +1,30 @@
+#pragma once
+
+#include "ContractionProblem.hpp"
+
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
+
+namespace onnx_mlir {
+
+enum class ContractionPlanKind {
+  StaticTiled,
+  BatchedDynamicVVD,
+  GroupedRowDynamicVVD,
+};
+
+struct ContractionPlan {
+  int64_t tileK = 1;
+  int64_t tileN = 1;
+  int64_t reductionSlices = 1;
+  int64_t outputTiles = 1;
+  int64_t laneCount = 0;
+};
+
+ContractionPlan makeContractionPlan(
+    const ContractionProblem& problem,
+    const spatial::SpatialTargetResources& target,
+    ContractionPlanKind kind,
+    int64_t laneCount = 0,
+    int64_t fragmentRows = 0);
+
+} // namespace onnx_mlir

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

@@ -0,0 +1,26 @@
+#pragma once
+
+#include "mlir/IR/BuiltinTypes.h"
+
+#include "llvm/ADT/SmallVector.h"
+
+#include <cstdint>
+
+namespace onnx_mlir {
+
+struct ContractionProblem {
+  llvm::SmallVector<int64_t> lhsBatchShape;
+  llvm::SmallVector<int64_t> rhsBatchShape;
+  llvm::SmallVector<int64_t> outputBatchShape;
+  int64_t lhsBatch = 1;
+  int64_t rhsBatch = 1;
+  int64_t batch = 1;
+  int64_t m = 0;
+  int64_t k = 0;
+  int64_t n = 0;
+  mlir::Type lhsElementType;
+  mlir::Type rhsElementType;
+  mlir::Type resultElementType;
+};
+
+} // namespace onnx_mlir

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

@@ -1,15 +1,70 @@
 #include "MatrixProductLowering.hpp"
 
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
 
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp"
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 using namespace mlir;
 
 namespace onnx_mlir {
 
+static bool isInsideSpatialCompute(Operation* op) {
+  for (Operation* parent = op; parent; parent = parent->getParentOp())
+    if (spatial::isAnySpatialComputeLike(parent))
+      return true;
+  return false;
+}
+
+static Value buildLinalgTranspose(Value value,
+                                  RankedTensorType resultType,
+                                  ArrayRef<int64_t> permutation,
+                                  PatternRewriter& rewriter,
+                                  Location loc) {
+  Value init = tensor::EmptyOp::create(
+      rewriter, loc, resultType.getShape(), resultType.getElementType());
+  return linalg::TransposeOp::create(
+    rewriter, loc, value, init, permutation).getResult()[0];
+}
+
+static Value materializeConstantTranspose(Value value,
+                                          RankedTensorType resultType,
+                                          ArrayRef<int64_t> permutation,
+                                          PatternRewriter& rewriter) {
+  auto denseAttr = getHostConstDenseElementsAttr(value);
+  if (!denseAttr)
+    return {};
+  auto transposedAttr = transposeDenseElementsAttr(denseAttr, permutation);
+  if (failed(transposedAttr) || transposedAttr->getType() != resultType)
+    return {};
+  return getOrCreateConstant(
+      rewriter, rewriter.getInsertionBlock()->getParentOp(), *transposedAttr, resultType);
+}
+
+Value createLinalgTranspose(Value value,
+                            RankedTensorType resultType,
+                            ArrayRef<int64_t> permutation,
+                            PatternRewriter& rewriter,
+                            Location loc) {
+  if (Value constant = materializeConstantTranspose(value, resultType, permutation, rewriter))
+    return constant;
+
+  if (isInsideSpatialCompute(rewriter.getInsertionBlock()->getParentOp()))
+    return buildLinalgTranspose(value, resultType, permutation, rewriter, loc);
+
+  auto compute = createSpatCompute<1>(
+    rewriter, loc, TypeRange {resultType}, {}, ValueRange {value},
+    [&](Value input) {
+      spatial::SpatYieldOp::create(
+        rewriter, loc, buildLinalgTranspose(input, resultType, permutation, rewriter, loc));
+    });
+  return compute.getResult(0);
+}
+
 Value createZeroPaddedTensor(Value value, RankedTensorType resultType, PatternRewriter& rewriter, Location loc) {
   auto sourceType = cast<RankedTensorType>(value.getType());
   SmallVector<OpFoldResult> lowPads(sourceType.getRank(), rewriter.getIndexAttr(0));

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

@@ -5,8 +5,16 @@
 #include "mlir/IR/Value.h"
 #include "mlir/Transforms/DialectConversion.h"
 
+#include "llvm/ADT/ArrayRef.h"
+
 namespace onnx_mlir {
 
+mlir::Value createLinalgTranspose(mlir::Value value,
+                                  mlir::RankedTensorType resultType,
+                                  llvm::ArrayRef<int64_t> permutation,
+                                  mlir::PatternRewriter& rewriter,
+                                  mlir::Location loc);
+
 mlir::Value createZeroPaddedTensor(mlir::Value value,
                                    mlir::RankedTensorType resultType,
                                    mlir::PatternRewriter& rewriter,

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

@@ -5,9 +5,9 @@
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/MatrixProductLowering.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-#include "src/Dialect/ONNX/ONNXOps.hpp"
 
 #include <numeric>
 
@@ -33,6 +33,16 @@ FailureOr<RowStripPhysicalValue> describeRowStripPhysicalValue(Value storage, Ra
     tilesPerRow};
 }
 
+FailureOr<RowStripPhysicalValue> getRowStripPhysicalValue(Value value) {
+  auto blueprint = value.getDefiningOp<spatial::SpatBlueprintOp>();
+  auto logicalType = dyn_cast<RankedTensorType>(value.getType());
+  if (!blueprint || !logicalType || blueprint.getOutput() != value
+      || blueprint.getPhysicalLayout() != spatial::PhysicalLayout::NHWCRowStrip
+      || !spatial::isPhysicalView(blueprint.getMode()))
+    return failure();
+  return describeRowStripPhysicalValue(blueprint.getInput(), logicalType);
+}
+
 RankedTensorType getRowStripFragmentType(RankedTensorType logicalType) {
   return RankedTensorType::get({logicalType.getDimSize(0), 1, logicalType.getDimSize(3),
                                 logicalType.getDimSize(1)},
@@ -144,6 +154,35 @@ FailureOr<Value> createRowStripStorageFromRows(Value rows,
   return batchOp->getResult(0);
 }
 
+FailureOr<Value> createRowStripStorageBlueprint(Value storage,
+                                                RankedTensorType logicalType,
+                                                PatternRewriter& rewriter,
+                                                Location loc) {
+  FailureOr<RowStripPhysicalValue> value = describeRowStripPhysicalValue(storage, logicalType);
+  if (failed(value))
+    return failure();
+
+  auto blueprint = spatial::SpatBlueprintOp::create(
+      rewriter,
+      loc,
+      logicalType,
+      storage,
+      ValueRange {},
+      spatial::getNCHWLayout(rewriter.getContext()),
+      spatial::getNHWCRowStripLayout(rewriter.getContext()),
+      rewriter.getDenseI64ArrayAttr({}),
+      rewriter.getDenseI64ArrayAttr({}),
+      rewriter.getStringAttr(kRowStripIndexMap),
+      spatial::getPhysicalViewMode(rewriter.getContext()),
+      nullptr,
+      nullptr,
+      nullptr,
+      nullptr,
+      nullptr,
+      nullptr);
+  return blueprint.getOutput();
+}
+
 FailureOr<Value> createRowStripAssemblyBlueprint(const RowStripPhysicalValue& value,
                                                  PatternRewriter& rewriter,
                                                  Location loc) {
@@ -160,8 +199,8 @@ FailureOr<Value> createRowStripAssemblyBlueprint(const RowStripPhysicalValue& va
         rewriter, loc, args.inputs.front(), args.lane, value.fragmentType);
       if (failed(fragment))
         return failure();
-      Value nchw = ONNXTransposeOp::create(
-        rewriter, loc, nchwFragmentType, *fragment, rewriter.getI64ArrayAttr({0, 3, 1, 2}));
+      Value nchw = createLinalgTranspose(
+        *fragment, nchwFragmentType, {0, 3, 1, 2}, rewriter, loc);
       publishGraphBatchPhysicalFragment(rewriter, loc, nchw, args.outputs.front(), args.lane);
       return success();
     });
@@ -176,32 +215,32 @@ FailureOr<Value> createRowStripAssemblyBlueprint(const RowStripPhysicalValue& va
         {1, std::min(tileChannels, value.logicalType.getDimSize(1) - channelOffset), 1,
          value.logicalType.getDimSize(3)}});
     }
-  return createFragmentAssemblyBlueprint(transposed->getResult(0), value.logicalType, entries, "nhwc_row_strip",
+  return createFragmentAssemblyBlueprint(transposed->getResult(0), value.logicalType, entries, "dense_nchw",
                                          kRowStripIndexMap, rewriter, loc);
 }
 
-FailureOr<Value> applyRowStripRelu(const RowStripPhysicalValue& value, PatternRewriter& rewriter, Location loc) {
+template <typename BuildActivation>
+static FailureOr<Value> applyRowStripActivation(const RowStripPhysicalValue& value,
+                                                PatternRewriter& rewriter,
+                                                Location loc,
+                                                BuildActivation buildActivation) {
   auto storageType = cast<RankedTensorType>(value.storage.getType());
-  const int64_t laneCount = storageType.getDimSize(0);
-  auto batchOp = createSpatComputeBatch(rewriter,
-                                        loc,
-                                        TypeRange {storageType},
-                                        laneCount,
-                                        {},
-                                        ValueRange {value.storage},
-                                        [&](detail::SpatComputeBatchBodyArgs args) {
-                                          FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(
-                                            rewriter, loc, args.inputs.front(), args.lane, value.fragmentType);
-                                          if (failed(fragment)) return failure();
-                                          Value relu = spatial::SpatReluOp::create(
-                                            rewriter, loc, value.fragmentType, *fragment).getResult();
-                                          publishGraphBatchPhysicalFragment(
-                                            rewriter, loc, relu, args.outputs.front(), args.lane);
-                                          return success();
-                                        });
-  if (failed(batchOp))
-    return failure();
-  return batchOp->getResult(0);
+  return mapGraphBatchFragments(value.storage, storageType, rewriter, loc, [&](Value fragment, RankedTensorType) {
+    return FailureOr<Value>(buildActivation(fragment));
+  });
+}
+
+FailureOr<Value> applyRowStripRelu(const RowStripPhysicalValue& value, PatternRewriter& rewriter, Location loc) {
+  return applyRowStripActivation(value, rewriter, loc, [&](Value fragment) {
+    return spatial::SpatReluOp::create(rewriter, loc, value.fragmentType, fragment).getResult();
+  });
+}
+
+FailureOr<Value> applyRowStripSilu(const RowStripPhysicalValue& value, PatternRewriter& rewriter, Location loc) {
+  return applyRowStripActivation(value, rewriter, loc, [&](Value fragment) {
+    Value sigmoid = spatial::SpatSigmoidOp::create(rewriter, loc, value.fragmentType, fragment).getResult();
+    return spatial::SpatVMulOp::create(rewriter, loc, value.fragmentType, fragment, sigmoid).getResult();
+  });
 }
 
 FailureOr<Value> applyRowStripBiasAdd(const RowStripPhysicalValue& value,

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

@@ -6,6 +6,12 @@
 
 namespace onnx_mlir {
 
+namespace spatial {
+class SpatBlueprintOp;
+class SpatFlattenPlanOp;
+struct SpatialTargetResources;
+} // namespace spatial
+
 inline constexpr llvm::StringLiteral kRowStripIndexMap = "nhwc_row_strip_fragments";
 
 struct RowStripPhysicalValue {
@@ -18,6 +24,8 @@ struct RowStripPhysicalValue {
 mlir::FailureOr<RowStripPhysicalValue> describeRowStripPhysicalValue(mlir::Value storage,
                                                                      mlir::RankedTensorType logicalType);
 
+mlir::FailureOr<RowStripPhysicalValue> getRowStripPhysicalValue(mlir::Value value);
+
 std::pair<llvm::SmallVector<int64_t>, llvm::SmallVector<int64_t>>
 buildRowStripMetadata(mlir::RankedTensorType type);
 
@@ -53,6 +61,11 @@ mlir::FailureOr<mlir::Value> createRowStripStorageFromRows(mlir::Value rows,
                                                            mlir::PatternRewriter& rewriter,
                                                            mlir::Location loc);
 
+mlir::FailureOr<mlir::Value> createRowStripStorageBlueprint(mlir::Value storage,
+                                                            mlir::RankedTensorType logicalType,
+                                                            mlir::PatternRewriter& rewriter,
+                                                            mlir::Location loc);
+
 mlir::FailureOr<mlir::Value> createRowStripAssemblyBlueprint(const RowStripPhysicalValue& value,
                                                                mlir::PatternRewriter& rewriter,
                                                                mlir::Location loc);
@@ -61,6 +74,10 @@ mlir::FailureOr<mlir::Value> applyRowStripRelu(const RowStripPhysicalValue& valu
                                                mlir::PatternRewriter& rewriter,
                                                mlir::Location loc);
 
+mlir::FailureOr<mlir::Value> applyRowStripSilu(const RowStripPhysicalValue& value,
+                                               mlir::PatternRewriter& rewriter,
+                                               mlir::Location loc);
+
 mlir::FailureOr<mlir::Value> applyRowStripBiasAdd(const RowStripPhysicalValue& value,
                                                   mlir::Value bias,
                                                   mlir::PatternRewriter& rewriter,
@@ -76,4 +93,14 @@ mlir::FailureOr<mlir::Value> applyRowStripConcat(llvm::ArrayRef<RowStripPhysical
                                                  mlir::PatternRewriter& rewriter,
                                                  mlir::Location loc);
 
+mlir::LogicalResult canLowerFlattenFromRowStrip(
+    spatial::SpatFlattenPlanOp flattenOp,
+    const spatial::SpatialTargetResources& target);
+
+mlir::LogicalResult lowerFlattenFromRowStrip(
+    const RowStripPhysicalValue& input,
+    spatial::SpatFlattenPlanOp flattenOp,
+    const spatial::SpatialTargetResources& target,
+    mlir::PatternRewriter& rewriter);
+
 } // namespace onnx_mlir

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

@@ -5,7 +5,6 @@
 
 #include "ShapeTilingUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
-#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 
 using namespace mlir;
@@ -67,11 +66,15 @@ sliceVector(const Value& vectorToSlice, int64_t sliceSize, PatternRewriter& rewr
 }
 
 DenseMap<CoreId, SmallVector<Value>>
-sliceVectorPerCrossbarPerCore(const Value& vectorToSlice, PatternRewriter& rewriter, Location loc) {
-  SmallVector<Value> slices = sliceVector(vectorToSlice, crossbarSize, rewriter, loc);
+sliceVectorPerCrossbarPerCore(const Value& vectorToSlice,
+                              PatternRewriter& rewriter,
+                              Location loc,
+                              const spatial::SpatialTargetResources& target) {
+  SmallVector<Value> slices = sliceVector(
+    vectorToSlice, static_cast<int64_t>(target.matrixShape.rows), rewriter, loc);
   DenseMap<CoreId, SmallVector<Value>> slicesPerCore;
   for (size_t sliceId = 0; sliceId < slices.size(); sliceId++) {
-    size_t coreId = sliceId / crossbarCountInCore;
+    size_t coreId = sliceId / target.matrixUnitsPerProcessor;
     slicesPerCore[coreId].push_back(slices[sliceId]);
   }
   return slicesPerCore;

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

@@ -7,6 +7,7 @@
 #include "llvm/ADT/SmallVector.h"
 
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
 
 namespace onnx_mlir {
 
@@ -24,8 +25,11 @@ llvm::SmallVector<mlir::Value> sliceVector(const mlir::Value& vectorToSlice,
                                            mlir::Location loc);
 
 /// Partitions one logical vector into per-core crossbar-sized slices using the
-/// current PIM target geometry.
+/// current Pim target geometry.
 llvm::DenseMap<CoreId, llvm::SmallVector<mlir::Value>> sliceVectorPerCrossbarPerCore(
-  const mlir::Value& vectorToSlice, mlir::PatternRewriter& rewriter, mlir::Location loc);
+  const mlir::Value& vectorToSlice,
+  mlir::PatternRewriter& rewriter,
+  mlir::Location loc,
+  const spatial::SpatialTargetResources& target);
 
 } // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp

@@ -1,688 +0,0 @@
-#include "mlir/Dialect/Affine/IR/AffineOps.h"
-#include "mlir/Dialect/Arith/IR/Arith.h"
-#include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/Dialect/Linalg/IR/Linalg.h"
-#include "mlir/Dialect/SCF/IR/SCF.h"
-#include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "mlir/Pass/Pass.h"
-#include "mlir/Transforms/DialectConversion.h"
-
-#include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/SmallPtrSet.h"
-
-#include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
-#include "mlir/Transforms/Passes.h"
-#include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp"
-#include "src/Accelerators/PIM/Pass/PIMPasses.h"
-#include "src/Dialect/ONNX/ONNXOps.hpp"
-
-using namespace mlir;
-
-namespace onnx_mlir {
-namespace {
-
-static constexpr StringLiteral kDenseLayout = "dense_nchw";
-static constexpr StringLiteral kRowStripLayout = "nhwc_row_strip";
-
-static FailureOr<RowStripPhysicalValue> getRowStripValue(llvm::DenseMap<Value, RowStripPhysicalValue>& rowStripValues,
-                                                         Value value) {
-  auto it = rowStripValues.find(value);
-  if (it == rowStripValues.end())
-    return failure();
-  return it->second;
-}
-
-static FailureOr<RowStripPhysicalValue> buildRowStripValue(spatial::SpatBlueprintOp blueprint,
-                                                           Value storage) {
-  auto logicalType = dyn_cast<RankedTensorType>(blueprint.getOutput().getType());
-  if (!logicalType)
-    return blueprint.emitOpError("requires ranked logical output type"), failure();
-  if (blueprint.getIndexMap() != kRowStripIndexMap)
-    return blueprint.emitOpError("requires the canonical row-strip index map"), failure();
-  FailureOr<RowStripPhysicalValue> value = describeRowStripPhysicalValue(storage, logicalType);
-  if (failed(value))
-    return blueprint.emitOpError("requires physical row-strip fragment storage"), failure();
-  return *value;
-}
-
-static FailureOr<Value>
-lowerRowStripRelu(const RowStripPhysicalValue& input, spatial::SpatReluPlanOp planOp, PatternRewriter& rewriter) {
-  return applyRowStripRelu(input, rewriter, planOp.getLoc());
-}
-
-static FailureOr<Value> lowerRowStripBiasAdd(const RowStripPhysicalValue& input,
-                                             spatial::SpatBiasAddPlanOp planOp,
-                                             PatternRewriter& rewriter) {
-  return applyRowStripBiasAdd(input, planOp.getBias(), rewriter, planOp.getLoc());
-}
-
-static FailureOr<Value> lowerRowStripAdd(const RowStripPhysicalValue& lhs,
-                                         const RowStripPhysicalValue& rhs,
-                                         spatial::SpatAddPlanOp planOp,
-                                         PatternRewriter& rewriter) {
-  return applyRowStripAdd(lhs, rhs, rewriter, planOp.getLoc());
-}
-
-static FailureOr<Value> lowerRowStripConcat(ArrayRef<RowStripPhysicalValue> inputs,
-                                            spatial::SpatConcatPlanOp planOp,
-                                            PatternRewriter& rewriter) {
-  auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
-  if (!outputType)
-    return failure();
-  return applyRowStripConcat(inputs, outputType, rewriter, planOp.getLoc());
-}
-
-static FailureOr<Value>
-materializeRowStripToDense(const RowStripPhysicalValue& rowStripValue, Location loc, PatternRewriter& rewriter) {
-  if (rowStripValue.logicalType.getRank() != 4 || !rowStripValue.logicalType.hasStaticShape())
-    return failure();
-  return createRowStripAssemblyBlueprint(rowStripValue, rewriter, loc);
-}
-
-static FailureOr<Value> lowerDenseBatchBiasAdd(Value input, Value bias, RankedTensorType resultType,
-                                               PatternRewriter& rewriter, Location loc) {
-  auto producer = input.getDefiningOp<spatial::SpatGraphComputeBatch>();
-  auto inputType = dyn_cast<RankedTensorType>(input.getType());
-  auto biasType = dyn_cast<RankedTensorType>(bias.getType());
-  if (!producer || !inputType || !biasType || !inputType.hasStaticShape() || !biasType.hasStaticShape()
-      || !resultType.hasStaticShape() || inputType.getDimSize(0) != producer.getLaneCount()
-      || biasType.getDimSize(0) != producer.getLaneCount() || resultType.getDimSize(0) != producer.getLaneCount())
-    return failure();
-  auto inputFragmentType = spatial::getGraphBatchFragmentType(inputType, producer.getLaneCount());
-  auto outputFragmentType = spatial::getGraphBatchFragmentType(resultType, producer.getLaneCount());
-  if (failed(inputFragmentType) || failed(outputFragmentType) || inputFragmentType->getRank() != biasType.getRank()
-      || inputFragmentType->getDimSize(0) != 1 || inputFragmentType->getShape().drop_front() != biasType.getShape().drop_front()
-      || inputFragmentType->getRank() != outputFragmentType->getRank() + 1)
-    return failure();
-  for (auto [inputDim, outputDim] : llvm::zip(inputFragmentType->getShape().drop_front(), outputFragmentType->getShape()))
-    if (outputDim > inputDim)
-      return failure();
-
-  auto batch = createSpatComputeBatch(rewriter, loc, TypeRange {resultType}, producer.getLaneCount(), {}, ValueRange {input, bias},
-    [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult {
-      FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[0], args.lane, *inputFragmentType);
-      if (failed(fragment))
-        return failure();
-      MixedSliceGeometry biasSlice;
-      for (int64_t dim : inputFragmentType->getShape()) {
-        biasSlice.offsets.push_back(biasSlice.offsets.empty() ? OpFoldResult(args.lane) : rewriter.getIndexAttr(0));
-        biasSlice.sizes.push_back(rewriter.getIndexAttr(dim));
-        biasSlice.strides.push_back(rewriter.getIndexAttr(1));
-      }
-      Value biasFragment = extractMixedSliceOrIdentity(rewriter, loc, args.inputs[1], *inputFragmentType, biasSlice);
-      if (!biasFragment)
-        return failure();
-      Value added = spatial::SpatVAddOp::create(rewriter, loc, *inputFragmentType, *fragment, biasFragment);
-      MixedSliceGeometry outputSlice;
-      outputSlice.offsets.assign(inputFragmentType->getRank(), rewriter.getIndexAttr(0));
-      outputSlice.sizes.push_back(rewriter.getIndexAttr(1));
-      outputSlice.strides.assign(inputFragmentType->getRank(), rewriter.getIndexAttr(1));
-      for (int64_t dim : outputFragmentType->getShape())
-        outputSlice.sizes.push_back(rewriter.getIndexAttr(dim));
-      Value output = extractMixedSliceOrIdentity(rewriter, loc, added, *outputFragmentType, outputSlice);
-      if (!output)
-        return failure();
-      publishGraphBatchPhysicalFragment(rewriter, loc, output, args.outputs.front(), args.lane);
-      return success();
-    });
-  if (failed(batch))
-    return failure();
-  return batch->getResult(0);
-}
-
-static LogicalResult lowerAddPlan(spatial::SpatAddPlanOp planOp,
-                                  llvm::DenseMap<Value, RowStripPhysicalValue>& rowStripValues,
-                                  llvm::SmallPtrSetImpl<Operation*>& eraseAfterLowering,
-                                  PatternRewriter& rewriter) {
-  FailureOr<RowStripPhysicalValue> lhs = getRowStripValue(rowStripValues, planOp.getLhs());
-  FailureOr<RowStripPhysicalValue> rhs = getRowStripValue(rowStripValues, planOp.getRhs());
-  if (succeeded(lhs) && succeeded(rhs)) {
-    auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
-      auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
-      return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
-    });
-    if (outputBlueprint == planOp.getResult().getUsers().end())
-      return planOp.emitOpError("row-strip add plan requires a row-strip blueprint result");
-
-    rewriter.setInsertionPoint(planOp);
-    FailureOr<Value> lowered = lowerRowStripAdd(*lhs, *rhs, planOp, rewriter);
-    if (failed(lowered))
-      return planOp.emitOpError("failed to lower selected row-strip Spatial add plan");
-    auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
-    FailureOr<RowStripPhysicalValue> output = buildRowStripValue(blueprint, *lowered);
-    if (failed(output))
-      return failure();
-    rowStripValues[blueprint.getResult()] = *output;
-    eraseAfterLowering.insert(planOp);
-    eraseAfterLowering.insert(blueprint);
-    return success();
-  }
-
-  rewriter.setInsertionPoint(planOp);
-  auto compute = createSpatCompute<2>(rewriter,
-                                      planOp.getLoc(),
-                                      planOp.getOutput().getType(),
-                                      {},
-                                      ValueRange {planOp.getLhs(), planOp.getRhs()},
-                                      [&](Value lhsValue, Value rhsValue) {
-                                        Value added = spatial::SpatVAddOp::create(
-                                          rewriter, planOp.getLoc(), planOp.getOutput().getType(), lhsValue, rhsValue);
-                                        spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), added);
-                                      });
-  rewriter.replaceOp(planOp, compute.getResults());
-  return success();
-}
-
-static LogicalResult lowerConcatPlan(spatial::SpatConcatPlanOp planOp,
-                                     llvm::DenseMap<Value, RowStripPhysicalValue>& rowStripValues,
-                                     llvm::SmallPtrSetImpl<Operation*>& eraseAfterLowering,
-                                     PatternRewriter& rewriter) {
-  SmallVector<RowStripPhysicalValue> inputs;
-  for (Value input : planOp.getInputs()) {
-    FailureOr<RowStripPhysicalValue> physical = getRowStripValue(rowStripValues, input);
-    if (failed(physical)) {
-      inputs.clear();
-      break;
-    }
-    inputs.push_back(*physical);
-  }
-  if (inputs.size() == planOp.getInputs().size()) {
-    auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
-      auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
-      return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
-    });
-    if (outputBlueprint == planOp.getResult().getUsers().end())
-      return planOp.emitOpError("row-strip concat plan requires a row-strip blueprint result");
-    rewriter.setInsertionPoint(planOp);
-    FailureOr<Value> lowered = lowerRowStripConcat(inputs, planOp, rewriter);
-    if (failed(lowered))
-      return planOp.emitOpError("failed to lower selected row-strip Spatial concat plan");
-    auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
-    FailureOr<RowStripPhysicalValue> output = buildRowStripValue(blueprint, *lowered);
-    if (failed(output))
-      return failure();
-    rowStripValues[blueprint.getResult()] = *output;
-    eraseAfterLowering.insert(planOp);
-    eraseAfterLowering.insert(blueprint);
-    return success();
-  }
-
-  rewriter.setInsertionPoint(planOp);
-  auto compute = createSpatCompute(
-    rewriter,
-    planOp.getLoc(),
-    TypeRange {planOp.getOutput().getType()},
-    {},
-    planOp.getInputs(),
-    [&](ValueRange values) {
-      Value concatenated = spatial::SpatConcatOp::create(
-        rewriter, planOp.getLoc(), planOp.getOutput().getType(), rewriter.getI64IntegerAttr(planOp.getAxis()), values);
-      spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), concatenated);
-    });
-  rewriter.replaceOp(planOp, compute.getResults());
-  return success();
-}
-
-struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, OperationPass<ModuleOp>> {
-  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerSpatialPlansPass)
-
-  StringRef getArgument() const override { return "lower-spatial-plans"; }
-  StringRef getDescription() const override { return "Lower selected Spatial planning ops to low-level Spatial IR."; }
-
-  void runOnOperation() override {
-    ModuleOp moduleOp = getOperation();
-    MLIRContext* ctx = moduleOp.getContext();
-    auto entryFunc = getPimEntryFunc(moduleOp);
-    if (failed(entryFunc)) {
-      moduleOp.emitError("failed to locate the PIM entry function during LowerSpatialPlans");
-      signalPassFailure();
-      return;
-    }
-    func::FuncOp funcOp = *entryFunc;
-    PatternRewriter rewriter(ctx);
-    llvm::DenseMap<Value, RowStripPhysicalValue> rowStripValues;
-    llvm::SmallPtrSet<Operation*, 16> eraseAfterLowering;
-    auto verifyLogicalPhase = [&](StringRef stage) -> bool {
-      if (succeeded(verifyLogicalSpatialGraphInvariants(*entryFunc)))
-        return true;
-      moduleOp.emitError() << "logical Spatial graph verification failed " << stage;
-      signalPassFailure();
-      return false;
-    };
-
-    if (!verifyLogicalPhase("at the start of LowerSpatialPlans"))
-      return;
-    for (Operation& op : llvm::make_early_inc_range(funcOp.getBody().front())) {
-      if (auto planOp = dyn_cast<spatial::SpatConv2DPlanOp>(&op)) {
-        FailureOr<RowStripPhysicalValue> rowStripInput = getRowStripValue(rowStripValues, planOp.getInput());
-        auto rowStripBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
-          auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
-          return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
-        });
-        if (rowStripBlueprint != planOp.getResult().getUsers().end()) {
-          rewriter.setInsertionPoint(planOp);
-          std::optional<Value> physicalInput;
-          if (succeeded(rowStripInput))
-            physicalInput = rowStripInput->storage;
-          FailureOr<Value> lowered = lowerSelectedConv2DPlan(
-            planOp,
-            physicalInput,
-            /*emitRowStripLayout=*/true,
-            rewriter);
-          if (failed(lowered)) {
-            auto diagnostic = planOp.emitOpError("failed to lower selected row-strip Spatial Conv plan with input ");
-            diagnostic << planOp.getInput().getType() << " and output " << planOp.getResult().getType();
-            if (physicalInput)
-              diagnostic << " from physical storage " << physicalInput->getType();
-            signalPassFailure();
-            return;
-          }
-          auto blueprint = cast<spatial::SpatBlueprintOp>(*rowStripBlueprint);
-          FailureOr<RowStripPhysicalValue> rowStripValue = buildRowStripValue(blueprint, *lowered);
-          if (failed(rowStripValue)) {
-            signalPassFailure();
-            return;
-          }
-          rowStripValues[blueprint.getResult()] = *rowStripValue;
-          eraseAfterLowering.insert(planOp);
-          eraseAfterLowering.insert(blueprint);
-          continue;
-        }
-        rewriter.setInsertionPoint(planOp);
-        FailureOr<Value> lowered =
-          lowerSelectedConv2DPlan(planOp, std::nullopt, /*emitRowStripLayout=*/false, rewriter);
-        if (failed(lowered)) {
-          planOp.emitOpError("failed to lower selected Spatial Conv plan");
-          signalPassFailure();
-          return;
-        }
-        rewriter.replaceOp(planOp, *lowered);
-        continue;
-      }
-
-      if (auto planOp = dyn_cast<spatial::SpatReluPlanOp>(&op)) {
-        if (succeeded(getRowStripValue(rowStripValues, planOp.getInput()))) {
-          auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
-            auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
-            return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
-          });
-          if (outputBlueprint == planOp.getResult().getUsers().end()) {
-            planOp.emitOpError("row-strip Relu plan requires a row-strip blueprint result");
-            signalPassFailure();
-            return;
-          }
-
-          FailureOr<RowStripPhysicalValue> input = getRowStripValue(rowStripValues, planOp.getInput());
-          rewriter.setInsertionPoint(planOp);
-          FailureOr<Value> lowered = lowerRowStripRelu(*input, planOp, rewriter);
-          if (failed(lowered)) {
-            planOp.emitOpError("failed to lower selected row-strip Spatial Relu plan");
-            signalPassFailure();
-            return;
-          }
-          auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
-          FailureOr<RowStripPhysicalValue> output = buildRowStripValue(blueprint, *lowered);
-          if (failed(output)) {
-            signalPassFailure();
-            return;
-          }
-          rowStripValues[blueprint.getResult()] = *output;
-          eraseAfterLowering.insert(planOp);
-          eraseAfterLowering.insert(blueprint);
-          continue;
-        }
-
-        rewriter.setInsertionPoint(planOp);
-        auto computeOp = createSpatCompute<1>(
-          rewriter, planOp.getLoc(), planOp.getOutput().getType(), {}, planOp.getInput(), [&](Value x) {
-            auto relu = spatial::SpatReluOp::create(rewriter, planOp.getLoc(), planOp.getOutput().getType(), x);
-            spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), relu.getResult());
-          });
-        rewriter.replaceOp(planOp, computeOp.getResults());
-        continue;
-      }
-      if (auto planOp = dyn_cast<spatial::SpatMaxPool2DPlanOp>(&op)) {
-        auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
-          auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
-          return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
-        });
-        if (outputBlueprint == planOp.getResult().getUsers().end()) {
-          planOp.emitOpError("selected MaxPool plan requires a row-strip blueprint result");
-          signalPassFailure();
-          return;
-        }
-
-        FailureOr<RowStripPhysicalValue> input = getRowStripValue(rowStripValues, planOp.getInput());
-        rewriter.setInsertionPoint(planOp);
-        std::optional<Value> physicalInput;
-        if (succeeded(input))
-          physicalInput = input->storage;
-        FailureOr<Value> lowered = lowerSelectedMaxPool2DPlan(
-          planOp, physicalInput, rewriter);
-        if (failed(lowered)) {
-          planOp.emitOpError("failed to lower selected row-strip Spatial MaxPool plan");
-          signalPassFailure();
-          return;
-        }
-        auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
-        FailureOr<RowStripPhysicalValue> output = buildRowStripValue(blueprint, *lowered);
-        if (failed(output)) {
-          signalPassFailure();
-          return;
-        }
-        rowStripValues[blueprint.getResult()] = *output;
-        eraseAfterLowering.insert(planOp);
-        eraseAfterLowering.insert(blueprint);
-        continue;
-      }
-      if (auto planOp = dyn_cast<spatial::SpatGlobalAveragePoolPlanOp>(&op)) {
-        auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
-          auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
-          return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
-        });
-        if (outputBlueprint == planOp.getResult().getUsers().end()) {
-          planOp.emitOpError("selected global AveragePool plan requires a row-strip blueprint result");
-          signalPassFailure();
-          return;
-        }
-
-        FailureOr<RowStripPhysicalValue> input = getRowStripValue(rowStripValues, planOp.getInput());
-        rewriter.setInsertionPoint(planOp);
-        std::optional<Value> physicalInput;
-        if (succeeded(input))
-          physicalInput = input->storage;
-        FailureOr<Value> lowered =
-          lowerSelectedGlobalAveragePoolPlan(planOp, physicalInput, rewriter);
-        if (failed(lowered)) {
-          planOp.emitOpError("failed to lower selected row-strip Spatial global AveragePool plan");
-          signalPassFailure();
-          return;
-        }
-        auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
-        FailureOr<RowStripPhysicalValue> output = buildRowStripValue(blueprint, *lowered);
-        if (failed(output)) {
-          signalPassFailure();
-          return;
-        }
-        rowStripValues[blueprint.getResult()] = *output;
-        eraseAfterLowering.insert(planOp);
-        eraseAfterLowering.insert(blueprint);
-        continue;
-      }
-      if (auto planOp = dyn_cast<spatial::SpatBiasAddPlanOp>(&op)) {
-        if (succeeded(getRowStripValue(rowStripValues, planOp.getInput()))) {
-          auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
-            auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
-            return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
-          });
-          if (outputBlueprint == planOp.getResult().getUsers().end()) {
-            planOp.emitOpError("row-strip bias_add plan requires a row-strip blueprint result");
-            signalPassFailure();
-            return;
-          }
-
-          FailureOr<RowStripPhysicalValue> input = getRowStripValue(rowStripValues, planOp.getInput());
-          rewriter.setInsertionPoint(planOp);
-          FailureOr<Value> lowered = lowerRowStripBiasAdd(*input, planOp, rewriter);
-          if (failed(lowered)) {
-            planOp.emitOpError("failed to lower selected row-strip Spatial bias_add plan");
-            signalPassFailure();
-            return;
-          }
-          auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
-          FailureOr<RowStripPhysicalValue> output = buildRowStripValue(blueprint, *lowered);
-          if (failed(output)) {
-            signalPassFailure();
-            return;
-          }
-          rowStripValues[blueprint.getResult()] = *output;
-          eraseAfterLowering.insert(planOp);
-          eraseAfterLowering.insert(blueprint);
-          continue;
-        }
-
-        auto resultType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
-        if (!resultType) {
-          planOp.emitOpError("requires ranked output type");
-          signalPassFailure();
-          return;
-        }
-        rewriter.setInsertionPoint(planOp);
-        FailureOr<Value> denseBias = materializeDenseBiasAddTensor(planOp.getBias(), resultType, rewriter, planOp.getLoc());
-        if (failed(denseBias)) {
-          planOp.emitOpError("failed to materialize dense Conv-style bias");
-          signalPassFailure();
-          return;
-        }
-        if (planOp.getInput().getDefiningOp<spatial::SpatGraphComputeBatch>()) {
-          FailureOr<Value> lowered = lowerDenseBatchBiasAdd(planOp.getInput(), *denseBias, resultType, rewriter, planOp.getLoc());
-          if (succeeded(lowered)) {
-            rewriter.replaceOp(planOp, *lowered);
-            continue;
-          }
-        }
-        auto computeOp = createSpatCompute<2>(rewriter,
-                                              planOp.getLoc(),
-                                              planOp.getOutput().getType(),
-                                              {},
-                                              ValueRange {planOp.getInput(), *denseBias},
-                                              [&](Value x, Value y) {
-                                                auto added = spatial::SpatVAddOp::create(
-                                                  rewriter, planOp.getLoc(), planOp.getOutput().getType(), x, y);
-                                                spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), added.getResult());
-                                              });
-        rewriter.replaceOp(planOp, computeOp.getResults());
-        continue;
-      }
-      if (auto planOp = dyn_cast<spatial::SpatAddPlanOp>(&op)) {
-        if (failed(lowerAddPlan(planOp, rowStripValues, eraseAfterLowering, rewriter))) {
-          signalPassFailure();
-          return;
-        }
-        continue;
-      }
-      if (auto planOp = dyn_cast<spatial::SpatConcatPlanOp>(&op)) {
-        if (failed(lowerConcatPlan(planOp, rowStripValues, eraseAfterLowering, rewriter))) {
-          signalPassFailure();
-          return;
-        }
-        continue;
-      }
-      if (auto flattenOp = dyn_cast<spatial::SpatGraphCompute>(&op)) {
-        if (flattenOp.getInputs().size() == 1) {
-          FailureOr<RowStripPhysicalValue> input =
-            getRowStripValue(rowStripValues, flattenOp.getInputs().front());
-          if (succeeded(input) && succeeded(canLowerFlattenFromRowStrip(flattenOp))) {
-            rewriter.setInsertionPoint(flattenOp);
-            if (failed(lowerFlattenFromRowStrip(*input, flattenOp, rewriter))) {
-              flattenOp.emitOpError("failed to preserve row-strip layout through Flatten");
-              signalPassFailure();
-              return;
-            }
-            continue;
-          }
-        }
-      }
-      if (auto materializeOp = dyn_cast<spatial::SpatMaterializeLayoutOp>(&op)) {
-        if (materializeOp.getSourcePhysicalLayout() == kDenseLayout
-            && materializeOp.getTargetPhysicalLayout() == kDenseLayout) {
-          rewriter.replaceOp(materializeOp, materializeOp.getInput());
-          continue;
-        }
-        if (materializeOp.getSourcePhysicalLayout() != kRowStripLayout
-            || materializeOp.getTargetPhysicalLayout() != kDenseLayout) {
-          materializeOp.emitOpError("non-dense materialize_layout lowering is not supported yet");
-          signalPassFailure();
-          return;
-        }
-        FailureOr<RowStripPhysicalValue> rowStripValue = getRowStripValue(rowStripValues, materializeOp.getInput());
-        if (failed(rowStripValue)) {
-          materializeOp.emitOpError("expected a row-strip blueprint input during row-strip materialization");
-          signalPassFailure();
-          return;
-        }
-        rewriter.setInsertionPoint(materializeOp);
-        FailureOr<Value> dense = materializeRowStripToDense(*rowStripValue, materializeOp.getLoc(), rewriter);
-        if (failed(dense)) {
-          materializeOp.emitOpError("failed to materialize selected row-strip layout back to dense NCHW");
-          signalPassFailure();
-          return;
-        }
-        rewriter.replaceOp(materializeOp, *dense);
-        continue;
-      }
-      if (auto blueprintOp = dyn_cast<spatial::SpatBlueprintOp>(&op)) {
-        if (std::optional<StringRef> mode = blueprintOp.getMode(); mode && *mode == "fragment_assembly")
-          continue;
-        if (blueprintOp.getPhysicalLayout() == kDenseLayout) {
-          rewriter.replaceOp(blueprintOp, blueprintOp.getInput());
-          continue;
-        }
-        if (blueprintOp.getPhysicalLayout() != kRowStripLayout) {
-          blueprintOp.emitOpError("non-dense blueprint lowering is not supported yet");
-          signalPassFailure();
-          return;
-        }
-        if (!eraseAfterLowering.contains(blueprintOp)) {
-          blueprintOp.emitOpError("unhandled row-strip blueprint remained during LowerSpatialPlans");
-          signalPassFailure();
-          return;
-        }
-      }
-    }
-    bool erasedAny = true;
-    while (erasedAny) {
-      erasedAny = false;
-      for (Operation& op : llvm::make_early_inc_range(funcOp.getBody().front())) {
-        if (!eraseAfterLowering.contains(&op))
-          continue;
-        if (!op.use_empty())
-          continue;
-        eraseAfterLowering.erase(&op);
-        rewriter.eraseOp(&op);
-        erasedAny = true;
-      }
-    }
-    if (!eraseAfterLowering.empty()) {
-      for (Operation& op : funcOp.getBody().front())
-        if (eraseAfterLowering.contains(&op))
-          op.emitOpError("selected row-strip planning op could not be fully eliminated during LowerSpatialPlans");
-      signalPassFailure();
-      return;
-    }
-    ConversionTarget helperTarget(*ctx);
-    helperTarget.addLegalDialect<spatial::SpatialDialect,
-                                 tensor::TensorDialect,
-                                 linalg::LinalgDialect,
-                                 affine::AffineDialect,
-                                 arith::ArithDialect,
-                                 scf::SCFDialect,
-                                 func::FuncDialect>();
-    helperTarget.addLegalOp<spatial::SpatGraphCompute, spatial::SpatGraphComputeBatch>();
-    helperTarget.addIllegalOp<ONNXGemmOp, ONNXTransposeOp>();
-    helperTarget.markOpRecursivelyLegal<spatial::SpatGraphCompute, spatial::SpatGraphComputeBatch>();
-
-    RewritePatternSet helperPatterns(ctx);
-    populateGemmPatterns(helperPatterns, ctx);
-    populateTransposePatterns(helperPatterns, ctx);
-    FrozenRewritePatternSet frozenHelperPatterns(
-      std::move(helperPatterns));
-    SmallVector<Operation*> topLevelHelperOps;
-    funcOp.walk([&](Operation* op) {
-      if (isa<spatial::SpatGraphCompute,
-              spatial::SpatGraphComputeBatch>(op))
-        return WalkResult::skip();
-      if (isa<ONNXGemmOp, ONNXTransposeOp>(op))
-        topLevelHelperOps.push_back(op);
-      return WalkResult::advance();
-    });
-    for (Operation *helper : topLevelHelperOps) {
-      if (failed(applyPartialConversion(
-            helper, helperTarget, frozenHelperPatterns))) {
-        moduleOp.emitError("failed to lower helper ONNX ops emitted by selected Spatial plan lowering");
-        signalPassFailure();
-        return;
-      }
-    }
-    ConversionTarget nestedHelperTarget(*ctx);
-    nestedHelperTarget.addLegalDialect<spatial::SpatialDialect,
-                                       tensor::TensorDialect,
-                                       linalg::LinalgDialect,
-                                       affine::AffineDialect,
-                                       arith::ArithDialect,
-                                       scf::SCFDialect,
-                                       func::FuncDialect>();
-    nestedHelperTarget.addIllegalOp<ONNXGemmOp, ONNXTransposeOp>();
-    SmallVector<Operation*> computeLikeOps;
-    funcOp.walk([&](Operation* op) {
-      if (isa<spatial::SpatGraphCompute, spatial::SpatGraphComputeBatch>(op))
-        computeLikeOps.push_back(op);
-    });
-    for (Operation* op : computeLikeOps) {
-      if (failed(applyFullConversion(
-            op, nestedHelperTarget, frozenHelperPatterns))) {
-        op->emitOpError("failed to lower nested helper ONNX ops emitted by selected Spatial plan lowering");
-        signalPassFailure();
-        return;
-      }
-    }
-    if (!verifyLogicalPhase("after nested helper conversions"))
-      return;
-    bool hasIllegalOps = false;
-    moduleOp.walk([&](Operation* op) {
-      if (isa<ONNXEntryPointOp>(op))
-        return;
-      if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
-        if (std::optional<StringRef> mode = blueprint.getMode(); mode && *mode == "fragment_assembly")
-          return;
-        op->emitOpError("planning blueprint must not remain after LowerSpatialPlans");
-        hasIllegalOps = true;
-      }
-      else if (isa<spatial::SpatConv2DPlanOp,
-                   spatial::SpatBiasAddPlanOp,
-                   spatial::SpatAddPlanOp,
-                   spatial::SpatReluPlanOp,
-                   spatial::SpatMaxPool2DPlanOp,
-                   spatial::SpatGlobalAveragePoolPlanOp,
-                   spatial::SpatMaterializeLayoutOp>(op)
-               || op->getDialect()->getNamespace() == "onnx") {
-        op->emitOpError("operation must not remain after LowerSpatialPlans");
-        hasIllegalOps = true;
-      }
-    });
-
-    PassManager canonicalizationPM(ctx);
-    canonicalizationPM.addPass(createCanonicalizerPass());
-    if (failed(canonicalizationPM.run(moduleOp)))
-      moduleOp.emitWarning("failed to run LowerSpatialPlansPass canonicalization; continuing");
-
-    if (hasIllegalOps) {
-      signalPassFailure();
-    } else {
-      dumpModule(moduleOp, "spatial1_graph");
-      spatial::SpatialDataflowExportStage exportMode = spatial::getSpatialDataflowExportStage();
-      if (spatial::shouldExportSpatialDataflowStage(exportMode, spatial::SpatialDataflowExportStage::Spatial1)
-          && failed(spatial::exportSpatialDataflowCsvGraph(funcOp, "spatial1_graph"))) {
-        signalPassFailure();
-        return;
-      }
-    }
-
-    if (!verifyLogicalPhase("at the end of LowerSpatialPlans"))
-      return;
-  }
-};
-
-} // namespace
-
-std::unique_ptr<Pass> createLowerSpatialPlansPass() { return std::make_unique<LowerSpatialPlansPass>(); }
-
-} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp

@@ -0,0 +1,32 @@
+#pragma once
+
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
+
+#include <cstdint>
+
+namespace onnx_mlir::spatial {
+
+enum class ConvLoweringStrategy : uint8_t {
+  Auto,
+  Legacy,
+  Depthwise,
+  PackedIm2Col,
+  StreamedPatch,
+  StreamedPacked,
+  OutputChannelTiled,
+  InputKTiled,
+  Tiled2D,
+};
+
+} // namespace onnx_mlir::spatial
+
+namespace onnx_mlir {
+
+struct ONNXToSpatialPlanningOptions {
+  uint64_t convIm2colMaxElements = 0;
+  uint64_t convStreamChunkPositions = 0;
+  spatial::ConvLoweringStrategy forcedConvStrategy = spatial::ConvLoweringStrategy::Auto;
+  bool reportConvLowering = true;
+};
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.cpp

@@ -6,7 +6,7 @@
 #include "Common/IR/WeightUtils.hpp"
 #include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 using namespace mlir;
@@ -108,7 +108,9 @@ void verifyScheduledInputs(ComputeOpTy compute,
   for (auto [inputIndex, input] : llvm::enumerate(compute.getInputs())) {
     size_t currentInputIndex = inputIndex;
     Operation* definingOp = input.getDefiningOp();
-    if (allowChannelReceiveInputs && isa_and_nonnull<spatial::SpatChannelReceiveOp>(definingOp))
+    if (allowChannelReceiveInputs
+        && isa_and_nonnull<spatial::SpatChannelReceiveOp,
+                           spatial::SpatHostWaitLoadOp>(definingOp))
       continue;
     if (isScheduledPhase1Value(input))
       continue;
@@ -130,8 +132,7 @@ template <typename ComputeOpTy>
 void verifyNoNestedFragmentAssemblyBlueprints(ComputeOpTy compute,
                                                   pim::CappedDiagnosticReporter& diagnostics) {
   compute.getBody().walk([&](spatial::SpatBlueprintOp blueprint) {
-    std::optional<StringRef> mode = blueprint.getMode();
-    if (!mode || *mode != "fragment_assembly")
+    if (!spatial::isFragmentAssembly(blueprint.getMode()))
       return;
     diagnostics.report(blueprint.getOperation(), [&](Operation* illegalOp) {
       illegalOp->emitOpError("fragment assembly blueprint must be host-level after merge materialization");
@@ -148,7 +149,10 @@ void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter
             spatial::SpatBiasAddPlanOp,
             spatial::SpatAddPlanOp,
             spatial::SpatConcatPlanOp,
+            spatial::SpatFlattenPlanOp,
             spatial::SpatReluPlanOp,
+            spatial::SpatSiluPlanOp,
+            spatial::SpatResizeNearestPlanOp,
             spatial::SpatMaxPool2DPlanOp,
             spatial::SpatGlobalAveragePoolPlanOp,
             spatial::SpatBlueprintOp,
@@ -161,7 +165,8 @@ void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter
       });
       continue;
     }
-    if (isa<spatial::SpatChannelReceiveOp, spatial::SpatChannelSendOp>(&op)) {
+    if (isa<spatial::SpatChannelReceiveOp, spatial::SpatChannelSendOp,
+            spatial::SpatHostStoreSyncOp, spatial::SpatHostWaitLoadOp>(&op)) {
       diagnostics.report(&op, [&](Operation* illegalOp) {
         illegalOp->emitOpError() << kPhaseMarker
                                  << " explicit channel communication is not expected before merge materialization";
@@ -180,7 +185,8 @@ void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter
 
 void verifyScheduledTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter& diagnostics) {
   for (Operation& op : funcOp.getOps()) {
-    if (isa<spatial::SpatChannelSendOp, spatial::SpatChannelReceiveOp>(&op)) {
+    if (isa<spatial::SpatChannelSendOp, spatial::SpatChannelReceiveOp,
+            spatial::SpatHostStoreSyncOp, spatial::SpatHostWaitLoadOp>(&op)) {
       diagnostics.report(&op, [&](Operation* illegalOp) {
         illegalOp->emitOpError() << kPhaseMarker << " real channel communication is not allowed in scheduled phase 1";
       });

src/PIM/Conversion/ONNXToSpatial/Passes/Analyses/SpatialLayoutCapabilities.cpp

@@ -0,0 +1,152 @@
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Transforms/PlanLowering.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir::spatial {
+
+static LayoutAlternative denseAlternative(Operation *op) {
+  LayoutAlternative alternative;
+  alternative.operandLayouts.assign(op->getNumOperands(), PhysicalLayout::DenseNCHW);
+  alternative.resultLayout = PhysicalLayout::DenseNCHW;
+  return alternative;
+}
+
+static LayoutAlternative rowStripAlternative(Operation *op,
+                                             ArrayRef<PhysicalLayout> operandLayouts) {
+  LayoutAlternative alternative;
+  alternative.operandLayouts.assign(operandLayouts.begin(), operandLayouts.end());
+  alternative.resultLayout = PhysicalLayout::NHWCRowStrip;
+  alternative.intrinsicCost = -2;
+  return alternative;
+}
+
+static bool hasRowStripInput(ArrayRef<PhysicalLayout> operandLayouts, unsigned index) {
+  return index < operandLayouts.size()
+      && operandLayouts[index] == PhysicalLayout::NHWCRowStrip;
+}
+
+SmallVector<LayoutAlternative> SpatConv2DPlanOp::getLayoutAlternatives(
+    const SpatialTargetResources& target, ArrayRef<PhysicalLayout> operandLayouts) {
+  SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
+  if (hasRowStripInput(operandLayouts, 0)) {
+    if (succeeded(canConsumeAndProduceRowStrip(*this, target)))
+      alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts));
+  }
+  else if (succeeded(canLowerConvPlanToRowStrip(*this, target))) {
+    LayoutAlternative alternative = denseAlternative(getOperation());
+    alternative.resultLayout = PhysicalLayout::NHWCRowStrip;
+    alternative.intrinsicCost = -2;
+    alternatives.push_back(std::move(alternative));
+  }
+  return alternatives;
+}
+
+SmallVector<LayoutAlternative> SpatFlattenPlanOp::getLayoutAlternatives(
+    const SpatialTargetResources& target, ArrayRef<PhysicalLayout> operandLayouts) {
+  SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
+  if (!operandLayouts.empty()
+      && operandLayouts[0] == PhysicalLayout::Fragmented) {
+    LayoutAlternative alternative = denseAlternative(getOperation());
+    alternative.operandLayouts[0] = PhysicalLayout::Fragmented;
+    alternatives.push_back(std::move(alternative));
+  }
+  if (hasRowStripInput(operandLayouts, 0)
+      && succeeded(canLowerFlattenFromRowStrip(*this, target))) {
+    LayoutAlternative alternative = rowStripAlternative(getOperation(), operandLayouts);
+    alternative.resultLayout = PhysicalLayout::DenseNCHW;
+    alternatives.push_back(std::move(alternative));
+  }
+  return alternatives;
+}
+
+SmallVector<LayoutAlternative> SpatReluPlanOp::getLayoutAlternatives(
+    const SpatialTargetResources&, ArrayRef<PhysicalLayout> operandLayouts) {
+  SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
+  if (hasRowStripInput(operandLayouts, 0))
+    alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts));
+  return alternatives;
+}
+
+SmallVector<LayoutAlternative> SpatSiluPlanOp::getLayoutAlternatives(
+    const SpatialTargetResources&, ArrayRef<PhysicalLayout> operandLayouts) {
+  SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
+  if (hasRowStripInput(operandLayouts, 0)) {
+    LayoutAlternative alternative = rowStripAlternative(getOperation(), operandLayouts);
+    alternative.intrinsicCost = -3;
+    alternatives.push_back(std::move(alternative));
+  }
+  return alternatives;
+}
+
+SmallVector<LayoutAlternative> SpatResizeNearestPlanOp::getLayoutAlternatives(
+    const SpatialTargetResources& target, ArrayRef<PhysicalLayout> operandLayouts) {
+  SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
+  if (hasRowStripInput(operandLayouts, 0)
+      && succeeded(canLowerResizeNearestPlanToRowStrip(*this, target)))
+    alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts));
+  return alternatives;
+}
+
+SmallVector<LayoutAlternative> SpatMaxPool2DPlanOp::getLayoutAlternatives(
+    const SpatialTargetResources& target, ArrayRef<PhysicalLayout> operandLayouts) {
+  SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
+  if (succeeded(canLowerMaxPoolPlanToRowStrip(*this, target))) {
+    LayoutAlternative alternative = denseAlternative(getOperation());
+    if (hasRowStripInput(operandLayouts, 0))
+      alternative = rowStripAlternative(getOperation(), operandLayouts);
+    alternative.resultLayout = PhysicalLayout::NHWCRowStrip;
+    alternative.intrinsicCost = -2;
+    alternatives.push_back(std::move(alternative));
+  }
+  return alternatives;
+}
+
+SmallVector<LayoutAlternative> SpatGlobalAveragePoolPlanOp::getLayoutAlternatives(
+    const SpatialTargetResources& target, ArrayRef<PhysicalLayout> operandLayouts) {
+  SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
+  if (succeeded(canLowerGlobalAveragePoolPlanToRowStrip(*this, target))) {
+    LayoutAlternative alternative = denseAlternative(getOperation());
+    if (hasRowStripInput(operandLayouts, 0))
+      alternative = rowStripAlternative(getOperation(), operandLayouts);
+    alternative.resultLayout = PhysicalLayout::NHWCRowStrip;
+    alternative.intrinsicCost = -2;
+    alternatives.push_back(std::move(alternative));
+  }
+  return alternatives;
+}
+
+SmallVector<LayoutAlternative> SpatBiasAddPlanOp::getLayoutAlternatives(
+    const SpatialTargetResources&, ArrayRef<PhysicalLayout> operandLayouts) {
+  SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
+  auto resultType = dyn_cast<RankedTensorType>(getOutput().getType());
+  if (resultType && hasRowStripInput(operandLayouts, 0)
+      && isSupportedBiasAddValue(getBias(), resultType))
+    alternatives.push_back(rowStripAlternative(getOperation(),
+                                                {PhysicalLayout::NHWCRowStrip,
+                                                 PhysicalLayout::DenseNCHW}));
+  return alternatives;
+}
+
+SmallVector<LayoutAlternative> SpatAddPlanOp::getLayoutAlternatives(
+    const SpatialTargetResources&, ArrayRef<PhysicalLayout> operandLayouts) {
+  SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
+  if (operandLayouts.size() >= 2 && hasRowStripInput(operandLayouts, 0)
+      && hasRowStripInput(operandLayouts, 1))
+    alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts));
+  return alternatives;
+}
+
+SmallVector<LayoutAlternative> SpatConcatPlanOp::getLayoutAlternatives(
+    const SpatialTargetResources&, ArrayRef<PhysicalLayout> operandLayouts) {
+  SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
+  if (!operandLayouts.empty() && llvm::all_of(operandLayouts, [](PhysicalLayout layout) {
+        return layout == PhysicalLayout::NHWCRowStrip;
+      }))
+    alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts));
+  return alternatives;
+}
+
+} // namespace onnx_mlir::spatial

src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/LowerSpatialPlansPass.cpp

@@ -0,0 +1,136 @@
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+#include "Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.hpp"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Transforms/SpatialPlanLoweringPatterns.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp"
+#include "src/Accelerators/PIM/Passes/PIMPasses.h"
+#include "src/Dialect/ONNX/ONNXOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace {
+
+struct LowerSpatialPlansPass final
+    : PassWrapper<LowerSpatialPlansPass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerSpatialPlansPass)
+
+  StringRef getArgument() const override { return "lower-spatial-plans"; }
+  StringRef getDescription() const override {
+    return "Lower selected Spatial planning ops to low-level Spatial IR.";
+  }
+
+  LowerSpatialPlansPass() = default;
+  LowerSpatialPlansPass(const spatial::SpatialTargetResources& target,
+                        const ONNXToSpatialPlanningOptions& options,
+                        spatial::SpatialDataflowExportStage exportStage)
+      : target(target), planningOptions(options), exportStage(exportStage), hasTarget(true) {}
+
+  void runOnOperation() override {
+    ModuleOp moduleOp = getOperation();
+    if (!hasTarget) {
+      moduleOp.emitError("Spatial plan lowering requires an injected SpatialTargetResources");
+      signalPassFailure();
+      return;
+    }
+    auto entryFunc = getPimEntryFunc(moduleOp);
+    if (failed(entryFunc)) {
+      moduleOp.emitError("failed to locate the Pim entry function during LowerSpatialPlans");
+      signalPassFailure();
+      return;
+    }
+    func::FuncOp funcOp = *entryFunc;
+    auto verifyLogicalPhase = [&](StringRef stage) -> bool {
+      if (succeeded(verifyLogicalSpatialGraphInvariants(funcOp)))
+        return true;
+      moduleOp.emitError() << "logical Spatial graph verification failed " << stage;
+      signalPassFailure();
+      return false;
+    };
+
+    if (!verifyLogicalPhase("at the start of LowerSpatialPlans"))
+      return;
+    if (failed(verifySelectedSpatialLayouts(funcOp, target))) {
+      moduleOp.emitError("selected Spatial layout verification failed");
+      signalPassFailure();
+      return;
+    }
+
+    MLIRContext* ctx = moduleOp.getContext();
+    RewritePatternSet patterns(ctx);
+    populateSpatialPlanLoweringPatterns(patterns, ctx, target, planningOptions);
+
+    ConversionTarget conversionTarget(*ctx);
+    conversionTarget.addLegalDialect<spatial::SpatialDialect,
+                                     tensor::TensorDialect,
+                                     linalg::LinalgDialect,
+                                     affine::AffineDialect,
+                                     arith::ArithDialect,
+                                     scf::SCFDialect,
+                                     func::FuncDialect>();
+    conversionTarget.addIllegalDialect<ONNXDialect>();
+    conversionTarget.addLegalOp<ONNXEntryPointOp>();
+    conversionTarget.addIllegalOp<spatial::SpatConv2DPlanOp,
+                                  spatial::SpatFlattenPlanOp,
+                                  spatial::SpatReluPlanOp,
+                                  spatial::SpatSiluPlanOp,
+                                  spatial::SpatResizeNearestPlanOp,
+                                  spatial::SpatMaxPool2DPlanOp,
+                                  spatial::SpatGlobalAveragePoolPlanOp,
+                                  spatial::SpatBiasAddPlanOp,
+                                  spatial::SpatAddPlanOp,
+                                  spatial::SpatConcatPlanOp,
+                                  spatial::SpatMaterializeLayoutOp>();
+    conversionTarget.addDynamicallyLegalOp<spatial::SpatBlueprintOp>(
+        [](spatial::SpatBlueprintOp blueprint) {
+          return spatial::isFragmentAssembly(blueprint.getMode());
+        });
+
+    if (failed(applyFullConversion(funcOp, conversionTarget,
+                                   std::move(patterns)))) {
+      moduleOp.emitError("failed to lower Spatial plans and layout materialization");
+      signalPassFailure();
+      return;
+    }
+
+    dumpModule(moduleOp, "spatial1_graph");
+    if (spatial::shouldExportSpatialDataflowStage(
+            exportStage, spatial::SpatialDataflowExportStage::Spatial1)
+        && failed(spatial::exportSpatialDataflowCsvGraph(funcOp, "spatial1_graph"))) {
+      signalPassFailure();
+      return;
+    }
+    verifyLogicalPhase("at the end of LowerSpatialPlans");
+  }
+
+  spatial::SpatialTargetResources target;
+  ONNXToSpatialPlanningOptions planningOptions;
+  spatial::SpatialDataflowExportStage exportStage = spatial::SpatialDataflowExportStage::None;
+  bool hasTarget = false;
+};
+
+} // namespace
+
+std::unique_ptr<Pass> createLowerSpatialPlansPass() {
+  return std::make_unique<LowerSpatialPlansPass>();
+}
+
+std::unique_ptr<Pass> createLowerSpatialPlansPass(
+    const spatial::SpatialTargetResources& target,
+    const ONNXToSpatialPlanningOptions& options,
+    spatial::SpatialDataflowExportStage exportStage) {
+  return std::make_unique<LowerSpatialPlansPass>(target, options, exportStage);
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/ONNXToSpatialPass.cpp

@@ -8,18 +8,19 @@
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/Passes.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 #include "llvm/ADT/SmallVector.h"
 
-#include "Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
-#include "ONNXToSpatialVerifier.hpp"
 
 using namespace mlir;
 
@@ -33,9 +34,17 @@ struct ONNXToSpatialPass : PassWrapper<ONNXToSpatialPass, OperationPass<ModuleOp
   StringRef getDescription() const override { return "Lower ONNX ops to Spatial ops."; }
 
   ONNXToSpatialPass() = default;
-  ONNXToSpatialPass(const ONNXToSpatialPass& pass) {}
+  ONNXToSpatialPass(const spatial::SpatialTargetResources& target,
+                    const ONNXToSpatialPlanningOptions& options)
+      : target(target), planningOptions(options), hasTarget(true) {}
+  ONNXToSpatialPass(const ONNXToSpatialPass& pass)
+      : target(pass.target), planningOptions(pass.planningOptions), hasTarget(pass.hasTarget) {}
 
   void runOnOperation() override;
+
+  spatial::SpatialTargetResources target;
+  ONNXToSpatialPlanningOptions planningOptions;
+  bool hasTarget = false;
 };
 
 } // namespace
@@ -49,14 +58,19 @@ static void populateEmptyFunction(func::FuncOp funcOp) {
   SmallVector<spatial::SpatBiasAddPlanOp> biasAddPlans(funcOp.getOps<spatial::SpatBiasAddPlanOp>());
   SmallVector<spatial::SpatAddPlanOp> addPlans(funcOp.getOps<spatial::SpatAddPlanOp>());
   SmallVector<spatial::SpatConcatPlanOp> concatPlans(funcOp.getOps<spatial::SpatConcatPlanOp>());
+  SmallVector<spatial::SpatFlattenPlanOp> flattenPlans(funcOp.getOps<spatial::SpatFlattenPlanOp>());
   SmallVector<spatial::SpatReluPlanOp> reluPlans(funcOp.getOps<spatial::SpatReluPlanOp>());
+  SmallVector<spatial::SpatSiluPlanOp> siluPlans(funcOp.getOps<spatial::SpatSiluPlanOp>());
+  SmallVector<spatial::SpatResizeNearestPlanOp> resizePlans(
+    funcOp.getOps<spatial::SpatResizeNearestPlanOp>());
   SmallVector<spatial::SpatMaxPool2DPlanOp> maxPoolPlans(funcOp.getOps<spatial::SpatMaxPool2DPlanOp>());
   SmallVector<spatial::SpatGlobalAveragePoolPlanOp> globalAveragePoolPlans(
     funcOp.getOps<spatial::SpatGlobalAveragePoolPlanOp>());
   SmallVector<spatial::SpatBlueprintOp> blueprints(funcOp.getOps<spatial::SpatBlueprintOp>());
   SmallVector<spatial::SpatMaterializeLayoutOp> materializers(funcOp.getOps<spatial::SpatMaterializeLayoutOp>());
   if (!computes.empty() || !computeBatches.empty() || !convPlans.empty() || !biasAddPlans.empty() || !addPlans.empty()
-      || !concatPlans.empty() || !reluPlans.empty() || !maxPoolPlans.empty() || !blueprints.empty()
+      || !concatPlans.empty() || !flattenPlans.empty() || !reluPlans.empty() || !siluPlans.empty() || !resizePlans.empty()
+      || !maxPoolPlans.empty() || !blueprints.empty()
       || !globalAveragePoolPlans.empty() || !materializers.empty()) {
     return;
   }
@@ -101,6 +115,11 @@ static void populateEmptyFunction(func::FuncOp funcOp) {
 
 void ONNXToSpatialPass::runOnOperation() {
   ModuleOp moduleOp = getOperation();
+  if (!hasTarget) {
+    moduleOp.emitError("ONNX-to-Spatial lowering requires an injected SpatialTargetResources");
+    signalPassFailure();
+    return;
+  }
   MLIRContext* ctx = &getContext();
 
   ConversionTarget preTarget(*ctx);
@@ -121,9 +140,25 @@ void ONNXToSpatialPass::runOnOperation() {
     return;
   }
 
+  RewritePatternSet matmulPatterns(ctx);
+  populateMatMulFusionPatterns(matmulPatterns, ctx, target);
+  if (failed(applyPatternsGreedily(moduleOp, std::move(matmulPatterns)))) {
+    moduleOp.emitError("failed to lower MatMul before producer conversion");
+    signalPassFailure();
+    return;
+  }
+
+  RewritePatternSet fusionPatterns(ctx);
+  populateElementwiseFusionPatterns(fusionPatterns, ctx);
+  if (failed(applyPatternsGreedily(moduleOp, std::move(fusionPatterns)))) {
+    moduleOp.emitError("failed to fuse layout-aware ONNX elementwise patterns");
+    signalPassFailure();
+    return;
+  }
+
   auto entryFunc = getPimEntryFunc(moduleOp);
   if (failed(entryFunc)) {
-    moduleOp.emitError("failed to locate the PIM entry function during ONNX-to-Spatial lowering");
+    moduleOp.emitError("failed to locate the Pim entry function during ONNX-to-Spatial lowering");
     signalPassFailure();
     return;
   }
@@ -161,7 +196,7 @@ void ONNXToSpatialPass::runOnOperation() {
   target.addIllegalOp<ONNXSplitOp>();
 
   RewritePatternSet conversionPatterns(ctx);
-  populateConversionPatterns(conversionPatterns, ctx);
+  populateConversionPatterns(conversionPatterns, ctx, this->target, planningOptions);
   if (failed(applyPartialConversion(moduleOp, target, std::move(conversionPatterns)))) {
     moduleOp.emitError("failed to convert required ONNX ops to Spatial ops");
     signalPassFailure();
@@ -210,7 +245,7 @@ void ONNXToSpatialPass::runOnOperation() {
   RewritePatternSet postPatterns(ctx);
   populatePostPatterns(postPatterns, ctx);
   if (failed(applyPartialConversion(*entryFunc, postTarget, std::move(postPatterns)))) {
-    moduleOp.emitError("failed to normalize weight-like Spatial compute operands before Spatial-to-PIM lowering");
+    moduleOp.emitError("failed to normalize weight-like Spatial compute operands before Spatial-to-Pim lowering");
     signalPassFailure();
     return;
   }
@@ -237,4 +272,10 @@ void ONNXToSpatialPass::runOnOperation() {
 
 std::unique_ptr<Pass> createONNXToSpatialPass() { return std::make_unique<ONNXToSpatialPass>(); }
 
+std::unique_ptr<Pass> createONNXToSpatialPass(
+    const spatial::SpatialTargetResources& target,
+    const ONNXToSpatialPlanningOptions& options) {
+  return std::make_unique<ONNXToSpatialPass>(target, options);
+}
+
 } // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/PlanLowering.hpp

@@ -0,0 +1,92 @@
+#pragma once
+
+#include <optional>
+
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Support/LogicalResult.h"
+
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+namespace onnx_mlir {
+
+struct RowStripPhysicalValue;
+struct ONNXToSpatialPlanningOptions;
+
+inline spatial::PhysicalLayout getSpatialPlanOperandLayout(mlir::Value value) {
+  if (auto materialize = value.getDefiningOp<spatial::SpatMaterializeLayoutOp>())
+    return materialize.getTargetPhysicalLayout();
+  if (auto blueprint = value.getDefiningOp<spatial::SpatBlueprintOp>())
+    return blueprint.getPhysicalLayout();
+  if (mlir::Operation* producer = value.getDefiningOp())
+    if (auto selected = spatial::getSelectedPhysicalLayout(producer))
+      return *selected;
+  return spatial::PhysicalLayout::DenseNCHW;
+}
+
+mlir::FailureOr<mlir::Value>
+lowerDenseFlattenPlan(spatial::SpatFlattenPlanOp planOp,
+                      mlir::Value input,
+                      mlir::PatternRewriter& rewriter);
+
+mlir::FailureOr<mlir::Value>
+lowerSelectedConv2DPlan(spatial::SpatConv2DPlanOp planOp,
+                        mlir::Value input,
+                        mlir::Value weight,
+                        mlir::Value bias,
+                        std::optional<mlir::Value> rowStripInput,
+                        bool emitRowStripLayout,
+                        const spatial::SpatialTargetResources& target,
+                        const ONNXToSpatialPlanningOptions& options,
+                        mlir::PatternRewriter& rewriter);
+
+mlir::LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp,
+                                               const spatial::SpatialTargetResources& target,
+                                               const ONNXToSpatialPlanningOptions* options = nullptr);
+mlir::LogicalResult canConsumeAndProduceRowStrip(spatial::SpatConv2DPlanOp planOp,
+                                                 const spatial::SpatialTargetResources& target,
+                                                 const ONNXToSpatialPlanningOptions* options = nullptr);
+
+mlir::LogicalResult canLowerResizeNearestPlanToRowStrip(
+  spatial::SpatResizeNearestPlanOp planOp, const spatial::SpatialTargetResources& target);
+
+mlir::FailureOr<mlir::Value> lowerSelectedResizeNearestPlan(
+  spatial::SpatResizeNearestPlanOp planOp,
+  mlir::Value input,
+  std::optional<mlir::Value> rowStripInput,
+  const spatial::SpatialTargetResources& target,
+  mlir::PatternRewriter& rewriter);
+
+mlir::LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp,
+                                                 const spatial::SpatialTargetResources& target);
+
+mlir::FailureOr<mlir::Value>
+lowerDenseMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
+                        mlir::Value input,
+                        const spatial::SpatialTargetResources& target,
+                       mlir::PatternRewriter& rewriter);
+
+mlir::FailureOr<mlir::Value>
+lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
+                           mlir::Value input,
+                           std::optional<mlir::Value> rowStripInput,
+                           const spatial::SpatialTargetResources& target,
+                           mlir::PatternRewriter& rewriter);
+
+mlir::LogicalResult
+canLowerGlobalAveragePoolPlanToRowStrip(spatial::SpatGlobalAveragePoolPlanOp planOp,
+                                        const spatial::SpatialTargetResources& target);
+
+mlir::FailureOr<mlir::Value>
+lowerDenseGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp,
+                                mlir::Value input,
+                                const spatial::SpatialTargetResources& target,
+                                mlir::PatternRewriter& rewriter);
+
+mlir::FailureOr<mlir::Value>
+lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp,
+                                   mlir::Value input,
+                                   std::optional<mlir::Value> rowStripInput,
+                                   const spatial::SpatialTargetResources& target,
+                                   mlir::PatternRewriter& rewriter);
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/SpatialLayoutPlanningPass.cpp

@@ -0,0 +1,325 @@
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+
+#include "llvm/ADT/DenseMap.h"
+
+#include "Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.hpp"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Transforms/PlanLowering.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Accelerators/PIM/Passes/PIMPasses.h"
+
+#include <algorithm>
+#include <limits>
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace {
+
+struct SpatialLayoutSelection {
+  llvm::DenseMap<Operation*, unsigned> selectedAlternative;
+  llvm::DenseMap<Value, spatial::PhysicalLayout> resultLayouts;
+};
+
+static spatial::PhysicalLayout getKnownLayout(
+    const SpatialLayoutSelection& selection, Value value) {
+  if (auto it = selection.resultLayouts.find(value); it != selection.resultLayouts.end())
+    return it->second;
+  return getSpatialPlanOperandLayout(value);
+}
+
+static SmallVector<spatial::PhysicalLayout> getOperandLayouts(
+    Operation* op, const SpatialLayoutSelection& selection) {
+  SmallVector<spatial::PhysicalLayout> operandLayouts;
+  operandLayouts.reserve(op->getNumOperands());
+  for (Value operand : op->getOperands())
+    operandLayouts.push_back(getKnownLayout(selection, operand));
+  return operandLayouts;
+}
+
+class SpatialLayoutAnalysis {
+public:
+  SpatialLayoutAnalysis(func::FuncOp funcOp,
+                        const spatial::SpatialTargetResources& target,
+                        bool selectTrivialPlan)
+      : funcOp(funcOp), target(target), selectTrivialPlan(selectTrivialPlan) {}
+
+  FailureOr<SpatialLayoutSelection> run() {
+    SpatialLayoutSelection selection;
+    SmallVector<Operation*> planOps;
+    for (Operation& op : funcOp.getBody().front()) {
+      if (!isa<spatial::SpatialLayoutCapabilityInterface>(&op))
+        continue;
+      planOps.push_back(&op);
+      selection.resultLayouts[op.getResult(0)] = spatial::PhysicalLayout::DenseNCHW;
+      selection.selectedAlternative[&op] = 0;
+    }
+
+    if (selectTrivialPlan)
+      return selection;
+
+    const size_t maxRounds = 2 * planOps.size() + 1;
+    for (size_t round = 0; round < maxRounds; ++round) {
+      bool changed = false;
+      SmallVector<Operation*> order(planOps);
+      if (round % 2)
+        std::reverse(order.begin(), order.end());
+
+      for (Operation* op : order) {
+        FailureOr<SmallVector<spatial::LayoutAlternative>> alternatives =
+            getAlternatives(op, selection);
+        if (failed(alternatives))
+          return failure();
+
+        unsigned currentIndex = selection.selectedAlternative.lookup(op);
+        if (currentIndex >= alternatives->size())
+          currentIndex = 0;
+        if (selection.selectedAlternative.lookup(op) != currentIndex) {
+          selection.selectedAlternative[op] = currentIndex;
+          changed = true;
+        }
+        Value result = op->getResult(0);
+        if (selection.resultLayouts.lookup(result) !=
+            (*alternatives)[currentIndex].resultLayout) {
+          selection.resultLayouts[result] = (*alternatives)[currentIndex].resultLayout;
+          changed = true;
+        }
+        int64_t bestCost = alternativeCost(op, (*alternatives)[currentIndex], selection);
+        unsigned bestIndex = currentIndex;
+        for (auto [index, alternative] : llvm::enumerate(*alternatives)) {
+          int64_t cost = alternativeCost(op, alternative, selection);
+          if (cost < bestCost) {
+            bestCost = cost;
+            bestIndex = index;
+          }
+        }
+        if (bestIndex == currentIndex)
+          continue;
+        selection.selectedAlternative[op] = bestIndex;
+        selection.resultLayouts[result] = (*alternatives)[bestIndex].resultLayout;
+        changed = true;
+      }
+      if (!changed)
+        return selection;
+    }
+    funcOp.emitError("Spatial layout selection did not converge within its bounded iteration budget");
+    return failure();
+  }
+
+  FailureOr<SmallVector<spatial::LayoutAlternative>> getAlternatives(
+      Operation* op, const SpatialLayoutSelection& selection) {
+    auto capability = dyn_cast<spatial::SpatialLayoutCapabilityInterface>(op);
+    if (!capability)
+      return failure();
+    SmallVector<spatial::LayoutAlternative> alternatives =
+      capability.getLayoutAlternatives(target, getOperandLayouts(op, selection));
+    if (alternatives.empty())
+      return op->emitOpError("does not advertise a legal Spatial layout alternative"), failure();
+    for (const spatial::LayoutAlternative& alternative : alternatives) {
+      if (alternative.operandLayouts.size() != op->getNumOperands())
+        return op->emitOpError("advertises a layout alternative with the wrong operand count"), failure();
+    }
+    if (llvm::any_of(op->getResult(0).getUses(), [](OpOperand& use) {
+          return isa<func::ReturnOp>(use.getOwner());
+        })
+        && llvm::none_of(alternatives, [](const spatial::LayoutAlternative& alternative) {
+             return alternative.resultLayout == spatial::PhysicalLayout::DenseNCHW;
+           }))
+      return op->emitOpError("does not provide the required DenseNCHW function-result layout"), failure();
+    return alternatives;
+  }
+
+private:
+  int64_t alternativeCost(Operation* op,
+                          const spatial::LayoutAlternative& alternative,
+                          const SpatialLayoutSelection& selection) {
+    if (llvm::any_of(op->getResult(0).getUses(), [](OpOperand& use) {
+          return isa<func::ReturnOp>(use.getOwner());
+        })
+        && alternative.resultLayout != spatial::PhysicalLayout::DenseNCHW)
+      return std::numeric_limits<int64_t>::max() / 4;
+
+    int64_t cost = alternative.intrinsicCost;
+    SmallVector<spatial::PhysicalLayout> operandLayouts = getOperandLayouts(op, selection);
+    for (auto [actual, required] : llvm::zip(operandLayouts, alternative.operandLayouts))
+      cost += actual != required;
+
+    Value result = op->getResult(0);
+    for (OpOperand& use : result.getUses()) {
+      auto user = dyn_cast<spatial::SpatialLayoutCapabilityInterface>(use.getOwner());
+      if (!user)
+        continue;
+      SmallVector<spatial::PhysicalLayout> userOperandLayouts =
+          getOperandLayouts(use.getOwner(), selection);
+      for (auto [index, operand] : llvm::enumerate(use.getOwner()->getOperands()))
+        if (operand == result)
+          userOperandLayouts[index] = alternative.resultLayout;
+      SmallVector<spatial::LayoutAlternative> userAlternatives =
+          user.getLayoutAlternatives(target, userOperandLayouts);
+      if (llvm::none_of(userAlternatives,
+            [&](const spatial::LayoutAlternative& userAlternative) {
+              return userAlternative.operandLayouts.size()
+                       == use.getOwner()->getNumOperands()
+                  && userAlternative.operandLayouts[use.getOperandNumber()]
+                       == alternative.resultLayout;
+            }))
+        ++cost;
+    }
+    return cost;
+  }
+
+  func::FuncOp funcOp;
+  const spatial::SpatialTargetResources& target;
+  bool selectTrivialPlan;
+};
+
+static LogicalResult materializeMismatchedUses(
+    IRRewriter& rewriter, const SpatialLayoutSelection& selection,
+    Operation* op, SpatialLayoutAnalysis& analysis) {
+  Value value = op->getResult(0);
+  spatial::PhysicalLayout sourceLayout = getKnownLayout(selection, value);
+  SmallVector<std::pair<OpOperand*, spatial::PhysicalLayout>> mismatches;
+  for (OpOperand& use : value.getUses()) {
+    Operation* userOp = use.getOwner();
+    auto capability = dyn_cast<spatial::SpatialLayoutCapabilityInterface>(userOp);
+    if (!capability) {
+      if (isa<func::ReturnOp>(userOp) || sourceLayout == spatial::PhysicalLayout::DenseNCHW)
+        continue;
+      mismatches.push_back({&use, spatial::PhysicalLayout::DenseNCHW});
+      continue;
+    }
+    FailureOr<SmallVector<spatial::LayoutAlternative>> alternatives =
+        analysis.getAlternatives(userOp, selection);
+    if (failed(alternatives))
+      return failure();
+    unsigned selectedIndex = selection.selectedAlternative.lookup(userOp);
+    if (selectedIndex >= alternatives->size())
+      return userOp->emitOpError()
+                 << "has no selected Spatial layout alternative (index " << selectedIndex
+                                   << ", alternatives " << alternatives->size() << ")",
+             failure();
+    spatial::PhysicalLayout required =
+        (*alternatives)[selectedIndex].operandLayouts[use.getOperandNumber()];
+    if (required != sourceLayout)
+      mismatches.push_back({&use, required});
+  }
+
+  for (auto [use, required] : mismatches) {
+    Operation* userOp = use->getOwner();
+    rewriter.setInsertionPoint(userOp);
+    auto materialized = spatial::SpatMaterializeLayoutOp::create(
+        rewriter, userOp->getLoc(), use->get().getType(), use->get(),
+        spatial::LogicalLayoutAttr::get(
+            rewriter.getContext(), spatial::LogicalLayout::NCHW),
+      spatial::PhysicalLayoutAttr::get(rewriter.getContext(), sourceLayout),
+      spatial::PhysicalLayoutAttr::get(rewriter.getContext(), required));
+    use->set(materialized.getResult());
+  }
+  return success();
+}
+
+static LogicalResult verifySelectedLayouts(
+    const SpatialLayoutSelection& selection,
+    ArrayRef<Operation*> planOps,
+    SpatialLayoutAnalysis& analysis) {
+  for (Operation* op : planOps) {
+    auto selected = spatial::getSelectedPhysicalLayout(op);
+    if (!selected)
+      return op->emitOpError("requires a selected physical layout"), failure();
+    FailureOr<SmallVector<spatial::LayoutAlternative>> alternatives =
+        analysis.getAlternatives(op, selection);
+    if (failed(alternatives))
+      return failure();
+    unsigned selectedIndex = selection.selectedAlternative.lookup(op);
+    if (selectedIndex >= alternatives->size())
+      return op->emitOpError()
+                 << "has no selected Spatial layout alternative (index " << selectedIndex
+                               << ", alternatives " << alternatives->size() << ")",
+             failure();
+    const spatial::LayoutAlternative& alternative = (*alternatives)[selectedIndex];
+    if (alternative.resultLayout != *selected
+        || getOperandLayouts(op, selection) != alternative.operandLayouts)
+      return op->emitOpError("selected physical layout does not satisfy its exact layout contract"), failure();
+  }
+  return success();
+}
+
+struct SpatialLayoutPlanningPass final
+    : PassWrapper<SpatialLayoutPlanningPass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SpatialLayoutPlanningPass)
+
+  StringRef getArgument() const override { return "spatial-layout-planning"; }
+  StringRef getDescription() const override {
+    return "Select Spatial layout alternatives and insert explicit reconciliation barriers.";
+  }
+
+  SpatialLayoutPlanningPass() = default;
+  SpatialLayoutPlanningPass(const spatial::SpatialTargetResources& target,
+                            bool selectTrivialPlan)
+  : target(target), selectTrivialPlan(selectTrivialPlan), hasTarget(true) {}
+
+  void runOnOperation() override {
+    ModuleOp moduleOp = getOperation();
+    if (!hasTarget) {
+      moduleOp.emitError("Spatial layout planning requires an injected SpatialTargetResources");
+      signalPassFailure();
+      return;
+    }
+    auto entryFunc = getPimEntryFunc(moduleOp);
+    if (failed(entryFunc)) {
+      moduleOp.emitError("failed to locate the Pim entry function during Spatial layout planning");
+      signalPassFailure();
+      return;
+    }
+
+    func::FuncOp funcOp = *entryFunc;
+    SpatialLayoutAnalysis analysis(funcOp, target, selectTrivialPlan);
+    FailureOr<SpatialLayoutSelection> selection = analysis.run();
+    if (failed(selection)) {
+      signalPassFailure();
+      return;
+    }
+
+    SmallVector<Operation*> planOps;
+    for (Operation& op : funcOp.getBody().front())
+      if (isa<spatial::SpatialLayoutCapabilityInterface>(&op))
+        planOps.push_back(&op);
+
+    IRRewriter rewriter(&getContext());
+    for (Operation* op : planOps) {
+      op->setAttr(spatial::kSelectedLayoutAttrName,
+                  spatial::PhysicalLayoutAttr::get(
+              rewriter.getContext(), selection->resultLayouts.lookup(op->getResult(0))));
+      if (failed(materializeMismatchedUses(
+              rewriter, *selection, op, analysis))) {
+        signalPassFailure();
+        return;
+      }
+    }
+    if (failed(verifySelectedLayouts(*selection, planOps, analysis))
+        || failed(verifyLogicalSpatialGraphInvariants(*entryFunc))) {
+      moduleOp.emitError("Spatial layout planning verification failed");
+      signalPassFailure();
+      return;
+    }
+  }
+
+  spatial::SpatialTargetResources target;
+  bool selectTrivialPlan = false;
+  bool hasTarget = false;
+};
+
+} // namespace
+
+std::unique_ptr<Pass> createSpatialLayoutPlanningPass() {
+  return std::make_unique<SpatialLayoutPlanningPass>();
+}
+
+std::unique_ptr<Pass> createSpatialLayoutPlanningPass(
+    const spatial::SpatialTargetResources& target, bool selectTrivialPlan) {
+  return std::make_unique<SpatialLayoutPlanningPass>(target, selectTrivialPlan);
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/SpatialPlanLoweringPatterns.cpp

@@ -0,0 +1,845 @@
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+#include "Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.hpp"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/MatrixProductLowering.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Transforms/PlanLowering.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp"
+#include "src/Accelerators/PIM/Passes/PIMPasses.h"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace {
+
+static FailureOr<RowStripPhysicalValue> getRowStripValue(Value value) {
+  return getRowStripPhysicalValue(value);
+}
+
+static FailureOr<Value> publishRowStripValue(Operation* planOp,
+                                             Value storage,
+                                             PatternRewriter& rewriter) {
+  auto logicalType = dyn_cast<RankedTensorType>(planOp->getResult(0).getType());
+  if (!logicalType)
+    return planOp->emitOpError("requires ranked logical output type"), failure();
+  FailureOr<RowStripPhysicalValue> value = describeRowStripPhysicalValue(storage, logicalType);
+  if (failed(value))
+    return planOp->emitOpError("lowering produced invalid row-strip physical storage"), failure();
+  FailureOr<Value> blueprint = createRowStripStorageBlueprint(
+      storage, logicalType, rewriter, planOp->getLoc());
+  if (failed(blueprint))
+    return planOp->emitOpError("failed to create row-strip storage Blueprint"), failure();
+  rewriter.replaceOp(planOp, *blueprint);
+  return *blueprint;
+}
+
+static bool isRowStripSelected(Operation* op) {
+  auto selected = spatial::getSelectedPhysicalLayout(op);
+  return selected && *selected == spatial::PhysicalLayout::NHWCRowStrip;
+}
+
+static bool isDenseSelected(Operation* op) {
+  auto selected = spatial::getSelectedPhysicalLayout(op);
+  return selected && *selected == spatial::PhysicalLayout::DenseNCHW;
+}
+
+static spatial::PhysicalLayout getKnownPhysicalLayout(Value value) {
+  return getSpatialPlanOperandLayout(value);
+}
+
+static LogicalResult verifySelectedLayouts(
+    func::FuncOp funcOp, const spatial::SpatialTargetResources& target) {
+  LogicalResult result = success();
+  funcOp.walk([&](Operation* op) {
+    auto capability = dyn_cast<spatial::SpatialLayoutCapabilityInterface>(op);
+    if (!capability)
+      return;
+    auto selected = spatial::getSelectedPhysicalLayout(op);
+    if (!selected) {
+      op->emitOpError("requires a selected physical layout from SpatialLayoutPlanning");
+      result = failure();
+      return;
+    }
+    if (*selected != spatial::PhysicalLayout::DenseNCHW
+        && *selected != spatial::PhysicalLayout::NHWCRowStrip) {
+      op->emitOpError("has an unsupported selected physical layout");
+      result = failure();
+      return;
+    }
+    SmallVector<spatial::PhysicalLayout> operandLayouts;
+    operandLayouts.reserve(op->getNumOperands());
+    for (Value operand : op->getOperands())
+      operandLayouts.push_back(getKnownPhysicalLayout(operand));
+    auto alternatives = capability.getLayoutAlternatives(target, operandLayouts);
+    if (llvm::none_of(alternatives, [&](const spatial::LayoutAlternative& alternative) {
+          return alternative.resultLayout == *selected
+              && alternative.operandLayouts == operandLayouts;
+        })) {
+      op->emitOpError("selected physical layout is not lowerable for its explicit operand layouts");
+      result = failure();
+    }
+  });
+  return result;
+}
+
+static FailureOr<Value>
+lowerRowStripRelu(const RowStripPhysicalValue& input, spatial::SpatReluPlanOp planOp, PatternRewriter& rewriter) {
+  return applyRowStripRelu(input, rewriter, planOp.getLoc());
+}
+
+static FailureOr<Value>
+lowerRowStripSilu(const RowStripPhysicalValue& input, spatial::SpatSiluPlanOp planOp, PatternRewriter& rewriter) {
+  return applyRowStripSilu(input, rewriter, planOp.getLoc());
+}
+
+static FailureOr<Value> lowerRowStripAdd(const RowStripPhysicalValue& lhs,
+                                         const RowStripPhysicalValue& rhs,
+                                         spatial::SpatAddPlanOp planOp,
+                                         PatternRewriter& rewriter) {
+  return applyRowStripAdd(lhs, rhs, rewriter, planOp.getLoc());
+}
+
+static FailureOr<Value> lowerRowStripConcat(ArrayRef<RowStripPhysicalValue> inputs,
+                                            spatial::SpatConcatPlanOp planOp,
+                                            PatternRewriter& rewriter) {
+  auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
+  if (!outputType)
+    return failure();
+  return applyRowStripConcat(inputs, outputType, rewriter, planOp.getLoc());
+}
+
+static FailureOr<Value>
+materializeRowStripToDense(const RowStripPhysicalValue& rowStripValue, Location loc, PatternRewriter& rewriter) {
+  if (rowStripValue.logicalType.getRank() != 4 || !rowStripValue.logicalType.hasStaticShape())
+    return failure();
+  return createRowStripAssemblyBlueprint(rowStripValue, rewriter, loc);
+}
+
+static FailureOr<Value> materializeDenseToRowStrip(
+    Value input, RankedTensorType logicalType, Location loc, PatternRewriter& rewriter) {
+  if (!logicalType || !logicalType.hasStaticShape() || logicalType.getRank() != 4
+      || logicalType.getDimSize(0) != 1)
+    return failure();
+  auto nhwcType = RankedTensorType::get(
+    {1, logicalType.getDimSize(2), logicalType.getDimSize(3), logicalType.getDimSize(1)},
+    logicalType.getElementType(), logicalType.getEncoding());
+  auto rowsType = RankedTensorType::get(
+    {logicalType.getDimSize(2) * logicalType.getDimSize(3), logicalType.getDimSize(1)},
+    logicalType.getElementType(), logicalType.getEncoding());
+  auto rowsCompute = createSpatCompute<1>(
+    rewriter, loc, rowsType, {}, input, [&](Value denseInput) {
+      Value nhwc = createLinalgTranspose(
+        denseInput, nhwcType, {0, 2, 3, 1}, rewriter, loc);
+      Value rows = tensor::CollapseShapeOp::create(
+        rewriter, loc, rowsType, nhwc,
+        SmallVector<ReassociationIndices> {{0, 1, 2}, {3}});
+      spatial::SpatYieldOp::create(rewriter, loc, rows);
+    });
+  Value rows = rowsCompute->getResult(0);
+  FailureOr<Value> storage = createRowStripStorageFromRows(rows, logicalType, rewriter, loc);
+  if (failed(storage))
+    return failure();
+  return createRowStripStorageBlueprint(*storage, logicalType, rewriter, loc);
+}
+
+static FailureOr<Value> lowerDenseBatchBiasAdd(Value input, Value bias, RankedTensorType resultType,
+                                               PatternRewriter& rewriter, Location loc) {
+  auto producer = input.getDefiningOp<spatial::SpatGraphComputeBatch>();
+  auto inputType = dyn_cast<RankedTensorType>(input.getType());
+  auto biasType = dyn_cast<RankedTensorType>(bias.getType());
+  if (!producer || !inputType || !biasType || !inputType.hasStaticShape() || !biasType.hasStaticShape()
+      || !resultType.hasStaticShape() || inputType.getDimSize(0) != producer.getLaneCount()
+      || biasType.getDimSize(0) != producer.getLaneCount() || resultType.getDimSize(0) != producer.getLaneCount())
+    return failure();
+  auto inputFragmentType = spatial::getGraphBatchFragmentType(inputType, producer.getLaneCount());
+  auto outputFragmentType = spatial::getGraphBatchFragmentType(resultType, producer.getLaneCount());
+  if (failed(inputFragmentType) || failed(outputFragmentType) || inputFragmentType->getRank() != biasType.getRank()
+      || inputFragmentType->getDimSize(0) != 1 || inputFragmentType->getShape().drop_front() != biasType.getShape().drop_front()
+      || inputFragmentType->getRank() != outputFragmentType->getRank() + 1)
+    return failure();
+  for (auto [inputDim, outputDim] : llvm::zip(inputFragmentType->getShape().drop_front(), outputFragmentType->getShape()))
+    if (outputDim > inputDim)
+      return failure();
+
+  auto batch = createSpatComputeBatch(rewriter, loc, TypeRange {resultType}, producer.getLaneCount(), {}, ValueRange {input, bias},
+    [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult {
+      FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[0], args.lane, *inputFragmentType);
+      if (failed(fragment))
+        return failure();
+      MixedSliceGeometry biasSlice;
+      for (int64_t dim : inputFragmentType->getShape()) {
+        biasSlice.offsets.push_back(biasSlice.offsets.empty() ? OpFoldResult(args.lane) : rewriter.getIndexAttr(0));
+        biasSlice.sizes.push_back(rewriter.getIndexAttr(dim));
+        biasSlice.strides.push_back(rewriter.getIndexAttr(1));
+      }
+      Value biasFragment = extractMixedSliceOrIdentity(rewriter, loc, args.inputs[1], *inputFragmentType, biasSlice);
+      if (!biasFragment)
+        return failure();
+      Value added = spatial::SpatVAddOp::create(rewriter, loc, *inputFragmentType, *fragment, biasFragment);
+      MixedSliceGeometry outputSlice;
+      outputSlice.offsets.assign(inputFragmentType->getRank(), rewriter.getIndexAttr(0));
+      outputSlice.sizes.push_back(rewriter.getIndexAttr(1));
+      outputSlice.strides.assign(inputFragmentType->getRank(), rewriter.getIndexAttr(1));
+      for (int64_t dim : outputFragmentType->getShape())
+        outputSlice.sizes.push_back(rewriter.getIndexAttr(dim));
+      Value output = extractMixedSliceOrIdentity(rewriter, loc, added, *outputFragmentType, outputSlice);
+      if (!output)
+        return failure();
+      publishGraphBatchPhysicalFragment(rewriter, loc, output, args.outputs.front(), args.lane);
+      return success();
+    });
+  if (failed(batch))
+    return failure();
+  return batch->getResult(0);
+}
+
+struct LowerDenseReluPlan final : OpConversionPattern<spatial::SpatReluPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatReluPlanOp planOp,
+                                spatial::SpatReluPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    auto selected = spatial::getSelectedPhysicalLayout(planOp.getOperation());
+    if (!selected || *selected != spatial::PhysicalLayout::DenseNCHW)
+      return failure();
+
+    auto computeOp = createSpatCompute<1>(
+      rewriter, planOp.getLoc(), planOp.getOutput().getType(), {}, adaptor.getInput(), [&](Value x) {
+        auto relu = spatial::SpatReluOp::create(rewriter, planOp.getLoc(), planOp.getOutput().getType(), x);
+        spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), relu.getResult());
+      });
+    rewriter.replaceOp(planOp, computeOp.getResults());
+    return success();
+  }
+};
+
+struct LowerDenseSiluPlan final : OpConversionPattern<spatial::SpatSiluPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatSiluPlanOp planOp,
+                                spatial::SpatSiluPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    auto selected = spatial::getSelectedPhysicalLayout(planOp.getOperation());
+    if (!selected || *selected != spatial::PhysicalLayout::DenseNCHW)
+      return failure();
+
+    auto computeOp = createSpatCompute<1>(
+      rewriter, planOp.getLoc(), planOp.getOutput().getType(), {}, adaptor.getInput(), [&](Value x) {
+        Value sigmoid = spatial::SpatSigmoidOp::create(
+          rewriter, planOp.getLoc(), planOp.getOutput().getType(), x).getResult();
+        Value silu = spatial::SpatVMulOp::create(
+          rewriter, planOp.getLoc(), planOp.getOutput().getType(), x, sigmoid).getResult();
+        spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), silu);
+      });
+    rewriter.replaceOp(planOp, computeOp.getResults());
+    return success();
+  }
+};
+
+struct LowerDenseResizePlan final : OpConversionPattern<spatial::SpatResizeNearestPlanOp> {
+  explicit LowerDenseResizePlan(MLIRContext* ctx, const spatial::SpatialTargetResources& target)
+      : OpConversionPattern<spatial::SpatResizeNearestPlanOp>(ctx), target(target) {}
+
+  LogicalResult matchAndRewrite(spatial::SpatResizeNearestPlanOp planOp,
+                                spatial::SpatResizeNearestPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (!isDenseSelected(planOp.getOperation()))
+      return failure();
+    FailureOr<Value> lowered = lowerSelectedResizeNearestPlan(
+        planOp, adaptor.getInput(), std::nullopt, target, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected dense nearest Resize plan");
+    rewriter.replaceOp(planOp, *lowered);
+    return success();
+  }
+
+  const spatial::SpatialTargetResources& target;
+};
+
+struct LowerDenseBiasAddPlan final : OpConversionPattern<spatial::SpatBiasAddPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatBiasAddPlanOp planOp,
+                                spatial::SpatBiasAddPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (!isDenseSelected(planOp.getOperation()))
+      return failure();
+    auto resultType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
+    if (!resultType)
+      return planOp.emitOpError("requires ranked output type");
+
+    FailureOr<Value> denseBias = materializeDenseBiasAddTensor(
+      adaptor.getBias(), resultType, rewriter, planOp.getLoc());
+    if (failed(denseBias))
+      return planOp.emitOpError("failed to materialize dense Conv-style bias");
+    if (adaptor.getInput().getDefiningOp<spatial::SpatGraphComputeBatch>()) {
+      FailureOr<Value> lowered = lowerDenseBatchBiasAdd(
+        adaptor.getInput(), *denseBias, resultType, rewriter, planOp.getLoc());
+      if (succeeded(lowered)) {
+        rewriter.replaceOp(planOp, *lowered);
+        return success();
+      }
+    }
+    auto computeOp = createSpatCompute<2>(
+      rewriter,
+      planOp.getLoc(),
+      planOp.getOutput().getType(),
+      {},
+      ValueRange {adaptor.getInput(), *denseBias},
+      [&](Value x, Value y) {
+        auto added = spatial::SpatVAddOp::create(
+          rewriter, planOp.getLoc(), planOp.getOutput().getType(), x, y);
+        spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), added.getResult());
+      });
+    rewriter.replaceOp(planOp, computeOp.getResults());
+    return success();
+  }
+};
+
+struct LowerDenseAddPlan final : OpConversionPattern<spatial::SpatAddPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatAddPlanOp planOp,
+                                spatial::SpatAddPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (!isDenseSelected(planOp.getOperation()))
+      return failure();
+    auto compute = createSpatCompute<2>(
+      rewriter,
+      planOp.getLoc(),
+      planOp.getOutput().getType(),
+      {},
+      ValueRange {adaptor.getLhs(), adaptor.getRhs()},
+      [&](Value lhsValue, Value rhsValue) {
+        Value added = spatial::SpatVAddOp::create(
+          rewriter, planOp.getLoc(), planOp.getOutput().getType(), lhsValue, rhsValue);
+        spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), added);
+      });
+    rewriter.replaceOp(planOp, compute.getResults());
+    return success();
+  }
+};
+
+struct LowerDenseConcatPlan final : OpConversionPattern<spatial::SpatConcatPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatConcatPlanOp planOp,
+                                spatial::SpatConcatPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (!isDenseSelected(planOp.getOperation()))
+      return failure();
+    auto compute = createSpatCompute(
+      rewriter,
+      planOp.getLoc(),
+      TypeRange {planOp.getOutput().getType()},
+      {},
+      adaptor.getInputs(),
+      [&](ValueRange values) {
+        Value concatenated = spatial::SpatConcatOp::create(
+          rewriter,
+          planOp.getLoc(),
+          planOp.getOutput().getType(),
+          rewriter.getI64IntegerAttr(planOp.getAxis()),
+          values);
+        spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), concatenated);
+      });
+    rewriter.replaceOp(planOp, compute.getResults());
+    return success();
+  }
+};
+
+static LogicalResult lowerAddPlan(spatial::SpatAddPlanOp planOp,
+                                  Value lhsValue, Value rhsValue,
+                                  PatternRewriter& rewriter) {
+  FailureOr<RowStripPhysicalValue> lhs = getRowStripValue(lhsValue);
+  FailureOr<RowStripPhysicalValue> rhs = getRowStripValue(rhsValue);
+  if (isRowStripSelected(planOp.getOperation()) && failed(lhs)) {
+    if (getKnownPhysicalLayout(lhsValue) == spatial::PhysicalLayout::NHWCRowStrip)
+      return failure();
+    return planOp.emitOpError("selected row-strip Add plan requires row-strip inputs");
+  }
+  if (isRowStripSelected(planOp.getOperation()) && failed(rhs)) {
+    if (getKnownPhysicalLayout(rhsValue) == spatial::PhysicalLayout::NHWCRowStrip)
+      return failure();
+    return planOp.emitOpError("selected row-strip Add plan requires row-strip inputs");
+  }
+  if (isRowStripSelected(planOp.getOperation())) {
+    rewriter.setInsertionPoint(planOp);
+    FailureOr<Value> lowered = lowerRowStripAdd(*lhs, *rhs, planOp, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected row-strip Spatial add plan");
+    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
+      return failure();
+    return success();
+  }
+  return planOp.emitOpError("dense Add plan was not lowered by the selected-plan patterns");
+}
+
+static LogicalResult lowerConcatPlan(spatial::SpatConcatPlanOp planOp,
+                                     ValueRange inputValues,
+                                     PatternRewriter& rewriter) {
+  SmallVector<RowStripPhysicalValue> inputs;
+  for (Value input : inputValues) {
+    FailureOr<RowStripPhysicalValue> physical = getRowStripValue(input);
+    if (failed(physical)) {
+      inputs.clear();
+      break;
+    }
+    inputs.push_back(*physical);
+  }
+  if (isRowStripSelected(planOp.getOperation()) && inputs.size() != inputValues.size()) {
+    if (llvm::any_of(inputValues, [](Value input) {
+          return getKnownPhysicalLayout(input) == spatial::PhysicalLayout::NHWCRowStrip;
+        }))
+      return failure();
+    return planOp.emitOpError("selected row-strip Concat plan requires row-strip inputs");
+  }
+  if (isRowStripSelected(planOp.getOperation())) {
+    rewriter.setInsertionPoint(planOp);
+    FailureOr<Value> lowered = lowerRowStripConcat(inputs, planOp, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected row-strip Spatial concat plan");
+    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
+      return failure();
+    return success();
+  }
+  return planOp.emitOpError("dense Concat plan was not lowered by the selected-plan patterns");
+}
+
+struct LowerSelectedConvPlan final : OpConversionPattern<spatial::SpatConv2DPlanOp> {
+  explicit LowerSelectedConvPlan(MLIRContext* ctx,
+                                 const spatial::SpatialTargetResources& target,
+                                 const ONNXToSpatialPlanningOptions& options)
+      : OpConversionPattern<spatial::SpatConv2DPlanOp>(ctx), target(target), options(options) {}
+
+  LogicalResult matchAndRewrite(spatial::SpatConv2DPlanOp planOp,
+                                spatial::SpatConv2DPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (isDenseSelected(planOp.getOperation())) {
+      FailureOr<Value> lowered = lowerSelectedConv2DPlan(
+          planOp, adaptor.getInput(), adaptor.getWeight(), adaptor.getBias(),
+          std::nullopt, /*emitRowStripLayout=*/false, target, options, rewriter);
+      if (failed(lowered))
+        return planOp.emitOpError("failed to lower selected dense Spatial Conv plan");
+      rewriter.replaceOp(planOp, *lowered);
+      return success();
+    }
+    if (!isRowStripSelected(planOp.getOperation()))
+      return failure();
+
+    FailureOr<RowStripPhysicalValue> rowStripInput = getRowStripValue(adaptor.getInput());
+    if (failed(rowStripInput)
+        && getKnownPhysicalLayout(adaptor.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+      return failure();
+    std::optional<Value> physicalInput;
+    if (succeeded(rowStripInput))
+      physicalInput = rowStripInput->storage;
+    FailureOr<Value> lowered = lowerSelectedConv2DPlan(
+        planOp, adaptor.getInput(), adaptor.getWeight(), adaptor.getBias(),
+        physicalInput, /*emitRowStripLayout=*/true, target, options, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected row-strip Spatial Conv plan");
+    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
+      return failure();
+    return success();
+  }
+
+  const spatial::SpatialTargetResources& target;
+  const ONNXToSpatialPlanningOptions& options;
+};
+
+struct LowerRowStripReluPlan final : OpConversionPattern<spatial::SpatReluPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatReluPlanOp planOp,
+                                spatial::SpatReluPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (!isRowStripSelected(planOp.getOperation()))
+      return failure();
+    FailureOr<RowStripPhysicalValue> input = getRowStripValue(adaptor.getInput());
+    if (failed(input)) {
+      if (getKnownPhysicalLayout(adaptor.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+        return failure();
+      return planOp.emitOpError("selected row-strip ReLU plan requires a row-strip input");
+    }
+    FailureOr<Value> lowered = lowerRowStripRelu(*input, planOp, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected row-strip Spatial ReLU plan");
+    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
+      return failure();
+    return success();
+  }
+};
+
+struct LowerRowStripSiluPlan final : OpConversionPattern<spatial::SpatSiluPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatSiluPlanOp planOp,
+                                spatial::SpatSiluPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (!isRowStripSelected(planOp.getOperation()))
+      return failure();
+    FailureOr<RowStripPhysicalValue> input = getRowStripValue(adaptor.getInput());
+    if (failed(input)) {
+      if (getKnownPhysicalLayout(adaptor.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+        return failure();
+      return planOp.emitOpError("selected row-strip SiLU plan requires a row-strip input");
+    }
+    FailureOr<Value> lowered = lowerRowStripSilu(*input, planOp, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected row-strip Spatial SiLU plan");
+    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
+      return failure();
+    return success();
+  }
+};
+
+struct LowerRowStripResizePlan final : OpConversionPattern<spatial::SpatResizeNearestPlanOp> {
+  explicit LowerRowStripResizePlan(MLIRContext* ctx, const spatial::SpatialTargetResources& target)
+      : OpConversionPattern<spatial::SpatResizeNearestPlanOp>(ctx), target(target) {}
+
+  LogicalResult matchAndRewrite(spatial::SpatResizeNearestPlanOp planOp,
+                                spatial::SpatResizeNearestPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (!isRowStripSelected(planOp.getOperation()))
+      return failure();
+    FailureOr<RowStripPhysicalValue> input = getRowStripValue(adaptor.getInput());
+    if (failed(input)) {
+      if (getKnownPhysicalLayout(adaptor.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+        return failure();
+      return planOp.emitOpError("selected row-strip Resize plan requires a row-strip input");
+    }
+    FailureOr<Value> lowered = lowerSelectedResizeNearestPlan(
+        planOp, adaptor.getInput(), input->storage, target, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected row-strip Resize plan");
+    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
+      return failure();
+    return success();
+  }
+
+  const spatial::SpatialTargetResources& target;
+};
+
+struct LowerDenseMaxPoolPlan final : OpConversionPattern<spatial::SpatMaxPool2DPlanOp> {
+  explicit LowerDenseMaxPoolPlan(MLIRContext* ctx, const spatial::SpatialTargetResources& target)
+      : OpConversionPattern<spatial::SpatMaxPool2DPlanOp>(ctx), target(target) {}
+
+  LogicalResult matchAndRewrite(spatial::SpatMaxPool2DPlanOp planOp,
+                                spatial::SpatMaxPool2DPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (!isDenseSelected(planOp.getOperation()))
+      return failure();
+    FailureOr<Value> lowered = lowerDenseMaxPool2DPlan(
+        planOp, adaptor.getInput(), target, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected dense Spatial MaxPool plan");
+    rewriter.replaceOp(planOp, *lowered);
+    return success();
+  }
+
+  const spatial::SpatialTargetResources& target;
+};
+
+struct LowerRowStripMaxPoolPlan final : OpConversionPattern<spatial::SpatMaxPool2DPlanOp> {
+  explicit LowerRowStripMaxPoolPlan(MLIRContext* ctx, const spatial::SpatialTargetResources& target)
+      : OpConversionPattern<spatial::SpatMaxPool2DPlanOp>(ctx), target(target) {}
+
+  LogicalResult matchAndRewrite(spatial::SpatMaxPool2DPlanOp planOp,
+                                spatial::SpatMaxPool2DPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (!isRowStripSelected(planOp.getOperation()))
+      return failure();
+    FailureOr<RowStripPhysicalValue> input = getRowStripValue(adaptor.getInput());
+    if (failed(input)
+        && getKnownPhysicalLayout(adaptor.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+      return failure();
+    std::optional<Value> physicalInput;
+    if (succeeded(input))
+      physicalInput = input->storage;
+    FailureOr<Value> lowered = lowerSelectedMaxPool2DPlan(
+        planOp, adaptor.getInput(), physicalInput, target, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected row-strip Spatial MaxPool plan");
+    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
+      return failure();
+    return success();
+  }
+
+  const spatial::SpatialTargetResources& target;
+};
+
+struct LowerRowStripGlobalAveragePoolPlan
+    final : OpConversionPattern<spatial::SpatGlobalAveragePoolPlanOp> {
+  explicit LowerRowStripGlobalAveragePoolPlan(MLIRContext* ctx, const spatial::SpatialTargetResources& target)
+      : OpConversionPattern<spatial::SpatGlobalAveragePoolPlanOp>(ctx), target(target) {}
+
+  LogicalResult matchAndRewrite(spatial::SpatGlobalAveragePoolPlanOp planOp,
+                                spatial::SpatGlobalAveragePoolPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (!isRowStripSelected(planOp.getOperation()))
+      return failure();
+    FailureOr<RowStripPhysicalValue> input = getRowStripValue(adaptor.getInput());
+    if (failed(input)
+        && getKnownPhysicalLayout(adaptor.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+      return failure();
+    std::optional<Value> physicalInput;
+    if (succeeded(input))
+      physicalInput = input->storage;
+    FailureOr<Value> lowered = lowerSelectedGlobalAveragePoolPlan(
+        planOp, adaptor.getInput(), physicalInput, target, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected row-strip Spatial global AveragePool plan");
+    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
+      return failure();
+    return success();
+  }
+
+  const spatial::SpatialTargetResources& target;
+};
+
+struct LowerDenseGlobalAveragePoolPlan
+    final : OpConversionPattern<spatial::SpatGlobalAveragePoolPlanOp> {
+  explicit LowerDenseGlobalAveragePoolPlan(MLIRContext* ctx,
+                                           const spatial::SpatialTargetResources& target)
+      : OpConversionPattern<spatial::SpatGlobalAveragePoolPlanOp>(ctx), target(target) {}
+
+  LogicalResult matchAndRewrite(spatial::SpatGlobalAveragePoolPlanOp planOp,
+                                spatial::SpatGlobalAveragePoolPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (!isDenseSelected(planOp.getOperation()))
+      return failure();
+    FailureOr<Value> lowered = lowerDenseGlobalAveragePoolPlan(
+        planOp, adaptor.getInput(), target, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected dense Spatial global AveragePool plan");
+    rewriter.replaceOp(planOp, *lowered);
+    return success();
+  }
+
+  const spatial::SpatialTargetResources& target;
+};
+
+struct LowerRowStripBiasAddPlan final : OpConversionPattern<spatial::SpatBiasAddPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatBiasAddPlanOp planOp,
+                                spatial::SpatBiasAddPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (!isRowStripSelected(planOp.getOperation()))
+      return failure();
+    FailureOr<RowStripPhysicalValue> input = getRowStripValue(adaptor.getInput());
+    if (failed(input)) {
+      if (getKnownPhysicalLayout(adaptor.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+        return failure();
+      return planOp.emitOpError("selected row-strip bias_add plan requires a row-strip input");
+    }
+    FailureOr<Value> lowered = applyRowStripBiasAdd(
+        *input, adaptor.getBias(), rewriter, planOp.getLoc());
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected row-strip Spatial bias_add plan");
+    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
+      return failure();
+    return success();
+  }
+};
+
+struct LowerRowStripAddPlan final : OpConversionPattern<spatial::SpatAddPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatAddPlanOp planOp,
+                                spatial::SpatAddPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (!isRowStripSelected(planOp.getOperation()))
+      return failure();
+    return lowerAddPlan(planOp, adaptor.getLhs(), adaptor.getRhs(), rewriter);
+  }
+};
+
+struct LowerRowStripConcatPlan final : OpConversionPattern<spatial::SpatConcatPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatConcatPlanOp planOp,
+                                spatial::SpatConcatPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (!isRowStripSelected(planOp.getOperation()))
+      return failure();
+    return lowerConcatPlan(planOp, adaptor.getInputs(), rewriter);
+  }
+};
+
+struct LowerMaterializeLayout final
+    : OpConversionPattern<spatial::SpatMaterializeLayoutOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatMaterializeLayoutOp materializeOp,
+                                spatial::SpatMaterializeLayoutOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    auto source = materializeOp.getSourcePhysicalLayout();
+    auto target = materializeOp.getTargetPhysicalLayout();
+    if (source == spatial::PhysicalLayout::DenseNCHW
+        && target == spatial::PhysicalLayout::DenseNCHW) {
+      rewriter.replaceOp(materializeOp, adaptor.getInput());
+      return success();
+    }
+    if (source == spatial::PhysicalLayout::DenseNCHW
+        && target == spatial::PhysicalLayout::NHWCRowStrip) {
+      auto logicalType = dyn_cast<RankedTensorType>(adaptor.getInput().getType());
+      if (!logicalType)
+        return materializeOp.emitOpError("requires a ranked dense input"), failure();
+      FailureOr<Value> rowStrip = materializeDenseToRowStrip(
+        adaptor.getInput(), logicalType, materializeOp.getLoc(), rewriter);
+      if (failed(rowStrip))
+        return materializeOp.emitOpError(
+          "failed to materialize dense NCHW storage to row-strip layout"), failure();
+      rewriter.replaceOp(materializeOp, *rowStrip);
+      return success();
+    }
+    if (source != spatial::PhysicalLayout::NHWCRowStrip
+        || target != spatial::PhysicalLayout::DenseNCHW)
+      return materializeOp.emitOpError(
+          "unsupported Spatial layout materialization direction"), failure();
+    auto inputType = dyn_cast<RankedTensorType>(adaptor.getInput().getType());
+    if (!inputType)
+      return materializeOp.emitOpError("requires a ranked row-strip input"), failure();
+    FailureOr<RowStripPhysicalValue> rowStripValue =
+        getRowStripValue(adaptor.getInput());
+    if (failed(rowStripValue))
+      return failure();
+    FailureOr<Value> dense = materializeRowStripToDense(
+        *rowStripValue, materializeOp.getLoc(), rewriter);
+    if (failed(dense))
+      return materializeOp.emitOpError(
+          "failed to materialize row-strip storage to dense NCHW"), failure();
+    rewriter.replaceOp(materializeOp, *dense);
+    return success();
+  }
+};
+
+struct LowerSelectedFlattenPlan final
+    : OpConversionPattern<spatial::SpatFlattenPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatFlattenPlanOp planOp,
+                                spatial::SpatFlattenPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (!isDenseSelected(planOp.getOperation()))
+      return failure();
+    FailureOr<RowStripPhysicalValue> rowStripInput = getRowStripValue(adaptor.getInput());
+    if (succeeded(rowStripInput)) {
+      if (failed(canLowerFlattenFromRowStrip(planOp, target))
+          || failed(lowerFlattenFromRowStrip(*rowStripInput, planOp, target, rewriter)))
+        return planOp.emitOpError("failed to lower selected Spatial Flatten plan"), failure();
+      return success();
+    }
+    FailureOr<Value> lowered = lowerDenseFlattenPlan(planOp, adaptor.getInput(), rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected dense Spatial Flatten plan"), failure();
+    rewriter.replaceOp(planOp, *lowered);
+    return success();
+  }
+
+  explicit LowerSelectedFlattenPlan(MLIRContext* context,
+                                    const spatial::SpatialTargetResources& target)
+      : OpConversionPattern<spatial::SpatFlattenPlanOp>(context), target(target) {}
+
+  const spatial::SpatialTargetResources& target;
+};
+
+struct EraseDeadPhysicalViewBlueprint final
+    : OpRewritePattern<spatial::SpatBlueprintOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatBlueprintOp blueprint,
+                                PatternRewriter& rewriter) const override {
+    if (!spatial::isPhysicalView(blueprint.getMode()) || !blueprint.use_empty())
+      return failure();
+    rewriter.eraseOp(blueprint);
+    return success();
+  }
+};
+
+static void populateConvPlanLoweringPatterns(
+    RewritePatternSet& patterns, MLIRContext* ctx,
+    const spatial::SpatialTargetResources& target,
+    const ONNXToSpatialPlanningOptions& options) {
+  patterns.add<LowerSelectedConvPlan>(ctx, target, options);
+}
+
+static void populateElementwisePlanLoweringPatterns(
+    RewritePatternSet& patterns, MLIRContext* ctx) {
+  patterns.add<LowerDenseReluPlan,
+               LowerRowStripReluPlan,
+               LowerDenseSiluPlan,
+               LowerRowStripSiluPlan,
+               LowerDenseBiasAddPlan,
+               LowerRowStripBiasAddPlan,
+               LowerDenseAddPlan,
+               LowerRowStripAddPlan>(ctx);
+}
+
+static void populatePoolPlanLoweringPatterns(
+    RewritePatternSet& patterns, MLIRContext* ctx,
+    const spatial::SpatialTargetResources& target) {
+  patterns.add<LowerDenseMaxPoolPlan,
+               LowerRowStripMaxPoolPlan,
+               LowerDenseGlobalAveragePoolPlan,
+               LowerRowStripGlobalAveragePoolPlan>(ctx, target);
+}
+
+static void populateResizePlanLoweringPatterns(
+    RewritePatternSet& patterns, MLIRContext* ctx,
+    const spatial::SpatialTargetResources& target) {
+  patterns.add<LowerDenseResizePlan, LowerRowStripResizePlan>(ctx, target);
+}
+
+static void populateConcatPlanLoweringPatterns(
+    RewritePatternSet& patterns, MLIRContext* ctx) {
+  patterns.add<LowerDenseConcatPlan, LowerRowStripConcatPlan>(ctx);
+}
+
+static void populateFlattenPlanLoweringPatterns(
+    RewritePatternSet& patterns, MLIRContext* ctx,
+    const spatial::SpatialTargetResources& target) {
+  patterns.add<LowerSelectedFlattenPlan>(ctx, target);
+}
+
+static void populateLayoutMaterializationPatterns(
+    RewritePatternSet& patterns, MLIRContext* ctx) {
+  patterns.add<LowerMaterializeLayout, EraseDeadPhysicalViewBlueprint>(ctx);
+}
+
+} // namespace
+
+void populateSpatialPlanLoweringPatterns(
+    RewritePatternSet& patterns, MLIRContext* ctx,
+    const spatial::SpatialTargetResources& target,
+    const ONNXToSpatialPlanningOptions& options) {
+  populateConvPlanLoweringPatterns(patterns, ctx, target, options);
+  populateElementwisePlanLoweringPatterns(patterns, ctx);
+  populatePoolPlanLoweringPatterns(patterns, ctx, target);
+  populateResizePlanLoweringPatterns(patterns, ctx, target);
+  populateConcatPlanLoweringPatterns(patterns, ctx);
+  populateFlattenPlanLoweringPatterns(patterns, ctx, target);
+  populateLayoutMaterializationPatterns(patterns, ctx);
+}
+
+LogicalResult verifySelectedSpatialLayouts(
+    func::FuncOp funcOp, const spatial::SpatialTargetResources& target) {
+  return verifySelectedLayouts(funcOp, target);
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/SpatialPlanLoweringPatterns.hpp

@@ -0,0 +1,20 @@
+#pragma once
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Support/LogicalResult.h"
+
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
+
+namespace onnx_mlir {
+
+void populateSpatialPlanLoweringPatterns(
+    mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx,
+    const spatial::SpatialTargetResources& target,
+    const ONNXToSpatialPlanningOptions& options);
+
+mlir::LogicalResult verifySelectedSpatialLayouts(
+    mlir::func::FuncOp funcOp, const spatial::SpatialTargetResources& target);
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Patterns.cpp

@@ -7,12 +7,15 @@ namespace onnx_mlir {
 
 void populatePrePatterns(RewritePatternSet& patterns, MLIRContext* ctx) { populateGeneratedPrePatterns(patterns, ctx); }
 
-void populateConversionPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
+void populateConversionPatterns(RewritePatternSet& patterns,
+                                MLIRContext* ctx,
+                                const spatial::SpatialTargetResources& target,
+                                const ONNXToSpatialPlanningOptions& options) {
   populateElementwisePatterns(patterns, ctx);
-  populateMatMulRewritePatterns(patterns, ctx);
-  populateGemmPatterns(patterns, ctx);
-  populateConvPatterns(patterns, ctx);
-  populatePoolPatterns(patterns, ctx);
+  populateMatMulRewritePatterns(patterns, ctx, target);
+  populateGemmPatterns(patterns, ctx, target);
+  populateConvPatterns(patterns, ctx, target, options);
+  populatePoolPatterns(patterns, ctx, target);
   populateReduceMeanPatterns(patterns, ctx);
   populateReluPatterns(patterns, ctx);
   populateSigmoidPatterns(patterns, ctx);

src/PIM/Conversion/ONNXToSpatial/Patterns.hpp

@@ -4,22 +4,43 @@
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/Transforms/DialectConversion.h"
 
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 namespace onnx_mlir {
 
+namespace spatial {
+struct SpatialTargetResources;
+}
+
 void populatePrePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
-void populateConversionPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
+void populateConversionPatterns(mlir::RewritePatternSet& patterns,
+                                mlir::MLIRContext* ctx,
+                                const spatial::SpatialTargetResources& target,
+                                const ONNXToSpatialPlanningOptions& options);
 void populatePostPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 
 void populateGeneratedPrePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateWeightPromotionPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 
-void populateConvPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
+void populateConvPatterns(mlir::RewritePatternSet& patterns,
+                          mlir::MLIRContext* ctx,
+                          const spatial::SpatialTargetResources& target,
+                          const ONNXToSpatialPlanningOptions& options);
 void populateElementwisePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
-void populateGemmPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
-void populateMatMulRewritePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
-void populatePoolPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
+void populateElementwiseFusionPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
+void populateGemmPatterns(mlir::RewritePatternSet& patterns,
+                          mlir::MLIRContext* ctx,
+                          const spatial::SpatialTargetResources& target);
+void populateMatMulRewritePatterns(mlir::RewritePatternSet& patterns,
+                                   mlir::MLIRContext* ctx,
+                                   const spatial::SpatialTargetResources& target);
+void populateMatMulFusionPatterns(mlir::RewritePatternSet& patterns,
+                                  mlir::MLIRContext* ctx,
+                                  const spatial::SpatialTargetResources& target);
+void populatePoolPatterns(mlir::RewritePatternSet& patterns,
+                          mlir::MLIRContext* ctx,
+                          const spatial::SpatialTargetResources& target);
 void populateReduceMeanPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateReluPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateSigmoidPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);

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

@@ -1,49 +1,133 @@
 #include "ConvGeometry.hpp"
 
 #include <algorithm>
+#include <limits>
 
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
-#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 
 namespace onnx_mlir {
 
+namespace {
+
+static const ONNXToSpatialPlanningOptions& defaultPlanningOptions() {
+  static const ONNXToSpatialPlanningOptions options {
+    std::numeric_limits<uint64_t>::max(),
+    std::numeric_limits<uint64_t>::max(),
+    spatial::ConvLoweringStrategy::Auto,
+    false,
+  };
+  return options;
+}
+
+} // namespace
+
+const ONNXToSpatialPlanningOptions& ConvLoweringState::planningOptions() const {
+  return options ? *options : defaultPlanningOptions();
+}
+
 bool isDepthwiseConv(int64_t group, int64_t numChannelsIn, int64_t numChannelsOut, int64_t numChannelsInPerGroup) {
   return group == numChannelsIn && numChannelsInPerGroup == 1 && numChannelsOut % group == 0;
 }
 
-ConvGeometry buildConvGeometry(const ConvLoweringState& state) {
+void classifyConvProblem(ConvProblem& problem) {
+  problem.isDepthwise = isDepthwiseConv(
+    problem.group, problem.numChannelsIn, problem.numChannelsOut,
+    problem.numChannelsInPerGroup);
+  problem.isGrouped = problem.group > 1;
+  problem.isPointwise = problem.wHeight == 1 && problem.wWidth == 1
+    && problem.strideHeight == 1 && problem.strideWidth == 1
+    && problem.dilationHeight == 1 && problem.dilationWidth == 1
+    && problem.padHeightBegin == 0 && problem.padHeightEnd == 0
+    && problem.padWidthBegin == 0 && problem.padWidthEnd == 0;
+}
+
+ConvGeometry buildConvGeometry(const ConvProblem& problem,
+                               const spatial::SpatialTargetResources& target) {
   ConvGeometry geo {
-    state.batchSize,
-    state.numChannelsIn,
-    state.xHeight,
-    state.xWidth,
-    state.numChannelsOut,
-    state.wHeight,
-    state.wWidth,
-    state.outHeight,
-    state.outWidth,
-    state.group,
-    state.numChannelsInPerGroup,
-    state.numChannelsOutPerGroup,
-    state.numChannelsInPerGroup * state.wHeight * state.wWidth,
-    state.numChannelsOutPerGroup,
-    state.batchSize * state.outHeight * state.outWidth,
-    static_cast<int64_t>(crossbarSize.getValue()),
+    problem.numChannelsInPerGroup * problem.wHeight * problem.wWidth,
+    problem.numChannelsOutPerGroup,
+    problem.batchSize * problem.outHeight * problem.outWidth,
+    static_cast<int64_t>(target.matrixShape.rows),
+    static_cast<int64_t>(target.matrixUnitsPerProcessor),
     1,
     0,
-    state.hasBias,
-    isDepthwiseConv(state.group, state.numChannelsIn, state.numChannelsOut, state.numChannelsInPerGroup),
   };
   geo.pack = std::max<int64_t>(1, geo.xbarSize / std::max<int64_t>(geo.k, geo.c));
   geo.im2colElements = static_cast<uint64_t>(std::max<int64_t>(0, geo.p)) * static_cast<uint64_t>(std::max<int64_t>(0, geo.k));
   return geo;
 }
 
-uint64_t chooseStreamChunkPositions(const ConvGeometry& geo, int64_t packFactor) {
+static ConvMaterializationKind getMaterializationKind(
+    spatial::ConvLoweringStrategy strategy) {
+  switch (strategy) {
+  case spatial::ConvLoweringStrategy::Depthwise:
+    return ConvMaterializationKind::StructuredDepthwise;
+  case spatial::ConvLoweringStrategy::Legacy:
+  case spatial::ConvLoweringStrategy::PackedIm2Col:
+    return ConvMaterializationKind::PackedIm2Col;
+  case spatial::ConvLoweringStrategy::StreamedPatch:
+  case spatial::ConvLoweringStrategy::OutputChannelTiled:
+  case spatial::ConvLoweringStrategy::Tiled2D:
+    return ConvMaterializationKind::StreamedPatch;
+  case spatial::ConvLoweringStrategy::StreamedPacked:
+    return ConvMaterializationKind::StreamedPacked;
+  case spatial::ConvLoweringStrategy::InputKTiled:
+    return ConvMaterializationKind::InputKTiled;
+  case spatial::ConvLoweringStrategy::Auto:
+    break;
+  }
+  llvm_unreachable("auto is not a Conv materialization kind");
+}
+
+static bool fitsSingleCrossbar(const ConvGeometry& geo) {
+  return geo.k <= geo.xbarSize && geo.c <= geo.xbarSize;
+}
+
+static bool fitsPackedIm2Col(const ConvGeometry& geo,
+                             const ONNXToSpatialPlanningOptions& options) {
+  return fitsSingleCrossbar(geo) && geo.pack >= 2
+    && geo.im2colElements <= options.convIm2colMaxElements;
+}
+
+mlir::FailureOr<ConvPlan> makeConvPlan(const ConvProblem& problem,
+                                       spatial::ConvLoweringStrategy strategy,
+                                       const spatial::SpatialTargetResources& target,
+                                       const ONNXToSpatialPlanningOptions& options) {
+  ConvGeometry geo = buildConvGeometry(problem, target);
+  auto plan = [&]() { return ConvPlan {getMaterializationKind(strategy)}; };
+  auto ifApplicable = [&](bool applicable) -> mlir::FailureOr<ConvPlan> {
+    return applicable ? mlir::FailureOr<ConvPlan>(plan()) : mlir::FailureOr<ConvPlan>(mlir::failure());
+  };
+  switch (strategy) {
+  case spatial::ConvLoweringStrategy::Auto:
+    return mlir::failure();
+  case spatial::ConvLoweringStrategy::Legacy:
+    return plan();
+  case spatial::ConvLoweringStrategy::Depthwise:
+    return ifApplicable(problem.isDepthwise);
+  case spatial::ConvLoweringStrategy::PackedIm2Col:
+    return ifApplicable(fitsPackedIm2Col(geo, options));
+  case spatial::ConvLoweringStrategy::StreamedPatch:
+    return ifApplicable(fitsSingleCrossbar(geo));
+  case spatial::ConvLoweringStrategy::StreamedPacked:
+    return ifApplicable(fitsSingleCrossbar(geo) && geo.pack >= 2);
+  case spatial::ConvLoweringStrategy::OutputChannelTiled:
+    return ifApplicable(geo.k <= geo.xbarSize && geo.c > geo.xbarSize);
+  case spatial::ConvLoweringStrategy::InputKTiled:
+    return ifApplicable(geo.k > geo.xbarSize && geo.c <= geo.xbarSize);
+  case spatial::ConvLoweringStrategy::Tiled2D:
+    return ifApplicable(geo.k > geo.xbarSize && geo.c > geo.xbarSize);
+  }
+  llvm_unreachable("unknown Conv lowering strategy");
+}
+
+uint64_t chooseStreamChunkPositions(const ConvGeometry& geo,
+                                    int64_t packFactor,
+                                    const ONNXToSpatialPlanningOptions& options) {
   const uint64_t patchElements = static_cast<uint64_t>(std::max<int64_t>(1, geo.k));
-  uint64_t chunkPositions = std::max<uint64_t>(1, pimConvIm2colMaxElements / patchElements);
+  uint64_t chunkPositions = std::max<uint64_t>(1, options.convIm2colMaxElements / patchElements);
   chunkPositions = std::min<uint64_t>(chunkPositions, static_cast<uint64_t>(std::max<int64_t>(1, geo.p)));
-  chunkPositions = std::min<uint64_t>(chunkPositions, std::max<uint64_t>(1, pimConvStreamChunkPositions));
+  chunkPositions = std::min<uint64_t>(chunkPositions, std::max<uint64_t>(1, options.convStreamChunkPositions));
 
   if (packFactor > 1 && chunkPositions > static_cast<uint64_t>(packFactor)) {
     chunkPositions -= chunkPositions % static_cast<uint64_t>(packFactor);
@@ -52,24 +136,26 @@ uint64_t chooseStreamChunkPositions(const ConvGeometry& geo, int64_t packFactor)
   return std::max<uint64_t>(1, chunkPositions);
 }
 
-RowInterval computeConvInputRowsForOutputRows(RowInterval outputRows, const ConvLoweringState& state) {
-  const int64_t rawBegin = outputRows.begin * state.strideHeight - state.padHeightBegin;
+RowInterval computeConvInputRowsForOutputRows(RowInterval outputRows, const ConvProblem& problem) {
+  const int64_t rawBegin = outputRows.begin * problem.strideHeight - problem.padHeightBegin;
   const int64_t rawEnd =
-    (outputRows.end - 1) * state.strideHeight - state.padHeightBegin + state.dilationHeight * (state.wHeight - 1) + 1;
-  return {std::max<int64_t>(0, rawBegin), std::min<int64_t>(state.xHeight, rawEnd)};
+    (outputRows.end - 1) * problem.strideHeight - problem.padHeightBegin
+    + problem.dilationHeight * (problem.wHeight - 1) + 1;
+  return {std::max<int64_t>(0, rawBegin), std::min<int64_t>(problem.xHeight, rawEnd)};
 }
 
-ConvRowDemand buildConvRowDemand(RowInterval outputRows, const ConvLoweringState& state) {
+ConvRowDemand buildConvRowDemand(RowInterval outputRows, const ConvProblem& problem) {
   ConvRowDemand demand;
   demand.outputRows = outputRows;
-  demand.neededInputRows = computeConvInputRowsForOutputRows(outputRows, state);
+  demand.neededInputRows = computeConvInputRowsForOutputRows(outputRows, problem);
   demand.acquiredInputRows = demand.neededInputRows;
 
-  const int64_t rawBegin = outputRows.begin * state.strideHeight - state.padHeightBegin;
+  const int64_t rawBegin = outputRows.begin * problem.strideHeight - problem.padHeightBegin;
   const int64_t rawEnd =
-    (outputRows.end - 1) * state.strideHeight - state.padHeightBegin + state.dilationHeight * (state.wHeight - 1) + 1;
+    (outputRows.end - 1) * problem.strideHeight - problem.padHeightBegin
+    + problem.dilationHeight * (problem.wHeight - 1) + 1;
   demand.topHaloRows = std::max<int64_t>(0, -rawBegin);
-  demand.bottomHaloRows = std::max<int64_t>(0, rawEnd - state.xHeight);
+  demand.bottomHaloRows = std::max<int64_t>(0, rawEnd - problem.xHeight);
   demand.acquiredInputRows = demand.neededInputRows;
   return demand;
 }

src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ConvGeometry.hpp

@@ -3,14 +3,19 @@
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Value.h"
 
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
 #include <cstdint>
 
+namespace mlir {
+class Operation;
+} // namespace mlir
+
 namespace onnx_mlir {
 
-struct ConvLoweringState {
-  mlir::Value x;
-  mlir::Value w;
-  mlir::Value b;
+struct ConvProblem {
   mlir::RankedTensorType xType;
   mlir::RankedTensorType wType;
   mlir::RankedTensorType outType;
@@ -35,29 +40,32 @@ struct ConvLoweringState {
   int64_t dilationHeight;
   int64_t dilationWidth;
   bool hasBias;
+  bool isDepthwise = false;
+  bool isGrouped = false;
+  bool isPointwise = false;
+};
+
+struct ConvLoweringState {
+  ConvProblem problem;
+  mlir::Operation* diagnosticAnchor = nullptr;
+  mlir::Value x;
+  mlir::Value w;
+  mlir::Value b;
+  const spatial::SpatialTargetResources* target = nullptr;
+  const ONNXToSpatialPlanningOptions* options = nullptr;
+
+  const spatial::SpatialTargetResources& targetInfo() const { return *target; }
+  const ONNXToSpatialPlanningOptions& planningOptions() const;
 };
 
 struct ConvGeometry {
-  int64_t batchSize;
-  int64_t numChannelsIn;
-  int64_t xHeight;
-  int64_t xWidth;
-  int64_t numChannelsOut;
-  int64_t wHeight;
-  int64_t wWidth;
-  int64_t outHeight;
-  int64_t outWidth;
-  int64_t group;
-  int64_t numChannelsInPerGroup;
-  int64_t numChannelsOutPerGroup;
   int64_t k;
   int64_t c;
   int64_t p;
   int64_t xbarSize;
+  int64_t matrixUnitsPerProcessor;
   int64_t pack;
   uint64_t im2colElements;
-  bool hasBias;
-  bool isDepthwise;
 };
 
 struct RowInterval {
@@ -73,14 +81,36 @@ struct ConvRowDemand {
   int64_t bottomHaloRows = 0;
 };
 
+enum class ConvMaterializationKind : uint8_t {
+  StructuredDepthwise,
+  PackedIm2Col,
+  StreamedPatch,
+  StreamedPacked,
+  InputKTiled,
+};
+
+struct ConvPlan {
+  ConvMaterializationKind kind = ConvMaterializationKind::PackedIm2Col;
+};
+
 bool isDepthwiseConv(int64_t group, int64_t numChannelsIn, int64_t numChannelsOut, int64_t numChannelsInPerGroup);
 
-ConvGeometry buildConvGeometry(const ConvLoweringState& state);
+void classifyConvProblem(ConvProblem& problem);
 
-uint64_t chooseStreamChunkPositions(const ConvGeometry& geo, int64_t packFactor);
+ConvGeometry buildConvGeometry(const ConvProblem& problem,
+                               const spatial::SpatialTargetResources& target);
 
-RowInterval computeConvInputRowsForOutputRows(RowInterval outputRows, const ConvLoweringState& state);
+mlir::FailureOr<ConvPlan> makeConvPlan(const ConvProblem& problem,
+                                       spatial::ConvLoweringStrategy strategy,
+                                       const spatial::SpatialTargetResources& target,
+                                       const ONNXToSpatialPlanningOptions& options);
 
-ConvRowDemand buildConvRowDemand(RowInterval outputRows, const ConvLoweringState& state);
+uint64_t chooseStreamChunkPositions(const ConvGeometry& geo,
+                                    int64_t packFactor,
+                                    const ONNXToSpatialPlanningOptions& options);
+
+RowInterval computeConvInputRowsForOutputRows(RowInterval outputRows, const ConvProblem& problem);
+
+ConvRowDemand buildConvRowDemand(RowInterval outputRows, const ConvProblem& problem);
 
 } // namespace onnx_mlir

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

@@ -16,6 +16,28 @@ using namespace mlir;
 namespace onnx_mlir {
 namespace {
 
+struct SiluToSpatialPlan : OpRewritePattern<ONNXMulOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ONNXMulOp mulOp, PatternRewriter& rewriter) const override {
+    ONNXSigmoidOp sigmoidOp = mulOp.getA().getDefiningOp<ONNXSigmoidOp>();
+    Value input = mulOp.getB();
+    if (!sigmoidOp) {
+      sigmoidOp = mulOp.getB().getDefiningOp<ONNXSigmoidOp>();
+      input = mulOp.getA();
+    }
+    if (!sigmoidOp || sigmoidOp.getX() != input || !sigmoidOp->hasOneUse()
+        || sigmoidOp.getResult().getType() != input.getType() || mulOp.getResult().getType() != input.getType())
+      return failure();
+
+    auto plan = spatial::SpatSiluPlanOp::create(
+      rewriter, mulOp.getLoc(), mulOp.getResult().getType(), input, spatial::getNCHWLayout(rewriter.getContext()));
+    rewriter.replaceOp(mulOp, plan.getResult());
+    rewriter.eraseOp(sigmoidOp);
+    return success();
+  }
+};
+
 static DenseElementsAttr getDenseConstantAttr(Value value) {
   if (auto constantOp = value.getDefiningOp<arith::ConstantOp>())
     return dyn_cast<DenseElementsAttr>(constantOp.getValue());
@@ -26,6 +48,56 @@ static DenseElementsAttr getDenseConstantAttr(Value value) {
   return nullptr;
 }
 
+struct BlueprintSplatMulToSpatial : OpConversionPattern<ONNXMulOp> {
+  explicit BlueprintSplatMulToSpatial(MLIRContext* ctx) : OpConversionPattern(ctx, 2) {}
+
+  LogicalResult
+  matchAndRewrite(ONNXMulOp op, ONNXMulOpAdaptor adaptor, ConversionPatternRewriter& rewriter) const override {
+    auto blueprint = op.getA().getDefiningOp<spatial::SpatBlueprintOp>();
+    Value scalar = adaptor.getB();
+    if (!blueprint) {
+      blueprint = op.getB().getDefiningOp<spatial::SpatBlueprintOp>();
+      scalar = adaptor.getA();
+    }
+    auto scalarAttr = getDenseConstantAttr(scalar);
+    auto resultType = dyn_cast<RankedTensorType>(op.getResult().getType());
+    auto storageType = blueprint ? dyn_cast<RankedTensorType>(blueprint.getInput().getType()) : RankedTensorType();
+    if (!blueprint || !blueprint.getFragments().empty() || !scalarAttr || !scalarAttr.isSplat() || !resultType
+        || resultType != blueprint.getOutput().getType() || !storageType)
+      return failure();
+
+    auto mapped = mapGraphBatchFragments(
+      blueprint.getInput(), storageType, rewriter, op.getLoc(), [&](Value fragment, RankedTensorType fragmentType) {
+        auto splat = DenseElementsAttr::get(fragmentType, scalarAttr.getSplatValue<Attribute>());
+        Value constant = arith::ConstantOp::create(rewriter, op.getLoc(), fragmentType, splat);
+        return FailureOr<Value>(
+          spatial::SpatVMulOp::create(rewriter, op.getLoc(), fragmentType, fragment, constant).getResult());
+      });
+    if (failed(mapped))
+      return failure();
+
+    auto result = spatial::SpatBlueprintOp::create(rewriter,
+                                                   op.getLoc(),
+                                                   resultType,
+                                                   *mapped,
+                                                   ValueRange {},
+                                                   blueprint.getLogicalLayoutAttr(),
+                                                   blueprint.getPhysicalLayoutAttr(),
+                                                   blueprint.getFragmentOffsetsAttr(),
+                                                   blueprint.getFragmentSizesAttr(),
+                                                   blueprint.getIndexMapAttr(),
+                                                   blueprint.getModeAttr(),
+                                                   blueprint.getFragmentOperandIndicesAttr(),
+                                                   blueprint.getFragmentSourceSlotsAttr(),
+                                                   blueprint.getFragmentSourceOffsetsAttr(),
+                                                   blueprint.getFragmentStridesAttr(),
+                                                   blueprint.getConflictPolicyAttr(),
+                                                   blueprint.getCoveragePolicyAttr());
+    rewriter.replaceOp(op, result.getOutput());
+    return success();
+  }
+};
+
 static FailureOr<Value> materializeBroadcastedConstantTensor(Value value,
                                                              RankedTensorType resultType,
                                                              ConversionPatternRewriter& rewriter,
@@ -188,14 +260,16 @@ struct AddToSpatialCompute : OpConversionPattern<ONNXAddOp> {
       classifyBiasAddPlanCandidate(adaptor.getA(), adaptor.getB(), resultType);
     if (succeeded(candidate)) {
       auto plan = spatial::SpatBiasAddPlanOp::create(
-        rewriter, op.getLoc(), resultType, candidate->data, candidate->bias, rewriter.getStringAttr("nchw"));
+        rewriter, op.getLoc(), resultType, candidate->data, candidate->bias,
+        spatial::getNCHWLayout(rewriter.getContext()));
       rewriter.replaceOp(op, plan.getResult());
       return success();
     }
 
     if (resultType.getRank() == 4 && adaptor.getA().getType() == resultType && adaptor.getB().getType() == resultType) {
       auto plan = spatial::SpatAddPlanOp::create(
-        rewriter, op.getLoc(), resultType, adaptor.getA(), adaptor.getB(), rewriter.getStringAttr("nchw"));
+        rewriter, op.getLoc(), resultType, adaptor.getA(), adaptor.getB(),
+        spatial::getNCHWLayout(rewriter.getContext()));
       rewriter.replaceOp(op, plan.getResult());
       return success();
     }
@@ -219,7 +293,12 @@ struct AddToSpatialCompute : OpConversionPattern<ONNXAddOp> {
 
 } // namespace
 
+void populateElementwiseFusionPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
+  patterns.add<SiluToSpatialPlan>(ctx);
+}
+
 void populateElementwisePatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
+  patterns.add<BlueprintSplatMulToSpatial>(ctx);
   patterns.add<AddToSpatialCompute>(ctx);
   patterns.add<BinaryElementwiseToSpatialCompute<ONNXSubOp, spatial::SpatVSubOp>>(ctx);
   patterns.add<BinaryElementwiseToSpatialCompute<ONNXMulOp, spatial::SpatVMulOp>>(ctx);

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

@@ -1,5 +1,6 @@
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/BuiltinTypes.h"
@@ -21,6 +22,9 @@
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ContractionProblem.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ContractionPlanning.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns/Math/Gemm.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -31,7 +35,7 @@ namespace onnx_mlir {
 namespace {
 
 static FailureOr<Value>
-materializeScaledConstantTensor(Value value, float factor, ConversionPatternRewriter& rewriter, Location loc) {
+materializeScaledConstantTensor(Value value, float factor, PatternRewriter& rewriter, Location loc) {
   if (factor == 1.0f)
     return value;
 
@@ -57,7 +61,12 @@ materializeScaledConstantTensor(Value value, float factor, ConversionPatternRewr
 }
 
 static Value createGemmBatchKOffset(
-  Value lane, int64_t numOutRows, int64_t numKSlices, ConversionPatternRewriter& rewriter, Location loc) {
+  Value lane,
+  int64_t numOutRows,
+  int64_t numKSlices,
+  int64_t xbarSize,
+  PatternRewriter& rewriter,
+  Location loc) {
   if (numKSlices == 1)
     return getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
 
@@ -65,7 +74,7 @@ static Value createGemmBatchKOffset(
   AffineExpr d0 = getAffineDimExpr(0, context);
   return createOrFoldAffineApply(rewriter,
                                  loc,
-                                 (d0.floorDiv(numOutRows) % numKSlices) * crossbarSize.getValue(),
+                                 (d0.floorDiv(numOutRows) % numKSlices) * xbarSize,
                                  ValueRange {lane},
                                  rewriter.getInsertionBlock()->getParentOp());
 }
@@ -74,7 +83,8 @@ static Value createGemmBatchHOffset(Value lane,
                                     int64_t numOutRows,
                                     int64_t numKSlices,
                                     int64_t numOutHSlices,
-                                    ConversionPatternRewriter& rewriter,
+                                    int64_t xbarSize,
+                                    PatternRewriter& rewriter,
                                     Location loc) {
   if (numOutHSlices == 1)
     return getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
@@ -83,14 +93,14 @@ static Value createGemmBatchHOffset(Value lane,
   AffineExpr d0 = getAffineDimExpr(0, context);
   return createOrFoldAffineApply(rewriter,
                                  loc,
-                                 d0.floorDiv(numOutRows * numKSlices) * crossbarSize.getValue(),
+                                 d0.floorDiv(numOutRows * numKSlices) * xbarSize,
                                  ValueRange {lane},
                                  rewriter.getInsertionBlock()->getParentOp());
 }
 
 static FailureOr<Value> materializePaddedConstantMatrix(Value value,
                                                         RankedTensorType resultType,
-                                                        ConversionPatternRewriter& rewriter,
+                                                        PatternRewriter& rewriter,
                                                         Location loc) {
   auto sourceType = cast<RankedTensorType>(value.getType());
   if (sourceType == resultType)
@@ -121,7 +131,7 @@ static FailureOr<Value> materializePaddedConstantMatrix(Value value,
 static FailureOr<Value> materializePaddedBroadcastedConstantTensor(Value value,
                                                                    RankedTensorType resultType,
                                                                    int64_t unpaddedColumns,
-                                                                   ConversionPatternRewriter& rewriter,
+                                                                   PatternRewriter& rewriter,
                                                                    Location loc) {
   auto denseAttr = getHostConstDenseElementsAttr(value);
   if (!denseAttr)
@@ -187,7 +197,7 @@ static FailureOr<Value> materializePaddedBroadcastedConstantTensor(Value value,
 static FailureOr<Value> prepareBias(Value c,
                                     RankedTensorType outType,
                                     RankedTensorType paddedOutType,
-                                    ConversionPatternRewriter& rewriter,
+                                    PatternRewriter& rewriter,
                                     Location loc) {
   auto cType = cast<RankedTensorType>(c.getType());
   if (!cType.hasStaticShape())
@@ -203,9 +213,15 @@ static FailureOr<Value> prepareBias(Value c,
 }
 
 static Value extractATile(
-  Value a, Value row, Value kOffset, RankedTensorType aTileType, ConversionPatternRewriter& rewriter, Location loc) {
+  Value a,
+  Value row,
+  Value kOffset,
+  RankedTensorType aTileType,
+  int64_t xbarSize,
+  PatternRewriter& rewriter,
+  Location loc) {
   SmallVector<OpFoldResult> offsets {row, kOffset};
-  SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(crossbarSize.getValue())};
+  SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarSize)};
   SmallVector<OpFoldResult> strides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
 
   return tensor::ExtractSliceOp::create(rewriter, loc, aTileType, a, offsets, sizes, strides).getResult();
@@ -219,7 +235,8 @@ static FailureOr<spatial::SpatComputeBatch> createVmmBatch(Value a,
                                                            int64_t numOutRows,
                                                            int64_t numKSlices,
                                                            int64_t numOutHSlices,
-                                                           ConversionPatternRewriter& rewriter,
+                                                           int64_t xbarSize,
+                                                           PatternRewriter& rewriter,
                                                            Location loc) {
   const int64_t laneCount = partialPiecesType.getDimSize(0);
   auto batchOp = createSpatComputeBatch(
@@ -232,21 +249,21 @@ static FailureOr<spatial::SpatComputeBatch> createVmmBatch(Value a,
     [&](detail::SpatComputeBatchBodyArgs args) {
       Value row =
         onnx_mlir::affineModConst(rewriter, loc, args.lane, numOutRows, rewriter.getInsertionBlock()->getParentOp());
-      Value kOffset = createGemmBatchKOffset(args.lane, numOutRows, numKSlices, rewriter, loc);
-      Value hOffset = createGemmBatchHOffset(args.lane, numOutRows, numKSlices, numOutHSlices, rewriter, loc);
+      Value kOffset = createGemmBatchKOffset(args.lane, numOutRows, numKSlices, xbarSize, rewriter, loc);
+      Value hOffset = createGemmBatchHOffset(
+        args.lane, numOutRows, numKSlices, numOutHSlices, xbarSize, rewriter, loc);
 
       auto aTileType =
-        RankedTensorType::get({1, static_cast<int64_t>(crossbarSize.getValue())}, aType.getElementType());
+        RankedTensorType::get({1, xbarSize}, aType.getElementType());
       auto bTileType = RankedTensorType::get(
-        {static_cast<int64_t>(crossbarSize.getValue()), static_cast<int64_t>(crossbarSize.getValue())},
+        {xbarSize, xbarSize},
         paddedBType.getElementType());
       auto pieceType =
-        RankedTensorType::get({1, static_cast<int64_t>(crossbarSize.getValue())}, partialPiecesType.getElementType());
-      Value aTile = extractATile(args.inputs.front(), row, kOffset, aTileType, rewriter, loc);
+        RankedTensorType::get({1, xbarSize}, partialPiecesType.getElementType());
+      Value aTile = extractATile(args.inputs.front(), row, kOffset, aTileType, xbarSize, rewriter, loc);
 
       SmallVector<OpFoldResult> bOffsets {kOffset, hOffset};
-      SmallVector<OpFoldResult> bSizes {rewriter.getIndexAttr(crossbarSize.getValue()),
-                                        rewriter.getIndexAttr(crossbarSize.getValue())};
+      SmallVector<OpFoldResult> bSizes {rewriter.getIndexAttr(xbarSize), rewriter.getIndexAttr(xbarSize)};
       SmallVector<OpFoldResult> unitStrides = getUnitStrides(rewriter, 2);
       Value bTile = extractStaticSliceOrIdentity(
         rewriter, loc, args.weights.front(), bTileType, bOffsets, bSizes, unitStrides);
@@ -259,19 +276,8 @@ static FailureOr<spatial::SpatComputeBatch> createVmmBatch(Value a,
   return *batchOp;
 }
 
-static Value
-createDynamicGemmBatchRow(Value lane, int64_t numOutCols, ConversionPatternRewriter& rewriter, Location loc) {
-  if (numOutCols == 1)
-    return lane;
-
-  MLIRContext* context = rewriter.getContext();
-  AffineExpr d0 = getAffineDimExpr(0, context);
-  return createOrFoldAffineApply(
-    rewriter, loc, d0.floorDiv(numOutCols), ValueRange {lane}, rewriter.getInsertionBlock()->getParentOp());
-}
-
 static Value extractDynamicGemmBColumn(
-  Value matrix, Value column, RankedTensorType vectorType, ConversionPatternRewriter& rewriter, Location loc) {
+  Value matrix, Value column, RankedTensorType vectorType, PatternRewriter& rewriter, Location loc) {
   SmallVector<OpFoldResult> offsets {rewriter.getIndexAttr(0), column};
   SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(vectorType.getDimSize(1)), rewriter.getIndexAttr(1)};
   SmallVector<OpFoldResult> strides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
@@ -291,7 +297,7 @@ static Value extractDynamicGemmBColumn(
 }
 
 static Value extractDynamicGemmRowVector(
-  Value matrix, Value row, RankedTensorType vectorType, ConversionPatternRewriter& rewriter, Location loc) {
+  Value matrix, Value row, RankedTensorType vectorType, PatternRewriter& rewriter, Location loc) {
   SmallVector<OpFoldResult> offsets {row, rewriter.getIndexAttr(0)};
   SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(vectorType.getDimSize(1))};
   SmallVector<OpFoldResult> strides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
@@ -322,13 +328,15 @@ static FailureOr<RankedTensorType> verifyDynamicGemmBiasType(RankedTensorType cT
 }
 
 static bool hasGemmBias(Value c) {
+  if (!c)
+    return false;
   Operation* definingOp = c.getDefiningOp();
   return (!definingOp || !isa<ONNXNoneOp>(definingOp)) && !isZeroSplatHostConstant(c);
 }
 
 static Value createScalarTensorConstant(RankedTensorType scalarType,
                                         float value,
-                                        ConversionPatternRewriter& rewriter,
+                                        PatternRewriter& rewriter,
                                         Location loc) {
   auto elementType = scalarType.getElementType();
   auto scalarAttr = rewriter.getFloatAttr(elementType, value);
@@ -341,7 +349,7 @@ static Value createBroadcastedBiasScalar(Value bias,
                                          Value row,
                                          Value column,
                                          RankedTensorType scalarType,
-                                         ConversionPatternRewriter& rewriter,
+                                         PatternRewriter& rewriter,
                                          Location loc) {
   SmallVector<OpFoldResult> unitStrides(biasType.getRank(), rewriter.getIndexAttr(1));
   if (biasType.getRank() == 1) {
@@ -373,33 +381,56 @@ static FailureOr<spatial::SpatComputeBatch> createVvdmulBatch(Value a,
                                                               Value b,
                                                               RankedTensorType aType,
                                                               RankedTensorType bType,
-                                                              RankedTensorType scalarPiecesType,
+                                                              RankedTensorType columnPiecesType,
                                                               RankedTensorType outType,
-                                                              ConversionPatternRewriter& rewriter,
+                                                              bool transposeB,
+                                                              PatternRewriter& rewriter,
                                                               Location loc) {
   const int64_t numOutRows = outType.getDimSize(0);
   const int64_t numOutCols = outType.getDimSize(1);
   const int64_t reductionSize = aType.getDimSize(1);
-  const int64_t laneCount = numOutRows * numOutCols;
+  auto vectorType = RankedTensorType::get({1, reductionSize}, aType.getElementType());
+  auto scalarType = RankedTensorType::get({1, 1}, outType.getElementType());
+  auto columnType = RankedTensorType::get({numOutRows, 1}, outType.getElementType());
   auto batchOp = createSpatComputeBatch(
     rewriter,
     loc,
-    TypeRange {scalarPiecesType},
-    laneCount,
+    TypeRange {columnPiecesType},
+    numOutCols,
     ValueRange {},
     ValueRange {a, b},
     [&](detail::SpatComputeBatchBodyArgs args) {
-      Value row = createDynamicGemmBatchRow(args.lane, numOutCols, rewriter, loc);
-      Value column =
-        onnx_mlir::affineModConst(rewriter, loc, args.lane, numOutCols, rewriter.getInsertionBlock()->getParentOp());
-
-      auto vectorType = RankedTensorType::get({1, reductionSize}, aType.getElementType());
-      auto scalarType = RankedTensorType::get({1, 1}, outType.getElementType());
-      Value aVector = extractDynamicGemmRowVector(args.inputs[0], row, vectorType, rewriter, loc);
-      Value bVector = extractDynamicGemmBColumn(args.inputs[1], column, vectorType, rewriter, loc);
-      Value scalar = spatial::SpatVVDMulOp::create(rewriter, loc, scalarType, aVector, bVector).getResult();
-
-      publishGraphBatchPhysicalFragment(rewriter, loc, scalar, args.outputs.front(), args.lane);
+      Value bVector = transposeB
+                        ? extractDynamicGemmRowVector(args.inputs[1], args.lane, vectorType, rewriter, loc)
+                        : extractDynamicGemmBColumn(args.inputs[1], args.lane, vectorType, rewriter, loc);
+      Value columnInit = tensor::EmptyOp::create(rewriter, loc, columnType.getShape(), columnType.getElementType());
+      Value c0 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
+      Value c1 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 1);
+      Value cNumOutRows =
+        getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), numOutRows);
+      auto loop = buildNormalizedScfFor(
+        rewriter,
+        loc,
+        c0,
+        cNumOutRows,
+        c1,
+        ValueRange {columnInit},
+        [&](OpBuilder&, Location nestedLoc, Value row, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
+          Value aVector = extractDynamicGemmRowVector(args.inputs[0], row, vectorType, rewriter, nestedLoc);
+          Value scalar = spatial::SpatVVDMulOp::create(rewriter, nestedLoc, scalarType, aVector, bVector).getResult();
+          Value next = tensor::InsertSliceOp::create(rewriter,
+                                                     nestedLoc,
+                                                     scalar,
+                                                     iterArgs.front(),
+                                                     SmallVector<OpFoldResult> {row, rewriter.getIndexAttr(0)},
+                                                     SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
+                                                                               rewriter.getIndexAttr(1)},
+                                                     getUnitStrides(rewriter, 2));
+          yielded.push_back(next);
+          return success();
+        });
+      assert(succeeded(loop) && "dynamic Gemm row loop construction must succeed");
+      publishGraphBatchPhysicalFragment(rewriter, loc, loop->results.front(), args.outputs.front(), args.lane);
     });
   if (failed(batchOp))
     return failure();
@@ -413,9 +444,9 @@ static FailureOr<spatial::SpatCompute> createDynamicGemmOutputCompute(Value scal
                                                                       RankedTensorType outType,
                                                                       float alpha,
                                                                       float beta,
-                                                                      ConversionPatternRewriter& rewriter,
+                                                                      PatternRewriter& rewriter,
                                                                       Location loc) {
-  const int64_t laneCount = scalarPiecesType.getDimSize(0);
+  const int64_t numOutRows = outType.getDimSize(0);
   const int64_t numOutCols = outType.getDimSize(1);
   SmallVector<Value> inputs {scalarPieces};
   if (bias)
@@ -428,43 +459,62 @@ static FailureOr<spatial::SpatCompute> createDynamicGemmOutputCompute(Value scal
     Value outputInit = tensor::EmptyOp::create(rewriter, loc, outType.getShape(), outType.getElementType()).getResult();
     Value c0 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
     Value c1 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 1);
-    Value cLaneCount = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), laneCount);
+    Value cNumOutCols = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), numOutCols);
+    auto columnType = RankedTensorType::get({numOutRows, 1}, outType.getElementType());
     auto loop = buildNormalizedScfFor(
       rewriter,
       loc,
       c0,
-      cLaneCount,
+      cNumOutCols,
       c1,
       ValueRange {outputInit},
-      [&](OpBuilder&, Location nestedLoc, Value lane, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
+      [&](OpBuilder&, Location nestedLoc, Value column, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
         Value outputAcc = iterArgs.front();
-        Value row = createDynamicGemmBatchRow(lane, numOutCols, rewriter, nestedLoc);
-        Value column =
-          onnx_mlir::affineModConst(rewriter, nestedLoc, lane, numOutCols, rewriter.getInsertionBlock()->getParentOp());
-        SmallVector<OpFoldResult> scalarSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-        SmallVector<OpFoldResult> unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-        FailureOr<Value> scalar = extractGraphBatchPhysicalFragment(rewriter, nestedLoc, pieces, lane, scalarType);
-        if (failed(scalar))
+        FailureOr<Value> columnPiece =
+          extractGraphBatchPhysicalFragment(rewriter, nestedLoc, pieces, column, columnType);
+        if (failed(columnPiece))
           return failure();
-        if (alpha != 1.0f) {
-          Value alphaTensor = createScalarTensorConstant(scalarType, alpha, rewriter, nestedLoc);
-          *scalar = spatial::SpatVMulOp::create(rewriter, nestedLoc, scalarType, *scalar, alphaTensor).getResult();
-        }
-        if (biasArg) {
-          Value biasScalar =
-            createBroadcastedBiasScalar(biasArg, biasType, row, column, scalarType, rewriter, nestedLoc);
-          if (beta != 1.0f) {
-            Value betaTensor = createScalarTensorConstant(scalarType, beta, rewriter, nestedLoc);
-            biasScalar =
-              spatial::SpatVMulOp::create(rewriter, nestedLoc, scalarType, biasScalar, betaTensor).getResult();
-          }
-          *scalar = spatial::SpatVAddOp::create(rewriter, nestedLoc, scalarType, *scalar, biasScalar).getResult();
-        }
-        SmallVector<OpFoldResult> outputOffsets {row, column};
-        Value outputNext =
-          tensor::InsertSliceOp::create(rewriter, nestedLoc, *scalar, outputAcc, outputOffsets, scalarSizes, unitStrides)
-            .getResult();
-        yielded.push_back(outputNext);
+        Value cNumOutRows =
+          getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), numOutRows);
+        auto rowLoop = buildNormalizedScfFor(
+          rewriter,
+          nestedLoc,
+          c0,
+          cNumOutRows,
+          c1,
+          ValueRange {outputAcc},
+          [&](OpBuilder&, Location rowLoc, Value row, ValueRange rowIterArgs, SmallVectorImpl<Value>& rowYielded) {
+            SmallVector<OpFoldResult> scalarOffsets {row, rewriter.getIndexAttr(0)};
+            SmallVector<OpFoldResult> scalarSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
+            SmallVector<OpFoldResult> unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
+            Value scalar = tensor::ExtractSliceOp::create(
+              rewriter, rowLoc, scalarType, *columnPiece, scalarOffsets, scalarSizes, unitStrides);
+            if (alpha != 1.0f) {
+              Value alphaTensor = createScalarTensorConstant(scalarType, alpha, rewriter, rowLoc);
+              scalar = spatial::SpatVMulOp::create(rewriter, rowLoc, scalarType, scalar, alphaTensor).getResult();
+            }
+            if (biasArg) {
+              Value biasScalar = createBroadcastedBiasScalar(biasArg, biasType, row, column, scalarType, rewriter, rowLoc);
+              if (beta != 1.0f) {
+                Value betaTensor = createScalarTensorConstant(scalarType, beta, rewriter, rowLoc);
+                biasScalar =
+                  spatial::SpatVMulOp::create(rewriter, rowLoc, scalarType, biasScalar, betaTensor).getResult();
+              }
+              scalar = spatial::SpatVAddOp::create(rewriter, rowLoc, scalarType, scalar, biasScalar).getResult();
+            }
+            Value next = tensor::InsertSliceOp::create(rewriter,
+                                                       rowLoc,
+                                                       scalar,
+                                                       rowIterArgs.front(),
+                                                       SmallVector<OpFoldResult> {row, column},
+                                                       scalarSizes,
+                                                       unitStrides);
+            rowYielded.push_back(next);
+            return success();
+          });
+        if (failed(rowLoop))
+          return failure();
+        yielded.push_back(rowLoop->results.front());
         return success();
       });
     if (failed(loop))
@@ -479,7 +529,7 @@ static Value createPartialGroupOffset(Value hSlice,
                                       int64_t kSlice,
                                       int64_t numKSlices,
                                       int64_t numOutRows,
-                                      ConversionPatternRewriter& rewriter,
+                                      PatternRewriter& rewriter,
                                       Location loc) {
   MLIRContext* context = rewriter.getContext();
   AffineExpr d0 = getAffineDimExpr(0, context);
@@ -496,10 +546,12 @@ static Value extractReductionPiece(Value partialPiecesArg,
                                    RankedTensorType pieceType,
                                    int64_t numKSlices,
                                    int64_t numOutRows,
-                                   ConversionPatternRewriter& rewriter,
+                                   int64_t xbarSize,
+                                   PatternRewriter& rewriter,
                                    Location loc) {
   SmallVector<OpFoldResult> unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-  SmallVector<OpFoldResult> pieceSizes {rewriter.getIndexAttr(numOutRows), rewriter.getIndexAttr(1), rewriter.getIndexAttr(crossbarSize.getValue())};
+  SmallVector<OpFoldResult> pieceSizes {
+    rewriter.getIndexAttr(numOutRows), rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarSize)};
   SmallVector<OpFoldResult> pieceOffsets {
     createPartialGroupOffset(hSlice, kSlice, numKSlices, numOutRows, rewriter, loc),
     rewriter.getIndexAttr(0),
@@ -514,13 +566,15 @@ static Value reducePartialPiecesForHSlice(Value partialPiecesArg,
                                           RankedTensorType pieceType,
                                           int64_t numKSlices,
                                           int64_t numOutRows,
-                                          ConversionPatternRewriter& rewriter,
+                                          int64_t xbarSize,
+                                          PatternRewriter& rewriter,
                                           Location loc) {
   SmallVector<Value> activePieces;
   activePieces.reserve(numKSlices);
   for (int64_t kSlice = 0; kSlice < numKSlices; ++kSlice)
     activePieces.push_back(
-      extractReductionPiece(partialPiecesArg, hSlice, kSlice, pieceType, numKSlices, numOutRows, rewriter, loc));
+      extractReductionPiece(
+        partialPiecesArg, hSlice, kSlice, pieceType, numKSlices, numOutRows, xbarSize, rewriter, loc));
 
   while (activePieces.size() > 1) {
     SmallVector<Value> nextPieces;
@@ -543,11 +597,12 @@ static FailureOr<Value> createReductionOutput(Value partialPieces,
                                               RankedTensorType outType,
                                               RankedTensorType paddedOutType,
                                               int64_t numKSlices,
-                                              ConversionPatternRewriter& rewriter,
+                                              int64_t xbarSize,
+                                              PatternRewriter& rewriter,
                                               Location loc) {
   const int64_t numOutRows = outType.getDimSize(0);
-  const int64_t numOutHSlices = ceilIntegerDivide(outType.getDimSize(1), crossbarSize.getValue());
-  auto pieceType = RankedTensorType::get({numOutRows, static_cast<int64_t>(crossbarSize.getValue())},
+  const int64_t numOutHSlices = ceilIntegerDivide(outType.getDimSize(1), xbarSize);
+  auto pieceType = RankedTensorType::get({numOutRows, xbarSize},
                                          partialPiecesType.getElementType());
 
   if (bias && cast<RankedTensorType>(bias.getType()) != paddedOutType)
@@ -559,20 +614,20 @@ static FailureOr<Value> createReductionOutput(Value partialPieces,
   SmallVector<Value> outputSlices;
   outputSlices.reserve(numOutHSlices);
   for (int64_t hSlice = 0; hSlice < numOutHSlices; ++hSlice) {
-    const int64_t columnOffset = hSlice * crossbarSize.getValue();
+    const int64_t columnOffset = hSlice * xbarSize;
     const int64_t columns =
-      std::min(static_cast<int64_t>(crossbarSize.getValue()), outType.getDimSize(1) - columnOffset);
+      std::min(xbarSize, outType.getDimSize(1) - columnOffset);
     auto outputSliceType = RankedTensorType::get({numOutRows, columns}, outType.getElementType());
     auto computeOp = createSpatCompute(
       rewriter, loc, TypeRange {outputSliceType}, {}, inputs, [&](ValueRange blockArgs) -> LogicalResult {
       Value hSliceValue =
         getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), hSlice);
       Value reduced = reducePartialPiecesForHSlice(
-        blockArgs[0], hSliceValue, pieceType, numKSlices, numOutRows, rewriter, loc);
+        blockArgs[0], hSliceValue, pieceType, numKSlices, numOutRows, xbarSize, rewriter, loc);
       if (bias) {
         SmallVector<OpFoldResult> biasOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(columnOffset)};
         SmallVector<OpFoldResult> pieceSizes {rewriter.getIndexAttr(numOutRows),
-                                              rewriter.getIndexAttr(crossbarSize.getValue())};
+                                              rewriter.getIndexAttr(xbarSize)};
         SmallVector<OpFoldResult> unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
         Value biasSlice =
           tensor::ExtractSliceOp::create(rewriter, loc, pieceType, blockArgs[1], biasOffsets, pieceSizes, unitStrides)
@@ -606,191 +661,245 @@ static FailureOr<Value> createReductionOutput(Value partialPieces,
 }
 
 struct GemmToSpatialComputes : OpConversionPattern<ONNXGemmOp> {
-  using OpConversionPattern::OpConversionPattern;
+  explicit GemmToSpatialComputes(MLIRContext* ctx, const spatial::SpatialTargetResources& target)
+      : OpConversionPattern<ONNXGemmOp>(ctx), target(target) {}
 
   LogicalResult matchAndRewrite(ONNXGemmOp gemmOp,
                                 ONNXGemmOpAdaptor gemmOpAdaptor,
                                 ConversionPatternRewriter& rewriter) const override;
+
+  const spatial::SpatialTargetResources& target;
 };
 
 } // namespace
 
-LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
-                                                     ONNXGemmOpAdaptor gemmOpAdaptor,
-                                                     ConversionPatternRewriter& rewriter) const {
-  Location loc = gemmOp.getLoc();
-  Value a = gemmOpAdaptor.getA();
-  Value b = gemmOpAdaptor.getB();
-  Value c = gemmOpAdaptor.getC();
-
+FailureOr<Value> lowerGemmToSpatial(
+    Operation* diagnosticAnchor,
+    Value a,
+    Value b,
+    Value c,
+    RankedTensorType outType,
+    bool transA,
+    bool transB,
+    float alpha,
+    float beta,
+    const spatial::SpatialTargetResources& target,
+    PatternRewriter& rewriter,
+    Location loc) {
   auto aType = dyn_cast<RankedTensorType>(a.getType());
   auto bType = dyn_cast<RankedTensorType>(b.getType());
-  auto outType = dyn_cast<RankedTensorType>(gemmOp.getY().getType());
-  if (!aType || !bType || !outType)
+  if (!diagnosticAnchor || !aType || !bType || !outType)
     return failure();
   if (!aType.hasStaticShape()) {
-    pim::emitUnsupportedStaticShapeDiagnostic(gemmOp, "Gemm input A");
+    pim::emitUnsupportedStaticShapeDiagnostic(diagnosticAnchor, "Gemm input A");
     return failure();
   }
   if (!bType.hasStaticShape()) {
-    pim::emitUnsupportedStaticShapeDiagnostic(gemmOp, "Gemm input B");
+    pim::emitUnsupportedStaticShapeDiagnostic(diagnosticAnchor, "Gemm input B");
     return failure();
   }
   if (!outType.hasStaticShape()) {
-    pim::emitUnsupportedStaticShapeDiagnostic(gemmOp, "Gemm result");
+    pim::emitUnsupportedStaticShapeDiagnostic(diagnosticAnchor, "Gemm result");
     return failure();
   }
   if (aType.getRank() != 2) {
-    pim::emitUnsupportedRankDiagnostic(gemmOp, "Gemm input A", aType.getRank(), {2});
+    pim::emitUnsupportedRankDiagnostic(diagnosticAnchor, "Gemm input A", aType.getRank(), {2});
     return failure();
   }
   if (bType.getRank() != 2) {
-    pim::emitUnsupportedRankDiagnostic(gemmOp, "Gemm input B", bType.getRank(), {2});
+    pim::emitUnsupportedRankDiagnostic(diagnosticAnchor, "Gemm input B", bType.getRank(), {2});
     return failure();
   }
   if (outType.getRank() != 2) {
-    pim::emitUnsupportedRankDiagnostic(gemmOp, "Gemm result", outType.getRank(), {2});
+    pim::emitUnsupportedRankDiagnostic(diagnosticAnchor, "Gemm result", outType.getRank(), {2});
     return failure();
   }
 
-  if (gemmOpAdaptor.getTransA()) {
+  if (transA) {
     auto aShape = aType.getShape();
     auto transposedType = RankedTensorType::get({aShape[1], aShape[0]}, aType.getElementType(), aType.getEncoding());
-    a = ONNXTransposeOp::create(rewriter, loc, transposedType, a, rewriter.getI64ArrayAttr({1, 0})).getResult();
+    a = createLinalgTranspose(a, transposedType, {1, 0}, rewriter, loc);
     aType = transposedType;
   }
 
-  if (gemmOpAdaptor.getTransB()) {
-    auto bShape = bType.getShape();
-    auto transposedType = RankedTensorType::get({bShape[1], bShape[0]}, bType.getElementType(), bType.getEncoding());
-    b = ONNXTransposeOp::create(rewriter, loc, transposedType, b, rewriter.getI64ArrayAttr({1, 0})).getResult();
-    bType = transposedType;
-  }
-
-  const int64_t numOutRows = outType.getDimSize(0);
-  const int64_t numOutCols = outType.getDimSize(1);
-  const int64_t reductionSize = aType.getDimSize(1);
+  ContractionProblem problem;
+  problem.lhsBatchShape = {};
+  problem.rhsBatchShape = {};
+  problem.outputBatchShape = {};
+  problem.lhsBatch = 1;
+  problem.rhsBatch = 1;
+  problem.batch = 1;
+  problem.m = outType.getDimSize(0);
+  problem.k = aType.getDimSize(1);
+  problem.n = outType.getDimSize(1);
+  problem.lhsElementType = aType.getElementType();
+  problem.rhsElementType = bType.getElementType();
+  problem.resultElementType = outType.getElementType();
+  const bool transposeB = transB;
 
   if (!isCompileTimeComputable(b)) {
+    ContractionPlan plan = makeContractionPlan(
+        problem, target, ContractionPlanKind::BatchedDynamicVVD);
     bool hasC = hasGemmBias(c);
-    float alpha = gemmOpAdaptor.getAlpha().convertToFloat();
-    float beta = gemmOpAdaptor.getBeta().convertToFloat();
     RankedTensorType biasType;
     if (hasC) {
       auto cType = dyn_cast<RankedTensorType>(c.getType());
       if (!cType || !cType.hasStaticShape()) {
-        pim::emitUnsupportedStaticShapeDiagnostic(gemmOp, "Gemm bias");
+        pim::emitUnsupportedStaticShapeDiagnostic(diagnosticAnchor, "Gemm bias");
         return failure();
       }
       auto verifiedBiasType = verifyDynamicGemmBiasType(cType, outType);
       if (failed(verifiedBiasType)) {
-        gemmOp.emitOpError("requires Gemm bias C to be broadcastable to the output shape");
+        diagnosticAnchor->emitOpError("requires Gemm bias C to be broadcastable to the output shape");
         return failure();
       }
       biasType = *verifiedBiasType;
     }
 
-    if (aType.getDimSize(0) != numOutRows || bType.getDimSize(0) != reductionSize
-        || bType.getDimSize(1) != numOutCols) {
-      gemmOp.emitOpError("has inconsistent A, B, and output shapes");
+    const int64_t bReductionSize = bType.getDimSize(transposeB ? 1 : 0);
+    const int64_t bOutputColumns = bType.getDimSize(transposeB ? 0 : 1);
+    if (aType.getDimSize(0) != problem.m || bReductionSize != problem.k || bOutputColumns != problem.n) {
+      diagnosticAnchor->emitOpError("has inconsistent A, B, and output shapes");
       return failure();
     }
 
-    const int64_t laneCount64 = numOutRows * numOutCols;
+    const int64_t laneCount64 = plan.laneCount;
     if (laneCount64 > std::numeric_limits<int32_t>::max()) {
-      gemmOp.emitOpError("requires Gemm dynamic batch lane count to fit in i32");
+      diagnosticAnchor->emitOpError("requires Gemm dynamic batch lane count to fit in i32");
       return failure();
     }
 
-    auto scalarPiecesType = spatial::getGraphBatchPhysicalResultType(laneCount64, RankedTensorType::get({1, 1}, outType.getElementType()));
-    auto batchOp = createVvdmulBatch(a, b, aType, bType, scalarPiecesType, outType, rewriter, loc);
+    auto columnType = RankedTensorType::get({problem.m, 1}, outType.getElementType());
+    auto scalarPiecesType = spatial::getGraphBatchPhysicalResultType(problem.n, columnType);
+    auto batchOp = createVvdmulBatch(a, b, aType, bType, scalarPiecesType, outType, transposeB, rewriter, loc);
     if (failed(batchOp))
       return failure();
     auto outputCompute = createDynamicGemmOutputCompute(
       batchOp->getResult(0), hasC ? c : Value(), scalarPiecesType, biasType, outType, alpha, beta, rewriter, loc);
     if (failed(outputCompute))
       return failure();
-    rewriter.replaceOp(gemmOp, outputCompute->getResults());
-    return success();
+    return outputCompute->getResult(0);
   }
 
-  auto scaledB = materializeScaledConstantTensor(b, gemmOpAdaptor.getAlpha().convertToFloat(), rewriter, loc);
+  if (transposeB) {
+    auto bShape = bType.getShape();
+    auto transposedType = RankedTensorType::get({bShape[1], bShape[0]}, bType.getElementType(), bType.getEncoding());
+    if (isCompileTimeComputable(b)) {
+      auto denseAttr = getHostConstDenseElementsAttr(b);
+      auto inputType = denseAttr ? dyn_cast<RankedTensorType>(denseAttr.getType()) : nullptr;
+      auto transposedAttr = inputType && inputType.hasStaticShape() && transposedType.hasStaticShape()
+        ? transposeDenseElementsAttr(denseAttr, {1, 0})
+        : FailureOr<DenseElementsAttr>(failure());
+      if (failed(transposedAttr) || transposedAttr->getType() != transposedType) {
+        diagnosticAnchor->emitOpError("requires Gemm input B transpose to remain statically materializable");
+        return failure();
+      }
+      b = getOrCreateConstant(rewriter,
+                              rewriter.getInsertionBlock()->getParentOp(),
+                              *transposedAttr,
+                              transposedType);
+    } else {
+      b = createLinalgTranspose(b, transposedType, {1, 0}, rewriter, loc);
+    }
+    bType = transposedType;
+  }
+
+  auto scaledB = materializeScaledConstantTensor(b, alpha, rewriter, loc);
   if (failed(scaledB)) {
-    gemmOp.emitOpError("requires constant Gemm input B when alpha is not 1.0");
+    diagnosticAnchor->emitOpError("requires constant Gemm input B when alpha is not 1.0");
     return failure();
   }
   b = *scaledB;
   bType = cast<RankedTensorType>(b.getType());
 
-  if (aType.getDimSize(0) != numOutRows || bType.getDimSize(0) != reductionSize || bType.getDimSize(1) != numOutCols) {
-    gemmOp.emitOpError("has inconsistent A, B, and output shapes after transpose handling");
+  if (aType.getDimSize(0) != problem.m || bType.getDimSize(0) != problem.k || bType.getDimSize(1) != problem.n) {
+    diagnosticAnchor->emitOpError("has inconsistent A, B, and output shapes after transpose handling");
     return failure();
   }
 
-  const int64_t numKSlices = ceilIntegerDivide(reductionSize, crossbarSize.getValue());
-  const int64_t numOutHSlices = ceilIntegerDivide(numOutCols, crossbarSize.getValue());
-  const int64_t paddedReductionSize = numKSlices * static_cast<int64_t>(crossbarSize.getValue());
-  const int64_t paddedOutCols = numOutHSlices * static_cast<int64_t>(crossbarSize.getValue());
+  ContractionPlan plan = makeContractionPlan(
+      problem, target, ContractionPlanKind::StaticTiled);
+  const int64_t xbarSize = plan.tileK;
+  const int64_t numKSlices = plan.reductionSlices;
+  const int64_t numOutHSlices = plan.outputTiles;
+  const int64_t paddedReductionSize = numKSlices * plan.tileK;
+  const int64_t paddedOutCols = numOutHSlices * plan.tileN;
 
   auto paddedBType = RankedTensorType::get({paddedReductionSize, paddedOutCols}, bType.getElementType());
   auto paddedB = materializePaddedConstantMatrix(b, paddedBType, rewriter, loc);
   if (failed(paddedB)) {
-    gemmOp.emitOpError("requires constant Gemm input B so tiled weights can be padded statically");
+    diagnosticAnchor->emitOpError("requires constant Gemm input B so tiled weights can be padded statically");
     return failure();
   }
   b = *paddedB;
-  auto paddedAType = RankedTensorType::get({numOutRows, paddedReductionSize}, aType.getElementType());
+  auto paddedAType = RankedTensorType::get({problem.m, paddedReductionSize}, aType.getElementType());
   a = createPaddedInputCompute(a, paddedAType, rewriter, loc);
   aType = paddedAType;
 
   Value bias;
   bool hasC = hasGemmBias(c);
-  auto paddedOutType = RankedTensorType::get({numOutRows, paddedOutCols}, outType.getElementType());
+  auto paddedOutType = RankedTensorType::get({problem.m, paddedOutCols}, outType.getElementType());
   if (hasC) {
     auto cType = dyn_cast<RankedTensorType>(c.getType());
     if (!cType || !cType.hasStaticShape()) {
-      pim::emitUnsupportedStaticShapeDiagnostic(gemmOp, "Gemm bias");
+      pim::emitUnsupportedStaticShapeDiagnostic(diagnosticAnchor, "Gemm bias");
       return failure();
     }
 
-    auto scaledC = materializeScaledConstantTensor(c, gemmOpAdaptor.getBeta().convertToFloat(), rewriter, loc);
+    auto scaledC = materializeScaledConstantTensor(c, beta, rewriter, loc);
     if (failed(scaledC)) {
-      gemmOp.emitOpError("requires constant Gemm bias C when beta is not 1.0");
+      diagnosticAnchor->emitOpError("requires constant Gemm bias C when beta is not 1.0");
       return failure();
     }
     c = *scaledC;
 
     auto preparedBias = prepareBias(c, outType, paddedOutType, rewriter, loc);
     if (failed(preparedBias)) {
-      gemmOp.emitOpError("requires Gemm bias C to be broadcastable to the output shape");
+      diagnosticAnchor->emitOpError("requires Gemm bias C to be broadcastable to the output shape");
       return failure();
     }
     bias = *preparedBias;
   }
 
-  const int64_t laneCount64 = numOutHSlices * numKSlices * numOutRows;
+  const int64_t laneCount64 = plan.laneCount;
   if (laneCount64 > std::numeric_limits<int32_t>::max()) {
-    gemmOp.emitOpError("requires Gemm tiled batch lane count to fit in i32");
+    diagnosticAnchor->emitOpError("requires Gemm tiled batch lane count to fit in i32");
     return failure();
   }
 
   auto partialPiecesType = spatial::getGraphBatchPhysicalResultType(
-    laneCount64, RankedTensorType::get({1, static_cast<int64_t>(crossbarSize.getValue())}, outType.getElementType()));
+    laneCount64, RankedTensorType::get({1, xbarSize}, outType.getElementType()));
   auto batchOp =
-    createVmmBatch(a, b, aType, paddedBType, partialPiecesType, numOutRows, numKSlices, numOutHSlices, rewriter, loc);
+    createVmmBatch(
+      a, b, aType, paddedBType, partialPiecesType, problem.m, numKSlices, numOutHSlices, xbarSize, rewriter, loc);
   if (failed(batchOp))
     return failure();
   auto reductionOutput = createReductionOutput(
-    batchOp->getResult(0), bias, partialPiecesType, outType, paddedOutType, numKSlices, rewriter, loc);
+    batchOp->getResult(0), bias, partialPiecesType, outType, paddedOutType, numKSlices, xbarSize, rewriter, loc);
   if (failed(reductionOutput))
     return failure();
 
-  rewriter.replaceOp(gemmOp, *reductionOutput);
+  return *reductionOutput;
+}
+
+LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
+                                                     ONNXGemmOpAdaptor gemmOpAdaptor,
+                                                     ConversionPatternRewriter& rewriter) const {
+  FailureOr<Value> result = lowerGemmToSpatial(
+      gemmOp.getOperation(), gemmOpAdaptor.getA(), gemmOpAdaptor.getB(), gemmOpAdaptor.getC(),
+      cast<RankedTensorType>(gemmOp.getY().getType()), gemmOpAdaptor.getTransA(),
+      gemmOpAdaptor.getTransB(), gemmOpAdaptor.getAlpha().convertToFloat(),
+      gemmOpAdaptor.getBeta().convertToFloat(), target, rewriter, gemmOp.getLoc());
+  if (failed(result))
+    return failure();
+  rewriter.replaceOp(gemmOp, *result);
   return success();
 }
 
-void populateGemmPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
-  patterns.insert<GemmToSpatialComputes>(ctx);
+void populateGemmPatterns(RewritePatternSet& patterns,
+                          MLIRContext* ctx,
+                          const spatial::SpatialTargetResources& target) {
+  patterns.insert<GemmToSpatialComputes>(ctx, target);
 }
 
 } // namespace onnx_mlir

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

@@ -0,0 +1,27 @@
+#pragma once
+
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Location.h"
+#include "mlir/IR/Value.h"
+#include "mlir/IR/PatternMatch.h"
+
+namespace onnx_mlir {
+namespace spatial {
+struct SpatialTargetResources;
+}
+
+mlir::FailureOr<mlir::Value> lowerGemmToSpatial(
+    mlir::Operation* diagnosticAnchor,
+    mlir::Value a,
+    mlir::Value b,
+    mlir::Value c,
+    mlir::RankedTensorType outputType,
+    bool transA,
+    bool transB,
+    float alpha,
+    float beta,
+    const spatial::SpatialTargetResources& target,
+    mlir::PatternRewriter& rewriter,
+    mlir::Location loc);
+
+} // namespace onnx_mlir

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

@@ -280,12 +280,12 @@ static FailureOr<Value> buildReduceMeanKeepdimsBlueprint(
   SmallVector<int64_t> fragmentStrides(fragmentOffsets.size(), 1);
   return spatial::SpatBlueprintOp::create(
            rewriter, loc, keepdimsType, batchValue, ValueRange {},
-           rewriter.getStringAttr("nchw"),
-           rewriter.getStringAttr("fragmented"),
+           spatial::getNCHWLayout(rewriter.getContext()),
+           spatial::getFragmentedLayout(rewriter.getContext()),
            rewriter.getDenseI64ArrayAttr(fragmentOffsets),
            rewriter.getDenseI64ArrayAttr(fragmentSizes),
            rewriter.getStringAttr("reduce_mean_keepdims_fragments"),
-           rewriter.getStringAttr("fragment_assembly"),
+           spatial::getFragmentAssemblyMode(rewriter.getContext()),
            rewriter.getDenseI64ArrayAttr(operandIndices),
            rewriter.getDenseI64ArrayAttr(sourceSlots),
            rewriter.getDenseI64ArrayAttr(sourceOffsets),
@@ -370,6 +370,14 @@ struct ReduceMeanToSpatialCompute : OpConversionPattern<ReduceMeanOp> {
     Location loc = reduceMeanOp.getLoc();
     RankedTensorType leafType = getAllOnesType(inputType, resultType.getElementType());
     RankedTensorType keepdimsType = getKeepdimsType(inputType, resultType.getElementType(), reducedAxes);
+    if (semantics->keepdims != 0 && inputType.getRank() == 4
+        && inputType.getDimSize(0) == 1 && semantics->axes == ArrayRef<int64_t>({2, 3})
+        && resultType == keepdimsType) {
+      auto plan = spatial::SpatGlobalAveragePoolPlanOp::create(
+        rewriter, loc, resultType, adaptor.getData(), spatial::getNCHWLayout(rewriter.getContext()));
+      rewriter.replaceOp(reduceMeanOp, plan.getResult());
+      return success();
+    }
     int64_t laneCount = 1;
     for (auto [dim, isReduced] : llvm::zip_equal(keepdimsType.getShape(), reducedAxes)) {
       if (isReduced)

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

@@ -14,10 +14,10 @@
 
 #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/MatrixProductLowering.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Transforms/PlanLowering.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
@@ -32,8 +32,10 @@ static Value materializeTileTensor(PatternRewriter& rewriter, Location loc, Valu
   return insertStaticSlice(rewriter, loc, tile, empty, getZeroOffsets(rewriter, tileType.getRank()));
 }
 
-static Value
-createPoolFillElement(ConversionPatternRewriter& rewriter, Location loc, Type elementType, bool useMinimumValue) {
+static Value createPoolFillElement(OpBuilder& rewriter,
+                                   Location loc,
+                                   Type elementType,
+                                   bool useMinimumValue) {
   Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
   if (!useMinimumValue)
     return getOrCreateConstant(rewriter, anchorOp, rewriter.getZeroAttr(elementType), elementType);
@@ -51,7 +53,7 @@ createPoolFillElement(ConversionPatternRewriter& rewriter, Location loc, Type el
   llvm_unreachable("unsupported pool element type");
 }
 
-static Value createPoolFillTensor(ConversionPatternRewriter& rewriter,
+static Value createPoolFillTensor(OpBuilder& rewriter,
                                   Location loc,
                                   RankedTensorType tensorType,
                                   bool useMinimumValue) {
@@ -59,16 +61,15 @@ static Value createPoolFillTensor(ConversionPatternRewriter& rewriter,
   return tensor::SplatOp::create(rewriter, loc, tensorType, fillElement);
 }
 
-template <typename PoolOp>
-static Value createPaddedPoolInput(ConversionPatternRewriter& rewriter,
+static Value createPaddedPoolInput(OpBuilder& rewriter,
                                    Location loc,
-                                   PoolOp poolOp,
                                    Value input,
                                    RankedTensorType inputType,
                                    int64_t padTop,
                                    int64_t padLeft,
                                    int64_t padBottom,
-                                   int64_t padRight) {
+                                   int64_t padRight,
+                                   bool useMinimumValue) {
   if (padTop == 0 && padLeft == 0 && padBottom == 0 && padRight == 0)
     return input;
 
@@ -90,8 +91,8 @@ static Value createPaddedPoolInput(ConversionPatternRewriter& rewriter,
     padBlock->addArgument(rewriter.getIndexType(), loc);
   padOp.getRegion().push_back(padBlock);
   rewriter.setInsertionPointToStart(padBlock);
-  Value padValue =
-    createPoolFillElement(rewriter, loc, inputType.getElementType(), std::is_same_v<PoolOp, ONNXMaxPoolSingleOutOp>);
+  Value padValue = createPoolFillElement(
+    rewriter, loc, inputType.getElementType(), useMinimumValue);
   tensor::YieldOp::create(rewriter, loc, padValue);
   rewriter.setInsertionPointAfter(padOp);
   return padOp.getResult();
@@ -160,7 +161,10 @@ struct PoolToSpatialCompute;
 
 template <typename PoolOp, typename PoolOpAdaptor, typename ReduceOp>
 struct PoolToSpatialComputeBase : public OpConversionPattern<PoolOp> {
-  using OpConversionPattern<PoolOp>::OpConversionPattern;
+  PoolToSpatialComputeBase(MLIRContext* ctx, const spatial::SpatialTargetResources& target)
+      : OpConversionPattern<PoolOp>(ctx), target(target) {}
+
+  const spatial::SpatialTargetResources& target;
 
   LogicalResult matchAndRewrite(PoolOp poolOp, PoolOpAdaptor adaptor, ConversionPatternRewriter& rewriter) const final {
     Location loc = poolOp.getLoc();
@@ -241,7 +245,7 @@ struct PoolToSpatialComputeBase : public OpConversionPattern<PoolOp> {
           rewriter.getDenseI64ArrayAttr({padTop, padLeft, padBottom, padRight}),
           rewriter.getDenseI64ArrayAttr({strideHeight, strideWidth}),
           rewriter.getDenseI64ArrayAttr({dilationHeight, dilationWidth}),
-          rewriter.getStringAttr("nchw"));
+          spatial::getNCHWLayout(rewriter.getContext()));
         rewriter.replaceOp(poolOp, plan.getResult());
         return success();
       }
@@ -251,12 +255,12 @@ struct PoolToSpatialComputeBase : public OpConversionPattern<PoolOp> {
              && dilationHeight == 1 && dilationWidth == 1 && padTop == 0
              && padLeft == 0 && padBottom == 0 && padRight == 0) {
       auto plan = spatial::SpatGlobalAveragePoolPlanOp::create(
-        rewriter, loc, outType, x, rewriter.getStringAttr("nchw"));
+        rewriter, loc, outType, x, spatial::getNCHWLayout(rewriter.getContext()));
       rewriter.replaceOp(poolOp, plan.getResult());
       return success();
     }
 
-    const int64_t xbarSize = static_cast<int64_t>(crossbarSize.getValue());
+    const int64_t xbarSize = static_cast<int64_t>(target.matrixShape.rows);
     const int64_t channelTileCount = (channels + xbarSize - 1) / xbarSize;
     const int64_t outputPatchCount = batchSize * outputHeight * outputWidth;
     const bool countIncludePad = [&]() {
@@ -292,7 +296,9 @@ struct PoolToSpatialComputeBase : public OpConversionPattern<PoolOp> {
     auto computeOp =
       createSpatCompute<numInputs>(rewriter, loc, outType, {}, ValueRange {x}, [&](Value xArg) -> LogicalResult {
         Value paddedInput =
-          createPaddedPoolInput(rewriter, loc, poolOp, xArg, xType, padTop, padLeft, padBottom, padRight);
+          createPaddedPoolInput(rewriter, loc, xArg, xType, padTop, padLeft,
+                                padBottom, padRight,
+                                std::is_same_v<PoolOp, ONNXMaxPoolSingleOutOp>);
         Value pooledOutputInit = tensor::EmptyOp::create(rewriter, loc, outType.getShape(), outType.getElementType());
 
         Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
@@ -424,7 +430,8 @@ struct PoolToSpatialCompute<ONNXAveragePoolOp>
 
 } // namespace
 
-LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp) {
+LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp,
+                                            const spatial::SpatialTargetResources&) {
   auto inputType = dyn_cast<RankedTensorType>(planOp.getInput().getType());
   auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
   if (!inputType || !outputType || !inputType.hasStaticShape() || !outputType.hasStaticShape())
@@ -439,6 +446,119 @@ LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp)
   return success();
 }
 
+FailureOr<Value> lowerDenseMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
+                                         Value input,
+                                         const spatial::SpatialTargetResources& target,
+                                         PatternRewriter& rewriter) {
+  auto inputType = dyn_cast<RankedTensorType>(input.getType());
+  auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
+  if (!inputType || !outputType || !inputType.hasStaticShape() || !outputType.hasStaticShape()
+      || inputType.getRank() != 4 || outputType.getRank() != 4)
+    return planOp.emitOpError("dense MaxPool lowering requires static rank-4 tensors"), failure();
+
+  auto kernel = planOp.getKernelShape();
+  auto pads = planOp.getPads();
+  auto strides = planOp.getStrides();
+  auto dilations = planOp.getDilations();
+  if (kernel.size() != 2 || pads.size() != 4 || strides.size() != 2 || dilations.size() != 2
+      || llvm::any_of(kernel, [](int64_t value) { return value <= 0; })
+      || llvm::any_of(strides, [](int64_t value) { return value <= 0; })
+      || llvm::any_of(dilations, [](int64_t value) { return value <= 0; })
+      || llvm::any_of(pads, [](int64_t value) { return value < 0; }))
+    return planOp.emitOpError("dense MaxPool lowering requires valid kernel, padding, stride, and dilation attributes"),
+           failure();
+
+  const int64_t batchSize = inputType.getDimSize(0);
+  const int64_t channels = inputType.getDimSize(1);
+  const int64_t outputHeight = outputType.getDimSize(2);
+  const int64_t outputWidth = outputType.getDimSize(3);
+  const int64_t tileWidth = std::max<int64_t>(1, target.matrixShape.rows);
+  const int64_t channelTileCount = (channels + tileWidth - 1) / tileWidth;
+  const int64_t outputPatchCount = batchSize * outputHeight * outputWidth;
+
+  auto compute = createSpatCompute<1>(
+    rewriter, planOp.getLoc(), outputType, {}, input,
+    [&](Value input) -> LogicalResult {
+      Value paddedInput = createPaddedPoolInput(
+        rewriter, planOp.getLoc(), input, inputType,
+        pads[0], pads[1], pads[2], pads[3], /*useMinimumValue=*/true);
+      Value outputInit = tensor::EmptyOp::create(
+        rewriter, planOp.getLoc(), outputType.getShape(), outputType.getElementType());
+      Operation* anchor = rewriter.getInsertionBlock()->getParentOp();
+      Value zero = getOrCreateIndexConstant(rewriter, anchor, 0);
+      Value one = getOrCreateIndexConstant(rewriter, anchor, 1);
+      Value patchCount = getOrCreateIndexConstant(rewriter, anchor, outputPatchCount);
+      Value pixelsPerBatch = getOrCreateIndexConstant(
+        rewriter, anchor, outputHeight * outputWidth);
+      Value outputWidthValue = getOrCreateIndexConstant(rewriter, anchor, outputWidth);
+      Value strideHeight = getOrCreateIndexConstant(rewriter, anchor, strides[0]);
+      Value strideWidth = getOrCreateIndexConstant(rewriter, anchor, strides[1]);
+
+      auto loop = buildNormalizedScfFor(
+        rewriter, planOp.getLoc(), zero, patchCount, one, ValueRange {outputInit},
+        [&](OpBuilder&, Location loc, Value patch, ValueRange iterArgs,
+            SmallVectorImpl<Value>& yielded) {
+          Value batch = arith::DivUIOp::create(rewriter, loc, patch, pixelsPerBatch);
+          Value batchPatch = arith::RemUIOp::create(rewriter, loc, patch, pixelsPerBatch);
+          Value outputRow = arith::DivUIOp::create(rewriter, loc, batchPatch, outputWidthValue);
+          Value outputColumn = arith::RemUIOp::create(rewriter, loc, batchPatch, outputWidthValue);
+          Value windowRow = arith::MulIOp::create(rewriter, loc, outputRow, strideHeight);
+          Value windowColumn = arith::MulIOp::create(rewriter, loc, outputColumn, strideWidth);
+          Value updated = iterArgs.front();
+
+          for (int64_t tile = 0; tile < channelTileCount; ++tile) {
+            const int64_t tileChannels = std::min<int64_t>(tileWidth, channels - tile * tileWidth);
+            auto tileType = RankedTensorType::get(
+              {1, tileChannels, 1, 1}, outputType.getElementType());
+            Value reduced = createPoolFillTensor(
+              rewriter, loc, tileType, /*useMinimumValue=*/true);
+            for (int64_t kernelRow = 0; kernelRow < kernel[0]; ++kernelRow) {
+              Value sourceRow = windowRow;
+              if (kernelRow * dilations[0] != 0)
+                sourceRow = arith::AddIOp::create(
+                  rewriter, loc, sourceRow,
+                  getOrCreateIndexConstant(rewriter, anchor, kernelRow * dilations[0]));
+              for (int64_t kernelColumn = 0; kernelColumn < kernel[1]; ++kernelColumn) {
+                Value sourceColumn = windowColumn;
+                if (kernelColumn * dilations[1] != 0)
+                  sourceColumn = arith::AddIOp::create(
+                    rewriter, loc, sourceColumn,
+                    getOrCreateIndexConstant(rewriter, anchor, kernelColumn * dilations[1]));
+                Value point = tensor::ExtractSliceOp::create(
+                  rewriter, loc, tileType, paddedInput,
+                  SmallVector<OpFoldResult> {
+                    batch, rewriter.getIndexAttr(tile * tileWidth), sourceRow, sourceColumn},
+                  SmallVector<OpFoldResult> {
+                    rewriter.getIndexAttr(1), rewriter.getIndexAttr(tileChannels),
+                    rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)},
+                  getUnitStrides(rewriter, 4));
+                point = materializeTileTensor(rewriter, loc, point);
+                reduced = spatial::SpatVMaxOp::create(
+                  rewriter, loc, tileType, reduced, point);
+              }
+            }
+            updated = tensor::InsertSliceOp::create(
+              rewriter, loc, reduced, updated,
+              SmallVector<OpFoldResult> {
+                batch, rewriter.getIndexAttr(tile * tileWidth), outputRow, outputColumn},
+              SmallVector<OpFoldResult> {
+                rewriter.getIndexAttr(1), rewriter.getIndexAttr(tileChannels),
+                rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)},
+              getUnitStrides(rewriter, 4));
+          }
+          yielded.push_back(updated);
+          return success();
+        });
+      if (failed(loop))
+        return failure();
+      spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), loop->results.front());
+      return success();
+    });
+  if (failed(compute))
+    return failure();
+  return compute->getResult(0);
+}
+
 static Value createClampedPoolIndexTable(PatternRewriter& rewriter,
                                          Operation* anchorOp,
                                          int64_t outputSize,
@@ -496,13 +616,15 @@ static Value extractPoolIndex(PatternRewriter& rewriter,
 }
 
 FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
+                                            Value input,
                                             std::optional<Value> rowStripInput,
+                                            const spatial::SpatialTargetResources& target,
                                             PatternRewriter& rewriter) {
-  if (failed(canLowerMaxPoolPlanToRowStrip(planOp)))
+  if (failed(canLowerMaxPoolPlanToRowStrip(planOp, target)))
     return failure();
 
   Location loc = planOp.getLoc();
-  auto inputType = cast<RankedTensorType>(planOp.getInput().getType());
+  auto inputType = cast<RankedTensorType>(input.getType());
   auto outputType = cast<RankedTensorType>(planOp.getOutput().getType());
   const int64_t channels = inputType.getDimSize(1);
   const int64_t inputHeight = inputType.getDimSize(2);
@@ -511,9 +633,9 @@ FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
   const int64_t outputWidth = outputType.getDimSize(3);
   const int64_t kernelHeight = planOp.getKernelShape()[0];
   const int64_t kernelWidth = planOp.getKernelShape()[1];
-  Value input = rowStripInput.value_or(planOp.getInput());
-  auto actualInputType = dyn_cast<RankedTensorType>(input.getType());
-  FailureOr<RowStripPhysicalValue> physicalValue = describeRowStripPhysicalValue(input, inputType);
+  Value actualInput = rowStripInput.value_or(input);
+  auto actualInputType = dyn_cast<RankedTensorType>(actualInput.getType());
+  FailureOr<RowStripPhysicalValue> physicalValue = describeRowStripPhysicalValue(actualInput, inputType);
   const bool physicalInput = succeeded(physicalValue);
   if (!physicalInput && actualInputType != inputType)
     return failure();
@@ -561,7 +683,7 @@ FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
     TypeRange {outputStorageType},
     outputHeight * tilesPerRow,
     {},
-    ValueRange {input},
+    ValueRange {actualInput},
     [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult {
       SmallVector<Value> inputRows;
       inputRows.reserve(kernelHeight);
@@ -590,8 +712,8 @@ FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
                                                                                 rewriter.getIndexAttr(1),
                                                                                 rewriter.getIndexAttr(inputWidth)},
                                                      getUnitStrides(rewriter, 4));
-          inputRows.push_back(ONNXTransposeOp::create(
-            rewriter, loc, inputFragmentType, nchw, rewriter.getI64ArrayAttr({0, 2, 3, 1})));
+          inputRows.push_back(createLinalgTranspose(
+            nchw, inputFragmentType, {0, 2, 3, 1}, rewriter, loc));
         }
       }
 
@@ -685,7 +807,8 @@ FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
   return batch->getResult(0);
 }
 
-LogicalResult canLowerGlobalAveragePoolPlanToRowStrip(spatial::SpatGlobalAveragePoolPlanOp planOp) {
+LogicalResult canLowerGlobalAveragePoolPlanToRowStrip(
+  spatial::SpatGlobalAveragePoolPlanOp planOp, const spatial::SpatialTargetResources&) {
   auto inputType = dyn_cast<RankedTensorType>(planOp.getInput().getType());
   auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
   if (!inputType || !outputType || !inputType.hasStaticShape() || !outputType.hasStaticShape())
@@ -697,22 +820,101 @@ LogicalResult canLowerGlobalAveragePoolPlanToRowStrip(spatial::SpatGlobalAverage
   return success();
 }
 
+FailureOr<Value> lowerDenseGlobalAveragePoolPlan(
+    spatial::SpatGlobalAveragePoolPlanOp planOp,
+    Value input,
+    const spatial::SpatialTargetResources& target,
+    PatternRewriter& rewriter) {
+  auto inputType = dyn_cast<RankedTensorType>(input.getType());
+  auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
+  if (!inputType || !outputType || !inputType.hasStaticShape()
+      || !outputType.hasStaticShape() || inputType.getRank() != 4
+      || outputType.getRank() != 4 || inputType.getDimSize(0) != 1
+      || outputType.getDimSize(0) != 1 || inputType.getDimSize(1) != outputType.getDimSize(1)
+      || outputType.getDimSize(2) != 1 || outputType.getDimSize(3) != 1)
+    return planOp.emitOpError("dense global AveragePool lowering requires static rank-4 floating-point tensors"),
+           failure();
+  auto elementType = dyn_cast<FloatType>(inputType.getElementType());
+  if (!elementType)
+    return planOp.emitOpError("dense global AveragePool lowering requires floating-point tensors"),
+           failure();
+
+  const int64_t channels = inputType.getDimSize(1);
+  const int64_t height = inputType.getDimSize(2);
+  const int64_t width = inputType.getDimSize(3);
+  const int64_t tileWidth = std::max<int64_t>(1, target.matrixShape.rows);
+  const int64_t channelTileCount = (channels + tileWidth - 1) / tileWidth;
+  const double scaleValue = 1.0 / static_cast<double>(height * width);
+
+  auto compute = createSpatCompute<1>(
+    rewriter, planOp.getLoc(), outputType, {}, input,
+    [&](Value input) -> LogicalResult {
+      Value output = tensor::EmptyOp::create(
+        rewriter, planOp.getLoc(), outputType.getShape(), outputType.getElementType());
+      Operation* anchor = rewriter.getInsertionBlock()->getParentOp();
+      for (int64_t tile = 0; tile < channelTileCount; ++tile) {
+        const int64_t tileChannels = std::min<int64_t>(tileWidth, channels - tile * tileWidth);
+        auto tileType = RankedTensorType::get(
+          {1, tileChannels, 1, 1}, outputType.getElementType());
+        Value reduced = createPoolFillTensor(
+          rewriter, planOp.getLoc(), tileType, /*useMinimumValue=*/false);
+        for (int64_t row = 0; row < height; ++row) {
+          for (int64_t column = 0; column < width; ++column) {
+            Value point = tensor::ExtractSliceOp::create(
+              rewriter, planOp.getLoc(), tileType, input,
+              SmallVector<OpFoldResult> {
+                rewriter.getIndexAttr(0), rewriter.getIndexAttr(tile * tileWidth),
+                rewriter.getIndexAttr(row), rewriter.getIndexAttr(column)},
+              SmallVector<OpFoldResult> {
+                rewriter.getIndexAttr(1), rewriter.getIndexAttr(tileChannels),
+                rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)},
+              getUnitStrides(rewriter, 4));
+            point = materializeTileTensor(rewriter, planOp.getLoc(), point);
+            reduced = spatial::SpatVAddOp::create(
+              rewriter, planOp.getLoc(), tileType, reduced, point);
+          }
+        }
+        auto scaleAttr = DenseElementsAttr::get(
+          tileType, rewriter.getFloatAttr(elementType, scaleValue));
+        Value scale = getOrCreateConstant(rewriter, anchor, scaleAttr, tileType);
+        reduced = spatial::SpatVMulOp::create(
+          rewriter, planOp.getLoc(), tileType, reduced, scale);
+        output = tensor::InsertSliceOp::create(
+          rewriter, planOp.getLoc(), reduced, output,
+          SmallVector<OpFoldResult> {
+            rewriter.getIndexAttr(0), rewriter.getIndexAttr(tile * tileWidth),
+            rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
+          SmallVector<OpFoldResult> {
+            rewriter.getIndexAttr(1), rewriter.getIndexAttr(tileChannels),
+            rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)},
+          getUnitStrides(rewriter, 4));
+      }
+      spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), output);
+      return success();
+    });
+  if (failed(compute))
+    return failure();
+  return compute->getResult(0);
+}
+
 FailureOr<Value> lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp,
+                                                    Value input,
                                                     std::optional<Value> rowStripInput,
+                                                    const spatial::SpatialTargetResources& target,
                                                     PatternRewriter& rewriter) {
-  if (failed(canLowerGlobalAveragePoolPlanToRowStrip(planOp)))
+  if (failed(canLowerGlobalAveragePoolPlanToRowStrip(planOp, target)))
     return failure();
 
   Location loc = planOp.getLoc();
-  auto inputType = cast<RankedTensorType>(planOp.getInput().getType());
+  auto inputType = cast<RankedTensorType>(input.getType());
   auto outputType = cast<RankedTensorType>(planOp.getOutput().getType());
   auto elementType = dyn_cast<FloatType>(inputType.getElementType());
   if (!elementType)
     return failure();
 
-  Value input = rowStripInput.value_or(planOp.getInput());
-  auto actualInputType = dyn_cast<RankedTensorType>(input.getType());
-  FailureOr<RowStripPhysicalValue> physicalValue = describeRowStripPhysicalValue(input, inputType);
+  Value actualInput = rowStripInput.value_or(input);
+  auto actualInputType = dyn_cast<RankedTensorType>(actualInput.getType());
+  FailureOr<RowStripPhysicalValue> physicalValue = describeRowStripPhysicalValue(actualInput, inputType);
   const bool physicalInput = succeeded(physicalValue);
   if (!physicalInput && actualInputType != inputType)
     return failure();
@@ -742,7 +944,7 @@ FailureOr<Value> lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePo
     TypeRange {outputStorageType},
     tilesPerRow,
     ValueRange {zero, scale},
-    ValueRange {input},
+    ValueRange {actualInput},
     [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult {
       Value reduced = args.weights[0];
       for (int64_t row = 0; row < height; ++row) {
@@ -777,8 +979,8 @@ FailureOr<Value> lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePo
                                       rewriter.getIndexAttr(1),
                                       rewriter.getIndexAttr(width)},
             getUnitStrides(rewriter, 4));
-          fragment = ONNXTransposeOp::create(
-            rewriter, loc, inputFragmentType, nchw, rewriter.getI64ArrayAttr({0, 2, 3, 1}));
+          fragment = createLinalgTranspose(
+            nchw, inputFragmentType, {0, 2, 3, 1}, rewriter, loc);
         }
         for (int64_t column = 0; column < width; ++column) {
           Value point = tensor::ExtractSliceOp::create(
@@ -811,9 +1013,11 @@ FailureOr<Value> lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePo
   return batch->getResult(0);
 }
 
-void populatePoolPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
-  patterns.insert<PoolToSpatialCompute<ONNXMaxPoolSingleOutOp>>(ctx);
-  patterns.insert<PoolToSpatialCompute<ONNXAveragePoolOp>>(ctx);
+void populatePoolPatterns(RewritePatternSet& patterns,
+                          MLIRContext* ctx,
+                          const spatial::SpatialTargetResources& target) {
+  patterns.insert<PoolToSpatialCompute<ONNXMaxPoolSingleOutOp>>(ctx, target);
+  patterns.insert<PoolToSpatialCompute<ONNXAveragePoolOp>>(ctx, target);
 }
 
 } // namespace onnx_mlir

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

@@ -17,7 +17,7 @@ struct ReluToSpatialCompute : OpConversionPattern<ONNXReluOp> {
     Location loc = reluOp.getLoc();
     Type resultType = reluOp.getResult().getType();
     auto reluPlan = spatial::SpatReluPlanOp::create(
-      rewriter, loc, resultType, adaptor.getX(), rewriter.getStringAttr("nchw"));
+      rewriter, loc, resultType, adaptor.getX(), spatial::getNCHWLayout(rewriter.getContext()));
     rewriter.replaceOp(reluOp, reluPlan.getResult());
     return success();
   }

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

@@ -32,7 +32,8 @@ struct Concat : public OpConversionPattern<ONNXConcatOp> {
              return type && type.hasStaticShape() && type.getRank() == 4;
            })) {
       rewriter.replaceOpWithNewOp<spatial::SpatConcatPlanOp>(
-        maxpoolOp, resultType, inputs, rewriter.getI64IntegerAttr(axis), rewriter.getStringAttr("nchw"));
+        maxpoolOp, resultType, inputs, rewriter.getI64IntegerAttr(axis),
+        spatial::getNCHWLayout(rewriter.getContext()));
       return success();
     }
 

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

@@ -4,11 +4,10 @@
 #include "llvm/ADT/SmallVector.h"
 
 #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
-#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Transforms/PlanLowering.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -48,12 +47,12 @@ static SmallVector<ReassociationIndices> getExpandFrom1DReassociation(int64_t ra
   return reassociation;
 }
 
-static Value buildFlatten(Value input,
-                          RankedTensorType sourceType,
-                          RankedTensorType resultType,
-                          int64_t axis,
-                          ConversionPatternRewriter& rewriter,
-                          Location loc) {
+static Value buildFlattenBody(Value input,
+                              RankedTensorType sourceType,
+                              RankedTensorType resultType,
+                              int64_t axis,
+                              PatternRewriter& rewriter,
+                              Location loc) {
   if (sourceType == resultType)
     return input;
 
@@ -76,6 +75,25 @@ static Value buildFlatten(Value input,
     rewriter, loc, resultType, flattened, getExpandFrom1DReassociation(resultType.getRank()));
 }
 
+static Value buildFlatten(Value input,
+                          RankedTensorType sourceType,
+                          RankedTensorType resultType,
+                          int64_t axis,
+                          PatternRewriter& rewriter,
+                          Location loc) {
+  if (spatial::isAnySpatialComputeLike(rewriter.getInsertionBlock()->getParentOp()))
+    return buildFlattenBody(input, sourceType, resultType, axis, rewriter, loc);
+
+  auto compute = createSpatCompute<1>(
+      rewriter, loc, TypeRange {resultType}, {}, ValueRange {input},
+      [&](Value computeInput) {
+        spatial::SpatYieldOp::create(
+            rewriter, loc,
+            buildFlattenBody(computeInput, sourceType, resultType, axis, rewriter, loc));
+      });
+  return compute.getResult(0);
+}
+
 struct Flatten : OpConversionPattern<ONNXFlattenOp> {
   using OpConversionPattern::OpConversionPattern;
 
@@ -98,53 +116,53 @@ struct Flatten : OpConversionPattern<ONNXFlattenOp> {
     if (resultType.getShape()[0] != outerDim || resultType.getShape()[1] != innerDim)
       return failure();
 
-    auto replaceWithFlatten = [&](auto build) -> LogicalResult {
-      Value flattened = materializeOrComputeUnary(adaptor.getInput(), resultType, rewriter, flattenOp.getLoc(), build);
-      rewriter.replaceOp(flattenOp, flattened);
-      return success();
-    };
-
-    return replaceWithFlatten([&](Value input) {
-      return buildFlatten(input, sourceType, resultType, *axis, rewriter, flattenOp.getLoc());
-    });
+    auto plan = spatial::SpatFlattenPlanOp::create(
+        rewriter, flattenOp.getLoc(), resultType, adaptor.getInput(),
+        rewriter.getI64IntegerAttr(*axis),
+        spatial::getNCHWLayout(rewriter.getContext()));
+    rewriter.replaceOp(flattenOp, plan.getOutput());
+    return success();
   }
 };
 
 struct RowStripFlattenAnalysis {
   spatial::SpatGraphComputeBatch consumer;
-  tensor::CollapseShapeOp collapse;
   RankedTensorType sourceType;
   RankedTensorType resultType;
   RankedTensorType weightType;
   DenseElementsAttr weight;
 };
 
-static FailureOr<RowStripFlattenAnalysis> analyzeRowStripFlatten(spatial::SpatGraphCompute flattenOp) {
-  if (flattenOp.getWeights().size() != 0 || flattenOp.getInputs().size() != 1
-      || flattenOp.getOutputs().size() != 1)
+static FailureOr<RowStripFlattenAnalysis> analyzeRowStripFlatten(
+    spatial::SpatFlattenPlanOp flattenOp, const spatial::SpatialTargetResources& target) {
+  if (flattenOp.getAxis() != 1)
     return failure();
-  auto sourceType = dyn_cast<RankedTensorType>(flattenOp.getInputs().front().getType());
-  auto resultType = dyn_cast<RankedTensorType>(flattenOp.getOutputs().front().getType());
+  auto sourceType = dyn_cast<RankedTensorType>(flattenOp.getInput().getType());
+  auto resultType = dyn_cast<RankedTensorType>(flattenOp.getOutput().getType());
   if (!sourceType || !resultType || !sourceType.hasStaticShape() || !resultType.hasStaticShape()
       || sourceType.getRank() != 4 || resultType.getRank() != 2 || sourceType.getDimSize(0) != 1
       || resultType.getDimSize(0) != 1 || resultType.getDimSize(1) != sourceType.getNumElements())
     return failure();
   const int64_t channels = sourceType.getDimSize(1);
-  const int64_t xbarDim = static_cast<int64_t>(crossbarSize.getValue());
+  const int64_t xbarDim = static_cast<int64_t>(target.matrixShape.rows);
   if (channels > xbarDim && channels % xbarDim != 0)
     return failure();
 
-  auto yieldOp = dyn_cast<spatial::SpatYieldOp>(flattenOp.getBody().front().getTerminator());
-  if (!yieldOp || yieldOp.getOutputs().size() != 1)
+  Value consumerInput = flattenOp.getOutput();
+  Operation* consumerOp = nullptr;
+  while (consumerInput.hasOneUse()) {
+    Operation* user = *consumerInput.getUsers().begin();
+    if (auto materialize = dyn_cast<spatial::SpatMaterializeLayoutOp>(user)) {
+      consumerInput = materialize.getOutput();
+      continue;
+    }
+    consumerOp = user;
+    break;
+  }
+  if (!consumerOp)
     return failure();
-  auto collapse = yieldOp.getOutputs().front().getDefiningOp<tensor::CollapseShapeOp>();
-  if (!collapse || collapse.getSrc() != *flattenOp.getInputArgument(0))
-    return failure();
-
-  if (!flattenOp.getResult(0).hasOneUse())
-    return failure();
-  auto consumer = dyn_cast<spatial::SpatGraphComputeBatch>(*flattenOp.getResult(0).getUsers().begin());
-  if (!consumer || consumer.getInputs().size() != 1 || consumer.getInputs().front() != flattenOp.getResult(0)
+  auto consumer = dyn_cast<spatial::SpatGraphComputeBatch>(consumerOp);
+  if (!consumer || consumer.getInputs().size() != 1 || consumer.getInputs().front() != consumerInput
       || consumer.getWeights().size() != 1)
     return failure();
   auto weightType = dyn_cast<RankedTensorType>(consumer.getWeights().front().getType());
@@ -155,21 +173,34 @@ static FailureOr<RowStripFlattenAnalysis> analyzeRowStripFlatten(spatial::SpatGr
   if (llvm::none_of(consumer.getBody().getOps<spatial::SpatVMMOp>(),
                     [](spatial::SpatVMMOp) { return true; }))
     return failure();
-  return RowStripFlattenAnalysis {consumer, collapse, sourceType, resultType, weightType, weight};
+  return RowStripFlattenAnalysis {consumer, sourceType, resultType, weightType, weight};
 }
 
 } // namespace
 
 void populateFlattenPatterns(RewritePatternSet& patterns, MLIRContext* ctx) { patterns.add<Flatten>(ctx); }
 
-LogicalResult canLowerFlattenFromRowStrip(spatial::SpatGraphCompute flattenOp) {
-  return succeeded(analyzeRowStripFlatten(flattenOp)) ? success() : failure();
+FailureOr<Value> lowerDenseFlattenPlan(spatial::SpatFlattenPlanOp planOp,
+                                        Value input,
+                                        PatternRewriter& rewriter) {
+  auto sourceType = dyn_cast<RankedTensorType>(input.getType());
+  auto resultType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
+  if (!sourceType || !resultType || !sourceType.hasStaticShape() || !resultType.hasStaticShape())
+    return failure();
+  return buildFlatten(input, sourceType, resultType, planOp.getAxis(), rewriter,
+                      planOp.getLoc());
+}
+
+LogicalResult canLowerFlattenFromRowStrip(spatial::SpatFlattenPlanOp flattenOp,
+                                          const spatial::SpatialTargetResources& target) {
+  return succeeded(analyzeRowStripFlatten(flattenOp, target)) ? success() : failure();
 }
 
 LogicalResult lowerFlattenFromRowStrip(const RowStripPhysicalValue& input,
-                                       spatial::SpatGraphCompute flattenOp,
+                                       spatial::SpatFlattenPlanOp flattenOp,
+                                       const spatial::SpatialTargetResources& target,
                                        PatternRewriter& rewriter) {
-  FailureOr<RowStripFlattenAnalysis> analysis = analyzeRowStripFlatten(flattenOp);
+  FailureOr<RowStripFlattenAnalysis> analysis = analyzeRowStripFlatten(flattenOp, target);
   if (failed(analysis))
     return failure();
   auto storageType = dyn_cast<RankedTensorType>(input.storage.getType());
@@ -204,19 +235,20 @@ LogicalResult lowerFlattenFromRowStrip(const RowStripPhysicalValue& input,
                                               analysis->weightType);
   analysis->consumer->setOperand(0, reorderedWeight);
 
-  BlockArgument flattenInput = *flattenOp.getInputArgument(0);
-  flattenOp.getInputsMutable().assign(input.storage);
-  flattenInput.setType(storageType);
-
-  OpBuilder::InsertionGuard guard(rewriter);
-  rewriter.setInsertionPoint(analysis->collapse);
-  auto flatType = RankedTensorType::get(
-    {storageType.getNumElements()}, storageType.getElementType(), storageType.getEncoding());
-  Value flat = tensor::CollapseShapeOp::create(
-    rewriter, flattenOp.getLoc(), flatType, flattenInput, getCollapseTo1DReassociation(storageType.getRank()));
-  Value logicalInput = tensor::ExpandShapeOp::create(
-    rewriter, flattenOp.getLoc(), analysis->resultType, flat, getExpandFrom1DReassociation(2));
-  rewriter.replaceOp(analysis->collapse, logicalInput);
+  auto compute = createSpatCompute<1>(
+      rewriter, flattenOp.getLoc(), TypeRange {analysis->resultType}, {},
+      ValueRange {input.storage}, [&](Value storage) {
+        auto flatType = RankedTensorType::get(
+            {storageType.getNumElements()}, storageType.getElementType(), storageType.getEncoding());
+        Value flat = tensor::CollapseShapeOp::create(
+            rewriter, flattenOp.getLoc(), flatType, storage,
+            getCollapseTo1DReassociation(storageType.getRank()));
+        Value logicalInput = tensor::ExpandShapeOp::create(
+            rewriter, flattenOp.getLoc(), analysis->resultType, flat,
+            getExpandFrom1DReassociation(2));
+        spatial::SpatYieldOp::create(rewriter, flattenOp.getLoc(), logicalInput);
+      });
+  rewriter.replaceOp(flattenOp, compute.getResult(0));
   return success();
 }
 

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

@@ -5,8 +5,11 @@
 
 #include "llvm/ADT/STLExtras.h"
 
+#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Transforms/PlanLowering.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -17,126 +20,144 @@ namespace onnx_mlir {
 namespace {
 
 static Value buildNearestAsymmetricIndex(
-  Value outputIndex, int64_t inputDim, int64_t outputDim, ConversionPatternRewriter& rewriter, Location loc) {
+  Value outputIndex, int64_t inputDim, int64_t outputDim, PatternRewriter& rewriter, Location loc) {
+  if (inputDim == outputDim)
+    return outputIndex;
   Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+  if (outputDim % inputDim == 0)
+    return affineFloorDivConst(rewriter, loc, outputIndex, outputDim / inputDim, anchorOp);
+  if (inputDim % outputDim == 0)
+    return affineMulConst(rewriter, loc, outputIndex, inputDim / outputDim, anchorOp);
   Value cInputDim = getOrCreateIndexConstant(rewriter, anchorOp, inputDim);
   Value cOutputDim = getOrCreateIndexConstant(rewriter, anchorOp, outputDim);
-  Value cInputDimLast = getOrCreateIndexConstant(rewriter, anchorOp, inputDim - 1);
   Value scaledIndex = arith::MulIOp::create(rewriter, loc, outputIndex, cInputDim);
-  Value inputIndex = arith::DivUIOp::create(rewriter, loc, scaledIndex, cOutputDim);
-  return arith::MinUIOp::create(rewriter, loc, inputIndex, cInputDimLast);
+  return arith::DivUIOp::create(rewriter, loc, scaledIndex, cOutputDim);
 }
 
-static FailureOr<Value> buildNearestResizeLoop(Value input,
-                                               RankedTensorType inputType,
-                                               RankedTensorType resultType,
-                                               ConversionPatternRewriter& rewriter,
-                                               Location loc) {
-  auto elemType = resultType.getElementType();
-  SmallVector<int64_t> unitShape(resultType.getRank(), 1);
-  auto unitTensorType = RankedTensorType::get(unitShape, elemType);
-
-  SmallVector<OpFoldResult> unitSizes(resultType.getRank(), rewriter.getIndexAttr(1));
-  SmallVector<OpFoldResult> unitStrides(resultType.getRank(), rewriter.getIndexAttr(1));
-
-  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
-  Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
-  Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
-  Value cOutputN = getOrCreateIndexConstant(rewriter, anchorOp, resultType.getDimSize(0));
-  Value cOutputC = getOrCreateIndexConstant(rewriter, anchorOp, resultType.getDimSize(1));
-  Value cOutputH = getOrCreateIndexConstant(rewriter, anchorOp, resultType.getDimSize(2));
-  Value cOutputW = getOrCreateIndexConstant(rewriter, anchorOp, resultType.getDimSize(3));
-
-  Value outputInit = tensor::EmptyOp::create(rewriter, loc, resultType.getShape(), elemType);
-
-  auto batchLoop = buildNormalizedScfFor(
-    rewriter,
-    loc,
-    c0,
-    cOutputN,
-    c1,
-    ValueRange {outputInit},
-    [&](OpBuilder&, Location nestedLoc, Value outputN, ValueRange batchIterArgs, SmallVectorImpl<Value>& batchYielded) {
-      Value outputBatchAcc = batchIterArgs.front();
-      Value inputN =
-        buildNearestAsymmetricIndex(outputN, inputType.getDimSize(0), resultType.getDimSize(0), rewriter, nestedLoc);
-
-      auto channelLoop = buildNormalizedScfFor(
-        rewriter,
-        nestedLoc,
-        c0,
-        cOutputC,
-        c1,
-        ValueRange {outputBatchAcc},
-        [&](OpBuilder&,
-            Location channelLoc,
-            Value outputC,
-            ValueRange channelIterArgs,
-            SmallVectorImpl<Value>& channelYielded) {
-          Value outputChannelAcc = channelIterArgs.front();
-          Value inputC = buildNearestAsymmetricIndex(
-            outputC, inputType.getDimSize(1), resultType.getDimSize(1), rewriter, channelLoc);
-
-          auto heightLoop = buildNormalizedScfFor(
-            rewriter,
-            channelLoc,
-            c0,
-            cOutputH,
-            c1,
-            ValueRange {outputChannelAcc},
-            [&](OpBuilder&,
-                Location heightLoc,
-                Value outputH,
-                ValueRange heightIterArgs,
-                SmallVectorImpl<Value>& heightYielded) {
-              Value outputHeightAcc = heightIterArgs.front();
-              Value inputH = buildNearestAsymmetricIndex(
-                outputH, inputType.getDimSize(2), resultType.getDimSize(2), rewriter, heightLoc);
-
-              auto widthLoop = buildNormalizedScfFor(
-                rewriter,
-                heightLoc,
-                c0,
-                cOutputW,
-                c1,
-                ValueRange {outputHeightAcc},
-                [&](OpBuilder&,
-                    Location widthLoc,
-                    Value outputW,
-                    ValueRange widthIterArgs,
-                    SmallVectorImpl<Value>& widthYielded) {
-                  Value outputWidthAcc = widthIterArgs.front();
-                  Value inputW = buildNearestAsymmetricIndex(
-                    outputW, inputType.getDimSize(3), resultType.getDimSize(3), rewriter, widthLoc);
-
-                  SmallVector<OpFoldResult> inputOffsets = {inputN, inputC, inputH, inputW};
-                  Value inputSlice = tensor::ExtractSliceOp::create(
-                    rewriter, widthLoc, unitTensorType, input, inputOffsets, unitSizes, unitStrides);
-
-                  SmallVector<OpFoldResult> outputOffsets = {outputN, outputC, outputH, outputW};
-                  Value updatedOutput = tensor::InsertSliceOp::create(
-                    rewriter, widthLoc, inputSlice, outputWidthAcc, outputOffsets, unitSizes, unitStrides);
-                  widthYielded.push_back(updatedOutput);
-                  return success();
-                });
-              if (failed(widthLoop))
-                return failure();
-              heightYielded.push_back(widthLoop->results.front());
-              return success();
-            });
-          if (failed(heightLoop))
-            return failure();
-          channelYielded.push_back(heightLoop->results.front());
+static FailureOr<Value> buildDenseNearestResize(Value input,
+                                                RankedTensorType inputType,
+                                                RankedTensorType resultType,
+                                                PatternRewriter& rewriter,
+                                                Location loc) {
+  ArrayRef<int64_t> shape = resultType.getShape();
+  int64_t rowCount = shape[0] * shape[1] * shape[2];
+  auto scalarType = RankedTensorType::get({1, 1, 1, 1}, resultType.getElementType());
+  auto rowType = RankedTensorType::get({1, 1, 1, shape[3]}, resultType.getElementType());
+  auto rowsType = RankedTensorType::get({rowCount, 1, 1, 1, shape[3]}, resultType.getElementType());
+  auto batch = createSpatComputeBatch(
+    rewriter, loc, TypeRange {rowsType}, rowCount, {}, ValueRange {input},
+    [&](detail::SpatComputeBatchBodyArgs args) {
+      Operation* anchor = rewriter.getInsertionBlock()->getParentOp();
+      Value outputN = affineFloorDivConst(rewriter, loc, args.lane, shape[1] * shape[2], anchor);
+      Value channelRow = affineModConst(rewriter, loc, args.lane, shape[1] * shape[2], anchor);
+      Value outputC = affineFloorDivConst(rewriter, loc, channelRow, shape[2], anchor);
+      Value outputH = affineModConst(rewriter, loc, channelRow, shape[2], anchor);
+      Value inputN = buildNearestAsymmetricIndex(outputN, inputType.getDimSize(0), shape[0], rewriter, loc);
+      Value inputC = buildNearestAsymmetricIndex(outputC, inputType.getDimSize(1), shape[1], rewriter, loc);
+      Value inputH = buildNearestAsymmetricIndex(outputH, inputType.getDimSize(2), shape[2], rewriter, loc);
+      Value row = tensor::EmptyOp::create(rewriter, loc, rowType.getShape(), rowType.getElementType());
+      Value c0 = getOrCreateIndexConstant(rewriter, anchor, 0);
+      Value c1 = getOrCreateIndexConstant(rewriter, anchor, 1);
+      Value width = getOrCreateIndexConstant(rewriter, anchor, shape[3]);
+      auto loop = buildNormalizedScfFor(
+        rewriter, loc, c0, width, c1, ValueRange {row},
+        [&](OpBuilder&, Location nestedLoc, Value outputW, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
+          Value inputW = buildNearestAsymmetricIndex(
+            outputW, inputType.getDimSize(3), shape[3], rewriter, nestedLoc);
+          SmallVector<OpFoldResult> unitSizes(4, rewriter.getIndexAttr(1));
+          SmallVector<OpFoldResult> unitStrides(4, rewriter.getIndexAttr(1));
+          Value scalar = tensor::ExtractSliceOp::create(
+            rewriter, nestedLoc, scalarType, args.inputs.front(),
+            SmallVector<OpFoldResult> {inputN, inputC, inputH, inputW}, unitSizes, unitStrides);
+          yielded.push_back(tensor::InsertSliceOp::create(
+            rewriter, nestedLoc, scalar, iterArgs.front(),
+            SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0),
+                                       rewriter.getIndexAttr(0), outputW},
+            unitSizes, unitStrides));
           return success();
         });
-      if (failed(channelLoop))
+      assert(succeeded(loop) && "nearest Resize row loop construction must succeed");
+      publishGraphBatchPhysicalFragment(rewriter, loc, loop->results.front(), args.outputs.front(), args.lane);
+    });
+  if (failed(batch))
+    return failure();
+
+  SmallVector<FragmentAssemblyEntry> entries;
+  entries.reserve(rowCount);
+  for (int64_t n = 0; n < shape[0]; ++n)
+    for (int64_t c = 0; c < shape[1]; ++c)
+      for (int64_t h = 0; h < shape[2]; ++h)
+        entries.push_back({(n * shape[1] + c) * shape[2] + h, 0, {n, c, h, 0}, {1, 1, 1, shape[3]}});
+  return createFragmentAssemblyBlueprint(
+    batch->getResult(0), resultType, entries, "dense_nchw", spatial::kContiguousRowMajorFragments, rewriter, loc);
+}
+
+static FailureOr<Value> buildRowStripNearestResize(
+    Value storage, RankedTensorType inputType, RankedTensorType resultType,
+    PatternRewriter& rewriter, Location loc) {
+  auto input = describeRowStripPhysicalValue(storage, inputType);
+  if (failed(input))
+    return failure();
+  int64_t tilesPerRow = input->tilesPerRow;
+  int64_t outputHeight = resultType.getDimSize(2);
+  int64_t outputWidth = resultType.getDimSize(3);
+  int64_t tileChannels = input->fragmentType.getDimSize(3);
+  int64_t laneCount = outputHeight * tilesPerRow;
+  auto outputFragmentType = RankedTensorType::get(
+    {1, 1, outputWidth, tileChannels}, resultType.getElementType());
+  auto outputStorageType = spatial::getGraphBatchPhysicalResultType(
+    laneCount, outputFragmentType);
+  auto pixelType = RankedTensorType::get(
+    {1, 1, 1, tileChannels}, resultType.getElementType());
+  auto batch = createSpatComputeBatch(
+    rewriter, loc, TypeRange {outputStorageType}, laneCount, {}, ValueRange {storage},
+    [&](detail::SpatComputeBatchBodyArgs args) {
+      Operation* anchor = rewriter.getInsertionBlock()->getParentOp();
+      Value outputRow = affineFloorDivConst(rewriter, loc, args.lane, tilesPerRow, anchor);
+      Value tile = affineModConst(rewriter, loc, args.lane, tilesPerRow, anchor);
+      Value inputRow = buildNearestAsymmetricIndex(
+        outputRow, inputType.getDimSize(2), outputHeight, rewriter, loc);
+      Value inputSlot = arith::AddIOp::create(
+        rewriter, loc, affineMulConst(rewriter, loc, inputRow, tilesPerRow, anchor), tile);
+      auto source = extractGraphBatchPhysicalFragment(
+        rewriter, loc, args.inputs.front(), inputSlot, input->fragmentType);
+      if (failed(source))
         return failure();
-      batchYielded.push_back(channelLoop->results.front());
+      Value initial = tensor::EmptyOp::create(
+        rewriter, loc, outputFragmentType.getShape(), resultType.getElementType());
+      Value c0 = getOrCreateIndexConstant(rewriter, anchor, 0);
+      Value c1 = getOrCreateIndexConstant(rewriter, anchor, 1);
+      Value width = getOrCreateIndexConstant(rewriter, anchor, outputWidth);
+      auto loop = buildNormalizedScfFor(
+        rewriter, loc, c0, width, c1, ValueRange {initial},
+        [&](OpBuilder&, Location nestedLoc, Value outputColumn, ValueRange iterArgs,
+            SmallVectorImpl<Value>& yielded) {
+          Value inputColumn = buildNearestAsymmetricIndex(
+            outputColumn, inputType.getDimSize(3), outputWidth, rewriter, nestedLoc);
+          Value pixel = tensor::ExtractSliceOp::create(
+            rewriter, nestedLoc, pixelType, *source,
+            SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0),
+                                      inputColumn, rewriter.getIndexAttr(0)},
+            SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1),
+                                      rewriter.getIndexAttr(1), rewriter.getIndexAttr(tileChannels)},
+            getUnitStrides(rewriter, 4));
+          yielded.push_back(tensor::InsertSliceOp::create(
+            rewriter, nestedLoc, pixel, iterArgs.front(),
+            SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0),
+                                      outputColumn, rewriter.getIndexAttr(0)},
+            SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1),
+                                      rewriter.getIndexAttr(1), rewriter.getIndexAttr(tileChannels)},
+            getUnitStrides(rewriter, 4)));
+          return success();
+        });
+      if (failed(loop))
+        return failure();
+      publishGraphBatchPhysicalFragment(
+        rewriter, loc, loop->results.front(), args.outputs.front(), args.lane);
       return success();
     });
-  if (failed(batchLoop))
-    return failure();
-  return batchLoop->results.front();
+  return failed(batch) ? FailureOr<Value>(failure())
+                       : FailureOr<Value>(batch->getResult(0));
 }
 
 struct Resize : OpConversionPattern<ONNXResizeOp> {
@@ -161,23 +182,41 @@ struct Resize : OpConversionPattern<ONNXResizeOp> {
         || llvm::any_of(resultType.getShape(), [](int64_t dim) { return dim <= 0; }))
       return rewriter.notifyMatchFailure(resizeOp, "resize lowering requires positive static dimensions.");
 
-    auto computeOp = createSpatCompute<1>(
-      rewriter, resizeOp.getLoc(), TypeRange {resultType}, {}, adaptor.getX(), [&](Value x) -> LogicalResult {
-        auto result = buildNearestResizeLoop(x, inputType, resultType, rewriter, resizeOp.getLoc());
-        if (failed(result))
-          return failure();
-        spatial::SpatYieldOp::create(rewriter, resizeOp.getLoc(), *result);
-        return success();
-      });
-    if (failed(computeOp))
-      return failure();
-    rewriter.replaceOp(resizeOp, computeOp->getResults());
+    auto plan = spatial::SpatResizeNearestPlanOp::create(
+      rewriter, resizeOp.getLoc(), resultType, adaptor.getX(), spatial::getNCHWLayout(rewriter.getContext()));
+    rewriter.replaceOp(resizeOp, plan.getResult());
     return success();
   }
 };
 
 } // namespace
 
+LogicalResult canLowerResizeNearestPlanToRowStrip(
+    spatial::SpatResizeNearestPlanOp planOp,
+    const spatial::SpatialTargetResources&) {
+  auto inputType = dyn_cast<RankedTensorType>(planOp.getInput().getType());
+  auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
+  return success(inputType && outputType && inputType.hasStaticShape()
+    && outputType.hasStaticShape() && inputType.getRank() == 4
+    && outputType.getRank() == 4 && inputType.getDimSize(0) == 1
+    && outputType.getDimSize(0) == 1
+    && inputType.getDimSize(1) == outputType.getDimSize(1));
+}
+
+FailureOr<Value> lowerSelectedResizeNearestPlan(
+    spatial::SpatResizeNearestPlanOp planOp, Value input,
+    std::optional<Value> rowStripInput,
+    const spatial::SpatialTargetResources&,
+    PatternRewriter& rewriter) {
+  auto inputType = cast<RankedTensorType>(input.getType());
+  auto outputType = cast<RankedTensorType>(planOp.getOutput().getType());
+  if (rowStripInput)
+    return buildRowStripNearestResize(
+      *rowStripInput, inputType, outputType, rewriter, planOp.getLoc());
+  return buildDenseNearestResize(
+    input, inputType, outputType, rewriter, planOp.getLoc());
+}
+
 void populateResizePatterns(RewritePatternSet& patterns, MLIRContext* ctx) { patterns.add<Resize>(ctx); }
 
 } // namespace onnx_mlir

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

@@ -61,6 +61,74 @@ static FailureOr<Value> materializeTransposedConstant(Value input,
                              resultType);
 }
 
+static FailureOr<Value> transposeFragmentAssemblyBlueprint(spatial::SpatBlueprintOp blueprint,
+                                                           RankedTensorType resultType,
+                                                           ArrayRef<int64_t> permutation,
+                                                           ConversionPatternRewriter& rewriter,
+                                                           Location loc) {
+  auto storageType = dyn_cast<RankedTensorType>(blueprint.getInput().getType());
+  auto sourceOffsets = blueprint.getFragmentSourceOffsets();
+  auto fragmentStrides = blueprint.getFragmentStrides();
+  if (!storageType || !storageType.hasStaticShape() || !resultType.hasStaticShape()
+      || !blueprint.getFragments().empty() || !spatial::isFragmentAssembly(blueprint.getMode())
+      || !blueprint.getFragmentOperandIndices() || !sourceOffsets || !fragmentStrides
+      || llvm::any_of(*sourceOffsets, [](int64_t offset) { return offset != 0; })
+      || storageType.getRank() != resultType.getRank() + 1)
+    return failure();
+  if (blueprint.getIndexMap() == spatial::kContiguousRowMajorFragments
+      && !spatial::isCanonicalContiguousRowMajorFragmentAssembly(blueprint))
+    return blueprint.emitOpError("contiguous row-major fragment physical source order or storage is not canonical"), failure();
+
+  SmallVector<int64_t> outputStorageShape {storageType.getDimSize(0)};
+  for (int64_t sourceDim : permutation)
+    outputStorageShape.push_back(storageType.getDimSize(sourceDim + 1));
+  auto outputStorageType = RankedTensorType::get(outputStorageShape, storageType.getElementType());
+  auto mapped = mapGraphBatchFragments(
+    blueprint.getInput(), outputStorageType, rewriter, loc, [&](Value fragment, RankedTensorType fragmentType) {
+      Value init = createTransposeInit(fragment, fragmentType, permutation, rewriter, loc);
+      return FailureOr<Value>(
+        linalg::TransposeOp::create(rewriter, loc, fragment, init, permutation).getResult()[0]);
+    });
+  if (failed(mapped))
+    return failure();
+
+  const int64_t rank = resultType.getRank();
+  const int64_t fragmentCount = blueprint.getFragmentOperandIndices()->size();
+  SmallVector<int64_t> offsets, sizes, strides;
+  offsets.reserve(fragmentCount * rank);
+  sizes.reserve(fragmentCount * rank);
+  strides.reserve(fragmentCount * rank);
+  ArrayRef<int64_t> inputOffsets = blueprint.getFragmentOffsets();
+  ArrayRef<int64_t> inputSizes = blueprint.getFragmentSizes();
+  for (int64_t fragment = 0; fragment < fragmentCount; ++fragment)
+    for (int64_t sourceDim : permutation) {
+      const int64_t index = fragment * rank + sourceDim;
+      offsets.push_back(inputOffsets[index]);
+      sizes.push_back(inputSizes[index]);
+      strides.push_back((*fragmentStrides)[index]);
+    }
+  auto transposedBlueprint = spatial::SpatBlueprintOp::create(rewriter,
+                                                               loc,
+                                                               resultType,
+                                                               *mapped,
+                                                               ValueRange {},
+                                                               blueprint.getLogicalLayoutAttr(),
+                                                               spatial::getFragmentedLayout(rewriter.getContext()),
+                                                               rewriter.getDenseI64ArrayAttr(offsets),
+                                                               rewriter.getDenseI64ArrayAttr(sizes),
+                                                               rewriter.getStringAttr("permuted_fragments"),
+                                                               blueprint.getModeAttr(),
+                                                               blueprint.getFragmentOperandIndicesAttr(),
+                                                               blueprint.getFragmentSourceSlotsAttr(),
+                                                               blueprint.getFragmentSourceOffsetsAttr(),
+                                                               rewriter.getDenseI64ArrayAttr(strides),
+                                                               blueprint.getConflictPolicyAttr(),
+                                                               blueprint.getCoveragePolicyAttr());
+  if (spatial::isCanonicalContiguousRowMajorFragmentAssembly(transposedBlueprint))
+    transposedBlueprint.setIndexMapAttr(rewriter.getStringAttr(spatial::kContiguousRowMajorFragments));
+  return transposedBlueprint.getOutput();
+}
+
 struct TransposeToLinalgTranspose : OpConversionPattern<ONNXTransposeOp> {
   using OpConversionPattern::OpConversionPattern;
 
@@ -75,6 +143,14 @@ struct TransposeToLinalgTranspose : OpConversionPattern<ONNXTransposeOp> {
     auto permutation = getTransposePermutationChecked(transposeOp.getPermAttr(), inputType.getRank());
     if (failed(permutation))
       return failure();
+    if (auto blueprint = adaptor.getData().getDefiningOp<spatial::SpatBlueprintOp>()) {
+      auto transposed =
+        transposeFragmentAssemblyBlueprint(blueprint, resultType, *permutation, rewriter, transposeOp.getLoc());
+      if (succeeded(transposed)) {
+        rewriter.replaceOp(transposeOp, *transposed);
+        return success();
+      }
+    }
     if (isCompileTimeComputable(adaptor.getData())) {
       auto constantTranspose =
         materializeTransposedConstant(adaptor.getData(), resultType, *permutation, rewriter, transposeOp.getLoc());

src/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp

@@ -1,44 +0,0 @@
-#pragma once
-
-#include <optional>
-
-#include "mlir/IR/PatternMatch.h"
-#include "mlir/Support/LogicalResult.h"
-
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-
-namespace onnx_mlir {
-
-struct RowStripPhysicalValue;
-
-mlir::FailureOr<mlir::Value>
-lowerSelectedConv2DPlan(spatial::SpatConv2DPlanOp planOp,
-                        std::optional<mlir::Value> rowStripInput,
-                        bool emitRowStripLayout,
-                        mlir::PatternRewriter& rewriter);
-
-mlir::LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp);
-mlir::LogicalResult canConsumeAndProduceRowStrip(spatial::SpatConv2DPlanOp planOp);
-
-mlir::LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp);
-
-mlir::FailureOr<mlir::Value>
-lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
-                           std::optional<mlir::Value> rowStripInput,
-                           mlir::PatternRewriter& rewriter);
-
-mlir::LogicalResult
-canLowerGlobalAveragePoolPlanToRowStrip(spatial::SpatGlobalAveragePoolPlanOp planOp);
-
-mlir::FailureOr<mlir::Value>
-lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp,
-                                   std::optional<mlir::Value> rowStripInput,
-                                   mlir::PatternRewriter& rewriter);
-
-mlir::LogicalResult canLowerFlattenFromRowStrip(spatial::SpatGraphCompute flattenOp);
-
-mlir::LogicalResult lowerFlattenFromRowStrip(const RowStripPhysicalValue& input,
-                                             spatial::SpatGraphCompute flattenOp,
-                                             mlir::PatternRewriter& rewriter);
-
-} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp

@@ -1,330 +0,0 @@
-#include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/IR/PatternMatch.h"
-#include "mlir/Pass/Pass.h"
-
-#include "llvm/ADT/DenseMap.h"
-
-#include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
-#include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-#include "src/Accelerators/PIM/Pass/PIMPasses.h"
-
-using namespace mlir;
-
-namespace onnx_mlir {
-namespace {
-
-static constexpr StringLiteral kLogicalLayout = "nchw";
-static constexpr StringLiteral kDenseLayout = "dense_nchw";
-static constexpr StringLiteral kRowStripLayout = "nhwc_row_strip";
-
-enum class SelectedLayout {
-  DenseNchw,
-  PixelMajorRowStrip,
-};
-
-static SelectedLayout getSelectedLayout(llvm::DenseMap<Value, SelectedLayout>& layouts, Value value) {
-  auto it = layouts.find(value);
-  return it == layouts.end() ? SelectedLayout::DenseNchw : it->second;
-}
-
-static bool usesSelectedRowStrip(Operation* user, llvm::DenseMap<Value, SelectedLayout>& layouts) {
-  if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(user))
-    return getSelectedLayout(layouts, reluPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
-  if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(user))
-    return getSelectedLayout(layouts, biasAddPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
-  if (auto addPlan = dyn_cast<spatial::SpatAddPlanOp>(user))
-    return getSelectedLayout(layouts, addPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
-  if (auto concatPlan = dyn_cast<spatial::SpatConcatPlanOp>(user))
-    return getSelectedLayout(layouts, concatPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
-  if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(user))
-    return getSelectedLayout(layouts, convPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
-  if (auto maxPoolPlan = dyn_cast<spatial::SpatMaxPool2DPlanOp>(user))
-    return getSelectedLayout(layouts, maxPoolPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
-  if (auto averagePoolPlan = dyn_cast<spatial::SpatGlobalAveragePoolPlanOp>(user))
-    return getSelectedLayout(layouts, averagePoolPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
-  if (auto flattenCompute = dyn_cast<spatial::SpatGraphCompute>(user))
-    return succeeded(canLowerFlattenFromRowStrip(flattenCompute));
-  return false;
-}
-
-static bool allUsersCanHandleRowStrip(Value value, llvm::DenseMap<Value, SelectedLayout>& layouts) {
-  for (Operation* user : value.getUsers()) {
-    if (usesSelectedRowStrip(user, layouts))
-      continue;
-    // Dense-only users must be materialized explicitly.
-    continue;
-  }
-  return true;
-}
-
-static bool canConsumeRowStripAsUser(Operation* user) {
-  if (isa<spatial::SpatReluPlanOp>(user))
-    return true;
-  if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(user)) {
-    auto resultType = dyn_cast<RankedTensorType>(biasAddPlan.getOutput().getType());
-    return resultType && isSupportedBiasAddValue(biasAddPlan.getBias(), resultType);
-  }
-  if (isa<spatial::SpatAddPlanOp>(user))
-    return true;
-  if (isa<spatial::SpatConcatPlanOp>(user))
-    return true;
-  if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(user))
-    return succeeded(canConsumeAndProduceRowStrip(convPlan));
-  if (auto maxPoolPlan = dyn_cast<spatial::SpatMaxPool2DPlanOp>(user))
-    return succeeded(canLowerMaxPoolPlanToRowStrip(maxPoolPlan));
-  if (auto averagePoolPlan = dyn_cast<spatial::SpatGlobalAveragePoolPlanOp>(user))
-    return succeeded(canLowerGlobalAveragePoolPlanToRowStrip(averagePoolPlan));
-  return false;
-}
-
-static bool hasRowStripConsumer(Value value) {
-  for (Operation* user : value.getUsers())
-    if (canConsumeRowStripAsUser(user))
-      return true;
-  return false;
-}
-
-static bool canSelectConvRowStrip(spatial::SpatConv2DPlanOp convPlan,
-                                  llvm::DenseMap<Value, SelectedLayout>& layouts) {
-  SelectedLayout inputLayout = getSelectedLayout(layouts, convPlan.getInput());
-  if (inputLayout == SelectedLayout::PixelMajorRowStrip)
-    return succeeded(canConsumeAndProduceRowStrip(convPlan));
-  return succeeded(canLowerConvPlanToRowStrip(convPlan));
-}
-
-static SelectedLayout chooseConvLayout(spatial::SpatConv2DPlanOp convPlan,
-                                       llvm::DenseMap<Value, SelectedLayout>& layouts) {
-  if (!canSelectConvRowStrip(convPlan, layouts))
-    return SelectedLayout::DenseNchw;
-  if (!allUsersCanHandleRowStrip(convPlan.getResult(), layouts))
-    return SelectedLayout::DenseNchw;
-  return SelectedLayout::PixelMajorRowStrip;
-}
-
-static SelectedLayout chooseReluLayout(spatial::SpatReluPlanOp reluPlan,
-                                       llvm::DenseMap<Value, SelectedLayout>& layouts) {
-  if (getSelectedLayout(layouts, reluPlan.getInput()) != SelectedLayout::PixelMajorRowStrip)
-    return SelectedLayout::DenseNchw;
-  if (!hasRowStripConsumer(reluPlan.getResult()))
-    return SelectedLayout::DenseNchw;
-  if (!allUsersCanHandleRowStrip(reluPlan.getResult(), layouts))
-    return SelectedLayout::DenseNchw;
-  return SelectedLayout::PixelMajorRowStrip;
-}
-
-static SelectedLayout chooseBiasAddLayout(spatial::SpatBiasAddPlanOp biasAddPlan,
-                                          llvm::DenseMap<Value, SelectedLayout>& layouts) {
-  if (getSelectedLayout(layouts, biasAddPlan.getInput()) != SelectedLayout::PixelMajorRowStrip)
-    return SelectedLayout::DenseNchw;
-  auto resultType = dyn_cast<RankedTensorType>(biasAddPlan.getOutput().getType());
-  if (!resultType || !isSupportedBiasAddValue(biasAddPlan.getBias(), resultType))
-    return SelectedLayout::DenseNchw;
-  if (!hasRowStripConsumer(biasAddPlan.getResult()))
-    return SelectedLayout::DenseNchw;
-  if (!allUsersCanHandleRowStrip(biasAddPlan.getResult(), layouts))
-    return SelectedLayout::DenseNchw;
-  return SelectedLayout::PixelMajorRowStrip;
-}
-
-static SelectedLayout chooseAddLayout(spatial::SpatAddPlanOp addPlan, llvm::DenseMap<Value, SelectedLayout>& layouts) {
-  if (getSelectedLayout(layouts, addPlan.getLhs()) != SelectedLayout::PixelMajorRowStrip
-      || getSelectedLayout(layouts, addPlan.getRhs()) != SelectedLayout::PixelMajorRowStrip)
-    return SelectedLayout::DenseNchw;
-  if (!allUsersCanHandleRowStrip(addPlan.getResult(), layouts))
-    return SelectedLayout::DenseNchw;
-  return SelectedLayout::PixelMajorRowStrip;
-}
-
-static SelectedLayout chooseConcatLayout(spatial::SpatConcatPlanOp concatPlan,
-                                         llvm::DenseMap<Value, SelectedLayout>& layouts) {
-  if (llvm::any_of(concatPlan.getInputs(), [&](Value input) {
-        return getSelectedLayout(layouts, input) != SelectedLayout::PixelMajorRowStrip;
-      }))
-    return SelectedLayout::DenseNchw;
-  if (!allUsersCanHandleRowStrip(concatPlan.getResult(), layouts))
-    return SelectedLayout::DenseNchw;
-  return SelectedLayout::PixelMajorRowStrip;
-}
-
-static SelectedLayout chooseMaxPoolLayout(spatial::SpatMaxPool2DPlanOp maxPoolPlan) {
-  return succeeded(canLowerMaxPoolPlanToRowStrip(maxPoolPlan)) ? SelectedLayout::PixelMajorRowStrip
-                                                               : SelectedLayout::DenseNchw;
-}
-
-static SelectedLayout chooseGlobalAveragePoolLayout(
-  spatial::SpatGlobalAveragePoolPlanOp averagePoolPlan) {
-  return succeeded(canLowerGlobalAveragePoolPlanToRowStrip(averagePoolPlan))
-           ? SelectedLayout::PixelMajorRowStrip
-           : SelectedLayout::DenseNchw;
-}
-
-static spatial::SpatBlueprintOp insertRowStripBlueprint(IRRewriter& rewriter, Value value) {
-  auto outputType = cast<RankedTensorType>(value.getType());
-  auto [offsets, sizes] = buildRowStripMetadata(outputType);
-  return spatial::SpatBlueprintOp::create(rewriter,
-                                              value.getLoc(),
-                                              outputType,
-                                              value,
-                                              ValueRange {},
-                                              rewriter.getStringAttr(kLogicalLayout),
-                                              rewriter.getStringAttr(kRowStripLayout),
-                                              rewriter.getDenseI64ArrayAttr(offsets),
-                                              rewriter.getDenseI64ArrayAttr(sizes),
-                                              rewriter.getStringAttr(kRowStripIndexMap),
-                                              nullptr,
-                                              nullptr,
-                                              nullptr,
-                                              nullptr,
-                                              nullptr,
-                                              nullptr,
-                                              nullptr);
-}
-
-static void materializeDenseUses(IRRewriter& rewriter,
-                                 Value layoutValue,
-                                 llvm::DenseMap<Value, SelectedLayout>& layouts) {
-  SmallVector<OpOperand*> denseUses;
-  for (OpOperand& use : layoutValue.getUses()) {
-    if (usesSelectedRowStrip(use.getOwner(), layouts))
-      continue;
-    denseUses.push_back(&use);
-  }
-
-  for (OpOperand* use : denseUses) {
-    Operation* owner = use->getOwner();
-    rewriter.setInsertionPoint(owner);
-    auto materialized = spatial::SpatMaterializeLayoutOp::create(rewriter,
-                                                                 owner->getLoc(),
-                                                                 use->get().getType(),
-                                                                 use->get(),
-                                                                 rewriter.getStringAttr(kLogicalLayout),
-                                                                 rewriter.getStringAttr(kRowStripLayout),
-                                                                 rewriter.getStringAttr(kDenseLayout));
-    use->set(materialized.getResult());
-  }
-}
-
-struct SpatialLayoutPlanningPass final : PassWrapper<SpatialLayoutPlanningPass, OperationPass<ModuleOp>> {
-  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SpatialLayoutPlanningPass)
-
-  StringRef getArgument() const override { return "spatial-layout-planning"; }
-  StringRef getDescription() const override { return "Select conservative Spatial layouts and insert reconciliation barriers."; }
-
-  void runOnOperation() override {
-    auto entryFunc = getPimEntryFunc(getOperation());
-    if (failed(entryFunc)) {
-      getOperation().emitError("failed to locate the PIM entry function during Spatial layout planning");
-      signalPassFailure();
-      return;
-    }
-
-    func::FuncOp funcOp = *entryFunc;
-    IRRewriter rewriter(&getContext());
-    llvm::DenseMap<Value, SelectedLayout> layouts;
-
-    bool changed = true;
-    while (changed) {
-      changed = false;
-      for (Operation& op : llvm::make_early_inc_range(funcOp.getBody().front())) {
-        if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(&op)) {
-          SelectedLayout selected = chooseConvLayout(convPlan, layouts);
-          if (layouts[convPlan.getResult()] != selected) {
-            layouts[convPlan.getResult()] = selected;
-            changed = true;
-          }
-          continue;
-        }
-        if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(&op)) {
-          SelectedLayout selected = chooseReluLayout(reluPlan, layouts);
-          if (layouts[reluPlan.getResult()] != selected) {
-            layouts[reluPlan.getResult()] = selected;
-            changed = true;
-          }
-          continue;
-        }
-        if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(&op)) {
-          SelectedLayout selected = chooseBiasAddLayout(biasAddPlan, layouts);
-          if (layouts[biasAddPlan.getResult()] != selected) {
-            layouts[biasAddPlan.getResult()] = selected;
-            changed = true;
-          }
-          continue;
-        }
-        if (auto addPlan = dyn_cast<spatial::SpatAddPlanOp>(&op)) {
-          SelectedLayout selected = chooseAddLayout(addPlan, layouts);
-          if (layouts[addPlan.getResult()] != selected) {
-            layouts[addPlan.getResult()] = selected;
-            changed = true;
-          }
-          continue;
-        }
-        if (auto concatPlan = dyn_cast<spatial::SpatConcatPlanOp>(&op)) {
-          SelectedLayout selected = chooseConcatLayout(concatPlan, layouts);
-          if (layouts[concatPlan.getResult()] != selected) {
-            layouts[concatPlan.getResult()] = selected;
-            changed = true;
-          }
-          continue;
-        }
-        if (auto maxPoolPlan = dyn_cast<spatial::SpatMaxPool2DPlanOp>(&op)) {
-          SelectedLayout selected = chooseMaxPoolLayout(maxPoolPlan);
-          if (layouts[maxPoolPlan.getResult()] != selected) {
-            layouts[maxPoolPlan.getResult()] = selected;
-            changed = true;
-          }
-          continue;
-        }
-        if (auto averagePoolPlan = dyn_cast<spatial::SpatGlobalAveragePoolPlanOp>(&op)) {
-          SelectedLayout selected = chooseGlobalAveragePoolLayout(averagePoolPlan);
-          if (layouts[averagePoolPlan.getResult()] != selected) {
-            layouts[averagePoolPlan.getResult()] = selected;
-            changed = true;
-          }
-          continue;
-        }
-      }
-    }
-
-    for (Operation& op : llvm::make_early_inc_range(funcOp.getBody().front())) {
-      Value producedValue;
-      if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(&op))
-        producedValue = convPlan.getResult();
-      else if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(&op))
-        producedValue = biasAddPlan.getResult();
-      else if (auto addPlan = dyn_cast<spatial::SpatAddPlanOp>(&op))
-        producedValue = addPlan.getResult();
-      else if (auto concatPlan = dyn_cast<spatial::SpatConcatPlanOp>(&op))
-        producedValue = concatPlan.getResult();
-      else if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(&op))
-        producedValue = reluPlan.getResult();
-      else if (auto maxPoolPlan = dyn_cast<spatial::SpatMaxPool2DPlanOp>(&op))
-        producedValue = maxPoolPlan.getResult();
-      else if (auto averagePoolPlan = dyn_cast<spatial::SpatGlobalAveragePoolPlanOp>(&op))
-        producedValue = averagePoolPlan.getResult();
-      else
-        continue;
-
-      if (getSelectedLayout(layouts, producedValue) != SelectedLayout::PixelMajorRowStrip)
-        continue;
-
-      rewriter.setInsertionPointAfter(&op);
-      auto blueprint = insertRowStripBlueprint(rewriter, producedValue);
-      rewriter.replaceAllUsesExcept(producedValue, blueprint.getResult(), blueprint);
-      materializeDenseUses(rewriter, blueprint.getResult(), layouts);
-    }
-    if (failed(verifyLogicalSpatialGraphInvariants(*entryFunc))) {
-      getOperation().emitError("logical Spatial graph verification failed after SpatialLayoutPlanning");
-      signalPassFailure();
-    }
-  }
-};
-
-} // namespace
-
-std::unique_ptr<Pass> createSpatialLayoutPlanningPass() { return std::make_unique<SpatialLayoutPlanningPass>(); }
-
-} // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.cpp

@@ -149,11 +149,10 @@ collectTopLevelFragmentAssemblyCopies(OpResult result, RankedTensorType packedRe
     auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(use.getOwner());
     if (!blueprint || blueprint->getParentOp() != blueprint->getParentOfType<func::FuncOp>())
       return failure();
-    std::optional<StringRef> mode = blueprint.getMode();
     std::optional<ArrayRef<int64_t>> operandIndicesAttr = blueprint.getFragmentOperandIndices();
     std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = blueprint.getFragmentSourceOffsets();
     std::optional<ArrayRef<int64_t>> sourceSlotsAttr = blueprint.getFragmentSourceSlots();
-    if (!mode || *mode != "fragment_assembly" || !operandIndicesAttr || !sourceOffsetsAttr || !sourceSlotsAttr)
+    if (!spatial::isFragmentAssembly(blueprint.getMode()) || !operandIndicesAttr || !sourceOffsetsAttr || !sourceSlotsAttr)
       return failure();
     if (!blueprint.getOutput().hasOneUse() || !isa<func::ReturnOp>(*blueprint.getOutput().getUsers().begin()))
       return failure();
@@ -308,7 +307,7 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatSchedul
       "resultful compute_batch lowering currently requires a spat.in_parallel terminator");
   }
 
-  auto coreIds = getRequiredScheduledBatchCoreIds(computeBatchOp, "spatial compute_batch core id");
+  auto coreIds = getRequiredScheduledBatchCoreIds(computeBatchOp, "Spatial compute_batch core id");
   if (failed(coreIds))
     return failure();
   SmallVector<Value> batchWeights(computeBatchOp.getWeights().begin(), computeBatchOp.getWeights().end());
@@ -318,7 +317,7 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatSchedul
 
   rewriter.setInsertionPointAfter(computeBatchOp);
   auto laneCountAttr = pim::getCheckedI32Attr(
-    rewriter, computeBatchOp, static_cast<uint64_t>(computeBatchOp.getLaneCount()), "pim core_batch lane count");
+    rewriter, computeBatchOp, static_cast<uint64_t>(computeBatchOp.getLaneCount()), "Pim core_batch lane count");
   if (failed(laneCountAttr))
     return failure();
   auto coreBatchOp =
@@ -418,8 +417,7 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatSchedul
     rewriter.setInsertionPointToEnd(newBlock);
 
     if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
-      std::optional<StringRef> modeAttr = blueprint.getMode();
-      if (modeAttr && *modeAttr == "fragment_assembly") {
+      if (spatial::isFragmentAssembly(blueprint.getMode())) {
         for (Operation* user : blueprint.getOutput().getUsers()) {
           if (!isa<tensor::ParallelInsertSliceOp>(user))
             return blueprint.emitOpError(
@@ -483,8 +481,7 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatSchedul
         auto hostTargetType = cast<ShapedType>(hostTarget.getType());
         if (auto blueprint =
               insertSlice.getSource().getDefiningOp<spatial::SpatBlueprintOp>()) {
-          std::optional<StringRef> modeAttr = blueprint.getMode();
-          if (modeAttr && *modeAttr == "fragment_assembly") {
+          if (spatial::isFragmentAssembly(blueprint.getMode())) {
             FailureOr<SmallVector<FragmentAssemblyCopy, 8>> fragmentAssemblyCopies =
               collectFragmentAssemblyCopiesFromBlueprint(blueprint, mapper, /*lane=*/0, /*hostTargetIndex=*/0);
             if (failed(fragmentAssemblyCopies))

src/PIM/Conversion/SpatialToPim/Common.cpp

@@ -1,7 +1,10 @@
 #include "mlir/IR/ValueRange.h"
 
 #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/IR/BuiltinOps.h"
 
 #include "llvm/ADT/STLExtras.h"
 
@@ -28,6 +31,49 @@ FailureOr<IntegerAttr> getTensorSizeInBytesAttr(Builder& builder, Operation* anc
   return pim::getCheckedI32Attr(builder, anchor, *byteSize, "tensor byte size");
 }
 
+LogicalResult materializePipelineHostBuffer(
+    func::FuncOp funcOp, RewriterBase &rewriter) {
+  auto bytes = funcOp->getAttrOfType<IntegerAttr>(
+    kPipelineHostBufferBytesAttrName);
+  if (!bytes)
+    return success();
+  if (bytes.getInt() <= 0)
+    return funcOp.emitOpError(
+      "pipeline host transfer buffer must be positive");
+  ModuleOp moduleOp = funcOp->getParentOfType<ModuleOp>();
+  if (moduleOp.lookupSymbol<memref::GlobalOp>(kPipelineHostBufferName))
+    return funcOp.emitOpError(
+      "pipeline host transfer buffer symbol already exists");
+  auto type = MemRefType::get(
+    {bytes.getInt()}, rewriter.getI8Type());
+  OpBuilder::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPointToStart(moduleOp.getBody());
+  memref::GlobalOp::create(
+    rewriter, funcOp.getLoc(),
+    rewriter.getStringAttr(kPipelineHostBufferName),
+    rewriter.getStringAttr("private"), TypeAttr::get(type), Attribute(),
+    UnitAttr(), IntegerAttr());
+  return success();
+}
+
+FailureOr<mlir::Value> getPipelineHostBuffer(
+    OpBuilder &builder, Operation *anchor) {
+  auto funcOp = anchor->getParentOfType<func::FuncOp>();
+  auto moduleOp = anchor->getParentOfType<ModuleOp>();
+  auto bytes = funcOp
+    ? funcOp->getAttrOfType<IntegerAttr>(kPipelineHostBufferBytesAttrName)
+    : IntegerAttr();
+  auto global = moduleOp
+    ? moduleOp.lookupSymbol<memref::GlobalOp>(kPipelineHostBufferName)
+    : memref::GlobalOp();
+  if (!bytes || !global)
+    return anchor->emitOpError(
+      "requires the pipeline host transfer buffer"), failure();
+  auto type = MemRefType::get({bytes.getInt()}, builder.getI8Type());
+  return memref::GetGlobalOp::create(
+    builder, anchor->getLoc(), type, kPipelineHostBufferName).getResult();
+}
+
 Operation* getEarliestUserWithinBlock(mlir::Value value) {
   auto users = value.getUsers();
 
@@ -129,6 +175,32 @@ LogicalResult validateFragmentAssemblyMetadata(spatial::SpatBlueprintOp blueprin
   return success();
 }
 
+FailureOr<mlir::Value> reshapeContiguousRowMajorFragments(RewriterBase& rewriter,
+                                                          Location loc,
+                                                          mlir::Value source,
+                                                          RankedTensorType resultType) {
+  auto sourceType = dyn_cast<RankedTensorType>(source.getType());
+  if (!sourceType || !sourceType.hasStaticShape() || !resultType.hasStaticShape() || resultType.getRank() < 2
+      || sourceType.getRank() != resultType.getRank() + 1 || sourceType.getElementType() != resultType.getElementType()
+      || sourceType.getNumElements() != resultType.getNumElements()
+      || sourceType.getDimSize(0) != getStaticShapeElementCount(resultType.getShape().drop_back())
+      || sourceType.getDimSize(sourceType.getRank() - 1) != resultType.getDimSize(resultType.getRank() - 1)
+      || llvm::any_of(sourceType.getShape().slice(1, sourceType.getRank() - 2), [](int64_t dim) { return dim != 1; }))
+    return failure();
+
+  SmallVector<ReassociationIndices> collapse {{}, {sourceType.getRank() - 1}};
+  for (int64_t dim = 0; dim < sourceType.getRank() - 1; ++dim)
+    collapse.front().push_back(dim);
+  auto flatType = RankedTensorType::get(
+    {sourceType.getDimSize(0), sourceType.getDimSize(sourceType.getRank() - 1)}, resultType.getElementType());
+  mlir::Value flat = tensor::CollapseShapeOp::create(rewriter, loc, flatType, source, collapse);
+
+  SmallVector<ReassociationIndices> expand {{}, {resultType.getRank() - 1}};
+  for (int64_t dim = 0; dim < resultType.getRank() - 1; ++dim)
+    expand.front().push_back(dim);
+  return tensor::ExpandShapeOp::create(rewriter, loc, resultType, flat, expand).getResult();
+}
+
 static SmallVector<int64_t, 4> expandFlatElementIndex(int64_t flatIndex, ArrayRef<int64_t> shape) {
   SmallVector<int64_t, 4> indices(shape.size(), 0);
   for (int64_t dim = static_cast<int64_t>(shape.size()) - 1; dim >= 0; --dim) {

src/PIM/Conversion/SpatialToPim/Common.hpp

@@ -10,6 +10,7 @@
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Value.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Support/LogicalResult.h"
 
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
@@ -23,6 +24,12 @@ namespace onnx_mlir {
 mlir::FailureOr<mlir::IntegerAttr>
 getTensorSizeInBytesAttr(mlir::Builder& builder, mlir::Operation* anchor, mlir::Value value);
 
+mlir::LogicalResult materializePipelineHostBuffer(
+  mlir::func::FuncOp funcOp, mlir::RewriterBase &rewriter);
+
+mlir::FailureOr<mlir::Value> getPipelineHostBuffer(
+  mlir::OpBuilder &builder, mlir::Operation *anchor);
+
 template <class T>
 size_t rangeLength(const mlir::iterator_range<T> range) {
   return std::distance(range.begin(), range.end());
@@ -51,6 +58,11 @@ mlir::LogicalResult validateFragmentAssemblyMetadata(onnx_mlir::spatial::SpatBlu
                                                      llvm::ArrayRef<int64_t> flatSizes,
                                                      llvm::ArrayRef<int64_t> flatStrides);
 
+mlir::FailureOr<mlir::Value> reshapeContiguousRowMajorFragments(mlir::RewriterBase& rewriter,
+                                                                mlir::Location loc,
+                                                                mlir::Value source,
+                                                                mlir::RankedTensorType resultType);
+
 mlir::FailureOr<mlir::SmallVector<int64_t, 4>>
 getStaticSliceOffsetsForElementOffset(mlir::Operation* anchor,
                                       mlir::ShapedType sourceType,

src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp

@@ -42,12 +42,11 @@ static FailureOr<Value> lowerFragmentAssemblyBlueprint(IRRewriter& rewriter,
   if (!resultType || !resultType.hasStaticShape())
     return blueprint.emitOpError("fragment assembly lowering requires a static ranked tensor result");
 
-  std::optional<StringRef> modeAttr = blueprint.getMode();
   std::optional<ArrayRef<int64_t>> operandIndicesAttr = blueprint.getFragmentOperandIndices();
   std::optional<ArrayRef<int64_t>> sourceSlotsAttr = blueprint.getFragmentSourceSlots();
   std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = blueprint.getFragmentSourceOffsets();
   std::optional<ArrayRef<int64_t>> fragmentStridesAttr = blueprint.getFragmentStrides();
-  if (!modeAttr || *modeAttr != "fragment_assembly" || !operandIndicesAttr || !sourceSlotsAttr
+  if (!spatial::isFragmentAssembly(blueprint.getMode()) || !operandIndicesAttr || !sourceSlotsAttr
       || !sourceOffsetsAttr || !fragmentStridesAttr)
     return blueprint.emitOpError("fragment assembly lowering requires explicit fragment metadata");
 
@@ -71,6 +70,16 @@ static FailureOr<Value> lowerFragmentAssemblyBlueprint(IRRewriter& rewriter,
                                               flatStrides)))
     return failure();
 
+  if (blueprint.getIndexMap() == spatial::kContiguousRowMajorFragments) {
+    if (!spatial::isCanonicalContiguousRowMajorFragmentAssembly(blueprint))
+      return blueprint.emitOpError("contiguous row-major fragment physical source order or storage is not canonical"), failure();
+    Value source = mapping.lookupOrDefault(blueprint.getInput());
+    auto reshaped = reshapeContiguousRowMajorFragments(
+      rewriter, blueprint.getLoc(), source, cast<RankedTensorType>(resultType));
+    if (failed(reshaped))
+      return blueprint.emitOpError("contiguous row-major fragment storage does not match its logical result"), failure();
+    return *reshaped;
+  }
   SmallVector<int64_t> hostStrides = computeRowMajorStrides(resultType.getShape());
   SmallVector<FragmentAssemblyCopy, 8> copies;
   for (int64_t fragmentIndex = 0; fragmentIndex < static_cast<int64_t>(operandIndices.size()); ++fragmentIndex) {
@@ -193,8 +202,7 @@ static bool isHostMaterializableHelperOp(Operation* op) {
   if (isa<arith::ConstantOp>(op) || op->hasTrait<OpTrait::ConstantLike>())
     return true;
   if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
-    std::optional<StringRef> mode = blueprint.getMode();
-    return mode && *mode == "fragment_assembly";
+    return spatial::isFragmentAssembly(blueprint.getMode());
   }
   return isShapingOnlyOp(op) || isPureIndexComputationOp(op);
 }
@@ -281,8 +289,7 @@ static bool inlineInputlessHelperComputeForWeightLikeUsers(spatial::SpatSchedule
   }
   for (Operation& op : block.without_terminator()) {
     if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
-      std::optional<StringRef> modeAttr = blueprint.getMode();
-      if (modeAttr && *modeAttr == "fragment_assembly") {
+      if (spatial::isFragmentAssembly(blueprint.getMode())) {
         auto lowered = lowerFragmentAssemblyBlueprint(rewriter, blueprint, mapping);
         if (failed(lowered))
           return false;
@@ -338,18 +345,42 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatScheduledCom
     auto blockArg = computeOp.getInputArgument(inputIndex);
     if (!blockArg)
       return computeOp.emitOpError("expected compute input block arguments during lowering");
-    auto receiveOp = dyn_cast_or_null<spatial::SpatChannelReceiveOp>(input.getDefiningOp());
+    auto channelReceive = dyn_cast_or_null<spatial::SpatChannelReceiveOp>(
+      input.getDefiningOp());
+    auto hostWaitLoad = dyn_cast_or_null<spatial::SpatHostWaitLoadOp>(
+      input.getDefiningOp());
+    Operation *receiveOp = channelReceive
+      ? channelReceive.getOperation() : hostWaitLoad.getOperation();
     if (receiveOp && !blockArg->use_empty()) {
       rewriter.setInsertionPoint(getEarliestUserWithinBlock(*blockArg));
       auto outputType = cast<ShapedType>(blockArg->getType());
-      auto outputBuffer = createEmptyTensorFromShaped(rewriter, receiveOp.getLoc(), outputType);
+      auto outputBuffer = createEmptyTensorFromShaped(
+        rewriter, receiveOp->getLoc(), outputType);
       auto sizeAttr = getTensorSizeInBytesAttr(rewriter, computeOp.getOperation(), *blockArg);
       if (failed(sizeAttr))
         return failure();
-      Value received =
-        PimReceiveOp::create(
-          rewriter, receiveOp.getLoc(), outputBuffer.getType(), outputBuffer, *sizeAttr, receiveOp.getSourceCoreId())
+      Value zero = arith::ConstantIndexOp::create(
+        rewriter, receiveOp->getLoc(), 0);
+      Value received;
+      if (hostWaitLoad) {
+        auto hostBuffer = getPipelineHostBuffer(rewriter, hostWaitLoad);
+        if (failed(hostBuffer))
+          return failure();
+        PimWaitOp::create(
+          rewriter, receiveOp->getLoc(), hostWaitLoad.getEventRegister(),
+          hostWaitLoad.getWaitValue());
+        received = PimMemCopyHostToDevOp::create(
+          rewriter, receiveOp->getLoc(), outputBuffer.getType(), zero,
+          hostWaitLoad.getHostOffset(), outputBuffer, *hostBuffer, *sizeAttr)
           .getOutput();
+        PimSyncOp::create(
+          rewriter, receiveOp->getLoc(), hostWaitLoad.getSourceCoreId(),
+          hostWaitLoad.getAcknowledgementEventRegister());
+      } else {
+        received = PimReceiveOp::create(
+          rewriter, receiveOp->getLoc(), outputBuffer.getType(), outputBuffer,
+          zero, *sizeAttr, channelReceive.getSourceCoreId()).getOutput();
+      }
       blockArg->replaceAllUsesWith(received);
       markOpToRemove(receiveOp);
       continue;
@@ -374,11 +405,12 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatScheduledCom
     if (rangeLength(resultUses) == 1) {
       OpOperand& resultUse = *resultUses.begin();
       Operation* resultUser = resultUse.getOwner();
-      if (isa<spatial::SpatChannelSendOp>(resultUser))
+      if (isa<spatial::SpatChannelSendOp,
+              spatial::SpatHostStoreSyncOp>(resultUser))
         continue;
     }
 
-    return computeOp.emitOpError("has an unsupported remaining result use during Spatial-to-PIM lowering");
+    return computeOp.emitOpError("has an unsupported remaining result use during Spatial-to-Pim lowering");
   }
 
   rewriter.setInsertionPoint(yieldOp);
@@ -388,7 +420,7 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatScheduledCom
   if (!computeOp.getWeights().empty())
     computeWeights.append(computeOp.getWeights().begin(), computeOp.getWeights().end());
   rewriter.setInsertionPointAfter(computeOp);
-  auto checkedCoreId = getRequiredScheduledCoreId(computeOp, "spatial compute core id");
+  auto checkedCoreId = getRequiredScheduledCoreId(computeOp, "Spatial compute core id");
   if (failed(checkedCoreId))
     return failure();
   auto coreIdAttr = pim::getCheckedI32Attr(rewriter, computeOp, static_cast<int64_t>(*checkedCoreId), "pim core id");

src/PIM/Conversion/SpatialToPim/Patterns.cpp

@@ -22,8 +22,7 @@ struct LowerFragmentAssemblyBlueprintPattern
   LogicalResult matchAndRewrite(spatial::SpatBlueprintOp op,
                                 OpAdaptor adaptor,
                                 ConversionPatternRewriter& rewriter) const override {
-    std::optional<StringRef> modeAttr = op.getMode();
-    if (!modeAttr || *modeAttr != "fragment_assembly")
+    if (!spatial::isFragmentAssembly(op.getMode()))
       return failure();
 
     auto resultType = dyn_cast<ShapedType>(op.getOutput().getType());
@@ -49,6 +48,16 @@ struct LowerFragmentAssemblyBlueprintPattern
           op, rank, fragmentOperands.size(), operandIndices, sourceOffsets, flatOffsets, flatSizes, flatStrides)))
       return failure();
 
+    if (op.getIndexMap() == spatial::kContiguousRowMajorFragments) {
+      if (!spatial::isCanonicalContiguousRowMajorFragmentAssembly(op))
+        return op.emitOpError("contiguous row-major fragment physical source order or storage is not canonical");
+      auto reshaped = reshapeContiguousRowMajorFragments(
+        rewriter, op.getLoc(), adaptor.getInput(), cast<RankedTensorType>(resultType));
+      if (failed(reshaped))
+        return op.emitOpError("contiguous row-major fragment storage does not match its logical result");
+      rewriter.replaceOp(op, *reshaped);
+      return success();
+    }
     Value currentOutput =
       tensor::EmptyOp::create(rewriter, op.getLoc(), resultType.getShape(), resultType.getElementType()).getResult();
     for (int64_t fragmentIndex = 0; fragmentIndex < static_cast<int64_t>(operandIndices.size()); ++fragmentIndex) {