fix scheduling reasoning on physical cores directly

document sync experiment

.agents/invariants/PIMSIM_NN_ORACLE_INVARIANT.md

@@ -0,0 +1,26 @@
+# 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.
+- Validation must report an explicit unsupported-op diagnostic as
+  `UNSUPPORTED`, not as a simulator failure.
+- A compiled artifact with zero active PIM cores has no meaningful
+  non-functional simulation and must be reported as `SKIP`.
+- 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/PIM_SYNCHRONIZATION_INVARIANT.md

@@ -0,0 +1,105 @@
+# PIM Synchronization Invariant
+
+## Scope
+
+This invariant applies to PIM `SYNC`/`WAIT` planning, event-register
+allocation, repeating core bodies, host/global-memory transfers, and pipeline
+restart synchronization.
+
+## ISA contract
+
+`SYNC event, target` and `WAIT event, expected` have static operands. Event
+registers start at zero, persist across instruction-stream restart, and are
+incremented by `SYNC`. `WAIT` succeeds only when the register equals its
+static expected count, then resets it to zero. A plan must therefore guarantee
+exactly the expected number of increments before each wait. It must not use
+generation-dependent registers, treat `WAIT` as `>=`, permit overshoot, or
+consume one event register with independent waits in the same iteration.
+
+Event registers are local to the target core. Allocation must keep READY,
+FREE/reuse, stage-zero barrier, and downstream restart events disjoint on each
+physical core.
+
+## Host/global-memory lifetime
+
+For every host-routed dependency from writer `W` to reader `R`, the repeating
+program must establish:
+
+```text
+ST_N -> READY SYNC -> READY WAIT -> LD_N
+LD_N -> reader END_N -> reader start_N+1 -> FREE SYNC
+     -> writer FREE WAIT -> ST_N+1
+```
+
+The reader emits one FREE `SYNC` per unique physical `(W, R)` relation at body
+entry, before any instruction that can block. The writer consumes each FREE
+event at most once per iteration, with a static count equal to the unique
+readers in that reuse group. The wait is placed immediately before the first
+host store that can overwrite a member of the group, after all preceding
+useful computation. Multiple transfers on one `(W, R)` relation do not create
+additional FREE signals.
+
+Relations with the same writer and first dangerous overwrite boundary belong
+to one FREE group: splitting them cannot unblock an earlier store and only
+adds waits and register pressure. After preserving the READY registers required
+by a core, planning assigns independent groups to distinct overwrite boundaries
+as remaining capacity permits. Capacity-forced merged groups are ordered by
+first overwrite and waited before their earliest member. Unused READY capacity
+must not force all readers into one early aggregate wait.
+
+Host receives lower to READY `WAIT` followed by host load only. They must not
+emit a post-load acknowledgement. Writers must not wait for host readers at
+the end of their body. The repeated reader-entry FREE signal is both the
+initial-slot bootstrap and the release of the previous generation.
+
+## Pipeline restart independence
+
+Host lifetime synchronization does not replace pipeline restart
+synchronization. At the end of every repeating stage-zero body, all unique
+stage-zero physical cores execute a dissemination barrier with
+`ceil(log2(stageZeroCoreCount))` rounds. In round `r`, rank `i` sends one
+signal to rank `(i + 2^r) mod count` and waits for exactly one signal on that
+round's destination-local event register. Each round has its own register, so
+an early signal for a later round cannot overshoot or satisfy another wait.
+
+Every downstream core sends one restart-permission signal to its release-tree
+parent at body entry, after all host FREE signals. The root signals the
+stage-zero leader; other cores signal their binary-tree parent. After the barrier, the lowest-ranked
+stage-zero core waits for the root's permission before sending one restart
+signal to it. Every downstream parent waits for exactly one restart signal
+and for one permission from each existing child before forwarding the restart
+to those children and ending its current body. This reverse permission path
+prevents a fast parent from sending generation `N+1` before a child consumes
+generation `N`, which would overshoot an exact-count restart event.
+
+Thus no downstream core can restart and advance its input generation until
+every stage-zero core has completed the protected generation, and every
+restart event receives exactly one signal between waits. Stage-zero and
+downstream core sets are disjoint; barrier, restart, and restart-permission
+registers are reserved on their respective target cores. Host READY/FREE
+allocation must use only the remaining registers and must not move, merge,
+remove, or weaken this ordering.
+
+## Ownership and verification
+
+Deferred boundary planning owns relation deduplication, event allocation,
+expected counts, and the first dangerous overwrite action. Boundary
+realization owns body-entry FREE signals and materializes planned waits at
+their exact boundaries. Channel lowering preserves those explicit operations
+and lowers host receives as READY wait plus load; it must not reconstruct
+lifetime policy after scheduling information is lost.
+
+Verification must reject non-static lane/core mappings, duplicate physical
+reader representations, register collisions, oversubscribed event capacity,
+and any plan whose generated FREE signal count differs from its static wait
+count. Static and simulator deadlock detectors must model blocking WAITs,
+their remaining matching SYNC producers, missing signals, and exact-count
+overshoot in addition to SEND/RECV cycles. Both detectors must retain the
+per-source contributions accumulated since the last successful WAIT and must not
+add a wait-for edge to a source that already supplied its share of the current
+event value merely because that source has the same static SYNC in a future
+iteration. Structural tests must inspect generated instruction streams. Happens-
+before tests must unroll at least two logical iterations and add only
+same-core restart edges plus matched communication edges. Functional
+simulation must enforce exact-count waits and include stalled or randomized
+legal schedules so correctness never depends on relative core speed.

.agents/invariants/PIPELINE_SCHEDULING_INVARIANT.md

@@ -0,0 +1,56 @@
+# 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.
+
+PEFT assigns physical cores using the injected target topology and their actual
+resident-weight capacity. Pipeline stage capacities cover every physical core
+and may have different sizes when the partitioner predicts a lower maximum
+stage interval. Before stage-local packing, the layout groups physical cores
+by bidirectional link cost; stage cores need not have consecutive IDs. Packing
+scores transfers on those physical links, and communication realization must
+preserve the resulting core identities.
+Synchronization and deferred transfers consume the explicit stage identity;
+they must not infer it from a physical core number after placement.
+
+## Ownership
+
+Physical PEFT supplies the initial placement. Stage partitioning is the first
+phase of pipeline scheduling and owns stage adjacency. 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 physical 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,9 @@ 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/PIM_SYNCHRONIZATION_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`).
@@ -53,20 +53,19 @@ ONNX-MLIR -> Spatial -> Pim (tensor) -> Pim (bufferized) -> PIM artifacts
    elementwise Add/Mul/Div, ReduceMean, pooling, Relu, Sigmoid, Softmax,
    Concat, Gather, Reshape, Resize, and Split.
    The compiler-layer target adapter supplies the target-neutral
-   `SpatialTargetInfo`. Layout-aware plan ops advertise typed alternatives
+   `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.
+   separate layers. Passes and their invariant/layout analyses live under
+   `Passes/Transforms` and `Passes/Analyses`.
 
 2. **Merge, schedule, and realize Spatial communication**
-   (`src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes`).
-   `TrivialGraphComputeMerge` performs local graph merging, then
-   `ScheduleSpatialGraph` materializes scheduled computes and explicit deferred
-   communication. `VerifyScheduledSpatial` checks that intermediate contract;
-   `RealizeSpatialCommunication` resolves transfers and forwarding; and
-   `VerifyRealizedSpatial` checks the final scheduled graph. Supporting
+   (`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`).
@@ -74,7 +73,7 @@ ONNX-MLIR -> Spatial -> Pim (tensor) -> Pim (bufferized) -> PIM artifacts
    including `pim.core`, `pim.core_batch`, communication, tensor packing, global
    tensor materialization, and return-path normalization.
 
-4. **Bufferization** (`src/PIM/Dialect/Pim/Transforms/Bufferization`).
+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,
@@ -82,66 +81,118 @@ ONNX-MLIR -> Spatial -> Pim (tensor) -> Pim (bufferized) -> PIM artifacts
    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`.
-- `--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:
 
@@ -155,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:
 
@@ -165,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.
@@ -282,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.
@@ -290,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};
-use std::io::{BufReader, Write};
+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,
@@ -114,7 +326,7 @@ fn map_crossbars_to_cores<'c>(
                 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(),
             );
@@ -174,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(())
 }
@@ -197,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() {
@@ -237,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()
@@ -310,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/mod.rs

@@ -2,17 +2,48 @@ use crate::utility::AddressArg;
 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> {
@@ -25,9 +56,320 @@ impl<'a> CPU<'a> {
         }
         Self {
             cores: cores.into(),
+            batch_outputs: None,
+            provenance: None,
         }
     }
 
+    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(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,

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::{
@@ -200,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)
@@ -285,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 ---
 
@@ -323,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)
@@ -389,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)
 }
@@ -474,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)
 }
 
@@ -780,6 +809,15 @@ where
         );
     }
     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)
 }
 
@@ -799,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)
 }
@@ -819,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)
 }
@@ -850,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);
@@ -860,6 +908,8 @@ 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)
 }
@@ -881,7 +931,12 @@ pub fn isa_recv(functor: usize) -> bool {
 
 #[inline(never)]
 pub fn recv(cores: &mut CPU, data: InstructionData) -> Result<InstructionStatus> {
-    Ok(InstructionStatus::Reciving(data))
+    Ok(InstructionStatus::Receiving(data))
+}
+
+#[inline(never)]
+pub fn isa_wait(functor: usize) -> bool {
+    (wait as *const () as usize) == functor
 }
 
 #[inline(never)]
@@ -889,6 +944,11 @@ pub fn wait(cores: &mut CPU, data: InstructionData) -> Result<InstructionStatus>
     Ok(InstructionStatus::Waiting(data))
 }
 
+#[inline(never)]
+pub fn isa_sync(functor: usize) -> bool {
+    (sync as *const () as usize) == functor
+}
+
 #[inline(never)]
 pub fn sync(cores: &mut CPU, data: InstructionData) -> Result<InstructionStatus> {
     Ok(InstructionStatus::Sync(data))

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

@@ -22,7 +22,7 @@ pub enum InstructionStatus {
     Completed,
     Waiting(InstructionData),
     Sending(InstructionData),
-    Reciving(InstructionData),
+    Receiving(InstructionData),
     Sync(InstructionData),
     #[default]
     NotExecuted,
@@ -59,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,
 }
 
@@ -86,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(),
         }
@@ -97,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;
         }
@@ -107,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/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,10 +33,10 @@ 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(),
         }
     }
 }
@@ -73,18 +73,27 @@ where
                 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;
@@ -124,19 +133,19 @@ where
             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
@@ -153,17 +162,17 @@ where
             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/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/sync.rs

@@ -295,3 +295,122 @@ fn multiple_send_recv_test() {
         "send_recv failed to store"
     );
 }
+
+#[test]
+fn sync_wait_exact_count_resets_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();
+    instructions.make_inst(wait, data.set_offset_select_value(0, 2).build());
+    cores.set_core(2, instructions.build());
+
+    Executable::new(cpu, cores.build()).execute().unwrap();
+}
+
+#[test]
+fn sync_wait_rejects_overshoot() {
+    let cpu = common::empty_cpu(3);
+    let mut cores = CoreInstructionsBuilder::new(3);
+    let mut instructions = InstructionsBuilder::new();
+    let mut data = InstructionDataBuilder::new();
+
+    data.set_core_indx(1).fix_core_indx();
+    instructions.make_inst(
+        sync,
+        data.set_imm_core(3).set_offset_select_value(0, 0).build(),
+    );
+    cores.set_core(1, instructions.build());
+
+    data.set_core_indx(2).fix_core_indx();
+    instructions.make_inst(
+        sync,
+        data.set_imm_core(3).set_offset_select_value(0, 0).build(),
+    );
+    cores.set_core(2, instructions.build());
+
+    data.set_core_indx(3).fix_core_indx();
+    instructions.make_inst(wait, data.set_offset_select_value(0, 1).build());
+    cores.set_core(3, instructions.build());
+
+    let error = Executable::new(cpu, cores.build()).execute().unwrap_err();
+    assert!(error.to_string().contains("overshot exact value"));
+}
+
+#[test]
+fn sync_wait_deadlock_cycle_is_reported() {
+    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();
+    instructions.make_inst(wait, data.set_offset_select_value(0, 1).build());
+    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();
+    instructions.make_inst(wait, data.set_offset_select_value(0, 1).build());
+    instructions.make_inst(
+        sync,
+        data.set_imm_core(1).set_offset_select_value(0, 0).build(),
+    );
+    cores.set_core(2, instructions.build());
+
+    let error = Executable::new(cpu, cores.build()).execute().unwrap_err();
+    assert!(error.to_string().contains("wait event"));
+}
+
+#[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)
 

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

@@ -125,7 +125,7 @@ static bool isZeroSplatGlobal(mlir::Value value) {
   return false;
 }
 
-// PIM instruction immediates are serialized as signed int32_t fields today
+// 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) {
@@ -141,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";
@@ -187,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;
@@ -276,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;
@@ -352,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,
@@ -544,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,
@@ -664,13 +676,13 @@ void PimCodeGen::codeGenVMVOp(pim::PimVMVOp vmvOp, const StaticValueKnowledge& k
   auto sourceType = cast<ShapedType>(vmvOp.getSource().getType());
   int32_t bitwidth = getVectorElementBitwidthOrCrash(sourceType);
   ensureVectorBitwidth(bitwidth, bitwidth);
-  setupRdRs1Rs2(addressOf(vmvOp.getTarget(), knowledge), *targetOffset,
-                addressOf(vmvOp.getSource(), knowledge), *sourceOffset, 0, *sourceStride);
+  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 = 0;
-  instruction.r1 = 1;
-  instruction.r2OrImm = 2;
+  instruction.rd = registers[0];
+  instruction.r1 = registers[1];
+  instruction.r2OrImm = registers[2];
   instruction.generic3 = vmvOp.getLength();
   emitInstruction(instruction);
 }
@@ -692,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");
@@ -749,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);
 }
@@ -914,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();
       }
 
@@ -940,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();
       }
 
@@ -954,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();
@@ -991,6 +1039,8 @@ static LogicalResult executeCompiledCorePlan(
     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))
@@ -1138,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 +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();
@@ -1229,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());
@@ -1341,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)
@@ -1407,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

@@ -176,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,
@@ -217,6 +217,8 @@ public:
 
   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>

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"),
@@ -87,6 +84,11 @@ llvm::cl::opt<uint64_t> pimConvStreamChunkPositions(
   llvm::cl::init(1024),
   llvm::cl::cat(OnnxMlirOptions));
 
+llvm::cl::opt<bool> pimReportConvLowering("pim-report-conv-lowering",
+                                          llvm::cl::desc("Emit a bounded Conv lowering report"),
+                                          llvm::cl::init(true),
+                                          llvm::cl::cat(OnnxMlirOptions));
+
 llvm::cl::opt<bool> pimEmitJson("pim-emit-json",
                                 llvm::cl::desc("Also emit per-core JSON instruction files alongside binary .pim files"),
                                 llvm::cl::init(false),
@@ -94,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 statically expanded Pim SEND/RECV and exact-count SYNC/WAIT order at verification time and fail on a blocking deadlock"),
   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,28 +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,13 +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/SpatialTargetInfo.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledSpatialPasses.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"
@@ -77,6 +79,7 @@ spatial::SchedulingTarget getDefaultPimSchedulingTarget() {
   target.residentWeightCapacity = crossbarCountInCore.getValue();
   target.matrixRows = crossbarSize.getValue();
   target.matrixColumns = crossbarSize.getValue();
+  target.synchronizationRegisterCount = kPimEventRegisterCount;
 
   setDefaultPimInterProcessorLatencies(target);
   return target;
@@ -94,20 +97,40 @@ spatial::ConvLoweringStrategy getSpatialConvLoweringStrategy(PimConvLoweringType
   case PimConvLoweringInputKTiled: return spatial::ConvLoweringStrategy::InputKTiled;
   case PimConvLoweringTiled2D: return spatial::ConvLoweringStrategy::Tiled2D;
   }
-  llvm_unreachable("unknown PIM Conv lowering strategy");
+  llvm_unreachable("unknown Pim Conv lowering strategy");
 }
 
-spatial::SpatialTargetInfo getPimSpatialTargetInfo(const spatial::SchedulingTarget& target) {
-  spatial::SpatialTargetInfo info;
-  info.matrixShape = {target.matrixRows, target.matrixColumns};
-  info.matrixUnitsPerProcessor = target.residentWeightCapacity;
-  info.processorCount = target.processorCount;
-  info.vectorWidth = target.vectorWidth;
-  info.convIm2colMaxElements = pimConvIm2colMaxElements.getValue();
-  info.convStreamChunkPositions = pimConvStreamChunkPositions.getValue();
-  info.convLoweringStrategy = getSpatialConvLoweringStrategy(pimConvLowering.getValue());
-  info.useExperimentalConvImplementation = useExperimentalConvImpl.getValue();
-  return info;
+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,
@@ -115,7 +138,7 @@ const llvm::json::Object& requireObject(const llvm::json::Object& object,
                                         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;
 }
 
@@ -128,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));
 }
 
@@ -136,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)};
 }
 
@@ -151,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)) {
@@ -164,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);
@@ -187,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) {
@@ -197,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[
@@ -221,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");
@@ -244,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);
@@ -255,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);
 
@@ -308,12 +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;
   }
@@ -322,28 +347,27 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
     addONNXToMLIRPasses(pm, /*target CPU*/ false);
 
   if (pimEmissionTarget >= EmitSpatial) {
-    spatial::SchedulingTarget schedulingTarget = getPimSchedulingTarget();
-    spatial::SpatialTargetInfo targetInfo = getPimSpatialTargetInfo(schedulingTarget);
-    pm.addPass(createONNXToSpatialPass(targetInfo));
-    pm.addPass(createSpatialLayoutPlanningPass(targetInfo));
-    pm.addPass(createLowerSpatialPlansPass(targetInfo));
+    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));
-    auto scheduledState = std::make_shared<spatial::ScheduledSpatialState>();
-    pm.addPass(spatial::createScheduleSpatialGraphPass(schedulingTarget, scheduledState));
-    pm.addPass(spatial::createVerifyScheduledSpatialPass(scheduledState));
-    pm.addPass(spatial::createRealizeSpatialCommunicationPass(schedulingTarget, scheduledState));
-    pm.addPass(spatial::createVerifyRealizedSpatialPass(scheduledState));
+      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(createPimBufferizationPreparationPass());
+    pm.addPass(createPimBufferizationPreparationPass(pimVerifyBufferizationCopyFreedom.getValue()));
     pm.addPass(createPimOneShotBufferizationPass());
     pm.addPass(createPimMemoryNormalizationPass());
     pm.addPass(createPimBufferizationVerificationPass());
@@ -358,7 +382,7 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
     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

@@ -17,6 +17,8 @@ static FailureOr<CompiledCoreOpKind> classifyCompiledCoreOpKind(Operation& op) {
   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::PimVVAddOp>(op)) return CompiledCoreOpKind::VVAdd;
@@ -44,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;
@@ -61,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;
@@ -80,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;

src/PIM/Compiler/PimCoreProgram.hpp

@@ -17,6 +17,8 @@ enum class CompiledCoreOpKind : uint8_t {
   VMV,
   Receive,
   Send,
+  Wait,
+  Sync,
   Concat,
   Vmm,
   VVAdd,

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,14 +26,14 @@ add_pim_library(OMONNXToSpatial
   Patterns/Tensor/Slice.cpp
   Patterns/Tensor/Split.cpp
   Patterns/Tensor/Transpose.cpp
-  ONNXToSpatialPass.cpp
-  SpatialLayoutCapabilities.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/ContractionMaterialization.cpp
   Common/ContractionPlanning.cpp
   Common/MatrixProductLowering.cpp
   Common/RowStripLayoutUtils.cpp

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

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

@@ -1,39 +0,0 @@
-#include "ContractionMaterialization.hpp"
-
-#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
-#include "MatrixProductLowering.hpp"
-
-namespace onnx_mlir {
-
-mlir::Value materializePaddedContractionInput(
-    mlir::Value input,
-    mlir::RankedTensorType paddedType,
-    mlir::PatternRewriter& rewriter,
-    mlir::Location loc) {
-  return createPaddedInputCompute(input, paddedType, rewriter, loc);
-}
-
-mlir::FailureOr<mlir::Value> materializeTransposedContractionConstant(
-    mlir::Value input,
-    mlir::RankedTensorType resultType,
-    llvm::ArrayRef<int64_t> permutation,
-    mlir::PatternRewriter& rewriter,
-    mlir::Location loc) {
-  auto denseAttr = getHostConstDenseElementsAttr(input);
-  auto inputType = denseAttr ? mlir::dyn_cast<mlir::RankedTensorType>(denseAttr.getType()) : nullptr;
-  if (!inputType || !inputType.hasStaticShape() || !resultType || !resultType.hasStaticShape()
-      || inputType.getRank() != resultType.getRank())
-    return mlir::failure();
-
-  auto transposedAttr = transposeDenseElementsAttr(denseAttr, permutation);
-  if (mlir::failed(transposedAttr) || transposedAttr->getType() != resultType)
-    return mlir::failure();
-
-  return getOrCreateConstant(rewriter,
-                             rewriter.getInsertionBlock()->getParentOp(),
-                             *transposedAttr,
-                             resultType);
-}
-
-} // namespace onnx_mlir

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

@@ -1,23 +0,0 @@
-#pragma once
-
-#include "llvm/ADT/ArrayRef.h"
-
-#include "mlir/IR/BuiltinTypes.h"
-#include "mlir/IR/PatternMatch.h"
-
-namespace onnx_mlir {
-
-mlir::Value materializePaddedContractionInput(
-    mlir::Value input,
-    mlir::RankedTensorType paddedType,
-    mlir::PatternRewriter& rewriter,
-    mlir::Location loc);
-
-mlir::FailureOr<mlir::Value> materializeTransposedContractionConstant(
-    mlir::Value input,
-    mlir::RankedTensorType resultType,
-    llvm::ArrayRef<int64_t> permutation,
-    mlir::PatternRewriter& rewriter,
-    mlir::Location loc);
-
-} // namespace onnx_mlir

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

@@ -12,56 +12,30 @@ static int64_t ceilDivide(int64_t value, int64_t divisor) {
   return divisor == 0 ? 0 : (value + divisor - 1) / divisor;
 }
 
-static llvm::SmallVector<int64_t> buildBatchMap(
-    llvm::ArrayRef<int64_t> sourceShape,
-    llvm::ArrayRef<int64_t> outputShape) {
-  llvm::SmallVector<int64_t> map(outputShape.size(), -1);
-  const int64_t offset = outputShape.size() - sourceShape.size();
-  for (int64_t source = 0; source < static_cast<int64_t>(sourceShape.size()); ++source) {
-    const int64_t output = source + offset;
-    if (sourceShape[source] != 1)
-      map[output] = source;
-  }
-  return map;
-}
-
 } // namespace
 
 ContractionPlan makeContractionPlan(
     const ContractionProblem& problem,
-    const spatial::SpatialTargetInfo& target,
+    const spatial::SpatialTargetResources& target,
     ContractionPlanKind kind,
     int64_t laneCount,
     int64_t fragmentRows) {
   ContractionPlan plan;
-  plan.problem = problem;
-  plan.kind = kind;
-  plan.tileM = std::max<int64_t>(1, target.matrixShape.rows);
   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));
