Skip to content

Adding a new op

A new op means writing code in the six places below, and the table is in the order to work through them.

The spec goes first: it decides what the other five files contain, and in the end it is what they are checked against. The spec is this pipeline's input, and the other five are written from it.

# File Named in the spec by Contents
1 src/tileops/manifest/<family>.yaml the key is the op's class name the spec itself
2 src/tileops/ops/<family>/… source.op the op class, subclassing Op
2 src/tileops/ops/__init__.py — the op's name, exported
3 src/tileops/kernels/<family>/… source.kernel the kernel class, subclassing Kernel
4 tests/ops/test_<name>.py source.test the comparison against ref_api
5 benchmarks/ops/bench_<name>.py source.bench the benchmark

GemmFwdOp — the plainest matmul there is — runs through all six below.

Step 1: write the spec

What the fields mean and how to write them is in writing a spec. Two things are specific to a new op.

The first is the status. A new op starts at status: spec-only, meaning the interface is settled and there is no implementation yet, so validation runs L0 — the structure check — and does not fail over the missing code.

The second is source.kernel_map, the one field in the spec that nothing can derive.

An op may have more than one kernel behind it: GEMM uses a matmul kernel at general shapes, and at M = 1 the problem degenerates to a matrix-vector product that a different kernel does faster. kernel_map is the roster of those kernels — a name for each, against the Kernel class that implements it:

GemmFwdOp:
  ref_api: torch.matmul
  family: gemm
  status: spec-only
  signature:
    inputs:
      a: {dtype: "float16 | bfloat16"}
      b: {dtype: "same_as(a)"}
    outputs:
      d: {dtype: "same_as(a)"}
    params:
      trans_a: {type: bool, default: false}
      trans_b: {type: bool, default: true}
    shape_rules:
      - "d.shape == ((a.shape[1] if trans_a else a.shape[0]), (b.shape[0] if trans_b else b.shape[1]))"
  source:
    kernel: tileops/kernels/gemm/dense.py
    kernel_map:
      gemm_kernel: GemmKernel
      gemv_kernel: GemvKernel
    op: tileops/ops/gemm/gemm.py
    test: tests/ops/test_gemm.py
    bench: benchmarks/ops/bench_gemm.py

Those names are how a kernel is asked for at runtime: _eager_forward picks one, passes the name to get_or_build_kernel, and the op layer looks the class up in kernel_map and builds it (see step 2). An external backend registers against the same roster — whichever name it registers a build_kernel for is the kernel of the op it takes over.

The names are yours to choose, they should say what the kernel is for, and once the op code uses one it should not change: it is the word the spec, the op and any backend all agree on. That is also why nothing can derive it — only whoever writes the kernels knows how many cases the op splits into.

Step 2: write the op class

The op class subclasses Op and sits between the spec and the kernel: it validates the arguments against the spec, infers the output shapes, then fetches a kernel and launches it. It comes first because the spec dictates all of it, and the line where it builds a kernel is what fixes that kernel's constructor signature.

The class, and its four members

GemmFwdOp's skeleton, with the parts of each body that are beside the point elided:

class GemmFwdOp(Op):
    def __init__(self, trans_a=False, trans_b=True, kernel_map=None, tune=False):
        self.trans_a, self.trans_b, self.tune = trans_a, trans_b, tune
        self.dispatch_kernel(kernel_map)             # establishes this instance's kernel_map

    @property
    def default_kernel_map(self):                    # the spec's source.kernel_map
        return {"gemm_kernel": GemmKernel, "gemv_kernel": GemvKernel}

    def _infer_output_shapes(self, a_shape, b_shape):
        m = a_shape[1] if self.trans_a else a_shape[0]
        n = b_shape[0] if self.trans_b else b_shape[1]
        return {"d": (m, n)}                         # the spec's shape_rules

    def forward(self, a, b):
        self._validate_dtypes(a, b)                  # generated by the base class
        m, n, k = self._infer_mnk(a, b)
        a, b = a.contiguous(), b.contiguous()        # handed over as the spec declares it
        slot = "gemv_kernel" if m == 1 else "gemm_kernel"
        kernel = self.get_or_build_kernel(
            slot,                                    # a name from kernel_map
            (a, b),                                  # the memo key's tensors, and what a backend receives
            key=(m, n, k, a.dtype),                  # the cache key on the in-tree side
            build=lambda: self.kernel_map[slot](m, n, k, a.dtype, tune=self.tune),
        )
        return kernel(a, b)

Four members to write, each of them from the spec:

# Member Written from
1 __init__ the names and defaults in signature.params, plus kernel_map and tune, closing with self.dispatch_kernel(kernel_map) to establish this instance's kernel_map
2 default_kernel_map source.kernel_map: the same names, against the Kernel classes themselves
3 _infer_output_shapes the rules in signature.shape_rules that derive an output's shape
4 forward signature.inputs — its order and defaults, optional inputs last — plus the validation, the contiguity, fetching the kernel and launching it

Two more members arrive on their own. When the subclass is defined, the base class synthesises _validate_dtypes and eval_roofline from the spec's dtype declarations and its roofline, so they are there to call — and worth overriding only where the op needs something the spec cannot say.

get_or_build_kernel

A kernel is a compiled artefact, hundreds of milliseconds to seconds to build, while an op instance is called over and over at different shapes and dtypes. The op layer therefore keeps a memo table: a kernel this call needs and has built before comes straight back, and only otherwise is one built and stored. get_or_build_kernel is that table's only entrance, and the point where the in-tree implementation and an external backend part ways — the second layer of selection in the backend protocol.

Its four arguments:

name — which kernel this call wants, as a name from kernel_map.

slot = "gemv_kernel" if m == 1 else "gemm_kernel"

