In Part 1: Symbolic Tracing, mytriton learned how to trace a Python function into an expression-tree IR.

In Part 2: Typed SSA, it learned how to infer types and lower that tree into typed SSA-style operations.

In Part 3: CUDA Lowering, it learned how to turn SSA into CUDA C++ and launch the generated kernel.

In Part 4: Elementwise Ops, the language grew enough elementwise operations to write ReLU, leaky ReLU, and sigmoid.

In Part 5: Verification, the middle of the compiler became stricter: SSA is now verified, optimized, and verified again before code generation.

In Part 6: Reductions, vectors learned how to cooperate inside one CUDA block: row-wise sum, max, min, softmax, and a first naive matmul.

Version 7 does not add a new user-facing language feature.

Instead, it asks a different question: can the same optimized SSA feed more than one backend?

Until now, mytriton had one real target:

SSA -> CUDA C++ -> CuPy RawKernel

That path is still the default. But Version 7 adds a second, experimental backend:

SSA -> MLIR GPU dialect -> NVVM -> cubin -> CuPy Module

The code for this milestone is here: https://github.com/pbelevich/mytriton/tree/ver7.

This version is intentionally smaller than Version 6 in terms of language coverage. The CUDA backend can still lower reductions, softmax, tl.where, tl.exp, tl.static_range, two-dimensional program IDs, and the naive matmul kernel.

The first MLIR backend cannot.

It supports only a narrow one-dimensional elementwise subset: tl.program_id(0), tl.arange(0, BLOCK), basic arithmetic, <, pointer addition, masked loads, and masked stores over ptr<f32>.

That sounds like a step backward until you look at what changed architecturally. The compiler is no longer only a CUDA string generator. It now has a real backend boundary.

Selecting a backend

The public API barely changes. A launch still returns:

expression_ops, ssa_ops, src = add_kernel[grid](
    x,
    y,
    out,
    n,
    BLOCK=256,
)

The difference is that src no longer means “CUDA source.” It means “the source generated by the selected backend.”

The default backend is still CUDA:

MYTRITON_BACKEND=cuda python my_script.py

The new backend is selected with:

MYTRITON_BACKEND=mlir python my_script.py

The compiler validates the environment variable early:

Backend: TypeAlias = Literal["cuda", "mlir"]


def _resolve_backend(self) -> Backend:
    selected = os.environ.get("MYTRITON_BACKEND", "cuda")
    if selected not in {"cuda", "mlir"}:
        raise ValueError(
            f"unsupported backend {selected!r}; expected 'cuda' or 'mlir'"
        )
    return cast(Backend, selected)

That is not fancy. It is just enough to make the rest of the compiler stop assuming that the generated source is CUDA C++.

The compilation artifact now stores src:

@dataclass(frozen=True)
class _CompilationArtifact:
    ops: tuple[Any, ...]
    ssa_ops: tuple[SSAOp, ...]
    src: str
    threads_per_block: int
    cubin: bytes | None = None

For CUDA, src is CUDA C++.

For MLIR, src is MLIR text.

The optional cubin is only meaningful for the MLIR execution path. More on that later.

The same kernel, the same SSA

The main test kernel is still vector addition:

@triton.jit
def add_kernel(x, y, out, n, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)

    mask = offs < n

    a = tl.load(x + offs, mask=mask, other=0.0)
    b = tl.load(y + offs, mask=mask, other=0.0)

    tl.store(out + offs, a + b, mask=mask)

The frontend, type inference, SSA lowering, verification, and optimization are shared by both backends. That means both backends start from the same SSA:

%0 = program_id {axis=0} : i32
%1 = mul %0, 256 : i32
%2 = arange {start=0, end=256} : vector<256 x i32>
%3 = add %1, %2 : vector<256 x i32>
%4 = addptr x, %3 : vector<256 x ptr<f32>>
%5 = cmp_lt %3, n : vector<256 x bool>
%6 = load %4, %5, 0.0 : vector<256 x f32>
%7 = addptr y, %3 : vector<256 x ptr<f32>>
%8 = load %7, %5, 0.0 : vector<256 x f32>
%9 = add %6, %8 : vector<256 x f32>
%10 = addptr out, %3 : vector<256 x ptr<f32>>
store %10, %9, %5

