tensorforge

A mini ONNX inference compiler with a graph IR, an optimizing pass pipeline, and a C++/CUDA runtime. It lowers an ONNX model to a small SSA dataflow IR, runs optimization passes (constant folding, DCE, and -- on the roadmap -- operator fusion and liveness-based memory planning), emits an ExecutionPlan, and executes it through a CUDA runtime that calls cuBLAS for heavy GEMMs and generated kernels for fused glue.

The runtime exposes the same inference interface as the distributed-inference-engine-cpp worker (predict / batchPredict), so it drops in as an alternate backend and can be A/B-benchmarked against ONNX Runtime through that project's tail-latency harness.

ONNX Runtime is the correctness reference and parity baseline -- not a strawman to beat. The contribution is graph-level optimization + fused codegen + memory planning, benchmarked honestly. See design.md for the full spec and MILESTONES.md for status.

Status

Milestone	Scope	State
M0	Scaffold, model export, toolchain	done
M1	Walking skeleton: IR to passes to plan to CUDA runtime, unfused, ORT-green on MLP	done
M2	Elementwise + GEMM-epilogue fusion, NVRTC codegen, launch-count bench	done
M3	Conv (im2col) + BN + pooling runtime, liveness memory planning, conv-epilogue fusion (§6.1) + BN fold, ResNet50-v2 ORT-green	done
M4	Drop-in backend of the distributed inference engine, tail-latency A/B	done

Correctness gate (vs ORT): max_abs_err < 1.2e-6 on mlp.onnx / ew.onnx, 4.8e-6 on resnet50-v2 (fused and unfused; fusion changes the output by 0.0 bit-for-bit).

Fusion on/off (same runtime, bench/launch_count.py):

model	ops	kernel launches	arena bytes
mlp.onnx	5 -> 3	8 -> 6 (1.33x)	2560 -> 1536 (1.67x)
ew.onnx	6 -> 2	7 -> 3 (2.33x)	6400 -> 2304 (2.78x)
resnet50-v2	174 -> 91	194 -> 144 (1.35x)	~12 MB

ResNet fusion: conv-epilogue fusion (§6.1) folds Conv -> BN -> Relu into a single FusedConvEpilogue -- BatchNorm's inference affine is folded into the conv weights offline (W' = scale*W, bias = shift), so the runtime does conv + bias + activation with no separate BN op or buffer (ops 174 -> 91, weights 260 -> 161). Remaining pre-activation BN -> Relu pairs (not preceded by a fusible conv) fuse into one kernel. Launch count is neutral because the bias+activation epilogue replaces the BN kernel 1:1.

A cuBLASLt fused-GEMM epilogue (bias+ReLU inside the matmul) was evaluated to remove that last kernel: it is numerically correct but slower for this batch-1 workload -- per-call descriptor setup + algo heuristic outweigh the saved elementwise kernel on small per-image GEMMs -- so it was dropped in favour of the fused conv_bias_act kernel, which is exactly the "single fused kernel" §6.1 allows. It would likely pay off at larger batch, with cached descriptors.

Memory planning (liveness + interval coloring, bench/memory.py), peak arena bytes planned vs naive per-tensor:

model	planned	naive	reduction
mlp.onnx	1.25 KB	1.50 KB	1.20x
resnet50-v2	11.1 MB	140.2 MB	12.63x

Layout

compiler/          Python: offline compilation
  ir.py            IR dataclasses + verify()
  importer.py      ONNX -> IR
  shape_infer.py   shape/dtype propagation
  passes/          const_fold, dce, fusion, memory_plan, PassManager
  codegen/         plan.json + weights.bin + generated CUDA (cuda_emit)
  compile.py       driver: model.onnx -> build/  (--no-fusion toggles fusion)
runtime/           C++17 / CUDA
  include/         public headers (compiled_engine.h matches the DIE interface)
  src/             plan_loader, buffer_arena, compiled_engine, nvrtc_jit, kernels/
  external/        vendored nlohmann/json (single header)
