npudialect

An out-of-tree MLIR dialect for an NPU with a software-managed scratchpad.

A GPU dialect can leave data movement implicit, because the hardware has a cache. An NPU with a software-managed scratchpad cannot: nothing moves unless the compiler emits a transfer, and the scheduling problem that dominates real NPU backends is overlapping those transfers with compute.

So this models data movement explicitly, and the three passes here are the three decisions a backend actually has to make: what to keep resident, how big a tile to use, and when to issue the next transfer.

#sp = #npu.mem_space<scratchpad>

func.func @tile(%src: memref<64x64xf32>) {
  %c0 = arith.constant 0 : index
  %a = npu.dma_load %src[%c0, %c0] {channel = 0 : i64}
     : memref<64x64xf32> -> memref<16x16xf32, #sp>
  npu.wait 0
  npu.matmul %a, %b, %c : memref<16x16xf32, #sp>, ...
  return
}

Building

Needs LLVM/MLIR 18. On Ubuntu:

sudo apt install llvm-18-dev libmlir-18-dev mlir-18-tools clang-18
./build.sh                      # configures and builds
cd ~/build/npudialect && ninja check-npu

-- Testing: 3 tests, 3 workers --
Total Discovered Tests: 3
  Passed: 3 (100.00%)

The build directory deliberately defaults outside the source tree. If you are on WSL with the source on a Windows drive, /mnt/... is a 9p mount and linking MLIR through it is several times slower than linking on the native filesystem.

The dialect

Six operations and one attribute, which is as small as it can be and still express the problem.

npu.dma_load	DRAM to scratchpad. Returns a new value, so the dependency rides on the use-def edge rather than needing alias analysis.
npu.dma_store	Scratchpad back to DRAM.
npu.wait	Block on a DMA channel.
npu.matmul	Tile-level matrix multiply. Verifies its operands are in the scratchpad.
npu.relu	Elementwise, so there is something worth fusing.
#npu.mem_space	dram or scratchpad.

Memory space is an enum, not a string

Worth pointing at, because the first version got it wrong. With a free-form string parameter, #npu.mem_space<"scratchpd"> parses cleanly and means "not the scratchpad" — a typo becomes a silent performance bug that the verifier cannot tell apart from a deliberate DRAM operand.

As a closed I32EnumAttr, the same typo is a parse error at the point it was written. Same syntax, and the failure moved from runtime to compile time.

The verifiers are the point

An NPU program that reads DRAM where it should read the scratchpad still computes the right answer. It just does it at DRAM speed, and no functional test will ever notice. Those are the failures encoded here:

npu.matmul' op operand 'lhs' is not in the scratchpad (memory space '<none>').
  Stage it with npu.dma_load first.

npu.dma_load' op result must live in the scratchpad, but its memory space is
  '<none>'. Give it #npu.mem_space<scratchpad>.

npu.dma_load' op tile is larger than the source in dimension 0: 16 > 8

npu.dma_load' op element type changes across a DMA: 'f32' to 'i8'.
  A DMA moves bytes; it does not convert them.

npu.matmul' op inner dimensions disagree: 16 and 32

Each is a real diagnostic from test/invalid.mlir, and each describes a program that is structurally valid MLIR and wrong for this hardware.

The passes

-npu-fuse-elementwise

Fuses npu.relu into the npu.matmul that produced its input.

The saving is not arithmetic — a ReLU is nearly free — it is a round trip through memory. Unfused, the matmul writes its accumulator out and the ReLU reads it back; on a scratchpad machine that traffic is the dominant cost, and it is invisible in an operation count.

// before
npu.matmul %a, %b, %c : ...
%r = npu.relu %c : memref<16x16xf32, #sp>

// after
npu.matmul %a, %b, %c {fused_relu} : ...

Only fires when the ReLU is the accumulator's sole other consumer. If anything else reads the raw value it still needs the unfused result, and fusing would silently change what that reader sees.

-npu-tile-matmul

Chooses the tile size that fits a stated scratchpad budget.

// -npu-tile-matmul=scratchpad-bytes=65536, on a 512x512 f32 matmul
npu.matmul %a, %b, %c {tile_bytes = 49152 : i64, tile_size = 64 : i64} : ...

3 × 64² × 4 bytes = 49,152, which fits 65,536. A tile of 128 would need 196,608 and does not.

Two deliberate choices. Powers of two, because DMA engines and tile registers are sized that way, and a tile of 47 would be padded to 64 in hardware while the compiler believed it had saved memory. And failure rather than clamping when nothing fits: silently using a tile that overflows the scratchpad is precisely the bug this dialect exists to catch.

It annotates rather than generating loops, which keeps the decision visible in the IR — a wrong tile size shows up in a dump instead of as an unexplained regression. Loop emission consumes these attributes and is not implemented; see Limitations.

-npu-double-buffer

Hoists each npu.dma_load above the npu.wait before it, so the next transfer overlaps the current compute.

// before                          // after
%t0 = npu.dma_load ... channel 0   %t0 = npu.dma_load ... channel 0
npu.wait 0                         %t1 = npu.dma_load ... channel 1
%t1 = npu.dma_load ... channel 1   npu.wait 0
npu.wait 1                         npu.wait 1

This is why npu.wait is a separate operation rather than implied by data flow: a model where waiting is implicit cannot express "issue, do other work, then wait", which is the one optimisation that matters most on this hardware.

It refuses to move a load whose operands are defined after the wait, and refuses to move a load on the very channel being waited for. Without those checks it would be a reordering that happens to work on the test input.

Limitations

• The simulator is not written. There is no execution, so the passes are verified by their effect on the IR rather than by a measured speedup. The double-buffering pass produces the right schedule; nothing here proves what that schedule would be worth on silicon.
• Tiling annotates, it does not generate loops. Emitting the scf.for nest and the per-tile DMA from these attributes is the obvious next step.
• No lowering to LLVM. This is a middle-end dialect; there is no backend.
• Fusion handles one pattern. ReLU into matmul, single-consumer only.
• The cost model is a capacity check. It answers "does this tile fit", not "is this tile fastest", which would need a bandwidth and latency model.
• Targets LLVM 18. LLVM 15 was tried first and does not work: it lacks SameOperandsAndResultType as a TableGen record, has no BytecodeOpInterface.h, calls the dialect library MLIRArithmeticDialect, and requires the older member-function cast spelling.

Layout

include/NPU/NPUOps.td      the dialect, in TableGen
lib/NPUDialect.cpp         verifiers and memory-space helpers
lib/NPUPasses.cpp          the three passes
tools/npu-opt.cpp          mlir-opt with this dialect registered
test/*.mlir                lit tests: roundtrip, verifiers, passes

Related reading

• Lattner et al., MLIR: Scaling Compiler Infrastructure for Domain Specific Computation, CGO 2021.
• Chen et al., TVM: An Automated End-to-End Optimizing Compiler for Deep Learning, OSDI 2018.
• Jouppi et al., In-Datacenter Performance Analysis of a Tensor Processing Unit, ISCA 2017. The source for treating the scratchpad, not the ALU, as the thing the compiler is really scheduling.

License

MIT. See LICENSE.