The in-tree side looks up the Kernel class under that name; a backend looks up the build_kernel it registered under it. An op has as many names as it has cases.

inputs — the tensors the kernel is about to be handed, in signature.inputs order, one slot per input.

self.get_or_build_kernel(slot, (a, b), ...)                     # GEMM: two required inputs
self.get_or_build_kernel("group_norm", (x, weight, bias), ...)  # an absent optional input is None

The external path keys on it — the device, plus each slot's (dtype, shape). The device counts because an artefact compiled for one card may hold resources on it. A backend's build_kernel receives the same tensors as TensorSpecs: device, dtype, shape, no data.

An optional input that was not passed keeps its slot, as None; that is what a backend reads presence off, rather than counting slots. Squeeze the empty slots out, and a clamp with only a lower bound looks exactly like one with only an upper bound.

Omitting inputs raises nothing until a backend is installed, and then OpNotAvailableError: the op stays in-tree only, out of reach of any target (see after install: two states).

key — what the in-tree kernel specializes on; the in-tree path only. What becomes of these last two once a backend serves the op is in how one call reaches build_kernel.

key=(m, n, k, a.dtype)                             # GEMM: three dimensions and the dtype
key=(self._cache_key(*input_shapes), x.dtype)      # the general form

The default _cache_key takes the sizes of every non-static axis across the inputs — always correct, but it can over-fragment: one compile per distinct shape. Where the kernel depends on fewer quantities, override it to project the shape onto those, flattening the leading dims to one product when the kernel treats its input as 2-D.

build — how that in-tree kernel is constructed; the in-tree path only.

build=lambda: self.kernel_map[slot](m, n, k, a.dtype, tune=self.tune)

Called once per key, which is why compiling belongs here. It may return one Kernel, a sequence of Kernels built together, or a dataclass carrying them — the last two suit an op that launches several kernels per call.

An op with no in-tree implementation at all, one written to depend on a backend, may leave build out; a call on a device no target claims then raises OpNotAvailableError.

Finishing: the compile boundary, and registering

Two things to finish, a few lines each:

  • To support torch.compile, declare a compile boundary as well: forward only calls the opaque operator, and the validation, the kernel lookup and the launch move into _eager_forward. The op above declares none, so its forward holds all the work. How to declare it is in bringing an op into torch.compile.
  • Add the op's name to the imports and __all__ in src/tileops/ops/__init__.py, or from tileops.ops import ... will not find it.

Step 3: write the kernel

A kernel class subclasses Kernel, lives under src/tileops/kernels/, is written in TileLang, compiles at construction and launches on __call__. Its constructor and call signatures are the ones step 2 just used — the build lambda and the kernel(a, b) that follows it.

This is the one place of the six the spec does not constrain: a kernel neither reads the spec nor is checked against it, and the spec records only its path and class name.

How the constructor and the call divide their arguments is a hard requirement: only values compiled into the generated code go in the constructor. GemmKernel divides them like this:

class GemmKernel(Kernel):
    def __init__(self, m, n, k, dtype, config=None, tune=False, trans_a=False, trans_b=False):
        self.kernel = _gemm_kernel(m, n, k, trans_a, trans_b, self.dtype_str)  # this line compiles
        self.init_config(config, tune)      # block_m / block_n / block_k / num_stages

    def __call__(self, a, b):               # a call passes tensors, nothing else
        ...

m, n, k, the dtype and the two layout flags are constructor arguments because the generated code treats them as constants: loop bounds, TMA descriptors and the WGMMA shape all unroll from them, as do the tile sizes (block_m and the rest). The tensors belong to __call__, where each call swaps pointers.

Dividing them wrong costs a recompile. A decode step advances one token at a time, so seq_len grows by one every step and batch changes with the running set:

# wrong: seq_len in the constructor — every step is a new kernel
kernel = AttnKernel(batch, seq_len, num_heads, dtype)

# right: compile-time constants in the constructor, the varying sizes per call
kernel = AttnKernel(num_heads, head_dim, dtype)
out = kernel(q, k, v)                       # seq_len is read off the tensor shapes

With the first form, seq_len ends up in get_or_build_kernel's key, every step misses, every step compiles, and decode goes nowhere.

Step 4: write the test

Tests live in tests/ops/, and what they compare against is the spec's ref_api, point by point, over the shapes and dtypes the spec declares — small shapes marked smoke for the PR checks, large ones full for the nightly.

The scaffolding is TestBase and FixtureBase from tests/test_base.py, with the cases in PARAMS.

Where the op has an optional input, both sides need a case — passed and not passed often run different kernels.

Step 5: write the benchmark

Benchmarks live in benchmarks/ops/ and subclass ManifestBenchmark. The shapes are not written here: they come from the spec's workloads through load_workloads(<op>), and hand-written shapes fail L4 validation:

from benchmarks.benchmark_base import ManifestBenchmark, workload_params
from tileops.manifest import load_workloads

_OP_NAME = "GemmFwdOp"

Record at least one non-TileOPs baseline as well, or the row has nothing to compare against. Where a baseline needs its input converted, that conversion stays inside its own timed region. What the reported numbers mean is in how a benchmark is timed.

Step 6: flip the status, and let CI take over

With the other five written, check your own work with the three commands below:

python scripts/validate_manifest.py --check-op GemmFwdOp   # spec and code agree, all five levels
python -m pytest tests/ops/test_gemm.py -v                # numerics match ref_api
python -m pytest benchmarks/ops/bench_gemm.py             # the benchmark produces numbers

With all three passing, flip the spec's status from spec-only to implemented. That one edit takes validation from L0 to all five levels and puts the op inside CI's reach: every later change is held against the spec by the validator, the tests and the nightly benchmark.

Afterwards

Once the op runs, two optional things remain: