test pimcomp adversarial memory scheduling

normalize names and artifact paths

.agents/invariants/PIMSIM_NN_ORACLE_INVARIANT.md

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

.agents/invariants/PIPELINE_SCHEDULING_INVARIANT.md

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

AGENTS.md

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

README.md

@@ -5,7 +5,7 @@ targeting in-memory computing / processing-in-memory (PIM) architectures. It
 extends ONNX-MLIR with a PIM accelerator and progressively lowers ONNX-MLIR
 through custom MLIR dialects to simulator artifacts.
 
-The current target is the PIM simulator stack under `backend-simulators/pim`.
+The current target is the Pim simulator stack under `backend-simulators/pim`.
 Raptor emits binary per-core `.pim` instruction files by default, plus
 `memory.bin`, `config.json`, and weight binaries. It can also emit per-core JSON
 instruction files with `--pim-emit-json`.
@@ -29,9 +29,9 @@ lowering, scheduling, memory layout, and code-generation optimizations.
 - `backend-simulators/pim/pim-simulator` is the in-tree Rust functional
   simulator used by validation. It reads Raptor's `pim/` artifact directory and
   compares simulator output against native ONNX-MLIR execution.
-- `backend-simulators/pim/pimsim-nn` is the non-functional simulator submodule
-  used internally by validation for latency, power, and energy.
-  The helper scripts in `pimcomp_utils/` are for comparison with PIMCOMP-NN and
+- `backend-simulators/pim/pimsim-nn` contains the non-functional Pimsim
+  simulator used internally by validation for latency, power, and energy.
+  The helper scripts in `pimcomp_utils/` are for comparison with Pimcomp and
   contain local paths; treat them as local utilities, not portable workflows.
 
 ## Compilation pipeline
@@ -43,7 +43,7 @@ them to ONNX-MLIR through generated shim directories under
 High-level lowering flow:
 
 ```
-ONNX-MLIR -> Spatial -> Pim (tensor) -> Pim (bufferized) -> PIM artifacts
+ONNX-MLIR -> Spatial -> Pim (tensor) -> Pim (bufferized) -> Pim artifacts
 ```
 
 1. **ONNX -> Spatial** (`src/PIM/Conversion/ONNXToSpatial`).
@@ -81,20 +81,20 @@ 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**
+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. Reusable lifetime
+   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
+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
@@ -102,40 +102,97 @@ Supporting pieces:
 - `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`.
+  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.
+  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:
 
@@ -149,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:
 
@@ -159,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.
@@ -276,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.
@@ -284,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,7 @@ 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)]

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/pimcore.rs

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

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

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

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

@@ -33,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,68 @@ fn multiple_send_recv_test() {
         "send_recv failed to store"
     );
 }
+
+#[test]
+fn sync_wait_tokens_test() {
+    let cpu = common::empty_cpu(2);
+    let mut cores = CoreInstructionsBuilder::new(2);
+    let mut instructions = InstructionsBuilder::new();
+    let mut data = InstructionDataBuilder::new();
+
+    data.set_core_indx(1).fix_core_indx();
+    for _ in 0..2 {
+        instructions.make_inst(
+            sync,
+            data.set_imm_core(2).set_offset_select_value(0, 0).build(),
+        );
+    }
+    cores.set_core(1, instructions.build());
+
+    data.set_core_indx(2).fix_core_indx();
+    for _ in 0..2 {
+        instructions.make_inst(wait, data.set_offset_select_value(0, 1).build());
+    }
+    cores.set_core(2, instructions.build());
+
+    Executable::new(cpu, cores.build()).execute().unwrap();
+}
+
+#[test]
+fn blocked_transfers_do_not_starve_sync_producer() {
+    let cpu = common::empty_cpu(4);
+    let mut cores = CoreInstructionsBuilder::new(4);
+
+    let mut instructions = InstructionsBuilder::new();
+    let mut data = InstructionDataBuilder::new();
+    data.set_core_indx(1).fix_core_indx();
+    instructions.make_inst(sldi, data.set_rdimm(1, 0).build());
+    instructions.make_inst(recv, data.set_rd(1).set_imm_core(2).set_imm_len(1).build());
+    instructions.make_inst(send, data.set_r1(1).set_imm_core(3).set_imm_len(1).build());
+    cores.set_core(1, instructions.build());
+
+    let mut instructions = InstructionsBuilder::new();
+    let mut data = InstructionDataBuilder::new();
+    data.set_core_indx(2).fix_core_indx();
+    instructions.make_inst(sldi, data.set_rdimm(1, 0).build());
+    instructions.make_inst(wait, data.set_offset_select_value(0, 1).build());
+    instructions.make_inst(send, data.set_r1(1).set_imm_core(1).set_imm_len(1).build());
+    cores.set_core(2, instructions.build());
+
+    let mut instructions = InstructionsBuilder::new();
+    let mut data = InstructionDataBuilder::new();
+    data.set_core_indx(3).fix_core_indx();
+    instructions.make_inst(sldi, data.set_rdimm(1, 0).build());
+    instructions.make_inst(recv, data.set_rd(1).set_imm_core(1).set_imm_len(1).build());
+    cores.set_core(3, instructions.build());
+
+    let mut instructions = InstructionsBuilder::new();
+    let mut data = InstructionDataBuilder::new();
+    data.set_core_indx(4).fix_core_indx();
+    instructions.make_inst(
+        sync,
+        data.set_imm_core(2).set_offset_select_value(0, 0).build(),
+    );
+    cores.set_core(4, instructions.build());
+
+    Executable::new(cpu, cores.build()).execute().unwrap();
+}

src/PIM/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,20 +55,20 @@ 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));
 
@@ -94,39 +96,74 @@ llvm::cl::opt<bool> pimEmitJson("pim-emit-json",
 
 llvm::cl::opt<bool> pimDetectCommunicationDeadlock(
   "pim-detect-communication-deadlock",
-  llvm::cl::desc("Expensively simulate the statically expanded PIM send/receive order at verification time and fail if a blocking communication deadlock is found"),
+  llvm::cl::desc("Expensively simulate the statically expanded Pim send/receive order at verification time and fail if a blocking communication deadlock is found"),
   llvm::cl::init(false),
   llvm::cl::cat(OnnxMlirOptions));
 
 llvm::cl::opt<bool> pimVerifyBufferizationCopyFreedom(
   "pim-verify-bufferization-copy-freedom",
-  llvm::cl::desc("Run the expensive official PIM tensor-copy freedom proof before bufferization"),
+  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> 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));
 
-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

@@ -59,15 +59,17 @@ extern llvm::cl::opt<bool> pimEmitJson;
 extern llvm::cl::opt<bool> pimReportConvLowering;
 extern llvm::cl::opt<bool> pimDetectCommunicationDeadlock;
 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();
+void verifyPimCompilerOptions();
 
 } // namespace onnx_mlir

src/PIM/Compiler/PimCompilerUtils.cpp

@@ -12,6 +12,7 @@
 #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"
@@ -78,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;
@@ -95,7 +97,7 @@ 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::SpatialDataflowExportStage getPimSpatialDataflowExportStage(
@@ -108,7 +110,7 @@ spatial::SpatialDataflowExportStage getPimSpatialDataflowExportStage(
   case SpatialDataflowExportSpatial4: return spatial::SpatialDataflowExportStage::Spatial4;
   case SpatialDataflowExportAll:      return spatial::SpatialDataflowExportStage::All;
   }
-  llvm_unreachable("unknown PIM Spatial dataflow export stage");
+  llvm_unreachable("unknown Pim Spatial dataflow export stage");
 }
 
 spatial::SpatialTargetResources getPimSpatialTargetResources(const spatial::SchedulingTarget& target) {
@@ -118,7 +120,7 @@ spatial::SpatialTargetResources getPimSpatialTargetResources(const spatial::Sche
   resources.processorCount = target.processorCount;
   resources.vectorWidth = target.vectorWidth;
   if (failed(resources.verify()))
-    llvm::report_fatal_error("PIM target resources are incomplete");
+    llvm::report_fatal_error("Pim target resources are incomplete");
   return resources;
 }
 
@@ -136,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;
 }
 
@@ -149,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));
 }
 
@@ -157,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)};
 }
 
@@ -172,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)) {
@@ -185,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);
@@ -208,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) {
@@ -218,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[
@@ -242,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");
@@ -265,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);
@@ -276,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);
 
@@ -329,7 +331,7 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
                   PassManager& pm,
                   EmissionTargetType& emissionTarget,
                   std::string outputNameNoExt) {
-  verifyExplicitPimCoreCount();
+  verifyPimCompilerOptions();
   spatial::SchedulingTarget schedulingTarget = getPimSchedulingTarget();
   spatial::SpatialTargetResources targetResources = getPimSpatialTargetResources(schedulingTarget);
 
@@ -349,12 +351,13 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
     spatial::SpatialDataflowExportStage exportStage =
       getPimSpatialDataflowExportStage(pimExportSpatialDataflow.getValue());
     pm.addPass(createONNXToSpatialPass(targetResources, planningOptions));
-    pm.addPass(createSpatialLayoutPlanningPass(targetResources));
+    pm.addPass(createSpatialLayoutPlanningPass(
+      targetResources, pimDisableSpatialPlanning.getValue()));
     pm.addPass(createLowerSpatialPlansPass(targetResources, planningOptions, exportStage));
     pm.addPass(createTrivialGraphComputeMergePass(
       schedulingTarget.residentWeightCapacity, exportStage));
     pm.addPass(spatial::createScheduleAndRealizeSpatialPass(
-      schedulingTarget, exportStage));
+      schedulingTarget, exportStage, pipelineStages.getValue()));
     pm.addPass(createMessagePass("Onnx lowered to Spatial"));
   }
 

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/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/ShapeTilingUtils.hpp

@@ -25,7 +25,7 @@ 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,

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

@@ -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;
@@ -163,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";
@@ -182,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/Transforms/LowerSpatialPlansPass.cpp

@@ -46,7 +46,7 @@ struct LowerSpatialPlansPass final
     }
     auto entryFunc = getPimEntryFunc(moduleOp);
     if (failed(entryFunc)) {
-      moduleOp.emitError("failed to locate the PIM entry function during LowerSpatialPlans");
+      moduleOp.emitError("failed to locate the Pim entry function during LowerSpatialPlans");
       signalPassFailure();
       return;
     }

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

@@ -158,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;
   }
@@ -245,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;
   }

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

@@ -42,8 +42,9 @@ static SmallVector<spatial::PhysicalLayout> getOperandLayouts(
 class SpatialLayoutAnalysis {
 public:
   SpatialLayoutAnalysis(func::FuncOp funcOp,
-                        const spatial::SpatialTargetResources& target)
-      : funcOp(funcOp), target(target) {}
+                        const spatial::SpatialTargetResources& target,
+                        bool selectTrivialPlan)
+      : funcOp(funcOp), target(target), selectTrivialPlan(selectTrivialPlan) {}
 
   FailureOr<SpatialLayoutSelection> run() {
     SpatialLayoutSelection selection;
@@ -56,6 +57,9 @@ public:
       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;
@@ -168,6 +172,7 @@ private:
 
   func::FuncOp funcOp;
   const spatial::SpatialTargetResources& target;
+  bool selectTrivialPlan;
 };
 
 static LogicalResult materializeMismatchedUses(
@@ -251,8 +256,9 @@ struct SpatialLayoutPlanningPass final
   }
 
   SpatialLayoutPlanningPass() = default;
-  explicit SpatialLayoutPlanningPass(const spatial::SpatialTargetResources& target)
-  : target(target), hasTarget(true) {}
+  SpatialLayoutPlanningPass(const spatial::SpatialTargetResources& target,
+                            bool selectTrivialPlan)
+  : target(target), selectTrivialPlan(selectTrivialPlan), hasTarget(true) {}
 
   void runOnOperation() override {
     ModuleOp moduleOp = getOperation();
@@ -263,13 +269,13 @@ struct SpatialLayoutPlanningPass final
     }
     auto entryFunc = getPimEntryFunc(moduleOp);
     if (failed(entryFunc)) {
-      moduleOp.emitError("failed to locate the PIM entry function during Spatial layout planning");
+      moduleOp.emitError("failed to locate the Pim entry function during Spatial layout planning");
       signalPassFailure();
       return;
     }
 
     func::FuncOp funcOp = *entryFunc;
-    SpatialLayoutAnalysis analysis(funcOp, target);
+    SpatialLayoutAnalysis analysis(funcOp, target, selectTrivialPlan);
     FailureOr<SpatialLayoutSelection> selection = analysis.run();
     if (failed(selection)) {
       signalPassFailure();
@@ -301,6 +307,7 @@ struct SpatialLayoutPlanningPass final
   }
 
   spatial::SpatialTargetResources target;
+  bool selectTrivialPlan = false;
   bool hasTarget = false;
 };
 
@@ -311,8 +318,8 @@ std::unique_ptr<Pass> createSpatialLayoutPlanningPass() {
 }
 
 std::unique_ptr<Pass> createSpatialLayoutPlanningPass(
-    const spatial::SpatialTargetResources& target) {
-  return std::make_unique<SpatialLayoutPlanningPass>(target);
+    const spatial::SpatialTargetResources& target, bool selectTrivialPlan) {
+  return std::make_unique<SpatialLayoutPlanningPass>(target, selectTrivialPlan);
 }
 
 } // namespace onnx_mlir

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

@@ -199,7 +199,7 @@ static bool writeConvLoweringReport(const ConvLoweringReportEntry& entry,
     return false;
   }
 
-  reportFile << "# PIM Conv Lowering Report (bounded to 512 rows)\n\n";
+  reportFile << "# Pim conv lowering report (bounded to 512 rows)\n\n";
   reportFile << "## Plan selection\n";
   writeConvReportTableHeader(reportFile, "Selector");
   bool realizationSectionStarted = false;

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/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,42 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatScheduledCom
     auto blockArg = computeOp.getInputArgument(inputIndex);
     if (!blockArg)
       return computeOp.emitOpError("expected compute input block arguments during lowering");
-    auto receiveOp = dyn_cast_or_null<spatial::SpatChannelReceiveOp>(input.getDefiningOp());
+    auto channelReceive = dyn_cast_or_null<spatial::SpatChannelReceiveOp>(
+      input.getDefiningOp());
+    auto hostWaitLoad = dyn_cast_or_null<spatial::SpatHostWaitLoadOp>(
+      input.getDefiningOp());
+    Operation *receiveOp = channelReceive
+      ? channelReceive.getOperation() : hostWaitLoad.getOperation();
     if (receiveOp && !blockArg->use_empty()) {
       rewriter.setInsertionPoint(getEarliestUserWithinBlock(*blockArg));
       auto outputType = cast<ShapedType>(blockArg->getType());
-      auto outputBuffer = createEmptyTensorFromShaped(rewriter, receiveOp.getLoc(), outputType);
+      auto outputBuffer = createEmptyTensorFromShaped(
+        rewriter, receiveOp->getLoc(), outputType);
       auto sizeAttr = getTensorSizeInBytesAttr(rewriter, computeOp.getOperation(), *blockArg);
       if (failed(sizeAttr))
         return failure();
-      Value received =
-        PimReceiveOp::create(
-          rewriter, receiveOp.getLoc(), outputBuffer.getType(), outputBuffer,
-          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();
+        PimSyncOp::create(
+          rewriter, receiveOp->getLoc(), hostWaitLoad.getSourceCoreId(),
+          hostWaitLoad.getAcknowledgementEventRegister());
+      } else {
+        received = PimReceiveOp::create(
+          rewriter, receiveOp->getLoc(), outputBuffer.getType(), outputBuffer,
+          zero, *sizeAttr, channelReceive.getSourceCoreId()).getOutput();
+      }
       blockArg->replaceAllUsesWith(received);
       markOpToRemove(receiveOp);
       continue;
@@ -383,11 +405,12 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatScheduledCom
     if (rangeLength(resultUses) == 1) {
       OpOperand& resultUse = *resultUses.begin();
       Operation* resultUser = resultUse.getOwner();
-      if (isa<spatial::SpatChannelSendOp>(resultUser))
+      if (isa<spatial::SpatChannelSendOp,
+              spatial::SpatHostStoreSyncOp>(resultUser))
         continue;
     }
 
-    return computeOp.emitOpError("has an unsupported remaining result use during Spatial-to-PIM lowering");
+    return computeOp.emitOpError("has an unsupported remaining result use during Spatial-to-Pim lowering");
   }
 
   rewriter.setInsertionPoint(yieldOp);
@@ -397,7 +420,7 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatScheduledCom
   if (!computeOp.getWeights().empty())
     computeWeights.append(computeOp.getWeights().begin(), computeOp.getWeights().end());
   rewriter.setInsertionPointAfter(computeOp);
-  auto checkedCoreId = getRequiredScheduledCoreId(computeOp, "spatial compute core id");
+  auto checkedCoreId = getRequiredScheduledCoreId(computeOp, "Spatial compute core id");
   if (failed(checkedCoreId))
     return failure();
   auto coreIdAttr = pim::getCheckedI32Attr(rewriter, computeOp, static_cast<int64_t>(*checkedCoreId), "pim core id");

src/PIM/Conversion/SpatialToPim/Patterns/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,69 @@ 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();
+      auto sync = pim::PimSyncOp::create(
+        rewriter, op.getLoc(), op.getSourceCoreId(),
+        op.getAcknowledgementEventRegister());
+      copyRaptorDebugAttrs(op.getOperation(), sync.getOperation());
+      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 +225,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

@@ -88,7 +88,7 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
   operationsToRemove.clear();
   ModuleOp moduleOp = getOperation();
   if (!hasTarget || failed(targetResources.verify())) {
-    moduleOp.emitError("Spatial-to-PIM lowering requires valid injected target resources");
+    moduleOp.emitError("Spatial-to-Pim lowering requires valid injected target resources");
     signalPassFailure();
     return;
   }
@@ -96,7 +96,7 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
 
   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;
   }
@@ -126,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;
   }
@@ -140,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;
   }
@@ -182,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);
@@ -202,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;
@@ -251,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;
   }

src/PIM/Conversion/SpatialToPim/SpatialToPimPass.hpp

@@ -26,7 +26,7 @@ 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;
   explicit SpatialToPimPass(const spatial::SpatialTargetResources& target)

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/PimBufferizationPass.cpp

@@ -402,7 +402,7 @@ static LogicalResult verifyPimCoresNeedNoTensorCopies(
 
   bufferization::BufferizationState state;
   if (failed(bufferization::insertTensorCopies(*clone, options, state))) {
-    moduleOp.emitError("official one-shot analysis failed while verifying PIM core copy freedom");
+    moduleOp.emitError("official one-shot analysis failed while verifying Pim core copy freedom");
     return failure();
   }
 
@@ -415,10 +415,10 @@ static LogicalResult verifyPimCoresNeedNoTensorCopies(
     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");
+      op->emitOpError("official one-shot bufferization requires a tensor copy inside a Pim core");
     });
   });
-  diagnostics.emitSuppressedSummary(moduleOp, "required PIM core tensor copies");
+  diagnostics.emitSuppressedSummary(moduleOp, "required Pim core tensor copies");
   return success(!diagnostics.hasFailure());
 }
 
@@ -440,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();
@@ -478,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;
         };
 
@@ -552,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();
@@ -589,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);
 }
 
@@ -602,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();
       });
   });
@@ -622,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();
         });
     }
@@ -631,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);
@@ -646,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);
@@ -657,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;
@@ -674,12 +701,12 @@ struct PimBufferizationPreparationPass
 
   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;
@@ -698,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 {
@@ -713,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;
@@ -734,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;
       }
@@ -755,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);
 }
