Hexir

Heterogeneous MLIR Compiler for Neural Networks

One graph in — partitioned, lowered, and executed across CPU and GPU.

Documentation · Getting started · Architecture

---

Hexir (Heterogeneous EXecution IR) is an MLIR compiler with two custom dialects: a graph level of neural-network operators on tensors, and a kernel level where each loop carries its own schedule, so a scheduling decision is written in the IR instead of being implied by which pass ran.

A partitioning pass assigns each operator to CPU or CUDA before lowering and propagates that placement through every conversion — linalg → bufferization → SCF → LLVM IR on the host, GPU dialect → NVVM → CUBIN on the device.

The output is a self-describing artifact with the device code embedded, run by a standalone C99 runtime that links no MLIR or LLVM. On a GTX 1660 Ti a GPU matmul reaches 47–68 GFLOP/s in f64 against cuBLAS at 123–179, checked against a CPU reference.

It is small enough to read. That is deliberate: the gap between MLIR's Toy tutorial and a production compiler like IREE is enormous, and Hexir sits in the middle — a complete pipeline, in about five thousand lines.

%m = hexir.linear %a, %w : tensor<2x2xf64>   // placed on the GPU
%r = hexir.relu %m : tensor<2x2xf64>         // placed on the CPU
hexir.print %r : tensor<2x2xf64>

Compile, then run

hexir -emit=hxb -o model.hxb                 # compile to a file
hexir-run model.hxb                          # run it later, no compiler present

There is one way to run a program: compile it to a file, then run the file. hexir-run links no MLIR and no LLVM. A GPU-placed kernel is compiled to a CUBIN and embedded in the file, so the runtime loads and launches it with no compiler in the process:

$ hexir -emit=hxb -o gpu.hxb -placement=hexir.linear=cuda
$ hexir-run --device=cuda gpu.hxb
device        : cuda (NVIDIA GeForce GTX 1660 Ti)
--
8.000000 17.000000
12.000000 14.000000

Identical to the CPU answer, which is the point.

How it works

Five levels, each answering a different question about the same program.

      hexir dialect        what to compute
            │
      partitioning         where each operation runs   (device = cpu | cuda)
            │
      hextir dialect       how one device computes it  (loops, buffers)
            │
   linalg / memref / scf   the actual loop nests
            │
      LLVM  /  NVVM        machine code

Placement is decided before lowering and travels down with the operation, so by the time GPU code generation runs, the ops that should become kernels are already labelled. At the kernel level it stops being an annotation and becomes code: a CPU op gets parallel loops, a GPU op gets loops bound to blockIdx and threadIdx.

Every stage can be printed, which is the most useful thing about the project for learning:

Command	Shows
-emit=mlir	the graph you wrote
-emit=mlir-tir	each operation as a kernel
-emit=mlir-linalg	loops on tensors
-emit=mlir-hetero	placement, as device attributes
-emit=mlir-gpu	CUDA kernels as gpu.launch
-emit=llvm	LLVM IR
-emit=hxb	a deployable module

Move an operation between devices without recompiling the compiler:

hexir -emit=mlir-tir -placement=hexir.linear=cuda,hexir.relu=cpu

Building

Needs a build of LLVM with MLIR.

git clone git@github.com:hamzaqureshi5/hexir.git && cd hexir
mkdir -p build && cd build
cmake .. -DCMAKE_BUILD_TYPE=Release \
         -DLLVM_DIR=/path/to/llvm-project/build/lib/cmake/llvm \
         -DMLIR_DIR=/path/to/llvm-project/build/lib/cmake/mlir
make -j$(nproc)

For the GPU path, LLVM must be built with -DLLVM_TARGETS_TO_BUILD="X86;NVPTX", and a CUDA toolkit must be installed. Full commands are in docs/llvm-cuda-build.txt.

Testing

cd build && make check-hexir              # everything
cd build && make check-hexir-compiler     # compiler only
cd build && make check-hexir-runtime      # runtime only

The compiler and the runtime each own their tests, in compiler/test/ and runtime/test/.

The runtime suite also runs standalone, against a module checked in at runtime/test/fixtures/, so it needs no MLIR:

cmake -S runtime -B build-runtime && cmake --build build-runtime
lit runtime/test --param build_dir=$PWD/build-runtime

That is what CI runs. The compiler is not covered yet: it needs an unreleased LLVM, so no packaged MLIR can build it.

Contributing

Fork, branch, keep make check-hexir green, open a pull request. New tests should compile their own .mlir input and pass -placement explicitly rather than relying on defaults.

License

Apache 2.0.

Acknowledgments

Built on MLIR. The dialect scaffolding started from the MLIR Toy tutorial.

---

_{Research software under active development — APIs and IR may change.}