That is the important invariant in this version.

The backend changes. The SSA does not.

The CUDA backend is still there

With the default backend, the generated source is the familiar CUDA C++:

extern "C" __global__
void add_kernel(float* x, float* y, float* out, int n) {
    int v0 = blockIdx.x;
    int v1 = (v0 * 256);
    int v2 = threadIdx.x;
    int v3 = (v1 + v2);
    bool v5 = (v3 < n);
    float v6 = (v5 ? x[v3] : 0.0f);
    float v8 = (v5 ? y[v3] : 0.0f);
    float v9 = (v6 + v8);
    if (v5) {
        out[v3] = v9;
    }
}

This is still a useful target because it is easy to inspect. It is also useful because it already supports more of the toy language than the new MLIR backend.

Version 7 is not replacing this path. It is adding another one beside it.

Emitting MLIR GPU dialect

With MYTRITON_BACKEND=mlir, the same SSA lowers to MLIR:

module attributes {gpu.container_module} {
  gpu.module @kernels {
    gpu.func @add_kernel(%x: memref<?xf32>, %y: memref<?xf32>, %out: memref<?xf32>, %n: i32) kernel {
      %bid_x = gpu.block_id x
      %tid_x = gpu.thread_id x
      %block_id_x = arith.index_cast %bid_x : index to i32
      %thread_id_x = arith.index_cast %tid_x : index to i32
      %c_i32_256 = arith.constant 256 : i32
      %v1 = arith.muli %block_id_x, %c_i32_256 : i32
      %v3 = arith.addi %v1, %thread_id_x : i32
      %idx4 = arith.index_cast %v3 : i32 to index
      %v5 = arith.cmpi slt, %v3, %n : i32
      ...
      gpu.return
    }
  }
}

The top-level shape is a GPU container module:

module attributes {gpu.container_module} {
  gpu.module @kernels {
    gpu.func @add_kernel(...) kernel {
      ...
    }
  }
}

That gpu.func is the kernel entry point. The scalar runtime parameters become ordinary MLIR function arguments, and pointers become memrefs:

ptr<f32> -> memref<?xf32>
i32      -> i32
f32      -> f32
bool     -> i1

That first rule is a big simplification. The MLIR backend currently supports only ptr<f32> as memref<?xf32>. It is not trying to model address spaces, strides beyond one dimension, or non-f32 pointer element types yet.

Program IDs and lanes

The CUDA backend maps:

program_id(0) -> blockIdx.x
arange(0, BLOCK) -> threadIdx.x

The MLIR backend makes the same mapping with GPU dialect operations:

%bid_x = gpu.block_id x
%tid_x = gpu.thread_id x
%block_id_x = arith.index_cast %bid_x : index to i32
%thread_id_x = arith.index_cast %tid_x : index to i32

MLIR’s gpu.block_id and gpu.thread_id produce index values. The current SSA uses i32 for program IDs and offsets, so the backend immediately casts them:

arith.index_cast %bid_x : index to i32

The reverse cast happens when an offset is used to index a memref:

%idx4 = arith.index_cast %v3 : i32 to index

That is one of the first places where MLIR forces a distinction that CUDA C++ lets me blur. In CUDA source, int v3 can index an array directly. In MLIR, the memref index type is explicit.

Arithmetic is direct

Most arithmetic lowering is straightforward:

opcode = {
    ("add", "i32"): "arith.addi",
    ("sub", "i32"): "arith.subi",
    ("mul", "i32"): "arith.muli",
    ("div", "i32"): "arith.divsi",
    ("add", "f32"): "arith.addf",
    ("sub", "f32"): "arith.subf",
    ("mul", "f32"): "arith.mulf",
    ("div", "f32"): "arith.divf",
}[(op.opcode, result_ty)]