@@ -766,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 {

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

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

src/PIM/Dialect/Pim/Passes/Transforms/InstructionSelection/InstructionSelectionPass.cpp

@@ -40,7 +40,7 @@ struct LowerTransposePattern final : OpRewritePattern<pim::PimTransposeOp> {
     auto sourceType = dyn_cast<MemRefType>(op.getInput().getType());
     auto targetType = dyn_cast<MemRefType>(op.getOutputBuffer().getType());
     if (!sourceType || !targetType || !sourceType.hasStaticShape() || !targetType.hasStaticShape())
-      return op.emitOpError("requires static memref operands before PIM instruction selection");
+      return op.emitOpError("requires static memref operands before Pim instruction selection");
 
     ArrayRef<int64_t> sourceShape = sourceType.getShape();
     size_t rank = sourceShape.size();
@@ -147,7 +147,7 @@ struct InstructionSelectionPass : PassWrapper<InstructionSelectionPass, Operatio
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(InstructionSelectionPass)
 
   StringRef getArgument() const override { return "pim-instruction-selection"; }
-  StringRef getDescription() const override { return "Select explicit PIM ISA operations"; }
+  StringRef getDescription() const override { return "Select explicit Pim ISA operations"; }
 
   void runOnOperation() override {
     RewritePatternSet patterns(&getContext());

src/PIM/Dialect/Pim/Passes/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.cpp

@@ -36,7 +36,7 @@ struct PimLocalMemoryPlanningPass : PassWrapper<PimLocalMemoryPlanningPass, Oper
 
   StringRef getArgument() const override { return "pim-local-memory-planning"; }
   StringRef getDescription() const override {
-    return "Plan liveness-based addresses for PIM core-local memory";
+    return "Plan liveness-based addresses for Pim core-local memory";
   }
 
   void runOnOperation() override {
@@ -149,14 +149,14 @@ FailureOr<CoreMemoryPlan> buildCoreMemoryPlan(Operation* coreLikeOp) {
   plan.intervals = std::move(*intervals);
   auto placements = planLocalMemoryPlacements(plan.intervals, kPimLocalMemoryAddressLimit);
   if (failed(placements)) {
-    coreLikeOp->emitError("PIM local-memory plan exceeds the signed int32 address range");
+    coreLikeOp->emitError("Pim local-memory plan exceeds the signed int32 address range");
     return failure();
   }
   plan.placements = std::move(*placements);
   for (const LocalMemoryPlacement& placement : plan.placements) {
     auto end = alignedEnd(placement.address, placement.size, kPimLocalMemoryAddressLimit);
     if (failed(end)) {
-      coreLikeOp->emitError("PIM local-memory plan has invalid address arithmetic");
+      coreLikeOp->emitError("Pim local-memory plan has invalid address arithmetic");
       return failure();
     }
     plan.arenaSize = std::max(plan.arenaSize, *end);

src/PIM/Dialect/Pim/Passes/Transforms/Verification/VerificationPass.cpp

@@ -117,7 +117,7 @@ verifyLocalMemoryPlan(Operation* coreLikeOp, pim::CappedDiagnosticReporter& diag
   for (StringRef name : kRemovedLocalMemoryPlanAttrNames)
     if (coreLikeOp->hasAttr(name)) {
       diagnostics.report(coreLikeOp, [name](Operation* op) {
-        op->emitError() << "contains removed PIM local-memory planning attribute '" << name << "'";
+        op->emitError() << "contains removed Pim local-memory planning attribute '" << name << "'";
       });
       hasFailure = true;
     }
@@ -137,7 +137,7 @@ verifyLocalMemoryPlan(Operation* coreLikeOp, pim::CappedDiagnosticReporter& diag
   auto analyzed = pim::analyzeLocalMemoryLifetimes(coreLikeOp);
   if (failed(analyzed)) {
     diagnostics.report(coreLikeOp, [](Operation* op) {
-      op->emitError("cannot analyze PIM local-memory lifetimes for plan verification");
+      op->emitError("cannot analyze Pim local-memory lifetimes for plan verification");
     });
     return failure();
   }
@@ -156,7 +156,7 @@ verifyLocalMemoryPlan(Operation* coreLikeOp, pim::CappedDiagnosticReporter& diag
     for (StringRef name : kRemovedLocalMemoryPlanAttrNames)
       if (allocation->hasAttr(name)) {
         diagnostics.report(allocation, [name](Operation* op) {
-          op->emitOpError() << "contains removed PIM local-memory planning attribute '" << name << "'";
+          op->emitOpError() << "contains removed Pim local-memory planning attribute '" << name << "'";
         });
         hasFailure = true;
       }
@@ -171,7 +171,7 @@ verifyLocalMemoryPlan(Operation* coreLikeOp, pim::CappedDiagnosticReporter& diag
     uint64_t address = static_cast<uint64_t>(addressAttr.getInt());
     if (address % 4 != 0 || address > arenaSize || interval.size > arenaSize - address) {
       diagnostics.report(allocation, [&](Operation* op) {
-        op->emitOpError() << "has invalid PIM local-memory range [" << address << ", "
+          op->emitOpError() << "has invalid Pim local-memory range [" << address << ", "
                           << (address <= arenaSize && interval.size <= arenaSize - address
                                 ? address + interval.size
                                 : arenaSize)
@@ -221,7 +221,7 @@ verifyLocalMemoryPlan(Operation* coreLikeOp, pim::CappedDiagnosticReporter& diag
       memref::AllocOp otherAllocation = other.allocation;
       diagnostics.report(allocation, [&](Operation*) {
         auto diagnostic = allocation.emitOpError()
-          << "PIM local-memory plan assigns simultaneously live allocations to overlapping ranges; first range ["
+          << "Pim local-memory plan assigns simultaneously live allocations to overlapping ranges; first range ["
           << conflicting->first << ", " << conflicting->first + other.size << "), second range [" << address
           << ", " << address + interval.size << "), live positions overlap at ["
           << std::max(interval.start, other.start) << ", " << std::min(interval.end, other.end) << "]";
@@ -241,6 +241,8 @@ static bool isSupportedCoreInstructionOp(Operation* op) {
              pim::PimVMVOp,
              pim::PimReceiveOp,
              pim::PimSendOp,
+             pim::PimSyncOp,
+             pim::PimWaitOp,
              pim::PimConcatOp,
              pim::PimVMMOp,
              pim::PimVVAddOp,
@@ -469,7 +471,7 @@ static LogicalResult appendCoreCommunicationEvents(Block& block,
         auto targetCoreId = resolveIndexValue(sendOp.getTargetCoreId(), knowledge);
         if (failed(targetCoreId)) {
           diagnostics.report(&op, [](Operation* illegalOp) {
-            illegalOp->emitOpError("cannot statically resolve send target core for PIM communication deadlock check");
+            illegalOp->emitOpError("cannot statically resolve send target core for Pim communication deadlock check");
           });
           return failure();
         }
@@ -488,7 +490,7 @@ static LogicalResult appendCoreCommunicationEvents(Block& block,
         if (failed(sourceCoreId)) {
           diagnostics.report(&op, [](Operation* illegalOp) {
             illegalOp->emitOpError(
-              "cannot statically resolve receive source core for PIM communication deadlock check");
+              "cannot statically resolve receive source core for Pim communication deadlock check");
           });
           return failure();
         }
@@ -528,7 +530,7 @@ static void printCommunicationWindow(llvm::raw_ostream& os,
 static void printCommunicationDeadlockReport(const DenseMap<int64_t, CommunicationEventVector>& coreEvents,
                                              const DenseMap<int64_t, size_t>& programCounters,
                                              ArrayRef<int64_t> cycle) {
-  llvm::errs() << "\n=== PIM static communication deadlock report ===\n";
+  llvm::errs() << "\n=== Pim static communication deadlock report ===\n";
   llvm::errs() << "wait cycle:";
   for (int64_t coreId : cycle)
     llvm::errs() << " " << coreId;
@@ -563,7 +565,7 @@ static void printCommunicationDeadlockReport(const DenseMap<int64_t, Communicati
       continue;
     printCommunicationWindow(llvm::errs(), coreEvents, coreId, pcIt->second);
   }
-  llvm::errs() << "=== end PIM static communication deadlock report ===\n\n";
+  llvm::errs() << "=== end Pim static communication deadlock report ===\n\n";
 }
 
 static void emitCommunicationDeadlockCycle(ModuleOp moduleOp,
@@ -574,8 +576,8 @@ static void emitCommunicationDeadlockCycle(ModuleOp moduleOp,
 
   auto diagnostic =
     moduleOp.emitError()
-    << "PIM communication deadlock check found a blocking send/receive cycle while statically simulating the "
-       "expanded per-core communication streams; see the PIM static communication deadlock report above";
+    << "Pim communication deadlock check found a blocking send/receive cycle while statically simulating the "
+       "expanded per-core communication streams; see the Pim static communication deadlock report above";
 
   for (int64_t coreId : cycle) {
     auto eventsIt = coreEvents.find(coreId);
@@ -726,7 +728,7 @@ static LogicalResult verifyNoStaticCommunicationDeadlock(ModuleOp moduleOp,
 
     auto diagnostic =
       moduleOp.emitError()
-      << "PIM communication deadlock check stalled without finding a closed wait cycle; this usually means a "
+      << "Pim communication deadlock check stalled without finding a closed wait cycle; this usually means a "
          "send/receive peer is missing or ordered after a finished core";
     for (const auto& [coreId, events] : coreEvents) {
       size_t pc = programCounters[coreId];
@@ -744,7 +746,7 @@ struct VerificationPass : PassWrapper<VerificationPass, OperationPass<ModuleOp>>
 
   StringRef getArgument() const override { return "verify-pim-pass"; }
   StringRef getDescription() const override {
-    return "Verify that bufferized PIM IR contains only explicit host/device transfers";
+    return "Verify that bufferized Pim IR contains only explicit host/device transfers";
   }
 
   VerificationPass() {}
@@ -761,7 +763,7 @@ struct VerificationPass : PassWrapper<VerificationPass, OperationPass<ModuleOp>>
     pim::CappedDiagnosticReporter diagnostics;
 
     if (!hasTarget || failed(targetResources.verify())) {
-      moduleOp.emitError("PIM codegen verification requires valid injected target resources");
+      moduleOp.emitError("Pim codegen verification requires valid injected target resources");
       signalPassFailure();
       return;
     }
@@ -790,7 +792,7 @@ struct VerificationPass : PassWrapper<VerificationPass, OperationPass<ModuleOp>>
         return;
 
       diagnostics.report(op, [](Operation* illegalOp) {
-        illegalOp->emitError("illegal Spatial operation reached PIM codegen verification");
+        illegalOp->emitError("illegal Spatial operation reached Pim codegen verification");
       });
     });
 
@@ -831,7 +833,7 @@ struct VerificationPass : PassWrapper<VerificationPass, OperationPass<ModuleOp>>
 
         if (!isAddressOnlyHostOp(&op)) {
           diagnostics.report(&op, [](Operation* illegalOp) {
-            illegalOp->emitOpError("illegal host-side runtime op remains after PIM bufferization; "
+            illegalOp->emitOpError("illegal host-side runtime op remains after Pim bufferization; "
                                    "fold it to constants or lower it into pim.core");
           });
           continue;
@@ -847,7 +849,7 @@ struct VerificationPass : PassWrapper<VerificationPass, OperationPass<ModuleOp>>
 
     if (diagnostics.hasFailure()) {
       diagnostics.emitSuppressedSummary(moduleOp, "verification failures");
-      moduleOp.emitError("PIM codegen verification failed; see diagnostics above");
+      moduleOp.emitError("Pim codegen verification failed; see diagnostics above");
       hasFailure = true;
     }
 
@@ -926,7 +928,7 @@ private:
         bool hasFailure = false;
         if (!isSupportedCoreInstructionOp(&op)) {
           diagnostics.report(&op, [](Operation* illegalOp) {
-            illegalOp->emitOpError("unsupported executable op reached PIM codegen verification");
+            illegalOp->emitOpError("unsupported executable op reached Pim codegen verification");
           });
           hasFailure = true;
         }
@@ -988,7 +990,7 @@ private:
           if (failed(resolveIndexValue(storeOp.getHostTargetOffset(), knowledge))
               || failed(resolveIndexValue(storeOp.getDeviceSourceOffset(), knowledge))) {
             diagnostics.report(&op, [](Operation* illegalOp) {
-              illegalOp->emitOpError("offset operands must be statically evaluable for PIM codegen");
+            illegalOp->emitOpError("offset operands must be statically evaluable for Pim codegen");
             });
             hasFailure = true;
           }
@@ -1004,7 +1006,7 @@ private:
           if (failed(resolveIndexValue(loadOp.getDeviceTargetOffset(), knowledge))
               || failed(resolveIndexValue(loadOp.getHostSourceOffset(), knowledge))) {
             diagnostics.report(&op, [](Operation* illegalOp) {
-              illegalOp->emitOpError("offset operands must be statically evaluable for PIM codegen");
+            illegalOp->emitOpError("offset operands must be statically evaluable for Pim codegen");
             });
             hasFailure = true;
           }
@@ -1020,7 +1022,7 @@ private:
           if (failed(resolveIndexValue(copyOp.getTargetOffset(), knowledge))
               || failed(resolveIndexValue(copyOp.getSourceOffset(), knowledge))) {
             diagnostics.report(&op, [](Operation* illegalOp) {
-              illegalOp->emitOpError("offset operands must be statically evaluable for PIM codegen");
+            illegalOp->emitOpError("offset operands must be statically evaluable for Pim codegen");
             });
             hasFailure = true;
           }
@@ -1030,7 +1032,7 @@ private:
             && failed(resolveIndexValue(receiveOp.getOutputOffset(), knowledge))) {
           diagnostics.report(&op, [](Operation* illegalOp) {
             illegalOp->emitOpError(
-              "output offset must be statically evaluable for PIM codegen");
+              "output offset must be statically evaluable for Pim codegen");
           });
           hasFailure = true;
         }

src/PIM/Dialect/Pim/Pim.td

@@ -11,7 +11,7 @@ include "mlir/Dialect/Bufferization/IR/BufferViewFlowOpInterface.td"
 
 def PimDialect : Dialect {
   let name = "pim";
-  let summary = "A low-level dialect for the PIM coprocessors on ReRAM crossbars";
+  let summary = "A low-level dialect for the Pim coprocessors on ReRAM crossbars";
   let cppNamespace = "::onnx_mlir::pim";
 }
 
@@ -27,7 +27,7 @@ def PimTensor :
 
 def PimCoreOp : PimOp<"core", [SingleBlock,
     DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmBlockArgumentNames"]>]> {
-  let summary = "Execute a block on a PIM core";
+  let summary = "Execute a block on a Pim core";
 
   let regions = (region SizedRegion<1>:$body);
 
@@ -118,6 +118,32 @@ def PimReceiveOp : PimOp<"receive", [DestinationStyleOpInterface]> {
   }];
 }
 
+def PimSyncOp : PimOp<"sync", []> {
+  let summary = "Signal an event register on another core";
+
+  let arguments = (ins
+    Index:$targetCoreId,
+    Index:$eventRegister
+  );
+
+  let assemblyFormat = [{
+    $targetCoreId `event` $eventRegister attr-dict
+  }];
+}
+
+def PimWaitOp : PimOp<"wait", []> {
+  let summary = "Wait for an event register value";
+
+  let arguments = (ins
+    Index:$eventRegister,
+    Index:$waitValue
+  );
+
+  let assemblyFormat = [{
+    $eventRegister `value` $waitValue attr-dict
+  }];
+}
+
 def PimMemCopyHostToDevOp : PimOp<"memcp_hd", [DestinationStyleOpInterface]> {
   let summary = "Copy a memory region from host memory into device memory";
 

src/PIM/Dialect/Spatial/CMakeLists.txt

@@ -35,6 +35,7 @@ add_pim_library(SpatialOps
   Passes/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.cpp
   Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp
   Passes/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp
+  Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.cpp
   Passes/Transforms/TrivialGraphComputeMergePass.cpp
 
   EXCLUDE_FROM_OM_LIBS

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp

@@ -219,10 +219,12 @@ static void appendReceive(BoundaryProgram &boundary,
         run->entryOffsets[run->entryOffsets.size() - 2]].family->requirement;
       CollectionTarget previousTarget {run->collection, run->positions.back()};
       bool sameEntry = previous == requirement;
-      if (sameEntry
+      bool sameRoute = run->slices.back().family->hostRouted
+                    == slice.family->hostRouted;
+      if (sameRoute && (sameEntry
           || (sameCollectionEmissionContract(previousTarget, target)
               && previous->publicationFragmentType
-                   == requirement->publicationFragmentType)) {
+                   == requirement->publicationFragmentType))) {
         run->slices.push_back(slice);
         if (sameEntry) {
           run->entryOffsets.back() = run->slices.size();
@@ -240,11 +242,146 @@ static void appendReceive(BoundaryProgram &boundary,
     target.collection, {slice}, {0, 1}, {target.position}, {lanes}, lanes});
 }
 
+struct HostTransferRef {
+  ExternalTransferFamily *family = nullptr;
+  size_t index = 0;
+};
+
+static unsigned getBarrierRoundCount(size_t coreCount) {
+  unsigned rounds = 0;
+  for (size_t distance = 1; distance < coreCount; distance *= 2)
+    ++rounds;
+  return rounds;
+}
+
+static LogicalResult assignPipelineSynchronization(
+    DeferredTransferPlan &transfers,
+    ArrayRef<BoundaryProgram> boundaries,
+    size_t synchronizationRegisterCount) {
+  bool pipelined = false;
+  for (ScheduledInfo &scheduled : transfers.scheduled) {
+    if (scheduled.pipelineStages.empty())
+      continue;
+    pipelined = true;
+    llvm::append_range(transfers.downstreamCores, scheduled.cores);
+    for (auto [core, stage] :
+         llvm::zip_equal(scheduled.cores, scheduled.pipelineStages))
+      if (stage == 0)
+        transfers.stageZeroCores.push_back(core);
+  }
+  if (!pipelined)
+    return success();
+  transfers.synchronizationRegisterCount = synchronizationRegisterCount;
+  llvm::sort(transfers.stageZeroCores);
+  transfers.stageZeroCores.erase(
+    llvm::unique(transfers.stageZeroCores), transfers.stageZeroCores.end());
+  llvm::sort(transfers.downstreamCores);
+  transfers.downstreamCores.erase(
+    llvm::unique(transfers.downstreamCores),
+    transfers.downstreamCores.end());
+  llvm::erase_if(transfers.downstreamCores, [&](int64_t core) {
+    return llvm::is_contained(transfers.stageZeroCores, core);
+  });
+
+  DenseMap<int64_t, SmallVector<HostTransferRef>> incomingByCore;
+  DenseMap<ExternalTransferFamily *, SmallVector<int64_t>> eventRegisters;
+  DenseMap<ExternalTransferFamily *, SmallVector<int64_t>> waitValues;
+  DenseMap<ExternalTransferFamily *, SmallVector<int64_t>> acknowledgementRegisters;
+  auto initialize = [&](ExternalTransferFamily &family) {
+    size_t count = family.targetCores.size();
+    eventRegisters.try_emplace(&family, count, 0);
+    waitValues.try_emplace(&family, count, 0);
+    acknowledgementRegisters.try_emplace(&family, count, 0);
+  };
+  for (const BoundaryProgram &boundary : boundaries)
+    for (const BoundaryInstruction &instruction : boundary.instructions) {
+      auto *receive = std::get_if<EmitReceiveAssemblyRun>(&instruction);
+      if (!receive || receive->slices.empty()
+          || !receive->slices.front().family->hostRouted)
+        continue;
+      for (const ScheduledTransferSlice &slice : receive->slices) {
+        ExternalTransferFamily &family = *slice.family;
+        initialize(family);
+        for (size_t offset = 0; offset < slice.transferCount; ++offset) {
+          size_t index = slice.familyOffset + offset;
+          int64_t source = family.sourceCores.valueAt(index);
+          int64_t target = family.targetCores.valueAt(index);
+          incomingByCore[target].push_back({&family, index});
+          ++transfers.hostAcknowledgementCounts[source];
+        }
+      }
+    }
+
+  unsigned barrierRounds = getBarrierRoundCount(
+    transfers.stageZeroCores.size());
+  bool stageZeroNeedsAcknowledgements = llvm::any_of(
+    transfers.stageZeroCores, [&](int64_t core) {
+      return transfers.hostAcknowledgementCounts.contains(core);
+    });
+  for (auto &[target, incoming] : incomingByCore) {
+    bool needsAcknowledgementRegister =
+      transfers.hostAcknowledgementCounts.contains(target);
+    bool stageZero = llvm::is_contained(transfers.stageZeroCores, target);
+    size_t reserved = stageZero
+      ? barrierRounds + (stageZeroNeedsAcknowledgements ? 1 : 0)
+      : 1 + (needsAcknowledgementRegister ? 1 : 0);
+    if (reserved >= synchronizationRegisterCount) {
+      incoming.front().family->requirement->exchange->deferred.emitOpError(
+        "pipeline synchronization leaves no event register for incoming host transfers");
+      return failure();
+    }
+    size_t groupCount = std::min(
+      incoming.size(), synchronizationRegisterCount - reserved);
+    // One wait consumes a complete consecutive group of producer signals.
+    SmallVector<size_t> groupSizes(groupCount);
+    for (size_t ordinal = 0; ordinal < incoming.size(); ++ordinal)
+      ++groupSizes[ordinal * groupCount / incoming.size()];
+    SmallVector<bool> first(groupCount, true);
+    for (size_t ordinal = 0; ordinal < incoming.size(); ++ordinal) {
+      size_t group = ordinal * groupCount / incoming.size();
+      HostTransferRef transfer = incoming[ordinal];
+      eventRegisters[transfer.family][transfer.index] = group;
+      acknowledgementRegisters[transfer.family][transfer.index] =
+        synchronizationRegisterCount - 1;
+      if (first[group]) {
+        waitValues[transfer.family][transfer.index] = groupSizes[group];
+        first[group] = false;
+      }
+    }
+  }
+  for (auto &[family, values] : eventRegisters) {
+    family->eventRegisters = StaticIntSequence::fromValues(values);
+    family->waitValues = StaticIntSequence::fromValues(waitValues[family]);
+    family->acknowledgementEventRegisters =
+      StaticIntSequence::fromValues(acknowledgementRegisters[family]);
+  }
+
+  if (!transfers.stageZeroCores.empty()) {
+    size_t reserved = barrierRounds
+                    + (stageZeroNeedsAcknowledgements ? 1 : 0);
+    if (reserved > synchronizationRegisterCount)
+      return transfers.scheduled.front().op->emitOpError(
+        "pipeline stage-zero barrier requires more synchronization registers than the target provides");
+  }
+  if (!transfers.downstreamCores.empty()) {
+    bool needsAcknowledgements = llvm::any_of(
+      transfers.downstreamCores, [&](int64_t core) {
+        return transfers.hostAcknowledgementCounts.contains(core);
+      });
+    if (1 + (needsAcknowledgements ? 1 : 0)
+        > synchronizationRegisterCount)
+      return transfers.scheduled.front().op->emitOpError(
+        "pipeline stage-zero release requires more synchronization registers than the target provides");
+  }
+  return success();
+}
+
 } // namespace
 
 FailureOr<DeferredBoundaryPlan> buildDeferredBoundaryPlan(
     DeferredTransferPlan &transfers,
-    const ScheduledCommunicationPlan &schedule) {
+    const ScheduledCommunicationPlan &schedule,
+    size_t synchronizationRegisterCount) {
   DeferredBoundaryPlan result;
   SmallVector<BoundaryProgram> boundaries;
   DenseMap<BoundaryKey, unsigned> indices;
@@ -371,6 +508,9 @@ FailureOr<DeferredBoundaryPlan> buildDeferredBoundaryPlan(
     return std::tie(scheduledOrder[lhs.key.first], lhs.key.second)
          < std::tie(scheduledOrder[rhs.key.first], rhs.key.second);
   });
+  if (failed(assignPipelineSynchronization(
+        transfers, boundaries, synchronizationRegisterCount)))
+    return failure();
   result.boundaries = std::move(boundaries);
   return result;
 }

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.hpp

@@ -53,6 +53,7 @@ struct DeferredBoundaryPlan {
 };
 
 mlir::FailureOr<DeferredBoundaryPlan> buildDeferredBoundaryPlan(DeferredTransferPlan& transfers,
-                                                                const ScheduledCommunicationPlan& schedule);
+                                                                const ScheduledCommunicationPlan& schedule,
+                                                                size_t synchronizationRegisterCount);
 
 } // namespace onnx_mlir::spatial

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp

@@ -4,10 +4,12 @@
 #include "DeferredBoundaryRealization.hpp"
 #include "DeferredProjectionAnalysis.hpp"
 #include "DeferredResultRealization.hpp"
+#include "DeferredTransferPlanning.hpp"
 #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/StaticIntGrid.hpp"
 #include "src/Accelerators/PIM/Common/IR/StaticIntSequence.hpp"
 #include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include <array>
 namespace onnx_mlir::spatial {
 using namespace mlir;
@@ -18,6 +20,10 @@ struct LogicalTransferMetadataView {
   StaticIntSequenceChain parentCounts;
   StaticIntSequenceChain sourceCores;
   StaticIntSequenceChain targetCores;
+  StaticIntSequenceChain hostOffsets;
+  StaticIntSequenceChain eventRegisters;
+  StaticIntSequenceChain waitValues;
+  StaticIntSequenceChain acknowledgementEventRegisters;
   StaticIntSequenceChain targetLanes;
   StaticIntSequenceChain localOffsets;
   SmallVector<StaticIntSequenceChain> projectionOffsets;
@@ -28,7 +34,8 @@ struct LogicalTransferMetadataView {
 };
 using MetadataMember = StaticIntSequenceChain LogicalTransferMetadataView::*;
 static constexpr std::array<MetadataMember, 3> transferMetadataMembers{
-    &LogicalTransferMetadataView::channels, &LogicalTransferMetadataView::sourceCores, &LogicalTransferMetadataView::targetCores};
+    &LogicalTransferMetadataView::channels, &LogicalTransferMetadataView::sourceCores,
+    &LogicalTransferMetadataView::targetCores};
 struct TransferGrids {
   std::array<StaticIntGrid, 3> values;
   StaticIntGrid &channels() { return values[0]; }
@@ -41,7 +48,8 @@ template <typename Build> static FailureOr<TransferGrids> buildTransferGrids(Bui
   auto targetCores = build(transferMetadataMembers[2]);
   if (failed(channels) || failed(sourceCores) || failed(targetCores))
     return failure();
-  return TransferGrids{{std::move(*channels), std::move(*sourceCores), std::move(*targetCores)}};
+  return TransferGrids{{std::move(*channels), std::move(*sourceCores),
+                        std::move(*targetCores)}};
 }
 using GridGeometry = DeferredGridSliceGeometry;
 using StaticGeometryMember = SmallVector<StaticIntSequence> DeferredStaticSliceGeometry::*;
@@ -82,6 +90,14 @@ static void appendMetadata(const ScheduledTransferSlice &slice, LogicalTransferM
   metadata.parentCounts.append(StaticIntSequence::uniform(family.requirement->exchange->externalTransferCount, count));
   metadata.sourceCores.append(family.sourceCores, familyIndex, count);
   metadata.targetCores.append(family.targetCores, familyIndex, count);
+  if (family.hostRouted) {
+    metadata.hostOffsets.append(family.hostOffsets, familyIndex, count);
+    metadata.eventRegisters.append(
+      family.eventRegisters, familyIndex, count);
+    metadata.waitValues.append(family.waitValues, familyIndex, count);
+    metadata.acknowledgementEventRegisters.append(
+      family.acknowledgementEventRegisters, familyIndex, count);
+  }
   metadata.targetLanes.append(StaticIntSequence::affine(targetLane, 1, count));
   if (family.requirement->producerLocalOffsets)
     metadata.localOffsets.append(*family.requirement->producerLocalOffsets, targetLane - requirementLanes.begin, count);
@@ -172,6 +188,7 @@ static LogicalResult emitSendRun(const EmitSendRun &run, Value lane, unsigned la
     appendMetadata(slice, metadataByLane[sourceLane]);
   }
   LogicalTransferMetadataView logical = buildMetadataView(run.slices);
+  ExternalTransferFamily &firstFamily = *run.slices.front().family;
   size_t actionCount = 0;
   for (const LogicalTransferMetadataView &laneMetadata : metadataByLane)
     actionCount = std::max(actionCount, laneMetadata.size());
@@ -185,6 +202,20 @@ static LogicalResult emitSendRun(const EmitSendRun &run, Value lane, unsigned la
   FailureOr<StaticIntGrid> localOffsets = buildGrid(&LogicalTransferMetadataView::localOffsets, logical.localOffsets.valueAt(0));
   if (failed(transferGrids) || failed(localOffsets))
     return failure();
+  std::optional<StaticIntGrid> hostOffsets;
+  std::optional<StaticIntGrid> eventRegisters;
+  if (firstFamily.hostRouted) {
+    auto offsets = buildGrid(
+      &LogicalTransferMetadataView::hostOffsets,
+      logical.hostOffsets.valueAt(0));
+    auto events = buildGrid(
+      &LogicalTransferMetadataView::eventRegisters,
+      logical.eventRegisters.valueAt(0));
+    if (failed(offsets) || failed(events))
+      return failure();
+    hostOffsets = std::move(*offsets);
+    eventRegisters = std::move(*events);
+  }
   GridGeometry projectionGrids;
   for (auto [geometryIndex, sourceMember] : llvm::enumerate(metadataGeometryMembers)) {
     const auto &logicalValues = logical.*sourceMember;
@@ -207,7 +238,6 @@ static LogicalResult emitSendRun(const EmitSendRun &run, Value lane, unsigned la
     const LogicalTransferMetadataView &source = metadataByLane[sourceLane];
     counts[sourceLane] = source.size();
   }
-  ExternalTransferFamily &firstFamily = *run.slices.front().family;
   RequirementFamily &requirement = *firstFamily.requirement;
   Operation *anchor = requirement.exchange->deferred;
   Location loc = requirement.exchange->deferred.getLoc();
@@ -217,10 +247,25 @@ static LogicalResult emitSendRun(const EmitSendRun &run, Value lane, unsigned la
     auto payload = materializeSendPayload(requirement, localOffset, projectionGrids[0].empty() ? nullptr : &projection, context, loc);
     if (failed(payload))
       return failure();
-    auto send = SpatChannelSendOp::create(
-        context.rewriter, loc, transferGrids->channels().emitLookup(action, runtimeLane, anchor, context.constants, context.rewriter, loc),
-        transferGrids->sourceCores().emitLookup(action, runtimeLane, anchor, context.constants, context.rewriter, loc),
-        transferGrids->targetCores().emitLookup(action, runtimeLane, anchor, context.constants, context.rewriter, loc), *payload);
+    Value sourceCore = transferGrids->sourceCores().emitLookup(
+      action, runtimeLane, anchor, context.constants, context.rewriter, loc);
+    Value targetCore = transferGrids->targetCores().emitLookup(
+      action, runtimeLane, anchor, context.constants, context.rewriter, loc);
+    Operation *send;
+    if (firstFamily.hostRouted)
+      send = SpatHostStoreSyncOp::create(
+        context.rewriter, loc, sourceCore, targetCore,
+        hostOffsets->emitLookup(
+          action, runtimeLane, anchor, context.constants, context.rewriter, loc),
+        eventRegisters->emitLookup(
+          action, runtimeLane, anchor, context.constants, context.rewriter, loc),
+        *payload);
+    else
+      send = SpatChannelSendOp::create(
+        context.rewriter, loc,
+        transferGrids->channels().emitLookup(
+          action, runtimeLane, anchor, context.constants, context.rewriter, loc),
+        sourceCore, targetCore, *payload);
     setLogicalTransferMetadata(send, logical);
     return success();
   };
@@ -255,14 +300,57 @@ static FailureOr<Value> emitReceiveValue(ArrayRef<ScheduledTransferSlice> slices
   };
   auto grids = buildTransferGrids([&](MetadataMember member) { return buildGrid(metadata.*member); });
   if (failed(grids)) return failure();
+  std::optional<StaticIntGrid> hostOffsets;
+  std::optional<StaticIntGrid> eventRegisters;
+  std::optional<StaticIntGrid> waitValues;
+  std::optional<StaticIntGrid> acknowledgementEventRegisters;
+  if (slices.front().family->hostRouted) {
+    auto offsets = buildGrid(metadata.hostOffsets);
+    auto events = buildGrid(metadata.eventRegisters);
+    auto waits = buildGrid(metadata.waitValues);
+    auto acknowledgements = buildGrid(
+      metadata.acknowledgementEventRegisters);
+    if (failed(offsets) || failed(events) || failed(waits)
+        || failed(acknowledgements))
+      return failure();
+    hostOffsets = std::move(*offsets);
+    eventRegisters = std::move(*events);
+    waitValues = std::move(*waits);
+    acknowledgementEventRegisters = std::move(*acknowledgements);
+  }
   Value position = lane ? lane : context.constants.getIndex(0);
   Value row = context.constants.getIndex(0);
-  auto receive = SpatChannelReceiveOp::create(context.rewriter, anchor->getLoc(), requirement.publicationFragmentType,
-    grids->channels().emitLookup(row, position, anchor, context.constants, context.rewriter, anchor->getLoc()),
-    grids->sourceCores().emitLookup(row, position, anchor, context.constants, context.rewriter, anchor->getLoc()),
-    grids->targetCores().emitLookup(row, position, anchor, context.constants, context.rewriter, anchor->getLoc()));
+  Value sourceCore = grids->sourceCores().emitLookup(
+    row, position, anchor, context.constants, context.rewriter, anchor->getLoc());
+  Value targetCore = grids->targetCores().emitLookup(
+    row, position, anchor, context.constants, context.rewriter, anchor->getLoc());
+  Operation *receive;
+  Value output;
+  if (slices.front().family->hostRouted) {
+    auto op = SpatHostWaitLoadOp::create(
+      context.rewriter, anchor->getLoc(), requirement.publicationFragmentType,
+      sourceCore, targetCore,
+      hostOffsets->emitLookup(
+        row, position, anchor, context.constants, context.rewriter, anchor->getLoc()),
+      eventRegisters->emitLookup(
+        row, position, anchor, context.constants, context.rewriter, anchor->getLoc()),
+      waitValues->emitLookup(
+        row, position, anchor, context.constants, context.rewriter, anchor->getLoc()),
+      acknowledgementEventRegisters->emitLookup(
+        row, position, anchor, context.constants, context.rewriter, anchor->getLoc()));
+    receive = op;
+    output = op.getOutput();
+  } else {
+    auto op = SpatChannelReceiveOp::create(
+      context.rewriter, anchor->getLoc(), requirement.publicationFragmentType,
+      grids->channels().emitLookup(
+        row, position, anchor, context.constants, context.rewriter, anchor->getLoc()),
+      sourceCore, targetCore);
+    receive = op;
+    output = op.getOutput();
+  }
   setLogicalTransferMetadata(receive, metadata);
-  return receive.getOutput();
+  return output;
 }
 
 static FailureOr<SmallVector<LogicalTransferMetadataView, 0>>
@@ -315,6 +403,11 @@ static FailureOr<Value> emitReceiveAssembly(const EmitReceiveAssemblyRun &run, V
   SmallVector<int64_t> counts(laneCount);
   std::optional<TransferGrids> transferGrids;
   std::optional<StaticIntGrid> positions;
+  std::optional<StaticIntGrid> hostOffsets;
+  std::optional<StaticIntGrid> eventRegisters;
+  std::optional<StaticIntGrid> waitValues;
+  std::optional<StaticIntGrid> acknowledgementEventRegisters;
+  bool hostRouted = run.slices.front().family->hostRouted;
   auto metadataByEntry = buildRectangularReceiveMetadata(run, laneCount);
   if (succeeded(metadataByEntry)) {
     auto buildRows = [&](auto member) {
@@ -324,6 +417,23 @@ static FailureOr<Value> emitReceiveAssembly(const EmitReceiveAssemblyRun &run, V
       return StaticIntGrid::fromRows(rows);
     };
     auto grids = buildTransferGrids(buildRows);
+    if (hostRouted) {
+      auto offsets = buildRows(
+        &LogicalTransferMetadataView::hostOffsets);
+      auto events = buildRows(
+        &LogicalTransferMetadataView::eventRegisters);
+      auto waits = buildRows(
+        &LogicalTransferMetadataView::waitValues);
+      auto acknowledgements = buildRows(
+        &LogicalTransferMetadataView::acknowledgementEventRegisters);
+      if (failed(offsets) || failed(events) || failed(waits)
+          || failed(acknowledgements))
+        return failure();
+      hostOffsets = std::move(*offsets);
+      eventRegisters = std::move(*events);
+      waitValues = std::move(*waits);
+      acknowledgementEventRegisters = std::move(*acknowledgements);
+    }
     SmallVector<StaticIntSequence> positionRows;
     for (unsigned position : run.positions)
       positionRows.push_back(StaticIntSequence::uniform(position, laneCount));
@@ -368,6 +478,23 @@ static FailureOr<Value> emitReceiveAssembly(const EmitReceiveAssemblyRun &run, V
       return StaticIntGrid::fromColumns(actionCount, columns, defaultValue);
     };
     auto grids = buildTransferGrids(buildGrid);
+    if (hostRouted) {
+      auto offsets = buildGrid(
+        &LogicalTransferMetadataView::hostOffsets);
+      auto events = buildGrid(
+        &LogicalTransferMetadataView::eventRegisters);
+      auto waits = buildGrid(
+        &LogicalTransferMetadataView::waitValues);
+      auto acknowledgements = buildGrid(
+        &LogicalTransferMetadataView::acknowledgementEventRegisters);
+      if (failed(offsets) || failed(events) || failed(waits)
+          || failed(acknowledgements))
+        return failure();
+      hostOffsets = std::move(*offsets);
+      eventRegisters = std::move(*events);
+      waitValues = std::move(*waits);
+      acknowledgementEventRegisters = std::move(*acknowledgements);
+    }
     SmallVector<StaticIntSequence> positionColumns;
     for (const StaticIntSequenceChain &values : positionsByLane)
       positionColumns.push_back(
@@ -386,15 +513,38 @@ static FailureOr<Value> emitReceiveAssembly(const EmitReceiveAssemblyRun &run, V
   Value runtimeLane = lane ? lane : context.constants.getIndex(0);
   auto emitEntry = [&](Value entry, Value current) -> FailureOr<Value> {
     Type fragmentType = run.slices.front().family->requirement->publicationFragmentType;
-    auto receive =
-        SpatChannelReceiveOp::create(context.rewriter, loc, fragmentType,
-                                     transferGrids->channels().emitLookup(entry, runtimeLane, anchor, context.constants, context.rewriter, loc),
-                                     transferGrids->sourceCores().emitLookup(entry, runtimeLane, anchor, context.constants, context.rewriter, loc),
-                                     transferGrids->targetCores().emitLookup(entry, runtimeLane, anchor, context.constants, context.rewriter, loc));
+    Value sourceCore = transferGrids->sourceCores().emitLookup(
+      entry, runtimeLane, anchor, context.constants, context.rewriter, loc);
+    Value targetCore = transferGrids->targetCores().emitLookup(
+      entry, runtimeLane, anchor, context.constants, context.rewriter, loc);
+    Operation *receive;
+    Value output;
+    if (hostRouted) {
+      auto op = SpatHostWaitLoadOp::create(
+        context.rewriter, loc, fragmentType, sourceCore, targetCore,
+        hostOffsets->emitLookup(
+          entry, runtimeLane, anchor, context.constants, context.rewriter, loc),
+        eventRegisters->emitLookup(
+          entry, runtimeLane, anchor, context.constants, context.rewriter, loc),
+        waitValues->emitLookup(
+          entry, runtimeLane, anchor, context.constants, context.rewriter, loc),
+        acknowledgementEventRegisters->emitLookup(
+          entry, runtimeLane, anchor, context.constants, context.rewriter, loc));
+      receive = op;
+      output = op.getOutput();
+    } else {
+      auto op = SpatChannelReceiveOp::create(
+        context.rewriter, loc, fragmentType,
+        transferGrids->channels().emitLookup(
+          entry, runtimeLane, anchor, context.constants, context.rewriter, loc),
+        sourceCore, targetCore);
+      receive = op;
+      output = op.getOutput();
+    }
     setLogicalTransferMetadata(receive, logical);
     Value position = positions->emitLookup(
       entry, runtimeLane, anchor, context.constants, context.rewriter, loc);
-    return insert(receive.getOutput(), position, entry, runtimeLane, current);
+    return insert(output, position, entry, runtimeLane, current);
   };
   if (actionCount == 1 && llvm::all_of(counts, [](int64_t count) { return count == 1; }))
     return emitEntry(context.constants.getIndex(0), initial);
@@ -1046,9 +1196,199 @@ static LogicalResult emitBoundary(const BoundaryProgram &boundary, ArrayRef<Defe
   return failed(values) ? failure() : replaceResults(exchanges, *values, replacements);
 }
 
+static unsigned getBarrierRoundCount(size_t coreCount) {
+  unsigned rounds = 0;
+  for (size_t distance = 1; distance < coreCount; distance *= 2)
+    ++rounds;
+  return rounds;
+}
+
+static LogicalResult emitCompletionSynchronization(
+    DeferredTransferPlan &transfers, DeferredEmissionContext &context) {
+  if (transfers.synchronizationRegisterCount == 0)
+    return success();
+  size_t acknowledgementRegister =
+    transfers.synchronizationRegisterCount - 1;
+  unsigned barrierRounds = getBarrierRoundCount(
+    transfers.stageZeroCores.size());
+  bool stageZeroNeedsAcknowledgements = llvm::any_of(
+    transfers.stageZeroCores, [&](int64_t core) {
+      return transfers.hostAcknowledgementCounts.contains(core);
+    });
+  size_t firstBarrierRegister = acknowledgementRegister
+                              - (stageZeroNeedsAcknowledgements ? 1 : 0);
+  DenseMap<int64_t, unsigned> stageZeroRank;
+  for (auto [rank, core] : llvm::enumerate(transfers.stageZeroCores))
+    stageZeroRank[core] = rank;
+  DenseMap<int64_t, unsigned> downstreamRank;
+  for (auto [rank, core] : llvm::enumerate(transfers.downstreamCores))
+    downstreamRank[core] = rank;
+  auto getReleaseRegister = [&](int64_t core) {
+    return acknowledgementRegister
+         - (transfers.hostAcknowledgementCounts.contains(core) ? 1 : 0);
+  };
+
+  for (ScheduledInfo &scheduled : transfers.scheduled) {
+    Block *block = scheduled.blocks.front();
+    context.rewriter.setInsertionPoint(block->getTerminator());
+    Location loc = scheduled.op->getLoc();
+    Value lane;
+    if (auto batch = dyn_cast<SpatScheduledComputeBatch>(scheduled.op))
+      lane = *batch.getLaneArgument();
+
+    SmallVector<int64_t> acknowledgementCounts, releaseRegisters;
+    SmallVector<int64_t> releaseWaitValues, leftTargets, leftRegisters;
+    SmallVector<int64_t> rightTargets, rightRegisters;
+    LaneSet barrierLanes, leaderLanes, leftLanes, rightLanes;
+    for (auto [index, core] : llvm::enumerate(scheduled.cores)) {
+      acknowledgementCounts.push_back(
+        transfers.hostAcknowledgementCounts.lookup(core));
+      if (stageZeroRank.contains(core))
+        barrierLanes = barrierLanes.unite(
+          LaneSet::range(index, index + 1));
+      if (!transfers.stageZeroCores.empty()
+          && core == transfers.stageZeroCores.front())
+        leaderLanes = leaderLanes.unite(LaneSet::range(index, index + 1));
+
+      auto rank = downstreamRank.find(core);
+      if (rank == downstreamRank.end()) {
+        releaseRegisters.push_back(0);
+        releaseWaitValues.push_back(0);
+        leftTargets.push_back(core);
+        leftRegisters.push_back(0);
+        rightTargets.push_back(core);
+        rightRegisters.push_back(0);
+        continue;
+      }
+      releaseRegisters.push_back(getReleaseRegister(core));
+      releaseWaitValues.push_back(1);
+      size_t left = 2 * rank->second + 1;
+      size_t right = left + 1;
+      if (left < transfers.downstreamCores.size()) {
+        int64_t child = transfers.downstreamCores[left];
+        leftTargets.push_back(child);
+        leftRegisters.push_back(getReleaseRegister(child));
+        leftLanes = leftLanes.unite(LaneSet::range(index, index + 1));
+      } else {
+        leftTargets.push_back(core);
+        leftRegisters.push_back(0);
+      }
+      if (right < transfers.downstreamCores.size()) {
+        int64_t child = transfers.downstreamCores[right];
+        rightTargets.push_back(child);
+        rightRegisters.push_back(getReleaseRegister(child));
+        rightLanes = rightLanes.unite(LaneSet::range(index, index + 1));
+      } else {
+        rightTargets.push_back(core);
+        rightRegisters.push_back(0);
+      }
+    }
+    Value runtimeLane = lane ? lane : context.constants.getIndex(0);
+    auto emitForLanes = [&](const LaneSet &active, auto emit) -> LogicalResult {
+      if (active.empty())
+        return success();
+      if (!lane) {
+        if (active.contains(0))
+          emit();
+        return success();
+      }
+      auto condition = emitLaneCondition(
+        active, lane, scheduled.cores.size(), scheduled.op, context, loc);
+      if (failed(condition))
+        return failure();
+      auto conditional = scf::IfOp::create(
+        context.rewriter, loc, TypeRange {}, *condition, false);
+      OpBuilder::InsertionGuard guard(context.rewriter);
+      context.rewriter.setInsertionPoint(
+        conditional.getThenRegion().front().getTerminator());
+      emit();
+      return success();
+    };
+    Value acknowledgementCount = emitStaticIntLookup(
+      StaticIntSequence::fromValues(acknowledgementCounts),
+      runtimeLane, scheduled.op,
+      context.constants, context.rewriter, loc);
+    SpatWaitOp::create(
+      context.rewriter, loc,
+      context.constants.getIndex(acknowledgementRegister),
+      acknowledgementCount);
+
+    // Dissemination barrier: every round doubles the covered stage-zero peers.
+    auto emitBarrier = [&]() {
+      for (unsigned round = 0; round < barrierRounds; ++round) {
+        SmallVector<int64_t> targets;
+        targets.reserve(scheduled.cores.size());
+        size_t distance = size_t {1} << round;
+        for (int64_t core : scheduled.cores) {
+          auto rank = stageZeroRank.find(core);
+          targets.push_back(rank == stageZeroRank.end()
+            ? core
+            : transfers.stageZeroCores[
+                (rank->second + distance)
+                % transfers.stageZeroCores.size()]);
+        }
+        Value target = emitStaticIntLookup(
+          StaticIntSequence::fromValues(targets),
+          runtimeLane, scheduled.op,
+          context.constants, context.rewriter, loc);
+        Value eventRegister = context.constants.getIndex(
+          firstBarrierRegister - round);
+        SpatSyncOp::create(
+          context.rewriter, loc, target, eventRegister);
+        SpatWaitOp::create(
+          context.rewriter, loc, eventRegister,
+          context.constants.getIndex(1));
+      }
+    };
+    if (barrierRounds > 0
+        && failed(emitForLanes(barrierLanes, emitBarrier)))
+      return failure();
+
+    // Gate downstream restarts so no core advances the simulator input
+    // iteration ahead of stage zero.
+    if (!transfers.downstreamCores.empty()
+        && failed(emitForLanes(leaderLanes, [&]() {
+          int64_t root = transfers.downstreamCores.front();
+          SpatSyncOp::create(
+            context.rewriter, loc, context.constants.getIndex(root),
+            context.constants.getIndex(getReleaseRegister(root)));
+        })))
+      return failure();
+
+    Value releaseRegister = emitStaticIntLookup(
+      StaticIntSequence::fromValues(releaseRegisters), runtimeLane,
+      scheduled.op, context.constants, context.rewriter, loc);
+    Value releaseWaitValue = emitStaticIntLookup(
+      StaticIntSequence::fromValues(releaseWaitValues), runtimeLane,
+      scheduled.op, context.constants, context.rewriter, loc);
+    SpatWaitOp::create(
+      context.rewriter, loc, releaseRegister, releaseWaitValue);
+
+    auto emitChild = [&](ArrayRef<int64_t> targets,
+                         ArrayRef<int64_t> registers) {
+      Value target = emitStaticIntLookup(
+        StaticIntSequence::fromValues(targets), runtimeLane, scheduled.op,
+        context.constants, context.rewriter, loc);
+      Value eventRegister = emitStaticIntLookup(
+        StaticIntSequence::fromValues(registers), runtimeLane, scheduled.op,
+        context.constants, context.rewriter, loc);
+      SpatSyncOp::create(context.rewriter, loc, target, eventRegister);
+    };
+    if (failed(emitForLanes(leftLanes, [&]() {
+          emitChild(leftTargets, leftRegisters);
+        }))
+        || failed(emitForLanes(rightLanes, [&]() {
+          emitChild(rightTargets, rightRegisters);
+        })))
+      return failure();
+  }
+  return success();
+}
+
 } // namespace
 
-LogicalResult realizeDeferredBoundaries(ArrayRef<BoundaryProgram> boundaries, ArrayRef<DeferredResultPlan> results, DeferredEmissionContext &context,
+LogicalResult realizeDeferredBoundaries(ArrayRef<BoundaryProgram> boundaries, ArrayRef<DeferredResultPlan> results,
+                                        DeferredTransferPlan &transfers, DeferredEmissionContext &context,
                                         DeferredReplacementMap &replacements) {
   ScheduledInfo *scheduled = nullptr;
   for (const BoundaryProgram &boundary : boundaries) {
@@ -1059,7 +1399,7 @@ LogicalResult realizeDeferredBoundaries(ArrayRef<BoundaryProgram> boundaries, Ar
     if (failed(emitBoundary(boundary, results, context, replacements)))
       return boundary.key.first->op->emitOpError("phase 2 failed to realize a communication boundary");
   }
-  return success();
+  return emitCompletionSynchronization(transfers, context);
 }
 
 } // namespace onnx_mlir::spatial

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryRealization.hpp

@@ -42,6 +42,7 @@ using DeferredReplacementMap =
 
 mlir::LogicalResult realizeDeferredBoundaries(mlir::ArrayRef<BoundaryProgram> boundaries,
                                               mlir::ArrayRef<DeferredResultPlan> results,
+                                              DeferredTransferPlan& transfers,
                                               DeferredEmissionContext& context,
                                               DeferredReplacementMap& replacements);
 

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp

@@ -28,6 +28,11 @@ static std::optional<Event> getPlannedHead(
   while (cursor.slice < plan.slices.size()) {
     const ScheduledTransferSlice &slice = plan.slices[cursor.slice];
     ExternalTransferFamily &family = *slice.family;
+    if (family.hostRouted) {
+      ++cursor.slice;
+      cursor.offset = 0;
+      continue;
+    }
     size_t begin = slice.familyOffset + cursor.offset;
     size_t length = slice.transferCount - cursor.offset;
     auto source = family.sourceStreams.find(stream, begin, length);
@@ -243,6 +248,8 @@ LogicalResult verifyPlannedCommunicationDeadlockFree(
   DenseMap<ExternalTransferFamily *, unsigned> familyIndex;
   for (const ScheduledTransferSlice &slice : plan.slices) {
     ExternalTransferFamily *family = slice.family;
+    if (family->hostRouted)
+      continue;
     if (!familyIndex.try_emplace(family, familyIndex.size()).second)
       continue;
     size_t count = family->channelIds.size();
@@ -258,18 +265,6 @@ LogicalResult verifyPlannedCommunicationDeadlockFree(
     familyChannels.emplace_back(
       first, first + static_cast<int64_t>(count));
   }
-  llvm::sort(familyChannels);
-  int64_t nextChannel = 0;
-  for (auto [firstChannel, endChannel] : familyChannels) {
-    if (firstChannel != nextChannel)
-      return anchor->emitError(
-        "planned communication channels are not exactly contiguous");
-    nextChannel = endChannel;
-  }
-  if (static_cast<uint64_t>(nextChannel) != plan.logicalTransferCount)
-    return anchor->emitError(
-      "planned communication channel count is inconsistent");
-
   for (const ScheduledTransferSlice &slice : plan.slices) {
     ExternalTransferFamily &family = *slice.family;
     for (size_t offset = 0; offset < slice.transferCount; ++offset) {
@@ -296,6 +291,8 @@ LogicalResult verifyRealizedCommunicationDeadlockFree(
   DenseMap<ExternalTransferFamily *, unsigned> familyIndex;
   for (const ScheduledTransferSlice &slice : plan.slices) {
     ExternalTransferFamily *family = slice.family;
+    if (family->hostRouted)
+      continue;
     if (!familyIndex.try_emplace(family, familyIndex.size()).second)
       continue;
     for (size_t index = 0; index < family->channelIds.size(); ++index)
@@ -305,6 +302,8 @@ LogicalResult verifyRealizedCommunicationDeadlockFree(
   DenseMap<int64_t, StaticIntSequenceChain> expected;
   for (const ScheduledTransferSlice &slice : plan.slices) {
     ExternalTransferFamily &family = *slice.family;
+    if (family.hostRouted)
+      continue;
     appendEventsByCore(expected, family.channelIds, family.sourceCores,
                        slice.familyOffset, slice.transferCount, true);
     appendEventsByCore(expected, family.channelIds, family.targetCores,

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationModel.hpp

@@ -198,6 +198,7 @@ struct ScheduledInfo {
   llvm::SmallVector<mlir::Block*> blocks;
   llvm::SmallVector<mlir::Operation*> stepAnchors;
   llvm::SmallVector<int64_t> cores;
+  llvm::SmallVector<unsigned> pipelineStages;
   unsigned stepCount = 0;
   llvm::SmallVector<ProducedValue*> produced;
   llvm::SmallVector<unsigned> streamIds;
@@ -233,6 +234,12 @@ struct ExternalTransferFamily {
   StaticIntSequence sourceCores = StaticIntSequence::uniform(0, 1);
   StaticIntSequence targetCores = StaticIntSequence::uniform(0, 1);
   StaticIntSequence channelIds = StaticIntSequence::uniform(0, 1);
+  StaticIntSequence hostOffsets = StaticIntSequence::uniform(0, 1);
+  StaticIntSequence eventRegisters = StaticIntSequence::uniform(0, 1);
+  StaticIntSequence waitValues = StaticIntSequence::uniform(1, 1);
+  StaticIntSequence acknowledgementEventRegisters =
+    StaticIntSequence::uniform(0, 1);
+  bool hostRouted = false;
 };
 
 struct DeferredExchangePlan {

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp

@@ -34,7 +34,9 @@ static LogicalResult verifyNoEscapingRegionValues(Operation* owner, StringRef ph
                               << escapingUser->getName() << " at " << escapingUser->getLoc();
 }
 
-static LogicalResult placeLogicalProcessorsOnPhysicalCores(DeferredTransferPlan& plan, const SchedulingTarget& target) {
+static LogicalResult placeLogicalProcessorsOnPhysicalCores(
+    DeferredTransferPlan& plan, const SchedulingTarget& target,
+    size_t pipelineStages) {
   std::vector<Cost> logicalTrafficFlits(target.processorCount * target.processorCount, 0);
   for (const std::unique_ptr<DeferredExchangePlan>& exchange : plan.exchanges)
     for (const ExternalTransferFamily& transfer : exchange->external) {
@@ -55,8 +57,15 @@ static LogicalResult placeLogicalProcessorsOnPhysicalCores(DeferredTransferPlan&
       }
     }
 
+  std::vector<size_t> placementGroups;
+  if (pipelineStages > 1) {
+    if (plan.processorStages.size() != target.processorCount)
+      return failure();
+    placementGroups = plan.processorStages;
+  }
   std::vector<size_t> physicalCoreForLogicalProcessor =
-    mapLogicalProcessorsToPhysicalCores(logicalTrafficFlits, target);
+    mapLogicalProcessorsToPhysicalCores(
+      logicalTrafficFlits, target, placementGroups);
   auto getPhysicalCore = [&](int64_t logicalProcessor) {
     assert(logicalProcessor >= 0 && static_cast<size_t>(logicalProcessor) < physicalCoreForLogicalProcessor.size()
            && "logical processor is outside the scheduling target");
@@ -209,19 +218,32 @@ static LogicalResult verifyDominance(func::FuncOp funcOp) {
 
 LogicalResult realizeDeferredCommunication(func::FuncOp funcOp,
                                            const ScheduledComputeMaterializationResult& materialization,
-                                           const SchedulingTarget& target) {
+                                           const SchedulingTarget& target,
+                                           size_t pipelineStages) {
   IRRewriter rewriter(funcOp.getContext());
   eraseUnusedIdentityDeferredCommunications(funcOp, rewriter);
 
-  auto transfers = buildDeferredTransferPlan(funcOp, materialization);
+  auto transfers = buildDeferredTransferPlan(
+    funcOp, materialization, pipelineStages, target.processorCount);
   if (failed(transfers))
     return funcOp.emitOpError("phase 2 failed to build symbolic transfer families");
-  if (failed(placeLogicalProcessorsOnPhysicalCores(*transfers, target)))
+  if (failed(placeLogicalProcessorsOnPhysicalCores(
+        *transfers, target, pipelineStages)))
     return failure();
+  if (transfers->pipelineHostBufferBytes != 0) {
+    auto bytes = pim::checkedCast<int64_t>(
+      transfers->pipelineHostBufferBytes, funcOp,
+      "pipeline host transfer storage");
+    if (failed(bytes))
+      return failure();
+    funcOp->setAttr(kPipelineHostBufferBytesAttrName,
+                    rewriter.getI64IntegerAttr(*bytes));
+  }
   auto schedule = scheduleDeferredCommunication(funcOp, *transfers);
   if (failed(schedule) || failed(verifyPlannedCommunicationDeadlockFree(funcOp, transfers->stepCounts, *schedule)))
     return funcOp.emitOpError("phase 2 failed to schedule symbolic communication");
-  auto boundaries = buildDeferredBoundaryPlan(*transfers, *schedule);
+  auto boundaries = buildDeferredBoundaryPlan(
+    *transfers, *schedule, target.synchronizationRegisterCount);
   if (failed(boundaries))
     return funcOp.emitOpError("phase 2 failed to build sparse boundary programs");
 
@@ -231,7 +253,9 @@ LogicalResult realizeDeferredCommunication(func::FuncOp funcOp,
   ConstantPool constants(funcOp, rewriter);
   DeferredEmissionContext context(rewriter, constants);
   DeferredReplacementMap replacements;
-  if (failed(realizeDeferredBoundaries(boundaries->boundaries, boundaries->results, context, replacements)))
+  if (failed(realizeDeferredBoundaries(
+        boundaries->boundaries, boundaries->results, *transfers,
+        context, replacements)))
     return failure();
   for (auto [op, replacement] : replacements) {
     if (op->getResult(0) == replacement)

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationRealization.hpp

@@ -9,6 +9,7 @@ struct SchedulingTarget;
 
 mlir::LogicalResult realizeDeferredCommunication(mlir::func::FuncOp funcOp,
                                                  const ScheduledComputeMaterializationResult& materialization,
-                                                 const SchedulingTarget& target);
+                                                 const SchedulingTarget& target,
+                                                 size_t pipelineStages = 1);
 
 } // namespace onnx_mlir::spatial

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp

@@ -11,7 +11,7 @@ using namespace mlir;
 namespace {
 
 using TransferEmissionSignature =
-  std::tuple<ScheduledInfo*, Value, Type, bool, bool, bool>;
+  std::tuple<ScheduledInfo*, Value, Type, bool, bool, bool, bool>;
 
 static TransferEmissionSignature getTransferEmissionSignature(
     const ExternalTransferFamily& family) {
@@ -21,7 +21,8 @@ static TransferEmissionSignature getTransferEmissionSignature(
           family.requirement->publicationFragmentType,
           family.requirement->graphLanes.has_value(),
           family.requirement->producerProjection.has_value(),
-          producer->scheduled->isBatch()};
+          producer->scheduled->isBatch(),
+          family.hostRouted};
 }
 
 struct StreamThreshold {

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp

@@ -5,6 +5,7 @@
 #include "DeferredProjectionAnalysis.hpp"
 #include "DeferredTransferPlanning.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 
 namespace onnx_mlir::spatial {
 using namespace mlir;
@@ -28,7 +29,14 @@ static FailureOr<unsigned> getStepIndex(
 
 static LogicalResult collectScheduledOperations(
     const ScheduledComputeMaterializationResult &materialization,
-    DeferredTransferPlan &plan) {
+    DeferredTransferPlan &plan,
+    size_t pipelineStageCount,
+    size_t processorCount) {
+  if (pipelineStageCount == 0
+      || (pipelineStageCount > 1
+          && materialization.processorStages.size() != processorCount))
+    return failure();
+  plan.processorStages = materialization.processorStages;
   unsigned nextStream = 0;
   for (const ScheduledMaterializationRecord &record :
        materialization.materializedSchedules) {
@@ -46,8 +54,17 @@ static LogicalResult collectScheduledOperations(
     if (llvm::any_of(info.stepAnchors,
                      [](Operation *anchor) { return !anchor; }))
       return op.emitOpError("phase 2 scheduled step anchor is missing");
-    for (size_t core : record.cpus)
+    for (size_t core : record.cpus) {
+      if (core >= processorCount)
+        return op.emitOpError("phase 2 scheduled core is outside the target");
       info.cores.push_back(core);
+      if (pipelineStageCount > 1) {
+        size_t stage = materialization.processorStages[core];
+        if (stage >= pipelineStageCount)
+          return op.emitOpError("phase 2 scheduled core has an invalid pipeline stage");
+        info.pipelineStages.push_back(stage);
+      }
+    }
     for (size_t lane = 0; lane < info.cores.size(); ++lane)
       info.streamIds.push_back(nextStream++);
     plan.scheduled.push_back(std::move(info));
@@ -308,17 +325,21 @@ static LogicalResult buildRequirementFamilies(DeferredTransferPlan& plan,
   return success();
 }
 
-static void buildAvailabilityFamilies(DeferredExchangePlan& exchange, uint64_t& nextChannel) {
+static LogicalResult buildAvailabilityFamilies(
+    DeferredTransferPlan &plan,
+    DeferredExchangePlan& exchange,
+    uint64_t& nextChannel) {
+  enum class Availability { Local, Direct, Host };
   for (RequirementFamily& requirement : exchange.requirements) {
     for (LaneInterval interval : requirement.targetLanes.intervals()) {
       unsigned runBegin = interval.begin;
-      bool runLocal = false;
+      Availability runAvailability = Availability::Local;
       bool haveRun = false;
-      auto flush = [&](unsigned end) {
+      auto flush = [&](unsigned end) -> LogicalResult {
         if (!haveRun || runBegin == end)
-          return;
+          return success();
         LaneSet lanes = LaneSet::range(runBegin, end);
-        if (runLocal) {
+        if (runAvailability == Availability::Local) {
           exchange.local.push_back({&requirement, lanes});
         }
         else {
@@ -339,25 +360,60 @@ static void buildAvailabilityFamilies(DeferredExchangePlan& exchange, uint64_t&
           family.sourceCores = StaticIntSequence::uniform(requirement.producer->core, count);
           family.targetCores = StaticIntSequence::fromValues(targetCores);
           family.channelIds = StaticIntSequence::affine(nextChannel, 1, count);
+          family.hostRouted = runAvailability == Availability::Host;
+          if (family.hostRouted) {
+            auto fragmentType = dyn_cast<ShapedType>(
+              requirement.publicationFragmentType);
+            auto fragmentBytes = fragmentType
+              ? pim::getCheckedShapedTypeSizeInBytes(
+                  fragmentType, exchange.deferred,
+                  "pipeline host transfer fragment")
+              : FailureOr<uint64_t>(failure());
+            if (failed(fragmentBytes))
+              return failure();
+            auto bytes = pim::checkedMul<size_t>(
+              count, static_cast<size_t>(*fragmentBytes), exchange.deferred,
+              "pipeline host transfer storage");
+            if (failed(bytes))
+              return failure();
+            family.hostOffsets = StaticIntSequence::affine(
+              plan.pipelineHostBufferBytes, *fragmentBytes, count);
+            auto endOffset = pim::checkedAdd<size_t>(
+              plan.pipelineHostBufferBytes, *bytes, exchange.deferred,
+              "pipeline host transfer storage");
+            if (failed(endOffset))
+              return failure();
+            plan.pipelineHostBufferBytes = *endOffset;
+          }
           nextChannel += count;
           exchange.externalTransferCount += count;
           exchange.external.push_back(std::move(family));
         }
+        return success();
       };
       for (unsigned lane = interval.begin; lane < interval.end; ++lane) {
         unsigned sourceStream = requirement.producer->scheduled->streamIds[requirement.producer->scheduledLane];
         bool local =
           sourceStream == exchange.target->streamIds[lane] && requirement.producer->step < exchange.consumerStep;
-        if (haveRun && local != runLocal) {
-          flush(lane);
+        bool crossStage = !exchange.target->pipelineStages.empty()
+          && requirement.producer->scheduled->pipelineStages[
+               requirement.producer->scheduledLane]
+             != exchange.target->pipelineStages[lane];
+        Availability availability = local ? Availability::Local
+          : crossStage ? Availability::Host : Availability::Direct;
+        if (haveRun && availability != runAvailability) {
+          if (failed(flush(lane)))
+            return failure();
           runBegin = lane;
         }
-        runLocal = local;
+        runAvailability = availability;
         haveRun = true;
       }
-      flush(interval.end);
+      if (failed(flush(interval.end)))
+        return failure();
     }
   }
+  return success();
 }
 
 static LogicalResult buildExchanges(func::FuncOp funcOp, DeferredTransferPlan& plan) {
@@ -387,7 +443,8 @@ static LogicalResult buildExchanges(func::FuncOp funcOp, DeferredTransferPlan& p
     exchange->program = std::move(*program);
     if (failed(buildRequirementFamilies(plan, *exchange, publicationCache)))
       return failure();
-    buildAvailabilityFamilies(*exchange, nextChannel);
+    if (failed(buildAvailabilityFamilies(plan, *exchange, nextChannel)))
+      return failure();
     plan.exchanges.push_back(std::move(exchange));
   }
   return success();
@@ -464,9 +521,12 @@ retargetBlueprint(DeferredTransferPlan& plan, SpatBlueprintOp blueprint, GraphBa
 
 FailureOr<DeferredTransferPlan> buildDeferredTransferPlan(
     func::FuncOp funcOp,
-    const ScheduledComputeMaterializationResult &materialization) {
+    const ScheduledComputeMaterializationResult &materialization,
+    size_t pipelineStages,
+    size_t processorCount) {
   DeferredTransferPlan plan;
-  if (failed(collectScheduledOperations(materialization, plan))
+  if (failed(collectScheduledOperations(
+        materialization, plan, pipelineStages, processorCount))
       || failed(collectProducedValues(materialization, plan))
       || failed(buildExchanges(funcOp, plan)))
     return failure();

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.hpp

@@ -8,16 +8,24 @@
 namespace onnx_mlir::spatial {
 
 struct DeferredTransferPlan {
+  std::vector<size_t> processorStages;
   llvm::SmallVector<ScheduledInfo, 0> scheduled;
   llvm::SmallVector<std::unique_ptr<ProducedValue>> producedStorage;
   llvm::DenseMap<int64_t, llvm::SmallVector<ProducedValue*>> producedByGraph;
   llvm::SmallVector<std::unique_ptr<DeferredExchangePlan>> exchanges;
   llvm::SmallVector<unsigned> stepCounts;
+  llvm::DenseMap<int64_t, unsigned> hostAcknowledgementCounts;
+  llvm::SmallVector<int64_t> stageZeroCores;
+  llvm::SmallVector<int64_t> downstreamCores;
+  size_t synchronizationRegisterCount = 0;
+  size_t pipelineHostBufferBytes = 0;
 };
 
 mlir::FailureOr<DeferredTransferPlan>
 buildDeferredTransferPlan(mlir::func::FuncOp funcOp,
-                          const ScheduledComputeMaterializationResult &materialization);
+                          const ScheduledComputeMaterializationResult &materialization,
+                          size_t pipelineStages,
+                          size_t processorCount);
 
 mlir::LogicalResult retargetDeferredPublications(mlir::func::FuncOp funcOp, DeferredTransferPlan& plan);
 

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp

@@ -808,7 +808,9 @@ materializeScheduledCompute(func::FuncOp funcOp,
     }
   }
 
-  return ScheduledComputeMaterializationResult {std::move(peftClassPlans), std::move(materializedSchedules), std::move(graphComputeToBlockMap)};
+  return ScheduledComputeMaterializationResult {
+    std::move(peftClassPlans), std::move(materializedSchedules),
+    std::move(graphComputeToBlockMap), schedule.processorStages};
 }
 
 

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.hpp

@@ -14,6 +14,7 @@ struct ScheduledComputeMaterializationResult {
   llvm::MapVector<size_t, PeftClassPlan> peftClassPlans;
   std::vector<ScheduledMaterializationRecord> materializedSchedules;
   DenseMap<GraphComputeBlockKey, Block *> graphComputeToBlockMap;
+  std::vector<size_t> processorStages;
 };
 
 FailureOr<BatchFragmentSpec>

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.cpp

@@ -3,12 +3,15 @@
 #include "DeferredCommunicationRealization.hpp"
 #include "ScheduledComputeReport.hpp"
 #include "ScheduledComputeVerification.hpp"
+#include "Scheduling/PipelineScheduling.hpp"
 #include "SpatialDataflowCsvExporter.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.hpp"
 #include "src/Accelerators/PIM/Passes/PIMPasses.h"
 
+#include <limits>
+
 using namespace mlir;
 
 namespace onnx_mlir {
@@ -25,20 +28,54 @@ static bool hasValidTarget(const SchedulingTarget& target) {
 static FailureOr<func::FuncOp> requireEntry(ModuleOp moduleOp) {
   auto entry = getPimEntryFunc(moduleOp);
   if (failed(entry)) {
-    moduleOp.emitError("failed to locate the PIM entry function during Spatial scheduling and realization");
+    moduleOp.emitError("failed to locate the Pim entry function during Spatial scheduling and realization");
     return failure();
   }
   return *entry;
 }
 
+static SchedulingTarget getPipelineSchedulingTarget(
+    const SchedulingTarget& physicalTarget, size_t pipelineStages) {
+  if (pipelineStages == 1)
+    return physicalTarget;
+
+  PipelineCoreLayout layout(physicalTarget.processorCount, pipelineStages);
+  SchedulingTarget schedulingTarget = physicalTarget;
+  schedulingTarget.processorCount = layout.getLogicalProcessorCount();
+  schedulingTarget.residentWeightCapacity = checkedMultiply(
+    physicalTarget.residentWeightCapacity, pipelineStages);
+  schedulingTarget.interProcessorLatencyNs.assign(
+    schedulingTarget.processorCount * schedulingTarget.processorCount, 0);
+  Cost latencySum = 0;
+  size_t pairCount = 0;
+  for (size_t source = 0; source < schedulingTarget.processorCount; ++source)
+    for (size_t destination = 0;
+         destination < schedulingTarget.processorCount; ++destination) {
+      Cost latency = physicalTarget.getInterProcessorLatencyNs(
+        source, destination);
+      schedulingTarget.interProcessorLatencyNs[
+        source * schedulingTarget.processorCount + destination] = latency;
+      if (source != destination) {
+        latencySum = checkedAdd(latencySum, latency);
+        ++pairCount;
+      }
+    }
+  schedulingTarget.averageInterProcessorLatencyNs = pairCount == 0
+    ? 0
+    : (latencySum + pairCount - 1) / pairCount;
+  return schedulingTarget;
+}
+
 struct ScheduleAndRealizeSpatialPass final
     : PassWrapper<ScheduleAndRealizeSpatialPass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ScheduleAndRealizeSpatialPass)
 
   ScheduleAndRealizeSpatialPass() = default;
   ScheduleAndRealizeSpatialPass(const SchedulingTarget& target,
-                                SpatialDataflowExportStage exportStage)
-      : target(target), exportStage(exportStage), hasTarget(true) {}
+                                SpatialDataflowExportStage exportStage,
+                                size_t pipelineStages)
+      : target(target), exportStage(exportStage),
+        pipelineStages(pipelineStages), hasTarget(true) {}
 
   StringRef getArgument() const override { return "schedule-and-realize-spatial"; }
   StringRef getDescription() const override {
@@ -52,6 +89,16 @@ struct ScheduleAndRealizeSpatialPass final
       signalPassFailure();
       return;
     }
+    PipelineCoreLayout pipelineLayout(target.processorCount, pipelineStages);
+    if (!pipelineLayout.isValid()
+        || (pipelineStages > 1
+            && target.synchronizationRegisterCount == 0)
+        || target.residentWeightCapacity
+             > std::numeric_limits<size_t>::max() / pipelineStages) {
+      moduleOp.emitError("ScheduleAndRealizeSpatial requires valid pipeline stages and resource counts");
+      signalPassFailure();
+      return;
+    }
     auto entry = requireEntry(moduleOp);
     if (failed(entry)) {
       signalPassFailure();
@@ -59,8 +106,36 @@ struct ScheduleAndRealizeSpatialPass final
     }
     func::FuncOp entryFunc = *entry;
 
-    MergeSchedulingAnalysis analysis(entryFunc, target);
-    MergeScheduleResult schedule = std::move(analysis.getResult());
+    SchedulingTarget schedulingTarget = getPipelineSchedulingTarget(
+      target, pipelineStages);
+    ComputeGraph scheduledGraph;
+    MergeScheduleResult schedule;
+    for (;;) {
+      MergeSchedulingAnalysis analysis(
+        entryFunc, schedulingTarget,
+        pipelineStages > 1 ? target.processorCount : 0);
+      scheduledGraph = analysis.getGraph();
+      schedule = std::move(analysis.getResult());
+      std::string pipelineError;
+      if (pipelineStages > 1) {
+        FailureOr<PipelineWorkloadPreparation> preparation =
+          preparePipelineWorkload(
+            scheduledGraph, schedule, pipelineStages, target, pipelineError);
+        if (failed(preparation)) {
+          moduleOp.emitError() << pipelineError;
+          signalPassFailure();
+          return;
+        }
+        if (*preparation == PipelineWorkloadPreparation::Changed)
+          continue;
+      }
+      if (succeeded(applyPipelineScheduling(
+            scheduledGraph, schedule, pipelineStages, target, pipelineError)))
+        break;
+      moduleOp.emitError() << pipelineError;
+      signalPassFailure();
+      return;
+    }
     PatternRewriter rewriter(moduleOp.getContext());
     FailureOr<ScheduledComputeMaterializationResult> materialization =
         materializeScheduledCompute(entryFunc, schedule, rewriter);
@@ -94,7 +169,8 @@ struct ScheduleAndRealizeSpatialPass final
         moduleOp, entryFunc, schedule, materializationResult.peftClassPlans,
         materializationResult.materializedSchedules);
 
-    if (failed(realizeDeferredCommunication(entryFunc, materializationResult, target))) {
+    if (failed(realizeDeferredCommunication(
+          entryFunc, materializationResult, target, pipelineStages))) {
       moduleOp.emitError("Spatial communication realization failed");
       signalPassFailure();
       return;
@@ -126,6 +202,7 @@ struct ScheduleAndRealizeSpatialPass final
 private:
   SchedulingTarget target;
   SpatialDataflowExportStage exportStage = SpatialDataflowExportStage::None;
+  size_t pipelineStages = 1;
   bool hasTarget = false;
 };
 
@@ -136,8 +213,11 @@ std::unique_ptr<Pass> createScheduleAndRealizeSpatialPass() {
 }
 
 std::unique_ptr<Pass> createScheduleAndRealizeSpatialPass(
-    const SchedulingTarget& target, SpatialDataflowExportStage exportStage) {
-  return std::make_unique<ScheduleAndRealizeSpatialPass>(target, exportStage);
+    const SchedulingTarget& target,
+    SpatialDataflowExportStage exportStage,
+    size_t pipelineStages) {
+  return std::make_unique<ScheduleAndRealizeSpatialPass>(
+    target, exportStage, pipelineStages);
 }
 
 } // namespace spatial

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

@@ -772,6 +772,11 @@ std::vector<ComputeGraphEdge> aggregateEdges(llvm::ArrayRef<ComputeGraphEdge> ed
 
 } // namespace
 
+TransferCost getTransferCostFromBytes(Cost bytes,
+                                      const SchedulingTarget& target) {
+  return SchedulerCostModel {target}.getTransferCostFromBytes(bytes);
+}
+
 uint64_t countComputeBodyInstructions(Region& body) {
   uint64_t numOperations = 0;
   body.walk([&](Operation* op) { numOperations = checkedAdd(numOperations, static_cast<uint64_t>(1)); });
@@ -875,9 +880,13 @@ ResidentWeightSet getComputeInstanceResidentWeights(const ComputeInstance& insta
   return tiled;
 }
 
-ComputeGraph buildComputeGraph(Operation* entryOp, const SchedulingTarget& target) {
+ComputeGraph buildComputeGraph(Operation* entryOp,
+                               const SchedulingTarget& target,
+                               size_t computePartitionCount) {
   ComputeGraph graph;
   SchedulerCostModel costModel {target};
+  if (computePartitionCount == 0)
+    computePartitionCount = target.processorCount;
 
   for (Region& region : entryOp->getRegions()) {
     for (Block& block : region) {
@@ -898,10 +907,10 @@ ComputeGraph buildComputeGraph(Operation* entryOp, const SchedulingTarget& targe
           if (isUsedAsWeightOnly(batch.getOperation()))
             continue;
           size_t chunkCount =
-            getBatchChunkTargetCount(batch, target.processorCount);
+            getBatchChunkTargetCount(batch, computePartitionCount);
           for (size_t chunkIndex = 0; chunkIndex < chunkCount; ++chunkIndex) {
             ComputeInstance instance = getBatchChunkForIndex(
-              batch, chunkIndex, target.processorCount);
+              batch, chunkIndex, computePartitionCount);
             size_t index = graph.nodes.size();
             graph.nodes.push_back({instance,
                                    getComputeInstanceCost(instance, target),
@@ -920,7 +929,7 @@ ComputeGraph buildComputeGraph(Operation* entryOp, const SchedulingTarget& targe
     for (Value input : inputs) {
       for (const ProducerValueRef& producerRef :
            collectProducerValueRefs(input, node.instance,
-                                    target.processorCount)) {
+                                    computePartitionCount)) {
         auto producerIt = graph.instanceToIndex.find(producerRef.instance);
         if (producerIt == graph.instanceToIndex.end())
           continue;

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

@@ -61,9 +61,13 @@ struct ComputeGraph {
   llvm::DenseMap<ComputeInstance, size_t> instanceToIndex;
 };
 
-ComputeGraph buildComputeGraph(mlir::Operation* entryOp, const SchedulingTarget& target);
+ComputeGraph buildComputeGraph(mlir::Operation* entryOp,
+                               const SchedulingTarget& target,
+                               size_t computePartitionCount = 0);
 bool verifyAcyclic(const ComputeGraph& graph);
 
+TransferCost getTransferCostFromBytes(Cost bytes,
+                                      const SchedulingTarget& target);
 uint64_t countComputeBodyInstructions(mlir::Region& body);
 uint64_t countComputeBodyOperationInstances(mlir::Region& body);
 Cost getComputeInstanceCost(const ComputeInstance& instance, const SchedulingTarget& target);

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedule.hpp

@@ -14,6 +14,7 @@ namespace spatial {
 
 struct MergeScheduleResult {
   size_t processorCount = 0;
+  std::vector<size_t> processorStages;
   std::vector<ComputeInstance> dominanceOrderCompute;
   llvm::DenseMap<ComputeInstance, size_t> computeToCpuMap;
   llvm::DenseMap<ComputeInstance, size_t> computeToCpuSlotMap;

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp

@@ -89,13 +89,14 @@ void verifySchedule(const ComputeGraph& graph,
 } // namespace
 
 MergeSchedulingAnalysis::MergeSchedulingAnalysis(mlir::Operation* op,
-                                                 const SchedulingTarget& schedulingTarget)
-: entryOp(op), target(schedulingTarget) {
+                                                 const SchedulingTarget& schedulingTarget,
+                                                 size_t partitionCount)
+: entryOp(op), target(schedulingTarget), computePartitionCount(partitionCount) {
   result = run();
 }
 
 MergeScheduleResult MergeSchedulingAnalysis::run() {
-  ComputeGraph graph = buildComputeGraph(entryOp, target);
+  graph = buildComputeGraph(entryOp, target, computePartitionCount);
   if (!verifyAcyclic(graph))
     llvm::report_fatal_error("merge scheduling: compute graph is cyclic");
 

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

@@ -3,6 +3,7 @@
 #include "mlir/IR/Operation.h"
 
 #include "MergeSchedule.hpp"
+#include "ComputeGraph.hpp"
 #include "SchedulingTarget.hpp"
 
 namespace onnx_mlir {
@@ -10,12 +11,17 @@ namespace spatial {
 
 class MergeSchedulingAnalysis {
 public:
-  MergeSchedulingAnalysis(mlir::Operation* op, const SchedulingTarget& target);
+  MergeSchedulingAnalysis(mlir::Operation* op,
+                          const SchedulingTarget& target,
+                          size_t computePartitionCount = 0);
   MergeScheduleResult& getResult() { return result; }
+  const ComputeGraph& getGraph() const { return graph; }
 
 private:
   mlir::Operation* entryOp = nullptr;
   const SchedulingTarget& target;
+  size_t computePartitionCount = 0;
+  ComputeGraph graph;
   MergeScheduleResult result;
 
   MergeScheduleResult run();

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

@@ -244,11 +244,13 @@ FailureOr<LanePublicationSignatures> buildLanePublicationSignatures(SpatComputeB
 } // namespace
 
 std::vector<size_t> mapLogicalProcessorsToPhysicalCores(ArrayRef<Cost> logicalTrafficFlits,
-                                                        const SchedulingTarget& target) {
+                                                        const SchedulingTarget& target,
+                                                        ArrayRef<size_t> placementGroups) {
   const size_t processorCount = target.processorCount;
   assert(logicalTrafficFlits.size() == processorCount * processorCount
          && "logical traffic matrix must cover every processor pair");
-
+  assert((placementGroups.empty() || placementGroups.size() == processorCount)
+         && "physical placement groups must cover every processor");
   std::vector<size_t> physicalCoreForLogicalProcessor(processorCount);
   std::iota(physicalCoreForLogicalProcessor.begin(), physicalCoreForLogicalProcessor.end(), 0);
 
@@ -266,6 +268,10 @@ std::vector<size_t> mapLogicalProcessorsToPhysicalCores(ArrayRef<Cost> logicalTr
     for (size_t peerLogicalProcessor = 0; peerLogicalProcessor < processorCount; ++peerLogicalProcessor) {
       if (peerLogicalProcessor == logicalProcessor)
         continue;
+      if (!placementGroups.empty()
+          && placementGroups[peerLogicalProcessor]
+               != placementGroups[logicalProcessor])
+        continue;
       size_t physicalCore = physicalCoreForLogicalProcessor[logicalProcessor];
       size_t peerPhysicalCore = physicalCoreForLogicalProcessor[peerLogicalProcessor];
       Cost currentCost = 0;

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

@@ -29,7 +29,8 @@ inline Time getPeftTransferTime(const TransferCost& transferCost,
 MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftScheduleOptions& options);
 
 std::vector<size_t> mapLogicalProcessorsToPhysicalCores(llvm::ArrayRef<Cost> logicalTrafficFlits,
-                                                        const SchedulingTarget& target);
+                                                        const SchedulingTarget& target,
+                                                        llvm::ArrayRef<size_t> placementGroups = {});
 
 } // namespace spatial
 } // namespace onnx_mlir

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.hpp

@@ -0,0 +1,99 @@
+#pragma once
+
+#include "mlir/Support/LogicalResult.h"
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/STLExtras.h"
+
+#include <algorithm>
+#include <cstddef>
+#include <numeric>
+#include <optional>
+#include <string>
+#include <vector>
+
+#include "ComputeGraph.hpp"
+#include "MergeSchedule.hpp"
+#include "SchedulingTarget.hpp"
+
+namespace onnx_mlir::spatial {
+
+struct PipelineStageRange {
+  size_t begin;
+  size_t size;
+};
+
+class PipelineCoreLayout {
+public:
+  PipelineCoreLayout(size_t processorCount, size_t stageCount)
+      : processorCount(processorCount), stageSizes(stageCount) {
+    if (stageCount == 0)
+      return;
+    size_t baseSize = processorCount / stageCount;
+    size_t largerStageCount = processorCount % stageCount;
+    for (size_t stage = 0; stage < stageCount; ++stage)
+      stageSizes[stage] = baseSize + (stage < largerStageCount);
+  }
+
+  explicit PipelineCoreLayout(llvm::ArrayRef<size_t> stageSizes)
+      : processorCount(std::accumulate(
+          stageSizes.begin(), stageSizes.end(), size_t {0})),
+        stageSizes(stageSizes.begin(), stageSizes.end()) {}
+
+  bool isValid() const {
+    return !stageSizes.empty()
+        && llvm::none_of(stageSizes, [](size_t size) { return size == 0; });
+  }
+
+  size_t getLogicalProcessorCount() const {
+    return isValid()
+      ? *std::min_element(stageSizes.begin(), stageSizes.end())
+      : 0;
+  }
+
+  size_t getStageCount() const { return stageSizes.size(); }
+
+  size_t getProcessorCount() const { return processorCount; }
+
+  llvm::ArrayRef<size_t> getStageSizes() const { return stageSizes; }
+
+  PipelineStageRange getStageRange(size_t stage) const {
+    return {std::accumulate(
+              stageSizes.begin(), stageSizes.begin() + stage, size_t {0}),
+            stageSizes[stage]};
+  }
+
+  std::optional<size_t> getStageForCore(size_t core) const {
+    if (!isValid() || core >= processorCount)
+      return std::nullopt;
+    size_t end = 0;
+    for (auto [stage, size] : llvm::enumerate(stageSizes)) {
+      end += size;
+      if (core < end)
+        return stage;
+    }
+    return std::nullopt;
+  }
+
+private:
+  size_t processorCount;
+  std::vector<size_t> stageSizes;
+};
+
+mlir::LogicalResult applyPipelineScheduling(const ComputeGraph& graph,
+                                            MergeScheduleResult& schedule,
+                                            size_t pipelineStages,
+                                            const SchedulingTarget& physicalTarget,
+                                            std::string& error);
+
+enum class PipelineWorkloadPreparation {
+  Ready,
+  Changed,
+};
+
+mlir::FailureOr<PipelineWorkloadPreparation> preparePipelineWorkload(
+    const ComputeGraph& graph, const MergeScheduleResult& schedule,
+    size_t pipelineStages, const SchedulingTarget& physicalTarget,
+    std::string& error);
+
+} // namespace onnx_mlir::spatial

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

@@ -23,6 +23,7 @@ struct SchedulingTarget {
   Cost transferWidthBytes = 8;
   Cost vectorWidth = 16;
   Cost vectorLatencyCycles = 4;
+  size_t synchronizationRegisterCount = 0;
 
   Cost matrixRows = 128;
   Cost matrixColumns = 128;

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.cpp

@@ -193,7 +193,7 @@ FailureOr<TopLevelOpInfo> buildTopLevelOpInfo(Operation& op, bool isScheduled, s
 
   if constexpr (std::is_same_v<ComputeOpTy, SpatScheduledCompute>) {
     if (auto compute = dyn_cast<ComputeOpTy>(&op)) {
-      auto coreId = getOptionalScheduledCoreId(compute, "spatial dataflow export core id");
+      auto coreId = getOptionalScheduledCoreId(compute, "Spatial dataflow export core id");
       if (failed(coreId))
         return failure();
       if (*coreId)
@@ -207,7 +207,7 @@ FailureOr<TopLevelOpInfo> buildTopLevelOpInfo(Operation& op, bool isScheduled, s
 template <typename BatchOpTy>
 FailureOr<SmallVector<int32_t, 8>> getBatchLaneCoreIds(BatchOpTy batch) {
   if constexpr (std::is_same_v<BatchOpTy, SpatScheduledComputeBatch>) {
-    auto coreIds = getOptionalScheduledBatchCoreIds(batch, "spatial dataflow export core ids");
+    auto coreIds = getOptionalScheduledBatchCoreIds(batch, "Spatial dataflow export core ids");
     if (failed(coreIds))
       return failure();
     if (!*coreIds)

src/PIM/Dialect/Spatial/Spatial.td

@@ -13,7 +13,7 @@ include "mlir/Interfaces/SideEffectInterfaces.td"
 
 def SpatialDialect : Dialect {
   let name = "spat";
-  let summary = "Dialect designed for deep learning computation in a spatial architecture";
+  let summary = "Dialect designed for deep learning computation in a Spatial architecture";
   let cppNamespace = "::onnx_mlir::spatial";
   let useDefaultAttributePrinterParser = 0;
   let extraClassDeclaration = [{
@@ -550,7 +550,8 @@ def SpatChannelSendOp : SpatOp<"channel_send", []> {
   );
 
   let assemblyFormat = [{
-    $input `channel` $channelId `from` $sourceCoreId `to` $targetCoreId attr-dict `:` type($input)
+    $input `channel` $channelId `from` $sourceCoreId `to` $targetCoreId
+      attr-dict `:` type($input)
   }];
 }
 
@@ -568,7 +569,74 @@ def SpatChannelReceiveOp : SpatOp<"channel_receive", []> {
   );
 
   let assemblyFormat = [{
-    `channel` $channelId `from` $sourceCoreId `to` $targetCoreId attr-dict `:` type($output)
+    `channel` $channelId `from` $sourceCoreId `to` $targetCoreId
+      attr-dict `:` type($output)
+  }];
+}
+
+def SpatHostStoreSyncOp : SpatOp<"host_store_sync", []> {
+  let summary = "Store a tensor to host memory and signal its consumer";
+
+  let arguments = (ins
+    Index:$sourceCoreId,
+    Index:$targetCoreId,
+    Index:$hostOffset,
+    Index:$eventRegister,
+    SpatTensor:$input
+  );
+
+  let assemblyFormat = [{
+    $input `from` $sourceCoreId `to` $targetCoreId
+      `host_offset` $hostOffset `event` $eventRegister attr-dict `:` type($input)
+  }];
+}
+
+def SpatHostWaitLoadOp : SpatOp<"host_wait_load", []> {
+  let summary = "Wait for producers, load from host memory, and acknowledge consumption";
+
+  let arguments = (ins
+    Index:$sourceCoreId,
+    Index:$targetCoreId,
+    Index:$hostOffset,
+    Index:$eventRegister,
+    Index:$waitValue,
+    Index:$acknowledgementEventRegister
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+
+  let assemblyFormat = [{
+    `from` $sourceCoreId `to` $targetCoreId
+      `host_offset` $hostOffset `event` $eventRegister `count` $waitValue
+      `ack` $acknowledgementEventRegister attr-dict `:` type($output)
+  }];
+}
+
+def SpatSyncOp : SpatOp<"sync", []> {
+  let summary = "Signal a synchronization register on another processor";
+
+  let arguments = (ins
+    Index:$targetCoreId,
+    Index:$eventRegister
+  );
+
+  let assemblyFormat = [{
+    $targetCoreId `event` $eventRegister attr-dict
+  }];
+}
+
+def SpatWaitOp : SpatOp<"wait", []> {
+  let summary = "Wait for a synchronization register value";
+
+  let arguments = (ins
+    Index:$eventRegister,
+    Index:$waitValue
+  );
+
+  let assemblyFormat = [{
+    $eventRegister `value` $waitValue attr-dict
   }];
 }
 

src/PIM/Passes/PIMPasses.h

@@ -16,7 +16,8 @@ enum class SpatialDataflowExportStage;
 std::unique_ptr<mlir::Pass> createScheduleAndRealizeSpatialPass();
 std::unique_ptr<mlir::Pass> createScheduleAndRealizeSpatialPass(
   const SchedulingTarget& target,
-  SpatialDataflowExportStage exportStage);
+  SpatialDataflowExportStage exportStage,
+  size_t pipelineStages = 1);
 }
 
 std::unique_ptr<mlir::Pass> createONNXToSpatialPass();
@@ -24,7 +25,8 @@ std::unique_ptr<mlir::Pass> createONNXToSpatialPass(
   const spatial::SpatialTargetResources& target,
   const ONNXToSpatialPlanningOptions& options);
 std::unique_ptr<mlir::Pass> createSpatialLayoutPlanningPass();
-std::unique_ptr<mlir::Pass> createSpatialLayoutPlanningPass(const spatial::SpatialTargetResources& target);
+std::unique_ptr<mlir::Pass> createSpatialLayoutPlanningPass(
+  const spatial::SpatialTargetResources& target, bool selectTrivialPlan = false);
 std::unique_ptr<mlir::Pass> createLowerSpatialPlansPass();
 std::unique_ptr<mlir::Pass> createLowerSpatialPlansPass(
   const spatial::SpatialTargetResources& target,

src/PIM/Passes/PimCodegen/EmitPimCodePass.cpp

@@ -12,7 +12,7 @@ namespace {
 struct EmitPimCodePass : PassWrapper<EmitPimCodePass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(EmitPimCodePass);
   StringRef getArgument() const override { return "emit-pim-code-pass"; }
-  StringRef getDescription() const override { return "Emit PIM simulator code artifacts"; }
+  StringRef getDescription() const override { return "Emit Pim simulator code artifacts"; }
 
   EmitPimCodePass() {}
   EmitPimCodePass(const EmitPimCodePass& pass) {}
@@ -25,7 +25,7 @@ struct EmitPimCodePass : PassWrapper<EmitPimCodePass, OperationPass<ModuleOp>> {
 
     int compiler_error_code = compileToPimCode(moduleOp, pimDir);
     if (compiler_error_code != CompilerSuccess) {
-      moduleOp.emitError() << "failed to emit PIM simulator code artifacts; compiler error code "
+      moduleOp.emitError() << "failed to emit Pim simulator code artifacts; compiler error code "
                            << compiler_error_code;
       signalPassFailure();
     }

test/PIM/SpatialSchedulingTargetTest.cpp

@@ -1,12 +1,44 @@
 #include <cassert>
 #include <cstdlib>
+#include <string>
 #include <vector>
 
 #include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.hpp"
 
 using namespace onnx_mlir::spatial;
 
 int main() {
+  PipelineCoreLayout unevenLayout(138, 4);
+  assert(unevenLayout.isValid());
+  assert(unevenLayout.getLogicalProcessorCount() == 34);
+  assert(unevenLayout.getStageRange(0).begin == 0);
+  assert(unevenLayout.getStageRange(0).size == 35);
+  assert(unevenLayout.getStageRange(1).begin == 35);
+  assert(unevenLayout.getStageRange(1).size == 35);
+  assert(unevenLayout.getStageRange(2).begin == 70);
+  assert(unevenLayout.getStageRange(2).size == 34);
+  assert(unevenLayout.getStageRange(3).begin == 104);
+  assert(unevenLayout.getStageRange(3).size == 34);
+  assert(unevenLayout.getStageForCore(34) == 0);
+  assert(unevenLayout.getStageForCore(35) == 1);
+  assert(unevenLayout.getStageForCore(69) == 1);
+  assert(unevenLayout.getStageForCore(70) == 2);
+  assert(unevenLayout.getStageForCore(137) == 3);
+  assert(!unevenLayout.getStageForCore(138));
+
+  PipelineCoreLayout dynamicLayout(std::vector<size_t> {2, 4, 1, 3});
+  assert(dynamicLayout.isValid());
+  assert(dynamicLayout.getProcessorCount() == 10);
+  assert(dynamicLayout.getStageRange(0).begin == 0);
+  assert(dynamicLayout.getStageRange(1).begin == 2);
+  assert(dynamicLayout.getStageRange(2).begin == 6);
+  assert(dynamicLayout.getStageRange(3).begin == 7);
+  assert(dynamicLayout.getStageForCore(1) == 0);
+  assert(dynamicLayout.getStageForCore(2) == 1);
+  assert(dynamicLayout.getStageForCore(6) == 2);
+  assert(dynamicLayout.getStageForCore(9) == 3);
+
   TransferCost transfer {.fixed = 50, .networkFlits = 4};
 
   SchedulingTarget fast;
@@ -54,5 +86,144 @@ int main() {
     0,
   };
   assert(mapLogicalProcessorsToPhysicalCores(logicalTrafficFlits, alreadyPlaced) == std::vector<size_t>({0, 1, 2}));
+
+  std::vector<size_t> placementGroups {0, 1, 1};
+  std::vector<size_t> groupedPlacement = mapLogicalProcessorsToPhysicalCores(
+    logicalTrafficFlits, line, placementGroups);
+  for (size_t processor = 0; processor < groupedPlacement.size(); ++processor)
+    assert(placementGroups[processor]
+           == placementGroups[groupedPlacement[processor]]);
+
+  ComputeGraph graph;
+  graph.successors.resize(6);
+  graph.predecessors.resize(6);
+  graph.successors[1].push_back({2, TransferCost {.fixed = 1, .networkFlits = 1}});
+  graph.predecessors[2].push_back({1, TransferCost {.fixed = 1, .networkFlits = 1}});
+  const Cost costs[] = {6, 4, 6, 4, 8, 8};
+  for (uint32_t task = 0; task < 6; ++task) {
+    ComputeInstance instance {nullptr, task, 1};
+    ResidentWeight weight;
+    weight.opaqueLane = task;
+    graph.nodes.push_back({instance, costs[task], {weight}, task});
+    graph.instanceToIndex[instance] = task;
+  }
+
+  MergeScheduleResult logicalSchedule;
+  logicalSchedule.processorCount = 2;
+  logicalSchedule.dominanceOrderCompute.reserve(graph.nodes.size());
+  for (size_t task = 0; task < graph.nodes.size(); ++task) {
+    const ComputeInstance& instance = graph.nodes[task].instance;
+    logicalSchedule.dominanceOrderCompute.push_back(instance);
+    size_t cpu = task < 4 ? 0 : 1;
+    logicalSchedule.computeToCpuMap[instance] = cpu;
+    logicalSchedule.computeToCpuSlotMap[instance] = task < 4 ? task : task - 4;
+    logicalSchedule.computeToAestMap[instance] = task;
+  }
+
+  SchedulingTarget physical = fast;
+  physical.processorCount = 4;
+  physical.residentWeightCapacity = 2;
+  physical.interProcessorLatencyNs = {
+    0, 3, 3, 3,
+    3, 0, 3, 3,
+    3, 3, 0, 3,
+    3, 3, 3, 0,
+  };
+  std::string pipelineError;
+  ComputeGraph emptyGraph;
+  MergeScheduleResult emptySchedule;
+  emptySchedule.processorCount = 2;
+  assert(mlir::succeeded(applyPipelineScheduling(
+    emptyGraph, emptySchedule, 2, physical, pipelineError)));
+  assert(emptySchedule.processorCount == physical.processorCount);
+  assert(emptySchedule.processorStages == std::vector<size_t>({0, 0, 1, 1}));
+
+  MergeScheduleResult pipelineSchedule = logicalSchedule;
+  assert(mlir::succeeded(applyPipelineScheduling(
+    graph, pipelineSchedule, 2, physical, pipelineError)));
+  assert(pipelineSchedule.processorCount == 4);
+  size_t predecessorCore = pipelineSchedule.computeToCpuMap.lookup(
+    graph.nodes[1].instance);
+  size_t successorCore = pipelineSchedule.computeToCpuMap.lookup(
+    graph.nodes[2].instance);
+  assert(pipelineSchedule.computeToAestMap.lookup(graph.nodes[2].instance)
+         >= pipelineSchedule.computeToAestMap.lookup(graph.nodes[1].instance)
+              + graph.nodes[1].cost
+              + getPeftTransferTime(
+                  TransferCost {.fixed = 1, .networkFlits = 1},
+                  predecessorCore, successorCore, physical));
+  assert(pipelineSchedule.processorStages[predecessorCore]
+         <= pipelineSchedule.processorStages[successorCore]);
+  assert(pipelineSchedule.processorStages[successorCore]
+         <= pipelineSchedule.processorStages[predecessorCore] + 1);
+  assert(pipelineSchedule.equivalentClass.empty());
+
+  graph.successors[0].push_back(
+    {5, TransferCost {.fixed = 1, .networkFlits = 1}});
+  graph.predecessors[5].push_back(
+    {0, TransferCost {.fixed = 1, .networkFlits = 1}});
+  SchedulingTarget fourStagePhysical = physical;
+  fourStagePhysical.processorCount = 8;
+  fourStagePhysical.interProcessorLatencyNs.assign(64, 3);
+  for (size_t core = 0; core < 8; ++core)
+    fourStagePhysical.interProcessorLatencyNs[core * 8 + core] = 0;
+  MergeScheduleResult fourStageSchedule = logicalSchedule;
+  assert(mlir::succeeded(applyPipelineScheduling(
+    graph, fourStageSchedule, 4, fourStagePhysical, pipelineError)));
+  for (size_t task = 0; task < graph.nodes.size(); ++task)
+    for (const auto &[predecessor, cost] : graph.predecessors[task]) {
+      (void)cost;
+      size_t sourceStage = fourStageSchedule.processorStages[
+        fourStageSchedule.computeToCpuMap.lookup(
+          graph.nodes[predecessor].instance)];
+      size_t targetStage = fourStageSchedule.processorStages[
+        fourStageSchedule.computeToCpuMap.lookup(graph.nodes[task].instance)];
+      assert(sourceStage <= targetStage);
+      assert(targetStage <= sourceStage + 1);
+    }
+
+  ComputeGraph communicationGraph;
+  communicationGraph.successors.resize(5);
+  communicationGraph.predecessors.resize(5);
+  communicationGraph.successors[4].push_back(
+    {3, TransferCost {.fixed = 0, .networkFlits = 1}});
+  communicationGraph.predecessors[3].push_back(
+    {4, TransferCost {.fixed = 0, .networkFlits = 1}});
+  const Cost communicationCosts[] = {6, 4, 6, 4, 1};
+  MergeScheduleResult communicationSchedule;
+  communicationSchedule.processorCount = 2;
+  for (uint32_t task = 0; task < 5; ++task) {
+    ComputeInstance instance {nullptr, task, 1};
+    ResidentWeight weight;
+    weight.opaqueLane = task;
+    communicationGraph.nodes.push_back(
+      {instance, communicationCosts[task], {weight}, task});
+    communicationGraph.instanceToIndex[instance] = task;
+    communicationSchedule.dominanceOrderCompute.push_back(instance);
+    size_t cpu = task < 4 ? 0 : 1;
+    communicationSchedule.computeToCpuMap[instance] = cpu;
+    communicationSchedule.computeToCpuSlotMap[instance] = task < 4 ? task : 0;
+    communicationSchedule.computeToAestMap[instance] = task;
+  }
+
+  SchedulingTarget fastPipeline = physical;
+  fastPipeline.residentWeightCapacity = 4;
+  MergeScheduleResult fastCommunicationSchedule = communicationSchedule;
+  assert(mlir::succeeded(applyPipelineScheduling(
+    communicationGraph, fastCommunicationSchedule, 2, fastPipeline, pipelineError)));
+
+  SchedulingTarget slowPipeline = fastPipeline;
+  slowPipeline.averageInterProcessorLatencyNs = 10;
+  MergeScheduleResult slowCommunicationSchedule = communicationSchedule;
+  assert(mlir::succeeded(applyPipelineScheduling(
+    communicationGraph, slowCommunicationSchedule, 2, slowPipeline, pipelineError)));
+  size_t sourceCore = slowCommunicationSchedule.computeToCpuMap.lookup(
+    communicationGraph.nodes[4].instance);
+  size_t targetCore = slowCommunicationSchedule.computeToCpuMap.lookup(
+    communicationGraph.nodes[3].instance);
+  assert(slowCommunicationSchedule.processorStages[sourceCore]
+         <= slowCommunicationSchedule.processorStages[targetCore]);
+  assert(slowCommunicationSchedule.processorStages[targetCore]
+         <= slowCommunicationSchedule.processorStages[sourceCore] + 1);
   return EXIT_SUCCESS;
 }

tools/analyze_spatial_ir_cardinality.py

@@ -238,7 +238,7 @@ def print_report(path: Path, counts: Counter, groups: dict[tuple[str, str], Chai
 
 
 def main() -> None:
-    parser = argparse.ArgumentParser(description="Analyze repeated Spatial/PIM tensor IR cardinality patterns.")
+    parser = argparse.ArgumentParser(description="Analyze repeated Spatial/Pim tensor IR cardinality patterns.")
     parser.add_argument("paths", nargs="+", help="MLIR files to analyze.")
     parser.add_argument("--limit", type=int, default=12, help="Maximum number of hot chains to print per file.")
     args = parser.parse_args()

tools/raptor_graph_explorer/raptor_graph_explorer.egg-info/PKG-INFO

@@ -169,11 +169,11 @@ Motifs are not inferred from rendered geometry. For each operation graph the too
 
 ## Viewer and API
 
-Except for spatial4, the viewer initially fetches an aggregate operation graph unless the browser has a saved view choice. Sigma renders the graph with WebGL. Controls cover report/view/metric selection, text and tensor search, mapping filters, self edges, relayout, fitting, and motif selection. Raw nodes and edges have their own selection details and never call aggregate-only detail or mapping endpoints. Mapping panels remain available for operation aggregate edges.
+The viewer initially fetches an aggregate operation graph, so a previously selected raw view cannot make a new report exceed the display safety cap before the first render. Browser responses disable caching so HTML and JavaScript from different tool revisions cannot be mixed. Sigma renders the graph with WebGL. Controls cover report/view selection, text and tensor search, mapping filters, self edges, relayout, fitting, and motif selection. Edges use a fixed width. Raw nodes and edges have their own selection details and never call aggregate-only detail or mapping endpoints. Nodes without an SSA name fall back to their source identity. Mapping panels remain available for operation aggregate edges.
 
 Operation expansion is a deterministic projection of the source graph. Expand selected, Expand all operations, Collapse selected operation, and Collapse all rebuild the complete display from the current expanded-operation set. An expanded operation's raw nodes replace its aggregate node, including isolated raw nodes. An aggregate edge is retained only when both endpoint operations are collapsed; otherwise its raw edges replace it with endpoints calculated from the complete expansion set. Expansion order therefore cannot leave stale endpoints or simultaneous aggregate/raw representations.
 
-Operation ranks are calculated from the complete operation graph, including isolated operations. Collapsed nodes use stable operation anchors. Expanded nodes are sorted by lane and node ID; lane numbers form perpendicular rows, equal lanes align across adjacent operations, and lane-less nodes use distinct deterministic scalar rows. This same model is used by Rerun layout, so unrelated operations do not jump during expansion or collapse and disconnected nodes never pile up at `(0, 0)`.
+Operation ranks are calculated from the complete operation graph, including isolated operations, and every view flows top-to-bottom. A collapsed operation reserves one displayed row. Expanded nodes use compact rows for distinct lanes in numeric order, followed by deterministic lane-less rows, so sparse lane numbers create no empty geometric space. Raw nodes sharing a lane receive centered deterministic horizontal offsets. Operation and raw views use stable server coordinates, and `Reset layout` restores them; core and node-kind views use ELK, and `Rerun layout` runs it again. Operation labels use the first SSA name; Sigma's collision grid thins normal labels, and a label is skipped when its measured screen rectangle intersects another visible node. Selected or hovered nodes keep the original white bubble with black text. Selecting a motif frames its member nodes on the first click. Aggregate-node sizes remain bounded.
 
 Browser dependencies are pinned in one place, `static/index.html`:
 
@@ -206,8 +206,10 @@ Raw pages default to 100 and cannot exceed 500. Subgraph depth cannot exceed fiv
 ## Performance behavior
 
 - CSV readers stream rows and insert them in bounded batches.
-- Temporary SQLite staging and bulk joins resolve endpoints; ingestion performs no per-edge node query.
-- Secondary indexes are created after raw insertion.
+- Temporary SQLite staging tables are function-scoped and dropped immediately; bulk joins resolve endpoints without per-edge node queries.
+- CSVs without additional columns use a constant empty JSON representation.
+- Ingestion closes its write connection before derived indexes, aggregation, motifs, and diagnostics reopen the database.
+- Secondary indexes are created after raw insertion and focus on serving and aggregation queries.
 - Raw edges are never retained as a Python object graph or loaded into NetworkX.
 - Only aggregate operation nodes/edges enter NetworkX.
 - Mapping statistics are grouped in SQL; exact stencil comparison streams one operation pair.

tools/raptor_graph_explorer/raptor_graph_explorer.egg-info/SOURCES.txt

@@ -1,6 +1,7 @@
 README.md
 pyproject.toml
 raptor_graph_explorer/__init__.py
+raptor_graph_explorer/__main__.py
 raptor_graph_explorer/aggregate.py
 raptor_graph_explorer/api.py
 raptor_graph_explorer/cli.py

tools/raptor_graph_explorer/raptor_graph_explorer/__main__.py

@@ -0,0 +1,3 @@
+from .cli import main
+
+raise SystemExit(main())

validation/.gitignore

@@ -1,22 +1,12 @@
-operations/**/inputs
-operations/**/outputs
-operations/**/raptor
-operations/**/runner
-operations/**/simulation
+operations/**/artifacts
 operations/**/*.csv
 !operations/validation_results.csv
 
-networks/**/inputs
-networks/**/outputs
-networks/**/raptor
-networks/**/pimcomp
-networks/**/runner
-networks/**/simulation
-networks/**/real_image_val
+networks/**/artifacts
 networks/**/*.png
 networks/**/*.jpg
 networks/**/*.csv
 !networks/validation_results.csv
 !networks/full_net/validation_results.csv
-!networks/pimcomp_models/validation_results.csv
-!networks/pimcomp_models/results.csv
+!networks/pimcomp_models/results_comparison.csv
+!networks/pimcomp_models/results_ablation.csv

validation/README.md

@@ -1,14 +1,14 @@
-# Raptor Validation
+# Raptor validation
 
 `validate.py` validates every ONNX model below a selected directory. For each
 model it can:
 
 1. compile an ONNX-MLIR reference library and runner;
 2. generate deterministic random inputs;
-3. compile PIM artifacts with Raptor;
-4. run the reference implementation and functional PIM simulator;
+3. compile Pim artifacts with Raptor;
+4. run the reference implementation and functional Pim simulator;
 5. compare their outputs;
-6. run `pimsim-nn` to report latency, power, and energy.
+6. run Pimsim to report latency, throughput, power, and energy.
 
 Run the script from the repository root with the repository Python environment.
 
@@ -59,34 +59,23 @@ Validate a network or network slice:
 
 `--operations-dir` may point to any directory tree containing `.onnx` files.
 The script discovers them recursively and writes `validation_results.csv` in
-that directory while retaining the terminal table.
+that directory while retaining separate latency and throughput terminal tables.
 
-## Raptor vs PIMCOMP comparison
+## Pim validation tools
 
-The PIMCOMP paper-model suite has a one-command Arch-A comparison:
-
-```bash
-.venv/bin/python validation/tools/pimcomp/run_pimcomp_paper_latency.py
-```
-
-The runner verifies PIMCOMP's population-200, 1000-iteration GA settings,
-builds Raptor and the existing `third_party/PIMCOMP-NN/build` tree, and compares
-the four paper models one at a time. Each PIMCOMP GA run evaluates candidates
-in parallel; set `OMP_NUM_THREADS` to control its worker count. Use
-`--models vgg8` for one model or `--dry-run` to print the commands.
-
-Generated artifacts are stored beside each model and ignored by Git. The
-comparison reuses the same model-level `inputs/`, `outputs/`, `runner/`,
-`raptor/`, and `simulation/` paths as regular validation. PIMCOMP-only
-artifacts and `comparison_report.{md,json}` live under `pimcomp/`. See
-[`networks/pimcomp_models/README.md`](networks/pimcomp_models/README.md) for
-profiles, model provenance, limitations, and remote execution.
+- [Pimcomp model suite](networks/pimcomp_models/README.md)
+- [Pimcomp model comparison tools](tools/pim/pimcomp/compare/README.md)
+- [Pimcomp correctness study](tools/pim/pimcomp/correctness/README.md)
+- [Raptor compiler ablation study](tools/pim/ablation/README.md)
 
 ## Validation modes
 
-The default mode performs the complete workflow.
+The default mode runs latency and throughput in one validation job. Latency
+uses one input, while throughput uses `--pipeline=4` with four distinct inputs.
+Both modes reuse the generated input batch, native runner, and reference
+outputs, and every throughput output is compared with its own reference.
 
-Use `--compile-only` to build the reference runner and PIM artifacts without
+Use `--compile-only` to build the reference runner and Pim artifacts without
 executing either implementation:
 
 ```bash
@@ -133,22 +122,23 @@ count with `-j` or `--jobs`:
 | `--onnx-include-dir PATH` | ONNX-MLIR runtime include directory. Required unless `--clean` is used. |
 | `--operations-dir PATH` | Directory tree containing models. Defaults to `validation/operations`. |
 | `--simulator-dir PATH` | Functional `pim-simulator` crate directory. Defaults to the in-tree simulator. |
-| `--non-functional-simulator-build-dir PATH` | `pimsim-nn` build directory. Defaults to the in-tree build. |
+| `--non-functional-simulator-build-dir PATH` | Pimsim build directory. Defaults to the in-tree build. |
 | `--pimcomp-config {arch-a,arch-b,arch-c}` | Non-functional hardware/timing profile. Defaults to `arch-a`. |
-| `--skip-non-functional-simulation` | Skip `pimsim-nn` latency, power, and energy measurement. |
+| `--skip-non-functional-simulation` | Skip Pimsim latency, throughput, power, and energy measurement. |
+| `--no-fast` | Disable fast throughput convergence for authoritative full-duration Pimsim experiments. |
 | `--threshold FLOAT` | Absolute output-comparison tolerance. Defaults to `1e-3`. |
 | `--relative-threshold FLOAT` | Relative output-comparison tolerance. Defaults to `1e-5`. |
 | `--seed INT` | Seed for generated inputs. Defaults to `0`. |
 | `--crossbar-size INT` | Crossbar dimensions passed to Raptor. Defaults to the Arch-A value, `128`. |
 | `--crossbar-count INT` | Crossbars per core passed to Raptor. Defaults to the Arch-A value, `96`. |
-| `--core-count INT` | PIM core count passed to Raptor. Defaults to the Arch-A value, `168`. |
+| `--core-count INT` | Pim core count passed to Raptor. Defaults to the Arch-A value, `168`. |
 | `--raptor-extra-arg=ARG` | Additional Raptor compiler argument. Repeat for multiple arguments. |
 | `--command-timeout-seconds FLOAT` | Timeout for each compiler, runner, and simulator subprocess. Defaults to `1000000.0`. |
 | `-j INT`, `--jobs INT` | Parallel validation workers. Defaults to all available CPUs and must be at least one. |
 | `--clean` | Remove generated validation artifacts and exit. |
-| `--compile-only` | Compile reference and PIM artifacts without execution or comparison. |
+| `--compile-only` | Compile reference and Pim artifacts without execution or comparison. |
 | `--run-only` | Reuse compiled artifacts and perform execution, simulation, and comparison. |
-| `--verbose` | Print passing per-stage and subprocess logs, plus average PIM pass timings. |
+| `--verbose` | Print passing per-stage and subprocess logs, plus average Pim pass timings. |
 
 Arguments beginning with `--` that are passed through to Raptor should use the
 equals form:
@@ -159,40 +149,55 @@ equals form:
 
 ## Hardware profiles and non-functional simulation
 
-The selected PIMCOMP profile must match `--core-count`, `--crossbar-count`, and
+The selected Pimcomp profile must match `--core-count`, `--crossbar-count`, and
 `--crossbar-size`. A mismatch disables only non-functional simulation and
 prints the incompatible values; functional validation still runs.
 
 The checked-in profiles are under
-`validation/pimsim_configs/pimcomp/<profile>/latency_config.json`.
+`validation/pimsim_configs/pimcomp/<profile>/`; latency uses
+`latency_config.json`, while throughput uses
+`throughput_config_<time>ms.json` and the mesh beside it.
 
-Use `--skip-non-functional-simulation` when latency, power, and energy are not required.
+Throughput measurement defaults to `pimsim-nn --fast` with a 1000 ms
+convergence deadline. Fast mode compares consecutive two-round windows with a
+fixed 1% tolerance and falls back to the legacy full-duration result if it
+does not converge. Use `--no-fast` for authoritative experiments.
+
+Use `--skip-non-functional-simulation` when latency, throughput, power, and energy are not required.
 The summary reports non-functional results as measured, failed, unsupported, or
 skipped.
 
 Overall PASS/FAIL is determined by compilation and functional output
 comparison. A non-functional simulation failure remains visible as `ERROR` in
-the latency, power, and energy columns but does not change a functional PASS.
+the corresponding latency, throughput, power, or energy columns but does not
+change a functional PASS.
 
-`pimsim-nn` does not currently implement the `vsoftmax` instruction. When its
+Pimsim does not currently implement the `vsoftmax` instruction. When its
 explicit unsupported-op diagnostic is encountered, Softmax validations retain
 their functional PASS and show `UNSUPPORTED` in the non-functional columns.
-Other `pimsim-nn` failures remain `ERROR`.
+Other Pimsim failures remain `ERROR`.
 
 ## Generated artifacts
 
-Artifacts are written beside each model:
+Generated files are grouped below an `artifacts/` directory beside each model
+or operation case. This keeps checked-in ONNX files and generated trees
+separate and lets `--clean` remove the complete workspace, including stale
+validation lock files:
 
 | Path | Contents |
 |---|---|
-| `inputs/` | Generated input CSV files. |
-| `outputs/` | ONNX-MLIR reference output CSV files. |
-| `raptor/` | Exported MLIR, dialect snapshots, reports, final FP32 `pim/` artifacts, and the int8-equivalent `pimsim_nn/` latency view. |
-| `runner/` | Generated reference runner source, build tree, and shared library. |
-| `simulation/` | Functional simulator outputs used for numerical comparison. |
-| `pimcomp/` | PIMCOMP graph, instruction, simulator, and comparison-report artifacts. |
+| `artifacts/inputs.csv` | Generated inputs, one batch entry per line. |
+| `artifacts/inputs/`, `artifacts/outputs/`, `artifacts/runner/` | Inputs, reference outputs, and the runner shared by latency and throughput validation. |
+| `artifacts/raptor/pim/`, `artifacts/simulation/latency/` | Latency Pim artifacts and functional simulator outputs. |
+| `artifacts/raptor/throughput/pim/`, `artifacts/simulation/throughput/` | Pipeline-4, batch-4 throughput Pim artifacts and functional simulator outputs. |
+| `artifacts/common/inputs/` | Shared generated input CSV files. |
+| `artifacts/common/outputs/` | Shared ONNX-MLIR reference output CSV files. |
+| `artifacts/common/runner/` | Shared reference runner source, build tree, and library. |
+| `artifacts/<arch>/<mode>[/pipelineN]/raptor/` | Architecture- and pipeline-specific Raptor MLIR and Pim artifacts. |
+| `artifacts/<arch>/<mode>[/pipelineN]/simulation/` | Functional simulator outputs for that comparison. |
+| `artifacts/<arch>/<mode>[/pipelineN]/pimcomp/` | Pimcomp graph, instruction, simulator, and comparison-report artifacts. |
 
-The `raptor/` directory may include `spatial0.mlir`,
+Each comparison's `raptor/` directory may include `spatial0.mlir`,
 `spatial1_graph.mlir`, `spatial2_trivial_merged.mlir`,
 `spatial3_scheduled_no_comm.mlir`, `spatial4_scheduled.mlir`, `pim0.mlir`,
 `pim1_buff.mlir`, `pim2_folded.mlir`, and `pim3_memory_planned.mlir`.
@@ -221,27 +226,36 @@ The generated operation inventory is documented in
 
 ## Manual functional simulator tracing
 
-After validation has produced a `raptor/pim/` directory, rerun the functional
+After validation has produced an `artifacts/raptor/pim/` directory, rerun the functional
 simulator with tracing from its crate directory:
 
 ```bash
 cd backend-simulators/pim/pim-simulator
 cargo run --no-default-features --features tracing --release \
   --package pim-simulator --bin pim-simulator -- \
-  -f /path/to/workspace/raptor/pim \
-  -o /path/to/workspace/simulation/out.bin \
-  -d <addr0>,<size0>,<addr1>,<size1>,...
+  -f /path/to/workspace/artifacts/raptor/pim \
+  -o /path/to/workspace/artifacts/simulation/out.bin \
+  -d <addr0>,<size0>,<addr1>,<size1>,... \
+  --mode latency \
+  --batch-size 1 \
+  --input-dir /path/to/workspace/artifacts/simulation/inputs
 ```
 
+Throughput mode additionally requires `--batch-size N` and at least `N`
+`input_<index>.bin` files in `--input-dir`. Each input binary concatenates the model tensors in graph
+input order. The comparison validator also writes one native reference and one
+`simulation/*_iterations/output_*.bin` dump per batch entry, and checks every
+entry rather than only the final output.
+
 Tracing writes `TraceCore0`, `TraceCore1`, and so on beside `out.bin`. The
 validator normally derives the `-d` address and byte ranges from
 `raptor/pim/config.json` and the model output shapes.
 
 ## Results and exit status
 
-The final table reports functional pass/fail state and non-functional latency
-and power. The summary includes pass/fail totals, non-functional simulation
-counts, total measured latency, and average PIM pass timings when `--verbose`
+The final table reports latency and throughput functional pass/fail state plus
+non-functional latency, throughput, power, and energy. The summary includes pass/fail totals, non-functional simulation
+counts, total measured latency, and average Pim pass timings when `--verbose`
 is enabled.
 
 - Exit status `0`: all discovered models passed, or cleanup completed.

validation/networks/pimcomp_models/README.md

@@ -1,34 +1,38 @@
-# PIMCOMP comparison models
+# Pimcomp comparison models
 
 This directory contains the four networks evaluated in
-[PIMCOMP: An End-to-End DNN Compiler for Processing-In-Memory Accelerators](https://arxiv.org/pdf/2411.09159):
+[Pimcomp: An End-to-End DNN Compiler for Processing-In-Memory Accelerators](https://arxiv.org/pdf/2411.09159):
 VGG-8, ResNet-18, ResNet-34, and GoogLeNet. It also contains YOLO11n as an
 additional compiler comparison model.
 
-See the runner-generated [results.csv](results.csv) for the current latency
-and energy results.
+See the runner-generated [results_comparison.csv](results_comparison.csv) for the current comparison
+results. Rows are retained separately for each model, architecture, mode, and
+pipeline. It records separate `PASS`/`FAIL` functional-validation fields for
+the Raptor and Pimcomp artifacts; rows without a generated report contain `NA`.
+Use `--archs` to select the architecture rows to generate; existing rows for
+other architectures remain unchanged.
 
 ## Models and provenance
 
 | Directory    | Model                | Input         | Provenance                                                                                                                                                                                                 |
 |--------------|----------------------|---------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
 | `resnet18/`  | ResNet-18 v1         | `1x3x224x224` | Symlink to the complete [ONNX Model Zoo `resnet18-v1-7`](https://huggingface.co/onnxmodelzoo/resnet18-v1-7) model already present at `../resnet18/depth_68/resnet18_depth_68.onnx`.                        |
-| `resnet34/`  | ResNet-34 v1         | `1x3x224x224` | [ONNX Model Zoo `resnet34-v1-7`](https://huggingface.co/onnxmodelzoo/resnet34-v1-7), with its symbolic batch fixed to 1 as PIMCOMP's frontend does.                                                        |
+| `resnet34/`  | ResNet-34 v1         | `1x3x224x224` | [ONNX Model Zoo `resnet34-v1-7`](https://huggingface.co/onnxmodelzoo/resnet34-v1-7), with its symbolic batch fixed to 1 as Pimcomp's frontend does.                                                        |
 | `googlenet/` | GoogLeNet            | `1x3x224x224` | Unmodified [ONNX Model Zoo `googlenet-12`](https://huggingface.co/onnxmodelzoo/googlenet-12).                                                                                                              |
-| `vgg8/`      | VGG-8 reconstruction | `1x1x28x28`   | Reconstruction of the [PIMCOMP VGG-8 benchmark](https://arxiv.org/html/2411.09159#S8.SS1), with six convolution and two fully connected layers.                                                            |
+| `vgg8/`      | VGG-8 reconstruction | `1x1x28x28`   | Reconstruction of the [Pimcomp VGG-8 benchmark](https://arxiv.org/html/2411.09159#S8.SS1), with six convolution and two fully connected layers.                                                            |
 | `yolo11n/`   | YOLO11n detection    | `1x3x640x640` | Derived from the canonical local model at `../yolo11n/depth_51/yolo11n_depth_51.onnx`, exported from [Ultralytics YOLO11n](https://github.com/ultralytics/ultralytics/blob/main/docs/en/models/yolo11.md). |
 
-`googlenet/googlenet-12-latency.onnx` is the explicit pimsim-nn-ready GoogLeNet model.
+`googlenet/googlenet-12-pimsim-nn.onnx` is the explicit Pimsim-ready GoogLeNet model.
 It removes the two LRN nodes and terminal Softmax from the original model,
-so that the comparison covers only operations scheduled by PIMCOMP and supported by pimsim-nn.
+so that the comparison covers only operations scheduled by Pimcomp and supported by Pimsim.
 
-`yolo11n/yolo11n-latency.onnx` is the explicit pimsim-nn-ready YOLO11n model.
+`yolo11n/yolo11n-pimsim-nn.onnx` is the explicit Pimsim-ready YOLO11n model.
 It removes the Softmax nodes from the original model,
-so that the compiled artifact can be simulated in pimsim-nn.
+so that the compiled artifact can be simulated in Pimsim.
 
 ## Unsupported and ignored operations
 
-PIMCOMP's frontend accepts exactly these ONNX operations:
+Pimcomp's frontend accepts exactly these ONNX operations:
 
 ```text
 Add, AveragePool, BatchNormalization, Clip, Concat, Conv, Dropout, Flatten,
@@ -36,14 +40,14 @@ Gather, Gemm, GlobalAveragePool, LRN, MatMul, MaxPool, Mul, Pad, Relu, Reshape,
 Shape, Sigmoid, Softmax, Squeeze, Sub, Sum, Tanh, Transpose, Unsqueeze
 ```
 
-`Constant` is consumed as frontend metadata rather than emitted as a PIMCOMP
+`Constant` is consumed as frontend metadata rather than emitted as a Pimcomp
 node. Every other ONNX operation is unsupported: the frontend prints
 `operation: <type> not considered` and stops at the first occurrence. Thus the
 complete unsupported set is the complement of the allowlist above for the
 model's ONNX opset. In particular, YOLO11n contains unsupported `Split` and
 `Resize` nodes.
 
-PIMCOMP's low-latency scheduler, hierarchy mapper, and genetic algorithm use
+Pimcomp's low-latency scheduler, hierarchy mapper, and genetic algorithm use
 this complete explicit no-consider set:
 
 ```text
@@ -59,7 +63,7 @@ specific Shape-Gather-Unsqueeze-Concat shape chain, and merge Pad into its
 consumer. These transformations do not make an otherwise standalone ignored
 operation timed.
 
-`pimsim-nn` consumes PIM ISA instructions. It supports every named opcode
+Pimsim consumes Pim ISA instructions. It supports every named opcode
 in the shared serialized range except `vsoftmax` (opcode 21),
 which is rejected explicitly in both JSON and binary input. It
 silently ignores no opcode; unknown names and numbers are errors.
@@ -67,12 +71,12 @@ silently ignores no opcode; unknown names and numbers are errors.
 These boundaries explain the dedicated artifacts:
 
 - GoogLeNet's two LRN nodes and terminal Softmax perform real computation but
-  are ignored by PIMCOMP, so the common latency artifact removes them. Its
+  are ignored by Pimcomp, so the common latency artifact removes them. Its
   inference Dropout and shape-only Reshape can remain without adding compute.
 - YOLO11n's latency artifact bypasses exactly its two Softmax nodes so it can
-  run in `pimsim-nn`. Every other node, including MatMul, Transpose, and the
+  run in Pimsim. Every other node, including MatMul, Transpose, and the
   final detection-decoding tail, remains present and timed by Raptor. No
-  PIMCOMP latency is reported because its frontend stops at `Split` and also
+  Pimcomp latency is reported because its frontend stops at `Split` and also
   lacks `Resize`; compiling that prefix would not represent YOLO11n.
 
 The authoritative lists are in
@@ -81,16 +85,16 @@ The authoritative lists are in
 [`ISA.h`](../../../backend-simulators/pim/pimsim-nn/src/isa/ISA.h), and
 [`Instruction.cpp`](../../../backend-simulators/pim/pimsim-nn/src/isa/Instruction.cpp).
 
-The PIMCOMP authors did not publish the ONNX checkpoints used by the paper.
-Running PIMCOMP's frontend on the three Model Zoo files above produces JSON
-graphs exactly equal to PIMCOMP-NN's bundled `resnet18.json`, `resnet34.json`,
+The Pimcomp authors did not publish the ONNX checkpoints used by the paper.
+Running Pimcomp's frontend on the three Model Zoo files above produces JSON
+graphs exactly equal to Pimcomp's bundled `resnet18.json`, `resnet34.json`,
 and `googlenet.json`.
 
-There is no VGG-8 artifact in the ONNX Model Zoo or any PIMCOMP-NN revision.
+There is no VGG-8 artifact in the ONNX Model Zoo or any Pimcomp revision.
 The included VGG-8 therefore has deterministic random weights and is suitable
 for compiler and simulator comparison, not paper-accuracy reproduction. The
 paper also says that VGG-8 and ResNet-18 were trained on MNIST, while the
-published PIMCOMP graphs and ResNet Model Zoo artifacts use ImageNet shapes.
+published Pimcomp graphs and ResNet Model Zoo artifacts use ImageNet shapes.
 
 Current SHA-256 checksums:
 
@@ -98,9 +102,9 @@ Current SHA-256 checksums:
 788088b908e233d924c7c26b997e89ee861290c7bc56783a306e8201d79aac8f  resnet18/resnet18-v1-7.onnx
 c3231061d081bdd47884137b02134f85142752a39e87263c529cd14ed242b096  resnet34/resnet34-v1-7.onnx
 c99c507058eaf41de8723408fdda7db8325cb57f0a89f2ee07a716d6e963e14e  googlenet/googlenet-12.onnx
-a26f9e33901c573e60c34a3f0abbb4744fff83e4e0f21b18fc66e20395e72982  googlenet/googlenet-12-latency.onnx
+a26f9e33901c573e60c34a3f0abbb4744fff83e4e0f21b18fc66e20395e72982  googlenet/googlenet-12-pimsim-nn.onnx
 396cdea21e5e7d02c3f26f14d22ef20975171702493f5c5e79b8e0d896e541ef  vgg8/vgg8-mnist-reconstructed.onnx
-229f3975af8933d39aee8d9031d969bff074b69c33e304a78abb35ff0c5f445f  yolo11n/yolo11n-latency.onnx
+229f3975af8933d39aee8d9031d969bff074b69c33e304a78abb35ff0c5f445f  yolo11n/yolo11n-pimsim-nn.onnx
 ```
 
 ## Paper hardware profiles
@@ -109,10 +113,10 @@ The files in
 [`../../pimsim_configs/pimcomp/`](../../pimsim_configs/pimcomp/)
 encode Table V's explicit resource parameters.
 Each profile subdirectory contains pre-generated latency and throughput
-`pimsim-nn` configs plus its matching mesh; comparison and validation select
-these checked-in artifacts without generating configs at runtime.
+Pimsim configs plus its matching mesh; comparison and validation reference
+these canonical artifacts directly.
 
-| Config                       |             Cores | Crossbars/core | Crossbar   |  Cell | PIMCOMP layout  |
+| Config                       |             Cores | Crossbars/core | Crossbar   |  Cell | Pimcomp layout  |
 |------------------------------|------------------:|---------------:|------------|------:|-----------------|
 | `arch-a/latency_config.json` |               168 |             96 | `128x128`  | 2-bit | `12x14`         |
 | `arch-b/latency_config.json` |               138 |            128 | `128x128`  | 2-bit | `6x23`          |
@@ -121,10 +125,10 @@ these checked-in artifacts without generating configs at runtime.
 `adc_count` is 16, matching the paper's 16-bit fixed-point weight precision.
 The paper does not give a two-dimensional core topology for Arch-A/B, so the
 factorizations above preserve core count but cannot reproduce unpublished NoC
-placement details. Released PIMCOMP-NN has no chip-count field; Arch-C is
+placement details. Released Pimcomp has no chip-count field; Arch-C is
 therefore flattened to 64 cores and does not model chip boundaries.
 
-The remaining latency and power values come from PIMCOMP-NN's released default
+The remaining latency and power values come from Pimcomp's released default
 configuration. Consequently, instruction/resource comparisons are
 reproducible, but absolute paper power and energy numbers are not.
 
@@ -143,9 +147,9 @@ cmake --build third_party/PIMCOMP-NN/build --target PIMCOMP-NN
 
 Do not build either project with `ninja` directly.
 
-## Compile with PIMCOMP
+## Compile with Pimcomp
 
-PIMCOMP-NN reads `third_party/PIMCOMP-NN/config.json` directly. Back it up,
+Pimcomp reads `third_party/PIMCOMP-NN/config.json` directly. Back it up,
 select one paper profile, and restore it when the shell exits:
 
 ```bash
@@ -158,7 +162,7 @@ trap 'cp "$CONFIG_BACKUP" "$PIMCOMP/config.json"' EXIT
 cp "$PIMCOMP_CONFIGS/arch-a/latency_config.json" "$PIMCOMP/config.json"
 ```
 
-The Model Zoo files map exactly to PIMCOMP's bundled model names, so compile
+The Model Zoo files map exactly to Pimcomp's bundled model names, so compile
 them directly:
 
 ```bash
@@ -175,7 +179,7 @@ cd "$PIMCOMP/build"
 ./PIMCOMP-NN -m=googlenet -r=balance -p=element -o=YES -v=YES -s=YES
 ```
 
-VGG-8 first needs PIMCOMP's JSON frontend. Use a temporary ONNX copy because
+VGG-8 first needs Pimcomp's JSON frontend. Use a temporary ONNX copy because
 the released frontend rewrites the input batch dimension in place:
 
 ```bash
@@ -197,49 +201,77 @@ random placement code occasionally segfaults; an unchanged retry succeeded in
 the observed cases.
 
 The paper's optimizer uses a genetic algorithm with population 200 and up to
-1000 iterations. The checked-out PIMCOMP submodule already has both paper
+1000 iterations. The checked-out Pimcomp submodule already has both paper
 settings in `backend/GeneticAlgorithm.h`; select them with `-r=GA`. Fitness
 evaluation uses OpenMP and bounded bandwidth timelines. Set `OMP_NUM_THREADS`
 to control its parallelism; otherwise OpenMP uses the available CPUs. The GA
 uses the fixed seed `1`, so repeated serial and parallel runs are reproducible.
 
-## Compare Raptor and PIMCOMP
+## Compare Raptor and Pimcomp
 
 The comparison driver uses one random input and one native ONNX-MLIR reference,
-compiles both instruction streams, runs both through `pimsim-nn`, runs
+compiles both instruction streams, runs both through Pimsim, runs
 functional validation through `pim-simulator`, and writes Markdown and JSON
 reports.
 
-To reproduce the complete Arch-A latency experiment, use the model-by-model
-runner. It verifies the paper GA settings, builds Raptor and the existing
-`third_party/PIMCOMP-NN/build` tree, then runs the `element`/batch-1 comparison
-for one model at a time and regenerates `results.csv` from the JSON reports:
+To reproduce the default `arch-a`/`arch-b` architectures and both
+latency/throughput modes, use the model-by-model runner. Use `--archs` to
+specify a different architecture set. It verifies the paper GA settings, expects Raptor
+and the existing `third_party/PIMCOMP-NN/build` tree to already be built, then
+runs the comparisons in parallel and regenerates `results_comparison.csv` from the JSON
+reports:
 
 ```bash
-.venv/bin/python validation/tools/pimcomp/run_pimcomp_paper_latency.py
+.venv/bin/python validation/tools/pim/pimcomp/compare/run_pimcomp_models.py
 ```
 
-Each model directory reuses regular validation's ignored `inputs/`, `outputs/`,
-`runner/`, `raptor/`, and `simulation/` paths. PIMCOMP-only artifacts and
-`comparison_report.{md,json}` live under `pimcomp/`. The frontend regenerates
-one isolated `models/JSON/` graph because PIMCOMP requires that relative
+Use `--archs arch-a --mode latency` for only the Arch-A latency experiment, or
+`--archs arch-a arch-b arch-c` to run all three architectures.
+
+Each model directory has a shared ignored `artifacts/common/` directory containing
+`inputs/`, `outputs/`, and the native `runner/`.
+Pimsim configs and network meshes remain canonical under
+`validation/pimsim_configs/pimcomp/` and are referenced in place.
+Architecture- and pipeline-specific `raptor/`, `simulation/`, and Pimcomp
+artifacts remain under each `artifacts/<arch>/<mode>[/pipelineN]/` comparison
+directory; ablation variants use
+`artifacts/<arch>/<mode>[/pipelineN]/ablation/<variant>/`. Pimcomp outputs are prepared
+once per model/architecture/mode and linked into the other pipeline directories;
+`comparison_report.{md,json}`
+live under its `pimcomp/`. The frontend regenerates
+one isolated `models/JSON/` graph because Pimcomp requires that relative
 layout; it is removed after a successful backend run and the shared submodule
-model directory is never modified. Models requiring BatchNormalization folding
-also receive a prepared ONNX file; other models use the original ONNX directly.
-PIMCOMP's source tree and build directory remain unchanged at runtime. Use
-`--models vgg8` to run one model, `--resume` after an interruption, `--dry-run`
+model directory is never modified. The original ONNX model is passed to the
+frontend unchanged.
+Pimcomp's source tree and build directory remain unchanged at runtime. Use
+`--models vgg8` to run one model, `--mode throughput` to select one mode,
+`--pipeline 4` to select one throughput pipeline, `--only raptor` or
+`--only pimcomp` to reuse the other compiler's existing artifacts, `--dry-run`
 to inspect every command, or `--out-dir PATH` to keep results outside
-`validation/`. The runner continues after a failed model so all reports are
-produced.
+`validation/`. Use `--clean` to remove generated comparison artifacts and
+summaries, including stale reference lock files. Selecting a subset replaces only those comparison rows and
+recomputes the aggregate `results_comparison.csv`; missing shared inputs, outputs, or the
+reference runner are generated even for an isolated run. Use `--jobs 4` to cap
+parallel comparisons. The per-stage timeout is unlimited by default; pass a
+positive `--timeout-seconds` value to impose one.
+Throughput comparisons default to `pimsim-nn --fast` with a 1000 ms
+convergence deadline. Add `--no-fast` for authoritative full-duration runs.
+Pimcomp receives the original ONNX model, and its frontend applies native
+BatchNormalization fusion when the graph matches its supported Conv/Gemm pattern.
+The runner continues after a failed model so all reports are produced.
+
+The known Pimcomp batch-scheduling correctness issue and a reproducible
+reference-intermediate prefill experiment are documented in
+[`validation/tools/pim/pimcomp/correctness/README.md`](../../tools/pim/pimcomp/correctness/README.md).
 
 Arch-A low-latency example:
 
 ```bash
 RAPTOR_ROOT=$PWD
 
-"$RAPTOR_ROOT/.venv/bin/python" "$RAPTOR_ROOT/validation/tools/pimcomp/compare_raptor_pimcomp.py" \
+"$RAPTOR_ROOT/.venv/bin/python" "$RAPTOR_ROOT/validation/tools/pim/pimcomp/compare/compare_raptor_pimcomp_model.py" \
   --model "$RAPTOR_ROOT/validation/networks/pimcomp_models/resnet34/resnet34-v1-7.onnx" \
-  --out-dir "$RAPTOR_ROOT/validation/networks/pimcomp_models/resnet34" \
+  --out-dir "$RAPTOR_ROOT/validation/networks/pimcomp_models/resnet34/artifacts/arch-a/latency" \
   --pimcomp-config "$RAPTOR_ROOT/validation/pimsim_configs/pimcomp/arch-a/latency_config.json" \
   --core-count 168 \
   --crossbar-count 96 \
@@ -251,10 +283,10 @@ RAPTOR_ROOT=$PWD
 ```
 
 Use the same command with
-`yolo11n/yolo11n-latency.onnx` to probe YOLO11n. Released PIMCOMP-NN cannot
+`yolo11n/yolo11n-pimsim-nn.onnx` to probe YOLO11n. Released Pimcomp cannot
 compile it: the frontend stops at `/model.2/Split`, and it also has no mapping
 for YOLO11n's two nearest-neighbor `Resize` nodes. Treating the emitted prefix
-as YOLO11n would produce a misleading latency, so no PIMCOMP number is
+as YOLO11n would produce a misleading latency, so no Pimcomp number is
 reported for this model.
 
 For Arch-A high throughput, use `--pimsim-mode throughput
@@ -275,18 +307,18 @@ generated report.
 The functional and non-functional simulators intentionally consume different
 artifacts:
 
-- Raptor and PIMCOMP are validated against the native ONNX-MLIR reference as
+- Raptor and Pimcomp are validated against the native ONNX-MLIR reference as
   FP32 programs in the Rust simulator. Raptor's emitted program is already
-  FP32. The PIMCOMP-to-Rust export expands its element-addressed storage and
+  FP32. The Pimcomp-to-Rust export expands its element-addressed storage and
   byte-sized transfers to FP32, emits `setbw 32, 32`, and keeps vector
   `imm_len` fields as element counts.
-- PIMCOMP's original `SimulationInfo.gz` is copied unchanged for `pimsim-nn`.
-  PIMCOMP hardcodes `setbw 8, 8` and one byte per element without performing
+- Pimcomp's original `SimulationInfo.gz` is copied unchanged for Pimsim.
+  Pimcomp hardcodes `setbw 8, 8` and one byte per element without performing
   numerical quantization; this artifact is used only for latency estimation.
 - Raptor's original FP32 artifact remains unchanged for functional validation.
   A separate `raptor/pimsim_nn/` view uses `setbw 8, 8` and scales its
   byte-addressed storage and transfer sizes from four bytes to one byte per
-  element. Vector `imm_len` fields remain element counts. Like PIMCOMP's
+  element. Vector `imm_len` fields remain element counts. Like Pimcomp's
   artifact, this view is not numerically valid and is used only for a fair
   non-functional comparison.
 
@@ -297,11 +329,11 @@ either latency-only artifact for semantic validation.
 Current Raptor status:
 
 - VGG-8, ResNet-18, fixed-batch ResNet-34, and GoogLeNet compile on Arch-A.
-- Use `googlenet-12-latency.onnx` for the paper-matched latency comparison.
-  It removes the two LRN nodes and terminal softmax that PIMCOMP does not
+- Use `googlenet-12-pimsim-nn.onnx` for the paper-matched latency comparison.
+  It removes the two LRN nodes and terminal softmax that Pimcomp does not
   schedule.
 - Raptor currently accepts one square `--crossbar-size`; Arch-C's rectangular
-  `512x1024` arrays can therefore be compiled by PIMCOMP but not compared
+  `512x1024` arrays can therefore be compiled by Pimcomp but not compared
   exactly with Raptor.
 
 Do not change the hardware profile to bypass either limitation; that would no
@@ -319,8 +351,8 @@ rsync -azL validation/networks/pimcomp_models/ \
   "monolith:$REMOTE_REPO/validation/networks/pimcomp_models/"
 rsync -az validation/pimsim_configs/pimcomp/ \
   "monolith:$REMOTE_REPO/validation/pimsim_configs/pimcomp/"
-rsync -az validation/tools/pimcomp/ \
-  "monolith:$REMOTE_REPO/validation/tools/pimcomp/"
+rsync -az validation/tools/pim/ \
+  "monolith:$REMOTE_REPO/validation/tools/pim/"
 rsync -az --exclude=.git --exclude=build --exclude=output \
   third_party/PIMCOMP-NN/ \
   "monolith:$REMOTE_REPO/third_party/PIMCOMP-NN/"
@@ -334,10 +366,10 @@ cd /home/gmagnani/Project/Raptor
 
 # One-time setup if the repository virtual environment is absent.
 python3 -m venv .venv
-.venv/bin/python -m pip install numpy onnx onnxruntime onnxsim colorama
+.venv/bin/python -m pip install numpy onnx onnxruntime colorama
 
-# Run every latency comparison serially.
-.venv/bin/python validation/tools/pimcomp/run_pimcomp_paper_latency.py
+# Run every configured comparison in parallel.
+.venv/bin/python validation/tools/pim/pimcomp/compare/run_pimcomp_models.py
 ```
 
 Copy reports back without transferring large compiler artifacts:

validation/networks/pimcomp_models/results.csv

@@ -1,5 +0,0 @@
-model,raptor_latency_ms,pimcomp_latency_ms,raptor_energy_pj,pimcomp_energy_pj,faster_compiler,speedup
-vgg8,1.465778,7.985074,477298145.040001,1597904071.120000,raptor,5.45
-resnet18,28.099952,58.853733,8781611766.119984,13983168748.119974,raptor,2.09
-resnet34,45.781486,91.607980,14962940227.679951,22722922369.680016,raptor,2.00
-googlenet,13.371204,62.923463,6117835798.919991,14547526780.240000,raptor,4.71

validation/operations/README.md

@@ -1,7 +1,7 @@
-# Operation Validation Suite
+# Operation validation suite
 
 This directory contains the ONNX models used by `validation/validate.py` to
-validate individual operations through compilation, PIM simulation, and
+validate individual operations through compilation, Pim simulation, and
 comparison with the ONNX-MLIR reference runtime.
 
 ## Naming
@@ -38,12 +38,12 @@ Run the complete suite with deadlock detection:
 
 Use `--compile-only` for compiler and deadlock checks, then `--run-only` to
 reuse those artifacts for reference execution, simulation, and comparison.
-The validator prints the complete operation results table before its summary
-and writes the same rows to `validation_results.csv`.
+The validator prints separate latency and throughput operation tables before
+its summary and writes all of their rows to `validation_results.csv`.
 
 ## Complete inventory
 
-The suite contains 168 models. Tensor shapes, attributes, and constants are
+The suite contains 177 models. Tensor shapes, attributes, and constants are
 defined in `gen_tests.py` and in the checked-in ONNX models.
 
 ### Add (5)
@@ -64,12 +64,13 @@ defined in `gen_tests.py` and in the checked-in ONNX models.
 | `negative_axis` | Concatenates tensors using a negative axis. |
 | `three_inputs_channel_axis` | Concatenates three runtime NCHW tensors along the channel axis. |
 
-### Conv (34)
+### Conv (42)
 
 | Case | Description |
 |---|---|
 | `batch_2` | Batched Conv with SAME_UPPER padding and bias. |
 | `batch_4_pointwise` | Pointwise Conv with batch size four. |
+| `input_224_7x7_stride2` | 224x224 RGB Conv with 64 output channels, a 7x7 kernel, stride two, and bias. |
 | `depthwise_1024_channels` | Depthwise pointwise Conv with 1024 groups. |
 | `depthwise_grouped` | Depthwise-style grouped Conv with one input channel per group. |
 | `dilated_3x3` | Conv with a dilated 3x3 kernel. |
@@ -88,14 +89,21 @@ defined in `gen_tests.py` and in the checked-in ONNX models.
 | `non_square_kernel_1x3` | Conv with a non-square 1x3 kernel. |
 | `non_square_kernel_3x1` | Conv with a non-square 3x1 kernel. |
 | `non_uniform_stride` | Conv with different height and width strides. |
+| `output_channel_grouping_minimal` | Minimal 64-to-256 pointwise Conv for output-channel grouping. |
 | `pointwise_1x1` | Basic pointwise channel-mixing Conv. |
 | `pointwise_tiled_chain` | Relu and chained pointwise Convs with a tiled intermediate. |
 | `real_asymmetric_padding` | Conv with asymmetric explicit padding. |
 | `relu_conv_store` | Conv, Relu, and a second Conv to validate an intermediate stored result. |
 | `same_lower_3x3` | 3x3 Conv with SAME_LOWER padding. |
 | `same_padding_3x3` | 3x3 Conv with SAME_UPPER padding. |
-| `simple` | Hand-authored basic 2x2 Conv. |
+| `kernel_2x2` | Hand-authored Conv with a 2x2 kernel. |
 | `stride_2` | 3x3 Conv with stride two. |
+| `strategy_depthwise_16` | 16-channel depthwise 3x3 Conv for depthwise lowering. |
+| `strategy_input_k_tiled` | 32-channel 3x3 Conv sized to exercise input-K tiling. |
+| `strategy_output_channel_tiled` | 8-to-192 3x3 Conv sized to exercise output-channel tiling. |
+| `strategy_streamed_packed` | 3-to-16 3x3 Conv on 128x128 input for streamed packed lowering. |
+| `strategy_streamed_patch` | 3-to-16 3x3 Conv on 64x64 input for streamed patch lowering. |
+| `strategy_tiled_2d` | 32-to-192 3x3 Conv sized to exercise 2D tiling. |
 | `with_bias_3x3` | Multi-channel 3x3 Conv with bias. |
 | `with_constant` | Hand-authored SAME_UPPER Conv with constant weight and bias. |
 | `without_kernel_shape_attr` | Conv whose kernel shape is inferred from its weight tensor. |
@@ -135,30 +143,30 @@ defined in `gen_tests.py` and in the checked-in ONNX models.
 | `dynamic_beta` | Uses runtime operands and bias with non-default beta scaling. |
 | `dynamic_bias` | Uses runtime matrix operands and runtime bias. |
 | `dynamic_bias_alpha_beta` | Combines runtime operands and bias with alpha and beta scaling. |
-| `dynamic_transB` | Transposes a runtime right-hand matrix. |
+| `dynamic_transpose_b` | Transposes a runtime right-hand matrix. |
 | `huge_1024` | Uses 1024-wide inner and output dimensions. |
 | `large` | Exercises larger rectangular matrices. |
 | `large_k_small_n` | Uses a large reduction dimension and narrow output. |
 | `non_square` | Uses different reduction and output widths. |
 | `scalar_bias` | Broadcasts a scalar bias to the full output. |
-| `simple` | Basic Gemm with square weights. |
+| `square_weights` | Basic Gemm with square weights. |
 | `small` | Tiny Gemm for fast focused validation. |
 | `small_k_large_n` | Uses a modest reduction dimension and wide output. |
-| `transA` | Transposes the left-hand matrix. |
-| `transA_transB` | Transposes both matrix operands. |
-| `transB` | Transposes the right-hand weight matrix. |
-| `transB_with_bias` | Combines a transposed weight matrix with bias. |
+| `transpose_a` | Transposes the left-hand matrix. |
+| `transpose_a_and_b` | Transposes both matrix operands. |
+| `transpose_b` | Transposes the right-hand weight matrix. |
+| `transpose_b_with_bias` | Combines a transposed weight matrix with bias. |
 | `with_bias` | Basic matrix product with vector bias. |
 
 ### Gemv (5)
 
 | Case | Description |
 |---|---|
-| `constant` | Vector-matrix product with all inputs constant. |
-| `simple` | Basic single-row vector-matrix product. |
-| `with_heterogeneous_constant` | Adds a non-uniform constant bias pattern. |
-| `with_homogeneous_constant` | Adds a constant bias matching the output shape. |
-| `with_scalar_constant` | Adds a scalar broadcast bias. |
+| `all_constant` | Vector-matrix product with all inputs constant. |
+| `constant_weight` | Basic single-row vector-matrix product with constant weights. |
+| `non_uniform_bias` | Adds a non-uniform constant bias pattern. |
+| `uniform_bias` | Adds a uniform constant bias pattern. |
+| `scalar_bias` | Adds a scalar broadcast bias. |
 
 ### MatMul (12)
 
@@ -177,11 +185,12 @@ defined in `gen_tests.py` and in the checked-in ONNX models.
 | `vector_matrix` | Vector-matrix multiplication producing a 1D output. |
 | `yolo_attention` | YOLO11n rank-4 dynamic MatMul-scale-transpose-MatMul attention chain. |
 
-### Mul (5)
+### Mul (6)
 
 | Case | Description |
 |---|---|
 | `after_conv` | Conv followed by per-channel scaling. |
+| `after_conv_scalar_constant` | Conv followed by Mul with a scalar constant. |
 | `basic` | Elementwise Mul on two inputs with identical shapes. |
 | `channel_broadcast_1024` | Mul with NCHW per-channel broadcasting over 1024 channels. |
 | `leading_dimension_broadcast` | Mul with trailing-dimension broadcasting. |
@@ -254,7 +263,7 @@ defined in `gen_tests.py` and in the checked-in ONNX models.
 | `height_only` | Nearest-neighbor resize of only the height dimension. |
 | `nearest_2x` | Nearest-neighbor upsampling by a factor of two. |
 | `nearest_downsample` | Nearest-neighbor downsampling. |
-| `non_uniform` | Nearest-neighbor resize with different spatial scales. |
+| `non_uniform_scales` | Nearest-neighbor resize with different spatial scales. |
 | `width_only` | Nearest-neighbor resize of only the width dimension. |
 | `with_sizes` | Resize using explicit output sizes instead of scales. |