Raptor/README.md

# Raptor

Raptor is a domain-specific MLIR compiler for neural networks in ONNX format,
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`.
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`.

## Overview

PIM architectures perform most computation directly in memory. The supported
target models a chip with:
- shared host memory,
- multiple PIM cores,
- ReRAM crossbars for vector-matrix / matrix-vector work,
- explicit communication between cores,
- no hardware branch or loop support in emitted simulator code.

Because repeated work such as convolutions is eventually made explicit, emitted
instruction counts can grow quickly. Most compiler work therefore focuses on
lowering, scheduling, memory layout, and code-generation optimizations.

### Targets and simulators

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

The PIM sources live under `src/PIM` and tests under `test/PIM`. CMake exposes
them to ONNX-MLIR through generated shim directories under
`onnx-mlir/src/Accelerators/PIM` and `onnx-mlir/test/accelerators/PIM`.

High-level lowering flow:

```
ONNX-MLIR -> Spatial -> Pim (tensor) -> Pim (bufferized) -> Pim artifacts
```

1. **ONNX -> Spatial** (`src/PIM/Conversion/ONNXToSpatial`).
   Lowers supported ONNX ops into the `spat` dialect
   (`src/PIM/Dialect/Spatial`). Conversion patterns are split by op family under
   `Patterns/{Math,NN,Tensor}` and currently cover Conv, Gemm, MatMul,
   elementwise Add/Mul/Div, ReduceMean, pooling, Relu, Sigmoid, Softmax,
   Concat, Gather, Reshape, Resize, and Split.
   The compiler-layer target adapter supplies the target-neutral
   `SpatialTargetResources`. Layout-aware plan ops advertise typed alternatives
   through the Spatial layout interface; the layout planner records the
   selected layout and explicit materialization edges. `LowerSpatialPlans`
   then pattern-lowers those selected plans. Contraction and Conv lowering
   keep semantic problems, target-dependent plans, and IR materializers in
   separate layers. Passes and their invariant/layout analyses live under
   `Passes/Transforms` and `Passes/Analyses`.

2. **Merge, schedule, and realize Spatial communication**
   (`src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes`).
   `TrivialGraphComputeMerge` performs local graph merging. One
   `ScheduleAndRealizeSpatial` pass then owns scheduling, intermediate
   verification, communication realization, and final verification. Supporting
   scheduling code lives under `MergeComputeNodes/Scheduling`.

3. **Spatial -> Pim** (`src/PIM/Conversion/SpatialToPim`).
   Lowers Spatial operations to the `pim` dialect (`src/PIM/Dialect/Pim`),
   including `pim.core`, `pim.core_batch`, communication, tensor packing, global
   tensor materialization, and return-path normalization.

4. **Bufferization** (`src/PIM/Dialect/Pim/Passes/Transforms/Bufferization`).
   `PimBufferizationPreparation` establishes writable destinations without
   duplicating the one-shot copy analysis, `PimOneShotBufferization` runs
   MLIR's one-shot analysis,
   `PimMemoryNormalization` forwards/removes redundant copies and normalizes
   addressable accesses, and `PimBufferizationVerification` checks tensor
   absence, contiguity, and copy address spaces.

5. **Pim local-memory planning**
   (`src/PIM/Dialect/Pim/Passes/Transforms/LocalMemoryPlanning`).
   Computes whole-core lifetimes, reuses addresses for non-overlapping
   allocations, and records the explicit plan in Pim IR. Reusable lifetime
   analysis lives under `src/PIM/Dialect/Pim/Passes/Analyses`.
6. **Pim verification and code generation** (`src/PIM/Passes/PimCodegen` and
   `src/PIM/Compiler`).
   Verifies the memory plan and other Pim invariants, then emits `.pim` core
   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
  instruction format, artifact writing, weight emission, and codegen entry
  points.
- `src/PIM/Conversion/SpatialToGraphviz` - optional Spatial graphviz conversion
  pass.
- `src/PIM/Passes` - pass registration and auxiliary passes.
- `src/PIM/PimAccelerator.{cpp,hpp}` - ONNX-MLIR accelerator entry point.

## Pim compiler options

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. Default: no Pim accelerator.
- `--EmitSpatial`, `--EmitPim`, `--EmitPimBufferized`,
  `--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: `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`. 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: `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: `auto`.
- `--pim-conv-im2col-max-elements=<N>` - maximum globally materialized im2col
  elements per convolution before streaming. Default: `1048576`.
- `--pim-conv-stream-chunk-positions=<N>` - maximum output positions per
  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: off.
- `--pim-verify-bufferization-copy-freedom` - run the expensive official Pim
  tensor-copy freedom proof before bufferization. Default: off.

## Standard Pim hardware profile

Raptor's standard development and YOLO validation profile is:

