# 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. 2. **Merge compute nodes** (`src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes`). Builds a compute graph, schedules it with the PEFT scheduler, and materializes the merge schedule into Spatial IR. 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`). Converts tensor-semantics PIM IR into memref-semantics PIM IR using MLIR's bufferization interfaces. 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=` - required positive core count for PIM compilation. - `--crossbar-size=` - crossbar width/height. Default in code is `128`. - `--crossbar-count=` - crossbars per core. Default in code is `64`. - `--pim-memory-report=` - 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=` - control Spatial dataflow CSV reports for the graph, trivially merged graph, scheduled, and realized snapshots under `reports/`. Default is `none`. - `--pim-conv-lowering=` - select the convolution lowering strategy. Default is `auto`. - `--pim-conv-im2col-max-elements=` - maximum globally materialized im2col elements per convolution before streaming. Default is `1048576`. - `--pim-conv-stream-chunk-positions=` - maximum output positions per streamed convolution chunk. Default is `1024`. - `--pim-report-conv-lowering=` - emit the bounded convolution lowering report. Default is `true`. - `--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 and power. 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 `. ## 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.