-  plan.rowTiles = std::max<int64_t>(1, ceilDivide(problem.m, plan.tileM));
-  plan.fragmentRows = std::max<int64_t>(
-      1, fragmentRows != 0 ? fragmentRows : plan.tileM);
-  plan.lhsBatchMap = buildBatchMap(problem.lhsBatchShape, problem.outputBatchShape);
-  plan.rhsBatchMap = buildBatchMap(problem.rhsBatchShape, problem.outputBatchShape);
+  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, plan.fragmentRows);
+    plan.laneCount = problem.batch * ceilDivide(problem.m, rowsPerLane);
   else
     plan.laneCount = problem.batch * problem.m * problem.n;
-
-  plan.expectedMvmCount = kind == ContractionPlanKind::StaticTiled ? plan.laneCount : 0;
-  plan.expectedVvdCount = kind == ContractionPlanKind::StaticTiled ? 0 : plan.laneCount;
-  plan.expectedVectorCount = plan.laneCount * plan.reductionSlices;
-  if (problem.resultElementType && problem.n > 0)
-    plan.physicalFragmentType = mlir::RankedTensorType::get(
-        {plan.fragmentRows, problem.n}, problem.resultElementType);
   return plan;
 }
 

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

@@ -2,7 +2,7 @@
 
 #include "ContractionProblem.hpp"
 
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetInfo.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
 
 namespace onnx_mlir {
 
@@ -13,27 +13,16 @@ enum class ContractionPlanKind {
 };
 
 struct ContractionPlan {
-  ContractionProblem problem;
-  ContractionPlanKind kind = ContractionPlanKind::StaticTiled;
-  int64_t tileM = 1;
   int64_t tileK = 1;
   int64_t tileN = 1;
-  int64_t fragmentRows = 1;
   int64_t reductionSlices = 1;
   int64_t outputTiles = 1;
-  int64_t rowTiles = 1;
   int64_t laneCount = 0;
-  int64_t expectedMvmCount = 0;
-  int64_t expectedVvdCount = 0;
-  int64_t expectedVectorCount = 0;
-  llvm::SmallVector<int64_t> lhsBatchMap;
-  llvm::SmallVector<int64_t> rhsBatchMap;
-  mlir::RankedTensorType physicalFragmentType;
 };
 
 ContractionPlan makeContractionPlan(
     const ContractionProblem& problem,
-    const spatial::SpatialTargetInfo& target,
+    const spatial::SpatialTargetResources& target,
     ContractionPlanKind kind,
     int64_t laneCount = 0,
     int64_t fragmentRows = 0);

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

@@ -8,8 +8,6 @@
 
 namespace onnx_mlir {
 
-enum class ContractionOrigin { Gemm, MatMul };
-
 struct ContractionProblem {
   llvm::SmallVector<int64_t> lhsBatchShape;
   llvm::SmallVector<int64_t> rhsBatchShape;
@@ -20,16 +18,9 @@ struct ContractionProblem {
   int64_t m = 0;
   int64_t k = 0;
   int64_t n = 0;
-  ContractionOrigin origin = ContractionOrigin::MatMul;
   mlir::Type lhsElementType;
   mlir::Type rhsElementType;
   mlir::Type resultElementType;
-  bool lhsTransposed = false;
-  bool rhsTransposed = false;
-  bool lhsWasVector = false;
-  bool rhsWasVector = false;
-  float alpha = 1.0f;
-  float beta = 1.0f;
 };
 
 } // namespace onnx_mlir

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

@@ -8,8 +8,8 @@ namespace onnx_mlir {
 
 namespace spatial {
 class SpatBlueprintOp;
-class SpatGraphCompute;
-struct SpatialTargetInfo;
+class SpatFlattenPlanOp;
+struct SpatialTargetResources;
 } // namespace spatial
 
 inline constexpr llvm::StringLiteral kRowStripIndexMap = "nhwc_row_strip_fragments";
@@ -94,13 +94,13 @@ mlir::FailureOr<mlir::Value> applyRowStripConcat(llvm::ArrayRef<RowStripPhysical
                                                  mlir::Location loc);
 
 mlir::LogicalResult canLowerFlattenFromRowStrip(
-    spatial::SpatGraphCompute flattenOp,
-    const spatial::SpatialTargetInfo& target);
+    spatial::SpatFlattenPlanOp flattenOp,
+    const spatial::SpatialTargetResources& target);
 
 mlir::LogicalResult lowerFlattenFromRowStrip(
     const RowStripPhysicalValue& input,
-    spatial::SpatGraphCompute flattenOp,
-    const spatial::SpatialTargetInfo& target,
+    spatial::SpatFlattenPlanOp flattenOp,
+    const spatial::SpatialTargetResources& target,
     mlir::PatternRewriter& rewriter);
 
 } // namespace onnx_mlir

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

@@ -69,7 +69,7 @@ DenseMap<CoreId, SmallVector<Value>>
 sliceVectorPerCrossbarPerCore(const Value& vectorToSlice,
                               PatternRewriter& rewriter,
                               Location loc,
-                              const spatial::SpatialTargetInfo& target) {
+                              const spatial::SpatialTargetResources& target) {
   SmallVector<Value> slices = sliceVector(
     vectorToSlice, static_cast<int64_t>(target.matrixShape.rows), rewriter, loc);
   DenseMap<CoreId, SmallVector<Value>> slicesPerCore;

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

@@ -7,7 +7,7 @@
 #include "llvm/ADT/SmallVector.h"
 
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetInfo.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
 
 namespace onnx_mlir {
 
@@ -25,11 +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 spatial::SpatialTargetInfo& target);
+  const spatial::SpatialTargetResources& target);
 
 } // 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;
@@ -147,6 +149,7 @@ void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter
             spatial::SpatBiasAddPlanOp,
             spatial::SpatAddPlanOp,
             spatial::SpatConcatPlanOp,
+            spatial::SpatFlattenPlanOp,
             spatial::SpatReluPlanOp,
             spatial::SpatSiluPlanOp,
             spatial::SpatResizeNearestPlanOp,
@@ -162,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";
@@ -181,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

