Kernel Tile IR (KTIR) Frontend (ktir_mlir_frontend)

KTIR is a tile-based, block-structured intermediate representation (IR) designed to express programs targeting multi-core accelerator architectures. It embodies a data-parallel abstraction, where the accelerator contains multiple cores, with each core comprising a compute engine and an on-chip scratchpad memory associated with the compute unit. The cores are attached together through an on-chip interconnect fabric, which also interfaces with one or more off-chip memory banks.

KTIR programs allow tensors to be placed in a distributed fashion across multiple memory elements (on-chip and/or off-chip memory). In the same vein, compute operations can be split into compute tiles and assigned to the compute engine within each core. KTIR allows each compute tile to have global view of all on-chip and off-chip memory elements via the on-chip interconnect fabric.

KTIR is built to represent complex computation kernels (e.g., attention in LLMs) mapped onto multi-core accelerators, going beyond a single operation (e.g., tensor Add). Kernels span multiple operations that are composed sequentially, interleaved with structured control-flow such as loops, conditionals, and multi-stage pipelines. KTIR includes constructs to capture kernel execution semantics including dependencies, synchronization points, and tensor liveliness. Intermediate tensors within the kernel could be placed on on-chip memory elements and reused across producer-consumer operations. With each compute tile having global access to all memory elements, individual compute operations can be flexibly tiled along different dimensions.

More details about KTIR can be found in the RFC: KTIR (KernelTileIR)

Dialect Overview

The ktdp dialect models tile-based, data-parallel kernels targeting multi-core accelerator architectures (e.g., Spyre). It provides:

• Memory views: construct_memory_view, construct_distributed_memory_view
• Access tiles: construct_access_tile (direct), construct_indirect_access_tile (gather/scatter)
• Data movement: load, store (driven by access tiles)
• Tile identity: get_compute_tile_id
• Runtime args: runtime_arg type + runtime_arg_extract op
• Inter-tile communication: inter_tile_produce + inter_tile_reduce (all-reduce; yield_partial / yield_reduced terminators). Broadcast (inter_tile_consume), reduce-scatter, per-tile sync, gather, and scatter are future work.

Custom types: !ktdp.access_tile<NxMxindex>, !ktdp.runtime_arg<type, granularity=N, upperbound=M>, !ktdp.tile_future<T_p_1, ...>

Memory spaces: #ktdp.memory_space<global>, #ktdp.memory_space<ct_local>, #ktdp.memory_space<ct_local, ct_id=7>

Building attributes and types from Python

The bindings expose builders for the custom types and attributes, so IR is constructed through verifier-checked APIs rather than assembled as text:

import mlir_ktdp.dialects.ktdp as ktdp

# #ktdp.memory_space<global>
ktdp.MemorySpaceAttr.get(ktdp.MemorySpaceKind.global_)

# #ktdp.memory_space<ct_local, ct_id = 7>
ktdp.MemorySpaceAttr.get(ktdp.MemorySpaceKind.ct_local, ct_id=7)

kind takes (and reads back) the tablegen-generated MemorySpaceKind enum — note the trailing underscore on global_, since global is a Python keyword. ct_id is optional and only valid for ct_local; an invalid combination raises ir.MLIRError carrying the verifier diagnostic.

Like the upstream MLIR type builders, get is verified and needs a Location, taken from the surrounding with Location... or an explicit loc=. Diagnostics are emitted there, so a rejected attribute points at the code that built it.

How to Build

Prerequisites

• CMake >= 3.20, Ninja, C++17 compiler
• uv (for Python environment)
• Python >= 3.12 (only for Python bindings)

CMake (C++ only)

cmake -S . -B build -GNinja -DMLIR_DIR=/path/to/llvm-build/lib/cmake/mlir
cmake --build build -j$(nproc)

The following CMake variables can be configured:

Name	Type	Description
MLIR_DIR (required)	STRING	Path to the CMake directory of an MLIR installation. e.g. ~/tools/llvm-12.0.1/lib/cmake/mlir
LLVM_EXTERNAL_LIT (optional)	STRING	Path to a lit executable, required for testing.
KTIR_ENABLE_PYTHON_BINDINGS (default: OFF)	BOOL	Whether to build the Python bindings.
KTIR_BUILD_TOOLS (default: ON)	BOOL	Whether to build the tool executables.
MLIR_LINK_MLIR_DYLIB (optional)	BOOL	Whether to link against a shared MLIR library.