The vector shape has already been interpreted as “one value per CUDA thread.” Just like the CUDA backend, the MLIR backend lowers each SSA vector operation as a scalar operation in the current thread.

For example:

%1 = mul %0, 256 : i32
%3 = add %1, %2 : vector<256 x i32>

becomes:

%c_i32_256 = arith.constant 256 : i32
%v1 = arith.muli %block_id_x, %c_i32_256 : i32
%v3 = arith.addi %v1, %thread_id_x : i32

The SSA says %3 is a vector. The backend knows that each thread owns one lane of that vector, so the emitted MLIR value is scalar.

Masked loads become scf.if

CUDA lowering emits a ternary expression for a masked load:

float v6 = (v5 ? x[v3] : 0.0f);

The MLIR backend emits structured control flow:

%c_f32_0_0 = arith.constant 0.000000e+00 : f32
%v6 = scf.if %v5 -> (f32) {
  %loaded6 = memref.load %x[%idx4] : memref<?xf32>
  scf.yield %loaded6 : f32
} else {
  scf.yield %c_f32_0_0 : f32
}

The result of the scf.if is the loaded value in the true branch and the fallback value in the false branch.

That mirrors the source language nicely:

tl.load(x + offs, mask=mask, other=0.0)

is not a raw load. It is a guarded value expression.

Masked stores use scf.if too:

scf.if %v5 {
  memref.store %v9, %out[%idx10] : memref<?xf32>
}

The same idea existed in CUDA as an if statement. MLIR just makes the control flow operation explicit.

Lowering MLIR to cubin

Generating MLIR text is useful by itself. It means a machine without MLIR bindings can still run the codegen tests and inspect the emitted source.

Execution needs more.

For CuPy inputs, the MLIR source has to be lowered to a GPU binary. Version 7 does that with a small pass pipeline:

def gpu_to_cubin_stages(*, chip: str = "sm_80", opt_level: int = 3):
    return [
        ("cleanup", "builtin.module(any(cse,canonicalize))"),
        (
            "attach-nvvm-target",
            "builtin.module("
            f"nvvm-attach-target{{module=.* chip={chip} O={opt_level}}}"
            ")",
        ),
        (
            "convert-scf-to-cf",
            "builtin.module(gpu.module(gpu.func(convert-scf-to-cf)))",
        ),
        (
            "convert-gpu-to-nvvm",
            "builtin.module(gpu.module(convert-gpu-to-nvvm))",
        ),
        ("gpu-module-to-binary", "builtin.module(gpu-module-to-binary)"),
    ]

The stages are intentionally exposed as a list of named steps. If lowering fails, the error says which stage failed:

raise RuntimeError(f"MLIR stage {name!r} failed:\n{error}") from error

This is not a polished MLIR driver. It is a small visible bridge from gpu.func to a cubin.

After the pass pipeline runs, MLIR prints a module containing a gpu.binary. For this tiny version, I extract the embedded object from the textual MLIR assembly and decode the byte string:

def extract_gpu_binary(mlir_text: str, *, symbol_name: str = "kernels") -> bytes:
    marker = f"gpu.binary @{symbol_name}"
    start = mlir_text.find(marker)
    if start < 0:
        raise ValueError(f"MLIR module does not contain {marker}")

    object_start = mlir_text.find(', "', start)
    if object_start < 0:
        raise ValueError(f"MLIR gpu.binary @{symbol_name} does not contain an object")

    return _decode_mlir_byte_string(mlir_text, object_start + 3)

This is deliberately humble. It is enough for the milestone, not a claim that string-parsing MLIR assembly is a general binary extraction API.

NumPy codegen versus CuPy execution

One important detail changed while building this version: MLIR source generation and MLIR cubin generation are separate.