| Parameter | Value |
| --- | ---: |
| Cores | 144 |
| Crossbars per core | 64 |
| Crossbar size | 128 × 128 |

Canonical compiler flags:

`--crossbar-count=64 --crossbar-size=128 --core-count=144`

`--crossbar-size`, `--crossbar-count`, and `--core-count` remain mandatory and
must be passed explicitly to the compiler.

Example:

```bash
./build_release/Release/bin/onnx-mlir model.onnx -o /tmp/raptor/model \
  --maccel=PIM --EmitPimCodegen \
  --crossbar-count=64 --crossbar-size=128 --core-count=144
```

This writes Pim artifacts under `/tmp/raptor/pim/`.

## Validation

Functional validation compiles ONNX models, compares native ONNX-MLIR and Pim
simulator outputs, and optionally reports latency, power, and energy. See
[`validation/README.md`](validation/README.md) for prerequisites, usage,
options, artifacts, and results.

## Build

Initialize submodules first:

```bash
git submodule update --init --recursive
```

The project follows ONNX-MLIR's build requirements. The CI workflow documents
the currently used versions and setup:
- CMake 4.3.0 in CI,
- LLVM/MLIR checked out under `onnx-mlir/llvm-project`,
- Protobuf `v34.0`,
- Rust stable for `pim-simulator`,
- Python packages `numpy`, `onnx`, `colorama` for validation.

### Protobuf

Install Protobuf if your system does not already provide a compatible version:

```bash
git clone --depth 1 --branch v34.0 https://github.com/protocolbuffers/protobuf
cmake -S protobuf -B protobuf/build -G Ninja \
  -DCMAKE_BUILD_TYPE=Release \
  -Dprotobuf_BUILD_TESTS=OFF
cmake --build protobuf/build
sudo cmake --install protobuf/build
```

You can then remove the temporary checkout:

```bash
rm -rf protobuf
```

### MLIR

Follow the ONNX-MLIR instructions in
`onnx-mlir/docs/BuildOnLinuxOSX.md` to build LLVM/MLIR. The local Raptor build
expects `MLIR_DIR` to point at the MLIR CMake package, for example:

```bash
MLIR_DIR=$(pwd)/onnx-mlir/llvm-project/build_release/lib/cmake/mlir
```

If your LLVM build directory is named `build` instead of `build_release`, adjust
the path accordingly.

### Raptor

Configure a release build:

```bash
MLIR_DIR=$(pwd)/onnx-mlir/llvm-project/build_release/lib/cmake/mlir
cmake -S . -B build_release -G Ninja \
  -DCMAKE_BUILD_TYPE=Release \
  -DONNX_MLIR_ACCELERATORS=PIM \
  -DLLVM_ENABLE_ASSERTIONS=ON \
  -DMLIR_DIR=${MLIR_DIR}
```

Configure a debug build similarly:

```bash
MLIR_DIR=$(pwd)/onnx-mlir/llvm-project/build_debug/lib/cmake/mlir
cmake -S . -B build_debug -G Ninja \
  -DCMAKE_BUILD_TYPE=Debug \
  -DONNX_MLIR_ACCELERATORS=PIM \
  -DLLVM_ENABLE_ASSERTIONS=ON \
  -DMLIR_DIR=${MLIR_DIR}
```

For debug development, using `mold` can reduce link time and memory use:

```bash
cmake -S . -B build_debug -G Ninja \
  -DCMAKE_BUILD_TYPE=Debug \
  -DONNX_MLIR_ACCELERATORS=PIM \
  -DLLVM_ENABLE_ASSERTIONS=ON \
  -DMLIR_DIR=${MLIR_DIR} \
  -DCMAKE_EXE_LINKER_FLAGS="-fuse-ld=mold" \
  -DCMAKE_SHARED_LINKER_FLAGS="-fuse-ld=mold" \
  -DCMAKE_MODULE_LINKER_FLAGS="-fuse-ld=mold"
```

Build the compiler with CMake:

```bash
cmake --build ./build_release
cmake --build ./build_debug
```

Do not invoke `ninja` directly for this project; use `cmake --build` so CMake's
configuration and generated shims stay consistent.

If a build fails because Protobuf headers are missing fixed-width integer
definitions, patch the affected Protobuf-generated files by adding
`#include <cstdint>`.

## Tests

The Rust simulator has its own tests:

```bash
cd backend-simulators/pim/pim-simulator
cargo test
```

## Repository layout

- `src/PIM/` - PIM accelerator implementation.
- `test/PIM/` - PIM C++ unit tests.
- `validation/` - functional validation scripts, ONNX operation tests, network
  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.
- `.github/actions/` and `.github/workflows/validate_operations.yml` - CI setup
  for MLIR/Protobuf caching, building Raptor, and validation.