Beyond SPMD: Compiling Triton Kernels for Dataflow AI Architectures

Triton compiles for GPUs, assuming SIMT parallelism over a flat device memory hierarchy. Dataflow AI architectures share tile-based abstractions but differ: many spatial cores, a small per-core scratchpad, and a core-to-core communication fabric instead of shared device memory.

We generate Triton kernels for dataflow hardware, keeping the PyTorch frontend and operator dispatch unchanged. At its center is KTIR (Kernel Tile IR), a new IR derived from Triton IR that abstracts the dataflow architecture. Triton cannot express core-to-core transfer, so we extend it with inter-tile communication APIs that lower to KTIR, and implement optimizations in torch.inductor.

This poster shows optimizations such as:

(1) Tiling that splits each operation into loops over tiles sized to fit the per-core scratchpad, keeping working sets on-core instead of spilling to device memory.

(2) Fusing operations with different iteration spaces, or work divisions, by inserting inter-tile communication that moves data directly core-to-core — avoiding device-memory round trips.

Visitors can see the generated kernels and codegen.