296 lines
11 KiB
Markdown
296 lines
11 KiB
Markdown
# 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` 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
|
||
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
|
||
`SpatialTargetInfo`. 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.
|
||
|
||
2. **Merge, schedule, and realize Spatial communication**
|
||
(`src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes`).
|
||
`TrivialGraphComputeMerge` performs local graph merging, then
|
||
`ScheduleSpatialGraph` materializes scheduled computes and explicit deferred
|
||
communication. `VerifyScheduledSpatial` checks that intermediate contract;
|
||
`RealizeSpatialCommunication` resolves transfers and forwarding; and
|
||
`VerifyRealizedSpatial` checks the final scheduled graph. Supporting
|
||
scheduling code lives under `MergeComputeNodes/Scheduling`.
|
||
|
||
3. **Spatial -> Pim** (`src/PIM/Conversion/SpatialToPim`).
|
||
Lowers Spatial operations to the `pim` dialect (`src/PIM/Dialect/Pim`),
|
||
including `pim.core`, `pim.core_batch`, communication, tensor packing, global
|
||
tensor materialization, and return-path normalization.
|
||
|
||
4. **Bufferization** (`src/PIM/Dialect/Pim/Transforms/Bufferization`).
|
||
`PimBufferizationPreparation` establishes writable destinations without
|
||
duplicating the one-shot copy analysis, `PimOneShotBufferization` runs
|
||
MLIR's one-shot analysis,
|
||
`PimMemoryNormalization` forwards/removes redundant copies and normalizes
|
||
addressable accesses, and `PimBufferizationVerification` checks tensor
|
||
absence, contiguity, and copy address spaces.
|
||
|
||
5. **PIM local-memory planning**
|
||
(`src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning`).
|
||
Computes whole-core lifetimes, reuses addresses for non-overlapping
|
||
allocations, and records the explicit plan in PIM IR.
|
||
6. **PIM verification and code generation** (`src/PIM/Pass/PimCodegen` and
|
||
`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/Pass` - 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.
|
||
- `--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
|
||
target adapter to construct the target-neutral Spatial scheduling cost and
|
||
topology model. Resource values must match the explicit core/crossbar flags.
|
||
- `--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.
|
||
- `--pim-emit-json` - also emit `core_*.json` instruction files alongside
|
||
`core_*.pim`.
|
||
- `--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`.
|
||
- `--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`.
|
||
- `--pim-conv-im2col-max-elements=<N>` - maximum globally materialized im2col
|
||
elements per convolution before streaming. Default is `1048576`.
|
||
- `--pim-conv-stream-chunk-positions=<N>` - maximum output positions per
|
||
streamed convolution chunk. Default is `1024`.
|
||
- `--use-experimental-conv-impl` - use the alternate convolution lowering.
|
||
- `--pim-detect-communication-deadlock` - statically simulate expanded
|
||
send/receive ordering and reject blocking deadlocks. Default is off.
|
||
- `--pim-materialize-scalar-fanout-global-order` - use the experimental,
|
||
expensive globally ordered scalar-fanout materializer. Default is off.
|
||
- `--pim-trace-communication-materialization` - emit verbose communication
|
||
materialization diagnostics and provenance attributes. Default is off.
|
||
- `--ignore-concat-error` - soft-fail a ConcatOp corner case.
|
||
|
||
## 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`
|
||
|
||
`--core-count` remains 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-NN.
|
||
- `.github/actions/` and `.github/workflows/validate_operations.yml` - CI setup
|
||
for MLIR/Protobuf caching, building Raptor, and validation.
|