With NumPy inputs, mytriton stops after generating MLIR source:

Python kernel -> SSA -> MLIR text

That path does not require MLIR Python bindings. It lets GitHub Actions run MLIR codegen tests on a normal CPU runner.

With CuPy inputs, mytriton needs an executable kernel:

Python kernel -> SSA -> MLIR text -> cubin -> CuPy launch

So cubin generation is lazy. The artifact initially stores only src:

artifact = _CompilationArtifact(
    ops=tuple(ops),
    ssa_ops=tuple(ssa_ops),
    src=src,
    threads_per_block=threads_per_block,
    cubin=cubin,
)

Then, at launch time, the MLIR backend checks whether execution is actually needed:

needs_execution = cuda_execution_required(runtime_args, backend_name="MLIR")
if artifact.cubin is None and needs_execution:
    artifact = replace(
        artifact,
        cubin=compile_mlir_source_to_cubin(artifact.src, chip=chip),
    )

That keeps the compile-only path light and keeps MLIR as an optional external toolchain instead of a required package dependency.

Calling the cubin

CuPy can load a cubin through cp.cuda.function.Module:

module = cp.cuda.function.Module()
module.load(cubin)
kernel = module.get_function(kernel_name)

The launch shape is the same one used by the CUDA backend:

kernel(
    launch_grid,
    (threads_per_block,),
    converted_args,
)

The interesting part is converted_args.

The CUDA backend can pass a CuPy array directly to a float* kernel parameter. The MLIR kernel parameter is a memref:

%x: memref<?xf32>

The ranked memref ABI expects more than one value. For the one-dimensional arrays in this MVP, each CuPy array becomes:

[
    array,              # allocated ptr
    array,              # aligned ptr
    np.int64(0),         # offset
    np.int64(array.shape[0]),
    np.int64(1),
]

Those five values are the runtime descriptor for memref<?xf32>:

allocated pointer
aligned pointer
offset
size for dimension 0
stride for dimension 0

The first two entries are both the CuPy array object. CuPy already knows how to pass the device pointer behind an array to a CUDA kernel. For a simple C-contiguous array, the allocated pointer and the aligned pointer are the same address, so mytriton passes the same object twice.

The offset is zero because this MVP does not try to pass array views into the middle of an allocation.

The size is the number of elements:

np.int64(array.shape[0])

The stride is one because the supported layout is contiguous:

np.int64(1)

This is very different from the CUDA C++ backend. CUDA source sees:

void add_kernel(float* x, float* y, float* out, int n)

and CuPy can pass each array as one pointer-like argument.

The MLIR kernel sees:

gpu.func @add_kernel(
  %x: memref<?xf32>,
  %y: memref<?xf32>,
  %out: memref<?xf32>,
  %n: i32
) kernel

so each memref parameter expands into its descriptor pieces at the ABI boundary.

So the Python runtime argument list:

x, y, out, n

becomes:

x_alloc, x_aligned, x_offset, x_size, x_stride,
y_alloc, y_aligned, y_offset, y_size, y_stride,
out_alloc, out_aligned, out_offset, out_size, out_stride,
n

If this backend supported a two-dimensional memref, the descriptor would need more shape information:

allocated pointer
aligned pointer
offset
size for dimension 0
size for dimension 1
stride for dimension 0
stride for dimension 1

That would also have to line up with how the kernel indexes the memref and with the actual strides of the CuPy array. A contiguous (rows, cols) array does not have stride one in both dimensions; its logical strides are usually (cols, 1).

That is why the MLIR execution helper currently accepts only one-dimensional C-contiguous CuPy arrays. The memref descriptor conversion is intentionally small and explicit.

Backend-parametrized tests

The tests now exercise both backends from the same test bodies.

The shared fixture is simple:

BACKENDS = ["cuda", "mlir"]


@pytest.fixture(params=BACKENDS)
def backend(request, monkeypatch):
    selected = request.param
    monkeypatch.setenv("MYTRITON_BACKEND", selected)
    return selected