Linking against shared LLVM/MLIR

This project supports linking against a version of LLVM/MLIR that was build using the LLVM_BUILD_LLVM_DYLIB etc. flags. Since MLIR does not export the MLIR_LINK_MLIR_DYLIB flag, it is inferred from the presence of the MLIR library target when it is not specified at configure time. When linking against shared MLIR, CMake will set up an RPATH to your MLIR install destination.

CMake with Python bindings

Note: The recommended way to build Python bindings is via uv sync, see the next section.

Python bindings are disabled by default. To enable them, add -DKTIR_ENABLE_PYTHON_BINDINGS=ON and ensure nanobind is installed in the active Python environment:

cmake -S . -B build -GNinja \
    -DMLIR_DIR=/path/to/llvm-build/lib/cmake/mlir \
    -DKTIR_ENABLE_PYTHON_BINDINGS=ON \
    -DPython3_EXECUTABLE=$(which python3) \
    -DLLVM_EXTERNAL_LIT=$(which lit)
cmake --build build -j$(nproc)

The built package lands at build/python_packages/ktdp/ — use PYTHONPATH=build/python_packages/ktdp to import mlir_ktdp.

Note: Add -DMLIR_PYTHON_STUBGEN_ENABLED=ON to generate .pyi type stubs for IDE support. This is set automatically when building via uv pip install (see below).

Note: -DLLVM_EXTERNAL_LIT is only needed to run LIT tests. mlir_wheel does not ship llvm-lit, so cmake cannot find it automatically.

pip install (scikit-build-core)

The project uses scikit-build-core to drive CMake and produce a wheel. First, create a venv:

uv venv --python 3.12

Use scripts/setup_mlir.py to obtain the MLIR installation, then build:

# Install Python test dependencies (pytest and lit)
uv sync --no-install-project --extra test

# Resolve MLIR.
# Skip this if you already have a local LLVM/MLIR build — just set MLIR_DIR
# directly: export MLIR_DIR=/path/to/llvm-build/lib/cmake/mlir
#
# Otherwise, setup_mlir.py downloads the pinned artifact (cached at
# ~/.cache/ktir-mlir/) or falls back to mlir_wheel.
# A GIT_PAT or GITHUB_TOKEN is required on the first download; subsequent
# runs use the cache and need no token.
MLIR_DIR=$(uv run --no-project python scripts/setup_mlir.py)

# Force mlir_wheel (no token required, always fetches latest):
MLIR_DIR=$(uv run --no-project python scripts/setup_mlir.py --wheel)

# Build and install the Python wheel (stubgen runs automatically for IDE stubs)
CMAKE_ARGS="-DMLIR_DIR=$MLIR_DIR" uv pip install .

See docs/ci.md for the full local development workflow including cache behaviour and fork setup.

Running Tests

LIT tests (C++ / IR)

cmake --build build --target check-ktir

Or directly — lit must be pointed at the build directory where cmake generates lit.site.cfg.py:

llvm-lit -sv build/test/Ktdp/
# or, with uv (lit is installed with --extra test):
uv run lit -sv build/test/Ktdp/

Python tests

uv run pytest python/test/

For bare cmake builds (no pip install), conftest.py adds build/python_packages/ktdp to sys.path automatically.

Python Tools

ktdp-walk

Parse an MLIR file and print a recursive walk of all operations with indentation showing nesting depth:

uv run ktdp-walk path/to/file.mlir
# or from stdin:
cat file.mlir | uv run ktdp-walk

ktdp-walk is also usable as a library:

from tools_ktdp import ktdp_context
from tools_ktdp.ir_utils import walk_module

# context manager
with ktdp_context() as ctx:
    ...

# walk returns [(op, depth), ...]
for op, depth in walk_module(source):
    print(f"{'  ' * depth}{op.name}")

Use ktir-opt

build/bin/ktir-opt your_file.mlir
build/bin/ktir-opt --show-dialects
build/bin/ktir-opt --verify-roundtrip file.mlir

Project Structure

include/
  Ktdp/                    # TableGen definitions + C++ headers
lib/
  Ktdp/                    # Dialect implementation
tools/
  ktir-lsp-server/         # MLIR LSP server
  ktir-opt/                # Optimizer driver tool
test/
  Ktdp/                    # 17 LIT test files