@@ -1,5 +1,6 @@
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.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;
@@ -28,7 +29,7 @@ static bool hasRowStripInput(ArrayRef<PhysicalLayout> operandLayouts, unsigned i
 }
 
 SmallVector<LayoutAlternative> SpatConv2DPlanOp::getLayoutAlternatives(
-    const SpatialTargetInfo& target, ArrayRef<PhysicalLayout> operandLayouts) {
+    const SpatialTargetResources& target, ArrayRef<PhysicalLayout> operandLayouts) {
   SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
   if (hasRowStripInput(operandLayouts, 0)) {
     if (succeeded(canConsumeAndProduceRowStrip(*this, target)))
@@ -43,8 +44,26 @@ SmallVector<LayoutAlternative> SpatConv2DPlanOp::getLayoutAlternatives(
   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 SpatialTargetInfo&, ArrayRef<PhysicalLayout> operandLayouts) {
+    const SpatialTargetResources&, ArrayRef<PhysicalLayout> operandLayouts) {
   SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
   if (hasRowStripInput(operandLayouts, 0))
     alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts));
@@ -52,7 +71,7 @@ SmallVector<LayoutAlternative> SpatReluPlanOp::getLayoutAlternatives(
 }
 
 SmallVector<LayoutAlternative> SpatSiluPlanOp::getLayoutAlternatives(
-    const SpatialTargetInfo&, ArrayRef<PhysicalLayout> operandLayouts) {
+    const SpatialTargetResources&, ArrayRef<PhysicalLayout> operandLayouts) {
   SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
   if (hasRowStripInput(operandLayouts, 0)) {
     LayoutAlternative alternative = rowStripAlternative(getOperation(), operandLayouts);
@@ -63,7 +82,7 @@ SmallVector<LayoutAlternative> SpatSiluPlanOp::getLayoutAlternatives(
 }
 
 SmallVector<LayoutAlternative> SpatResizeNearestPlanOp::getLayoutAlternatives(
-    const SpatialTargetInfo& target, ArrayRef<PhysicalLayout> operandLayouts) {
+    const SpatialTargetResources& target, ArrayRef<PhysicalLayout> operandLayouts) {
   SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
   if (hasRowStripInput(operandLayouts, 0)
       && succeeded(canLowerResizeNearestPlanToRowStrip(*this, target)))
@@ -72,7 +91,7 @@ SmallVector<LayoutAlternative> SpatResizeNearestPlanOp::getLayoutAlternatives(
 }
 
 SmallVector<LayoutAlternative> SpatMaxPool2DPlanOp::getLayoutAlternatives(
-    const SpatialTargetInfo& target, ArrayRef<PhysicalLayout> operandLayouts) {
+    const SpatialTargetResources& target, ArrayRef<PhysicalLayout> operandLayouts) {
   SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
   if (succeeded(canLowerMaxPoolPlanToRowStrip(*this, target))) {
     LayoutAlternative alternative = denseAlternative(getOperation());
@@ -86,7 +105,7 @@ SmallVector<LayoutAlternative> SpatMaxPool2DPlanOp::getLayoutAlternatives(
 }
 
 SmallVector<LayoutAlternative> SpatGlobalAveragePoolPlanOp::getLayoutAlternatives(
-    const SpatialTargetInfo& target, ArrayRef<PhysicalLayout> operandLayouts) {
+    const SpatialTargetResources& target, ArrayRef<PhysicalLayout> operandLayouts) {
   SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
   if (succeeded(canLowerGlobalAveragePoolPlanToRowStrip(*this, target))) {
     LayoutAlternative alternative = denseAlternative(getOperation());
@@ -100,7 +119,7 @@ SmallVector<LayoutAlternative> SpatGlobalAveragePoolPlanOp::getLayoutAlternative
 }
 
 SmallVector<LayoutAlternative> SpatBiasAddPlanOp::getLayoutAlternatives(
-    const SpatialTargetInfo&, ArrayRef<PhysicalLayout> operandLayouts) {
+    const SpatialTargetResources&, ArrayRef<PhysicalLayout> operandLayouts) {
   SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
   auto resultType = dyn_cast<RankedTensorType>(getOutput().getType());
   if (resultType && hasRowStripInput(operandLayouts, 0)
@@ -112,7 +131,7 @@ SmallVector<LayoutAlternative> SpatBiasAddPlanOp::getLayoutAlternatives(
 }
 
 SmallVector<LayoutAlternative> SpatAddPlanOp::getLayoutAlternatives(
-    const SpatialTargetInfo&, ArrayRef<PhysicalLayout> operandLayouts) {
+    const SpatialTargetResources&, ArrayRef<PhysicalLayout> operandLayouts) {
   SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
   if (operandLayouts.size() >= 2 && hasRowStripInput(operandLayouts, 0)
       && hasRowStripInput(operandLayouts, 1))
@@ -121,7 +140,7 @@ SmallVector<LayoutAlternative> SpatAddPlanOp::getLayoutAlternatives(
 }
 
 SmallVector<LayoutAlternative> SpatConcatPlanOp::getLayoutAlternatives(
-    const SpatialTargetInfo&, ArrayRef<PhysicalLayout> operandLayouts) {
+    const SpatialTargetResources&, ArrayRef<PhysicalLayout> operandLayouts) {
   SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
   if (!operandLayouts.empty() && llvm::all_of(operandLayouts, [](PhysicalLayout layout) {
         return layout == PhysicalLayout::NHWCRowStrip;

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

@@ -12,15 +12,15 @@
 
 #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;
 
@@ -34,14 +34,16 @@ struct ONNXToSpatialPass : PassWrapper<ONNXToSpatialPass, OperationPass<ModuleOp
   StringRef getDescription() const override { return "Lower ONNX ops to Spatial ops."; }
 
   ONNXToSpatialPass() = default;
-  explicit ONNXToSpatialPass(const spatial::SpatialTargetInfo& target)
-      : target(target), hasTarget(true) {}
+  ONNXToSpatialPass(const spatial::SpatialTargetResources& target,
+                    const ONNXToSpatialPlanningOptions& options)
+      : target(target), planningOptions(options), hasTarget(true) {}
   ONNXToSpatialPass(const ONNXToSpatialPass& pass)
-      : target(pass.target), hasTarget(pass.hasTarget) {}
+      : target(pass.target), planningOptions(pass.planningOptions), hasTarget(pass.hasTarget) {}
 
   void runOnOperation() override;
 
-  spatial::SpatialTargetInfo target;
+  spatial::SpatialTargetResources target;
+  ONNXToSpatialPlanningOptions planningOptions;
   bool hasTarget = false;
 };
 
@@ -56,6 +58,7 @@ 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(
@@ -66,7 +69,7 @@ static void populateEmptyFunction(func::FuncOp funcOp) {
   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() || !siluPlans.empty() || !resizePlans.empty()
+      || !concatPlans.empty() || !flattenPlans.empty() || !reluPlans.empty() || !siluPlans.empty() || !resizePlans.empty()
       || !maxPoolPlans.empty() || !blueprints.empty()
       || !globalAveragePoolPlans.empty() || !materializers.empty()) {
     return;
@@ -113,7 +116,7 @@ static void populateEmptyFunction(func::FuncOp funcOp) {
 void ONNXToSpatialPass::runOnOperation() {
   ModuleOp moduleOp = getOperation();
   if (!hasTarget) {
-    moduleOp.emitError("ONNX-to-Spatial lowering requires an injected SpatialTargetInfo");
+    moduleOp.emitError("ONNX-to-Spatial lowering requires an injected SpatialTargetResources");
     signalPassFailure();
     return;
   }
@@ -155,7 +158,7 @@ void ONNXToSpatialPass::runOnOperation() {
 
   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;
   }
@@ -193,7 +196,7 @@ void ONNXToSpatialPass::runOnOperation() {
   target.addIllegalOp<ONNXSplitOp>();
 
   RewritePatternSet conversionPatterns(ctx);
-  populateConversionPatterns(conversionPatterns, ctx, this->target);
+  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();
@@ -242,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;
   }
@@ -269,8 +272,10 @@ void ONNXToSpatialPass::runOnOperation() {
 
 std::unique_ptr<Pass> createONNXToSpatialPass() { return std::make_unique<ONNXToSpatialPass>(); }
 
-std::unique_ptr<Pass> createONNXToSpatialPass(const spatial::SpatialTargetInfo& target) {
-  return std::make_unique<ONNXToSpatialPass>(target);
+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

@@ -10,55 +10,83 @@
 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::SpatialTargetInfo& target,
+                        const spatial::SpatialTargetResources& target,
+                        const ONNXToSpatialPlanningOptions& options,
                         mlir::PatternRewriter& rewriter);
 
 mlir::LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp,
-                                               const spatial::SpatialTargetInfo& target);
+                                               const spatial::SpatialTargetResources& target,
+                                               const ONNXToSpatialPlanningOptions* options = nullptr);
 mlir::LogicalResult canConsumeAndProduceRowStrip(spatial::SpatConv2DPlanOp planOp,
-                                                 const spatial::SpatialTargetInfo& target);
+                                                 const spatial::SpatialTargetResources& target,
+                                                 const ONNXToSpatialPlanningOptions* options = nullptr);
 
 mlir::LogicalResult canLowerResizeNearestPlanToRowStrip(
-  spatial::SpatResizeNearestPlanOp planOp, const spatial::SpatialTargetInfo& target);
+  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::SpatialTargetInfo& target,
+  const spatial::SpatialTargetResources& target,
   mlir::PatternRewriter& rewriter);
 
 mlir::LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp,
-                                                 const spatial::SpatialTargetInfo& target);
+                                                 const spatial::SpatialTargetResources& target);
 
 mlir::FailureOr<mlir::Value>
 lowerDenseMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
-                        const spatial::SpatialTargetInfo& target,
+                        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::SpatialTargetInfo& target,
+                           const spatial::SpatialTargetResources& target,
                            mlir::PatternRewriter& rewriter);
 
 mlir::LogicalResult
 canLowerGlobalAveragePoolPlanToRowStrip(spatial::SpatGlobalAveragePoolPlanOp planOp,
-                                        const spatial::SpatialTargetInfo& target);
+                                        const spatial::SpatialTargetResources& target);
 
 mlir::FailureOr<mlir::Value>
 lowerDenseGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp,
-                                const spatial::SpatialTargetInfo& target,
+                                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::SpatialTargetInfo& target,
+                                   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

@@ -4,23 +4,21 @@
 #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/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/DialectConversion.h"
 
-#include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
-#include "mlir/Transforms/Passes.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/PlanLowering.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/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp"
-#include "src/Accelerators/PIM/Pass/PIMPasses.h"
+#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp"
+#include "src/Accelerators/PIM/Passes/PIMPasses.h"
 
 using namespace mlir;
 
@@ -59,19 +57,11 @@ static bool isDenseSelected(Operation* op) {
 }
 
 static spatial::PhysicalLayout getKnownPhysicalLayout(Value value) {
-  if (auto materialize = value.getDefiningOp<spatial::SpatMaterializeLayoutOp>())
-    return materialize.getTargetPhysicalLayout();
-  if (auto blueprint = value.getDefiningOp<spatial::SpatBlueprintOp>())
-    return blueprint.getPhysicalLayout();
-  if (Operation* producer = value.getDefiningOp()) {
-    if (auto selected = spatial::getSelectedPhysicalLayout(producer))
-      return *selected;
-  }
-  return spatial::PhysicalLayout::DenseNCHW;
+  return getSpatialPlanOperandLayout(value);
 }
 
 static LogicalResult verifySelectedLayouts(
-    func::FuncOp funcOp, const spatial::SpatialTargetInfo& target) {
+    func::FuncOp funcOp, const spatial::SpatialTargetResources& target) {
   LogicalResult result = success();
   funcOp.walk([&](Operation* op) {
     auto capability = dyn_cast<spatial::SpatialLayoutCapabilityInterface>(op);
@@ -115,12 +105,6 @@ lowerRowStripSilu(const RowStripPhysicalValue& input, spatial::SpatSiluPlanOp pl
   return applyRowStripSilu(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,
@@ -222,17 +206,18 @@ static FailureOr<Value> lowerDenseBatchBiasAdd(Value input, Value bias, RankedTe
   return batch->getResult(0);
 }
 
-struct LowerDenseReluPlan final : OpRewritePattern<spatial::SpatReluPlanOp> {
-  using OpRewritePattern::OpRewritePattern;
+struct LowerDenseReluPlan final : OpConversionPattern<spatial::SpatReluPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
 
   LogicalResult matchAndRewrite(spatial::SpatReluPlanOp planOp,
-                                PatternRewriter& rewriter) const override {
+                                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(), {}, planOp.getInput(), [&](Value x) {
+      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());
       });
@@ -241,17 +226,18 @@ struct LowerDenseReluPlan final : OpRewritePattern<spatial::SpatReluPlanOp> {
   }
 };
 
-struct LowerDenseSiluPlan final : OpRewritePattern<spatial::SpatSiluPlanOp> {
-  using OpRewritePattern::OpRewritePattern;
+struct LowerDenseSiluPlan final : OpConversionPattern<spatial::SpatSiluPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
 
   LogicalResult matchAndRewrite(spatial::SpatSiluPlanOp planOp,
-                                PatternRewriter& rewriter) const override {
+                                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(), {}, planOp.getInput(), [&](Value x) {
+      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(
@@ -263,29 +249,32 @@ struct LowerDenseSiluPlan final : OpRewritePattern<spatial::SpatSiluPlanOp> {
   }
 };
 
-struct LowerDenseResizePlan final : OpRewritePattern<spatial::SpatResizeNearestPlanOp> {
-  explicit LowerDenseResizePlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
-      : OpRewritePattern<spatial::SpatResizeNearestPlanOp>(ctx), target(target) {}
+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,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatResizeNearestPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
     if (!isDenseSelected(planOp.getOperation()))
       return failure();
-    FailureOr<Value> lowered = lowerSelectedResizeNearestPlan(planOp, std::nullopt, target, rewriter);
+    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::SpatialTargetInfo& target;
+  const spatial::SpatialTargetResources& target;
 };
 
-struct LowerDenseBiasAddPlan final : OpRewritePattern<spatial::SpatBiasAddPlanOp> {
-  using OpRewritePattern::OpRewritePattern;
+struct LowerDenseBiasAddPlan final : OpConversionPattern<spatial::SpatBiasAddPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
 
   LogicalResult matchAndRewrite(spatial::SpatBiasAddPlanOp planOp,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatBiasAddPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
     if (!isDenseSelected(planOp.getOperation()))
       return failure();
     auto resultType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
@@ -293,12 +282,12 @@ struct LowerDenseBiasAddPlan final : OpRewritePattern<spatial::SpatBiasAddPlanOp
       return planOp.emitOpError("requires ranked output type");
 
     FailureOr<Value> denseBias = materializeDenseBiasAddTensor(
-      planOp.getBias(), resultType, rewriter, planOp.getLoc());
+      adaptor.getBias(), resultType, rewriter, planOp.getLoc());
     if (failed(denseBias))
       return planOp.emitOpError("failed to materialize dense Conv-style bias");
-    if (planOp.getInput().getDefiningOp<spatial::SpatGraphComputeBatch>()) {
+    if (adaptor.getInput().getDefiningOp<spatial::SpatGraphComputeBatch>()) {
       FailureOr<Value> lowered = lowerDenseBatchBiasAdd(
-        planOp.getInput(), *denseBias, resultType, rewriter, planOp.getLoc());
+        adaptor.getInput(), *denseBias, resultType, rewriter, planOp.getLoc());
       if (succeeded(lowered)) {
         rewriter.replaceOp(planOp, *lowered);
         return success();
@@ -309,7 +298,7 @@ struct LowerDenseBiasAddPlan final : OpRewritePattern<spatial::SpatBiasAddPlanOp
       planOp.getLoc(),
       planOp.getOutput().getType(),
       {},
-      ValueRange {planOp.getInput(), *denseBias},
+      ValueRange {adaptor.getInput(), *denseBias},
       [&](Value x, Value y) {
         auto added = spatial::SpatVAddOp::create(
           rewriter, planOp.getLoc(), planOp.getOutput().getType(), x, y);
@@ -320,11 +309,12 @@ struct LowerDenseBiasAddPlan final : OpRewritePattern<spatial::SpatBiasAddPlanOp
   }
 };
 
-struct LowerDenseAddPlan final : OpRewritePattern<spatial::SpatAddPlanOp> {
-  using OpRewritePattern::OpRewritePattern;
+struct LowerDenseAddPlan final : OpConversionPattern<spatial::SpatAddPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
 
   LogicalResult matchAndRewrite(spatial::SpatAddPlanOp planOp,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatAddPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
     if (!isDenseSelected(planOp.getOperation()))
       return failure();
     auto compute = createSpatCompute<2>(
@@ -332,7 +322,7 @@ struct LowerDenseAddPlan final : OpRewritePattern<spatial::SpatAddPlanOp> {
       planOp.getLoc(),
       planOp.getOutput().getType(),
       {},
-      ValueRange {planOp.getLhs(), planOp.getRhs()},
+      ValueRange {adaptor.getLhs(), adaptor.getRhs()},
       [&](Value lhsValue, Value rhsValue) {
         Value added = spatial::SpatVAddOp::create(
           rewriter, planOp.getLoc(), planOp.getOutput().getType(), lhsValue, rhsValue);
@@ -343,11 +333,12 @@ struct LowerDenseAddPlan final : OpRewritePattern<spatial::SpatAddPlanOp> {
   }
 };
 
-struct LowerDenseConcatPlan final : OpRewritePattern<spatial::SpatConcatPlanOp> {
-  using OpRewritePattern::OpRewritePattern;
+struct LowerDenseConcatPlan final : OpConversionPattern<spatial::SpatConcatPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
 
   LogicalResult matchAndRewrite(spatial::SpatConcatPlanOp planOp,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatConcatPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
     if (!isDenseSelected(planOp.getOperation()))
       return failure();
     auto compute = createSpatCompute(
@@ -355,7 +346,7 @@ struct LowerDenseConcatPlan final : OpRewritePattern<spatial::SpatConcatPlanOp>
       planOp.getLoc(),
       TypeRange {planOp.getOutput().getType()},
       {},
-      planOp.getInputs(),
+      adaptor.getInputs(),
       [&](ValueRange values) {
         Value concatenated = spatial::SpatConcatOp::create(
           rewriter,
@@ -371,16 +362,17 @@ struct LowerDenseConcatPlan final : OpRewritePattern<spatial::SpatConcatPlanOp>
 };
 
 static LogicalResult lowerAddPlan(spatial::SpatAddPlanOp planOp,
+                                  Value lhsValue, Value rhsValue,
                                   PatternRewriter& rewriter) {
-  FailureOr<RowStripPhysicalValue> lhs = getRowStripValue(planOp.getLhs());
-  FailureOr<RowStripPhysicalValue> rhs = getRowStripValue(planOp.getRhs());
+  FailureOr<RowStripPhysicalValue> lhs = getRowStripValue(lhsValue);
+  FailureOr<RowStripPhysicalValue> rhs = getRowStripValue(rhsValue);
   if (isRowStripSelected(planOp.getOperation()) && failed(lhs)) {
-    if (getKnownPhysicalLayout(planOp.getLhs()) == spatial::PhysicalLayout::NHWCRowStrip)
+    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(planOp.getRhs()) == spatial::PhysicalLayout::NHWCRowStrip)
+    if (getKnownPhysicalLayout(rhsValue) == spatial::PhysicalLayout::NHWCRowStrip)
       return failure();
     return planOp.emitOpError("selected row-strip Add plan requires row-strip inputs");
   }
@@ -397,9 +389,10 @@ static LogicalResult lowerAddPlan(spatial::SpatAddPlanOp planOp,
 }
 
 static LogicalResult lowerConcatPlan(spatial::SpatConcatPlanOp planOp,
+                                     ValueRange inputValues,
                                      PatternRewriter& rewriter) {
   SmallVector<RowStripPhysicalValue> inputs;
-  for (Value input : planOp.getInputs()) {
+  for (Value input : inputValues) {
     FailureOr<RowStripPhysicalValue> physical = getRowStripValue(input);
     if (failed(physical)) {
       inputs.clear();
@@ -407,8 +400,8 @@ static LogicalResult lowerConcatPlan(spatial::SpatConcatPlanOp planOp,
     }
     inputs.push_back(*physical);
   }
-  if (isRowStripSelected(planOp.getOperation()) && inputs.size() != planOp.getInputs().size()) {
-    if (llvm::any_of(planOp.getInputs(), [](Value input) {
+  if (isRowStripSelected(planOp.getOperation()) && inputs.size() != inputValues.size()) {
+    if (llvm::any_of(inputValues, [](Value input) {
           return getKnownPhysicalLayout(input) == spatial::PhysicalLayout::NHWCRowStrip;
         }))
       return failure();
@@ -426,15 +419,19 @@ static LogicalResult lowerConcatPlan(spatial::SpatConcatPlanOp planOp,
   return planOp.emitOpError("dense Concat plan was not lowered by the selected-plan patterns");
 }
 
-struct LowerSelectedConvPlan final : OpRewritePattern<spatial::SpatConv2DPlanOp> {
-  explicit LowerSelectedConvPlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
-      : OpRewritePattern<spatial::SpatConv2DPlanOp>(ctx), target(target) {}
+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,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatConv2DPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
     if (isDenseSelected(planOp.getOperation())) {
       FailureOr<Value> lowered = lowerSelectedConv2DPlan(
-        planOp, std::nullopt, /*emitRowStripLayout=*/false, target, rewriter);
+          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);
@@ -443,15 +440,16 @@ struct LowerSelectedConvPlan final : OpRewritePattern<spatial::SpatConv2DPlanOp>
     if (!isRowStripSelected(planOp.getOperation()))
       return failure();
 
-    FailureOr<RowStripPhysicalValue> rowStripInput = getRowStripValue(planOp.getInput());
+    FailureOr<RowStripPhysicalValue> rowStripInput = getRowStripValue(adaptor.getInput());
     if (failed(rowStripInput)
-        && getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+        && getKnownPhysicalLayout(adaptor.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
       return failure();
     std::optional<Value> physicalInput;
     if (succeeded(rowStripInput))
       physicalInput = rowStripInput->storage;
     FailureOr<Value> lowered = lowerSelectedConv2DPlan(
-      planOp, physicalInput, /*emitRowStripLayout=*/true, target, rewriter);
+        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)))
@@ -459,19 +457,21 @@ struct LowerSelectedConvPlan final : OpRewritePattern<spatial::SpatConv2DPlanOp>
     return success();
   }
 
-  const spatial::SpatialTargetInfo& target;
+  const spatial::SpatialTargetResources& target;
+  const ONNXToSpatialPlanningOptions& options;
 };
 
-struct LowerRowStripReluPlan final : OpRewritePattern<spatial::SpatReluPlanOp> {
-  using OpRewritePattern::OpRewritePattern;
+struct LowerRowStripReluPlan final : OpConversionPattern<spatial::SpatReluPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
 
   LogicalResult matchAndRewrite(spatial::SpatReluPlanOp planOp,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatReluPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
     if (!isRowStripSelected(planOp.getOperation()))
       return failure();
-    FailureOr<RowStripPhysicalValue> input = getRowStripValue(planOp.getInput());
+    FailureOr<RowStripPhysicalValue> input = getRowStripValue(adaptor.getInput());
     if (failed(input)) {
-      if (getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+      if (getKnownPhysicalLayout(adaptor.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
         return failure();
       return planOp.emitOpError("selected row-strip ReLU plan requires a row-strip input");
     }
@@ -484,16 +484,17 @@ struct LowerRowStripReluPlan final : OpRewritePattern<spatial::SpatReluPlanOp> {
   }
 };
 
-struct LowerRowStripSiluPlan final : OpRewritePattern<spatial::SpatSiluPlanOp> {
-  using OpRewritePattern::OpRewritePattern;
+struct LowerRowStripSiluPlan final : OpConversionPattern<spatial::SpatSiluPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
 
   LogicalResult matchAndRewrite(spatial::SpatSiluPlanOp planOp,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatSiluPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
     if (!isRowStripSelected(planOp.getOperation()))
       return failure();
-    FailureOr<RowStripPhysicalValue> input = getRowStripValue(planOp.getInput());
+    FailureOr<RowStripPhysicalValue> input = getRowStripValue(adaptor.getInput());
     if (failed(input)) {
-      if (getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+      if (getKnownPhysicalLayout(adaptor.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
         return failure();
       return planOp.emitOpError("selected row-strip SiLU plan requires a row-strip input");
     }
@@ -506,21 +507,23 @@ struct LowerRowStripSiluPlan final : OpRewritePattern<spatial::SpatSiluPlanOp> {
   }
 };
 
-struct LowerRowStripResizePlan final : OpRewritePattern<spatial::SpatResizeNearestPlanOp> {
-  explicit LowerRowStripResizePlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
-      : OpRewritePattern<spatial::SpatResizeNearestPlanOp>(ctx), target(target) {}
+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,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatResizeNearestPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
     if (!isRowStripSelected(planOp.getOperation()))
       return failure();
-    FailureOr<RowStripPhysicalValue> input = getRowStripValue(planOp.getInput());
+    FailureOr<RowStripPhysicalValue> input = getRowStripValue(adaptor.getInput());
     if (failed(input)) {
-      if (getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+      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, input->storage, target, rewriter);
+    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)))
@@ -528,43 +531,47 @@ struct LowerRowStripResizePlan final : OpRewritePattern<spatial::SpatResizeNeare
     return success();
   }
 
-  const spatial::SpatialTargetInfo& target;
+  const spatial::SpatialTargetResources& target;
 };
 
-struct LowerDenseMaxPoolPlan final : OpRewritePattern<spatial::SpatMaxPool2DPlanOp> {
-  explicit LowerDenseMaxPoolPlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
-      : OpRewritePattern<spatial::SpatMaxPool2DPlanOp>(ctx), target(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,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatMaxPool2DPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
     if (!isDenseSelected(planOp.getOperation()))
       return failure();
-    FailureOr<Value> lowered = lowerDenseMaxPool2DPlan(planOp, target, rewriter);
+    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::SpatialTargetInfo& target;
+  const spatial::SpatialTargetResources& target;
 };
 
-struct LowerRowStripMaxPoolPlan final : OpRewritePattern<spatial::SpatMaxPool2DPlanOp> {
-  explicit LowerRowStripMaxPoolPlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
-      : OpRewritePattern<spatial::SpatMaxPool2DPlanOp>(ctx), target(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,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatMaxPool2DPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
     if (!isRowStripSelected(planOp.getOperation()))
       return failure();
-    FailureOr<RowStripPhysicalValue> input = getRowStripValue(planOp.getInput());
+    FailureOr<RowStripPhysicalValue> input = getRowStripValue(adaptor.getInput());
     if (failed(input)
-        && getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+        && getKnownPhysicalLayout(adaptor.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
       return failure();
     std::optional<Value> physicalInput;
     if (succeeded(input))
       physicalInput = input->storage;
-    FailureOr<Value> lowered = lowerSelectedMaxPool2DPlan(planOp, physicalInput, target, rewriter);
+    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)))
@@ -572,26 +579,28 @@ struct LowerRowStripMaxPoolPlan final : OpRewritePattern<spatial::SpatMaxPool2DP
     return success();
   }
 
-  const spatial::SpatialTargetInfo& target;
+  const spatial::SpatialTargetResources& target;
 };
 
 struct LowerRowStripGlobalAveragePoolPlan
-    final : OpRewritePattern<spatial::SpatGlobalAveragePoolPlanOp> {
-  explicit LowerRowStripGlobalAveragePoolPlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
-      : OpRewritePattern<spatial::SpatGlobalAveragePoolPlanOp>(ctx), target(target) {}
+    final : OpConversionPattern<spatial::SpatGlobalAveragePoolPlanOp> {
+  explicit LowerRowStripGlobalAveragePoolPlan(MLIRContext* ctx, const spatial::SpatialTargetResources& target)
+      : OpConversionPattern<spatial::SpatGlobalAveragePoolPlanOp>(ctx), target(target) {}
 
   LogicalResult matchAndRewrite(spatial::SpatGlobalAveragePoolPlanOp planOp,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatGlobalAveragePoolPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
     if (!isRowStripSelected(planOp.getOperation()))
       return failure();
-    FailureOr<RowStripPhysicalValue> input = getRowStripValue(planOp.getInput());
+    FailureOr<RowStripPhysicalValue> input = getRowStripValue(adaptor.getInput());
     if (failed(input)
-        && getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+        && getKnownPhysicalLayout(adaptor.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
       return failure();
     std::optional<Value> physicalInput;
     if (succeeded(input))
       physicalInput = input->storage;
-    FailureOr<Value> lowered = lowerSelectedGlobalAveragePoolPlan(planOp, physicalInput, target, rewriter);
+    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)))
@@ -599,43 +608,47 @@ struct LowerRowStripGlobalAveragePoolPlan
     return success();
   }
 
-  const spatial::SpatialTargetInfo& target;
+  const spatial::SpatialTargetResources& target;
 };
 
 struct LowerDenseGlobalAveragePoolPlan
-    final : OpRewritePattern<spatial::SpatGlobalAveragePoolPlanOp> {
+    final : OpConversionPattern<spatial::SpatGlobalAveragePoolPlanOp> {
   explicit LowerDenseGlobalAveragePoolPlan(MLIRContext* ctx,
-                                           const spatial::SpatialTargetInfo& target)
-      : OpRewritePattern<spatial::SpatGlobalAveragePoolPlanOp>(ctx), target(target) {}
+                                           const spatial::SpatialTargetResources& target)
+      : OpConversionPattern<spatial::SpatGlobalAveragePoolPlanOp>(ctx), target(target) {}
 
   LogicalResult matchAndRewrite(spatial::SpatGlobalAveragePoolPlanOp planOp,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatGlobalAveragePoolPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
     if (!isDenseSelected(planOp.getOperation()))
       return failure();
-    FailureOr<Value> lowered = lowerDenseGlobalAveragePoolPlan(planOp, target, rewriter);
+    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::SpatialTargetInfo& target;
+  const spatial::SpatialTargetResources& target;
 };
 
-struct LowerRowStripBiasAddPlan final : OpRewritePattern<spatial::SpatBiasAddPlanOp> {
-  using OpRewritePattern::OpRewritePattern;
+struct LowerRowStripBiasAddPlan final : OpConversionPattern<spatial::SpatBiasAddPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
 
   LogicalResult matchAndRewrite(spatial::SpatBiasAddPlanOp planOp,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatBiasAddPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
     if (!isRowStripSelected(planOp.getOperation()))
       return failure();
-    FailureOr<RowStripPhysicalValue> input = getRowStripValue(planOp.getInput());
+    FailureOr<RowStripPhysicalValue> input = getRowStripValue(adaptor.getInput());
     if (failed(input)) {
-      if (getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+      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 = lowerRowStripBiasAdd(*input, planOp, rewriter);
+    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)))
@@ -644,48 +657,51 @@ struct LowerRowStripBiasAddPlan final : OpRewritePattern<spatial::SpatBiasAddPla
   }
 };
 
-struct LowerRowStripAddPlan final : OpRewritePattern<spatial::SpatAddPlanOp> {
-  using OpRewritePattern::OpRewritePattern;
+struct LowerRowStripAddPlan final : OpConversionPattern<spatial::SpatAddPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
 
   LogicalResult matchAndRewrite(spatial::SpatAddPlanOp planOp,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatAddPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
     if (!isRowStripSelected(planOp.getOperation()))
       return failure();
-    return lowerAddPlan(planOp, rewriter);
+    return lowerAddPlan(planOp, adaptor.getLhs(), adaptor.getRhs(), rewriter);
   }
 };
 
-struct LowerRowStripConcatPlan final : OpRewritePattern<spatial::SpatConcatPlanOp> {
-  using OpRewritePattern::OpRewritePattern;
+struct LowerRowStripConcatPlan final : OpConversionPattern<spatial::SpatConcatPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
 
   LogicalResult matchAndRewrite(spatial::SpatConcatPlanOp planOp,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatConcatPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
     if (!isRowStripSelected(planOp.getOperation()))
       return failure();
-    return lowerConcatPlan(planOp, rewriter);
+    return lowerConcatPlan(planOp, adaptor.getInputs(), rewriter);
   }
 };
 
 struct LowerMaterializeLayout final
-    : OpRewritePattern<spatial::SpatMaterializeLayoutOp> {
-  using OpRewritePattern::OpRewritePattern;
+    : OpConversionPattern<spatial::SpatMaterializeLayoutOp> {
+  using OpConversionPattern::OpConversionPattern;
 
   LogicalResult matchAndRewrite(spatial::SpatMaterializeLayoutOp materializeOp,
-                                PatternRewriter& rewriter) const override {
+                                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, materializeOp.getInput());
+      rewriter.replaceOp(materializeOp, adaptor.getInput());
       return success();
     }
     if (source == spatial::PhysicalLayout::DenseNCHW
         && target == spatial::PhysicalLayout::NHWCRowStrip) {
-      auto logicalType = dyn_cast<RankedTensorType>(materializeOp.getInput().getType());
+      auto logicalType = dyn_cast<RankedTensorType>(adaptor.getInput().getType());
       if (!logicalType)
         return materializeOp.emitOpError("requires a ranked dense input"), failure();
       FailureOr<Value> rowStrip = materializeDenseToRowStrip(
-        materializeOp.getInput(), logicalType, materializeOp.getLoc(), rewriter);
+        adaptor.getInput(), logicalType, materializeOp.getLoc(), rewriter);
       if (failed(rowStrip))
         return materializeOp.emitOpError(
           "failed to materialize dense NCHW storage to row-strip layout"), failure();
@@ -696,14 +712,13 @@ struct LowerMaterializeLayout final
         || target != spatial::PhysicalLayout::DenseNCHW)
       return materializeOp.emitOpError(
           "unsupported Spatial layout materialization direction"), failure();
-    auto inputType = dyn_cast<RankedTensorType>(materializeOp.getInput().getType());
+    auto inputType = dyn_cast<RankedTensorType>(adaptor.getInput().getType());
     if (!inputType)
       return materializeOp.emitOpError("requires a ranked row-strip input"), failure();
     FailureOr<RowStripPhysicalValue> rowStripValue =
-        getRowStripValue(materializeOp.getInput());
+        getRowStripValue(adaptor.getInput());
     if (failed(rowStripValue))
-      return materializeOp.emitOpError(
-          "requires an explicitly defining row-strip physical value"), failure();
+      return failure();
     FailureOr<Value> dense = materializeRowStripToDense(
         *rowStripValue, materializeOp.getLoc(), rewriter);
     if (failed(dense))
@@ -714,189 +729,117 @@ struct LowerMaterializeLayout final
   }
 };
 
-struct LowerRowStripFlatten final
-    : OpRewritePattern<spatial::SpatGraphCompute> {
-  explicit LowerRowStripFlatten(MLIRContext* context,
-                                const spatial::SpatialTargetInfo& target)
-      : OpRewritePattern<spatial::SpatGraphCompute>(context), target(target) {}
+struct LowerSelectedFlattenPlan final
+    : OpConversionPattern<spatial::SpatFlattenPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
 
-  LogicalResult matchAndRewrite(spatial::SpatGraphCompute flattenOp,
-                                PatternRewriter& rewriter) const override {
-    if (flattenOp.getInputs().size() != 1)
+  LogicalResult matchAndRewrite(spatial::SpatFlattenPlanOp planOp,
+                                spatial::SpatFlattenPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (!isDenseSelected(planOp.getOperation()))
       return failure();
-    FailureOr<RowStripPhysicalValue> input =
-        getRowStripValue(flattenOp.getInputs().front());
-    if (failed(input) || failed(canLowerFlattenFromRowStrip(flattenOp, target)))
-      return failure();
-    if (failed(lowerFlattenFromRowStrip(*input, flattenOp, target, rewriter)))
-      return flattenOp.emitOpError(
-          "failed to preserve row-strip layout through Flatten"), 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();
   }
 
-  const spatial::SpatialTargetInfo& target;
+  explicit LowerSelectedFlattenPlan(MLIRContext* context,
+                                    const spatial::SpatialTargetResources& target)
+      : OpConversionPattern<spatial::SpatFlattenPlanOp>(context), target(target) {}
+
+  const spatial::SpatialTargetResources& target;
 };
 
-struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, OperationPass<ModuleOp>> {
-  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerSpatialPlansPass)
+struct EraseDeadPhysicalViewBlueprint final
+    : OpRewritePattern<spatial::SpatBlueprintOp> {
+  using OpRewritePattern::OpRewritePattern;
 
-  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;
-  explicit LowerSpatialPlansPass(const spatial::SpatialTargetInfo& target)
-      : target(target), hasTarget(true) {}
-
-  void runOnOperation() override {
-    ModuleOp moduleOp = getOperation();
-    if (!hasTarget) {
-      moduleOp.emitError("Spatial plan lowering requires an injected SpatialTargetInfo");
-      signalPassFailure();
-      return;
-    }
-    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);
-    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;
-    if (failed(verifySelectedLayouts(funcOp, target))) {
-      moduleOp.emitError("selected Spatial layout verification failed");
-      signalPassFailure();
-      return;
-    }
-
-    RewritePatternSet selectedPlanPatterns(ctx);
-    selectedPlanPatterns.add<LowerDenseReluPlan,
-                             LowerRowStripReluPlan,
-                             LowerDenseSiluPlan,
-                             LowerRowStripSiluPlan,
-                             LowerDenseBiasAddPlan,
-                             LowerRowStripBiasAddPlan,
-                             LowerDenseAddPlan,
-                             LowerRowStripAddPlan,
-                             LowerDenseConcatPlan,
-                             LowerRowStripConcatPlan>(ctx);
-    selectedPlanPatterns.add<LowerSelectedConvPlan,
-                             LowerDenseResizePlan,
-                             LowerRowStripResizePlan,
-                             LowerDenseMaxPoolPlan,
-                             LowerRowStripMaxPoolPlan,
-                             LowerDenseGlobalAveragePoolPlan,
-                             LowerRowStripGlobalAveragePoolPlan>(ctx, target);
-    if (failed(applyPatternsGreedily(funcOp, std::move(selectedPlanPatterns)))) {
-      moduleOp.emitError("failed to lower selected Spatial plans");
-      signalPassFailure();
-      return;
-    }
-
-    RewritePatternSet layoutPatterns(ctx);
-    layoutPatterns.add<LowerMaterializeLayout>(ctx);
-    layoutPatterns.add<LowerRowStripFlatten>(ctx, target);
-    ConversionTarget layoutTarget(*ctx);
-    layoutTarget.addLegalDialect<spatial::SpatialDialect,
-                                 tensor::TensorDialect,
-                                 linalg::LinalgDialect,
-                                 affine::AffineDialect,
-                                 arith::ArithDialect,
-                                 scf::SCFDialect,
-                                 func::FuncDialect>();
-    layoutTarget.addIllegalDialect<ONNXDialect>();
-    layoutTarget.addIllegalOp<spatial::SpatMaterializeLayoutOp>();
-    layoutTarget.addDynamicallyLegalOp<spatial::SpatGraphCompute>(
-        [&](spatial::SpatGraphCompute computeOp) {
-          if (computeOp.getInputs().size() != 1)
-            return true;
-          FailureOr<RowStripPhysicalValue> input =
-              getRowStripValue(computeOp.getInputs().front());
-          return failed(input) || failed(canLowerFlattenFromRowStrip(computeOp, target));
-        });
-    FrozenRewritePatternSet frozenLayoutPatterns(std::move(layoutPatterns));
-    if (failed(applyFullConversion(funcOp, layoutTarget,
-                                   frozenLayoutPatterns))) {
-      moduleOp.emitError("failed to lower explicit Spatial layout materialization");
-      signalPassFailure();
-      return;
-    }
-
-    if (!verifyLogicalPhase("after selected-plan conversion"))
-      return;
-    SmallVector<spatial::SpatBlueprintOp> deadPhysicalViews;
-    funcOp.walk([&](spatial::SpatBlueprintOp blueprint) {
-      if (spatial::isPhysicalView(blueprint.getMode()) && blueprint.use_empty())
-        deadPhysicalViews.push_back(blueprint);
-    });
-    for (spatial::SpatBlueprintOp blueprint : deadPhysicalViews)
-      rewriter.eraseOp(blueprint);
-    bool hasIllegalOps = false;
-    moduleOp.walk([&](Operation* op) {
-      if (isa<ONNXEntryPointOp>(op))
-        return;
-      if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
-        if (spatial::isFragmentAssembly(blueprint.getMode()))
-          return;
-        op->emitOpError("planning blueprint must not remain after LowerSpatialPlans");
-        hasIllegalOps = true;
-      }
-      else if (isa<spatial::SpatConv2DPlanOp,
-                   spatial::SpatBiasAddPlanOp,
-                   spatial::SpatAddPlanOp,
-                   spatial::SpatReluPlanOp,
-                   spatial::SpatSiluPlanOp,
-                   spatial::SpatResizeNearestPlanOp,
-                   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;
+  LogicalResult matchAndRewrite(spatial::SpatBlueprintOp blueprint,
+                                PatternRewriter& rewriter) const override {
+    if (!spatial::isPhysicalView(blueprint.getMode()) || !blueprint.use_empty())
+      return failure();
+    rewriter.eraseOp(blueprint);
+    return success();
   }
-
-  spatial::SpatialTargetInfo target;
-  bool hasTarget = false;
 };
 
+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
 
-std::unique_ptr<Pass> createLowerSpatialPlansPass() { return std::make_unique<LowerSpatialPlansPass>(); }
+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);
+}
 
-std::unique_ptr<Pass> createLowerSpatialPlansPass(const spatial::SpatialTargetInfo& target) {
-  return std::make_unique<LowerSpatialPlansPass>(target);
+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

@@ -9,11 +9,12 @@ void populatePrePatterns(RewritePatternSet& patterns, MLIRContext* ctx) { popula
 
 void populateConversionPatterns(RewritePatternSet& patterns,
                                 MLIRContext* ctx,
-                                const spatial::SpatialTargetInfo& target) {
+                                const spatial::SpatialTargetResources& target,
+                                const ONNXToSpatialPlanningOptions& options) {
   populateElementwisePatterns(patterns, ctx);
   populateMatMulRewritePatterns(patterns, ctx, target);
   populateGemmPatterns(patterns, ctx, target);
-  populateConvPatterns(patterns, ctx, target);
+  populateConvPatterns(patterns, ctx, target, options);
   populatePoolPatterns(patterns, ctx, target);
   populateReduceMeanPatterns(patterns, ctx);
   populateReluPatterns(patterns, ctx);

src/PIM/Conversion/ONNXToSpatial/Patterns.hpp

@@ -4,18 +4,20 @@
 #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 SpatialTargetInfo;
+struct SpatialTargetResources;
 }
 
 void populatePrePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateConversionPatterns(mlir::RewritePatternSet& patterns,
                                 mlir::MLIRContext* ctx,
-                                const spatial::SpatialTargetInfo& target);
+                                const spatial::SpatialTargetResources& target,
+                                const ONNXToSpatialPlanningOptions& options);
 void populatePostPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 
 void populateGeneratedPrePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
@@ -23,21 +25,22 @@ void populateWeightPromotionPatterns(mlir::RewritePatternSet& patterns, mlir::ML
 
 void populateConvPatterns(mlir::RewritePatternSet& patterns,
                           mlir::MLIRContext* ctx,
-                          const spatial::SpatialTargetInfo& target);
+                          const spatial::SpatialTargetResources& target,
+                          const ONNXToSpatialPlanningOptions& options);
 void populateElementwisePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateElementwiseFusionPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateGemmPatterns(mlir::RewritePatternSet& patterns,
                           mlir::MLIRContext* ctx,
-                          const spatial::SpatialTargetInfo& target);
+                          const spatial::SpatialTargetResources& target);
 void populateMatMulRewritePatterns(mlir::RewritePatternSet& patterns,
                                    mlir::MLIRContext* ctx,
-                                   const spatial::SpatialTargetInfo& target);
+                                   const spatial::SpatialTargetResources& target);
 void populateMatMulFusionPatterns(mlir::RewritePatternSet& patterns,
                                   mlir::MLIRContext* ctx,
-                                  const spatial::SpatialTargetInfo& target);
+                                  const spatial::SpatialTargetResources& target);
 void populatePoolPatterns(mlir::RewritePatternSet& patterns,
                           mlir::MLIRContext* ctx,
-                          const spatial::SpatialTargetInfo& target);
+                          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,6 +1,7 @@
 #include "ConvGeometry.hpp"
 
 #include <algorithm>
+#include <limits>
 
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
 
@@ -8,12 +9,22 @@ namespace onnx_mlir {
 
 namespace {
 
-static int64_t ceilDivide(int64_t value, int64_t divisor) {
-  return divisor == 0 ? 0 : (value + divisor - 1) / divisor;
+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;
 }
@@ -31,20 +42,8 @@ void classifyConvProblem(ConvProblem& problem) {
 }
 
 ConvGeometry buildConvGeometry(const ConvProblem& problem,
-                               const spatial::SpatialTargetInfo& target) {
+                               const spatial::SpatialTargetResources& target) {
   ConvGeometry geo {
-    problem.batchSize,
-    problem.numChannelsIn,
-    problem.xHeight,
-    problem.xWidth,
-    problem.numChannelsOut,
-    problem.wHeight,
-    problem.wWidth,
-    problem.outHeight,
-    problem.outWidth,
-    problem.group,
-    problem.numChannelsInPerGroup,
-    problem.numChannelsOutPerGroup,
     problem.numChannelsInPerGroup * problem.wHeight * problem.wWidth,
     problem.numChannelsOutPerGroup,
     problem.batchSize * problem.outHeight * problem.outWidth,
@@ -52,11 +51,6 @@ ConvGeometry buildConvGeometry(const ConvProblem& problem,
     static_cast<int64_t>(target.matrixUnitsPerProcessor),
     1,
     0,
-    problem.hasBias,
-    isDepthwiseConv(problem.group,
-                    problem.numChannelsIn,
-                    problem.numChannelsOut,
-                    problem.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));
@@ -64,216 +58,76 @@ ConvGeometry buildConvGeometry(const ConvProblem& problem,
 }
 
 static ConvMaterializationKind getMaterializationKind(
-    const ConvProblem& problem, spatial::ConvLoweringStrategy strategy) {
-  if (strategy == spatial::ConvLoweringStrategy::Depthwise)
-    return ConvMaterializationKind::StructuredDepthwise;
-  if (problem.isPointwise)
-    return ConvMaterializationKind::PointwiseContraction;
+    spatial::ConvLoweringStrategy strategy) {
   switch (strategy) {
   case spatial::ConvLoweringStrategy::Depthwise:
     return ConvMaterializationKind::StructuredDepthwise;
-  case spatial::ConvLoweringStrategy::PackedIm2Col:
   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::OutputChannelTiled:
-    return ConvMaterializationKind::OutputChannelTiled;
   case spatial::ConvLoweringStrategy::InputKTiled:
     return ConvMaterializationKind::InputKTiled;
-  case spatial::ConvLoweringStrategy::Tiled2D:
-    return ConvMaterializationKind::Tiled2D;
   case spatial::ConvLoweringStrategy::Auto:
     break;
   }
   llvm_unreachable("auto is not a Conv materialization kind");
 }
 
-static ConvPlan makeCandidatePlan(const ConvProblem& problem,
-                                  spatial::ConvLoweringStrategy strategy,
-                                  const spatial::SpatialTargetInfo& target) {
-  ConvPlan plan;
-  plan.geometry = buildConvGeometry(problem, target);
-  plan.strategy = strategy;
-  plan.materializationKind = getMaterializationKind(problem, strategy);
-  plan.laneCount = plan.geometry.p;
-  plan.reductionCount = std::max<int64_t>(
-    1, (plan.geometry.k + plan.geometry.xbarSize - 1) / plan.geometry.xbarSize);
-  plan.mvmCount = plan.laneCount * plan.reductionCount;
-  plan.vectorCount = plan.mvmCount;
-  plan.weightElements = static_cast<uint64_t>(std::max<int64_t>(0, problem.numChannelsOut))
-    * static_cast<uint64_t>(std::max<int64_t>(0, plan.geometry.k));
-  plan.scratchElements = plan.geometry.im2colElements;
-  plan.materializationElements = strategy == spatial::ConvLoweringStrategy::Depthwise
-    ? 0
-    : std::min<uint64_t>(plan.geometry.im2colElements, target.convIm2colMaxElements);
-  plan.requiresInputMaterialization = strategy != spatial::ConvLoweringStrategy::Depthwise;
-  plan.producesRowStrip = strategy != spatial::ConvLoweringStrategy::InputKTiled
-    && ceilDivide(plan.geometry.k, plan.geometry.xbarSize) <= plan.geometry.matrixUnitsPerProcessor;
-  plan.consumesRowStrip = plan.producesRowStrip;
-  // Conv materializers emit local compute and leave inter-core communication
-  // to Spatial scheduling; zero is an explicit ownership statement here.
-  plan.communicationElements = 0;
-  plan.usesContraction = problem.isPointwise || strategy != spatial::ConvLoweringStrategy::Depthwise;
-  if (problem.isPointwise) {
-    ContractionProblem contraction;
-    contraction.origin = ContractionOrigin::Gemm;
-    contraction.batch = 1;
-    contraction.m = plan.geometry.p;
-    contraction.k = plan.geometry.c;
-    contraction.n = problem.numChannelsOutPerGroup;
-    contraction.lhsElementType = problem.xType.getElementType();
-    contraction.rhsElementType = problem.wType.getElementType();
-    contraction.resultElementType = problem.outType.getElementType();
-    plan.contraction = makeContractionPlan(
-      contraction, target, ContractionPlanKind::StaticTiled);
-    plan.hasContractionPlan = true;
-    plan.laneCount = plan.contraction.laneCount;
-    plan.mvmCount = plan.contraction.expectedMvmCount;
-    plan.vectorCount = plan.contraction.expectedVectorCount;
-    plan.reductionCount = plan.contraction.reductionSlices;
-  }
-  return plan;
-}
-
 static bool fitsSingleCrossbar(const ConvGeometry& geo) {
   return geo.k <= geo.xbarSize && geo.c <= geo.xbarSize;
 }
 
 static bool fitsPackedIm2Col(const ConvGeometry& geo,
-                             const spatial::SpatialTargetInfo& target) {
+                             const ONNXToSpatialPlanningOptions& options) {
   return fitsSingleCrossbar(geo) && geo.pack >= 2
-    && geo.im2colElements <= target.convIm2colMaxElements;
-}
-
-static mlir::FailureOr<ConvPlan> buildDepthwiseCandidate(
-    const ConvProblem& problem, const spatial::SpatialTargetInfo& target) {
-  if (!problem.isDepthwise)
-    return mlir::failure();
-  return makeCandidatePlan(problem, spatial::ConvLoweringStrategy::Depthwise, target);
-}
-
-static mlir::FailureOr<ConvPlan> buildPackedIm2ColCandidate(
-    const ConvProblem& problem, const spatial::SpatialTargetInfo& target) {
-  ConvGeometry geo = buildConvGeometry(problem, target);
-  if (!fitsPackedIm2Col(geo, target))
-    return mlir::failure();
-  return makeCandidatePlan(problem, spatial::ConvLoweringStrategy::PackedIm2Col, target);
-}
-
-static mlir::FailureOr<ConvPlan> buildStreamedPatchCandidate(
-    const ConvProblem& problem, const spatial::SpatialTargetInfo& target) {
-  if (!fitsSingleCrossbar(buildConvGeometry(problem, target)))
-    return mlir::failure();
-  return makeCandidatePlan(problem, spatial::ConvLoweringStrategy::StreamedPatch, target);
-}
-
-static mlir::FailureOr<ConvPlan> buildStreamedPackedCandidate(
-    const ConvProblem& problem, const spatial::SpatialTargetInfo& target) {
-  ConvGeometry geo = buildConvGeometry(problem, target);
-  if (!fitsSingleCrossbar(geo) || geo.pack < 2)
-    return mlir::failure();
-  return makeCandidatePlan(problem, spatial::ConvLoweringStrategy::StreamedPacked, target);
-}
-
-static mlir::FailureOr<ConvPlan> buildOutputChannelTiledCandidate(
-    const ConvProblem& problem, const spatial::SpatialTargetInfo& target) {
-  ConvGeometry geo = buildConvGeometry(problem, target);
-  if (geo.k > geo.xbarSize || geo.c <= geo.xbarSize)
-    return mlir::failure();
-  return makeCandidatePlan(problem, spatial::ConvLoweringStrategy::OutputChannelTiled, target);
-}
-
-static mlir::FailureOr<ConvPlan> buildInputKTiledCandidate(
-    const ConvProblem& problem, const spatial::SpatialTargetInfo& target) {
-  ConvGeometry geo = buildConvGeometry(problem, target);
-  if (geo.k <= geo.xbarSize || geo.c > geo.xbarSize)
-    return mlir::failure();
-  return makeCandidatePlan(problem, spatial::ConvLoweringStrategy::InputKTiled, target);
-}
-
-static mlir::FailureOr<ConvPlan> buildTiled2DCandidate(
-    const ConvProblem& problem, const spatial::SpatialTargetInfo& target) {
-  ConvGeometry geo = buildConvGeometry(problem, target);
-  if (geo.k <= geo.xbarSize || geo.c <= geo.xbarSize)
-    return mlir::failure();
-  return makeCandidatePlan(problem, spatial::ConvLoweringStrategy::Tiled2D, target);
-}
-
-static mlir::FailureOr<ConvPlan> buildLegacyCandidate(
-    const ConvProblem& problem, const spatial::SpatialTargetInfo& target) {
-  // Legacy is retained as the explicit compatibility/debug materializer and
-  // as the safe fallback when structured depthwise lowering is unavailable.
-  return makeCandidatePlan(problem, spatial::ConvLoweringStrategy::Legacy, target);
+    && geo.im2colElements <= options.convIm2colMaxElements;
 }
 
 mlir::FailureOr<ConvPlan> makeConvPlan(const ConvProblem& problem,
                                        spatial::ConvLoweringStrategy strategy,
-                                       const spatial::SpatialTargetInfo& target) {
+                                       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 buildLegacyCandidate(problem, target);
+    return plan();
   case spatial::ConvLoweringStrategy::Depthwise:
-    return buildDepthwiseCandidate(problem, target);
+    return ifApplicable(problem.isDepthwise);
   case spatial::ConvLoweringStrategy::PackedIm2Col:
-    return buildPackedIm2ColCandidate(problem, target);
+    return ifApplicable(fitsPackedIm2Col(geo, options));
   case spatial::ConvLoweringStrategy::StreamedPatch:
-    return buildStreamedPatchCandidate(problem, target);
+    return ifApplicable(fitsSingleCrossbar(geo));
   case spatial::ConvLoweringStrategy::StreamedPacked:
-    return buildStreamedPackedCandidate(problem, target);
+    return ifApplicable(fitsSingleCrossbar(geo) && geo.pack >= 2);
   case spatial::ConvLoweringStrategy::OutputChannelTiled:
-    return buildOutputChannelTiledCandidate(problem, target);
+    return ifApplicable(geo.k <= geo.xbarSize && geo.c > geo.xbarSize);
   case spatial::ConvLoweringStrategy::InputKTiled:
-    return buildInputKTiledCandidate(problem, target);
+    return ifApplicable(geo.k > geo.xbarSize && geo.c <= geo.xbarSize);
   case spatial::ConvLoweringStrategy::Tiled2D:
-    return buildTiled2DCandidate(problem, target);
+    return ifApplicable(geo.k > geo.xbarSize && geo.c > geo.xbarSize);
   }
   llvm_unreachable("unknown Conv lowering strategy");
 }
 
-llvm::SmallVector<ConvPlan, 8> buildConvPlanCandidates(
-    const ConvProblem& problem, const spatial::SpatialTargetInfo& target) {
-  ConvGeometry geo = buildConvGeometry(problem, target);
-  llvm::SmallVector<ConvPlan, 8> candidates;
-  auto append = [&](spatial::ConvLoweringStrategy strategy) {
-    mlir::FailureOr<ConvPlan> candidate = makeConvPlan(problem, strategy, target);
-    if (succeeded(candidate))
-      candidates.push_back(*candidate);
-  };
-
-  if (problem.isDepthwise) {
-    append(spatial::ConvLoweringStrategy::Depthwise);
-    append(spatial::ConvLoweringStrategy::Legacy);
-    return candidates;
-  }
-  if (fitsPackedIm2Col(geo, target))
-    append(spatial::ConvLoweringStrategy::PackedIm2Col);
-  if (fitsSingleCrossbar(geo) && geo.pack >= 2)
-    append(spatial::ConvLoweringStrategy::StreamedPacked);
-  if (fitsSingleCrossbar(geo))
-    append(spatial::ConvLoweringStrategy::StreamedPatch);
-  if (geo.k <= geo.xbarSize && geo.c > geo.xbarSize)
-    append(spatial::ConvLoweringStrategy::OutputChannelTiled);
-  if (geo.k > geo.xbarSize && geo.c <= geo.xbarSize)
-    append(spatial::ConvLoweringStrategy::Legacy);
-  if (geo.k > geo.xbarSize && geo.c <= geo.xbarSize)
-    append(spatial::ConvLoweringStrategy::InputKTiled);
-  if (geo.k > geo.xbarSize && geo.c > geo.xbarSize)
-    append(spatial::ConvLoweringStrategy::Tiled2D);
-  return candidates;
-}
-
 uint64_t chooseStreamChunkPositions(const ConvGeometry& geo,
                                     int64_t packFactor,
-                                    const spatial::SpatialTargetInfo& target) {
+                                    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, target.convIm2colMaxElements / 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, target.convStreamChunkPositions));
+  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);

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

@@ -3,8 +3,8 @@
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Value.h"
 
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ContractionPlanning.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetInfo.hpp"
+#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>
@@ -45,29 +45,20 @@ struct ConvProblem {
   bool isPointwise = false;
 };
 
-struct ConvLoweringState : ConvProblem {
+struct ConvLoweringState {
+  ConvProblem problem;
   mlir::Operation* diagnosticAnchor = nullptr;
   mlir::Value x;
   mlir::Value w;
   mlir::Value b;
-  const spatial::SpatialTargetInfo* target = nullptr;
+  const spatial::SpatialTargetResources* target = nullptr;
+  const ONNXToSpatialPlanningOptions* options = nullptr;
 
-  const spatial::SpatialTargetInfo& targetInfo() const { return *target; }
+  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;
@@ -75,8 +66,6 @@ struct ConvGeometry {
   int64_t matrixUnitsPerProcessor;
   int64_t pack;
   uint64_t im2colElements;
-  bool hasBias;
-  bool isDepthwise;
 };
 
 struct RowInterval {
@@ -94,35 +83,14 @@ struct ConvRowDemand {
 
 enum class ConvMaterializationKind : uint8_t {
   StructuredDepthwise,
-  PointwiseContraction,
   PackedIm2Col,
   StreamedPatch,
   StreamedPacked,
-  OutputChannelTiled,
   InputKTiled,
-  Tiled2D,
 };
 
 struct ConvPlan {
-  ConvGeometry geometry;
-  spatial::ConvLoweringStrategy strategy = spatial::ConvLoweringStrategy::Auto;
-  ConvMaterializationKind materializationKind = ConvMaterializationKind::PackedIm2Col;
-  int64_t laneCount = 0;
-  int64_t mvmCount = 0;
-  int64_t vectorCount = 0;
-  int64_t reductionCount = 0;
-  uint64_t weightElements = 0;
-  uint64_t scratchElements = 0;
-  uint64_t materializationElements = 0;
-  uint64_t communicationElements = 0;
-  spatial::PhysicalLayout resultLayout = spatial::PhysicalLayout::DenseNCHW;
-  bool consumesRowStrip = false;
-  bool producesRowStrip = false;
-  bool requiresInputMaterialization = false;
-  bool requiresOutputMaterialization = false;
-  bool usesContraction = false;
-  bool hasContractionPlan = false;
-  ContractionPlan contraction;
+  ConvMaterializationKind kind = ConvMaterializationKind::PackedIm2Col;
 };
 
 bool isDepthwiseConv(int64_t group, int64_t numChannelsIn, int64_t numChannelsOut, int64_t numChannelsInPerGroup);
@@ -130,18 +98,16 @@ bool isDepthwiseConv(int64_t group, int64_t numChannelsIn, int64_t numChannelsOu
 void classifyConvProblem(ConvProblem& problem);
 
 ConvGeometry buildConvGeometry(const ConvProblem& problem,
-                               const spatial::SpatialTargetInfo& target);
+                               const spatial::SpatialTargetResources& target);
 
 mlir::FailureOr<ConvPlan> makeConvPlan(const ConvProblem& problem,
                                        spatial::ConvLoweringStrategy strategy,
-                                       const spatial::SpatialTargetInfo& target);
-
-llvm::SmallVector<ConvPlan, 8> buildConvPlanCandidates(
-    const ConvProblem& problem, const spatial::SpatialTargetInfo& target);
+                                       const spatial::SpatialTargetResources& target,
+                                       const ONNXToSpatialPlanningOptions& options);
 
 uint64_t chooseStreamChunkPositions(const ConvGeometry& geo,
                                     int64_t packFactor,
-                                    const spatial::SpatialTargetInfo& target);
+                                    const ONNXToSpatialPlanningOptions& options);
 
 RowInterval computeConvInputRowsForOutputRows(RowInterval outputRows, const ConvProblem& problem);
 

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

@@ -53,10 +53,10 @@ struct BlueprintSplatMulToSpatial : OpConversionPattern<ONNXMulOp> {
 
   LogicalResult
   matchAndRewrite(ONNXMulOp op, ONNXMulOpAdaptor adaptor, ConversionPatternRewriter& rewriter) const override {
-    auto blueprint = adaptor.getA().getDefiningOp<spatial::SpatBlueprintOp>();
+    auto blueprint = op.getA().getDefiningOp<spatial::SpatBlueprintOp>();
     Value scalar = adaptor.getB();
     if (!blueprint) {
-      blueprint = adaptor.getB().getDefiningOp<spatial::SpatBlueprintOp>();
+      blueprint = op.getB().getDefiningOp<spatial::SpatBlueprintOp>();
       scalar = adaptor.getA();
     }
     auto scalarAttr = getDenseConstantAttr(scalar);

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

@@ -23,7 +23,6 @@
 #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/ContractionMaterialization.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"
@@ -329,6 +328,8 @@ 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);
 }
@@ -660,14 +661,14 @@ static FailureOr<Value> createReductionOutput(Value partialPieces,
 }
 
 struct GemmToSpatialComputes : OpConversionPattern<ONNXGemmOp> {
-  explicit GemmToSpatialComputes(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
+  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::SpatialTargetInfo& target;
+  const spatial::SpatialTargetResources& target;
 };
 
 } // namespace
@@ -682,7 +683,7 @@ FailureOr<Value> lowerGemmToSpatial(
     bool transB,
     float alpha,
     float beta,
-    const spatial::SpatialTargetInfo& target,
+    const spatial::SpatialTargetResources& target,
     PatternRewriter& rewriter,
     Location loc) {
   auto aType = dyn_cast<RankedTensorType>(a.getType());
@@ -731,14 +732,9 @@ FailureOr<Value> lowerGemmToSpatial(
   problem.m = outType.getDimSize(0);
   problem.k = aType.getDimSize(1);
   problem.n = outType.getDimSize(1);
-  problem.origin = ContractionOrigin::Gemm;
   problem.lhsElementType = aType.getElementType();
   problem.rhsElementType = bType.getElementType();
   problem.resultElementType = outType.getElementType();
-  problem.lhsTransposed = transA;
-  problem.rhsTransposed = transB;
-  problem.alpha = alpha;
-  problem.beta = beta;
   const bool transposeB = transB;
 
   if (!isCompileTimeComputable(b)) {
@@ -789,13 +785,19 @@ FailureOr<Value> lowerGemmToSpatial(
     auto bShape = bType.getShape();
     auto transposedType = RankedTensorType::get({bShape[1], bShape[0]}, bType.getElementType(), bType.getEncoding());
     if (isCompileTimeComputable(b)) {
-      auto transposedConstant = materializeTransposedContractionConstant(
-          b, transposedType, {1, 0}, rewriter, loc);
-      if (failed(transposedConstant)) {
+      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 = *transposedConstant;
+      b = getOrCreateConstant(rewriter,
+                              rewriter.getInsertionBlock()->getParentOp(),
+                              *transposedAttr,
+                              transposedType);
     } else {
       b = createLinalgTranspose(b, transposedType, {1, 0}, rewriter, loc);
     }
@@ -831,7 +833,7 @@ FailureOr<Value> lowerGemmToSpatial(
   }
   b = *paddedB;
   auto paddedAType = RankedTensorType::get({problem.m, paddedReductionSize}, aType.getElementType());
-  a = materializePaddedContractionInput(a, paddedAType, rewriter, loc);
+  a = createPaddedInputCompute(a, paddedAType, rewriter, loc);
   aType = paddedAType;
 
   Value bias;
@@ -896,7 +898,7 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
 
 void populateGemmPatterns(RewritePatternSet& patterns,
                           MLIRContext* ctx,
-                          const spatial::SpatialTargetInfo& target) {
+                          const spatial::SpatialTargetResources& target) {
   patterns.insert<GemmToSpatialComputes>(ctx, target);
 }
 

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

@@ -7,7 +7,7 @@
 
 namespace onnx_mlir {
 namespace spatial {
-struct SpatialTargetInfo;
+struct SpatialTargetResources;
 }
 
 mlir::FailureOr<mlir::Value> lowerGemmToSpatial(
@@ -20,7 +20,7 @@ mlir::FailureOr<mlir::Value> lowerGemmToSpatial(
     bool transB,
     float alpha,
     float beta,
-    const spatial::SpatialTargetInfo& target,
+    const spatial::SpatialTargetResources& target,
     mlir::PatternRewriter& rewriter,
     mlir::Location loc);
 

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

@@ -12,9 +12,10 @@
 #include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.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/ContractionMaterialization.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ContractionPlanning.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/MatrixProductLowering.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns/Math/Gemm.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -362,16 +363,16 @@ static Value transposeLastTwoDims(Value value, PatternRewriter& rewriter, Locati
     return input;
   auto type = cast<RankedTensorType>(value.getType());
   auto shape = type.getShape();
-  auto createONNXTranspose = [&](RankedTensorType resultType, ArrayRef<int64_t> permutation) {
-    return ONNXTransposeOp::create(rewriter, loc, resultType, value, rewriter.getI64ArrayAttr(permutation)).getResult();
+  auto createTranspose = [&](RankedTensorType resultType, ArrayRef<int64_t> permutation) {
+    return createLinalgTranspose(value, resultType, permutation, rewriter, loc);
   };
   if (type.getRank() == 2) {
     auto resultType = RankedTensorType::get({shape[1], shape[0]}, type.getElementType(), type.getEncoding());
-    return createONNXTranspose(resultType, {1, 0});
+    return createTranspose(resultType, {1, 0});
   }
 
   auto resultType = RankedTensorType::get({shape[0], shape[2], shape[1]}, type.getElementType(), type.getEncoding());
-  return createONNXTranspose(resultType, {0, 2, 1});
+  return createTranspose(resultType, {0, 2, 1});
 }
 
 static FailureOr<Value> materializePaddedBatchedWeight(Value value,
@@ -835,7 +836,7 @@ static FailureOr<Value> createBatchedReductionCompute(Value partialPieces,
   return computeOp->getResult(0);
 }
 
-struct NormalizedMatMulInfo : ContractionProblem {
+struct NormalizedMatMulInfo {
   NormalizedMatMulInfo(RankedTensorType lhsType,
                        RankedTensorType rhsType,
                        RankedTensorType outType,
@@ -844,7 +845,7 @@ struct NormalizedMatMulInfo : ContractionProblem {
                        ContractionProblem problem,
                        bool lhsWasVector,
                        bool rhsWasVector)
-      : ContractionProblem(std::move(problem)),
+      : problem(std::move(problem)),
         lhsType(lhsType),
         rhsType(rhsType),
         outType(outType),
@@ -853,6 +854,7 @@ struct NormalizedMatMulInfo : ContractionProblem {
         lhsWasVector(lhsWasVector),
         rhsWasVector(rhsWasVector) {}
 
+  ContractionProblem problem;
   RankedTensorType lhsType;
   RankedTensorType rhsType;
   RankedTensorType outType;
@@ -862,14 +864,15 @@ struct NormalizedMatMulInfo : ContractionProblem {
   bool rhsWasVector;
 };
 
-struct MatMulLoweringPlan : ContractionProblem {
+struct MatMulLoweringPlan {
   MatMulLoweringPlan(Value lhs, Value rhs, const NormalizedMatMulInfo& info)
-      : ContractionProblem(info),
+      : problem(info.problem),
         lhs(lhs),
         rhs(rhs),
         lhsType(cast<RankedTensorType>(lhs.getType())),
         rhsType(cast<RankedTensorType>(rhs.getType())) {}
 
+  ContractionProblem problem;
   Value lhs;
   Value rhs;
   RankedTensorType lhsType;
@@ -952,14 +955,9 @@ static FailureOr<NormalizedMatMulInfo> analyzeMatMulShape(ONNXMatMulOp matmulOp)
                         m,
                         k,
                         n,
-                        ContractionOrigin::MatMul,
                         lhsType.getElementType(),
                         rhsType.getElementType(),
-                        outType.getElementType(),
-                        false,
-                        false,
-                        lhsWasVector,
-                        rhsWasVector},
+                        outType.getElementType()},
     lhsWasVector,
     rhsWasVector);
 }
@@ -979,10 +977,10 @@ static MatMulLoweringPlan buildLoweringPlan(Value normalizedLhs,
   plan.rhs = transposeLastTwoDims(normalizedLhs, rewriter, loc);
   plan.lhsType = cast<RankedTensorType>(plan.lhs.getType());
   plan.rhsType = cast<RankedTensorType>(plan.rhs.getType());
-  std::swap(plan.lhsBatchShape, plan.rhsBatchShape);
-  std::swap(plan.lhsBatch, plan.rhsBatch);
-  plan.m = info.n;
-  plan.n = info.m;
+  std::swap(plan.problem.lhsBatchShape, plan.problem.rhsBatchShape);
+  std::swap(plan.problem.lhsBatch, plan.problem.rhsBatch);
+  plan.problem.m = info.problem.n;
+  plan.problem.n = info.problem.m;
   plan.transposedResult = true;
   return plan;
 }
@@ -1006,17 +1004,17 @@ static Value finalizeNormalizedMatMulResult(Value value,
   // and removing the synthetic unit matrix axes introduced for vector operands.
   Value result = value;
   RankedTensorType currentType = directOutType;
-  if (info.outputBatchShape.size() > 1) {
-    SmallVector<int64_t> expandedShape(info.outputBatchShape.begin(), info.outputBatchShape.end());
-    expandedShape.push_back(info.m);
-    expandedShape.push_back(info.n);
+  if (info.problem.outputBatchShape.size() > 1) {
+    SmallVector<int64_t> expandedShape(info.problem.outputBatchShape.begin(), info.problem.outputBatchShape.end());
+    expandedShape.push_back(info.problem.m);
+    expandedShape.push_back(info.problem.n);
     auto expandedType = RankedTensorType::get(expandedShape, info.outType.getElementType(), info.outType.getEncoding());
-    result = expandBatchDims(result, expandedType, info.outputBatchShape.size(), rewriter, loc);
+    result = expandBatchDims(result, expandedType, info.problem.outputBatchShape.size(), rewriter, loc);
     currentType = expandedType;
   }
 
   SmallVector<bool> removedAxes(currentType.getRank(), false);
-  if (info.outputBatchShape.empty())
+  if (info.problem.outputBatchShape.empty())
     removedAxes[0] = true;
   if (info.lhsWasVector)
     removedAxes[currentType.getRank() - 2] = true;
@@ -1026,79 +1024,77 @@ static Value finalizeNormalizedMatMulResult(Value value,
 }
 
 struct MatMulToGemm : OpRewritePattern<ONNXMatMulOp> {
-  using OpRewritePattern::OpRewritePattern;
+  explicit MatMulToGemm(MLIRContext* ctx,
+                        const spatial::SpatialTargetResources& target)
+      : OpRewritePattern<ONNXMatMulOp>(ctx), target(target) {}
 
   LogicalResult matchAndRewrite(ONNXMatMulOp matmulOp, PatternRewriter& rewriter) const override {
     auto shapeInfo = analyzeMatMulShape(matmulOp);
     if (failed(shapeInfo) || shapeInfo->lhsWasVector || shapeInfo->rhsWasVector)
       return failure();
 
-    if (!shapeInfo->outputBatchShape.empty())
+    if (!shapeInfo->problem.outputBatchShape.empty())
       return failure();
 
     Location loc = matmulOp.getLoc();
     bool useTransposedForm = isCompileTimeComputable(matmulOp.getA()) && !isCompileTimeComputable(matmulOp.getB());
 
-    Value lhs = collapseBatchDims(matmulOp.getA(), shapeInfo->lhsBatch, shapeInfo->m, shapeInfo->k, rewriter, loc);
-    Value rhs = collapseBatchDims(matmulOp.getB(), shapeInfo->rhsBatch, shapeInfo->k, shapeInfo->n, rewriter, loc);
-    int64_t lhsBatchForGemm = shapeInfo->lhsBatch;
-    int64_t rhsBatchForGemm = shapeInfo->rhsBatch;
-    int64_t gemmM = shapeInfo->m;
-    int64_t gemmK = shapeInfo->k;
-    int64_t gemmN = shapeInfo->n;
+    Value lhs = collapseBatchDims(matmulOp.getA(), shapeInfo->problem.lhsBatch, shapeInfo->problem.m, shapeInfo->problem.k, rewriter, loc);
+    Value rhs = collapseBatchDims(matmulOp.getB(), shapeInfo->problem.rhsBatch, shapeInfo->problem.k, shapeInfo->problem.n, rewriter, loc);
+    int64_t lhsBatchForGemm = shapeInfo->problem.lhsBatch;
+    int64_t rhsBatchForGemm = shapeInfo->problem.rhsBatch;
+    int64_t gemmM = shapeInfo->problem.m;
+    int64_t gemmK = shapeInfo->problem.k;
+    int64_t gemmN = shapeInfo->problem.n;
     if (useTransposedForm) {
       lhs = transposeLastTwoDims(matmulOp.getB(), rewriter, loc);
-      lhsBatchForGemm = shapeInfo->rhsBatch;
+      lhsBatchForGemm = shapeInfo->problem.rhsBatch;
       rhs = transposeLastTwoDims(matmulOp.getA(), rewriter, loc);
-      rhsBatchForGemm = shapeInfo->lhsBatch;
-      gemmM = shapeInfo->n;
-      gemmN = shapeInfo->m;
+      rhsBatchForGemm = shapeInfo->problem.lhsBatch;
+      gemmM = shapeInfo->problem.n;
+      gemmN = shapeInfo->problem.m;
     }
 
-    auto gemmType = RankedTensorType::get({gemmM, gemmN}, shapeInfo->outType.getElementType());
-    Value none = ONNXNoneOp::create(rewriter, loc, rewriter.getNoneType());
+    auto gemmType = RankedTensorType::get(
+        {gemmM, gemmN}, shapeInfo->outType.getElementType(), shapeInfo->outType.getEncoding());
     Value lhsMatrix = extractBatchMatrix(lhs, /*batchIndex=*/0, lhsBatchForGemm, gemmM, gemmK, rewriter, loc);
     Value rhsMatrix = extractBatchMatrix(rhs, /*batchIndex=*/0, rhsBatchForGemm, gemmK, gemmN, rewriter, loc);
-    Value gemmResult = ONNXGemmOp::create(rewriter,
-                                          loc,
-                                          gemmType,
-                                          lhsMatrix,
-                                          rhsMatrix,
-                                          none,
-                                          rewriter.getF32FloatAttr(1.0f),
-                                          rewriter.getF32FloatAttr(1.0f),
-                                          rewriter.getBoolAttr(false),
-                                          rewriter.getBoolAttr(false))
-                         .getY();
+    FailureOr<Value> gemmResult = lowerGemmToSpatial(
+        matmulOp.getOperation(), lhsMatrix, rhsMatrix, Value(), gemmType,
+        /*transA=*/false, /*transB=*/false, /*alpha=*/1.0f, /*beta=*/1.0f,
+        target, rewriter, loc);
+    if (failed(gemmResult))
+      return failure();
+    Value result = *gemmResult;
     if (useTransposedForm)
-      gemmResult =
-        ONNXTransposeOp::create(rewriter, loc, shapeInfo->outType, gemmResult, rewriter.getI64ArrayAttr({1, 0}))
-          .getResult();
+      result = transposeLastTwoDims(result, rewriter, loc);
 
-    if (shapeInfo->outputBatchShape.empty()) {
-      rewriter.replaceOp(matmulOp, gemmResult);
+    if (shapeInfo->problem.outputBatchShape.empty()) {
+      rewriter.replaceOp(matmulOp, result);
       return success();
     }
 
     auto directOutType =
-      RankedTensorType::get({1, shapeInfo->m, shapeInfo->n}, shapeInfo->outType.getElementType(), shapeInfo->outType.getEncoding());
-    Value batchedResult = ensureBatchedTensor(gemmResult, /*batchSize=*/1, shapeInfo->m, shapeInfo->n, rewriter, loc);
+      RankedTensorType::get({1, shapeInfo->problem.m, shapeInfo->problem.n}, shapeInfo->outType.getElementType(), shapeInfo->outType.getEncoding());
+    Value batchedResult = ensureBatchedTensor(result, /*batchSize=*/1, shapeInfo->problem.m, shapeInfo->problem.n, rewriter, loc);
     Value finalResult = finalizeNormalizedMatMulResult(batchedResult, directOutType, *shapeInfo, rewriter, loc);
     rewriter.replaceOp(matmulOp, finalResult);
     return success();
   }
+
+  const spatial::SpatialTargetResources& target;
 };
 
 struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
   explicit MatMulBatchedToSpatialComputes(MLIRContext* ctx,
-                                           const spatial::SpatialTargetInfo& target)
+                                           const spatial::SpatialTargetResources& target)
       : OpRewritePattern<ONNXMatMulOp>(ctx), target(target) {}
 
   LogicalResult matchAndRewrite(ONNXMatMulOp matmulOp, PatternRewriter& rewriter) const override {
     auto shapeInfo = analyzeMatMulShape(matmulOp);
     if (failed(shapeInfo))
       return failure();
-    if (!shapeInfo->lhsWasVector && !shapeInfo->rhsWasVector && shapeInfo->outputBatchShape.empty())
+    if (!shapeInfo->lhsWasVector && !shapeInfo->rhsWasVector && shapeInfo->problem.outputBatchShape.empty())
       return failure();
 
     Location loc = matmulOp.getLoc();
@@ -1114,62 +1110,63 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
     const bool rhsStoredAsRows = rhsRows && !useTransposedForm;
 
     Value lhs =
-      normalizeMatMulOperand(matmulOp.getA(), shapeInfo->normalizedLhsType, shapeInfo->lhsWasVector, rewriter, loc);
+      normalizeMatMulOperand(matmulOp.getA(), shapeInfo->normalizedLhsType,
+                             shapeInfo->lhsWasVector, rewriter, loc);
     Value rhs = normalizeMatMulOperand(
-      rhsStoredAsRows ? rhsRows : matmulOp.getB(), shapeInfo->normalizedRhsType, shapeInfo->rhsWasVector, rewriter, loc);
-    lhs = collapseBatchDims(lhs, shapeInfo->lhsBatch, shapeInfo->m, shapeInfo->k, rewriter, loc);
+      rhsStoredAsRows ? rhsRows : matmulOp.getB(), shapeInfo->normalizedRhsType,
+      shapeInfo->rhsWasVector, rewriter, loc);
+    lhs = collapseBatchDims(lhs, shapeInfo->problem.lhsBatch, shapeInfo->problem.m, shapeInfo->problem.k, rewriter, loc);
     rhs = collapseBatchDims(rhs,
-                            shapeInfo->rhsBatch,
-                            rhsStoredAsRows ? shapeInfo->n : shapeInfo->k,
-                            rhsStoredAsRows ? shapeInfo->k : shapeInfo->n,
+                            shapeInfo->problem.rhsBatch,
+                            rhsStoredAsRows ? shapeInfo->problem.n : shapeInfo->problem.k,
+                            rhsStoredAsRows ? shapeInfo->problem.k : shapeInfo->problem.n,
                             rewriter,
                             loc);
     MatMulLoweringPlan plan = buildLoweringPlan(
       lhs, rhs, *shapeInfo, useTransposedForm, rewriter, loc);
 
-    plan.lhs = ensureBatchedTensor(plan.lhs, plan.lhsBatch, plan.m, plan.k, rewriter, loc);
+    plan.lhs = ensureBatchedTensor(plan.lhs, plan.problem.lhsBatch, plan.problem.m, plan.problem.k, rewriter, loc);
     plan.rhs = ensureBatchedTensor(plan.rhs,
-                                   plan.rhsBatch,
-                                   rhsStoredAsRows ? plan.n : plan.k,
-                                   rhsStoredAsRows ? plan.k : plan.n,
+                                   plan.problem.rhsBatch,
+                                   rhsStoredAsRows ? plan.problem.n : plan.problem.k,
+                                   rhsStoredAsRows ? plan.problem.k : plan.problem.n,
                                    rewriter,
                                    loc);
     plan.lhsType = cast<RankedTensorType>(plan.lhs.getType());
     plan.rhsType = cast<RankedTensorType>(plan.rhs.getType());
     auto directOutType = RankedTensorType::get(
-      {plan.batch, plan.m, plan.n}, shapeInfo->outType.getElementType(), shapeInfo->outType.getEncoding());
+      {plan.problem.batch, plan.problem.m, plan.problem.n}, shapeInfo->outType.getElementType(), shapeInfo->outType.getEncoding());
 
     if (isCompileTimeComputable(plan.rhs)) {
       ContractionPlan contractionPlan = makeContractionPlan(
-        plan, target, ContractionPlanKind::StaticTiled);
+        plan.problem, target, ContractionPlanKind::StaticTiled);
       const int64_t numKSlices = contractionPlan.reductionSlices;
       const int64_t numOutHSlices = contractionPlan.outputTiles;
       const int64_t paddedReductionSize = numKSlices * xbarSize;
       const int64_t paddedOutCols = numOutHSlices * xbarSize;
       auto paddedLhsType = RankedTensorType::get(
-        {plan.lhsBatch, plan.m, paddedReductionSize}, plan.lhsType.getElementType(), plan.lhsType.getEncoding());
+        {plan.problem.lhsBatch, plan.problem.m, paddedReductionSize}, plan.lhsType.getElementType(), plan.lhsType.getEncoding());
       auto paddedRhsType = RankedTensorType::get(
-        {plan.batch, paddedReductionSize, paddedOutCols}, plan.rhsType.getElementType(), plan.rhsType.getEncoding());
+        {plan.problem.batch, paddedReductionSize, paddedOutCols}, plan.rhsType.getElementType(), plan.rhsType.getEncoding());
       auto paddedOutType =
-        RankedTensorType::get({plan.batch, plan.m, paddedOutCols}, shapeInfo->outType.getElementType());
+        RankedTensorType::get({plan.problem.batch, plan.problem.m, paddedOutCols}, shapeInfo->outType.getElementType());
 
       auto paddedRhs =
-        materializePaddedBatchedWeight(plan.rhs, plan.rhsBatchShape, plan.outputBatchShape, paddedRhsType, rewriter);
+        materializePaddedBatchedWeight(plan.rhs, plan.problem.rhsBatchShape, plan.problem.outputBatchShape, paddedRhsType, rewriter);
       if (succeeded(paddedRhs)) {
-        Value paddedLhs = materializePaddedContractionInput(
-          plan.lhs, paddedLhsType, rewriter, loc);
+        Value paddedLhs = createPaddedInputCompute(plan.lhs, paddedLhsType, rewriter, loc);
         const int64_t laneCount = contractionPlan.laneCount;
         auto partialPiecesType = spatial::getGraphBatchPhysicalResultType(
           laneCount, RankedTensorType::get({1, xbarSize}, shapeInfo->outType.getElementType()));
         auto batchOp = createBatchedVmmBatch(paddedLhs,
                                              *paddedRhs,
                                              paddedLhsType,
-                                             plan.lhsBatchShape,
+                                             plan.problem.lhsBatchShape,
                                              paddedRhsType,
-                                             plan.rhsBatchShape,
-                                             plan.outputBatchShape,
+                                             plan.problem.rhsBatchShape,
+                                             plan.problem.outputBatchShape,
                                              partialPiecesType,
-                                             plan.m,
+                                             plan.problem.m,
                                              numKSlices,
                                              numOutHSlices,
                                              xbarSize,
@@ -1181,7 +1178,7 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
                                                     partialPiecesType,
                                                     directOutType,
                                                     paddedOutType,
-                                                    plan.batch,
+                                                    plan.problem.batch,
                                                     numKSlices,
                                                     xbarSize,
                                                     rewriter,
@@ -1190,12 +1187,7 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
           return failure();
         Value finalResult = *result;
         if (plan.transposedResult) {
-          auto transposedOutType = RankedTensorType::get({plan.batch, shapeInfo->m, shapeInfo->n},
-                                                         shapeInfo->outType.getElementType(),
-                                                         shapeInfo->outType.getEncoding());
-          finalResult =
-            ONNXTransposeOp::create(rewriter, loc, transposedOutType, finalResult, rewriter.getI64ArrayAttr({0, 2, 1}))
-              .getResult();
+          finalResult = transposeLastTwoDims(finalResult, rewriter, loc);
         }
         finalResult = finalizeNormalizedMatMulResult(finalResult, directOutType, *shapeInfo, rewriter, loc);
         rewriter.replaceOp(matmulOp, finalResult);
@@ -1205,24 +1197,24 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
     RankedTensorType blueprintType = !shapeInfo->lhsWasVector && !shapeInfo->rhsWasVector
                                    ? shapeInfo->outType : directOutType;
     SmallVector<int64_t> blueprintBatchShape = !shapeInfo->lhsWasVector && !shapeInfo->rhsWasVector
-                                             ? shapeInfo->outputBatchShape : SmallVector<int64_t> {plan.batch};
-    const int64_t rowsPerLane = chooseDynamicMatMulRowsPerLane(plan.m, plan.k, plan.n, xbarSize);
+                                             ? shapeInfo->problem.outputBatchShape : SmallVector<int64_t> {plan.problem.batch};
+    const int64_t rowsPerLane = chooseDynamicMatMulRowsPerLane(plan.problem.m, plan.problem.k, plan.problem.n, xbarSize);
     ContractionPlan contractionPlan = makeContractionPlan(
-      plan, target, ContractionPlanKind::GroupedRowDynamicVVD,
-      /*laneCount=*/plan.batch * plan.m / rowsPerLane, rowsPerLane);
+      plan.problem, target, ContractionPlanKind::GroupedRowDynamicVVD,
+      /*laneCount=*/plan.problem.batch * plan.problem.m / rowsPerLane, rowsPerLane);
     const int64_t laneCount = contractionPlan.laneCount;
     SmallVector<int64_t> fragmentShape(blueprintType.getRank(), 1);
     fragmentShape[fragmentShape.size() - 2] = rowsPerLane;
-    fragmentShape.back() = plan.n;
+    fragmentShape.back() = plan.problem.n;
     auto fragmentType = RankedTensorType::get(fragmentShape, shapeInfo->outType.getElementType());
     auto rowPiecesType = spatial::getGraphBatchPhysicalResultType(laneCount, fragmentType);
     Value transposedRhs = rhsStoredAsRows ? plan.rhs : transposeLastTwoDims(plan.rhs, rewriter, loc);
     auto batchOp = createBatchedVvdmulBatch(plan.lhs,
-                                            plan.lhsBatchShape,
+                                            plan.problem.lhsBatchShape,
                                             transposedRhs,
-                                            plan.rhsBatchShape,
-                                            plan.outputBatchShape,
-                                            plan.k,
+                                            plan.problem.rhsBatchShape,
+                                            plan.problem.outputBatchShape,
+                                            plan.problem.k,
                                             rowsPerLane,
                                             rowPiecesType,
                                             directOutType,
@@ -1231,23 +1223,33 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
                                             loc);
     if (failed(batchOp))
       return failure();
+    Value rowPieces = batchOp->getResult(0);
+    if (outputScale) {
+      auto scaleAttr = getHostConstDenseElementsAttr(outputScale);
+      if (!scaleAttr || !scaleAttr.isSplat())
+        return failure();
+      auto scaled = mapGraphBatchFragments(
+        rowPieces, rowPiecesType, rewriter, loc, [&](Value fragment, RankedTensorType fragmentType) {
+          auto splat = DenseElementsAttr::get(fragmentType, scaleAttr.getSplatValue<Attribute>());
+          Value scale = arith::ConstantOp::create(rewriter, loc, fragmentType, splat);
+          return FailureOr<Value>(
+            spatial::SpatVMulOp::create(rewriter, loc, fragmentType, fragment, scale).getResult());
+        });
+      if (failed(scaled))
+        return failure();
+      rowPieces = *scaled;
+      if (auto scaleBatch = rowPieces.getDefiningOp<spatial::SpatGraphComputeBatch>())
+        scaleBatch->setAttr("spat.preserve_compute_boundary", rewriter.getUnitAttr());
+    }
     auto result = createBatchedRowOutputBlueprint(
-      batchOp->getResult(0), blueprintType, blueprintBatchShape, rowsPerLane, rewriter, loc);
+      rowPieces, blueprintType, blueprintBatchShape, rowsPerLane, rewriter, loc);
     if (failed(result))
       return failure();
     Value finalResult = *result;
     if (plan.transposedResult) {
-      auto transposedOutType = RankedTensorType::get({plan.batch, shapeInfo->m, shapeInfo->n},
-                                                     shapeInfo->outType.getElementType(),
-                                                     shapeInfo->outType.getEncoding());
-      finalResult =
-        ONNXTransposeOp::create(rewriter, loc, transposedOutType, finalResult, rewriter.getI64ArrayAttr({0, 2, 1}))
-          .getResult();
+      finalResult = transposeLastTwoDims(finalResult, rewriter, loc);
     }
     finalResult = finalizeNormalizedMatMulResult(finalResult, directOutType, *shapeInfo, rewriter, loc);
-    if (outputScale)
-      finalResult = ONNXMulOp::create(
-        rewriter, loc, shapeInfo->outType, finalResult, outputScale).getResult();
     rewriter.replaceOp(matmulOp, finalResult);
     if (foldedTranspose && foldedTranspose->use_empty())
       rewriter.eraseOp(foldedTranspose);
@@ -1256,7 +1258,7 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
     return success();
   }
 
-  const spatial::SpatialTargetInfo& target;
+  const spatial::SpatialTargetResources& target;
 };
 
 struct TransposedRhsMatMulToSpatial : MatMulBatchedToSpatialComputes {
@@ -1273,14 +1275,14 @@ struct TransposedRhsMatMulToSpatial : MatMulBatchedToSpatialComputes {
 
 void populateMatMulFusionPatterns(RewritePatternSet& patterns,
                                   MLIRContext* ctx,
-                                  const spatial::SpatialTargetInfo& target) {
+                                  const spatial::SpatialTargetResources& target) {
   patterns.add<TransposedRhsMatMulToSpatial>(ctx, target);
 }
 
 void populateMatMulRewritePatterns(RewritePatternSet& patterns,
                                    MLIRContext* ctx,
-                                   const spatial::SpatialTargetInfo& target) {
-  patterns.insert<MatMulToGemm>(ctx);
+                                   const spatial::SpatialTargetResources& target) {
+  patterns.insert<MatMulToGemm>(ctx, target);
   patterns.insert<MatMulBatchedToSpatialComputes>(ctx, target);
 }
 

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

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

@@ -17,7 +17,7 @@
 #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"
 
@@ -161,10 +161,10 @@ struct PoolToSpatialCompute;
 
 template <typename PoolOp, typename PoolOpAdaptor, typename ReduceOp>
 struct PoolToSpatialComputeBase : public OpConversionPattern<PoolOp> {
-  PoolToSpatialComputeBase(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
+  PoolToSpatialComputeBase(MLIRContext* ctx, const spatial::SpatialTargetResources& target)
       : OpConversionPattern<PoolOp>(ctx), target(target) {}
 
-  const spatial::SpatialTargetInfo& target;
+  const spatial::SpatialTargetResources& target;
 
   LogicalResult matchAndRewrite(PoolOp poolOp, PoolOpAdaptor adaptor, ConversionPatternRewriter& rewriter) const final {
     Location loc = poolOp.getLoc();
@@ -431,7 +431,7 @@ struct PoolToSpatialCompute<ONNXAveragePoolOp>
 } // namespace
 
 LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp,
-                                            const spatial::SpatialTargetInfo&) {
+                                            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())
@@ -447,9 +447,10 @@ LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp,
 }
 
 FailureOr<Value> lowerDenseMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
-                                         const spatial::SpatialTargetInfo& target,
+                                         Value input,
+                                         const spatial::SpatialTargetResources& target,
                                          PatternRewriter& rewriter) {
-  auto inputType = dyn_cast<RankedTensorType>(planOp.getInput().getType());
+  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)
@@ -476,7 +477,7 @@ FailureOr<Value> lowerDenseMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
   const int64_t outputPatchCount = batchSize * outputHeight * outputWidth;
 
   auto compute = createSpatCompute<1>(
-    rewriter, planOp.getLoc(), outputType, {}, planOp.getInput(),
+    rewriter, planOp.getLoc(), outputType, {}, input,
     [&](Value input) -> LogicalResult {
       Value paddedInput = createPaddedPoolInput(
         rewriter, planOp.getLoc(), input, inputType,
@@ -615,14 +616,15 @@ static Value extractPoolIndex(PatternRewriter& rewriter,
 }
 
 FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
+                                            Value input,
                                             std::optional<Value> rowStripInput,
-                                            const spatial::SpatialTargetInfo& target,
+                                            const spatial::SpatialTargetResources& target,
                                             PatternRewriter& rewriter) {
   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);
@@ -631,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();
@@ -681,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);
@@ -806,7 +808,7 @@ FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
 }
 
 LogicalResult canLowerGlobalAveragePoolPlanToRowStrip(
-  spatial::SpatGlobalAveragePoolPlanOp planOp, const spatial::SpatialTargetInfo&) {
+  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())
@@ -820,9 +822,10 @@ LogicalResult canLowerGlobalAveragePoolPlanToRowStrip(
 
 FailureOr<Value> lowerDenseGlobalAveragePoolPlan(
     spatial::SpatGlobalAveragePoolPlanOp planOp,
-    const spatial::SpatialTargetInfo& target,
+    Value input,
+    const spatial::SpatialTargetResources& target,
     PatternRewriter& rewriter) {
-  auto inputType = dyn_cast<RankedTensorType>(planOp.getInput().getType());
+  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
@@ -844,7 +847,7 @@ FailureOr<Value> lowerDenseGlobalAveragePoolPlan(
   const double scaleValue = 1.0 / static_cast<double>(height * width);
 
   auto compute = createSpatCompute<1>(
-    rewriter, planOp.getLoc(), outputType, {}, planOp.getInput(),
+    rewriter, planOp.getLoc(), outputType, {}, input,
     [&](Value input) -> LogicalResult {
       Value output = tensor::EmptyOp::create(
         rewriter, planOp.getLoc(), outputType.getShape(), outputType.getElementType());
@@ -895,22 +898,23 @@ FailureOr<Value> lowerDenseGlobalAveragePoolPlan(
 }
 
 FailureOr<Value> lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp,
+                                                    Value input,
                                                     std::optional<Value> rowStripInput,
-                                                    const spatial::SpatialTargetInfo& target,
+                                                    const spatial::SpatialTargetResources& target,
                                                     PatternRewriter& rewriter) {
   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();
@@ -940,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) {
@@ -1011,7 +1015,7 @@ FailureOr<Value> lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePo
 
 void populatePoolPatterns(RewritePatternSet& patterns,
                           MLIRContext* ctx,
-                          const spatial::SpatialTargetInfo& target) {
+                          const spatial::SpatialTargetResources& target) {
   patterns.insert<PoolToSpatialCompute<ONNXMaxPoolSingleOutOp>>(ctx, target);
   patterns.insert<PoolToSpatialCompute<ONNXAveragePoolOp>>(ctx, target);
 }

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

@@ -7,7 +7,7 @@
 #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"
@@ -47,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;
 
@@ -75,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;
 
@@ -97,21 +116,17 @@ 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;
@@ -119,12 +134,11 @@ struct RowStripFlattenAnalysis {
 };
 
 static FailureOr<RowStripFlattenAnalysis> analyzeRowStripFlatten(
-  spatial::SpatGraphCompute flattenOp, const spatial::SpatialTargetInfo& target) {
-  if (flattenOp.getWeights().size() != 0 || flattenOp.getInputs().size() != 1
-      || flattenOp.getOutputs().size() != 1)
+    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())
@@ -134,17 +148,21 @@ static FailureOr<RowStripFlattenAnalysis> analyzeRowStripFlatten(
   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,32 @@ static FailureOr<RowStripFlattenAnalysis> analyzeRowStripFlatten(
   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,
-                                          const spatial::SpatialTargetInfo& target) {
+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,
-                                       const spatial::SpatialTargetInfo& target,
+                                       spatial::SpatFlattenPlanOp flattenOp,
+                                       const spatial::SpatialTargetResources& target,
                                        PatternRewriter& rewriter) {
   FailureOr<RowStripFlattenAnalysis> analysis = analyzeRowStripFlatten(flattenOp, target);
   if (failed(analysis))
@@ -206,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

@@ -9,7 +9,7 @@
 #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/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"
@@ -193,7 +193,7 @@ struct Resize : OpConversionPattern<ONNXResizeOp> {
 
 LogicalResult canLowerResizeNearestPlanToRowStrip(
     spatial::SpatResizeNearestPlanOp planOp,
-    const spatial::SpatialTargetInfo&) {
+    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()
@@ -204,16 +204,17 @@ LogicalResult canLowerResizeNearestPlanToRowStrip(
 }
 
 FailureOr<Value> lowerSelectedResizeNearestPlan(
-    spatial::SpatResizeNearestPlanOp planOp, std::optional<Value> rowStripInput,
-    const spatial::SpatialTargetInfo&,
+    spatial::SpatResizeNearestPlanOp planOp, Value input,
+    std::optional<Value> rowStripInput,
+    const spatial::SpatialTargetResources&,
     PatternRewriter& rewriter) {
-  auto inputType = cast<RankedTensorType>(planOp.getInput().getType());
+  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(
-    planOp.getInput(), inputType, outputType, rewriter, planOp.getLoc());
+    input, inputType, outputType, rewriter, planOp.getLoc());
 }
 
 void populateResizePatterns(RewritePatternSet& patterns, MLIRContext* ctx) { patterns.add<Resize>(ctx); }

src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp

@@ -1,265 +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/RowStripLayoutUtils.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-#include "src/Accelerators/PIM/Pass/PIMPasses.h"
-
-#include <algorithm>
-
-using namespace mlir;
-
-namespace onnx_mlir {
-namespace {
-
-using LayoutMap = llvm::DenseMap<Value, spatial::PhysicalLayout>;
-
-static spatial::PhysicalLayout getSelectedLayout(const LayoutMap& layouts, Value value) {
-  if (auto it = layouts.find(value); it != layouts.end())
-    return it->second;
-  if (auto materialize = value.getDefiningOp<spatial::SpatMaterializeLayoutOp>())
-    return materialize.getTargetPhysicalLayout();
-  if (auto blueprint = value.getDefiningOp<spatial::SpatBlueprintOp>())
-    return blueprint.getPhysicalLayout();
-  return spatial::PhysicalLayout::DenseNCHW;
-}
-
-static SmallVector<spatial::PhysicalLayout> getOperandLayouts(
-    Operation* op, const LayoutMap& layouts) {
-  SmallVector<spatial::PhysicalLayout> operandLayouts;
-  operandLayouts.reserve(op->getNumOperands());
-  for (Value operand : op->getOperands())
-    operandLayouts.push_back(getSelectedLayout(layouts, operand));
-  return operandLayouts;
-}
-
-static FailureOr<SmallVector<spatial::LayoutAlternative>> getAlternatives(
-    Operation* op, const LayoutMap& layouts, const spatial::SpatialTargetInfo& target) {
-  auto capability = dyn_cast<spatial::SpatialLayoutCapabilityInterface>(op);
-  if (!capability)
-    return failure();
-  SmallVector<spatial::LayoutAlternative> alternatives =
-      capability.getLayoutAlternatives(target, getOperandLayouts(op, layouts));
-  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();
-  return alternatives;
-}
-
-static unsigned findCurrentAlternative(
-    Operation* op, ArrayRef<spatial::LayoutAlternative> alternatives,
-    spatial::PhysicalLayout selectedResult) {
-  for (auto [index, alternative] : llvm::enumerate(alternatives))
-    if (alternative.resultLayout == selectedResult)
-      return index;
-  return 0;
-}
-
-static int64_t alternativeCost(Operation* op,
-                               const spatial::LayoutAlternative& alternative,
-                               const LayoutMap& layouts,
-                               const LayoutMap& selectedResults,
-                               const spatial::SpatialTargetInfo& target) {
-  int64_t cost = alternative.intrinsicCost;
-  SmallVector<spatial::PhysicalLayout> operandLayouts = getOperandLayouts(op, layouts);
-  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) {
-      if (alternative.resultLayout != spatial::PhysicalLayout::DenseNCHW) {
-        auto flatten = dyn_cast<spatial::SpatGraphCompute>(use.getOwner());
-        if (!flatten || failed(canLowerFlattenFromRowStrip(flatten, target)))
-          ++cost;
-      }
-      continue;
-    }
-    auto userAlternatives = getAlternatives(use.getOwner(), selectedResults, target);
-    if (failed(userAlternatives))
-      continue;
-    spatial::PhysicalLayout userResult =
-        selectedResults.lookup(use.getOwner()->getResult(0));
-    unsigned userIndex = findCurrentAlternative(use.getOwner(), *userAlternatives, userResult);
-    if (use.getOperandNumber() < (*userAlternatives)[userIndex].operandLayouts.size()
-        && (*userAlternatives)[userIndex].operandLayouts[use.getOperandNumber()]
-             != alternative.resultLayout)
-      ++cost;
-  }
-  return cost;
-}
-
-static LogicalResult materializeMismatchedUses(
-    IRRewriter& rewriter, Value value, const LayoutMap& layouts,
-    const spatial::SpatialTargetInfo& target) {
-  spatial::PhysicalLayout sourceLayout = getSelectedLayout(layouts, value);
-  SmallVector<std::pair<OpOperand*, spatial::PhysicalLayout>> mismatches;
-  for (OpOperand& use : value.getUses()) {
-    Operation* userOp = use.getOwner();
-    spatial::PhysicalLayout required = spatial::PhysicalLayout::DenseNCHW;
-    if (auto capability = dyn_cast<spatial::SpatialLayoutCapabilityInterface>(userOp)) {
-      auto alternatives = getAlternatives(userOp, layouts, target);
-      if (failed(alternatives))
-        return failure();
-      spatial::PhysicalLayout selected =
-          getSelectedLayout(layouts, userOp->getResult(0));
-      unsigned selectedIndex = findCurrentAlternative(userOp, *alternatives, selected);
-      required = (*alternatives)[selectedIndex].operandLayouts[use.getOperandNumber()];
-    }
-    else if (auto flatten = dyn_cast<spatial::SpatGraphCompute>(userOp);
-             flatten && sourceLayout == spatial::PhysicalLayout::NHWCRowStrip
-                  && succeeded(canLowerFlattenFromRowStrip(flatten, target))) {
-      continue;
-    }
-    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(
-    ArrayRef<Operation*> planOps, const LayoutMap& layouts,
-    const spatial::SpatialTargetInfo& target) {
-  for (Operation* op : planOps) {
-    auto selected = spatial::getSelectedPhysicalLayout(op);
-    if (!selected)
-      return op->emitOpError("requires a selected physical layout"), failure();
-    auto alternatives = getAlternatives(op, layouts, target);
-    if (failed(alternatives))
-      return failure();
-    if (llvm::none_of(*alternatives, [&](const spatial::LayoutAlternative& alternative) {
-          return alternative.resultLayout == *selected;
-        }))
-      return op->emitOpError("selected physical layout is not advertised by its 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;
-  explicit SpatialLayoutPlanningPass(const spatial::SpatialTargetInfo& target)
-      : target(target), hasTarget(true) {}
-
-  void runOnOperation() override {
-    ModuleOp moduleOp = getOperation();
-    if (!hasTarget) {
-      moduleOp.emitError("Spatial layout planning requires an injected SpatialTargetInfo");
-      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;
-    SmallVector<Operation*> planOps;
-    for (Operation& op : funcOp.getBody().front())
-      if (isa<spatial::SpatialLayoutCapabilityInterface>(&op))
-        planOps.push_back(&op);
-
-    LayoutMap layouts;
-    for (Operation* op : planOps)
-      layouts[op->getResult(0)] = spatial::PhysicalLayout::DenseNCHW;
-
-    const size_t maxRounds = 2 * planOps.size() + 1;
-    bool converged = false;
-    for (size_t round = 0; round < maxRounds && !converged; ++round) {
-      converged = true;
-      SmallVector<Operation*> order(planOps);
-      if (round % 2)
-        std::reverse(order.begin(), order.end());
-      for (Operation* op : order) {
-        auto alternatives = getAlternatives(op, layouts, target);
-        if (failed(alternatives)) {
-          signalPassFailure();
-          return;
-        }
-        spatial::PhysicalLayout current = layouts.lookup(op->getResult(0));
-        unsigned currentIndex = findCurrentAlternative(op, *alternatives, current);
-        int64_t bestCost = alternativeCost(
-            op, (*alternatives)[currentIndex], layouts, layouts, target);
-        unsigned bestIndex = currentIndex;
-        for (auto [index, alternative] : llvm::enumerate(*alternatives)) {
-          int64_t cost = alternativeCost(op, alternative, layouts, layouts, target);
-          if (cost < bestCost) {
-            bestCost = cost;
-            bestIndex = index;
-          }
-        }
-        spatial::PhysicalLayout selected = (*alternatives)[bestIndex].resultLayout;
-        if (selected != current) {
-          layouts[op->getResult(0)] = selected;
-          converged = false;
-        }
-      }
-    }
-    if (!converged) {
-      moduleOp.emitError("Spatial layout selection did not converge within its bounded iteration budget");
-      signalPassFailure();
-      return;
-    }
-    IRRewriter rewriter(&getContext());
-    for (Operation* op : planOps) {
-      op->setAttr(spatial::kSelectedLayoutAttrName,
-                  spatial::PhysicalLayoutAttr::get(
-                      rewriter.getContext(), layouts.lookup(op->getResult(0))));
-      if (failed(materializeMismatchedUses(rewriter, op->getResult(0), layouts, target))) {
-        signalPassFailure();
-        return;
-      }
-    }
-    if (failed(verifySelectedLayouts(planOps, layouts, target))
-        || failed(verifyLogicalSpatialGraphInvariants(*entryFunc))) {
-      moduleOp.emitError("Spatial layout planning verification failed");
-      signalPassFailure();
-    }
-  }
-
-  spatial::SpatialTargetInfo target;
-  bool hasTarget = false;
-};
-
-} // namespace
-
-std::unique_ptr<Pass> createSpatialLayoutPlanningPass() {
-  return std::make_unique<SpatialLayoutPlanningPass>();
-}
-
-std::unique_ptr<Pass> createSpatialLayoutPlanningPass(
-    const spatial::SpatialTargetInfo& target) {
-  return std::make_unique<SpatialLayoutPlanningPass>(target);
-}
-
-} // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.cpp

@@ -307,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());
@@ -317,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 =

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

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

src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp

@@ -345,20 +345,39 @@ 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,
-          arith::ConstantIndexOp::create(rewriter, receiveOp.getLoc(), 0),
-          *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();
+      } else {
+        received = PimReceiveOp::create(
+          rewriter, receiveOp->getLoc(), outputBuffer.getType(), outputBuffer,
+          zero, *sizeAttr, channelReceive.getSourceCoreId()).getOutput();
+      }
       blockArg->replaceAllUsesWith(received);
       markOpToRemove(receiveOp);
       continue;
@@ -383,11 +402,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);
@@ -397,7 +417,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/ChannelLowering.cpp

@@ -57,10 +57,29 @@ struct ChannelSendLowering : OpRewritePattern<spatial::SpatChannelSendOp> {
   }
 };
 
-struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp> {
+struct HostStoreSyncLowering : OpRewritePattern<spatial::SpatHostStoreSyncOp> {
   using OpRewritePattern::OpRewritePattern;
 
-  LogicalResult matchAndRewrite(spatial::SpatChannelReceiveOp op, PatternRewriter& rewriter) const override {
+  LogicalResult matchAndRewrite(spatial::SpatHostStoreSyncOp op, PatternRewriter& rewriter) const override {
+    auto sizeAttr = getTensorSizeInBytesAttr(rewriter, op.getOperation(), op.getInput());
+    auto hostBuffer = getPipelineHostBuffer(rewriter, op);
+    if (failed(sizeAttr) || failed(hostBuffer))
+      return failure();
+    Value zero = arith::ConstantIndexOp::create(rewriter, op.getLoc(), 0);
+    pim::PimMemCopyDevToHostOp::create(
+      rewriter, op.getLoc(), hostBuffer->getType(), op.getHostOffset(), zero,
+      *hostBuffer, op.getInput(), *sizeAttr);
+    auto sync = pim::PimSyncOp::create(
+      rewriter, op.getLoc(), op.getTargetCoreId(), op.getEventRegister());
+    copyRaptorDebugAttrs(op.getOperation(), sync.getOperation());
+    rewriter.eraseOp(op);
+    return success();
+  }
+};
+
+template <typename ReceiveOp, typename CreateReceive>
+static LogicalResult lowerReceive(
+    ReceiveOp op, PatternRewriter& rewriter, CreateReceive createReceive) {
     if (op->use_empty()) {
       rewriter.eraseOp(op);
       return success();
@@ -86,12 +105,11 @@ struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp>
     if (failed(sizeAttr))
       return failure();
     Value zero = arith::ConstantIndexOp::create(rewriter, op.getLoc(), 0);
-    auto receive = pim::PimReceiveOp::create(
-      rewriter, op.getLoc(), op.getResult().getType(), outputBuffer, zero, *sizeAttr, op.getSourceCoreId());
-    copyRaptorDebugAttrs(op.getOperation(), receive.getOperation());
-    Value received = receive.getOutput();
+    auto received = createReceive(outputBuffer, zero, *sizeAttr);
+    if (failed(received))
+      return failure();
     if (!destinationInsert) {
-      rewriter.replaceOp(op, received);
+      rewriter.replaceOp(op, *received);
       return success();
     }
 
@@ -99,10 +117,65 @@ struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp>
     Value targetOffset = createDestinationByteOffset(rewriter, destinationInsert);
     auto copy = pim::PimMemCopyOp::create(
       rewriter, op.getLoc(), destinationInsert.getDestType(), targetOffset, zero,
-      destinationInsert.getDest(), received, *sizeAttr);
+      destinationInsert.getDest(), *received, *sizeAttr);
     rewriter.replaceOp(destinationInsert, copy.getOutput());
     rewriter.eraseOp(op);
     return success();
+}
+
+struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatChannelReceiveOp op, PatternRewriter& rewriter) const override {
+    return lowerReceive(op, rewriter, [&](Value outputBuffer, Value zero, IntegerAttr sizeAttr) -> FailureOr<Value> {
+      auto receive = pim::PimReceiveOp::create(
+        rewriter, op.getLoc(), op.getResult().getType(), outputBuffer, zero,
+        sizeAttr, op.getSourceCoreId());
+      copyRaptorDebugAttrs(op.getOperation(), receive.getOperation());
+      return receive.getOutput();
+    });
+  }
+};
+
+struct HostWaitLoadLowering : OpRewritePattern<spatial::SpatHostWaitLoadOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatHostWaitLoadOp op, PatternRewriter& rewriter) const override {
+    return lowerReceive(op, rewriter, [&](Value outputBuffer, Value zero, IntegerAttr sizeAttr) -> FailureOr<Value> {
+      auto hostBuffer = getPipelineHostBuffer(rewriter, op);
+      if (failed(hostBuffer))
+        return failure();
+      auto wait = pim::PimWaitOp::create(
+        rewriter, op.getLoc(), op.getEventRegister(),
+        op.getWaitValue());
+      copyRaptorDebugAttrs(op.getOperation(), wait.getOperation());
+      Value output = pim::PimMemCopyHostToDevOp::create(
+        rewriter, op.getLoc(), outputBuffer.getType(), zero,
+        op.getHostOffset(), outputBuffer, *hostBuffer, sizeAttr).getOutput();
+      return output;
+    });
+  }
+};
+
+struct SyncLowering : OpRewritePattern<spatial::SpatSyncOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatSyncOp op,
+                                PatternRewriter& rewriter) const override {
+    rewriter.replaceOpWithNewOp<pim::PimSyncOp>(
+      op, op.getTargetCoreId(), op.getEventRegister());
+    return success();
+  }
+};
+
+struct WaitLowering : OpRewritePattern<spatial::SpatWaitOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatWaitOp op,
+                                PatternRewriter& rewriter) const override {
+    rewriter.replaceOpWithNewOp<pim::PimWaitOp>(
+      op, op.getEventRegister(), op.getWaitValue());
+    return success();
   }
 };
 
@@ -148,7 +221,10 @@ struct ConcatLowering : OpRewritePattern<spatial::SpatConcatOp> {
 } // namespace
 
 void populateChannelLoweringPatterns(RewritePatternSet& patterns) {
-  patterns.add<ChannelSendLowering, ChannelReceiveLowering, ExtractRowsLowering, ConcatLowering>(patterns.getContext());
+  patterns.add<ChannelSendLowering, ChannelReceiveLowering,
+               HostStoreSyncLowering, HostWaitLoadLowering,
+               SyncLowering, WaitLowering, ExtractRowsLowering,
+               ConcatLowering>(patterns.getContext());
 }
 
 } // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/ReturnPathNormalization.cpp

@@ -734,7 +734,7 @@ raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::low
     auto storedType = dyn_cast<RankedTensorType>(storedValue.getType());
     if (!storedType) {
       producerOp->emitOpError(
-        "has an unsupported non-ranked concat-return helper yield during Spatial-to-PIM lowering");
+        "has an unsupported non-ranked concat-return helper yield during Spatial-to-Pim lowering");
       return ReturnPathLoweringResult::Failure;
     }
     rewriter.setInsertionPointAfterValue(storedValue);
@@ -748,7 +748,7 @@ raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::low
       SmallVector<int64_t> destinationIndices;
       if (failed(mapIndicesThroughHelperChain(
             sourceIndices, concatReturnUse->concatShape, concatReturnUse->helperChain, destinationIndices))) {
-        producerOp->emitOpError("has an unsupported concat-return helper chain during Spatial-to-PIM lowering");
+        producerOp->emitOpError("has an unsupported concat-return helper chain during Spatial-to-Pim lowering");
         return ReturnPathLoweringResult::Failure;
       }
 
@@ -859,6 +859,10 @@ void raptor::SpatialToPimPass::replaceReturnWithOutputBuffers(func::ReturnOp ret
       markOpToRemove(receiveOp);
       return;
     }
+    if (auto receiveOp = dyn_cast<spatial::SpatHostWaitLoadOp>(op)) {
+      markOpToRemove(receiveOp);
+      return;
+    }
   };
 
   SmallVector<Value> originalOperands(returnOp.getOperands().begin(), returnOp.getOperands().end());

src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp

@@ -29,13 +29,13 @@
 #include "Common/IR/ConstantUtils.hpp"
 #include "Common/PimCommon.hpp"
 #include "Common/Support/CheckedArithmetic.hpp"
-#include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
+#include "Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.hpp"
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "Conversion/SpatialToPim/Common.hpp"
 #include "Conversion/SpatialToPim/Patterns.hpp"
 #include "Dialect/Pim/PimOps.hpp"
 #include "Dialect/Spatial/SpatialOps.hpp"
-#include "Pass/PIMPasses.h"
+#include "Passes/PIMPasses.h"
 #include "SpatialToPimPass.hpp"
 
 using namespace mlir;
@@ -66,17 +66,20 @@ createZeroPaddedTensor(IRRewriter& rewriter, Location loc, Value value, RankedTe
   return padOp.getResult();
 }
 
-static FailureOr<Value> padHVectorInputToCrossbarSize(IRRewriter& rewriter, Location loc, Value vector) {
+static FailureOr<Value> padHVectorInputToCrossbarSize(IRRewriter& rewriter,
+                                                      Location loc,
+                                                      Value vector,
+                                                      int64_t crossbarSize) {
   auto vectorType = cast<RankedTensorType>(vector.getType());
   ArrayRef<int64_t> shape = vectorType.getShape();
   assert(isHVectorShape(shape) && "expected a horizontal vector");
-  assert(shape[1] <= static_cast<int64_t>(crossbarSize) && "vector width must fit in one crossbar");
+  assert(shape[1] <= crossbarSize && "vector width must fit in one crossbar");
 
-  if (shape[1] == static_cast<int64_t>(crossbarSize))
+  if (shape[1] == crossbarSize)
     return vector;
 
   auto paddedType = RankedTensorType::get(
-    {shape[0], static_cast<int64_t>(crossbarSize)}, vectorType.getElementType(), vectorType.getEncoding());
+    {shape[0], crossbarSize}, vectorType.getElementType(), vectorType.getEncoding());
   return createZeroPaddedTensor(rewriter, loc, vector, paddedType);
 }
 
@@ -84,11 +87,16 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
   outputTensors.clear();
   operationsToRemove.clear();
   ModuleOp moduleOp = getOperation();
+  if (!hasTarget || failed(targetResources.verify())) {
+    moduleOp.emitError("Spatial-to-Pim lowering requires valid injected target resources");
+    signalPassFailure();
+    return;
+  }
   MLIRContext* ctx = moduleOp.getContext();
 
   auto entryFunc = getPimEntryFunc(moduleOp);
   if (failed(entryFunc)) {
-    moduleOp.emitError("failed to locate the PIM entry function during Spatial-to-PIM lowering");
+    moduleOp.emitError("failed to locate the Pim entry function during Spatial-to-Pim lowering");
     signalPassFailure();
     return;
   }
@@ -118,12 +126,16 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
                     spatial::SpatConcatOp,
                     spatial::SpatChannelReceiveOp,
                     spatial::SpatChannelSendOp,
+                    spatial::SpatHostStoreSyncOp,
+                    spatial::SpatHostWaitLoadOp,
+                    spatial::SpatSyncOp,
+                    spatial::SpatWaitOp,
                     spatial::SpatExtractRowsOp>();
 
   RewritePatternSet initialPatterns(ctx);
   populateInitialPatterns(initialPatterns);
   if (failed(applyPartialConversion(moduleOp, target, std::move(initialPatterns)))) {
-    moduleOp.emitError("failed to lower required Spatial ops to the initial PIM form");
+    moduleOp.emitError("failed to lower required Spatial ops to the initial Pim form");
     signalPassFailure();
     return;
   }
@@ -132,10 +144,16 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
   populateGlobalTensorMaterializationPatterns(globalTensorPatterns);
   walkAndApplyPatterns(moduleOp, std::move(globalTensorPatterns));
 
+  if (funcOp->hasAttr(kPipelineHostBufferBytesAttrName)
+      && failed(materializePipelineHostBuffer(funcOp, rewriter))) {
+    signalPassFailure();
+    return;
+  }
+
   auto returnOp = cast<func::ReturnOp>(funcOp.front().getTerminator());
   addReturnOutputBuffers(returnOp, rewriter);
   if (failed(allocateAndInitializeCoreLocalVariables(funcOp, rewriter))) {
-    funcOp.emitOpError("failed to allocate or initialize core-local tensors during Spatial-to-PIM lowering");
+    funcOp.emitOpError("failed to allocate or initialize core-local tensors during Spatial-to-Pim lowering");
     signalPassFailure();
     return;
   }
@@ -174,6 +192,17 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
       continue;
     }
   }
+  SmallVector<spatial::SpatHostWaitLoadOp> hostWaitLoadOps;
+  for (auto op : funcOp.getOps<spatial::SpatHostWaitLoadOp>())
+    hostWaitLoadOps.push_back(op);
+  for (auto op : hostWaitLoadOps) {
+    bool onlyPendingRemovalUsers = llvm::all_of(
+      op->getUsers(), [&](Operation* user) {
+        return llvm::is_contained(operationsToRemove, user);
+      });
+    if (onlyPendingRemovalUsers)
+      markOpToRemove(op);
+  }
 
   RewritePatternSet coreBodyPatterns(ctx);
   populateCoreBodyPatterns(coreBodyPatterns);
@@ -194,6 +223,10 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
                             spatial::SpatConcatOp,
                             spatial::SpatChannelReceiveOp,
                             spatial::SpatChannelSendOp,
+                            spatial::SpatHostStoreSyncOp,
+                            spatial::SpatHostWaitLoadOp,
+                            spatial::SpatSyncOp,
+                            spatial::SpatWaitOp,
                             spatial::SpatExtractRowsOp>();
 
   SmallVector<pim::PimCoreOp> coreOps;
@@ -243,12 +276,16 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
   communicationTarget.addIllegalOp<spatial::SpatConcatOp,
                                    spatial::SpatChannelReceiveOp,
                                    spatial::SpatChannelSendOp,
+                                   spatial::SpatHostStoreSyncOp,
+                                   spatial::SpatHostWaitLoadOp,
+                                   spatial::SpatSyncOp,
+                                   spatial::SpatWaitOp,
                                    spatial::SpatExtractRowsOp>();
 
   RewritePatternSet communicationPatterns(ctx);
   populateChannelLoweringPatterns(communicationPatterns);
   if (failed(applyFullConversion(funcOp, communicationTarget, std::move(communicationPatterns)))) {
-    funcOp.emitOpError("failed to lower Spatial communication ops to PIM communication ops");
+    funcOp.emitOpError("failed to lower Spatial communication ops to Pim communication ops");
     signalPassFailure();
     return;
   }
@@ -265,15 +302,16 @@ LogicalResult raptor::SpatialToPimPass::enlargeVMMOutTensorsToCrossbarSize(func:
     ArrayRef<int64_t> outputShape = outputType.getShape();
     assert(isHVectorShape(outputShape) && "expected a horizontal vector output");
     auto weightType = cast<RankedTensorType>(vmmOp.getWeight().getType());
-    const int64_t xbarDim = static_cast<int64_t>(crossbarSize);
+    const int64_t xbarDim = static_cast<int64_t>(targetResources.matrixShape.columns);
     const int64_t paddedOutputWidth = ceilIntegerDivide(outputShape[1], xbarDim) * xbarDim;
     assert(weightType.getRank() == 2 && weightType.getDimSize(1) == paddedOutputWidth
            && "expected VMM weight width to match the padded output width");
-    assert(paddedOutputWidth / xbarDim <= static_cast<int64_t>(crossbarCountInCore)
+    assert(paddedOutputWidth / xbarDim <= static_cast<int64_t>(targetResources.matrixUnitsPerProcessor)
            && "output width must fit in one core");
 
     rewriter.setInsertionPoint(vmmOp);
-    auto paddedInput = padHVectorInputToCrossbarSize(rewriter, vmmOp.getLoc(), vmmOp.getInput());
+    auto paddedInput = padHVectorInputToCrossbarSize(
+      rewriter, vmmOp.getLoc(), vmmOp.getInput(), xbarDim);
     if (failed(paddedInput)) {
       hasFailure = true;
       return WalkResult::interrupt();
@@ -375,4 +413,9 @@ void raptor::SpatialToPimPass::eraseOpsToRemove() {
 
 std::unique_ptr<Pass> createSpatialToPimPass() { return std::make_unique<raptor::SpatialToPimPass>(); }
 
+std::unique_ptr<Pass> createSpatialToPimPass(
+    const spatial::SpatialTargetResources& target) {
+  return std::make_unique<raptor::SpatialToPimPass>(target);
+}
+
 } // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/SpatialToPimPass.hpp

@@ -18,6 +18,7 @@
 #include "Conversion/SpatialToPim/Common.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
 
 namespace onnx_mlir {
 namespace raptor {
@@ -25,10 +26,13 @@ namespace raptor {
 struct SpatialToPimPass : mlir::PassWrapper<SpatialToPimPass, mlir::OperationPass<mlir::ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SpatialToPimPass)
   llvm::StringRef getArgument() const override { return "convert-spatial-to-pim"; }
-  llvm::StringRef getDescription() const override { return "Lower Spatial ops to PIM-ready format"; }
+  llvm::StringRef getDescription() const override { return "Lower Spatial ops to Pim-ready format"; }
 
   SpatialToPimPass() = default;
-  SpatialToPimPass(const SpatialToPimPass& pass) {}
+  explicit SpatialToPimPass(const spatial::SpatialTargetResources& target)
+      : targetResources(target), hasTarget(true) {}
+  SpatialToPimPass(const SpatialToPimPass& pass)
+      : targetResources(pass.targetResources), hasTarget(pass.hasTarget) {}
 
   void runOnOperation() final;
 
@@ -37,6 +41,8 @@ private:
 
   llvm::SmallVector<OutputTensorFactory> outputTensors;
   llvm::SmallVector<mlir::Operation*> operationsToRemove;
+  spatial::SpatialTargetResources targetResources;
+  bool hasTarget = false;
 
   mlir::LogicalResult allocateAndInitializeCoreLocalVariables(mlir::func::FuncOp funcOp, mlir::IRRewriter& rewriter);
   mlir::LogicalResult

src/PIM/Dialect/Pim/CMakeLists.txt

@@ -1,12 +1,12 @@
 add_onnx_mlir_dialect(Pim pim)
 add_onnx_mlir_dialect_doc(pim Pim.td)
 
-add_subdirectory(Analysis)
-add_subdirectory(Transforms/Bufferization)
-add_subdirectory(Transforms/HostConstantFolding)
-add_subdirectory(Transforms/InstructionSelection)
-add_subdirectory(Transforms/LocalMemoryPlanning)
-add_subdirectory(Transforms/Verification)
+add_subdirectory(Passes/Analyses)
+add_subdirectory(Passes/Transforms/Bufferization)
+add_subdirectory(Passes/Transforms/HostConstantFolding)
+add_subdirectory(Passes/Transforms/InstructionSelection)
+add_subdirectory(Passes/Transforms/LocalMemoryPlanning)
+add_subdirectory(Passes/Transforms/Verification)
 
 add_pim_library(PimOps
   PimOps.hpp

src/PIM/Dialect/Pim/Passes/Analyses/LocalMemoryLifetimeAnalysis.cpp

@@ -8,7 +8,7 @@
 
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
-#include "src/Accelerators/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/Passes/Analyses/LocalMemoryLifetimeAnalysis.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 
 using namespace mlir;

src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/BufferizationUtils.cpp

@@ -3,8 +3,8 @@
 
 #include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/BufferizationUtils.hpp"
-#include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/Common.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/Passes/Transforms/Bufferization/BufferizationUtils.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/Passes/Transforms/Bufferization/Common.hpp"
 
 using namespace mlir;
 using namespace bufferization;

src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/Common.cpp

@@ -1,4 +1,4 @@
-#include "Dialect/Pim/Transforms/Bufferization/Common.hpp"
+#include "Dialect/Pim/Passes/Transforms/Bufferization/Common.hpp"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"

src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/ContiguityPatterns.cpp

@@ -302,12 +302,19 @@ static FailureOr<int64_t> getShapedByteSize(MemRefType type) {
   return static_cast<int64_t>(*byteSize);
 }
 
-static FailureOr<SmallVector<int64_t>>
+struct LogicalCopyShape {
+  SmallVector<int64_t> dimensions;
+  Type elementType;
+};
+
+static bool isPackedByteBuffer(MemRefType type) {
+  return type.getRank() == 1 && type.getElementType().isInteger(8);
+}
+
+static FailureOr<LogicalCopyShape>
 inferLogicalCopyShape(MemRefType targetType, MemRefType sourceType, int64_t size) {
   if (!targetType.hasStaticShape() || !sourceType.hasStaticShape())
     return failure();
-  if (targetType.getElementType() != sourceType.getElementType() || targetType.getRank() != sourceType.getRank())
-    return failure();
 
   auto targetBytes = getShapedByteSize(targetType);
   auto sourceBytes = getShapedByteSize(sourceType);
@@ -316,18 +323,37 @@ inferLogicalCopyShape(MemRefType targetType, MemRefType sourceType, int64_t size
 
   bool targetMatches = *targetBytes == size;
   bool sourceMatches = *sourceBytes == size;
-  if (targetMatches && sourceMatches && targetType.getShape() != sourceType.getShape())
+  bool matchingTypes = targetType.getElementType() == sourceType.getElementType()
+                    && targetType.getRank() == sourceType.getRank();
+  if (matchingTypes) {
+    if (targetMatches && sourceMatches
+        && targetType.getShape() != sourceType.getShape())
+      return failure();
+    MemRefType logicalType = targetMatches ? targetType : sourceType;
+    if (targetMatches || sourceMatches)
+      return LogicalCopyShape {
+        SmallVector<int64_t>(logicalType.getShape()),
+        logicalType.getElementType()};
     return failure();
-  if (targetMatches)
-    return SmallVector<int64_t>(targetType.getShape().begin(), targetType.getShape().end());
-  if (sourceMatches)
-    return SmallVector<int64_t>(sourceType.getShape().begin(), sourceType.getShape().end());
+  }
+  if (targetMatches && isPackedByteBuffer(sourceType))
+    return LogicalCopyShape {
+      SmallVector<int64_t>(targetType.getShape()),
+      targetType.getElementType()};
+  if (sourceMatches && isPackedByteBuffer(targetType))
+    return LogicalCopyShape {
+      SmallVector<int64_t>(sourceType.getShape()),
+      sourceType.getElementType()};
   return failure();
 }
 
-static FailureOr<int64_t> getContiguousSuffixRank(Value value, ArrayRef<int64_t> copyShape) {
+static FailureOr<int64_t> getContiguousSuffixRank(
+    Value value, ArrayRef<int64_t> copyShape, Type elementType = {}) {
   auto type = dyn_cast<MemRefType>(value.getType());
+  if (type && elementType && isPackedByteBuffer(type))
+    return copyShape.size();
   if (!type || !type.hasStaticShape() || !hasByteSizedElementType(type.getElementType())
+      || (elementType && type.getElementType() != elementType)
       || type.getRank() != static_cast<int64_t>(copyShape.size()))
     return failure();
   if (llvm::any_of(copyShape, [](int64_t dim) { return dim <= 0; }))
@@ -351,6 +377,30 @@ static FailureOr<int64_t> getContiguousSuffixRank(Value value, ArrayRef<int64_t>
   return contiguousSuffixRank;
 }
 
+static FailureOr<SmallVector<int64_t>> getOuterByteStrides(
+    Value value, const LogicalCopyShape &copyShape, size_t outerRank) {
+  auto type = cast<MemRefType>(value.getType());
+  SmallVector<int64_t> strides;
+  if (isPackedByteBuffer(type))
+    strides = computeRowMajorStrides(copyShape.dimensions);
+  else {
+    auto proven = getProvenMemRefStrides(value);
+    if (failed(proven))
+      return failure();
+    strides = std::move(*proven);
+  }
+  int64_t elementByteWidth = static_cast<int64_t>(
+    getElementTypeSizeInBytes(copyShape.elementType));
+  SmallVector<int64_t> result;
+  for (int64_t stride : ArrayRef<int64_t>(strides).take_front(outerRank)) {
+    auto byteStride = checkedPositiveMul(stride, elementByteWidth);
+    if (failed(byteStride))
+      return failure();
+    result.push_back(*byteStride);
+  }
+  return result;
+}
+
 static FailureOr<CopyEndpointPlan> analyzeCopyEndpoint(Value value, Value initialByteOffset, MemRefType logicalType) {
   if (!logicalType.hasStaticShape() || !hasByteSizedElementType(logicalType.getElementType()))
     return failure();
@@ -430,8 +480,7 @@ analyzeCopyRewrite(Value target, Value source, Value targetOffset, Value sourceO
 
   auto targetBytes = getShapedByteSize(targetType);
   auto sourceBytes = getShapedByteSize(sourceType);
-  if (targetType.getElementType() == sourceType.getElementType() && succeeded(targetBytes) && succeeded(sourceBytes)
-      && size <= *targetBytes && size <= *sourceBytes) {
+  if (succeeded(targetBytes) && succeeded(sourceBytes) && size <= *targetBytes && size <= *sourceBytes) {
     auto targetSuffixRank = getContiguousSuffixRank(target, targetType.getShape());
     auto sourceSuffixRank = getContiguousSuffixRank(source, sourceType.getShape());
     if (succeeded(targetSuffixRank) && succeeded(sourceSuffixRank)
@@ -449,8 +498,10 @@ analyzeCopyRewrite(Value target, Value source, Value targetOffset, Value sourceO
   if (failed(logicalCopyShape))
     return failure();
 
-  auto targetSuffixRank = getContiguousSuffixRank(target, *logicalCopyShape);
-  auto sourceSuffixRank = getContiguousSuffixRank(source, *logicalCopyShape);
+  auto targetSuffixRank = getContiguousSuffixRank(
+    target, logicalCopyShape->dimensions, logicalCopyShape->elementType);
+  auto sourceSuffixRank = getContiguousSuffixRank(
+    source, logicalCopyShape->dimensions, logicalCopyShape->elementType);
   if (failed(targetSuffixRank) || failed(sourceSuffixRank))
     return failure();
 
@@ -459,23 +510,24 @@ analyzeCopyRewrite(Value target, Value source, Value targetOffset, Value sourceO
   plan.source = *sourcePlan;
 
   int64_t contiguousSuffixRank = std::min(*targetSuffixRank, *sourceSuffixRank);
-  if (contiguousSuffixRank == static_cast<int64_t>(logicalCopyShape->size())) {
+  if (contiguousSuffixRank
+      == static_cast<int64_t>(logicalCopyShape->dimensions.size())) {
     plan.kind = CopyRewritePlan::Kind::Direct;
     plan.directBytes = size;
     return plan;
   }
 
-  auto targetStrides = getProvenMemRefStrides(target);
-  auto sourceStrides = getProvenMemRefStrides(source);
-  if (failed(targetStrides) || failed(sourceStrides))
-    return failure();
-
-  int64_t elementByteWidth = static_cast<int64_t>(getElementTypeSizeInBytes(targetType.getElementType()));
+  int64_t elementByteWidth = static_cast<int64_t>(
+    getElementTypeSizeInBytes(logicalCopyShape->elementType));
   plan.kind = CopyRewritePlan::Kind::Loop;
   plan.loop.targetBaseOffset = plan.target.offset;
   plan.loop.sourceBaseOffset = plan.source.offset;
-  plan.loop.outerShape.assign(logicalCopyShape->begin(), logicalCopyShape->end() - contiguousSuffixRank);
-  SmallVector<int64_t> chunkShape(logicalCopyShape->end() - contiguousSuffixRank, logicalCopyShape->end());
+  plan.loop.outerShape.assign(
+    logicalCopyShape->dimensions.begin(),
+    logicalCopyShape->dimensions.end() - contiguousSuffixRank);
+  SmallVector<int64_t> chunkShape(
+    logicalCopyShape->dimensions.end() - contiguousSuffixRank,
+    logicalCopyShape->dimensions.end());
   auto outerElements = checkedPositiveProduct(plan.loop.outerShape);
   auto chunkElements = checkedPositiveProduct(chunkShape);
   auto chunkBytes = failed(chunkElements)
@@ -485,18 +537,14 @@ analyzeCopyRewrite(Value target, Value source, Value targetOffset, Value sourceO
     return failure();
   plan.loop.outerElements = *outerElements;
   plan.loop.chunkBytes = *chunkBytes;
-  for (int64_t stride : ArrayRef<int64_t>(*targetStrides).take_front(plan.loop.outerShape.size())) {
-    auto byteStride = checkedPositiveMul(stride, elementByteWidth);
-    if (failed(byteStride))
-      return failure();
-    plan.loop.targetOuterByteStrides.push_back(*byteStride);
-  }
-  for (int64_t stride : ArrayRef<int64_t>(*sourceStrides).take_front(plan.loop.outerShape.size())) {
-    auto byteStride = checkedPositiveMul(stride, elementByteWidth);
-    if (failed(byteStride))
-      return failure();
-    plan.loop.sourceOuterByteStrides.push_back(*byteStride);
-  }
+  auto targetStrides = getOuterByteStrides(
+    target, *logicalCopyShape, plan.loop.outerShape.size());
+  auto sourceStrides = getOuterByteStrides(
+    source, *logicalCopyShape, plan.loop.outerShape.size());
+  if (failed(targetStrides) || failed(sourceStrides))
+    return failure();
+  plan.loop.targetOuterByteStrides = std::move(*targetStrides);
+  plan.loop.sourceOuterByteStrides = std::move(*sourceStrides);
   if (plan.loop.chunkBytes <= 0)
     return failure();
   return plan;

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

@@ -7,7 +7,7 @@
 #include "OpBufferizationInterfaces.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
-#include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/BufferizationUtils.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/Passes/Transforms/Bufferization/BufferizationUtils.hpp"
 
 using namespace mlir;
 using namespace bufferization;

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

@@ -20,11 +20,11 @@
 #include "Common/Support/Diagnostics.hpp"
 #include "Compiler/PimCodeGen.hpp"
 #include "Dialect/Pim/PimOps.hpp"
-#include "Dialect/Pim/Transforms/Bufferization/Common.hpp"
-#include "Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.hpp"
+#include "Dialect/Pim/Passes/Transforms/Bufferization/Common.hpp"
+#include "Dialect/Pim/Passes/Transforms/Bufferization/ContiguityPatterns.hpp"
 #include "src/Accelerators/PIM/Common/IR/CoreBlockUtils.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-#include "src/Accelerators/PIM/Pass/PIMPasses.h"
+#include "src/Accelerators/PIM/Passes/PIMPasses.h"
 #include "src/Compiler/CompilerOptions.hpp"
 
 using namespace mlir;
@@ -382,8 +382,51 @@ static bufferization::OneShotBufferizationOptions makePimBufferizationOptions()
   return options;
 }
 
-static LogicalResult preparePimBufferization(func::FuncOp funcOp) {
+static LogicalResult verifyPimCoresNeedNoTensorCopies(
+    ModuleOp moduleOp, const bufferization::OneShotBufferizationOptions& baseOptions) {
+  static constexpr StringLiteral kExistingAlloc = "raptor.existing_core_alloc";
+  OwningOpRef<ModuleOp> clone = moduleOp.clone();
+  clone->walk([&](bufferization::AllocTensorOp alloc) {
+    if (alloc->getParentOfType<pim::PimCoreOp>()
+        || alloc->getParentOfType<pim::PimCoreBatchOp>())
+      alloc->setAttr(kExistingAlloc, UnitAttr::get(moduleOp.getContext()));
+  });
+
+  auto options = baseOptions;
+  options.bufferizeFunctionBoundaries = false;
+  options.opFilter.allowOperation([](Operation* op) {
+    return isa<pim::PimCoreOp, pim::PimCoreBatchOp>(op)
+        || op->getParentOfType<pim::PimCoreOp>()
+        || op->getParentOfType<pim::PimCoreBatchOp>();
+  });
+
+  bufferization::BufferizationState state;
+  if (failed(bufferization::insertTensorCopies(*clone, options, state))) {
+    moduleOp.emitError("official one-shot analysis failed while verifying Pim core copy freedom");
+    return failure();
+  }
+
+  CappedDiagnosticReporter diagnostics;
+  clone->walk([&](bufferization::AllocTensorOp alloc) {
+    if (alloc->hasAttr(kExistingAlloc)
+        || (!alloc->getParentOfType<pim::PimCoreOp>()
+            && !alloc->getParentOfType<pim::PimCoreBatchOp>()))
+      return;
+    Operation* requiredBy = alloc->getUsers().empty()
+      ? alloc.getOperation() : *alloc->getUsers().begin();
+    diagnostics.report(requiredBy, [](Operation* op) {
+      op->emitOpError("official one-shot bufferization requires a tensor copy inside a Pim core");
+    });
+  });
+  diagnostics.emitSuppressedSummary(moduleOp, "required Pim core tensor copies");
+  return success(!diagnostics.hasFailure());
+}
+
+static LogicalResult preparePimBufferization(
+    ModuleOp moduleOp, func::FuncOp funcOp, bool verifyCopyFreedom) {
   materializeWritableConstantDestinations(funcOp);
+  if (verifyCopyFreedom)
+    return verifyPimCoresNeedNoTensorCopies(moduleOp, makePimBufferizationOptions());
   return success();
 }
 
@@ -397,7 +440,7 @@ static LogicalResult runOneShotPimBufferization(
   bufferization::BufferizationState state;
   if (failed(bufferization::insertTensorCopies(moduleOp, hostOptions, state))
       || failed(bufferization::bufferizeModuleOp(moduleOp, options, state))) {
-    moduleOp.emitError("Failed to bufferize PIM and Spatial ops");
+    moduleOp.emitError("Failed to bufferize Pim and Spatial ops");
     return failure();
   }
   return success();
@@ -435,7 +478,7 @@ static LogicalResult verifyContiguousRuntimeOperands(ModuleOp moduleOp) {
           if (succeeded(resolveContiguousAddress(operand, knowledge)) || succeeded(compileContiguousAddressExpr(operand)))
             return;
           op.emitOpError() << "operand #" << operandIndex
-                           << " is not backed by contiguous addressable storage after PIM bufferization";
+                           << " is not backed by contiguous addressable storage after Pim bufferization";
           hasFailure = true;
         };
 
@@ -509,7 +552,7 @@ static LogicalResult verifyContiguousRuntimeOperands(ModuleOp moduleOp) {
   });
 
   if (hasFailure) {
-    moduleOp.emitError("PIM bufferization must fully normalize executable runtime operand contiguity before codegen");
+    moduleOp.emitError("Pim bufferization must fully normalize executable runtime operand contiguity before codegen");
     return failure();
   }
   return success();
@@ -546,7 +589,7 @@ static LogicalResult verifyPimCopyAddressSpaces(ModuleOp moduleOp) {
   });
   if (failureCount != 0)
     moduleOp.emitError() << "found " << failureCount
-                         << " PIM copy address-space violation(s); the first is reported above";
+                         << " Pim copy address-space violation(s); the first is reported above";
   return success(failureCount == 0);
 }
 
@@ -559,18 +602,30 @@ static LogicalResult normalizePimMemory(ModuleOp moduleOp, func::FuncOp funcOp)
   PatternRewriter rewriter(ctx);
 
   SmallVector<MemRefCopyWorkItem> copyWorklist;
+  SmallVector<PimMemCopyDevToHostOp> hostToHostCopies;
   llvm::SmallPtrSet<Operation*, 16> seenCopyOps;
+  llvm::SmallPtrSet<Operation*, 4> seenHostToHostCopies;
   auto addCopyOp = [&](memref::CopyOp copyOp, const StaticValueKnowledge& knowledge) {
     if (seenCopyOps.insert(copyOp.getOperation()).second)
       copyWorklist.push_back({copyOp, knowledge});
   };
+  auto collectCopy = [&](Operation &op,
+                         const StaticValueKnowledge &knowledge) {
+    if (auto copyOp = dyn_cast<memref::CopyOp>(&op))
+      addCopyOp(copyOp, knowledge);
+    if (auto copyOp = dyn_cast<PimMemCopyDevToHostOp>(&op);
+        copyOp
+        && isHostBackedPimAddress(copyOp.getDeviceSource(), knowledge)
+        && isHostBackedPimAddress(copyOp.getHostTarget(), knowledge)
+        && seenHostToHostCopies.insert(copyOp).second)
+      hostToHostCopies.push_back(copyOp);
+  };
 
   moduleOp.walk([&](pim::PimCoreOp coreOp) {
     StaticValueKnowledge knowledge = seedCoreKnowledge(coreOp);
     (void) walkPimCoreBlockStructurally(
       coreOp.getBody().front(), knowledge, [&](Operation& op, const StaticValueKnowledge& opKnowledge) {
-        if (auto copyOp = dyn_cast<memref::CopyOp>(&op))
-          addCopyOp(copyOp, opKnowledge);
+        collectCopy(op, opKnowledge);
         return success();
       });
   });
@@ -579,8 +634,7 @@ static LogicalResult normalizePimMemory(ModuleOp moduleOp, func::FuncOp funcOp)
       StaticValueKnowledge knowledge = seedCoreBatchKnowledge(coreBatchOp, lane);
       (void) walkPimCoreBlockStructurally(
         coreBatchOp.getBody().front(), knowledge, [&](Operation& op, const StaticValueKnowledge& opKnowledge) {
-          if (auto copyOp = dyn_cast<memref::CopyOp>(&op))
-            addCopyOp(copyOp, opKnowledge);
+          collectCopy(op, opKnowledge);
           return success();
         });
     }
@@ -588,6 +642,22 @@ static LogicalResult normalizePimMemory(ModuleOp moduleOp, func::FuncOp funcOp)
 
   bool hasFailed = false;
   Value zeroOffset = getOrCreateIndexConstant(rewriter, funcOp, 0);
+  for (PimMemCopyDevToHostOp copyOp : hostToHostCopies) {
+    rewriter.setInsertionPoint(copyOp);
+    auto scratchType = MemRefType::get(
+      {copyOp.getSize()}, rewriter.getI8Type());
+    Value scratch = memref::AllocOp::create(
+      rewriter, copyOp.getLoc(), scratchType);
+    auto load = PimMemCopyHostToDevOp::create(
+      rewriter, copyOp.getLoc(), scratchType, zeroOffset,
+      copyOp.getDeviceSourceOffset(), scratch, copyOp.getDeviceSource(),
+      copyOp.getSizeAttr());
+    auto store = PimMemCopyDevToHostOp::create(
+      rewriter, copyOp.getLoc(), copyOp.getHostTarget().getType(),
+      copyOp.getHostTargetOffset(), zeroOffset, copyOp.getHostTarget(),
+      load.getOutput(), copyOp.getSizeAttr());
+    rewriter.replaceOp(copyOp, store.getOutput());
+  }
   for (const MemRefCopyWorkItem& workItem : copyWorklist) {
     memref::CopyOp copyOp = workItem.copyOp;
     rewriter.setInsertionPoint(copyOp);
@@ -603,7 +673,7 @@ static LogicalResult normalizePimMemory(ModuleOp moduleOp, func::FuncOp funcOp)
   GreedyRewriteConfig contiguityConfig;
   contiguityConfig.enableFolding(false);
   if (failed(applyPatternsGreedily(moduleOp, std::move(contiguityPatterns), contiguityConfig))) {
-    moduleOp.emitError("failed to normalize PIM copy contiguity during bufferization");
+    moduleOp.emitError("failed to normalize Pim copy contiguity during bufferization");
     return failure();
   }
   annotateWeightsMemrefs(moduleOp, funcOp);
@@ -614,7 +684,7 @@ static LogicalResult normalizePimMemory(ModuleOp moduleOp, func::FuncOp funcOp)
 static FailureOr<func::FuncOp> requirePimEntryFunc(ModuleOp moduleOp, StringRef phase) {
   auto entryFunc = getPimEntryFunc(moduleOp);
   if (failed(entryFunc)) {
-    moduleOp.emitError("failed to locate the PIM entry function during ") << phase;
+    moduleOp.emitError("failed to locate the Pim entry function during ") << phase;
     return failure();
   }
   return *entryFunc;
@@ -626,21 +696,27 @@ struct PimBufferizationPreparationPass
     : PassWrapper<PimBufferizationPreparationPass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimBufferizationPreparationPass)
 
+  explicit PimBufferizationPreparationPass(bool verifyCopyFreedom = false)
+      : verifyCopyFreedom(verifyCopyFreedom) {}
+
   StringRef getArgument() const override { return "pim-bufferization-preparation"; }
   StringRef getDescription() const override {
-    return "Prepare writable tensor destinations for PIM one-shot bufferization.";
+    return "Prepare writable tensor destinations for Pim one-shot bufferization.";
   }
 
   void runOnOperation() final {
     ModuleOp moduleOp = getOperation();
-    auto funcOp = requirePimEntryFunc(moduleOp, "PIM bufferization preparation");
+    auto funcOp = requirePimEntryFunc(moduleOp, "Pim bufferization preparation");
     if (failed(funcOp)) {
       signalPassFailure();
       return;
     }
-    if (failed(preparePimBufferization(*funcOp)))
+    if (failed(preparePimBufferization(moduleOp, *funcOp, verifyCopyFreedom)))
       signalPassFailure();
   }
+
+private:
+  bool verifyCopyFreedom;
 };
 
 struct PimOneShotBufferizationPass
@@ -649,7 +725,7 @@ struct PimOneShotBufferizationPass
 
   StringRef getArgument() const override { return "pim-one-shot-bufferization"; }
   StringRef getDescription() const override {
-    return "Run one-shot bufferization for PIM and Spatial tensors.";
+    return "Run one-shot bufferization for Pim and Spatial tensors.";
   }
 
   void runOnOperation() final {
@@ -664,12 +740,12 @@ struct PimMemoryNormalizationPass
 
   StringRef getArgument() const override { return "pim-memory-normalization"; }
   StringRef getDescription() const override {
-    return "Normalize PIM memory copies and verify addressable operands.";
+    return "Normalize Pim memory copies and verify addressable operands.";
   }
 
   void runOnOperation() final {
     ModuleOp moduleOp = getOperation();
-    auto funcOp = requirePimEntryFunc(moduleOp, "PIM memory normalization");
+    auto funcOp = requirePimEntryFunc(moduleOp, "Pim memory normalization");
     if (failed(funcOp)) {
       signalPassFailure();
       return;
@@ -685,20 +761,20 @@ static LogicalResult verifyNoTensorValues(ModuleOp moduleOp) {
     if (failureCount >= 8)
       return;
     if (op->getDialect()->getNamespace() == "tensor") {
-      op->emitOpError("tensor operation remains after PIM bufferization");
+      op->emitOpError("tensor operation remains after Pim bufferization");
       ++failureCount;
       return;
     }
     for (Value value : op->getOperands()) {
       if (isa<TensorType>(value.getType())) {
-        op->emitOpError("tensor operand remains after PIM bufferization");
+        op->emitOpError("tensor operand remains after Pim bufferization");
         ++failureCount;
         return;
       }
     }
     for (Value value : op->getResults()) {
       if (isa<TensorType>(value.getType())) {
-        op->emitOpError("tensor result remains after PIM bufferization");
+        op->emitOpError("tensor result remains after Pim bufferization");
         ++failureCount;
         return;
       }
@@ -706,7 +782,7 @@ static LogicalResult verifyNoTensorValues(ModuleOp moduleOp) {
   });
   if (failureCount != 0)
     moduleOp.emitError() << "found " << failureCount
-                         << " tensor value(s) after PIM bufferization"
+                         << " tensor value(s) after Pim bufferization"
                          << (failureCount == 8 ? " (first 8 reported)" : "");
   return success(failureCount == 0);
 }
@@ -717,7 +793,7 @@ struct PimBufferizationVerificationPass
 
   StringRef getArgument() const override { return "pim-bufferization-verification"; }
   StringRef getDescription() const override {
-    return "Verify tensor elimination, contiguity, and PIM copy address spaces.";
+    return "Verify tensor elimination, contiguity, and Pim copy address spaces.";
   }
 
   void runOnOperation() final {
@@ -735,6 +811,10 @@ std::unique_ptr<Pass> createPimBufferizationPreparationPass() {
   return std::make_unique<PimBufferizationPreparationPass>();
 }
 
+std::unique_ptr<Pass> createPimBufferizationPreparationPass(bool verifyCopyFreedom) {
+  return std::make_unique<PimBufferizationPreparationPass>(verifyCopyFreedom);
+}
+
 std::unique_ptr<Pass> createPimOneShotBufferizationPass() {
   return std::make_unique<PimOneShotBufferizationPass>();
 }

src/PIM/Dialect/Pim/Passes/Transforms/HostConstantFolding/HostConstantFoldingPass.cpp

@@ -5,7 +5,7 @@
 
 #include "Patterns.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/Passes/Transforms/Bufferization/ContiguityPatterns.hpp"
 
 using namespace mlir;
 
@@ -16,7 +16,7 @@ struct HostConstantFoldingPass : PassWrapper<HostConstantFoldingPass, OperationP
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(HostConstantFoldingPass)
 
   StringRef getArgument() const override { return "pim-host-constant-folding-pass"; }
-  StringRef getDescription() const override { return "Fold host-side constant expressions before PIM verification"; }
+  StringRef getDescription() const override { return "Fold host-side constant expressions before Pim verification"; }
 
   LogicalResult initialize(MLIRContext* context) override {
     RewritePatternSet owningPatterns(context);
@@ -38,7 +38,7 @@ struct HostConstantFoldingPass : PassWrapper<HostConstantFoldingPass, OperationP
     GreedyRewriteConfig config;
     config.enableFolding();
     if (failed(applyPatternsGreedily(moduleOp, *patterns, config))) {
-      moduleOp.emitError("PIM host constant folding failed in the greedy rewrite driver");
+      moduleOp.emitError("Pim host constant folding failed in the greedy rewrite driver");
       signalPassFailure();
       return;
     }

src/PIM/Dialect/Pim/Passes/Transforms/HostConstantFolding/Patterns/Constant.cpp

@@ -472,7 +472,7 @@ struct FoldConstantHostCopyPattern final : OpRewritePattern<memref::CopyOp> {
   }
 };
 
-// Converts PIM copies from dense globals into direct folded globals before codegen.
+// Converts Pim copies from dense globals into direct folded globals before codegen.
 struct FoldConstantMemCpPattern final : OpRewritePattern<pim::PimMemCopyOp> {
   using OpRewritePattern::OpRewritePattern;