The add_kernel and copy_kernel codegen tests assert:

  • expression-tree operations;
  • optimized SSA;
  • generated CUDA source when backend == "cuda";
  • generated MLIR source when backend == "mlir".

The execution tests are marked separately:

@pytest.mark.execution
def test_add_kernel_execution(cp, backend):
    ...

That lets normal CI run codegen and unit tests without launching kernels:

python -m pytest -m "not execution"

On a CUDA machine, the full test suite still runs execution tests too:

MYTRITON_REQUIRE_CUDA=1 python -m pytest

This split matters more now that there are two backends. I want the repository to verify CUDA and MLIR source generation on ordinary GitHub Actions runners, while keeping actual GPU execution as a local or specialized-runner concern.

Honest MVP boundaries

Version 6 added tl.program_id(1) and matrix-style kernels. The first MLIR backend does not support them.

That creates a dangerous possibility: silently generating wrong MLIR.

For example, if the MLIR backend blindly mapped every program ID to %block_id_x, a two-dimensional kernel would compile to something that looked plausible but used the wrong block coordinate.

So Version 7 rejects unsupported cases explicitly:

if op.opcode == "program_id":
    axis = op.attrs["axis"]
    if axis != 0:
        raise TypeError(f"MLIR MVP supports only program_id(0), got {axis}")

and:

if op.opcode == "arange":
    start = op.attrs["start"]
    if start != 0:
        raise TypeError(f"MLIR MVP supports only arange(0, N), got {start}")

Those checks are small, but they are important. A limited backend is fine. A backend that quietly lies is not.

There are unit tests for these boundaries now:

with pytest.raises(TypeError, match=r"MLIR MVP supports only program_id\(0\)"):
    MLIRCodegen().generate("kernel", ops, [])

That is a recurring pattern in this project: when a feature is not supported yet, make the boundary visible.

What changed conceptually

Before Version 7, mytriton had a frontend, a middle-end, and one backend.

That was enough to feel like a compiler, but the backend boundary was still a little theoretical. Since there was only one backend, it was easy for compiler APIs to accidentally mean “CUDA.”

The launch result was named cuda_src.

The artifact stored cuda_src.

The tests asserted CUDA source.

After Version 7, those names have to become more precise:

src = source for the selected backend

That seems like a naming cleanup, but it reflects a deeper change. The compiler now has a target-independent path up to optimized SSA and target-specific paths after that.

                  -> CUDA C++ -> NVRTC / RawKernel
Python -> SSA ----|
                  -> MLIR GPU -> NVVM -> cubin / Module

The MLIR backend is still tiny. But the second arrow changes the shape of the project.

What’s next

There are two natural directions from here.

One is to make the MLIR backend less tiny. The next easy operations would be the elementwise features from Part 4: Elementwise Ops:

  • unary negation;
  • tl.exp;
  • tl.maximum and tl.minimum;
  • tl.where.

Those would let the MLIR backend run ReLU, leaky ReLU, and sigmoid. That would be a good test of whether the new backend can grow feature by feature the same way the CUDA backend did.

The other direction is to bring in more of Version 6:

  • tl.program_id(1);
  • two-dimensional launch grids;
  • reductions;
  • eventually softmax.

That is a larger step, especially reductions. CUDA lowering implemented reductions with explicit shared memory and synchronization. MLIR has its own ways to represent GPU memory and barriers, and I do not want to rush that into an accidental translation.

Tiled matmul is still waiting too. Part 6: Reductions made it obvious that source-visible shared memory and synchronization are missing from the language. Version 7 did not solve that. It did something more structural first: it made mytriton stop having only one target.

That feels like the right detour. Before teaching the language more CUDA-shaped features, I want the compiler boundary between “this is my IR” and “this is one backend’s lowering choice” to stay visible.

All code for this milestone is available at https://github.com/pbelevich/mytriton/tree/ver7.