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:

./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 for prerequisites, usage, options, artifacts, and results.

Build

Initialize submodules first:

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:

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:

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:

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:

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:

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:

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:

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:

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.