bench/             correctness.py (ORT gate), launch_count.py, memory.py, latency.py
models/            export_models.py -> mlp.onnx, ew.onnx
tests/             pytest golden tests for the passes

Build and run

Toolchain: CUDA 13.x (/opt/cuda), a C++17 compiler, CMake >= 3.18, and a Python 3.12 venv (ONNX Runtime has no 3.14 wheels). Target device: RTX 4060 Laptop (Ada, sm_89).

# one-time: python deps
python3.12 -m venv .venv && . .venv/bin/activate
pip install numpy onnx onnxruntime pytest

# export the model, compile it to build/ artifacts
python models/export_models.py
python -m compiler.compile models/mlp.onnx -o build

# build the CUDA runtime (ensure /opt/cuda/bin is on PATH)
cmake -S runtime -B runtime/build -DCMAKE_BUILD_TYPE=Release
cmake --build runtime/build -j

# correctness gate vs ONNX Runtime + unit tests
python bench/correctness.py models/mlp.onnx
python -m pytest tests/ -q

Artifacts

compile.py writes a build/ directory:

• plan.json -- the ExecutionPlan (arena size, buffers, weights table, op schedule).
• weights.bin -- all constants concatenated; the plan references them by offset.

The runtime loads both, allocates the device arena once, and runs the op schedule with no per-inference device allocation on the steady-state path.

Drop-in backend of the distributed inference engine (M4)

CompiledEngine matches the worker's batchPredict interface, so it substitutes for ONNX Runtime behind an ENGINE_BACKEND=ort|compiled switch in distributed-inference-engine-cpp. Build the worker with -DWITH_COMPILED_BACKEND=ON and point TF_BUILD_DIR at a tensorforge build dir (plan.json + weights.bin):

python -m compiler.compile models/resnet50.onnx -o dist/resnet50
# in the DIE repo:
cmake -S . -B build_tf -DWITH_COMPILED_BACKEND=ON && cmake --build build_tf -j
ENGINE_BACKEND=compiled TF_BUILD_DIR=.../dist/resnet50 \
  ./build_tf/worker_node 8001 worker_1 models/resnet50-v2-7.onnx

Both backends pass DIE's full suite (diagnostics, benchmark + cache test, tail_latency_sweep.py). A/B on the unique-input inference path (3 workers + gateway, ResNet50-v2, RTX 4060, 300 req/level):

concurrency	compiled tput/s	compiled p50/p99 ms	ORT-CUDA tput/s	ORT-CUDA p50/p99 ms
8	91.6	82 / 140	54.6	156 / 255
16	162.2	92 / 167	95.6	165 / 259
32	194.2	144 / 306	154.5	193 / 281

Honest framing: a small single-machine sweep with a Python client in the loop and zero-padded probe inputs. The compiled engine leads on throughput and median here; ORT's tuned conv kernels close the gap on the p99 tail at higher concurrency. Not a blanket win -- the point is drop-in parity, measured honestly.

Engine latency (isolated from the client)

The HTTP sweep above is client-bound (a single Python process caps it ~500 req/s), so it can't measure true per-inference latency. bench/latency.py (via the tf_bench binary) loads the engine once and times predict() on the GPU:

model	fp32 p50	fp32 infer/s	TF32 p50	TF32 infer/s
resnet50-v2	4.12 ms	242	3.06 ms	324
mlp	40 us	22.7k	40 us	21k

Latency work (all correctness-preserving except TF32): a 1x1-stride-1 conv fast path (GEMM directly on the input, skipping an identity im2col), BatchNorm+Relu fusion, and dropping a redundant device sync cut ResNet kernel launches 227 -> 144 and fp32 latency ~14%. TF32 tensor cores (TF_TF32=1) add ~1.35x but cost ~1e-3 relative error that accumulates past the correctness gate over 50 conv layers, so it is opt-in, off by default -- the gate stays green.