diff --git a/.codespellrc b/.codespellrc new file mode 100644 index 00000000..71fb0584 --- /dev/null +++ b/.codespellrc @@ -0,0 +1,4 @@ +[codespell] +# Fixed vocabulary, not typos: `aranges` is the plural of tl.arange, `lits` +# names the walk's string-literal spans (tilelens/ir/_mlir_walk.py). +ignore-words-list = aranges,lits diff --git a/tests/end_to_end/test_ir_client.py b/tests/end_to_end/test_ir_client.py index e2b13a44..97d7eedd 100644 --- a/tests/end_to_end/test_ir_client.py +++ b/tests/end_to_end/test_ir_client.py @@ -272,3 +272,20 @@ def test_another_triton_release_compiles_nothing_and_goes_on(monkeypatch): ((), (failure,)) = ir.logs[0] unavailable = host_compile_unavailable(failure.error) assert "supports Triton 3.8.x only" in str(unavailable) + + +def test_the_real_reader_through_the_parse_cache(): + ir = _ToyIR() # ParseCache's default reader: tilelens.ir.ttir_reader.parse_ttir + traced = tilelens.trace(ir)(_make_add_one()) + x, out = _inputs() + + traced[(4,)](x, out, 64, BLOCK=16) + traced[(4,)](x, out, 64, BLOCK=16) + + first, second = ir.outcomes + assert (first.error, first.refusal) == (None, None) + assert first.graph.kernel_name == "add_one" + # The second launch is a cache hit. + assert second is first + (config,) = ir.last_verdict.per_config + assert config.status == "parsed" diff --git a/tests/unit/ir/reader_kernels.py b/tests/unit/ir/reader_kernels.py new file mode 100644 index 00000000..3d44416a --- /dev/null +++ b/tests/unit/ir/reader_kernels.py @@ -0,0 +1,281 @@ +"""Kernels whose TTIR the reader tests read as ``reader_ttir/.ttir``. + +``ttir_corpus.py`` compiles them at test time; ``test_ttir_reader.py`` +reads the result, and checks that a refusal's loc names the right line of +this file. + +Not a test module: pytest imports it (python_files = *.py) and finds nothing. +""" + +from __future__ import annotations + +from typing import Any + +import triton +import triton.language as tl + + +# ── p1: a loop-carried pointer advanced by a loop-VARIANT amount ── +@triton.jit +def p1_variant_delta(x_ptr, out_ptr, n): + p = x_ptr + 20 + for k in range(-8, n): # n is a runtime scalar + v = tl.load(p) + tl.store(out_ptr, v) + p += k # the advance depends on the induction variable + + +# ── p2: two pointer iter_args swapped in the scf.yield ── +@triton.jit +def p2_swap(a_ptr, b_ptr, n): + p = a_ptr + q = b_ptr + 100 + for i in range(0, n): + tl.store(p, 1.0) + p, q = q, p # ping-pong between two buffers + + +# ── p3: a noinline tt.call (the callee body is another tt.func) ── +@triton.jit(noinline=True) +def _p3_helper(x_ptr, pid): + tl.store(x_ptr + pid, 1.0) + + +@triton.jit +def p3_call_guarded(x_ptr, n): + # callee formal names match caller names; the call sits under `if pid < n` + pid = tl.program_id(0) + if pid < n: + _p3_helper(x_ptr, pid) + + +@triton.jit +def p3_call_offset(x_ptr, n): + # the ACTUAL argument differs from the caller value sharing the formal's name + pid = tl.program_id(0) + _p3_helper(x_ptr, pid + 100) + + +@triton.jit(noinline=True) +def _p3_helper_c(dst, i): + tl.store(dst + i, 1.0) + + +@triton.jit +def p3_call_formals(x_ptr, n): + # callee formals match no caller name + pid = tl.program_id(0) + _p3_helper_c(x_ptr, pid + 100) + + +# ── p4: an atomic observation reaching an address ── +@triton.jit +def p4_observed_direct(cnt_ptr, x_ptr, n): + old = tl.atomic_add(cnt_ptr, 1) + p = x_ptr + old % n + v = tl.load(p) + tl.store(p, v + 1.0) + + +@triton.jit +def p4_observed_loop(cnt_ptr, x_ptr, n): + # the same address, carried through an scf.for iter_arg (offset0) + old = tl.atomic_add(cnt_ptr, 1) + p = x_ptr + old % n + for i in range(0, 4): + v = tl.load(p) + tl.store(p, v + 1.0) + p += 1 + + +@triton.jit +def p4_observed_delta(cnt_ptr, x_ptr, n): + # the observation in the per-iteration DELTA instead of offset0 + old = tl.atomic_add(cnt_ptr, 1) + step = old % n + p = x_ptr + for i in range(0, 4): + v = tl.load(p) + tl.store(p, v + 1.0) + p += step + + +# ── integer widths ── +@triton.jit +def rv_trunci_alias(x_ptr): + # trunc_i32(pid_i64 * 2**32) is 0 for every pid: every program stores x[0] + pid = tl.program_id(0).to(tl.int64) + off = (pid * 4294967296).to(tl.int32) + tl.store(x_ptr + off, 1) + + +@triton.jit +def rv_i32_wrap(x_ptr, S): + # (pid * S) * S wraps to 0 in i32 for S = 65536 + pid = tl.program_id(0) + off = (pid * S) * S + tl.store(x_ptr + off, 1) + + +@triton.jit +def unsigned_index(x_ptr, n): + # divui, cmpi ult and extui read their operands unsigned + pid = tl.program_id(0).to(tl.uint32) + q = pid // 3 + m = pid < n.to(tl.uint32) + tl.store(x_ptr + q, 1.0, mask=m) + + +# ── inline asm ── +@triton.jit +def rv_inline_asm_store(x_ptr, OFF: tl.constexpr): + # an impure asm st.global through a pointer cast to i64: a memory access + # the graph cannot see + p = (x_ptr + OFF).to(tl.int64, bitcast=True) + v = tl.full((), 7, tl.int32) + tl.inline_asm_elementwise( + "st.global.b32 [$1], $2; mov.b32 $0, 0;", + "=r,l,r", + [p, v], + dtype=tl.int32, + is_pure=False, + pack=1, + ) + + +# ── shapes the reader models ── +@triton.jit +def loop_two_step_advance(x_ptr, n, s): + # two addptrs per iteration: the advance is their (loop-invariant) sum + p = x_ptr + for i in range(0, n): + tl.store(p, 1.0) + p += s + p += 2 + + +@triton.jit +def loop_observed_advance(cnt_ptr, x_ptr, n): + # the advance is an atomic observed inside the loop: loop-variant + p = x_ptr + for i in range(0, n): + tl.store(p, 1.0) + p += tl.atomic_add(cnt_ptr, 1) + + +@triton.jit +def where_pointer(x_ptr, n): + # arith.select over two pointers of one base + offs = tl.arange(0, 16) + p = tl.where(offs < n, x_ptr + offs, x_ptr + 100) + tl.store(p, 1.0) + + +@triton.jit +def tile3d_shared_arange(x_ptr, N: tl.constexpr): + # one make_range on all three dims of a tile: three lane variables + r = tl.arange(0, N) + off = r[:, None, None] * (N * N) + r[None, :, None] * N + r[None, None, :] + tl.store(x_ptr + off, 1.0) + + +@triton.jit +def expand_iterarg_3d(x_ptr, out_ptr, n, N: tl.constexpr): + # a loop-carried [N, N] pointer tile expanded to 3D inside the loop, next + # to the same make_range at dim 0: the load reads j*N + l - i*N + r = tl.arange(0, N) + p = x_ptr + r[:, None] * N + r[None, :] + for k in range(0, n): + q = p[None, :, :] - r[:, None, None] * N + v = tl.load(q) + o = out_ptr + r[:, None, None] * N * N + r[None, :, None] * N + r[None, None, :] + tl.store(o, v) + p += 1 + + +@triton.jit +def expand_iterarg_mask(x_ptr, out_ptr, n, M, N: tl.constexpr): + # a loop-carried 1D pointer expanded to 2D inside the loop, masked by the + # same make_range at the pointer's lane (dim 1) + r = tl.arange(0, N) + p = x_ptr + r + for k in range(0, n): + q = p[None, :] + r[:, None] * 0 + v = tl.load(q, mask=r[None, :] < M) + tl.store(out_ptr + r[:, None] * N + r[None, :], v) + p += N + + +@triton.jit +def int_iterarg_offset(x_ptr, n, B: tl.constexpr): + # an integer offset carried by the loop: a loop-variant address + offs = tl.arange(0, B) + for k in range(0, n): + tl.store(x_ptr + offs, 1.0) + offs += B + + +@triton.jit +def iv_wrap(x_ptr, lo, n): + # the induction variable's increment wraps in i32 when n is near INT32_MAX + for k in range(lo, n, 1 << 20): + tl.store(x_ptr + k, 1.0) + + +@triton.jit +def pure_asm_int_addr(x_ptr): + # a "pure" asm handed the address as an integer stores through it + a = x_ptr.to(tl.int64, bitcast=False) + r = tl.inline_asm_elementwise( + "st.global.b32 [$1], $2; mov.b32 $0, 0;", + "=r,l,r", + [a, tl.full([], 7, tl.int32)], + dtype=tl.int32, + is_pure=True, + pack=1, + ) + tl.store(x_ptr + 4096, r) + + +@triton.jit +def observed_lanes(cnt_ptr, x_ptr, N: tl.constexpr): + # two lanes of one tensor atomic's old values in one address + r = tl.arange(0, N) + old = tl.atomic_add(cnt_ptr + r, 1) + c = tl.minimum(tl.maximum(old, 0), 10) + tl.store(x_ptr + c[:, None] - c[None, :], 1.0) + + +_F32 = "*fp32" +# corpus name (reader_ttir/.ttir, the kernel of that name) -> +# (signature, constexprs), compiled for sm80 with the default options +SPECS: dict[str, tuple[dict[str, str], dict[str, Any]]] = { + "p1_variant_delta": ({"x_ptr": _F32, "out_ptr": _F32, "n": "i32"}, {}), + "p2_swap": ({"a_ptr": _F32, "b_ptr": _F32, "n": "i32"}, {}), + "p3_call_guarded": ({"x_ptr": _F32, "n": "i32"}, {}), + "p3_call_offset": ({"x_ptr": _F32, "n": "i32"}, {}), + "p3_call_formals": ({"x_ptr": _F32, "n": "i32"}, {}), + "p4_observed_direct": ({"cnt_ptr": "*i32", "x_ptr": _F32, "n": "i32"}, {}), + "p4_observed_loop": ({"cnt_ptr": "*i32", "x_ptr": _F32, "n": "i32"}, {}), + "p4_observed_delta": ({"cnt_ptr": "*i32", "x_ptr": _F32, "n": "i32"}, {}), + "rv_trunci_alias": ({"x_ptr": "*i32"}, {}), + "rv_i32_wrap": ({"x_ptr": "*i32", "S": "i32"}, {}), + "unsigned_index": ({"x_ptr": _F32, "n": "i32"}, {}), + "rv_inline_asm_store": ({"x_ptr": "*i32", "OFF": "constexpr"}, {"OFF": 4096}), + "loop_two_step_advance": ({"x_ptr": _F32, "n": "i32", "s": "i32"}, {}), + "loop_observed_advance": ({"cnt_ptr": "*i32", "x_ptr": _F32, "n": "i32"}, {}), + "where_pointer": ({"x_ptr": _F32, "n": "i32"}, {}), + "tile3d_shared_arange": ({"x_ptr": _F32, "N": "constexpr"}, {"N": 4}), + "expand_iterarg_3d": ( + {"x_ptr": _F32, "out_ptr": _F32, "n": "i32", "N": "constexpr"}, + {"N": 4}, + ), + "expand_iterarg_mask": ( + {"x_ptr": _F32, "out_ptr": _F32, "n": "i32", "M": "i32", "N": "constexpr"}, + {"N": 4}, + ), + "int_iterarg_offset": ({"x_ptr": _F32, "n": "i32", "B": "constexpr"}, {"B": 8}), + "iv_wrap": ({"x_ptr": _F32, "lo": "i32", "n": "i32"}, {}), + "pure_asm_int_addr": ({"x_ptr": "*i32"}, {}), + "observed_lanes": ({"cnt_ptr": "*i32", "x_ptr": _F32, "N": "constexpr"}, {"N": 4}), +} diff --git a/tests/unit/ir/test_mlir_walk.py b/tests/unit/ir/test_mlir_walk.py new file mode 100644 index 00000000..ff5de0bb --- /dev/null +++ b/tests/unit/ir/test_mlir_walk.py @@ -0,0 +1,527 @@ +"""tilelens.ir._mlir_walk: the bindings-first walk under the TTIR reader. + +The corpus is tests/unit/ir/ttir_corpus.py: the TTIR Triton 3.8 prints for +tests/unit/ir/ttir_kernels.py's kernels, compiled at test time, and +hand-written texts. +""" + +from __future__ import annotations + +import re +import subprocess +import sys +import tempfile +from pathlib import Path + +import pytest + +from tilelens.ir import _mlir_walk as W + +from . import ttir_corpus + +REPO = Path(__file__).resolve().parents[3] +FILES = [n.removeprefix("ttir/") for n in ttir_corpus.NAMES if n.startswith("ttir/")] + + +def _text(name: str) -> str: + return ttir_corpus.text(f"ttir/{name}") + + +def _ops(m: W.Module, name: str) -> list[W.Op]: + return [op for op in m.ops if op.name == name] + + +@pytest.fixture(autouse=True) +def _fresh_cache(): + W._CACHE.clear() + yield + W._CACHE.clear() + + +# ─────────────────────────── the corpus ─────────────────────────── + + +def test_generic_form_reads_like_the_custom_form(): + # the bindings read the attributes, whatever form the text prints + m = W._walk(_text("crafted_generic_form.ttir")) + assert _ops(m, "arith.cmpi")[0].attrs["predicate"] == "slt" + assert _ops(m, "tt.get_program_id")[0].attrs["axis"] == 1 + + +def test_structure_records_are_consistent(): + for name in FILES: + m = W._walk(_text(name)) + assert m.ops[0].name == "builtin.module" and m.ops[0].path == () + assert [op.index for op in m.ops] == list(range(len(m.ops))) + assert [b.index for b in m.blocks] == list(range(len(m.blocks))) + assert [v.index for v in m.values] == list(range(len(m.values))) + for op in m.ops: + for v in op.results: + assert m.values[v].op == op.index + for k, blocks in enumerate(op.regions): + for pos, b in enumerate(blocks): + blk = m.blocks[b] + assert (blk.op, blk.region, blk.position) == (op.index, k, pos) + for i in blk.ops: + assert m.ops[i].path == op.path + (b,) + if op.path: + assert m.blocks[op.block].ops[op.position] == op.index + for v in m.values: + assert (v.op is None) != (v.block is None) + # pre-order: an op's regions come after it, its uses are in-module values + for op in m.ops: + assert all(0 <= v < len(m.values) for v in op.operands) + assert op.operand_types == tuple(m.values[v].type for v in op.operands) + + +# ─────────────────────────── a misread fails closed ─────────────────────────── +# The bindings parse the original text while the text scan reads an edited +# copy: a scan that misreads the module must never align silently. + + +@pytest.mark.parametrize( + "misread, match", + [ + # the op name the line prints is not the op the bindings walked + (lambda ln: ln.replace("arith.cmpi", "arith.subi"), "text op 'arith.subi'"), + # the line's loc, the anchor to the bindings, names another site + ( + lambda ln: re.sub(r"loc\(#loc\d+\)$", 'loc("k.py":1:1)', ln), + "result 0 loc differs", + ), + ], + ids=["op-name", "result-loc"], +) +def test_a_misread_reports_its_line(misread, match): + text = _text("golden_add_sm80.ttir") + lines = text.splitlines() + i = next(k for k, ln in enumerate(lines) if "arith.cmpi slt" in ln) + lines[i] = misread(lines[i]) + with pytest.raises(W.MisalignedModule, match=match) as e: + W._walk(text, scan_text="\n".join(lines)) + assert e.value.line_no == i + 1 + + +# a zero-result op's loc has no second source: a garbled trailer must not +# read as "no loc" while the rest of the module prints locs +@pytest.mark.parametrize( + "old, new", [("tt.return loc(", "tt.return oc("), ("} loc(#loc)", "} loc#(#loc)")] +) +def test_a_garbled_loc_is_not_read_as_no_loc(old, new): + text = _text("golden_add_sm80.ttir") + assert old in text + with pytest.raises(W.MisalignedModule, match="prints no loc"): + W._walk(text, scan_text=text.replace(old, new, 1)) + + +# ─────────────────────────── what the walk reads ─────────────────────────── + + +def test_unicode_names_and_locs(): + m = W._walk(_text("crafted_unicode_strings.ttir")) + assert [f.sym_name for f in m.funcs] == ["核"] + assert [a.name for a in m.funcs[0].args] == ["π_ptr", "数"] + assert _ops(m, "tt.load")[0].loc_name == "值" + assert {op.loc.file for op in m.ops if op.loc} == {"/tmp/内核/k.py"} + + +@pytest.mark.parametrize( + "body, value", + [ + ("\\E9\\94\\99", "错"), + ("a\\22b\\\\c", 'a"b\\c'), + ("tab\\09nl\\0A", "tab\tnl\n"), + ("\\n\\t", "\n\t"), + ("plain", "plain"), + ], +) +def test_unescape_decodes_utf8_bytes(body, value): + assert W._unescape(body) == value + + +@pytest.mark.parametrize("body", ["\\E9\\94", "\\q", "\\"]) +def test_unescape_rejects_malformed(body): + with pytest.raises(ValueError): + W._unescape(body) + + +def test_constant_kinds(): + m = W._walk(_text("adv_consts.ttir")) + values = [op.attrs["value"] for op in _ops(m, "arith.constant")] + assert -9223372036854775807 in values and ("float", None) in values + text = _text("crafted_attr_dicts.ttir") + assert [op.attrs["value"] for op in _ops(W._walk(text), "arith.constant")] == [ + 16, + -3, + ] + dense = "%ds = arith.constant dense<[1, 2]> : tensor<2xi32>\n %r = tt.make_range" + m = W._walk(text.replace("%r = tt.make_range", dense, 1)) + assert ("dense", None) in [op.attrs["value"] for op in _ops(m, "arith.constant")] + + +@pytest.mark.parametrize( + "name, shapes", + [ + ("crafted_empty_else.ttir", {"scf.if": [(1, 1)]}), + ("crafted_empty_for.ttir", {"scf.for": [(1,)]}), + ( + "crafted_empty_bodies.ttir", + {"scf.if": [(1, 1), (1, 1), (1, 0)], "scf.for": [(1,)]}, + ), + ], +) +def test_empty_region_bodies(name, shapes): + m = W._walk(_text(name)) + for opname, want in shapes.items(): + assert [tuple(len(r) for r in op.regions) for op in _ops(m, opname)] == want + for blk in m.blocks: + if m.ops[blk.op].name in ("scf.if", "scf.for"): + last = m.ops[blk.ops[-1]] + assert last.name == "scf.yield" + if last.implicit: + assert last.line_no is None and last.operands == () and last.loc is None + + +def test_empty_for_body_keeps_its_induction_variable(): + m = W._walk(_text("crafted_empty_for.ttir")) + (loop,) = _ops(m, "scf.for") + (body,) = loop.regions[0] + assert len(m.blocks[body].args) == 1 and m.blocks[body].arg_types == ("i32",) + assert [m.ops[i].implicit for i in m.blocks[body].ops] == [True] + + +def test_descriptor_operand_order(): + m = W._walk(_text("adv_descs.ttir")) + tile, row = "!tt.tensordesc<32x32xf16>", "!tt.tensordesc<1x32xf16>" + (store,) = _ops(m, "tt.descriptor_store") + (red,) = _ops(m, "tt.descriptor_reduce") + (scatter,) = _ops(m, "tt.descriptor_scatter") + types = lambda op: [m.values[v].type for v in op.operands] # noqa: E731 + # ODS order (desc, src, indices...), printed `%desc[%i, %j], %src` + assert types(store) == [tile, "tensor<32x32xf16>", "i32", "i32"] + assert types(red) == types(store) + # descriptor_scatter prints in ODS order (desc, x_offsets, y_offset, src) + assert types(scatter) == [row, "tensor<32xi32>", "i32", "tensor<32x32xf16>"] + + +def test_dot_scaled_operand_order(): + m = W._walk(_text("kernel_dot_scaled.ttir")) + (d,) = _ops(m, "tt.dot_scaled") + # ODS (a, b, c, a_scale, b_scale), printed `%a scale %as, %b scale %bs, %c` + assert [m.values[v].type for v in d.operands] == [ + "tensor<128x64xf8E4M3FN>", + "tensor<64x128xf8E4M3FN>", + "tensor<128x128xf32>", + "tensor<128x2xi8>", + "tensor<128x2xi8>", + ] + assert [m.values[v].name for v in d.operands[3:]] == ["a_scale", "b_scale"] + + +def test_cf_blocks_align(): + m = W._walk(_text("golden_nested_guard_merge_sm80.ttir")) + (func,) = _ops(m, "tt.func") + labels = [m.blocks[b].label for b in func.regions[0]] + assert labels[0] is None and "^bb5" in labels and len(labels) == 6 + assert all(op.loc is not None for op in _ops(m, "cf.br") + _ops(m, "cf.cond_br")) + + +def test_parser_invented_locs_are_absent(): + m = W._walk(_text("spike_spin_while.ttir")) + whiles = _ops(m, "scf.while") + before = m.blocks[whiles[1].regions[0][0]] + assert before.arg_names == (None,) # `scf.while (%v_1 = %v)` prints no arg loc + for op in m.ops: + for loc in (op.loc, *op.callers): + assert loc is None or not loc.file.startswith( + ("/proc/self/fd", tempfile.gettempdir()) + ) + + +def test_callsite_and_fused_locs(): + m = W._walk(_text("crafted_locs.ttir")) + b = [op for op in m.ops if op.name == "arith.addi"][1] + assert b.loc == W.SourceLoc("y.py", 3, 4) and b.loc_name == "inner" + assert b.callers == (W.SourceLoc("k.py", 2, 5), W.SourceLoc("k.py", 3, 6)) + (store,) = _ops(m, "tt.store") + assert store.loc == W.SourceLoc("k.py", 2, 5) and store.callers == ( + W.SourceLoc("k.py", 2, 5), + ) + (ret,) = _ops(m, "tt.return") + assert ret.loc == W.SourceLoc("k.py", 12, 1) # alias defined after its use + + +def test_quoted_symbols_and_strings_with_syntax_chars(): + m = W._walk(_text("crafted_symbols_strings.ttir")) + assert [(f.sym_name, f.visibility) for f in m.funcs] == [ + ('f{%x} "q" (a)', "private"), + ("k", "public"), + ] + (call,) = _ops(m, "tt.call") + assert call.attrs["callee"] == 'f{%x} "q" (a)' and len(call.results) == 2 + (asm,) = _ops(m, "tt.elementwise_inline_asm") + assert asm.attrs["pure"] is True + + +def test_program_id_axes(): + m = W._walk(_text("golden_tile2d_sm80.ttir")) + assert sorted(op.attrs["axis"] for op in _ops(m, "tt.get_program_id")) == [0, 1] + + +def _const_module(lit: str, ty: str) -> str: + return ( + "module {\n" + " tt.func public @k() attributes {noinline = false} {\n" + f" %c = arith.constant {lit} : {ty}\n" + " tt.return\n" + " }\n" + "}\n" + ) + + +@pytest.mark.parametrize( + "lit, ty, value", + [ + # the parser accepts these and MLIR holds the bits: read as the + # printer prints them (signed; i1 as a bool) + ("4294967295", "i32", -1), + ("2147483648", "i32", -(2**31)), + ("255", "i8", -1), + ("dense<4294967295>", "tensor<4xi32>", -1), + ("18446744073709551615", "i64", -1), + ("dense<1>", "tensor<4xi1>", True), + ], +) +def test_constant_outside_the_printed_signed_range(lit, ty, value): + (c,) = _ops(W._walk(_const_module(lit, ty)), "arith.constant") + assert c.attrs["value"] == value + + +@pytest.mark.parametrize( + "lit, ty, value", + [ + ("-1", "i32", -1), + ("2147483647", "i32", 2**31 - 1), + ("-2147483648", "i32", -(2**31)), + ("-128", "i8", -128), + ("dense<-1>", "tensor<4xi32>", -1), + ("-9223372036854775808", "i64", -(2**63)), + ("9223372036854775807", "index", 2**63 - 1), + ], +) +def test_constant_in_the_printed_signed_range(lit, ty, value): + (c,) = _ops(W._walk(_const_module(lit, ty)), "arith.constant") + assert c.attrs["value"] == value + + +_READ_CONSTANTS = r""" +import sys +sys.path.insert(0, sys.argv[1]) +from tilelens.ir import _mlir_walk as W +for text in sys.argv[2:]: + (c,) = [op for op in W.walk_module(text).ops if op.name == "arith.constant"] + print(c.attrs["value"]) +""" + + +def test_a_constant_wider_than_64_bits_does_not_abort_the_process(): + # get_constant_value asserts on one (APInt::getSExtValue) and the + # assertion aborts the process; in a child, so a regression fails here + # instead of killing the test run + texts = [ + _const_module(str(2**127 - 1), "i128"), + _const_module(f"dense<{-(2**126)}>", "tensor<4xi128>"), + ] + r = subprocess.run( + [sys.executable, "-c", _READ_CONSTANTS, str(REPO), *texts], + capture_output=True, + text=True, + timeout=120, + ) + assert r.returncode == 0, (r.returncode, r.stderr[-2000:]) + assert r.stdout.splitlines() == ["('dense', None)", "('dense', None)"] + + +_FOR = """module { + tt.func public @k(%p: !tt.ptr, %n: i32) attributes {noinline = false} { + %c0 = arith.constant 0 : i32 + %c1 = arith.constant 1 : i32 + scf.for KW%i = %c0 to %n step %c1 : i32 { + %q = tt.addptr %p, %i : !tt.ptr, i32 + tt.store %q, %i : !tt.ptr + } + tt.return + } +} +""" + + +def test_scf_for_unsigned_compare_is_recorded(): + signed, unsigned = _FOR.replace("KW", ""), _FOR.replace("KW", "unsigned ") + (loop,) = _ops(W._walk(signed), "scf.for") + assert loop.attrs["unsignedCmp"] is False + (loop,) = _ops(W._walk(unsigned), "scf.for") + assert loop.attrs["unsignedCmp"] is True + # an unknown header keyword is a misread, never an ignored word + with pytest.raises(W.MisalignedModule, match="scf.for header"): + W._walk(unsigned, scan_text=_FOR.replace("KW", "signless ")) + + +def test_ttgir_is_refused(): + text = "#blocked = #ttg.blocked<{sizePerThread = [1], threadsPerWarp = [32], warpsPerCTA = [4], order = [0]}>\n" + text += _text("golden_add_sm80.ttir") + with pytest.raises((W.MisalignedModule, W.ModuleParseError)): + W._walk(text) + + +# ─────────────────────────── parse errors ─────────────────────────── + + +def test_parse_error_carries_the_diagnostic(capfd): + text = _text("golden_add_sm80.ttir") + lines = text.splitlines() + i = next(k for k, ln in enumerate(lines) if "tt.make_range {" in ln) + lines[i] = lines[i].replace("tt.make_range {", "tt.make_range_v2 {") + with pytest.raises(W.ModuleParseError) as e: + W.walk_module("\n".join(lines)) + assert ( + "tt.make_range_v2" in e.value.diagnostic and "is unknown" in e.value.diagnostic + ) + assert "" in e.value.diagnostic and "/proc/self/fd" not in e.value.diagnostic + assert e.value.line_no == i + 1 + out, err = capfd.readouterr() + assert "make_range_v2" not in err # captured, not leaked onto the process stderr + + +# ─────────────────────────── the release, barriers, block pointers ─────────────────────────── + + +def _barrier_module(barrier: str) -> str: + return ( + "module {\n" + " tt.func public @k(%p: !tt.ptr) attributes {noinline = false} {\n" + f" {barrier}\n" + " %v = tt.load %p : !tt.ptr\n" + " tt.return\n" + " }\n" + "}\n" + ) + + +def test_the_barrier_aligns(): + m = W._walk(_barrier_module("ttg.barrier all")) + (barrier,) = _ops(m, "ttg.barrier") + assert not barrier.results and barrier.loc is None + # a combination of address spaces prints as `a|b`; the text never reads it + assert _ops( + W._walk(_barrier_module("ttg.barrier local|global_read")), "ttg.barrier" + ) + # gpu.barrier still parses (the reader refuses it, see test_ttir_reader) + assert _ops(W._walk(_barrier_module("gpu.barrier")), "gpu.barrier") + + +def test_another_triton_release_is_refused_before_any_parse(monkeypatch): + import triton + + monkeypatch.setattr(W, "_walk", lambda *a: pytest.fail("walked")) + monkeypatch.setattr(triton, "__version__", "3.6.0") + with pytest.raises(W.ModuleParseError, match="Triton 3.8.x only.*is 3.6.0"): + W.walk_module(_barrier_module("")) + + +_SPLIT_NEWLINE = "!tt.ptr<\n tensor<32xf32>>" +_SPLIT_COMMENT = "!tt.ptr< // c\n tensor<32xf32>>" + +_BLOCK_PTR = """module { + tt.func public @k(%p: !tt.ptr, %b: TYPE) attributes {noinline = false} { + tt.return + } +} +""" + + +@pytest.mark.parametrize( + "text, line", + [ + (_BLOCK_PTR.replace("TYPE", "!tt.ptr>"), 2), + (_BLOCK_PTR.replace("TYPE", "!tt.ptr, 1>"), 2), + (_BLOCK_PTR.replace("TYPE", "!tt.ptr< tensor <4xf32>>"), 2), + (_BLOCK_PTR.replace("TYPE", "!tt>>"), 2), + (_BLOCK_PTR.replace("TYPE", "!tt.ptr>>"), 2), + # split over lines, or around a comment: no single line holds it + (_BLOCK_PTR.replace("TYPE", _SPLIT_NEWLINE), 2), + (_BLOCK_PTR.replace("TYPE", _SPLIT_COMMENT), 2), + (_BLOCK_PTR.replace("TYPE", "!tt.ptr\n>"), 2), + (_BLOCK_PTR.replace("TYPE", '!tt.ptr< // "q>\n // b\n tensor<4xf32>>'), 2), + (_BLOCK_PTR.replace("TYPE", "!tt>>"), 2), + # a type alias may hide the pointee + ("!t = tensor<4xf32>\n" + _BLOCK_PTR.replace("TYPE", "!tt.ptr"), 1), + ("!t // c\n = tensor<4xf32>\n" + _BLOCK_PTR.replace("TYPE", "!tt.ptr"), 1), + ( + "// c\n\n !t\n = tensor<4xf32>\n" + + _BLOCK_PTR.replace("TYPE", "!tt.ptr"), + 3, + ), + ], +) +def test_block_pointer_types_are_refused_before_the_parser(monkeypatch, text, line): + """The parser aborts the process on a block-pointer type.""" + + def no_parse(_data): + raise AssertionError("the bindings got a block-pointer type") + + monkeypatch.setattr(W, "_bind_walk", no_parse) + with pytest.raises(W.ModuleParseError, match="no block pointers") as e: + W._walk(text) + assert e.value.line_no == line and f"line {line}:" in e.value.diagnostic + + +def test_the_block_pointer_screen_reads_types_not_strings(): + text = _module_with_print('"ptr>")) + W._screen(_module_with_print('"ptr< // "')) + in_string = _BLOCK_PTR.replace("TYPE", "i32").replace( + "{noinline = false}", + '{noinline = false, s = "a // b", t = !tt.ptr<\n tensor<4xf32>>}', + ) + with pytest.raises(W.ModuleParseError, match="line 2: a block-pointer type"): + W._screen(in_string) + unterminated = text.replace('"ptr str: + return ( + "module {\n" + " tt.func public @k(%x: i32) attributes {noinline = false} {\n" + f" tt.print {prefix} {{hex = false, isSigned = array}} : %x : i32\n" + " tt.return\n" + " }\n" + "}\n" + ) + + +# ─────────────────────────── the cache ─────────────────────────── + + +def test_cache_returns_the_same_module_and_caches_misalignment(monkeypatch): + text = _text("golden_add_sm80.ttir") + assert W.walk_module(text) is W.walk_module(text) + # a comment on an op line is a text the scan does not read + bad = text.replace(" loc(#loc1)", " loc(#loc1) // c", 1) + with pytest.raises(W.MisalignedModule) as first: + W.walk_module(bad) + + def no_parse(_data): + raise AssertionError("parsed twice") + + monkeypatch.setattr(W, "_bind_walk", no_parse) + with pytest.raises(W.MisalignedModule) as again: + W.walk_module(bad) + assert ( + again.value.problems == first.value.problems and again.value is not first.value + ) + W.walk_module(text) diff --git a/tests/unit/ir/test_ttir_reader.py b/tests/unit/ir/test_ttir_reader.py new file mode 100644 index 00000000..a5cd6f1a --- /dev/null +++ b/tests/unit/ir/test_ttir_reader.py @@ -0,0 +1,889 @@ +"""tilelens.ir.ttir_reader: the TTIR -> AccessGraph reader on top of _mlir_walk. + +The corpus is tests/unit/ir/ttir_corpus.py, compiled at test time with +Triton 3.8: ttir/ (the walk layer's corpus) and reader_ttir/ (the kernels +of tests/unit/ir/reader_kernels.py). +""" + +from __future__ import annotations + +import itertools +import pickle +from pathlib import Path + +import pytest + +from tilelens.ir import ParseCache, Refusal, SourceLocation +from tilelens.ir import _mlir_walk as W +from tilelens.ir import ttir_reader as R +from tilelens.ir.ttir_reader import ( + AccessGraph, + Arange, + Bin, + Cmp, + Const, + IntCast, + IterArgOffset, + Param, + Pid, + Select, + TTIRKind, + UnsupportedTTIR, + mentions_observed, + parse_ttir, + width_obligations, +) + +from . import ttir_corpus + +KERNELS = Path(__file__).resolve().parent / "reader_kernels.py" +FILES = ttir_corpus.NAMES + + +def _text(name: str) -> str: + return ttir_corpus.text(name) + + +def _graph(name: str) -> AccessGraph: + return parse_ttir(_text(name)) + + +def _refusal(name_or_text: str) -> UnsupportedTTIR: + text = _text(name_or_text) if name_or_text.endswith(".ttir") else name_or_text + with pytest.raises(UnsupportedTTIR) as info: + parse_ttir(text) + return info.value + + +def _module(body: str, args: str = "%p: !tt.ptr, %n: i32", extra: str = "") -> str: + """A minimal TTIR module (no locs) around ``body``'s op lines.""" + lines = "\n ".join(line.strip() for line in body.strip().splitlines()) + return ( + f"module {{\n tt.func public @k({args}) attributes {{noinline = false}} {{\n" + f" {lines}\n tt.return\n }}\n{extra}}}\n" + ) + + +# Where a refusal's loc points: the frontend locates an op at its own AST +# node (scf.yield at the loop header, a call spanning lines at its first line). +_REFUSAL_SITE = { + "p1_variant_delta": "for k in range(-8, n):", + "loop_observed_advance": "for i in range(0, n):", + "rv_inline_asm_store": "tl.inline_asm_elementwise(", +} + + +def _source_line(loc) -> str: + """The reader_kernels.py source line a probe text's loc points at.""" + assert loc is not None and loc.file.endswith("reader_kernels.py"), loc + return KERNELS.read_text(encoding="utf-8").splitlines()[loc.line - 1] + + +@pytest.fixture(autouse=True) +def _fresh_cache(): + W._CACHE.clear() + yield + W._CACHE.clear() + + +# ─────────────────────────── refusal form ─────────────────────────── + + +def test_kinds_are_the_representational_ones(): + assert {k.value for k in TTIRKind} == { + "indirect-address", + "data-dependent-bound", + "nested-loop", + "control-flow", + "out-of-vocabulary", + "call", + "loop-variant-advance", + "inline-asm", + "reader-misalignment", + "unparsable", + "other", + } + # a kind is its string, also when formatted + assert ( + TTIRKind.CALL == "call" and f"{TTIRKind.CALL}" == str(TTIRKind.CALL) == "call" + ) + + +def test_unsupported_ttir_carries_structured_fields(): + loc = W.SourceLoc("k.py", 3, 4) + e = UnsupportedTTIR("control-flow", "scf.while", line_no=7, loc=loc) + assert e.kind is TTIRKind.CONTROL_FLOW + assert (e.message, e.line_no, e.loc, str(e)) == ("scf.while", 7, loc, "scf.while") + back = pickle.loads(pickle.dumps(e)) + assert (back.kind, back.message, back.line_no, back.loc) == ( + e.kind, + e.message, + 7, + loc, + ) + # client-owned kinds are not reader kinds + for kind in ("spin-shape", "cas-value", "data-dependent-mask"): + with pytest.raises(ValueError): + UnsupportedTTIR(kind, "no") + # the verdict record copies the fields, no string round-trip; the + # private walk loc becomes the public SourceLocation (a saved trace holds it) + r = Refusal.from_exception(e) + assert (r.kind, r.message, r.line_no, r.loc) == ( + "control-flow", + "scf.while", + 7, + SourceLocation("k.py", 3, 4), + ) + assert type(r.loc) is SourceLocation + + +def test_parse_cache_reads_with_this_reader(): + cache = ParseCache() + parsed = cache.get(_text("ttir/golden_add_sm80.ttir")) + assert isinstance(parsed.graph, AccessGraph) and parsed.error is None + refused = cache.get(_text("reader_ttir/p2_swap.ttir")) + assert isinstance(refused.refusal, UnsupportedTTIR) and refused.error is None + assert refused.refusal.kind is TTIRKind.LOOP_VARIANT_ADVANCE + + +def test_walk_failures_map_to_kinds(): + text = _text("ttir/golden_add_sm80.ttir") + commented = text.replace(" loc(#loc1)", " loc(#loc1) // c", 1) + e = _refusal(commented) + assert e.kind is TTIRKind.READER_MISALIGNMENT and "unexpected comment" in e.message + assert "// c" in commented.splitlines()[e.line_no - 1] + e = _refusal("module {\n tt.func public @k( {\n}\n") + assert e.kind is TTIRKind.UNPARSABLE + assert e.line_no == 2 and "" in e.message + + +def test_refusals_carry_the_op_line_and_source_loc(): + e = _refusal("reader_ttir/p3_call_offset.ttir") + assert e.kind is TTIRKind.CALL and "_p3_helper" in e.message + assert ( + "tt.call" + in _text("reader_ttir/p3_call_offset.ttir").splitlines()[e.line_no - 1] + ) + assert "_p3_helper(" in _source_line(e.loc) + + +# ─────────────────────────── every corpus text's outcome ─────────────────────────── + +# The reader's refusal kind for every refused corpus text (the rest parse). +REFUSED = { + "ttir/adv_descs.ttir": "out-of-vocabulary", + "ttir/adv_views.ttir": "other", # tt.reshape of a pointer tile + "ttir/crafted_attr_dicts.ttir": "other", + "ttir/crafted_same_dest.ttir": "control-flow", + "ttir/crafted_symbols_strings.ttir": "call", + "ttir/golden_gather_sm80.ttir": "indirect-address", + "ttir/golden_nested_guard_merge_sm80.ttir": "control-flow", + "ttir/spike_inline_asm.ttir": "inline-asm", + "ttir/spike_spin_while.ttir": "control-flow", + "reader_ttir/p1_variant_delta.ttir": "loop-variant-advance", + "reader_ttir/loop_observed_advance.ttir": "loop-variant-advance", + "reader_ttir/p2_swap.ttir": "loop-variant-advance", + "reader_ttir/p3_call_guarded.ttir": "call", + "reader_ttir/p3_call_offset.ttir": "call", + "reader_ttir/p3_call_formals.ttir": "call", + "reader_ttir/rv_inline_asm_store.ttir": "inline-asm", + "reader_ttir/int_iterarg_offset.ttir": "loop-variant-advance", + "reader_ttir/pure_asm_int_addr.ttir": "inline-asm", + "reader_ttir/observed_lanes.ttir": "indirect-address", +} + + +@pytest.mark.parametrize("name", FILES) +def test_each_text_reads_as_pinned(name): + try: + parse_ttir(_text(name)) + except UnsupportedTTIR as e: + assert e.kind == REFUSED.get(name), e.message + else: + assert name not in REFUSED + + +# ─────────────────────────── soundness edge cases ─────────────────────────── + + +def test_p1_loop_variant_advance_refuses(): + e = _refusal("reader_ttir/p1_variant_delta.ttir") + assert e.kind is TTIRKind.LOOP_VARIANT_ADVANCE and "loop-variant" in e.message + assert _REFUSAL_SITE["p1_variant_delta"] in _source_line(e.loc) + + +def test_p2_swapped_iter_args_refuse(): + e = _refusal("reader_ttir/p2_swap.ttir") + assert ( + e.kind is TTIRKind.LOOP_VARIANT_ADVANCE + and "not advanced from itself" in e.message + ) + + +def test_observation_inside_the_loop_is_a_variant_advance(): + e = _refusal("reader_ttir/loop_observed_advance.ttir") + assert e.kind is TTIRKind.LOOP_VARIANT_ADVANCE + assert _REFUSAL_SITE["loop_observed_advance"] in _source_line(e.loc) + + +def test_quoted_callee_and_a_second_function_refuse_as_call(): + e = _refusal("ttir/crafted_symbols_strings.ttir") + assert e.kind is TTIRKind.CALL and "'f{%x} \"q\" (a)'" in e.message + # a second tt.func refuses even when nothing calls it + extra = ( + " tt.func private @h(%q: !tt.ptr) attributes {noinline = true} {\n" + " tt.return\n }\n" + ) + e = _refusal(_module("", extra=extra)) + assert e.kind is TTIRKind.CALL and "'h'" in e.message + + +def test_p4_walkers_resolve_loop_carried_pointers(): + direct = _graph("reader_ttir/p4_observed_direct.ttir") + loop = _graph("reader_ttir/p4_observed_loop.ttir") + delta = _graph("reader_ttir/p4_observed_delta.ttir") + for g in (direct, loop, delta): + load = next(a for a in g.accesses if a.kind == "load") + assert mentions_observed(load.offset, g) + # the address reaches the observation only through the iter_arg (its + # offset0, resp. its delta), which a term-local walker misses + for g, part in ((loop, "offset0"), (delta, "delta")): + load = next(a for a in g.accesses if a.kind == "load") + assert isinstance(load.offset, IterArgOffset) + assert mentions_observed(getattr(g.iter_args[0], part), g) + + +def test_rv_inline_asm_refuses(): + e = _refusal("reader_ttir/rv_inline_asm_store.ttir") + assert e.kind is TTIRKind.INLINE_ASM and "side effects" in e.message + assert _REFUSAL_SITE["rv_inline_asm_store"] in _source_line(e.loc) + # a pure asm with a pointer operand is still an address handed to asm + text = _text("ttir/spike_inline_asm.ttir").replace("pure = false", "pure = true") + e = _refusal(text) + assert e.kind is TTIRKind.INLINE_ASM and "pointer operand" in e.message + # a pure asm over data is plain data + text = "\n".join( + line + for line in _text("ttir/spike_inline_asm.ttir").splitlines() + if "pure = false" not in line + ) + assert [a.kind for a in parse_ttir(text).accesses] == ["load"] + + +def _footprint(g, access, params: dict, iters: int) -> set[int]: + """The modeled element offsets of ``access`` over iterations + ``0 .. iters-1`` with every arange lane free, the lanes keyed by + (make_range, dim) as a consumer keys them; mask applied.""" + lanes = sorted( + { + (n.ssa, n.dim, n.start, n.end) + for n in R._nodes((access.offset, access.mask), g.iter_args) + if isinstance(n, Arange) + } + ) + + def ev(t, env): + if isinstance(t, Const): + return t.value + if isinstance(t, Param): + return env[t.name] + if isinstance(t, Arange): + return env[(t.ssa, t.dim)] + if isinstance(t, IterArgOffset): + info = g.iter_args[t.arg_id] + return ev(info.offset0, env) + env["k"] * ev(info.delta, env) + if isinstance(t, IntCast): + return ev(t.x, env) + a, b = ev(t.a, env), ev(t.b, env) + if isinstance(t, Cmp): + return int({"slt": a < b}[t.pred]) + return {"+": a + b, "-": a - b, "*": a * b}[t.op] + + out = set() + for k in range(iters): + for values in itertools.product(*(range(s, e) for _, _, s, e in lanes)): + env = dict(params, k=k) + env.update({(ssa, dim): v for (ssa, dim, _, _), v in zip(lanes, values)}) + if access.mask is None or ev(access.mask, env): + out.add(ev(access.offset, env)) + return out + + +def test_expanded_loop_carried_tiles_keep_their_lanes(): + # [N, N] pointer tile expanded to 3D in the loop, next to the same + # make_range at dim 0 + g = _graph("reader_ttir/expand_iterarg_3d.ttir") + load = g.accesses[0] + assert load.kind == "load" and load.in_loop + (tile,) = [ + n for n in R._nodes((load.offset,), None) if isinstance(n, IterArgOffset) + ] + expanded = g.iter_args[tile.arg_id] + assert expanded.base_param == "x_ptr" and expanded.delta == Const(1) + n = 4 + want = { + j * n + m - i * n + k + for i, j, m in itertools.product(range(n), repeat=3) + for k in range(2) + } + assert _footprint(g, load, {}, 2) == want # [-12, 16], every OOB offset kept + # 1D pointer expanded to 2D, masked at its own lane (dim 1) + g = _graph("reader_ttir/expand_iterarg_mask.ttir") + load = g.accesses[0] + assert load.offset == IterArgOffset(1) and g.iter_args[1].delta == Const(4) + assert _footprint(g, load, {"M": 2}, 2) == {0, 1, 4, 5} + + +def test_tensor_observations_do_not_meet_across_lanes(): + # c[:, None] - c[None, :] over one tensor atomic's old values + e = _refusal("reader_ttir/observed_lanes.ttir") + assert e.kind is TTIRKind.INDIRECT_ADDRESS and "across lanes" in e.message + assert "c[:, None] - c[None, :]" in _source_line(e.loc) + # an expanded tile of a pointer that advances by per-lane observations + body = """ + %c0 = arith.constant 0 : i32 + %c1 = arith.constant 1 : i32 + %r = tt.make_range {end = 4 : i32, start = 0 : i32} : tensor<4xi32> + %ps = tt.splat %p : !tt.ptr -> tensor<4x!tt.ptr> + %q = tt.addptr %ps, %r : tensor<4x!tt.ptr>, tensor<4xi32> + %one = arith.constant dense<1> : tensor<4xi32> + %old = tt.atomic_rmw add, acq_rel, gpu, %q, %one : (tensor<4x!tt.ptr>, tensor<4xi32>) -> tensor<4xi32> + %res = scf.for %k = %c0 to %n step %c1 iter_args(%a = %q) -> (tensor<4x!tt.ptr>) : i32 { + %e = tt.expand_dims %a {axis = 0 : i32} : tensor<4x!tt.ptr> -> tensor<1x4x!tt.ptr> + %v = tt.load %e : tensor<1x4x!tt.ptr> + %a2 = tt.addptr %a, %old : tensor<4x!tt.ptr>, tensor<4xi32> + scf.yield %a2 : tensor<4x!tt.ptr> + }""" + e = _refusal(_module(body, args="%p: !tt.ptr, %n: i32")) + assert e.kind is TTIRKind.OTHER and "per-lane atomic result" in e.message + + +def test_loop_carried_integers_make_addresses_loop_variant(): + # offs += B carried by the loop + e = _refusal("reader_ttir/int_iterarg_offset.ttir") + assert ( + e.kind is TTIRKind.LOOP_VARIANT_ADVANCE and "carried by the loop" in e.message + ) + # a pointer advanced by a loop-carried integer + body = """ + %c0 = arith.constant 0 : i32 + %c1 = arith.constant 1 : i32 + %r:2 = scf.for %i = %c0 to %n step %c1 iter_args(%a = %p, %s = %c0) -> (!tt.ptr, i32) : i32 { + %v = tt.load %a : !tt.ptr + %a2 = tt.addptr %a, %s : !tt.ptr, i32 + %s2 = arith.addi %s, %c1 : i32 + scf.yield %a2, %s2 : !tt.ptr, i32 + }""" + assert _refusal(_module(body)).kind is TTIRKind.LOOP_VARIANT_ADVANCE + + +def test_an_address_handed_to_opaque_ops_as_an_integer_refuses(): + # tl.inline_asm_elementwise(..., is_pure=True) given x_ptr.to(tl.int64) + e = _refusal("reader_ttir/pure_asm_int_addr.ttir") + assert e.kind is TTIRKind.INLINE_ASM and "tt.ptr_to_int" in e.message + # through arithmetic with loaded data, and through a loop-carried value + args = "%p: !tt.ptr, %n: i32" + head = """ + %x = tt.load %p : !tt.ptr + %i = tt.ptr_to_int %p : !tt.ptr -> i64 + %a = arith.addi %i, %x : i64""" + asm = ( + '%y = tt.elementwise_inline_asm "mov.b64 $0, $1;" {constraints = "=l,l", ' + "packed_element = 1 : i32, pure = true} %a : i64 -> i64" + ) + assert _refusal(_module(head + "\n" + asm, args=args)).kind is TTIRKind.INLINE_ASM + loop = """ + %c0 = arith.constant 0 : i32 + %c1 = arith.constant 1 : i32 + %z = arith.constant 0 : i64 + %i = tt.ptr_to_int %p : !tt.ptr -> i64 + %a = scf.for %k = %c0 to %n step %c1 iter_args(%s = %z) -> (i64) : i32 { + %s2 = arith.addi %s, %i : i64 + scf.yield %s2 : i64 + }""" + assert _refusal(_module(loop + "\n" + asm, args=args)).kind is TTIRKind.INLINE_ASM + extern = ( + '%y = tt.extern_elementwise %a {libname = "", libpath = "", pure = true, ' + 'symbol = "f"} : (i64) -> i64' + ) + e = _refusal(_module(head + "\n" + extern, args=args)) + assert e.kind is TTIRKind.OUT_OF_VOCABULARY and "tt.ptr_to_int" in e.message + # loaded data carries no address: an asm over it stays plain data + data = head.replace("%i, %x", "%x, %x") + "\n" + asm + assert [a.kind for a in parse_ttir(_module(data, args=args)).accesses] == ["load"] + + +def test_llvm_and_gpu_ops_outside_the_inert_ones_refuse(): + # memory effects through ops of the llvm dialect + cases = { + "llvm.inttoptr": """ + %i = tt.ptr_to_int %p : !tt.ptr -> i64 + %lp = llvm.inttoptr %i : i64 to !llvm.ptr<1> + %c = arith.constant 7 : i32 + %old = llvm.atomicrmw add %lp, %c monotonic : !llvm.ptr<1>, i32""", + "llvm.inline_asm": """ + %i = tt.ptr_to_int %p : !tt.ptr -> i64 + %c = arith.constant 7 : i32 + %r = llvm.inline_asm has_side_effects "st.global.b32 [$1], $2; mov.b32 $0, 0;", "=r,l,r" %i, %c : (i64, i32) -> i32""", + } + for name, body in cases.items(): + e = _refusal(_module(body, args="%p: !tt.ptr, %n: i32")) + assert e.kind is TTIRKind.OUT_OF_VOCABULARY and name in e.message + # inside a combine region too + body = """ + %r = tt.make_range {end = 4 : i32, start = 0 : i32} : tensor<4xi32> + %s = "tt.reduce"(%r) <{axis = 0 : i32}> ({ + ^bb0(%a: i32, %b: i32): + %i = tt.ptr_to_int %p : !tt.ptr -> i64 + %lp = llvm.inttoptr %i : i64 to !llvm.ptr<1> + %y = arith.addi %a, %b : i32 + tt.reduce.return %y : i32 + }) : (tensor<4xi32>) -> i32""" + e = _refusal(_module(body, args="%p: !tt.ptr, %n: i32")) + assert e.kind is TTIRKind.OUT_OF_VOCABULARY and "llvm.inttoptr" in e.message + # the inert ones stay inert: ttg.barrier is the one ttg op accepted; + # gpu.barrier still parses, but the frontend never emits it: refused + g = parse_ttir(_module("ttg.barrier all\n%v = tt.load %p : !tt.ptr")) + assert [a.kind for a in g.accesses] == ["load"] + e = _refusal(_module("gpu.barrier\n%v = tt.load %p : !tt.ptr")) + assert e.kind is TTIRKind.OUT_OF_VOCABULARY, e.message + assert "op gpu.barrier is not TTIR" in e.message + e = _refusal(_module("ttg.local_barrier\n%v = tt.load %p : !tt.ptr")) + assert e.kind in (TTIRKind.UNPARSABLE, TTIRKind.OUT_OF_VOCABULARY) + + +def test_pointers_outside_global_memory_refuse(): + # a shared-memory (address space 3) pointer argument + e = _refusal( + _module( + "%c1 = arith.constant 1 : i32\ntt.store %p, %c1 : !tt.ptr", + args="%p: !tt.ptr, %n: i32", + ) + ) + assert e.kind is TTIRKind.OUT_OF_VOCABULARY and "address space" in e.message + # a pointer made into another address space + body = """ + %c = arith.constant 0 : i64 + %q = tt.int_to_ptr %c : i64 -> !tt.ptr + %v = tt.load %q : !tt.ptr""" + e = _refusal(_module(body)) + assert e.kind is TTIRKind.OUT_OF_VOCABULARY and "address space" in e.message + assert e.line_no == 4 # the tt.int_to_ptr + + +def test_generic_load_reads_its_mask_by_operand_segment(): + # generic tt.load with `other` and no mask; an i1 `other` is no mask + load = ( + '%v = "tt.load"(%p, %f) <{boundaryCheck = array, cache = 1 : i32, ' + "evict = 1 : i32, isVolatile = false, operandSegmentSizes = array}> " + ": (!tt.ptr, i1) -> i1" + ) + for segments, mask in (("1, 0, 1", None), ("1, 1, 0", Const(0))): + g = parse_ttir( + _module( + "%f = arith.constant false\n" + load.replace("SEG", segments), + args="%p: !tt.ptr, %n: i32", + ) + ) + (access,) = g.accesses + assert (access.mask, access.mask_dropped) == (mask, False), segments + + +# ─────────────────────────── modeled shapes ─────────────────────────── + + +def test_two_step_advance_is_the_sum(): + g = _graph("reader_ttir/loop_two_step_advance.ttir") + (info,) = g.iter_args + assert (info.base_param, info.offset0, info.delta) == ( + "x_ptr", + Const(0), + Bin("+", Param("s"), Const(2)), + ) + assert g.accesses[0].offset == IterArgOffset(0) and g.accesses[0].in_loop + + +def test_same_base_pointer_selects(): + g = _graph("reader_ttir/where_pointer.ttir") + (store,) = g.accesses + assert store.base_param == "x_ptr" and isinstance(store.offset, Select) + g = _graph("ttir/spike_if_yield.ttir") + store = g.accesses[-1] + assert store.kind == "store" and store.base_param == "out_ptr" + assert isinstance(store.offset, Select) and isinstance(store.offset.cond, Cmp) + + +def test_three_dims_of_one_make_range_are_three_lanes(): + (store,) = _graph("reader_ttir/tile3d_shared_arange.ttir").accesses + ranges = [n for n in R._nodes((store.offset,), None) if isinstance(n, Arange)] + assert len({r.ssa for r in ranges}) == 1 + assert sorted(r.dim for r in ranges) == [0, 1, 2] + + +def test_same_width_pointer_bitcast_keeps_the_base(): + g = _graph("ttir/golden_atomic_fmax_sm80.ttir") + atomics = [a for a in g.accesses if a.kind == "atomic_rmw"] + assert len(atomics) == 2 + assert all(a.base_param == "out_ptr" and a.elem_bits == 32 for a in atomics) + # the atomics' mask reads loaded data + assert all(a.mask_dropped and a.mask is None for a in atomics) + # a bitcast to another element width changes what an element offset means + body = """ + %q = tt.bitcast %p : !tt.ptr -> !tt.ptr + %v = tt.load %q : !tt.ptr""" + e = _refusal(_module(body)) + assert e.kind is TTIRKind.OTHER and "element width" in e.message + + +def test_other_refusals(): + cases = { + # a block pointer: the dialect has none (its parser aborts on the + # type), so the walk refuses the text before parsing + "unparsable": """ + %c0 = arith.constant 0 : i32 + %c1 = arith.constant 1 : i64 + %c32 = arith.constant 32 : i64 + %bp = tt.make_tensor_ptr %p, [%c32, %c32], [%c32, %c1], [%c0, %c0] {order = array} : > + %v = tt.load %bp : !tt.ptr>""", + "out-of-vocabulary": """ + %r = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32> + %v = arith.sitofp %r : tensor<64xi32> to tensor<64xf32> + %s = "tt.reduce"(%v) <{axis = 0 : i32}> ({ + ^bb0(%a: f32, %b: f32): + %x = tt.load %p : !tt.ptr + %y = arith.addf %a, %x : f32 + tt.reduce.return %y : f32 + }) : (tensor<64xf32>) -> f32""", + "control-flow": "%c0 = arith.constant 0 : i32\n" + "%b = arith.cmpi slt, %n, %c0 : i32\n" + 'cf.assert %b, "neg"', + "nested-loop": """ + %c0 = arith.constant 0 : i32 + %c1 = arith.constant 1 : i32 + scf.for %i = %c0 to %n step %c1 : i32 { + scf.for %j = %i to %n step %c1 : i32 { + %q = tt.addptr %p, %j : !tt.ptr, i32 + %v = tt.load %q : !tt.ptr + } + }""", + } + for kind, body in cases.items(): + assert _refusal(_module(body)).kind == kind, kind + e = _refusal(_module(cases["unparsable"])) + assert "no block pointers" in e.message and e.line_no == 7 + e = _refusal( + _module( + "%x = tt.load %p : !tt.ptr\n%y = tt.extern_elementwise %x " + '{libname = "", libpath = "", pure = false, symbol = "foo"} : (f32) -> f32' + ) + ) + assert e.kind is TTIRKind.OUT_OF_VOCABULARY and "'foo'" in e.message + # deep structured nesting refuses instead of exhausting the recursion limit + depth = 250 + nest = ( + "".join("scf.if %c {\n" for _ in range(depth)) + + "tt.store %p, %f : !tt.ptr\n" + + "}\n" * depth + ) + e = _refusal( + _module( + "%f = arith.constant 1.0 : f32\n" + nest, args="%p: !tt.ptr, %c: i1" + ) + ) + assert e.kind is TTIRKind.OTHER and "nested deeper" in e.message + + +def test_pure_extern_and_signed_i1_compare_are_data(): + g = parse_ttir( + _module( + "%x = tt.load %p : !tt.ptr\n%y = tt.extern_elementwise %x " + '{libname = "", libpath = "", pure = true, symbol = "__nv_expf"} : (f32) -> f32\n' + "tt.store %p, %y : !tt.ptr" + ) + ) + assert [a.kind for a in g.accesses] == ["load", "store"] + # a signed compare of i1 reads true as -1: not the boolean model, so the + # mask is dropped (widened), never misread + g = parse_ttir( + _module( + """ + %c0 = arith.constant 0 : i32 + %b = arith.cmpi slt, %n, %c0 : i32 + %t = arith.constant true + %c = arith.cmpi slt, %b, %t : i1 + %v = tt.load %p, %c : !tt.ptr""" + ) + ) + (load,) = g.accesses + assert load.mask is None and load.mask_dropped + + +def test_matmul_loop_iter_args_and_params(): + g = _graph("ttir/golden_matmul_s3_sm80.ttir") + assert g.kernel_name == "matmul_kernel" and g.loop is not None + assert [i.base_param for i in g.iter_args] == ["a_ptr", "b_ptr"] + assert all( + i.arg_id == k and i.loop_ssa == g.loop.loop_ssa + for k, i in enumerate(g.iter_args) + ) + assert Param("K") in list(R._nodes((g.loop.upper,), None)) + loads = [a for a in g.accesses if a.kind == "load"] + assert [a.offset for a in loads] == [IterArgOffset(0), IterArgOffset(1)] + assert all(a.in_loop for a in loads) and not g.accesses[-1].in_loop + assert g.loop.bits == 32 and not g.loop.unsigned and g.loop.line_no is not None + assert g.arg("K").int_bits == 32 and g.arg("a_ptr").elem_bits == 16 + + +# ─────────────────────────── integer widths ─────────────────────────── + + +def _value(t, env: dict) -> int: + """The unbounded-integer reading of a (shallow) term: casts as identity.""" + if isinstance(t, Const): + return t.value + if isinstance(t, Pid): + return env["pid"] + if isinstance(t, Param): + return env[t.name] + if isinstance(t, IntCast): + return _value(t.x, env) + if isinstance(t, Cmp): + a, b = _value(t.a, env), _value(t.b, env) + return int({"slt": a < b, "ult": a < b, "eq": a == b}[t.pred]) + if isinstance(t, Bin): + a, b = _value(t.a, env), _value(t.b, env) + ops = {"+": a + b, "*": a * b, "-": a - b} + if t.op in ("u//", "//"): + return int(a / b) + if t.op == "%": + return a - b * int(a / b) + return ops[t.op] + raise AssertionError(type(t).__name__) + + +def _holds(ob, env: dict) -> bool: + v = _value(ob.term, env) + if ob.signed: + return -(1 << (ob.bits - 1)) <= v < (1 << (ob.bits - 1)) + return 0 <= v < (1 << ob.bits) + + +def _integer_nodes(g): + for a in g.accesses: + yield from R._nodes((a.offset, a.mask, a.path), g.iter_args) + if g.loop is not None: + yield from R._nodes((g.loop.lower, g.loop.upper, g.loop.step), None) + + +@pytest.mark.parametrize("name", [f for f in FILES if f not in REFUSED]) +def test_every_integer_op_carries_its_width(name): + for n in _integer_nodes(_graph(name)): + if isinstance(n, Bin): + # bits=None only for the element-offset sums of tt.addptr + assert isinstance(n.bits, int) or (n.bits is None and n.op == "+") + elif isinstance(n, Cmp): + assert isinstance(n.bits, int) and n.bits >= 1 + elif isinstance(n, IntCast): + assert n.kind in ("trunci", "extsi", "extui") and n.src_bits != n.dst_bits + + +def test_i32_wrap_obligations(): + g = _graph("reader_ttir/rv_i32_wrap.ttir") + (store,) = g.accesses + inner = Bin("*", Pid(0), Param("S"), 32) + outer = Bin("*", inner, Param("S"), 32) + assert store.offset == Bin("+", Const(0), outer) + obs = width_obligations(g, store) + assert [(o.term, o.bits, o.signed) for o in obs] == [ + (outer, 32, True), + (inner, 32, True), + ] + assert all("pid * S" in _source_line(o.loc) for o in obs) + text = _text("reader_ttir/rv_i32_wrap.ttir").splitlines() + assert all("arith.muli" in text[o.line_no - 1] for o in obs) + # (pid * 65536) * 65536 is 0 in i32: the unbounded model is exact only for pid 0 + assert all(_holds(o, {"pid": 0, "S": 65536}) for o in obs) + assert [_holds(o, {"pid": 1, "S": 65536}) for o in obs] == [False, True] + + +def test_trunci_obligation(): + g = _graph("reader_ttir/rv_trunci_alias.ttir") + (store,) = g.accesses + wide = Bin("*", IntCast("extsi", 32, 64, Pid(0)), Const(1 << 32), 64) + assert store.offset == Bin("+", Const(0), IntCast("trunci", 64, 32, wide)) + obs = width_obligations(g, store) + assert [(o.term, o.bits, o.signed) for o in obs] == [ + (wide, 32, True), + (wide, 64, True), + ] + trunc = obs[0] + assert ( + "arith.trunci" + in _text("reader_ttir/rv_trunci_alias.ttir").splitlines()[trunc.line_no - 1] + ) + assert ".to(tl.int32)" in _source_line(trunc.loc) + # trunc_i32(pid * 2**32) is 0 for every pid: exact only for pid 0 + assert _holds(trunc, {"pid": 0}) and not _holds(trunc, {"pid": 1}) + assert _holds(obs[1], {"pid": 1}) + + +def test_unsigned_reads_need_non_negative_operands(): + g = _graph("reader_ttir/unsigned_index.ttir") + (store,) = g.accesses + quotient = Bin("u//", Pid(0), Const(3), 32) + assert store.offset == Bin("+", Const(0), IntCast("extui", 32, 64, quotient)) + assert store.mask == Cmp("ult", Pid(0), Param("n"), 32) + obs = {(o.term, o.bits, o.signed) for o in width_obligations(g, store)} + assert obs == { + (quotient, 32, False), # the extui operand + (quotient, 32, True), # the divui result + (Pid(0), 32, False), + (Const(3), 32, False), + (Param("n"), 32, False), # cmpi ult reads n unsigned + } + obs = width_obligations(g, store) + assert all(_holds(o, {"pid": 5, "n": 8}) for o in obs) + # n = -1 reads as 2**32 - 1 under ult: the unbounded model is not exact + assert not all(_holds(o, {"pid": 5, "n": -1}) for o in obs) + + +def test_narrowing_casts_and_extui_in_spike_casts(): + g = _graph("ttir/spike_casts.ttir") + load = g.accesses[0] + casts = [n for n in R._nodes((load.offset,), None) if isinstance(n, IntCast)] + assert {(c.kind, c.src_bits, c.dst_bits) for c in casts} == { + ("trunci", 32, 16), + ("extsi", 16, 64), + ("trunci", 32, 8), + ("extui", 8, 64), + } + obs = [ + (o.bits, o.signed) for o in width_obligations(g, load) if o.line_no is not None + ] + assert (16, True) in obs and (8, True) in obs and (8, False) in obs + + +def test_i1_casts_and_unsigned_loops(): + # extsi of an i1 maps true to -1: read as 0 - extui(b) + g = parse_ttir( + _module( + """ + %c0 = arith.constant 0 : i32 + %b = arith.cmpi slt, %n, %c0 : i32 + %e = arith.extsi %b : i1 to i32 + %q = tt.addptr %p, %e : !tt.ptr, i32 + %v = tt.load %q : !tt.ptr""" + ) + ) + b = Cmp("slt", Param("arg1"), Const(0), 32) + assert g.accesses[0].offset == Bin( + "+", Const(0), Bin("-", Const(0), IntCast("extui", 1, 32, b), 32) + ) + # trunci to i1 keeps only bit 0: exact for 0 / 1 + g = parse_ttir( + _module("%b = arith.trunci %n : i32 to i1\n%v = tt.load %p, %b : !tt.ptr") + ) + ((ob,),) = [width_obligations(g, a) for a in g.accesses] + assert (ob.term, ob.bits, ob.signed) == (Param("arg1"), 1, False) + # an unsigned loop compare needs non-negative bounds + g = parse_ttir( + _module( + """ + %c0 = arith.constant 0 : i32 + %c1 = arith.constant 1 : i32 + scf.for unsigned %i = %c0 to %n step %c1 : i32 { + %q = tt.addptr %p, %i : !tt.ptr, i32 + %v = tt.load %q : !tt.ptr + }""" + ) + ) + assert g.loop.unsigned and g.loop.bits == 32 + obs = [ + (o.term, o.bits, o.signed, o.line_no) + for o in width_obligations(g, g.accesses[0]) + ] + latch = Bin("+", Bin("-", Param("arg1"), Const(1), 32), Const(1), 32) + assert obs == [ + (Const(0), 32, False, 5), + (Param("arg1"), 32, False, 5), + (Const(1), 32, False, 5), + (latch, 32, False, 5), # the increment does not wrap, read unsigned + ] + + +def test_loop_increment_obligation(): + # for k in range(lo, n, 1 << 20) wraps its induction variable when + # n is near INT32_MAX (the GPU then runs iterations at negative k) + g = _graph("reader_ttir/iv_wrap.ttir") + (store,) = g.accesses + latch = Bin("+", Bin("-", Param("n"), Const(1), 32), Const(1 << 20), 32) + obs = width_obligations(g, store) + assert [(o.term, o.bits, o.signed, o.role) for o in obs] == [ + (latch, 32, True, "loop") + ] + assert "range(lo, n, 1 << 20)" in _source_line(obs[0].loc) + assert ( + "scf.for" in _text("reader_ttir/iv_wrap.ttir").splitlines()[obs[0].line_no - 1] + ) + assert _holds(obs[0], {"n": 1000}) + assert not _holds(obs[0], {"n": (1 << 31) - (1 << 19)}) + # an access outside the loop gets no loop obligations + g = _graph("ttir/golden_matmul_s3_sm80.ttir") + assert not any(o.role == "loop" for o in width_obligations(g, g.accesses[-1])) + assert any(o.role == "loop" for o in width_obligations(g, g.accesses[0])) + + +def test_signed_remainder_needs_its_quotient_to_fit(): + # remsi(INT_MIN, -1) is undefined, though its value 0 fits + body = """ + %cm = arith.constant -2147483648 : i32 + %cn = arith.constant -1 : i32 + %r = arith.remsi %cm, %cn : i32 + %q = tt.addptr %p, %r : !tt.ptr, i32 + %v = tt.load %q : !tt.ptr""" + g = parse_ttir(_module(body)) + (load,) = g.accesses + lo = Const(-(1 << 31)) + obs = width_obligations(g, load) + assert [(o.term, o.bits, o.signed, o.line_no) for o in obs] == [ + (Bin("%", lo, Const(-1), 32), 32, True, 5), + (Bin("//", lo, Const(-1), 32), 32, True, 5), + ] + assert _holds(obs[0], {}) and not _holds(obs[1], {}) + + +def test_obligation_roles(): + # a node the path, mask and offset share is listed once, under the most + # restrictive role (path, then mask, then offset) + body = """ + %c4 = arith.constant 4 : i32 + %pid = tt.get_program_id x : i32 + %o = arith.muli %pid, %c4 : i32 + %s = arith.subi %n, %c4 : i32 + %b = arith.cmpi slt, %o, %s : i32 + scf.if %b { + %t = arith.addi %o, %c4 : i32 + %m = arith.cmpi slt, %t, %n : i32 + %o2 = arith.addi %t, %c4 : i32 + %q = tt.addptr %p, %o2 : !tt.ptr, i32 + %v = tt.load %q, %m : !tt.ptr + }""" + g = parse_ttir(_module(body)) + (load,) = g.accesses + o = Bin("*", Pid(0), Const(4), 32) + s = Bin("-", Param("arg1"), Const(4), 32) + t = Bin("+", o, Const(4), 32) + o2 = Bin("+", t, Const(4), 32) + assert [(ob.term, ob.role) for ob in width_obligations(g, load)] == [ + (o, "path"), + (s, "path"), + (t, "mask"), + (o2, "offset"), + ] + + +def test_obligation_sites_do_not_affect_term_equality(): + a = Bin("+", Pid(0), Const(1), 32, line_no=3, loc=W.SourceLoc("a.py", 1, 1)) + b = Bin("+", Pid(0), Const(1), 32, line_no=9, loc=None) + assert a == b and hash(a) == hash(b) and repr(a) == repr(b) + assert Bin("+", Pid(0), Const(1), 64) != a # the width is part of the value diff --git a/tests/unit/ir/ttir_corpus.py b/tests/unit/ir/ttir_corpus.py new file mode 100644 index 00000000..0a6e1f47 --- /dev/null +++ b/tests/unit/ir/ttir_corpus.py @@ -0,0 +1,270 @@ +"""The TTIR the IR tests read, compiled from kernels at test time. + +``text("ttir/.ttir")`` is the TTIR Triton 3.8 prints for spec +```` of ``ttir_kernels.py``, and ``text("reader_ttir/.ttir")`` +the TTIR of kernel ```` of ``reader_kernels.py``: host-compiled up to +the TTIR stage only (no GPU, no Triton cache; the whole corpus takes well +under a second), each loc's path cut back to ``tests/`` or ``triton/``, +once per process. The ``crafted_*`` texts are hand-written TTIR for shapes +no kernel prints reliably (empty region bodies, generic form, quoted +symbols, loc forms, cf edge cases). + +Not a test module: pytest imports it (python_files = *.py) and finds nothing. +""" + +from __future__ import annotations + +import functools +import importlib.util +import re +import sys +from pathlib import Path +from typing import Any + +import pytest + +KERNELS = Path(__file__).resolve().parent + +CRAFTED = { + "ttir/crafted_attr_dicts.ttir": r"""module { + tt.func public @k(%p: !tt.ptr) attributes {noinline = false} { + %c = arith.constant {tt.divisibility = dense<16> : tensor<1xi32>} 16 : i32 + %d = arith.constant {axis = 5 : i32} -3 : i32 + %r = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32> + %s = "tt.reduce"(%r) <{axis = 0 : i32}> ({ + ^bb0(%a: i32, %b: i32): + %t = arith.addi %a, %b : i32 + tt.reduce.return %t : i32 + }) {tt.divisibility = dense<16> : tensor<1xi32>, axis_note = 3 : i32} : (tensor<64xi32>) -> i32 + %e = tt.expand_dims %r {axis = 0 : i32, tt.note = "axis = 1"} : tensor<64xi32> -> tensor<1x64xi32> + %q = tt.addptr %p, %s : !tt.ptr, i32 + %q2 = tt.addptr %q, %c : !tt.ptr, i32 + %q3 = tt.addptr %q2, %d : !tt.ptr, i32 + tt.store %q3, %s : !tt.ptr + tt.return + } +} +""", + "ttir/crafted_empty_bodies.ttir": r"""#loc = loc("k.py":1:0) +module { + tt.func public @k(%p: !tt.ptr loc("p"(#loc)), %c: i1 loc("c"(#loc)), %n: i32 loc("n"(#loc))) attributes {noinline = false} { + %c0 = arith.constant 0 : i32 loc(#loc1) + %c1 = arith.constant 1 : i32 loc(#loc1) + scf.if %c { + } else { + tt.store %p, %c0 : !tt.ptr loc(#loc2) + } loc(#loc1) + scf.if %c { + tt.store %p, %c1 : !tt.ptr loc(#loc2) + } else { + } loc(#loc1) + scf.for %i = %c0 to %n step %c1 : i32 { + } loc(#loc3) + scf.if %c { + } loc(#loc1) + tt.return loc(#loc4) + } loc(#loc) +} loc(#loc) +#loc1 = loc("k.py":2:4) +#loc2 = loc("k.py":3:8) +#loc3 = loc("k.py":4:4) +#loc4 = loc("k.py":5:4) +""", + "ttir/crafted_empty_else.ttir": r"""module { + tt.func public @k(%p: !tt.ptr, %c: i1) attributes {noinline = false} { + %c0 = arith.constant 0 : i32 + scf.if %c { + tt.store %p, %c0 : !tt.ptr + } else { + } + tt.return + } +} +""", + "ttir/crafted_empty_for.ttir": r"""module { + tt.func public @k(%p: !tt.ptr, %n: i32) attributes {noinline = false} { + %c0 = arith.constant 0 : i32 + %c1 = arith.constant 1 : i32 + scf.for %i = %c0 to %n step %c1 : i32 { + } + tt.store %p, %c0 : !tt.ptr + tt.return + } +} +""", + "ttir/crafted_generic_form.ttir": r"""module { + tt.func public @k(%p: !tt.ptr, %a: i32, %b: i32) attributes {noinline = false} { + %c = "arith.cmpi"(%a, %b) <{predicate = 2 : i64}> : (i32, i32) -> i1 + %x = "arith.select"(%c, %a, %b) : (i1, i32, i32) -> i32 + %o = "tt.atomic_rmw"(%p, %x) <{atomic_rmw_op = 5 : i32, scope = 1 : i32, sem = 4 : i32}> : (!tt.ptr, i32) -> i32 + %pid = "tt.get_program_id"() <{axis = 1 : i32}> : () -> i32 + %q = tt.addptr %p, %pid : !tt.ptr, i32 + tt.store %q, %o : !tt.ptr + tt.return + } +} +""", + "ttir/crafted_locs.ttir": r"""#loc = loc("k.py":1:0) +#loc1 = loc("k.py":2:5) +#loc2 = loc("k.py":3:6) +#loc9 = loc(fused<"meta">[#loc1, #loc2]) +#loc10 = loc(callsite(#loc1 at #loc9)) +module { + tt.func public @k(%p: !tt.ptr loc("p"(#loc)), %n: i32 loc(unknown)) attributes {noinline = false} { + %c1 = arith.constant 1 : i32 loc(fused[#loc1, "x.py":7:3]) + %a = arith.addi %n, %c1 : i32 loc(#loc9) + %b = arith.addi %a, %c1 : i32 loc(callsite("inner"("y.py":3:4) at callsite(#loc1 at #loc2))) + %q = tt.addptr %p, %b : !tt.ptr, i32 loc("q.py":9:9) + tt.store %q, %a : !tt.ptr loc(#loc10) + tt.return loc(#loc11) + } loc(#loc) +} loc(#loc) +#loc11 = loc("k.py":12:1) +""", + "ttir/crafted_same_dest.ttir": r"""module { + tt.func public @k(%p: !tt.ptr, %a: i32, %b: i32) attributes {noinline = false} { + %c = arith.cmpi slt, %a, %b : i32 + cf.cond_br %c, ^bb1(%a : i32), ^bb1(%b : i32) + ^bb1(%x: i32): + %y = arith.addi %x, %a : i32 + cf.cond_br %c, ^bb2(%y, %x : i32, i32), ^bb3 + ^bb2(%u: i32, %v: i32): + %q = tt.addptr %p, %u : !tt.ptr, i32 + tt.store %q, %v : !tt.ptr + cf.br ^bb3 + ^bb3: + tt.return + } +} +""", + "ttir/crafted_symbols_strings.ttir": r"""module { + tt.func private @"f{%x} \22q\22 (a)"(%a: i32, %b: i32) -> (i32, i32) attributes {noinline = true} { + %s = arith.addi %a, %b : i32 + tt.return %s, %a : i32, i32 + } + tt.func public @k(%p: !tt.ptr, %n: i32) attributes {noinline = false} { + %r:2 = tt.call @"f{%x} \22q\22 (a)"(%n, %n) : (i32, i32) -> (i32, i32) + %y = tt.elementwise_inline_asm "{ mov.u32 $0, %tid.x; } // loc(\22x\22) }" {constraints = "=r,r", packed_element = 1 : i32, pure = true} %r#0 : i32 -> i32 + %c = arith.cmpi sgt, %y, %r#1 : i32 + tt.assert %c, "bad { } loc( \22 %x" : i1 + tt.print " p={%d} " {hex = false, isSigned = array} : %y : i32 + %q = tt.addptr %p, %y : !tt.ptr, i32 + tt.store %q, %r#1 : !tt.ptr + tt.return + } +} +""", + "ttir/crafted_unicode_strings.ttir": r"""#loc = loc("/tmp/\E5\86\85\E6\A0\B8/k.py":1:0) +#loc1 = loc("/tmp/\E5\86\85\E6\A0\B8/k.py":2:4) +module { + tt.func public @"\E6\A0\B8"(%p: !tt.ptr loc("\CF\80_ptr"(#loc)), %c: i1 loc("\E6\95\B0"(#loc))) attributes {noinline = false} { + %v = tt.load %p : !tt.ptr loc("\E5\80\BC"(#loc1)) + tt.assert %c, "\E9\94\99\E8\AF\AF \22quoted\22 \\ back" : i1 loc(#loc1) + tt.print "\CF\80=" {hex = false, isSigned = array} : %v : f32 loc(#loc1) + tt.return loc(#loc1) + } loc(#loc) +} loc(#loc) +""", +} + +# A loc's path up to this repository's tests/ or Triton's package directory: +# where the checkout and Triton happen to be on the machine that compiles. +_REPO_LOC = r'loc\("[^"]*/tests/(?=unit/ir/)' +_TRITON_LOC = r'loc\("[^"]*/triton/' + + +def portable(text: str) -> str: + """``text`` with each loc naming a file of this repository from + ``tests/`` and one of Triton's own sources (tl.cdiv, ...) from + ``triton/``, so a text reads the same on every machine.""" + return re.sub(_TRITON_LOC, 'loc("triton/', re.sub(_REPO_LOC, 'loc("tests/', text)) + + +def _load(name: str) -> Any: + """Kernel module ``name`` of this directory: its @triton.jit functions + built as real JITFunctions even after tests/unit/test_multithreading.py + set TRITON_INTERPRET=1 at import time.""" + from triton import knobs + + spec = importlib.util.spec_from_file_location(name, KERNELS / f"{name}.py") + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module # Triton reads the jit fn's module + missing = object() + previous = knobs.runtime.__dict__.get("interpret", missing) + knobs.runtime.__dict__["interpret"] = False + try: + spec.loader.exec_module(module) + finally: + if previous is missing: + knobs.runtime.__dict__.pop("interpret", None) + else: + knobs.runtime.__dict__["interpret"] = previous + return module + + +@functools.cache +def _kernels() -> tuple[Any, Any]: + return _load("ttir_kernels"), _load("reader_kernels") + + +NAMES = sorted( + [f"ttir/{n}.ttir" for n in _kernels()[0].SPECS] + + [f"reader_ttir/{n}.ttir" for n in _kernels()[1].SPECS] + + list(CRAFTED) +) + + +def _real_compiles_available() -> bool: + # Triton imported under TRITON_INTERPRET=1 builds its own standard library + # as InterpretedFunctions, so nothing can compile for real in-process. + import triton.language.standard as tl_standard + from triton.runtime.jit import JITFunction + + return isinstance(tl_standard.cdiv, JITFunction) + + +def _ttir(fn, sig, consts, cc=80, options=None, attrs=None) -> str: + """triton.compile's front half for a ``cuda:`` target, stopped + after the TTIR stage.""" + from triton._C.libtriton import ir + from triton.backends.compiler import GPUTarget + from triton.compiler import ASTSource + from triton.compiler.compiler import make_backend + + src = ASTSource(fn=fn, signature=sig, constexprs=consts, attrs=attrs or {}) + target = GPUTarget("cuda", cc, 32) + backend = make_backend(target) + opts = backend.parse_options({"num_warps": 4, **(options or {})}) + ctx = ir.context() + ir.load_dialects(ctx) + backend.load_dialects(ctx) + module = src.make_ir( + target, + opts, + backend.get_codegen_implementation(opts), + backend.get_module_map(), + ctx, + ) + stages: dict[str, Any] = {} + backend.add_stages(stages, opts, src.language) + return str(stages["ttir"](module, {"hash": "", "target": target, **opts.__dict__})) + + +@functools.cache +def text(name: str) -> str: + """The TTIR ``name`` (``ttir/.ttir`` or ``reader_ttir/.ttir``, one + of NAMES).""" + if name in CRAFTED: + return CRAFTED[name] + if not _real_compiles_available(): + pytest.skip( + "Triton was imported under TRITON_INTERPRET=1: nothing compiles", + allow_module_level=True, # a test module may read a text at import + ) + walk, reader = _kernels() + kind, stem = name.removesuffix(".ttir").split("/") + if kind == "ttir": + return portable(_ttir(*walk.SPECS[stem])) + sig, consts = reader.SPECS[stem] + return portable(_ttir(getattr(reader, stem), sig, consts)) diff --git a/tests/unit/ir/ttir_kernels.py b/tests/unit/ir/ttir_kernels.py new file mode 100644 index 00000000..c4559e49 --- /dev/null +++ b/tests/unit/ir/ttir_kernels.py @@ -0,0 +1,652 @@ +"""Kernels whose TTIR the walk and reader tests read as ``ttir/.ttir``. + +ttir_corpus.py compiles each spec below to TTIR at test time. + +Not a test module: pytest imports it (python_files = *.py) and finds nothing. +""" + +from __future__ import annotations + +from typing import Any + +import triton +import triton.language as tl + + +@triton.jit +def dot_precisions(a_ptr, b_ptr, c_ptr, BLOCK: tl.constexpr): + # fp32 inputs: `ieee` is the printer-elided default, tf32x3 prints + offs = tl.arange(0, BLOCK) + idx = offs[:, None] * BLOCK + offs[None, :] + a = tl.load(a_ptr + idx) + b = tl.load(b_ptr + idx) + c = tl.dot(a, b, input_precision="ieee") + d = tl.dot(a, b, input_precision="tf32x3") + tl.store(c_ptr + idx, c + d) + + +@triton.jit +def eps_consts(x_ptr, s_ptr, out_ptr, BLOCK: tl.constexpr): + # uppercase-E float literals: a scalar constant and a dense splat + offs = tl.arange(0, BLOCK) + x = tl.load(x_ptr + offs) + s = tl.load(s_ptr) + 1e-6 + tl.store(out_ptr + offs, x * s + 1e-12) + + +@triton.jit +def unicode_msgs(x_ptr, BLOCK: tl.constexpr): + # a non-ASCII tt.assert message (debug=True); device_print prefixes must be ASCII + offs = tl.arange(0, BLOCK) + x = tl.load(x_ptr + offs) + tl.device_assert(x > 0, "错误: π must be > 0") + tl.device_print("x=", x) + tl.store(x_ptr + offs, x + 1) + + +@triton.jit +def deep_chain(out_ptr, s, N: tl.constexpr): + # a def chain deeper than Python's default recursion limit + pid = tl.program_id(0) + off = pid + for _ in tl.static_range(N): + off = off * s + pid + tl.store(out_ptr + off, 1.0) + + +@triton.jit +def dot_scaled_k(a_ptr, as_ptr, b_ptr, bs_ptr, c_ptr, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr): # fmt: skip + # tt.dot_scaled prints `%a scale %as, %b scale %bs, %c`: ODS (a, b, c, as, bs) + rm = tl.arange(0, M) + rn = tl.arange(0, N) + rk = tl.arange(0, K) + rs = tl.arange(0, K // 32) + a = tl.load(a_ptr + rm[:, None] * K + rk[None, :]) + b = tl.load(b_ptr + rk[:, None] * N + rn[None, :]) + a_scale = tl.load(as_ptr + rm[:, None] * (K // 32) + rs[None, :]) + b_scale = tl.load(bs_ptr + rn[:, None] * (K // 32) + rs[None, :]) + c = tl.dot_scaled(a, a_scale, "e4m3", b, b_scale, "e4m3") + tl.store(c_ptr + rm[:, None] * N + rn[None, :], c) + + +@triton.jit +def descs(a_ptr, M, N, BM: tl.constexpr, BN: tl.constexpr): + # device-side tensor descriptors, descriptor_load / store / reduce / gather + # / scatter + d = tl.make_tensor_descriptor(a_ptr, [M, N], [N, 1], [BM, BN]) + x = d.load([0, BN]) + d.store([BM, 0], x) + d.atomic_add([BM, BN], x) + d1 = tl.make_tensor_descriptor(a_ptr, [M, N], [N, 1], [1, BN]) + rows = tl.arange(0, BM) + g = d1.gather(rows, 0) + d1.scatter(g, rows, BN) + + +@triton.jit +def add_kernel(x_ptr, y_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(0) + offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offs < n_elements + x = tl.load(x_ptr + offs, mask=mask) + y = tl.load(y_ptr + offs, mask=mask) + tl.store(out_ptr + offs, x + y, mask=mask) + + +@triton.jit +def matmul_kernel( + a_ptr, + b_ptr, + c_ptr, + M, + N, + K, + stride_am, + stride_bk, + stride_cm, # inner strides are 1 (row-major) + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + offs_k = tl.arange(0, BLOCK_K) + + a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] + b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_n[None, :] + + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_K)): + a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_K, other=0.0) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_K, other=0.0) + acc += tl.dot(a, b) + a_ptrs += BLOCK_K + b_ptrs += BLOCK_K * stride_bk + + c = acc.to(tl.float16) + c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] + c_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) + tl.store(c_ptrs, c, mask=c_mask) + + +@triton.jit +def tile2d_kernel( + in_ptr, + out_ptr, + M, + N, + stride_m, + stride_n, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + # 2D tile copy: independent row/col arange instances, per-axis masks ANDed + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + ptrs = in_ptr + offs_m[:, None] * stride_m + offs_n[None, :] * stride_n + mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) + vals = tl.load(ptrs, mask=mask, other=0.0) + optrs = out_ptr + offs_m[:, None] * stride_m + offs_n[None, :] * stride_n + tl.store(optrs, vals * 2.0, mask=mask) + + +@triton.jit +def atomic_fmax_kernel(x_ptr, out_ptr, n_elements, BLOCK: tl.constexpr): + # float tl.atomic_max lowers to a sign-trick dance: the pointer is + # tt.bitcast to i32 and the two RMWs' masks derive from the loaded value + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n_elements + v = tl.load(x_ptr + offs, mask=mask, other=0.0) + tl.atomic_max(out_ptr + offs, v, mask=mask) + + +@triton.jit +def nested_guard_merge_kernel(x_ptr, out_ptr, n, T, BLOCK: tl.constexpr): + # early returns under nested guards: cf.cond_br blocks that merge + pid = tl.program_id(0) + base = pid * BLOCK + if pid >= T: + return + if pid == 0: + base = 0 + if n < 0: + return + else: + base = base + n + offs = base + tl.arange(0, BLOCK) + v = tl.load(x_ptr + offs) + tl.store(out_ptr + offs, v) + + +@triton.jit +def pid_branch_kernel(x_ptr, out_ptr, n_elements, BLOCK: tl.constexpr): + # a pid-dependent branch without results: only program 0 stores (an + # scf.if with a then-region only) + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n_elements + v = tl.load(x_ptr + offs, mask=mask) + if pid == 0: + tl.store(out_ptr + offs, v, mask=mask) + + +@triton.jit +def grid_stride_kernel( + x_ptr, out_ptr, n_rows, stride, NUM_PRGMS: tl.constexpr, BLOCK: tl.constexpr +): + # a grid-stride loop: its lower bound is the program id + pid = tl.program_id(0) + cols = tl.arange(0, BLOCK) + for row in range(pid, n_rows, NUM_PRGMS): + v = tl.load(x_ptr + row * stride + cols) + tl.store(out_ptr + row * stride + cols, v) + + +@triton.jit +def cas_kernel(lock_ptr, out_ptr): + # a scalar tt.atomic_cas (no mask operand) on a pointer argument + old = tl.atomic_cas(lock_ptr, 0, 1) + tl.store(out_ptr, old) + + +@triton.jit +def gather_kernel(idx_ptr, src_ptr, out_ptr, n_elements, BLOCK: tl.constexpr): + # an indirect load: a loaded value feeds the second load's address + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n_elements + idx = tl.load(idx_ptr + offs, mask=mask, other=0) + vals = tl.load(src_ptr + idx, mask=mask, other=0.0) + tl.store(out_ptr + offs, vals, mask=mask) + + +@triton.jit +def if_yield(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr): + # scf.if with yields (the loaded value and the pointer both flow out) + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + m = offs < n + if pid % 2 == 0: + v = tl.load(x_ptr + offs, mask=m) + dst = out_ptr + offs + else: + v = tl.load(y_ptr + offs, mask=m) * 2.0 + dst = out_ptr + n + offs + tl.store(dst, v, mask=m) + + +@triton.jit +def spin_while(lock_ptr, flag_ptr, out_ptr): + # scf.while spin loop (ticket lock style): atomic_cas poll + volatile flag poll + while tl.atomic_cas(lock_ptr, 0, 1, sem="acquire", scope="gpu") == 1: + pass + v = tl.load(flag_ptr, volatile=True) + while v == 0: + v = tl.load(flag_ptr, volatile=True) + tl.store(out_ptr, v) + tl.atomic_xchg(lock_ptr, 0, sem="release", scope="gpu") + + +@triton.jit +def atomics(p_ptr, q_ptr, n, BLOCK: tl.constexpr): + # atomic_rmw / atomic_cas with every sem / scope spelling + offs = tl.arange(0, BLOCK) + m = offs < n + tl.atomic_add(p_ptr + offs, 1.0, mask=m, sem="relaxed", scope="cta") + tl.atomic_max(q_ptr + offs, 7, mask=m, sem="release", scope="sys") + tl.atomic_min(q_ptr + offs, -3, mask=m, sem="acquire", scope="gpu") + tl.atomic_and(q_ptr + offs, 0xF0, mask=m) + tl.atomic_or(q_ptr + offs, 1, mask=m, sem="acq_rel") + tl.atomic_xor(q_ptr + offs, 2, mask=m) + old = tl.atomic_xchg(q_ptr, 5, sem="relaxed", scope="sys") + tl.atomic_cas(q_ptr + 1, old, 9, sem="acq_rel", scope="cta") + tl.atomic_add(q_ptr + 2, 3) # scalar, default sem/scope + + +@triton.jit +def _combine_max_abs(a, b): + return tl.maximum(tl.abs(a), tl.abs(b)) + + +@triton.jit +def reduce_scan(x_ptr, out_ptr, BLOCK: tl.constexpr): + # tt.reduce with a combine region (builtin sum, custom combine, argmax = 2 + # results), tt.scan (cumsum) + offs = tl.arange(0, BLOCK) + x = tl.load(x_ptr + offs) + tl.store(out_ptr, tl.sum(x, axis=0)) + tl.store(out_ptr + 1, tl.reduce(x, 0, _combine_max_abs)) + tl.store(out_ptr + 2, tl.argmax(x, axis=0).to(tl.float32)) + tl.store(out_ptr + 3 + offs, tl.cumsum(x, axis=0)) + + +@triton.jit +def inline_asm(x_ptr, out_ptr, BLOCK: tl.constexpr): + # tl.inline_asm_elementwise: pure, impure (a store), braces / % in the asm + offs = tl.arange(0, BLOCK) + x = tl.load(x_ptr + offs) + y = tl.inline_asm_elementwise( + "{ .reg .u32 t; mov.u32 t, %tid.x; shl.b32 $0, $1, 3; }", + "=r,r", + [x], + dtype=tl.int32, + is_pure=True, + pack=1, + ) + tl.inline_asm_elementwise( + "st.global.b32 [$1], $2; mov.u32 $0, 0;", + "=r,l,r", + [out_ptr + offs, y], + dtype=tl.int32, + is_pure=False, + pack=1, + ) + + +@triton.jit +def casts(x_ptr, out_ptr, n, BLOCK: tl.constexpr): + # trunci / extsi / extui casts on the address path and on values + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + o16 = offs.to(tl.int16) + o64 = o16.to(tl.int64) + u8 = offs.to(tl.uint8) + idx = o64 + u8.to(tl.int32) + v = tl.load(x_ptr + idx, mask=idx < n) + tl.store(out_ptr + offs.to(tl.int64), v, mask=offs < n) + + +@triton.jit +def consts(x_ptr, i8_ptr, i64_ptr, u32_ptr, BLOCK: tl.constexpr): + # constants: negative / large / min-int / unsigned / narrow ints, special + # floats, dense splats of negative ints and bools + offs = tl.arange(0, BLOCK) + v = tl.load(x_ptr + offs) + neg = tl.full([BLOCK], -7, tl.int32) + mn = tl.full([BLOCK], -2147483648, tl.int32) + tl.store(x_ptr + offs, v + neg + mn.to(tl.float32)) + tl.store(x_ptr + offs + BLOCK, tl.where(v != v, float("nan"), float("-inf"))) + tl.store(x_ptr + offs + 2 * BLOCK, tl.full([BLOCK], -0.0, tl.float32) + 1e-30) + tl.store( + i8_ptr + offs, tl.full([BLOCK], -128, tl.int8) + tl.full([BLOCK], 127, tl.int8) + ) + big = tl.full([BLOCK], 9223372036854775807, tl.int64) + small = tl.full([BLOCK], -9223372036854775808, tl.int64) + tl.store(i64_ptr + offs, big + small + 4294967296) + tl.store(i64_ptr - 9223372036854775807 + offs, big) + tl.store(u32_ptr + offs, tl.full([BLOCK], 4294967295, tl.uint32)) + b = tl.full([BLOCK], True, tl.int1) + tl.store(u32_ptr + offs + BLOCK, tl.full([BLOCK], 1, tl.uint32), mask=b) + tl.store(u32_ptr + offs + 2 * BLOCK, tl.full([BLOCK], 1, tl.uint32), mask=offs < -1) + h = tl.full([BLOCK], 1.5, tl.float16) * tl.full([BLOCK], -2.0, tl.bfloat16).to( + tl.float16 + ) + tl.store(x_ptr + offs + 3 * BLOCK, h.to(tl.float32)) + + +@triton.jit +def views(x_ptr, out_ptr, BLOCK: tl.constexpr): + # reshape with can_reorder (tt.reshape allow_reorder), trans, broadcast, gather + offs = tl.arange(0, BLOCK) + p = x_ptr + offs + p2 = tl.reshape(p, [BLOCK // 4, 4], can_reorder=True) + m2 = tl.reshape(offs < 50, [BLOCK // 4, 4], can_reorder=True) + v = tl.load(p2, mask=m2) + t = tl.trans(v) + q = ( + out_ptr + + tl.arange(0, 4)[:, None] * (BLOCK // 4) + + tl.arange(0, BLOCK // 4)[None, :] + ) + tl.store(q, t) + g = tl.gather(offs, (BLOCK - 1) - offs, 0) + tl.store(out_ptr + BLOCK + g, tl.load(x_ptr + offs)) + + +# (kernel, signature, constexprs, compute capability, compile options, +# ASTSource attrs) +_Spec = tuple[Any, dict[str, str], dict[str, int], int, dict[str, Any], dict] + + +def _divisible_by_16(count: int) -> dict: + """ASTSource attrs: divisibility 16 on the first ``count`` arguments, as + the JIT specializes well-aligned pointers and sizes.""" + return {(i,): [["tt.divisibility", 16]] for i in range(count)} + + +# corpus name (ttir/.ttir) -> spec +SPECS: dict[str, _Spec] = { + "kernel_dot_precisions": ( + dot_precisions, + {"a_ptr": "*fp32", "b_ptr": "*fp32", "c_ptr": "*fp32", "BLOCK": "constexpr"}, + {"BLOCK": 16}, + 80, + {}, + {}, + ), + "kernel_eps_consts": ( + eps_consts, + {"x_ptr": "*fp32", "s_ptr": "*fp32", "out_ptr": "*fp32", "BLOCK": "constexpr"}, + {"BLOCK": 64}, + 80, + {}, + {}, + ), + "kernel_unicode_msgs": ( + unicode_msgs, + {"x_ptr": "*fp32", "BLOCK": "constexpr"}, + {"BLOCK": 64}, + 80, + {"debug": True}, + {}, + ), + "kernel_deep_chain": ( + deep_chain, + {"out_ptr": "*fp32", "s": "i32", "N": "constexpr"}, + {"N": 600}, + 80, + {}, + {}, + ), + "kernel_dot_scaled": ( + dot_scaled_k, + { + "a_ptr": "*fp8e4nv", + "as_ptr": "*u8", + "b_ptr": "*fp8e4nv", + "bs_ptr": "*u8", + "c_ptr": "*fp32", + "M": "constexpr", + "N": "constexpr", + "K": "constexpr", + }, + {"M": 128, "N": 128, "K": 64}, + 100, + {}, + {}, + ), + "adv_descs": ( + descs, + { + "a_ptr": "*fp16", + "M": "i32", + "N": "i32", + "BM": "constexpr", + "BN": "constexpr", + }, + {"BM": 32, "BN": 32}, + 100, + {}, + {}, + ), + "golden_add_sm80": ( + add_kernel, + { + "x_ptr": "*fp32", + "y_ptr": "*fp32", + "out_ptr": "*fp32", + "n_elements": "i32", + "BLOCK_SIZE": "constexpr", + }, + {"BLOCK_SIZE": 1024}, + 80, + {"num_stages": 3}, + _divisible_by_16(4), + ), + "golden_matmul_s3_sm80": ( + matmul_kernel, + { + "a_ptr": "*fp16", + "b_ptr": "*fp16", + "c_ptr": "*fp16", + "M": "i32", + "N": "i32", + "K": "i32", + "stride_am": "i32", + "stride_bk": "i32", + "stride_cm": "i32", + "BLOCK_M": "constexpr", + "BLOCK_N": "constexpr", + "BLOCK_K": "constexpr", + }, + {"BLOCK_M": 64, "BLOCK_N": 64, "BLOCK_K": 32}, + 80, + {"num_stages": 3}, + _divisible_by_16(9), + ), + "golden_tile2d_sm80": ( + tile2d_kernel, + { + "in_ptr": "*fp32", + "out_ptr": "*fp32", + "M": "i32", + "N": "i32", + "stride_m": "i32", + "stride_n": "i32", + "BLOCK_M": "constexpr", + "BLOCK_N": "constexpr", + }, + {"BLOCK_M": 32, "BLOCK_N": 32}, + 80, + {"num_stages": 1}, + _divisible_by_16(6), + ), + "golden_atomic_fmax_sm80": ( + atomic_fmax_kernel, + { + "x_ptr": "*fp32", + "out_ptr": "*fp32", + "n_elements": "i32", + "BLOCK": "constexpr", + }, + {"BLOCK": 256}, + 80, + {"num_stages": 1}, + _divisible_by_16(3), + ), + "golden_nested_guard_merge_sm80": ( + nested_guard_merge_kernel, + { + "x_ptr": "*fp32", + "out_ptr": "*fp32", + "n": "i32", + "T": "i32", + "BLOCK": "constexpr", + }, + {"BLOCK": 64}, + 80, + {"num_stages": 1}, + {}, + ), + "golden_pid_branch_sm80": ( + pid_branch_kernel, + { + "x_ptr": "*fp32", + "out_ptr": "*fp32", + "n_elements": "i32", + "BLOCK": "constexpr", + }, + {"BLOCK": 256}, + 80, + {"num_stages": 1}, + _divisible_by_16(3), + ), + "golden_grid_stride_sm80": ( + grid_stride_kernel, + { + "x_ptr": "*fp32", + "out_ptr": "*fp32", + "n_rows": "i32", + "stride": "i32", + "NUM_PRGMS": "constexpr", + "BLOCK": "constexpr", + }, + {"NUM_PRGMS": 4, "BLOCK": 64}, + 80, + {"num_stages": 1}, + {}, + ), + "golden_cas_sm80": ( + cas_kernel, + {"lock_ptr": "*i32", "out_ptr": "*i32"}, + {}, + 80, + {"num_stages": 1}, + _divisible_by_16(2), + ), + "golden_gather_sm80": ( + gather_kernel, + { + "idx_ptr": "*i32", + "src_ptr": "*fp32", + "out_ptr": "*fp32", + "n_elements": "i32", + "BLOCK": "constexpr", + }, + {"BLOCK": 256}, + 80, + {"num_stages": 1}, + _divisible_by_16(4), + ), + "spike_if_yield": ( + if_yield, + { + "x_ptr": "*fp32", + "y_ptr": "*fp32", + "out_ptr": "*fp32", + "n": "i32", + "BLOCK": "constexpr", + }, + {"BLOCK": 64}, + 80, + {"num_stages": 1}, + {}, + ), + "spike_spin_while": ( + spin_while, + {"lock_ptr": "*i32", "flag_ptr": "*i32", "out_ptr": "*i32"}, + {}, + 80, + {"num_stages": 1}, + {}, + ), + "spike_atomics": ( + atomics, + {"p_ptr": "*fp32", "q_ptr": "*i32", "n": "i32", "BLOCK": "constexpr"}, + {"BLOCK": 64}, + 80, + {"num_stages": 1}, + {}, + ), + "spike_reduce_scan": ( + reduce_scan, + {"x_ptr": "*fp32", "out_ptr": "*fp32", "BLOCK": "constexpr"}, + {"BLOCK": 64}, + 80, + {"num_stages": 1}, + {}, + ), + "spike_inline_asm": ( + inline_asm, + {"x_ptr": "*i32", "out_ptr": "*i32", "BLOCK": "constexpr"}, + {"BLOCK": 64}, + 80, + {"num_stages": 1}, + {}, + ), + "spike_casts": ( + casts, + {"x_ptr": "*fp32", "out_ptr": "*fp32", "n": "i32", "BLOCK": "constexpr"}, + {"BLOCK": 64}, + 80, + {"num_stages": 1}, + {}, + ), + "adv_consts": ( + consts, + { + "x_ptr": "*fp32", + "i8_ptr": "*i8", + "i64_ptr": "*i64", + "u32_ptr": "*u32", + "BLOCK": "constexpr", + }, + {"BLOCK": 64}, + 80, + {}, + {}, + ), + "adv_views": ( + views, + {"x_ptr": "*fp32", "out_ptr": "*fp32", "BLOCK": "constexpr"}, + {"BLOCK": 64}, + 80, + {}, + {}, + ), +} diff --git a/tilelens/core/config.py b/tilelens/core/config.py index 3a0eaf9e..cb886891 100644 --- a/tilelens/core/config.py +++ b/tilelens/core/config.py @@ -110,8 +110,8 @@ def reset(self) -> None: # The one Triton minor release IR mode runs on: the host compile -# (tilelens.core.host_compile) uses its private API, so on any other release -# it refuses. +# (tilelens.core.host_compile) and the TTIR walk (tilelens.ir._mlir_walk) use +# its private API, so on any other release both refuse. IR_TRITON_RELEASE = "3.8" diff --git a/tilelens/ir/__init__.py b/tilelens/ir/__init__.py index 0bc6c7ea..6bcba1c6 100644 --- a/tilelens/ir/__init__.py +++ b/tilelens/ir/__init__.py @@ -1,4 +1,4 @@ -"""Compiled-IR layer: what IR-mode clients get of the kernels Triton compiles (TTIR). +"""Compiled-IR layer: read Triton-compiled kernels (TTIR) for the IR-mode clients. Exports resolve on first access, so importing ``tilelens.ir`` imports neither Triton nor its MLIR bindings. diff --git a/tilelens/ir/_mlir_walk.py b/tilelens/ir/_mlir_walk.py new file mode 100644 index 00000000..8b3434b9 --- /dev/null +++ b/tilelens/ir/_mlir_walk.py @@ -0,0 +1,1639 @@ +"""Walk a printed TTIR module: MLIR bindings first, the text for lines and locs. + +Private to ``tilelens.ir``; only ``ttir_reader.py`` imports it. This is the one +module that touches Triton's private MLIR bindings (``triton._C.libtriton.ir``), +and every lifetime hazard of those bindings stays here. + +``walk_module(text)`` reads the same bytes twice: + +1. The bindings parse the text (``ir.parse_mlir_module``, TTIR dialects only) + and give the module: op names, operand / result values, full type + strings, region / block nesting, block arguments, value locs, and the + attributes the reader reads (``_ATTRS``: ``get_int_attr`` for the integer + and enum attributes, ``get_constant_value`` for integer constants, the + string / bool / symbol getters for the rest). +2. A line scan of the text rebuilds the op tree for what the bindings cannot + reach: each op's line, the loc of an op without results (the bindings read + locs off values only), a region op's closing line, block labels, + ``scf.for``'s ``unsigned`` keyword (a unit attribute no getter reads) and + a generic-form ``tt.load``'s operand segments. + +The two trees are zipped in pre-order and must agree on every op: name, +result / region / block / op counts, and the printed loc of every op with +results equals the loc the bindings report for it (the anchor that ties +each text line to its op). A module that prints locs must +print one on every op; an op line holding a comment, and an attribute alias +other than ``#loc`` (TTGIR), are refused. The only tolerated differences +are the printer's own elisions: a trailing zero-operand ``scf.yield`` of an +``scf.for`` / ``scf.if`` block (also when that yield is the block's only op +and the region prints as ``{ }``) and unprinted empty trailing regions. +Anything else raises ``MisalignedModule``; a text the MLIR parser rejects +raises ``ModuleParseError`` with the parser's own diagnostic. + +The result is frozen pure-Python data: ``Module`` holds ``Op`` / ``Block`` / +``Value`` / ``Func`` records, which compare by value, hash, pickle and +deep-copy. Values are dense local ints (indices into ``Module.values``), +never the bindings' ``value.id()`` pointers. No binding object ever leaves +``_bind_walk``. + +Binding lifetime (dropping the context first crashes the process): the +context is pinned on the module (``mod.context = ctx``), everything is +extracted inside one function, the module's body block is erased, the module +is dropped before the context, and nothing derived from either is returned. +No binding erases the module op itself, so each parse still leaks that +(empty) op; results are cached by sha256 of the text, so each text leaks +once. The parse window redirects the process's fd 2 under a lock; a forked +child gets both back. A failed assertion inside the bindings aborts the +process too (Triton's LLVM keeps assertions): a block-pointer type never +reaches the parser (``_screen``), an integer constant wider than 64 bits is +never handed to ``get_constant_value`` (``_constant_value``), and +``get_int_attr`` is asked only for attributes whose signless integer type +the op's verifier enforces. + +The getters and the enum numbering are Triton 3.8's; on any other Triton +release ``walk_module`` refuses (tilelens.core.config.IR_TRITON_RELEASE). +""" + +from __future__ import annotations + +import collections +import hashlib +import os +import re +import sys +import tempfile +import threading +from dataclasses import dataclass +from typing import Any, Iterator, Mapping, Sequence + +# ─────────────────────────── records ─────────────────────────── + + +class _FrozenMap(Mapping[str, Any]): + """The read-only mapping behind the records' mapping fields. Unlike + ``MappingProxyType`` it hashes (its values are plain hashable data) and + pickles, so the frozen records hash, pickle and deep-copy too.""" + + __slots__ = ("_d",) + + def __init__(self, items: Mapping[str, Any]) -> None: + self._d = dict(items) + + def __getitem__(self, key: str) -> Any: + return self._d[key] + + def __iter__(self) -> Iterator[str]: + return iter(self._d) + + def __len__(self) -> int: + return len(self._d) + + def __hash__(self) -> int: + return hash(frozenset(self._d.items())) + + def __repr__(self) -> str: + return repr(self._d) + + def __reduce__(self) -> tuple[Any, ...]: + return (_FrozenMap, (self._d,)) + + +@dataclass(frozen=True) +class SourceLoc: + file: str + line: int + col: int + + +@dataclass(frozen=True) +class Value: + index: int # position in Module.values + type: str # full type text: "tensor<64x!tt.ptr>", "i32", ... + op: int | None # defining op index (an op result), else None + block: int | None # owning block index (a block argument), else None + position: int # result number, or argument number + # NameLoc name of the value's own loc (the Python variable / parameter + # name), None when unnamed. Printed SSA names are never exposed: the + # printer sanitises and uniquifies them. + name: str | None + + +@dataclass(frozen=True) +class Block: + index: int # position in Module.blocks + op: int # owning op index + region: int # region number within the owning op + position: int # block number within the region + label: str | None # printed label ("^bb1"); None for an unlabeled entry block + args: tuple[int, ...] # value indices + arg_types: tuple[str, ...] + arg_names: tuple[str | None, ...] # NameLoc names (see Value.name) + ops: tuple[int, ...] # op indices in program order + + +@dataclass(frozen=True) +class Op: + index: int # pre-order position in Module.ops (the text order) + name: str # "tt.load", "scf.for", "builtin.module", ... + operands: tuple[int, ...] # value indices, in ODS operand order + operand_types: tuple[str, ...] + results: tuple[int, ...] # value indices + result_types: tuple[str, ...] + # What the reader reads (read-only, see _ATTRS): integers are ints, enums + # their printed keyword ("slt"), program-id axes 0-2; + # arith.constant "value" is an int (sign-extended) / bool for i1, else + # ("float", None) / ("dense", None); scf.for "unsignedCmp" is a bool. + attrs: Mapping[str, Any] + regions: tuple[tuple[int, ...], ...] # block indices, per walked region + # block indices from the module body down to the parent block; () for + # the module op itself + path: tuple[int, ...] + position: int # position within the parent block + line_no: int | None # header line (1-based); None for an elided terminator + end_line: int | None # closing line of a region op, else == line_no + loc: SourceLoc | None # the op's own site (callee frame of a callsite loc) + callers: tuple[SourceLoc, ...] # callsite chain, innermost caller first + loc_name: str | None # NameLoc label of the op's loc, if any + implicit: bool = False # a terminator the printer elided + + @property + def block(self) -> int | None: + return self.path[-1] if self.path else None + + +@dataclass(frozen=True) +class FuncArg: + index: int + value: int # value index of the entry-block argument + type: str + name: str | None # NameLoc name: the Python parameter name + + +@dataclass(frozen=True) +class Func: + op: int # the tt.func op index + sym_name: str + visibility: str + args: tuple[FuncArg, ...] # () for a body-less declaration + + +@dataclass(frozen=True) +class Module: + ops: tuple[Op, ...] + blocks: tuple[Block, ...] + values: tuple[Value, ...] + funcs: tuple[Func, ...] # tt.func ops in text order + stats: Mapping[str, int] # what the alignment checked (tests) + + +class MisalignedModule(Exception): + """The text scan and the bindings disagree (or the text holds a construct + the text layer cannot read faithfully). ``problems`` lists every mismatch + found, ``line_no`` is the first text line involved (None if unknown).""" + + def __init__(self, problems: Sequence[str], line_no: int | None = None) -> None: + self.problems = tuple(problems) or ("misaligned module",) + self.line_no = line_no + super().__init__(self.problems[0]) + + +class ModuleParseError(Exception): + """The MLIR parser rejected the text (``diagnostic`` is its own message, + with the temporary parse path replaced by ````), the parse input + could not be created, the text holds a construct the parser cannot be + handed (a block-pointer type, see _screen), or the installed Triton is + not the release the walk reads.""" + + def __init__(self, diagnostic: str, line_no: int | None = None) -> None: + self.diagnostic = diagnostic + self.line_no = line_no + super().__init__(diagnostic) + + +# ─────────────────────────── types ─────────────────────────── + + +@dataclass(frozen=True) +class TypeInfo: + """A printed TTIR type split into shape and element, with its bit widths.""" + + text: str + shape: tuple[int, ...] # () for scalars + elem: str # element type text: "i32", "f16", "!tt.ptr", ... + int_bits: int | None # iN element -> N (signless); index -> 64 + float_bits: int | None + pointee: str | None # element pointer -> pointee type text + pointee_bits: int | None + + +_FLOAT_BITS = { + "f64": 64, "f32": 32, "f16": 16, "bf16": 16, "tf32": 32, + "f8E4M3FN": 8, "f8E5M2": 8, "f8E4M3FNUZ": 8, "f8E5M2FNUZ": 8, + "f8E4M3B11FNUZ": 8, "f8E8M0FNU": 8, "f4E2M1FN": 4, +} # fmt: skip +_RE_TENSOR = re.compile(r"^tensor<((?:\d+x)*)(.*)>$") +_RE_INT = re.compile(r"^i(\d+)$") +_RE_PTR = re.compile(r"^!tt\.ptr<(.*?)(?:, \d+)?>$") + + +def _scalar_bits(t: str) -> int | None: + m = _RE_INT.match(t) + if m: + return int(m.group(1)) + return _FLOAT_BITS.get(t) + + +def parse_type(text: str) -> TypeInfo: + shape: tuple[int, ...] = () + elem = text + m = _RE_TENSOR.match(text) + if m: + shape = tuple(int(d) for d in m.group(1).split("x") if d) + elem = m.group(2) + im = _RE_INT.match(elem) + pm = _RE_PTR.match(elem) + pointee = pm.group(1) if pm else None + int_bits = int(im.group(1)) if im else (64 if elem == "index" else None) + return TypeInfo( + text=text, + shape=shape, + elem=elem, + int_bits=int_bits, + float_bits=_FLOAT_BITS.get(elem), + pointee=pointee, + pointee_bits=_scalar_bits(pointee) if pointee else None, + ) + + +# ─────────────────────────── strings and brackets ─────────────────────────── + + +def _string_end(s: str, start: int) -> int: + """Index of the quote closing the string literal that opens at ``start``.""" + i = start + 1 + while i < len(s): + c = s[i] + if c == "\\": + i += 2 + continue + if c == '"': + return i + i += 1 + raise ValueError("unterminated string literal") + + +_HEX = frozenset("0123456789abcdefABCDEF") +_SIMPLE_ESCAPES = {"\\": 0x5C, '"': 0x22, "n": 0x0A, "t": 0x09} + + +def _unescape(body: str) -> str: + """Decode an MLIR string literal body. The printer escapes every + non-printable or non-ASCII byte as ``\\XX``; the escapes of one character + are the bytes of its UTF-8 encoding, so collect bytes and decode once.""" + out = bytearray() + i = 0 + while i < len(body): + c = body[i] + if c != "\\": + out += c.encode("utf-8") + i += 1 + continue + nxt = body[i + 1 : i + 2] + if nxt in _SIMPLE_ESCAPES: + out.append(_SIMPLE_ESCAPES[nxt]) + i += 2 + elif len(body) >= i + 3 and body[i + 1] in _HEX and body[i + 2] in _HEX: + out.append(int(body[i + 1 : i + 3], 16)) + i += 3 + else: + raise ValueError(f"bad string escape {body[i : i + 3]!r}") + try: + return out.decode("utf-8") + except UnicodeDecodeError as e: + raise ValueError(f"string literal is not UTF-8: {e}") from None + + +def _mask(line: str) -> tuple[str, list[tuple[int, int, str]]]: + """Replace every string literal's body by '_' (same length, so indices into + the masked line index the raw line too). Returns the masked line and the + literals as (open quote index, close quote index, decoded text).""" + out: list[str] = [] + lits: list[tuple[int, int, str]] = [] + i = 0 + while i < len(line): + c = line[i] + if c == '"': + j = _string_end(line, i) + lits.append((i, j, _unescape(line[i + 1 : j]))) + out.append('"' + "_" * (j - i - 1) + '"') + i = j + 1 + continue + out.append(c) + i += 1 + return "".join(out), lits + + +_OPEN = {"(": ")", "[": "]", "{": "}", "<": ">"} + + +def _match_close(s: str, i: int) -> int: + """Index of the bracket closing ``s[i]`` in masked text ('->' is no '>').""" + stack = [_OPEN[s[i]]] + j = i + 1 + while j < len(s): + c = s[j] + if c in _OPEN: + stack.append(_OPEN[c]) + elif c == ">" and s[j - 1] == "-": + pass + elif c in ")]}>": + if c != stack[-1]: + raise ValueError(f"bracket mismatch at column {j + 1}") + stack.pop() + if not stack: + return j + j += 1 + raise ValueError("unclosed bracket") + + +def _trailing_loc(masked: str) -> tuple[int, int] | None: + """(start, end) of the ``loc(...)`` that ends the masked text, if any.""" + s = masked.rstrip() + if not s.endswith(")"): + return None + depth = 0 + for j in range(len(s) - 1, -1, -1): + c = s[j] + if c == ")": + depth += 1 + elif c == "(": + depth -= 1 + if depth == 0: + if s[max(0, j - 3) : j] == "loc" and ( + j < 4 or not (s[j - 4].isalnum() or s[j - 4] in "_.$") + ): + return j - 3, len(s) + return None + return None + + +# ─────────────────────────── locs ─────────────────────────── +# One grammar for both sources: the text's `loc(#locN)` trailers resolved +# through the `#locN = loc(...)` table, and the bindings' fully inlined +# `str(value.get_loc())`. Both normalise to nested tuples compared with ==. + +_RE_LOC_ALIAS_DEF = re.compile(r"^(#loc\d*)\s*=\s*loc\((.*)\)\s*$") +_RE_LOC_ALIAS = re.compile(r"#loc\d*") +_RE_FILE_POS = re.compile(r":(\d+):(\d+)") +_LOC_DEPTH_LIMIT = 128 + + +class _LocParser: + def __init__(self, aliases: Mapping[str, str], parse_path: str) -> None: + self._aliases = aliases # "#loc12" -> inner text of loc(...) + self._parse_path = parse_path + self._alias_memo: dict[str, tuple] = {} + self._text_memo: dict[str, tuple | None] = {} + + def parse(self, loc_text: str) -> tuple | None: + """Normalised tree of a ``loc(...)`` text; None for a loc the parser + invented (it names the temporary parse path: the text printed none).""" + got = self._text_memo.get(loc_text, _MISSING) + if got is not _MISSING: + return got # type: ignore[return-value] + s = loc_text.strip() + if not (s.startswith("loc(") and s.endswith(")")): + raise ValueError(f"not a loc: {loc_text[:80]!r}") + tree: tuple | None = self._full(s[4:-1], 0) + if _names_path(tree, self._parse_path): # type: ignore[arg-type] + tree = None + self._text_memo[loc_text] = tree + return tree + + def _full(self, s: str, depth: int) -> tuple: + tree, rest = self._expr(s.strip(), depth) + if rest.strip(): + raise ValueError(f"trailing loc text: {rest[:60]!r}") + return tree + + def _expr(self, s: str, depth: int) -> tuple[tuple, str]: + if depth > _LOC_DEPTH_LIMIT: + raise ValueError("loc nesting too deep") + if s.startswith("#loc"): + m = _RE_LOC_ALIAS.match(s) + assert m is not None + name = m.group(0) + if name not in self._alias_memo: + if name not in self._aliases: + raise ValueError(f"undefined loc alias {name}") + self._alias_memo[name] = ("pending",) + self._alias_memo[name] = self._full(self._aliases[name], depth + 1) + elif self._alias_memo[name] == ("pending",): + raise ValueError(f"cyclic loc alias {name}") + return self._alias_memo[name], s[m.end() :] + if s.startswith("unknown"): + return ("unknown",), s[len("unknown") :] + if s.startswith("callsite("): + callee, rest = self._expr(s[len("callsite(") :].lstrip(), depth + 1) + rest = rest.lstrip() + if not rest.startswith("at "): + raise ValueError("callsite loc without 'at'") + caller, rest = self._expr(rest[3:].lstrip(), depth + 1) + rest = rest.lstrip() + if not rest.startswith(")"): + raise ValueError("unclosed callsite loc") + return ("callsite", callee, caller), rest[1:] + if s.startswith("fused"): + rest = s[len("fused") :] + meta: str | None = None + if rest.startswith("<"): + close = _match_close(_mask(rest)[0], 0) + meta = rest[1:close] + rest = rest[close + 1 :] + if not rest.startswith("["): + raise ValueError("bad fused loc") + parts = [] + rest = rest[1:].lstrip() + while not rest.startswith("]"): + p, rest = self._expr(rest, depth + 1) + parts.append(p) + rest = rest.lstrip() + if rest.startswith(","): + rest = rest[1:].lstrip() + elif not rest.startswith("]"): + raise ValueError("bad fused loc list") + return ("fused", meta, tuple(parts)), rest[1:] + if s.startswith('"'): + end = _string_end(s, 0) + text = _unescape(s[1:end]) + rest = s[end + 1 :] + m = _RE_FILE_POS.match(rest) + if m: + rest = rest[m.end() :] + if rest.lstrip().startswith("to"): + raise ValueError("file range locs are not supported") + return ("file", text, int(m.group(1)), int(m.group(2))), rest + if rest.startswith("("): + child, rest = self._expr(rest[1:], depth + 1) + rest = rest.lstrip() + if not rest.startswith(")"): + raise ValueError("unclosed name loc") + return ("name", text, child), rest[1:] + return ("name", text, None), rest + raise ValueError(f"unrecognized loc: {s[:60]!r}") + + +_MISSING = object() + + +def _names_path(tree: tuple, path: str) -> bool: + """Does any file loc inside ``tree`` name ``path``? (iterative)""" + stack = [tree] + while stack: + t = stack.pop() + if t is None: + continue + kind = t[0] + if kind == "file": + if t[1] == path: + return True + elif kind == "name": + stack.append(t[2]) + elif kind == "callsite": + stack += [t[1], t[2]] + elif kind == "fused": + stack += list(t[2]) + return False + + +def _loc_site( + tree: tuple | None, +) -> tuple[SourceLoc | None, tuple[SourceLoc, ...], str | None]: + """(site, callers, name) of a normalised loc. The callee frame of a + callsite is the op's site (a memory op belongs to the callee); the + caller chain follows, innermost first. Recursion depth is bounded by the + parser's ``_LOC_DEPTH_LIMIT``.""" + if tree is None: + return None, (), None + kind = tree[0] + if kind == "file": + return SourceLoc(tree[1], tree[2], tree[3]), (), None + if kind == "name": + site, callers, _ = _loc_site(tree[2]) + return site, callers, tree[1] + if kind == "callsite": + site, callers, name = _loc_site(tree[1]) + csite, ccallers, _ = _loc_site(tree[2]) + return site, callers + ((csite,) if csite else ()) + ccallers, name + if kind == "fused": + for part in tree[2]: + got = _loc_site(part) + if got[0] is not None: + return got + return None, (), None + + +# ─────────────────────────── text scan ─────────────────────────── + +_RE_RESULTS = re.compile( + r"^((?:%[-\w.$]+(?::\d+)?)(?:\s*,\s*%[-\w.$]+(?::\d+)?)*)\s*=\s*" +) +_RE_OPNAME = re.compile(r"^[A-Za-z_][\w$.]*") +_RE_LABEL = re.compile(r"^\^[-\w.$]+") + + +class _TextError(Exception): + def __init__(self, line_no: int | None, msg: str) -> None: + super().__init__(msg) + self.line_no = line_no + self.msg = msg + + +@dataclass +class _TBlock: + label: str | None + ops: list["_TOp"] + line_no: int + + +@dataclass +class _TOp: + name: str + line_no: int + raw: str # stripped raw header line (comment removed) + masked: str # masked header line (same length) + name_end: int # index just past the op name + n_results: int + generic: bool # printed in the generic form: "tt.reduce"(...) + opens_region: bool + regions: list[list[_TBlock]] + close_line: int | None = None + close_raw: str = "" + close_masked: str = "" + + +@dataclass +class _TextTree: + root: _TOp + aliases: dict[str, str] + + +def _n_results(prefix: str) -> int: + """``%a, %b:2`` -> 3.""" + n = 0 + for tok in prefix.split(","): + _, _, k = tok.strip().partition(":") + n += int(k) if k else 1 + return n + + +def _parse_op_line( + ln: int, raw: str, masked: str, lits: list[tuple[int, int, str]] +) -> _TOp: + rm = _RE_RESULTS.match(masked) + start = rm.end() if rm else 0 + body = masked[start:] + generic = body.startswith('"') + if generic: + name_end = _string_end(raw, start) + 1 + name = next((t for a, _b, t in lits if a == start), "") + else: + nm = _RE_OPNAME.match(body) + if nm is None: + raise _TextError(ln, f"cannot read an op name: {raw[:80]!r}") + name = nm.group(0) + name_end = start + nm.end() + if "." not in name: + name = f"builtin.{name}" + n_results = _n_results(rm.group(1)) if rm else 0 + opens = masked.endswith("{") + return _TOp(name, ln, raw, masked, name_end, n_results, generic, opens, []) + + +def _scan_text(text: str) -> _TextTree: + """Rebuild the op tree from printed lines. Returns a synthetic file-level + op whose single region holds the top-level ops, and the ``#loc`` table.""" + root = _TOp("", 0, "", "", 0, 0, False, True, [[]]) + stack = [root] + aliases: dict[str, str] = {} + for ln, raw_line in enumerate(text.splitlines(), start=1): + raw = raw_line.strip() + if not raw or raw.startswith("//"): + continue + try: + masked, lits = _mask(raw) + except ValueError as e: + raise _TextError(ln, str(e)) from None + cut = masked.find("//") + if cut >= 0: # a block label's predecessor comment + raw, masked = raw[:cut].rstrip(), masked[:cut].rstrip() + lits = [x for x in lits if x[1] < cut] + top = stack[-1] + if masked.startswith("#"): + m = _RE_LOC_ALIAS_DEF.match(raw) + if len(stack) != 1 or m is None: + raise _TextError( + ln, + f"attribute alias {raw[:60]!r}: only #loc aliases are TTIR (TTGIR input?)", + ) + aliases[m.group(1)] = m.group(2) + continue + if masked.startswith("^"): # its arguments come from the bindings + lm = _RE_LABEL.match(masked) + if top is root or lm is None: + raise _TextError(ln, f"bad block label {raw[:60]!r}") + after = lm.end() + try: + if masked.startswith("(", after): + after = _match_close(masked, after) + 1 + except ValueError as e: + raise _TextError(ln, str(e)) from None + if masked[after:].strip() != ":": + raise _TextError(ln, f"bad block label {raw[:60]!r}") + top.regions[-1].append(_TBlock(lm.group(0), [], ln)) + continue + if cut >= 0: + comment = raw_line.strip()[_mask(raw_line.strip())[0].find("//") + 2 :] + raise _TextError(ln, f"unexpected comment {comment.strip()[:60]!r}") + if masked.startswith("}"): + if top is root: + raise _TextError(ln, "unbalanced '}'") + rest = masked[1:].lstrip() + if rest.endswith("{"): # "} else {", "} do {", "}, {" + top.regions.append([]) + continue + if rest.startswith(")"): # generic op closer "}) ..." + rest = rest[1:].lstrip() + off = len(masked) - len(rest) + top.close_line, top.close_raw, top.close_masked = ln, raw[off:], rest + stack.pop() + continue + try: + op = _parse_op_line(ln, raw, masked, lits) + except ValueError as e: + raise _TextError(ln, str(e)) from None + region = top.regions[-1] + if not region: + region.append(_TBlock(None, [], ln)) + region[-1].ops.append(op) + if op.opens_region: + op.regions = [[]] + stack.append(op) + if len(stack) != 1: + raise _TextError( + stack[-1].line_no, + f"region opened at line {stack[-1].line_no} is never closed", + ) + return _TextTree(root, aliases) + + +# what the text alone says: scf.for's `unsigned` keyword (unsignedCmp, a +# unit attribute: it changes the trip count), and the operand segments of a +# generic-form tt.load (its mask and other are both optional) +_RE_FOR_HEAD = re.compile(r"\s+(unsigned\s+)?%") +_RE_SEGMENTS = re.compile(r"\boperandSegmentSizes\s*=\s*array") + + +def _text_attrs(t: _TOp, name: str) -> dict[str, Any]: + if name == "scf.for": + m = None if t.generic else _RE_FOR_HEAD.match(t.masked, t.name_end) + if m is None: + raise ValueError("unrecognized scf.for header") + return {"unsignedCmp": m.group(1) is not None} + if name == "tt.load" and t.generic: + m = _RE_SEGMENTS.search(t.masked, t.name_end) + if m is None: + raise ValueError("generic tt.load without operandSegmentSizes") + return {"operandSegmentSizes": tuple(int(x) for x in m.group(1).split(","))} + return {} + + +# A block-pointer type anywhere in a text: `ptr>`, `!tt>>`, whitespace, +# newlines or `//` comments in between), and any type alias definition, which +# could hide the pointee. The _GAP forms read the raw text (a comment runs to +# the end of its line), the others the code view of _screen_view. +_GAP = r"(?:\s|//[^\n]*(?![^\n]))*" +_RE_BLOCK_PTR_TYPE = re.compile(r"\bptr\s*<\s*tensor\b") +_RE_BLOCK_PTR_TYPE_GAP = re.compile(rf"\bptr{_GAP}<{_GAP}tensor\b") +_RE_TYPE_ALIAS_DEF = re.compile(r"^\s*(![-\w.$]+)\s*=", re.M) +_RE_TYPE_ALIAS_DEF_GAP = re.compile(rf"^\s*![-\w.$]+{_GAP}=", re.M) + + +def _code_line(line: str) -> str: + """``line`` as the parser's tokens see it, same length: string literal + bodies masked ('_'), a ``//`` comment outside strings blanked to the end + of the line. After an unterminated string the rest stays raw (the + parser stops there; reading it can only add a refusal).""" + out: list[str] = [] + i = 0 + while i < len(line): + c = line[i] + if c == '"': + try: + j = _string_end(line, i) + except ValueError: + out.append(line[i:]) + break + out.append('"' + "_" * (j - i - 1) + '"') + i = j + 1 + continue + if line.startswith("//", i): + out.append(" " * (len(line) - i)) + break + out.append(c) + i += 1 + return "".join(out) + + +def _screen_view(text: str) -> str: + """The code view of ``text`` (``_code_line`` per line), offsets kept, so + a match's line is its count of newlines before it plus one.""" + return "\n".join(_code_line(line) for line in text.split("\n")) + + +def _screen(text: str) -> None: + """Refuse, before the bindings see it, a text the parser cannot be + handed: without block-pointer types the parser aborts the whole process + on one (an assertion in PointerType::get). The whole text is searched, + so a type split over lines or around a comment is found too; string + literals are not read as types, comments are skipped.""" + if not (_RE_BLOCK_PTR_TYPE_GAP.search(text) or _RE_TYPE_ALIAS_DEF_GAP.search(text)): + return + view = _screen_view(text) + found = [] + m = _RE_BLOCK_PTR_TYPE.search(view) + if m is not None: + found.append((m.start(), "a block-pointer type (!tt.ptr>)")) + m = _RE_TYPE_ALIAS_DEF.search(view) + if m is not None: + found.append( + (m.start(1), "a type alias definition (it may name a block-pointer type)") + ) + if not found: + return + at, what = min(found) + ln = view.count("\n", 0, at) + 1 + raise ModuleParseError( + f"line {ln}: {what}: Triton's TTIR has no " + "block pointers, and its parser aborts the process on one; the " + "text is refused before parsing", + ln, + ) + + +# ─────────────────────────── attributes (bindings) ─────────────────────────── +# What the reader reads off an op, per op name: (key in Op.attrs, getter, +# attribute name, names of an enum's integers or None, required). An enum +# reads as its integer and is named by the keyword the printer prints for it +# (arith::CmpIPredicate); a program-id axis stays 0-2. Every int attribute +# here is an inherent one the op's verifier types as a signless integer, so +# get_int_attr never meets the type its getInt() asserts on. + +_CMPI = ("eq", "ne", "slt", "sle", "sgt", "sge", "ult", "ule", "ugt", "uge") +_AXIS = (0, 1, 2) +_ATTRS: Mapping[str, tuple[tuple[str, str, str, tuple | None, bool], ...]] = { + "tt.get_program_id": (("axis", "int", "axis", _AXIS, True),), + "tt.get_num_programs": (("axis", "int", "axis", _AXIS, True),), + "tt.make_range": ( + ("start", "int", "start", None, True), + ("end", "int", "end", None, True), + ), + "tt.expand_dims": (("axis", "int", "axis", None, True),), + "arith.cmpi": (("predicate", "int", "predicate", _CMPI, True),), + "tt.elementwise_inline_asm": (("pure", "bool", "pure", None, False),), + "tt.extern_elementwise": ( + ("pure", "bool", "pure", None, False), + ("symbol", "str", "symbol", None, False), + ), + "tt.call": (("callee", "sym", "callee", None, False),), + "tt.func": ( + ("sym_name", "str", "sym_name", None, True), + ("visibility", "str", "sym_visibility", None, False), + ), +} +_GETTERS: Mapping[str, tuple[str, type]] = { + "int": ("get_int_attr", int), + "bool": ("get_bool_attr", bool), + "str": ("get_str_attr", str), + "sym": ("get_flat_symbol_ref_attr", str), +} + + +def _is_exactly(value: Any, want: type) -> bool: + """``isinstance`` without bool passing as int.""" + return isinstance(value, bool) == (want is bool) and isinstance(value, want) + + +def _read_attrs(op, name: str, result_types: Sequence[str]) -> tuple[dict, list]: + """(attrs, problems) of one binding op; called inside the parse.""" + attrs: dict[str, Any] = {} + problems: list[str] = [] + for key, kind, aname, names, required in _ATTRS.get(name, ()): + getter, want = _GETTERS[kind] + v = getattr(op, getter)(aname) + if v is None: + if required: + problems.append(f"attribute {aname!r} not recovered") + elif not _is_exactly(v, want): + problems.append(f"{aname}: {getter} returned {type(v).__name__}") + elif names is not None and not (0 <= v < len(names) and names[v] is not None): + problems.append(f"{aname} {v} outside the closed vocabulary") + else: + attrs[key] = v if names is None else names[v] + if name == "arith.constant": + try: + attrs["value"] = _constant_value(op, result_types) + except ValueError as e: + problems.append(str(e)) + return attrs, problems + + +def _constant_value(op, result_types: Sequence[str]) -> Any: + """arith.constant's value: get_constant_value reads an integer scalar + or splat as its sign-extended int (the printer's signed reading; i1 true + reads -1, kept as a bool) and anything else as None: a float or a + non-splat dense is ("float", None) / ("dense", None), as is an integer + wider than 64 bits, which is never handed to the getter (getSExtValue + asserts, and Triton's LLVM keeps assertions: the process would abort).""" + if len(result_types) != 1: + raise ValueError("arith.constant without one result") + ty = parse_type(result_types[0]) + elem = parse_type(ty.elem) + bits = elem.int_bits + if bits is None: + if elem.float_bits is None: + raise ValueError(f"constant of unsupported element type {elem.elem}") + return ("float", None) + if bits > 64: + return ("dense", None) + v = op.get_constant_value() + if v is None: + if not ty.shape: + raise ValueError("integer constant without a value") + return ("dense", None) + if not _is_exactly(v, int): + raise ValueError(f"get_constant_value returned {type(v).__name__}") + if bits == 1: + if v not in (0, -1): + raise ValueError(f"i1 constant {v}") + return v == -1 + return v + + +# ─────────────────────────── bindings side ─────────────────────────── + + +@dataclass +class _BOp: + name: str + operands: tuple[int, ...] # raw value ids (valid within one parse only) + operand_types: tuple[str, ...] + results: tuple[int, ...] + result_types: tuple[str, ...] + result_locs: tuple[str, ...] + region_ids: tuple[int, ...] + region_sizes: tuple[int, ...] + block_id: int | None + attrs: dict[str, Any] + problems: list[str] # attributes _read_attrs could not read + + +@dataclass +class _BBlock: + region_id: int + args: tuple[int, ...] + arg_types: tuple[str, ...] + arg_locs: tuple[str, ...] + ops: list[int] # indices into the walk list + + +@dataclass +class _BindWalk: + ops: list[_BOp] + blocks: dict[int, _BBlock] + region_blocks: dict[int, list[int]] # raw region id -> raw block ids in order + path: str # the parse path (parser-invented locs name it) + + +def _write_all(fd: int, data: bytes) -> None: + view = memoryview(data) + while view: + n = os.write(fd, view) + view = view[n:] + + +class _ParseInput: + """The text as a file path for ``parse_mlir_module``: an anonymous memfd + via ``/proc/self/fd`` when available, else a temporary file.""" + + def __init__(self, data: bytes) -> None: + self.path: str | None = None + self._fd: int | None = None + self._tmp: str | None = None + memfd_create = getattr(os, "memfd_create", None) + if memfd_create is not None: + try: + fd = memfd_create("tilelens-ttir", getattr(os, "MFD_CLOEXEC", 0)) + except OSError: + fd = None + if fd is not None: + path = f"/proc/self/fd/{fd}" + try: + _write_all(fd, data) + ok = os.path.exists(path) + except OSError: + ok = False + if ok: + self._fd, self.path = fd, path + return + os.close(fd) + try: + fd, tmp = tempfile.mkstemp(prefix="tilelens-", suffix=".ttir") + except OSError as e: + raise ModuleParseError( + f"cannot create a temporary file for the MLIR parser: {e}" + ) from None + try: + try: + _write_all(fd, data) + finally: + os.close(fd) + except OSError as e: + os.unlink(tmp) + raise ModuleParseError(f"cannot write the MLIR parser input: {e}") from None + self._tmp = self.path = tmp + + def close(self) -> None: + if self._fd is not None: + os.close(self._fd) + self._fd = None + if self._tmp is not None: + try: + os.unlink(self._tmp) + except OSError: + pass + self._tmp = None + + +# (saved fd 2, capture buffer) while a capture redirects fd 2; set and +# cleared under _PARSE_LOCK, read by the fork handler (_after_fork_in_child) +_REDIRECT: tuple[int, int] | None = None + + +class _StderrCapture: + """Redirect fd 2 (where the C++ parser prints its diagnostic) into an + anonymous file for the duration of a ``with`` block; only used under + ``_PARSE_LOCK``. Everything written to fd 2 in that window lands in + ``data``: on a failed parse that is the diagnostic, on a successful one + it belongs to someone else (parser warnings, other threads), so + ``replay()`` passes it on to the real fd 2.""" + + def __init__(self) -> None: + self._buf: int | None = None + self._saved: int | None = None + self.data = b"" + self.text = "" + + def __enter__(self) -> "_StderrCapture": + global _REDIRECT + try: + memfd_create = getattr(os, "memfd_create", None) + if memfd_create is not None: + buf = memfd_create("tilelens-diag", getattr(os, "MFD_CLOEXEC", 0)) + else: + with tempfile.TemporaryFile() as f: + buf = os.dup(f.fileno()) + except OSError: + return self # no capture: the diagnostic stays on stderr + try: + sys.stderr.flush() + except (AttributeError, OSError, ValueError): + pass + try: + saved = os.dup(2) + except OSError: + os.close(buf) + return self + # published before fd 2 moves (and cleared after it is back), so a + # fork at any point of the window can restore fd 2 in the child + _REDIRECT = (saved, buf) + try: + os.dup2(buf, 2) + except OSError: + _REDIRECT = None + os.close(saved) + os.close(buf) + return self + self._buf, self._saved = buf, saved + return self + + def __exit__(self, *exc: object) -> None: + global _REDIRECT + if self._saved is not None: + os.dup2(self._saved, 2) + _REDIRECT = None + os.close(self._saved) + self._saved = None + if self._buf is not None: + try: + os.lseek(self._buf, 0, os.SEEK_SET) + chunks = [] + while chunk := os.read(self._buf, 1 << 16): + chunks.append(chunk) + self.data = b"".join(chunks) + self.text = self.data.decode("utf-8", "replace") + finally: + os.close(self._buf) + self._buf = None + + def replay(self) -> None: + if self.data: + try: + _write_all(2, self.data) + except OSError: + pass + + +_PARSE_LOCK = threading.Lock() # fd-2 redirection is process-wide + + +def _bind_walk(data: bytes) -> _BindWalk: + """Parse ``data`` with the bindings and flatten the module into + pure-Python records. The context is pinned on the module for the + module's whole life, the module is dropped before the context, and no + binding object survives the call (dropping the context first + segfaults). After a successful parse, whatever else reached fd 2 in the + window is passed on.""" + from triton._C.libtriton import ir # TTIR dialects only: no backend (TTGIR) loading + + source = _ParseInput(data) + try: + path = source.path + assert path is not None + ctx = ir.context() + try: + ir.load_dialects(ctx) + capture = _StderrCapture() + mod = None + with _PARSE_LOCK, capture: + try: + mod = ir.parse_mlir_module(path, ctx) + except RuntimeError: + pass + if mod is None: + raise _parse_error(capture.text, path) + capture.replay() + try: + walk, failure = _pinned_extract(mod, ctx) + finally: + del mod # the module before its context + finally: + del ctx + finally: + source.close() + if walk is None: + raise MisalignedModule([f"bindings walk failed: {failure}"]) + ops, blocks, region_blocks = walk + return _BindWalk(ops, blocks, region_blocks, path) + + +def _pinned_extract(mod, ctx) -> tuple[Any, str | None]: + """Pin ``ctx`` on ``mod`` (proton's pattern: the module keeps its context + alive), then extract. Returns (records, None) or (None, reason); an + exception's traceback, whose frames hold binding objects, dies here while + the module and its context are still alive.""" + try: + mod.context = ctx + except Exception as e: # noqa: BLE001 + return ( + None, + f"cannot pin the MLIR context on the module: {type(e).__name__}: {e}", + ) + body: list[Any] = [] + try: + walk = _extract(mod, body) + except Exception as e: # noqa: BLE001 (bindings drift: an unexpected getter result) + return None, f"{type(e).__name__}: {e}" + # No binding erases a module, so each parse would leak its whole op tree + # (~20 KiB for a 12 KiB text); erasing the body block frees all but the + # empty module op. Safe here: extraction is over, ctx is pinned, and the + # body block is the one binding object still alive (dropped at once). + if body: + try: + body.pop().erase() + except Exception: # noqa: BLE001 (bindings drift: keep the bounded leak) + pass + return walk, None + + +_RE_DIAG_POS = re.compile(r'loc\("":(\d+):\d+\)') + + +def _parse_error(diag: str, path: str) -> ModuleParseError: + diag = diag.replace(path, "").strip() or "the MLIR parser rejected the text" + m = _RE_DIAG_POS.search(diag) + return ModuleParseError(diag, int(m.group(1)) if m else None) + + +def _extract( + mod, body: list[Any] +) -> tuple[list[_BOp], dict[int, _BBlock], dict[int, list[int]]]: + """Copy the walked module into pure-Python records (post-order walk: + ops of a block arrive in program order, blocks of a region in order). + The module's body block (the one binding object kept) goes to ``body``, + for ``_pinned_extract`` to erase.""" + ops: list[_BOp] = [] + blocks: dict[int, _BBlock] = {} + region_blocks: dict[int, list[int]] = {} + + def cb(op) -> None: + blk = op.get_block() + bid = None + if blk is not None: + bid = blk.id() + rec = blocks.get(bid) + if rec is None: + args = [blk.get_argument(i) for i in range(blk.get_num_arguments())] + parent = blk.get_parent() + rid = parent.id() + if parent.get_parent_region() is None: + body.append(blk) # the module's body block + rec = blocks[bid] = _BBlock( + rid, + tuple(x.id() for x in args), + tuple(str(x.get_type()) for x in args), + tuple(str(x.get_loc()) for x in args), + [], + ) + region_blocks.setdefault(rid, []).append(bid) + rec.ops.append(len(ops)) + name = op.get_name() + opnds = [op.get_operand(i) for i in range(op.get_num_operands())] + res = [op.get_result(i) for i in range(op.get_num_results())] + regs = [op.get_region(i) for i in range(op.get_num_regions())] + res_types = tuple(str(x.get_type()) for x in res) + attrs, problems = _read_attrs(op, name, res_types) + ops.append( + _BOp( + name, + tuple(x.id() for x in opnds), + tuple(str(x.get_type()) for x in opnds), + tuple(x.id() for x in res), + res_types, + tuple(str(x.get_loc()) for x in res), + tuple(x.id() for x in regs), + tuple(x.size() for x in regs), + bid, + attrs, + problems, + ) + ) + + mod.walk(cb) + return ops, blocks, region_blocks + + +# ─────────────────────────── alignment ─────────────────────────── + + +@dataclass +class _OpB: # mutable op builder + name: str + operands_raw: tuple[int, ...] + operand_types: tuple[str, ...] + results: tuple[int, ...] + result_types: tuple[str, ...] + attrs: dict[str, Any] + path: tuple[int, ...] + position: int + line_no: int | None + end_line: int | None + loc: SourceLoc | None + callers: tuple[SourceLoc, ...] + loc_name: str | None + implicit: bool + regions: list[tuple[int, ...]] + + +@dataclass +class _BlockB: + op: int + region: int + position: int + label: str | None + args: tuple[int, ...] + arg_types: tuple[str, ...] + arg_names: tuple[str | None, ...] + ops: list[int] + path: tuple[int, ...] # path of the ops inside this block + + +# what Module.stats counts +_STATS = ("ops", "implicit_ops", "blocks", "values", "funcs", "result_locs", "attrs") +_ELIDES_YIELD = frozenset({"scf.for", "scf.if"}) + + +class _Aligner: + def __init__(self, tree: _TextTree, bw: _BindWalk) -> None: + self.tree = tree + self.bw = bw + self.locp = _LocParser(tree.aliases, bw.path) + self.problems: list[tuple[int | None, str]] = [] + self.ops: list[_OpB] = [] + self.blocks: list[_BlockB] = [] + self.values: list[list[Any]] = [] # [type, op, block, position, name] + self.vmap: dict[int, int] = {} # raw value id -> value index + # op lines with / without a printed trailing loc: the printer (debug + # info on) gives every op one, so a module mixing both is misread + self.loc_lines: list[int] = [] + self.no_loc_lines: list[int] = [] + self.stats: collections.Counter[str] = collections.Counter( + dict.fromkeys(_STATS, 0) + ) + + def bad(self, line: int | None, msg: str) -> None: + self.problems.append((line, f"line {line}: {msg}" if line is not None else msg)) + + def new_value( + self, + raw: int, + type_: str, + op: int | None, + block: int | None, + pos: int, + loc: str, + ) -> int: + if raw in self.vmap: + raise ValueError("a value appears twice in the walk") + idx = len(self.values) + tree = self.locp.parse(loc) + name = _loc_site(tree)[2] if tree is not None else None + self.values.append([type_, op, block, pos, name]) + self.vmap[raw] = idx + return idx + + def text_loc(self, t: _TOp) -> tuple | None: + masked, raw = ( + (t.close_masked, t.close_raw) if t.opens_region else (t.masked, t.raw) + ) + span = _trailing_loc(masked) + (self.loc_lines if span is not None else self.no_loc_lines).append(t.line_no) + if span is None: + return None + return self.locp.parse(raw[span[0] : span[1]]) + + # ── traversal ── + def run(self) -> Module: + root_w = [i for i, b in enumerate(self.bw.ops) if b.block_id is None] + top = self.tree.root.regions[0][0].ops if self.tree.root.regions[0] else [] + if len(root_w) != 1 or len(top) != 1: + self.bad( + None, + f"expected one top-level op: text has {len(top)}, bindings {len(root_w)}", + ) + raise self.failure() + # task: ("op", text op, walk index, parent block index | None, position) + # or ("implicit", walk index, parent block index, position) + stack: list[tuple] = [("op", top[0], root_w[0], None, 0)] + while stack: + task = stack.pop() + if task[0] == "op": + children = self.visit(*task[1:]) + else: + children = [] + self.implicit(*task[1:]) + stack.extend(reversed(children)) + if len(self.ops) != len(self.bw.ops): + self.bad( + None, f"{len(self.ops)} aligned ops != {len(self.bw.ops)} walked ops" + ) + if self.loc_lines: + for line in self.no_loc_lines: + self.bad(line, "op prints no loc while the module prints locs") + if self.problems: + raise self.failure() + return self.freeze() + + def failure(self) -> MisalignedModule: + first = next((ln for ln, _ in self.problems if ln is not None), None) + return MisalignedModule([m for _, m in self.problems], first) + + def implicit(self, wi: int, block: int, pos: int) -> None: + b = self.bw.ops[wi] + idx = len(self.ops) + self.stats["implicit_ops"] += 1 + path = self.blocks[block].path + self.ops.append( + _OpB( + b.name, + (), + (), + (), + (), + {}, + path, + pos, + None, + None, + None, + (), + None, + True, + [], + ) + ) + self.blocks[block].ops.append(idx) + + def visit(self, t: _TOp, wi: int, block: int | None, pos: int) -> list: + b = self.bw.ops[wi] + line = t.line_no + idx = len(self.ops) + path = self.blocks[block].path if block is not None else () + if block is not None: + self.blocks[block].ops.append(idx) + structural_ok = True + if t.name != b.name: + self.bad(line, f"text op {t.name!r} != walked op {b.name!r}") + structural_ok = False + if t.n_results != len(b.results): + self.bad( + line, + f"{t.name}: {t.n_results} printed results != {len(b.results)} walked", + ) + structural_ok = False + # printed regions: trailing empty regions may be omitted + if len(t.regions) > len(b.region_ids) or any( + b.region_sizes[i] for i in range(len(t.regions), len(b.region_ids)) + ): + self.bad( + line, + f"{t.name}: {len(t.regions)} printed regions vs walked sizes {list(b.region_sizes)}", + ) + structural_ok = False + # the op's loc: the printed one, which must equal the loc the bindings + # report for each of its results (the one source for a zero-result op) + try: + tloc = self.text_loc(t) + for k, s in enumerate(b.result_locs): + tree = self.locp.parse(s) + if tree is not None: + self.stats["result_locs"] += 1 + if tree != tloc: + self.bad( + line, + f"{t.name}: result {k} loc differs: text {tloc} vs bindings {tree}", + ) + except ValueError as e: + self.bad(line, f"{t.name}: unreadable loc: {e}") + tloc = None + site, callers, lname = _loc_site(tloc) + attrs = dict(b.attrs) + if structural_ok: + for p in b.problems: + self.bad(line, f"{b.name}: {p}") + try: + attrs.update(_text_attrs(t, b.name)) + except ValueError as e: + self.bad(line, f"{b.name}: {e}") + self.stats["attrs"] += len(attrs) + rec = _OpB( + b.name, + b.operands, + b.operand_types, + (), + b.result_types, + attrs, + path, + pos, + line, + t.close_line if t.opens_region else line, + site, + callers, + lname, + False, + [], + ) + self.ops.append(rec) + rec.results = tuple( + self.new_value(raw, b.result_types[k], idx, None, k, b.result_locs[k]) + for k, raw in enumerate(b.results) + ) + if not structural_ok: + rec.regions = [() for _ in b.region_ids] + return [] + children: list[tuple] = [] + for ri, rid in enumerate(b.region_ids): + children += self.region(t, b, idx, ri, rid, rec) + return children + + def region(self, t: _TOp, b: _BOp, idx: int, ri: int, rid: int, rec: _OpB) -> list: + line = t.line_no + wblocks = self.bw.region_blocks.get(rid, []) + if len(wblocks) != b.region_sizes[ri]: + self.bad(line, f"{t.name}: region {ri} has a block without ops") + tblocks = t.regions[ri] if ri < len(t.regions) else [] + if not tblocks and len(wblocks) == 1 and b.name in _ELIDES_YIELD: + only = self.bw.blocks[wblocks[0]].ops + tail = self.bw.ops[only[-1]] if len(only) == 1 else None + if tail is not None and tail.name == "scf.yield" and not tail.operands: + # `{ }`: the region's one block held only the elided yield + tblocks = [_TBlock(None, [], line)] + if len(tblocks) != len(wblocks): + self.bad( + line, + f"{t.name}: region {ri}: {len(tblocks)} printed blocks != {len(wblocks)} walked", + ) + rec.regions.append(()) + return [] + keys = [] + children: list[tuple] = [] + for bi, (tb, rbid) in enumerate(zip(tblocks, wblocks)): + bb = self.bw.blocks[rbid] + bidx = len(self.blocks) + keys.append(bidx) + if tb.label is None and bi != 0: + self.bad(tb.line_no, f"{t.name}: unlabeled block {ri}.{bi}") + blk = _BlockB( + idx, ri, bi, tb.label, (), bb.arg_types, (), [], rec.path + (bidx,) + ) + self.blocks.append(blk) + blk.args = tuple( + self.new_value(raw, bb.arg_types[ai], None, bidx, ai, bb.arg_locs[ai]) + for ai, raw in enumerate(bb.args) + ) + blk.arg_names = tuple(self.values[v][4] for v in blk.args) + # ops, with the one tolerated gap: an elided trailing scf.yield + n_text, n_walk = len(tb.ops), len(bb.ops) + elided = False + if n_text != n_walk: + tail = self.bw.ops[bb.ops[-1]] if bb.ops else None + elided = ( + n_text == n_walk - 1 + and tail is not None + and tail.name == "scf.yield" + and not tail.operands + and b.name in _ELIDES_YIELD + ) + if not elided: + self.bad( + tb.line_no, + f"{t.name}: block {ri}.{bi}: {n_text} printed ops != {n_walk} walked", + ) + continue + for p, (tchild, wchild) in enumerate(zip(tb.ops, bb.ops)): + children.append(("op", tchild, wchild, bidx, p)) + if elided: + children.append(("implicit", bb.ops[-1], bidx, n_text)) + rec.regions.append(tuple(keys)) + return children + + def freeze(self) -> Module: + values = tuple(Value(i, *v) for i, v in enumerate(self.values)) + blocks = tuple( + Block( + i, + b.op, + b.region, + b.position, + b.label, + b.args, + b.arg_types, + b.arg_names, + tuple(b.ops), + ) + for i, b in enumerate(self.blocks) + ) + ops = [] + funcs = [] + for i, r in enumerate(self.ops): + operands = tuple(self.vmap[v] for v in r.operands_raw) + ops.append( + Op( + i, + r.name, + operands, + r.operand_types, + r.results, + r.result_types, + _FrozenMap(r.attrs), + tuple(r.regions), + r.path, + r.position, + r.line_no, + r.end_line, + r.loc, + r.callers, + r.loc_name, + r.implicit, + ) + ) + if r.name == "tt.func": + args: tuple[FuncArg, ...] = () + if r.regions and r.regions[0]: + entry = blocks[r.regions[0][0]] + args = tuple( + FuncArg(k, v, entry.arg_types[k], entry.arg_names[k]) + for k, v in enumerate(entry.args) + ) + visibility = r.attrs.get("visibility", "public") + funcs.append(Func(i, r.attrs["sym_name"], visibility, args)) + self.stats.update( + ops=len(ops), blocks=len(blocks), values=len(values), funcs=len(funcs) + ) + return Module( + tuple(ops), + blocks, + values, + tuple(funcs), + _FrozenMap(self.stats), + ) + + +# ─────────────────────────── entry point ─────────────────────────── + + +def _walk(text: str, scan_text: str | None = None) -> Module: + """Uncached walk. ``scan_text`` (tests only) feeds the text layer a + different string than the bindings parse.""" + try: + data = text.encode("utf-8") + except UnicodeEncodeError as e: + raise ModuleParseError(f"the text is not encodable as UTF-8: {e}") from None + _screen(text) + bw = _bind_walk(data) + try: + tree = _scan_text(text if scan_text is None else scan_text) + except _TextError as e: + raise MisalignedModule( + [f"line {e.line_no}: {e.msg}" if e.line_no else e.msg], e.line_no + ) from None + try: + return _Aligner(tree, bw).run() + except MisalignedModule: + raise + except (ValueError, KeyError, IndexError, AssertionError, TypeError) as e: + # a text shape the aligner does not model: fail closed + raise MisalignedModule([f"aligner: {type(e).__name__}: {e}"]) from None + + +_CACHE_SIZE = 32 +_CACHE: collections.OrderedDict[ + bytes, Module | tuple[tuple[str, ...], int | None] +] = collections.OrderedDict() +_CACHE_LOCK = threading.Lock() + + +def walk_module(text: str) -> Module: + """Walk one printed TTIR module (see the module docstring). + + Raises ``MisalignedModule`` when the text layer and the bindings + disagree and ``ModuleParseError`` when the MLIR parser rejects the text + (or the text holds a construct the parser cannot be handed, or the + installed Triton is not the release the walk reads). Results (and + misalignments) are cached by sha256 of the text, so a text is parsed + once while it stays among the last ``_CACHE_SIZE`` distinct texts. + """ + from ..core.config import ir_triton_unsupported + + unsupported = ir_triton_unsupported() + if unsupported is not None: + raise ModuleParseError(unsupported) + key = hashlib.sha256(text.encode("utf-8", "surrogatepass")).digest() + with _CACHE_LOCK: + hit = _CACHE.get(key) + if hit is not None: + _CACHE.move_to_end(key) + if hit is None: + try: + hit = _walk(text) + except MisalignedModule as e: + hit = (e.problems, e.line_no) + with _CACHE_LOCK: + _CACHE[key] = hit + while len(_CACHE) > _CACHE_SIZE: + _CACHE.popitem(last=False) + if isinstance(hit, Module): + return hit + raise MisalignedModule(*hit) + + +def _after_fork_in_child() -> None: + """A fork while another thread holds ``_PARSE_LOCK`` / ``_CACHE_LOCK`` + leaves the child a lock nobody releases (its next walk would hang), and a + fork inside a parse window leaves the child's fd 2 in the capture buffer. + The child gets fresh locks and its fd 2 back; the capture's two fds stay + open (their owner is the parent's thread, which does not run here).""" + global _PARSE_LOCK, _CACHE_LOCK, _REDIRECT + _PARSE_LOCK = threading.Lock() + _CACHE_LOCK = threading.Lock() + redirect, _REDIRECT = _REDIRECT, None + if redirect is not None: + try: + os.dup2(redirect[0], 2) + except OSError: + pass + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=_after_fork_in_child) diff --git a/tilelens/ir/ttir_reader.py b/tilelens/ir/ttir_reader.py new file mode 100644 index 00000000..c8ffbdbb --- /dev/null +++ b/tilelens/ir/ttir_reader.py @@ -0,0 +1,1596 @@ +"""TTIR reader shared by the compiled-mode clients. + +Reads the pre-optimization Triton IR (TTIR) of one kernel specialization +into an ``AccessGraph``: the kernel's function arguments, every global +memory access (``tt.load`` / ``tt.store`` / ``tt.atomic_rmw`` / +``tt.atomic_cas``) as an *element offset* expression relative to a base +pointer argument, the mask guarding it, and the loop structure. Scalar +arguments (``n_elements``, ``M``, strides, ...) stay symbolic (``Param`` +nodes) and are substituted with concrete launch values later; +``tl.constexpr`` values are already folded into TTIR constants. + +Why TTIR (not TTGIR): element addressing is cleanest here, before +layouts/pipelining add noise, and TTIR has no indirect loads unless the +kernel itself gathers — the data-dependent case, marked with ``DataDep``. + +This module is mechanism-only: it reads and flags (``DataDep`` markers, +``guarded`` accesses, width obligations, ``UnsupportedTTIR``); what to do +about a flagged or unsupported kernel is the policy of each client that +consumes the graph. It either represents the IR faithfully or raises an +``UnsupportedTTIR`` whose ``kind`` says what it cannot represent. + +Structure comes from ``_mlir_walk``: Triton's MLIR bindings parse the text +and give the op tree, values, types and the attributes the reader reads; a +line scan of the same text adds each op's line and the locs the bindings do +not expose. The reader walks that op tree over regions and blocks, and its +environment is keyed by the walk's value indices, never by printed SSA +names. + +Address model: ``tt.addptr(base, off)`` accumulates an ELEMENT offset; the +byte address is ``base.data_ptr() + offset * elem_size``. An access is OOB +iff, for some program id / arange lane / loop iteration with its mask true, +the element offset escapes ``[0, numel)`` of its base tensor. + +Integer model: every integer term denotes the IR value's SIGNED +reading as an unbounded integer (i1 terms are booleans, 0/1). That reading +is exact while the ``width_obligations`` of the access hold (they include +the loop's own increment); a consumer that evaluates terms with unbounded +integers must discharge (or report) them. ``Param`` values are the IR's +signed reading of the argument too. Division or remainder by zero is +undefined in the IR rather than a width condition, and so is an scf.for +step that is not positive: no obligation excludes them, that is the +consumer's call. + +Terms are frozen dataclasses with the generated ``==`` / ``hash`` / +``repr``, which recurse: on a term deeper than Python's recursion limit +(``kernel_deep_chain`` has more than 1000 levels) they, and +``copy.deepcopy``, raise ``RecursionError``; only ``pickle`` and this +module's own walkers are iterative. Consumers of such graphs key memo +tables by identity. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass, field, replace +from enum import Enum +from typing import Callable, Iterable, Iterator, NoReturn, Sequence + +from ._mlir_walk import ( + MisalignedModule, + Module, + ModuleParseError, + Op, + SourceLoc, + parse_type, + walk_module, +) + + +class TTIRKind(str, Enum): + """What the reader cannot represent. Only representational limits live + here; refusals a client makes about a graph it did receive (a + data-dependent mask, a CAS value, ...) are that client's own kinds.""" + + INDIRECT_ADDRESS = "indirect-address" + DATA_DEPENDENT_BOUND = "data-dependent-bound" + NESTED_LOOP = "nested-loop" + CONTROL_FLOW = "control-flow" + OUT_OF_VOCABULARY = "out-of-vocabulary" + CALL = "call" + LOOP_VARIANT_ADVANCE = "loop-variant-advance" + INLINE_ASM = "inline-asm" + READER_MISALIGNMENT = "reader-misalignment" + UNPARSABLE = "unparsable" + OTHER = "other" + + def __str__(self) -> str: + return self.value + + def __format__(self, spec: str) -> str: + return format(self.value, spec) + + +class UnsupportedTTIR(Exception): + """Raised for constructs outside the compiled-mode model (indirect or + data-dependent addressing, block pointers, nested loops, calls, ...). + + ``kind`` (a :class:`TTIRKind`) is the machine-readable class of the + limitation, ``message`` the human-readable detail, ``line_no`` the + refused op's line in the TTIR text and ``loc`` its user-source location + (None when unknown). Clients read these fields; ``str(exc)`` is the + message alone. + """ + + def __init__( + self, + kind: TTIRKind | str, + message: str, + *, + line_no: int | None = None, + loc: SourceLoc | None = None, + ) -> None: + super().__init__(message) + self.kind = TTIRKind(kind) + self.message = message + self.line_no = line_no + self.loc = loc + + def __reduce__(self): + return ( + type(self), + (self.kind, self.message), + {"line_no": self.line_no, "loc": self.loc}, + ) + + +# ─────────────────────────── address-expression terms ─────────────────────────── +# A small lazily-evaluated tree. Leaves that are only known at launch time +# (scalar kernel args) are Param nodes; pid / arange / loop variables become +# free variables with range constraints in a client's query. +# +# Bin, Cmp and IntCast also record the op they came from (``line_no``, +# ``loc``) for width obligations; those two fields take no part in +# equality, hashing or repr, so equal expressions compare equal wherever +# they were computed. + + +@dataclass(frozen=True) +class Const: + value: int + + +@dataclass(frozen=True) +class Pid: + axis: int # 0=x, 1=y, 2=z + + +@dataclass(frozen=True) +class NumPrograms: + """``tt.get_num_programs axis`` — the launch grid size along ``axis``. + Uniform across program instances, but it PARAMETERIZES the kernel's + behavior by the grid.""" + + axis: int + + +@dataclass(frozen=True) +class Arange: + ssa: str # unique per make_range site (the walk's result value index) + start: int + end: int + # Which tensor dimension this lane index varies along. -1 = 1D / not yet + # placed; set and kept current by expand_dims. Consumers key a lane + # variable by (dim, end - start), not by make_range site: every tensor + # one access combines has the access's shape, so all aranges along one + # dim with one extent index the SAME position there (tl.arange(0, 16) + + # tl.arange(16, 32) is 2i + 16, not i + j + 16), while a single + # make_range reused for several dimensions of a tile (triton does this) + # is one independent variable per dimension, or the modeled footprint + # would collapse to the diagonal. + dim: int = -1 + + +@dataclass(frozen=True) +class Param: + name: str # scalar kernel argument, substituted per launch + + +@dataclass(frozen=True) +class IterArgOffset: + """The element-offset contribution of a loop-carried pointer at the + current iteration: ``offset0 + k * delta`` (resolved from + ``graph.iter_args[arg_id]`` at eval time).""" + + arg_id: int + + +@dataclass(frozen=True) +class LoopVar: + """The scf.for induction variable; a free variable over the iterations + that run (e.g. it appears in masks like ``K - k*BLOCK_K``).""" + + loop_ssa: str + + +# Integer ops (Bin.op): the arith op each spelling reads, signed first. +# "//" and "%" truncate toward zero (divsi / remsi); the "u"-prefixed ops +# read their operands unsigned (divui / remui / minui / maxui). +_BIN_OPS = { + "arith.addi": "+", + "arith.subi": "-", + "arith.muli": "*", + "arith.divsi": "//", + "arith.remsi": "%", + "arith.minsi": "min", + "arith.maxsi": "max", + "arith.divui": "u//", + "arith.remui": "u%", + "arith.minui": "umin", + "arith.maxui": "umax", +} +UNSIGNED_BIN_OPS = frozenset({"u//", "u%", "umin", "umax"}) +UNSIGNED_PREDICATES = frozenset({"ult", "ule", "ugt", "uge"}) +_SIGNED_PREDICATES = frozenset({"slt", "sle", "sgt", "sge"}) + + +@dataclass(frozen=True) +class Bin: + op: str # + - * // % min max, or u// u% umin umax (see _BIN_OPS) + a: "Term" + b: "Term" + # Width of the integer result (the IR type). None for the element-offset + # sum a ``tt.addptr`` accumulates, which is address arithmetic, not an + # IR integer. + bits: int | None = None + line_no: int | None = field(default=None, compare=False, repr=False) + loc: SourceLoc | None = field(default=None, compare=False, repr=False) + + +@dataclass(frozen=True) +class Cmp: + pred: str # eq/ne, slt/sle/sgt/sge, ult/ule/ugt/uge (unsigned reads) + a: "Term" + b: "Term" + bits: int | None = None # operand width + line_no: int | None = field(default=None, compare=False, repr=False) + loc: SourceLoc | None = field(default=None, compare=False, repr=False) + + +@dataclass(frozen=True) +class BoolBin: + op: str # and / or + a: "Term" + b: "Term" + + +@dataclass(frozen=True) +class Select: + cond: "Term" + t: "Term" + f: "Term" + + +@dataclass(frozen=True) +class Not: + """Boolean negation — the path condition of an scf.if else-region.""" + + a: "Term" + + +@dataclass(frozen=True) +class IntCast: + """``arith.trunci`` / ``extsi`` / ``extui`` of ``x`` from ``src_bits`` to + ``dst_bits`` (a cast is never a value passthrough). Its value is + ``x`` exactly when the cast's width obligation holds (trunci: ``x`` + fits the destination; extui: ``x`` is non-negative); extsi always + preserves the signed reading. ``extsi`` from i1 (true -> -1) is read as + ``0 - extui(x)`` and never appears as an IntCast.""" + + kind: str # "trunci" | "extsi" | "extui" + src_bits: int + dst_bits: int + x: "Term" + line_no: int | None = field(default=None, compare=False, repr=False) + loc: SourceLoc | None = field(default=None, compare=False, repr=False) + + +# Sentinel for a value loaded from memory (tt.load result) or computed from +# loaded data (arith.*f, tt.dot, ...). If one ever reaches an address it +# means data-dependent addressing -> unsupported; in a mask it is dropped +# (``mask_dropped``), under an scf.if it leaves the branch ``guarded``. +@dataclass(frozen=True) +class DataDep: + why: str = "value derived from loaded data" + + +@dataclass(frozen=True) +class Observed: + """The OLD value observed by the atomic at ``graph.accesses[access_index]``: + a fresh per-program-instance symbol, NOT a function of other leaves. The + reader binds an INTEGER-typed ``tt.atomic_rmw`` / ``tt.atomic_cas`` + result to this instead of ``DataDep`` so downstream masks and branch + conditions stay modelable; float-typed atomic results keep the DataDep + fallback. What an observation means (a free variable, a modeled value, + a refusal in an address) is each consumer's policy; find them with the + graph-aware :func:`mentions_observed`. + + A tensor atomic observes one old value per lane, and the symbol stands + for the value at the lane the surrounding term is read at. It has no + lane placement of its own, so the reader never lets two lanes of one + tensor observation meet: an ``expand_dims`` of a term holding one + degrades to DataDep.""" + + access_index: int + + +Term = ( + Const + | Pid + | NumPrograms + | Arange + | Param + | IterArgOffset + | LoopVar + | Bin + | Cmp + | BoolBin + | Select + | Not + | IntCast + | DataDep + | Observed +) + + +def _children(t: object) -> tuple: + if isinstance(t, (Bin, Cmp, BoolBin)): + return (t.a, t.b) + if isinstance(t, Select): + return (t.cond, t.t, t.f) + if isinstance(t, (Not,)): + return (t.a,) + if isinstance(t, IntCast): + return (t.x,) + return () + + +def _nodes( + roots: Iterable[object], iter_args: "Sequence[IterArgInfo] | None" +) -> Iterator[object]: + """Every node reachable from ``roots`` in pre-order, each once (by + identity); with ``iter_args``, an IterArgOffset also reaches its + IterArgInfo's ``offset0`` and ``delta``. Iterative: terms can be deeper + than Python's recursion limit.""" + seen: set[int] = set() + stack = [r for r in reversed(list(roots)) if r is not None] + while stack: + t = stack.pop() + if id(t) in seen: + continue + seen.add(id(t)) + yield t + if isinstance(t, IterArgOffset): + if iter_args is not None: + info = iter_args[t.arg_id] + stack += [info.delta, info.offset0] + continue + stack.extend(reversed(_children(t))) + + +def mentions_observed(term: object, graph: "AccessGraph") -> bool: + """True when ``term`` reaches an :class:`Observed` leaf, through the + graph's loop-carried pointers (an IterArgOffset's ``offset0`` and + ``delta``) included.""" + return any(isinstance(n, Observed) for n in _nodes((term,), graph.iter_args)) + + +# DataDep is also the generic unknown-value top (loop accumulators, +# unmodeled ops, ...). Only these ``why`` prefixes mean the value truly +# derives from MEMORY CONTENTS — refusals classify just those as +# indirection; the rest are modeling gaps and keep the default kind. +_MEMORY_WHYS = ( + "loaded value", + "atomic result", + "arith over loaded data", + "cmpi over loaded data", + "select over loaded data", + "bool op over loaded data", +) + + +def _from_memory(v: object) -> bool: + return isinstance(v, DataDep) and v.why.startswith(_MEMORY_WHYS) + + +@dataclass(frozen=True) +class PtrValue: + """A pointer-typed value: base argument + accumulated element offset (a + single lane's offset; arange/loop free vars cover all lanes and + iterations in the query).""" + + base_param: str + offset: Term + + +# ─────────────────────────── graph structures ─────────────────────────── + + +@dataclass(frozen=True) +class FuncArg: + name: str # the Python parameter name (NameLoc), else "arg" + is_ptr: bool + elem_bits: int # for ptr args: pointee width; 0 for scalars + # For integer scalars: the IR width (i32 -> 32, i1 -> 1); 0 otherwise. + int_bits: int = 0 + + +@dataclass(frozen=True) +class AccessEvent: + kind: str # "load" | "store" | "atomic_rmw" | "atomic_cas" + base_param: str + offset: Term + mask: Term | None # None = unconditional access + elem_bits: int + loc: SourceLoc | None + line_no: int + # True when some enclosing scf.if condition could NOT be modeled (it + # derives from loaded data). The access is then checked as if + # unconditional: UNSAT stays a sound proof, but a SAT model may sit in a + # branch the launch never takes. Modeled conditions ride in ``path`` + # instead and do not set this flag. + guarded: bool = False + # Conjunction of the MODELED enclosing branch conditions, with + # else-regions negated (Not). The access executes iff path ∧ mask. + path: Term | None = None + # True when the access sits inside the scf.for body: it executes once + # per iteration — and NOT AT ALL when the launch's trip count is zero, + # which consumers must model (a zero-trip loop has no footprint). + in_loop: bool = False + # True when the printed mask operand derived from loaded data and was + # over-approximated as FREE (mask=None): dropping a constraint only + # widens the modeled footprint, so UNSAT stays a sound proof — but a SAT + # model may pick a lane the real mask disables. + mask_dropped: bool = False + + +@dataclass(frozen=True) +class IterArgInfo: + """A loop-carried pointer: ``offset0 + k * delta`` at iteration k. A + tile of one expanded (``tt.expand_dims``) inside the loop is an entry of + its own, the same pointer with ``offset0`` and ``delta`` expanded alike, + so its lanes sit at their positions in the expanded shape.""" + + arg_id: int + base_param: str + offset0: Term + delta: Term # per-iteration element advance (loop-invariant) + loop_ssa: str = "" # the scf.for this iter_arg belongs to (LoopInfo.loop_ssa) + + +@dataclass(frozen=True) +class LoopInfo: + loop_ssa: str + lower: Term + upper: Term + step: Term + # Width of the induction variable, and whether the loop compares its + # bounds unsigned (``scf.for unsigned``: the bounds' width obligations + # then require them non-negative). + bits: int | None = None + unsigned: bool = False + line_no: int | None = None + loc: SourceLoc | None = None + + +@dataclass(frozen=True) +class AccessGraph: + kernel_name: str + func_args: tuple[FuncArg, ...] + accesses: tuple[AccessEvent, ...] + loop: LoopInfo | None + # Loop-carried pointers, indexed by arg_id (``iter_args[k].arg_id == k``), + # expanded tiles of them included (see IterArgInfo). + iter_args: tuple[IterArgInfo, ...] = () + + def __post_init__(self) -> None: + for name in ("func_args", "accesses", "iter_args"): + object.__setattr__(self, name, tuple(getattr(self, name))) + + def arg(self, name: str) -> FuncArg | None: + for a in self.func_args: + if a.name == name: + return a + return None + + +# ─────────────────────────── width obligations ─────────────────────────── + + +@dataclass(frozen=True) +class WidthObligation: + """``term`` must fit the width it is read at: with ``signed``, + ``-2**(bits-1) <= term < 2**(bits-1)``; otherwise ``0 <= term < 2**bits``. + ``line_no`` / ``loc`` are the op that imposes it; ``role`` is the part + of the access it comes from (see :func:`width_obligations`).""" + + term: Term + bits: int + signed: bool + line_no: int | None + loc: SourceLoc | None + role: str # "loop" | "path" | "mask" | "offset" + + +def width_obligations( + graph: AccessGraph, access: AccessEvent +) -> tuple[WidthObligation, ...]: + """The conditions under which the unbounded-integer reading of + ``access`` (its offset, mask and path, loop-carried pointers resolved, + and the loop of an access in the loop) equals the IR's fixed-width + arithmetic: + + * every integer ``Bin`` result fits its width, signed; + * the quotient of a signed remainder fits too (``INT_MIN % -1`` is + undefined in the IR); + * the operands of an unsigned op or predicate are non-negative; + * a ``trunci`` operand fits the destination width (to i1: is 0 or 1); + * an ``extui`` operand is non-negative; + * the bounds of an ``unsigned`` loop are non-negative; + * the loop's increment does not wrap: ``upper - 1 + step``, which + bounds the last iterate plus ``step``, fits the induction variable + (sufficient, not necessary). + + ``role`` says where an obligation comes from: ``"loop"`` (the loop's + bounds and increment), ``"path"``, ``"mask"`` or ``"offset"``. A term + node several parts share is listed once, under the first role in that + order, and the order is also the discipline for discharging them: an + ``offset`` obligation only matters where the access executes (path and + mask hold), a ``mask`` one only where the path holds, and ``path`` and + ``loop`` ones hold unconditionally (the increment's only when the loop + runs at least once). Both arms of a ``Select`` are listed, without its + condition, so checking an arm unconditionally may report an overflow in + the arm the launch does not take. + + Mechanism only: listed role by role in walk order, each (term node, + width, signedness) once; nothing is evaluated.""" + loop = graph.loop if access.in_loop else None + out: list[WidthObligation] = [] + seen: set[tuple[int, int, bool]] = set() + + def need(term: object, bits: int, signed: bool, site: object, role: str) -> None: + key = (id(term), bits, signed) + if key not in seen: + seen.add(key) + out.append( + WidthObligation( + term, # type: ignore[arg-type] + bits, + signed, + getattr(site, "line_no", None), + getattr(site, "loc", None), + role, + ) + ) + + groups: list[tuple[str, tuple[object, ...]]] = [] + if loop is not None: + if loop.bits is not None: + if loop.unsigned: + for bound in (loop.lower, loop.upper, loop.step): + need(bound, loop.bits, False, loop, "loop") + latch = Bin( + "+", + Bin("-", loop.upper, Const(1), loop.bits, loop.line_no, loop.loc), + loop.step, + loop.bits, + loop.line_no, + loop.loc, + ) + need(latch, loop.bits, not loop.unsigned, loop, "loop") + groups.append(("loop", (loop.lower, loop.upper, loop.step))) + groups += [ + ("path", (access.path,)), + ("mask", (access.mask,)), + ("offset", (access.offset,)), + ] + for role, roots in groups: + for n in _nodes(roots, graph.iter_args): + if isinstance(n, Bin) and n.bits is not None: + need(n, n.bits, True, n, role) + if n.op in UNSIGNED_BIN_OPS: + need(n.a, n.bits, False, n, role) + need(n.b, n.bits, False, n, role) + elif n.op == "%": + quotient = Bin("//", n.a, n.b, n.bits, n.line_no, n.loc) + need(quotient, n.bits, True, n, role) + elif isinstance(n, Cmp) and n.pred in UNSIGNED_PREDICATES and n.bits: + need(n.a, n.bits, False, n, role) + need(n.b, n.bits, False, n, role) + elif isinstance(n, IntCast): + if n.kind == "trunci": + need(n.x, n.dst_bits, n.dst_bits > 1, n, role) + elif n.kind == "extui" and n.src_bits > 1: + need(n.x, n.src_bits, False, n, role) + return tuple(out) + + +# ─────────────────────────── the reader ─────────────────────────── + + +def parse_ttir(text: str) -> AccessGraph: + """Read one TTIR module into an AccessGraph (single-path model: at most + one scf.for, structured scf.if only). + + Raises :class:`UnsupportedTTIR` for anything the graph cannot represent: + a text the walk layer cannot align (``READER_MISALIGNMENT``) or does not + parse (``UNPARSABLE``: the MLIR parser rejects it, it holds a + block-pointer type, or the installed Triton is not 3.8), indirect + addressing, + pointers outside global memory, nested/while loops, unstructured + control flow, calls, loop-variant addresses, inline asm that is impure + or handed an address, or any op outside the address vocabulary that + could touch memory. + """ + try: + module = walk_module(text) + except MisalignedModule as e: + raise UnsupportedTTIR( + TTIRKind.READER_MISALIGNMENT, "; ".join(e.problems), line_no=e.line_no + ) from e + except ModuleParseError as e: + raise UnsupportedTTIR( + TTIRKind.UNPARSABLE, e.diagnostic, line_no=e.line_no + ) from e + return _Builder(module).build() + + +# Memory ops outside the modeled vocabulary (TMA descriptors, ...: targets +# without native TMA lower descriptors to pointer math before TTIR is +# printed; sm90+ and hip TDM targets such as gfx1250 keep them). +_MEMORY_PREFIXES = ("tt.descriptor_", "tt.experimental_") +_ACCESS_OPS = frozenset({"tt.load", "tt.store", "tt.atomic_rmw", "tt.atomic_cas"}) +# Dialects whose ops the reader reads as data or control; from any other +# dialect only the inert ops (below) are accepted. +_DIALECTS = frozenset({"tt", "arith", "math", "scf", "cf", "ub"}) + + +# Result-free ops without memory effects the reader skips (a barrier orders +# memory, it addresses none). tl.debug_barrier() is `ttg.barrier all`: the +# one ttg op accepted, not the dialect. gpu.barrier, which the frontend +# never emits, is not (the bindings still parse it). +_INERT = frozenset( + { + "tt.return", + "scf.yield", + "tt.reduce.return", + "tt.scan.return", + "tt.print", + "tt.assert", + "llvm.intr.assume", + "ttg.barrier", + } +) +# Ops that hand the reader values it cannot see into (their operands must +# carry no address, not even as an integer). +_OPAQUE = frozenset({"tt.elementwise_inline_asm", "tt.extern_elementwise"}) +# Structured-region nesting the reader follows (deeper refuses, instead of +# exhausting Python's recursion limit). +_MAX_DEPTH = 200 +_LOOP_SSA = "%loop" # the single-path model has at most one loop +# The DataDep reason of a value carried by the loop that is not a pointer: +# in an address it makes the address loop-variant. +_LOOP_CARRIED = "loop accumulator" +_RE_ADDR_SPACE = re.compile(r", (\d+)>$") + + +def _addr_space(type_text: str) -> int: + """Address space of a (tensor of) ``!tt.ptr`` type; 1 (global memory) + when unprinted, as the printer elides it.""" + m = _RE_ADDR_SPACE.search(parse_type(type_text).elem) + return int(m.group(1)) if m else 1 + + +@dataclass +class _IfFrame: + """Walker state for one scf.if region being read.""" + + cond: "Term | None" # modeled condition; None -> accesses stay `guarded` + branch: str = "then" + + +class _ForFrame: + """Marks the scf.for body being read.""" + + +def _pointee_bits(pointee: str | None) -> int: + if pointee is None: + return 0 + if pointee.startswith("!tt.ptr<"): + return 64 # a pointer to pointers + return parse_type(pointee).int_bits or parse_type(pointee).float_bits or 0 + + +class _Builder: + def __init__(self, module: Module) -> None: + self.m = module + # value index -> Term (int/bool), PtrValue, or DataDep + self.env: dict[int, object] = {} + self.func_args: list[FuncArg] = [] + self.accesses: list[AccessEvent] = [] + self.iter_args: list[IterArgInfo] = [] + self.loop: LoopInfo | None = None + self.loops_opened = 0 + self.frames: list[_IfFrame | _ForFrame] = [] + self.depth = 0 + # (arg_id, axis) -> the arg_id of that iter_arg's tile expanded at + # axis inside the loop (its delta is filled in when the loop closes) + self.expanded: dict[tuple[int, int], int] = {} + # access indices of the tensor atomics (one observation per lane) + self.lane_observed: set[int] = set() + + # ── helpers ── + def refuse(self, kind: TTIRKind, message: str, op: Op | None) -> NoReturn: + raise UnsupportedTTIR( + kind, + message, + line_no=op.line_no if op is not None else None, + loc=op.loc if op is not None else None, + ) + + def val(self, v: int) -> object: + got = self.env.get(v) + if got is None: + # A value the reader never bound (defined in a region it does not + # read): be conservative. + return DataDep(f"unresolved value {v}") + return got + + def bind(self, op: Op, value: object) -> None: + for r in op.results: + self.env[r] = value + + def as_term(self, v: object, ctx: str, op: Op) -> Term: + if isinstance(v, DataDep): + self.refuse(TTIRKind.OTHER, f"{ctx}: data-dependent ({v.why})", op) + if isinstance(v, PtrValue): + self.refuse(TTIRKind.OTHER, f"{ctx}: pointer used as integer", op) + return v # type: ignore[return-value] + + def int_bits(self, v: int) -> int | None: + return parse_type(parse_type(self.m.values[v].type).elem).int_bits + + def branch_state(self) -> tuple[bool, Term | None, bool]: + """(guarded, path, in_loop) for an access under the open frames: + ``guarded`` if any enclosing condition is unmodeled; ``path`` is the + conjunction of the modeled ones (else-regions negated), outermost + first; ``in_loop`` when the scf.for body encloses the access.""" + guarded = False + path: Term | None = None + in_loop = False + for f in self.frames: + if isinstance(f, _ForFrame): + in_loop = True + continue + if f.cond is None: + guarded = True + continue + c: Term = f.cond if f.branch == "then" else Not(f.cond) + path = c if path is None else BoolBin("and", path, c) + return guarded, path, in_loop + + def arg(self, name: str) -> FuncArg | None: + return next((a for a in self.func_args if a.name == name), None) + + def global_pointers(self, types: Iterable[str], op: Op) -> None: + """Element offsets are into global memory: refuse a pointer into any + other address space (shared memory, ...).""" + for t in types: + if "!tt.ptr<" in t and _addr_space(t) != 1: + self.refuse( + TTIRKind.OUT_OF_VOCABULARY, + f"pointer type {t} is not in global memory (address space 1)", + op, + ) + + # ── the module ── + def build(self) -> AccessGraph: + m = self.m + if not m.funcs: + self.refuse(TTIRKind.OTHER, "no tt.func found (not TTIR?)", None) + for op in m.ops: + if len(op.path) == 1 and op.name != "tt.func": + self.refuse( + TTIRKind.OUT_OF_VOCABULARY, + f"{op.name} at module level is not TTIR", + op, + ) + # A call's callee runs in its own frame with its own arguments; the + # graph has no call model, so any call, and any second function, + # refuses (the callee name comes from the symbol, quoted or not). + calls = [op for op in m.ops if op.name == "tt.call"] + if calls: + first = calls[0] + self.refuse( + TTIRKind.CALL, + f"tt.call to {first.attrs.get('callee')!r}: calls are not modeled", + first, + ) + if len(m.funcs) > 1: + extra = m.funcs[1] + self.refuse( + TTIRKind.CALL, + f"{len(m.funcs)} functions in the module (tt.func " + f"{extra.sym_name!r}): calls are not modeled", + m.ops[extra.op], + ) + func = m.funcs[0] + fop = m.ops[func.op] + if not fop.regions or not fop.regions[0]: + self.refuse(TTIRKind.OTHER, f"tt.func {func.sym_name!r} has no body", fop) + for fa in func.args: + self.bind_arg(fa, fop) + body = fop.regions[0] + self.block(body[0]) + if len(body) > 1: + self.refuse( + TTIRKind.CONTROL_FLOW, + f"tt.func {func.sym_name!r} has {len(body)} blocks", + fop, + ) + return AccessGraph( + kernel_name=func.sym_name, + func_args=tuple(self.func_args), + accesses=tuple(self.accesses), + loop=self.loop, + iter_args=tuple(self.iter_args), + ) + + def bind_arg(self, fa, fop: Op) -> None: + ti = parse_type(fa.type) + name = fa.name if fa.name is not None else f"arg{fa.index}" + if self.arg(name) is not None: + self.refuse(TTIRKind.OTHER, f"two parameters named {name!r}", fop) + self.global_pointers((fa.type,), fop) + is_ptr = ti.pointee is not None and not ti.shape + int_bits = ti.int_bits if not is_ptr and not ti.shape else None + self.func_args.append( + FuncArg( + name=name, + is_ptr=is_ptr, + elem_bits=_pointee_bits(ti.pointee) if is_ptr else 0, + int_bits=int_bits or 0, + ) + ) + if is_ptr: + self.env[fa.value] = PtrValue(name, Const(0)) + elif int_bits is not None: + self.env[fa.value] = Param(name) + else: + self.env[fa.value] = DataDep(f"{fa.type} argument") + + def block(self, bidx: int) -> None: + self.depth += 1 + try: + for oi in self.m.blocks[bidx].ops: + self.visit(self.m.ops[oi]) + finally: + self.depth -= 1 + + def visit(self, op: Op) -> None: + if self.depth > _MAX_DEPTH: + self.refuse( + TTIRKind.OTHER, + f"structured regions nested deeper than {_MAX_DEPTH}", + op, + ) + name = op.name + self.global_pointers(op.operand_types + op.result_types, op) + if name in _INERT: + return + handler = _HANDLERS.get(name) + if handler is not None: + handler(self, op) + else: + self.op_other(op) + + # ── value ops ── + def op_program_id(self, op: Op) -> None: + self.bind(op, Pid(op.attrs["axis"])) + + def op_num_programs(self, op: Op) -> None: + self.bind(op, NumPrograms(op.attrs["axis"])) + + def op_make_range(self, op: Op) -> None: + self.bind(op, Arange(f"%v{op.results[0]}", op.attrs["start"], op.attrs["end"])) + + def op_constant(self, op: Op) -> None: + v = op.attrs["value"] + if isinstance(v, bool): + # i1 constants (e.g. the dense mask of an unmasked atomic). + # Const(0/1) in a boolean position is coerced by the evaluator. + self.bind(op, Const(1 if v else 0)) + elif isinstance(v, int): + self.bind(op, Const(v)) + else: + self.bind(op, DataDep("float/array constant")) + + def op_passthrough(self, op: Op) -> None: + # tt.splat: replicate a scalar / seed a pointer tile; + # tt.broadcast: a shape change, value passthrough + self.bind(op, self.val(op.operands[0])) + + def op_expand_dims(self, op: Op) -> None: + v = self.val(op.operands[0]) + axis = op.attrs["axis"] + root = v.offset if isinstance(v, PtrValue) else v + if any( + isinstance(n, Observed) and n.access_index in self.lane_observed + for n in _nodes((root,), self.iter_args) + ): + # Observed has no lane placement: two lanes of one tensor + # observation must not meet in a term (see Observed). + self.bind(op, DataDep("atomic result expanded across lanes")) + return + self.bind(op, _expand_dims(v, axis, lambda t: self.expanded_iter_arg(t, axis))) + + def expanded_iter_arg(self, t: IterArgOffset, axis: int) -> IterArgOffset: + """The loop-carried pointer ``t`` with its lanes re-placed by an + expand_dims at ``axis``: an iter_args entry of its own (IterArgInfo), + one per (pointer, axis), whose delta the loop fills in on closing.""" + key = (t.arg_id, axis) + aid = self.expanded.get(key) + if aid is None: + src = self.iter_args[t.arg_id] + aid = len(self.iter_args) + offset0 = _expand_dims(src.offset0, axis) + self.iter_args.append( + IterArgInfo(aid, src.base_param, offset0, Const(0), _LOOP_SSA) # type: ignore[arg-type] + ) + self.expanded[key] = aid + return IterArgOffset(aid) + + def op_cast(self, op: Op) -> None: + x = self.val(op.operands[0]) + if isinstance(x, (DataDep, PtrValue)): + self.bind(op, x) + return + kind = op.name.split(".", 1)[1] + src = self.int_bits(op.operands[0]) + dst = self.int_bits(op.results[0]) + assert src is not None and dst is not None + term: Term = x # type: ignore[assignment] + if kind == "extsi" and src == 1: + # sign-extending an i1 maps true to -1 + ext = IntCast("extui", 1, dst, term, op.line_no, op.loc) + self.bind(op, Bin("-", Const(0), ext, dst, op.line_no, op.loc)) + return + self.bind(op, IntCast(kind, src, dst, term, op.line_no, op.loc)) + + def op_addptr(self, op: Op) -> None: + base, off = self.val(op.operands[0]), self.val(op.operands[1]) + if not isinstance(base, PtrValue): + self.refuse( + _address_kind(base), + f"addptr base is not a pointer{_why(base)}", + op, + ) + if isinstance(off, DataDep): + # A value in an address chain that cannot be modeled: a free + # address makes the query meaningless. Only offsets truly + # derived from MEMORY CONTENTS classify as indirection. + kind = _address_kind(off) + what = ( + "carried by the loop" + if kind is TTIRKind.LOOP_VARIANT_ADVANCE + else "data-dependent" + ) + self.refuse(kind, f"addptr offset: {what} ({off.why})", op) + off_t = self.as_term(off, "addptr offset", op) + self.bind( + op, + PtrValue( + base.base_param, # type: ignore[union-attr] + Bin("+", base.offset, off_t, None, op.line_no, op.loc), # type: ignore[union-attr] + ), + ) + + def op_bin(self, op: Op) -> None: + a, b = self.val(op.operands[0]), self.val(op.operands[1]) + bits = self.int_bits(op.results[0]) + if isinstance(a, DataDep) or isinstance(b, DataDep): + self.bind(op, _data_dep((a, b), "arith over loaded data")) + elif bits == 1: + self.bind(op, DataDep(f"{op.name} on i1")) + else: + self.bind( + op, + Bin( + _BIN_OPS[op.name], + self.as_term(a, "arith", op), + self.as_term(b, "arith", op), + bits, + op.line_no, + op.loc, + ), + ) + + def op_cmpi(self, op: Op) -> None: + a, b = self.val(op.operands[0]), self.val(op.operands[1]) + pred = op.attrs["predicate"] + bits = self.int_bits(op.operands[0]) + if isinstance(a, DataDep) or isinstance(b, DataDep): + self.bind(op, _data_dep((a, b), "cmpi over loaded data")) + elif bits == 1 and pred in _SIGNED_PREDICATES: + # a signed i1 reads true as -1; the boolean model reads it as 1 + self.bind(op, DataDep("signed comparison of i1 values")) + else: + self.bind( + op, + Cmp( + pred, + self.as_term(a, "cmpi", op), + self.as_term(b, "cmpi", op), + bits, + op.line_no, + op.loc, + ), + ) + + def op_boolbin(self, op: Op) -> None: + if self.int_bits(op.results[0]) != 1: + # Wide-int andi/ori is BITWISE arithmetic, not boolean logic; + # modeling it as And/Or would silently corrupt address math. + # Degrade to DataDep so an address use fails closed. + self.bind( + op, + DataDep(f"bitwise {op.name} on non-i1 type {op.result_types[0]}"), + ) + return + a, b = self.val(op.operands[0]), self.val(op.operands[1]) + is_and = op.name == "arith.andi" + if isinstance(a, DataDep) or isinstance(b, DataDep): + why = _data_dep((a, b), "bool op over loaded data").why + self.bind(op, DataDep(why)) + return + self.bind( + op, + BoolBin( + "and" if is_and else "or", + self.as_term(a, "bool", op), + self.as_term(b, "bool", op), + ), + ) + + def op_select(self, op: Op) -> None: + c, t, f = (self.val(v) for v in op.operands) + self.bind(op, self.merge(c, t, f, "select")) + + def merge(self, c: object, t: object, f: object, what: str) -> object: + """``c ? t : f`` for an arith.select or an scf.if result: pointers + of one base select their offsets; anything unmodelable is DataDep.""" + if ( + isinstance(c, (DataDep, PtrValue)) + or isinstance(t, DataDep) + or (isinstance(f, DataDep)) + ): + return _data_dep((c, t, f), "select over loaded data") + if isinstance(t, PtrValue) and isinstance(f, PtrValue): + if t.base_param != f.base_param: + return DataDep(f"{what} of pointers with different bases") + return PtrValue(t.base_param, Select(c, t.offset, f.offset)) # type: ignore[arg-type] + if isinstance(t, PtrValue) or isinstance(f, PtrValue): + return DataDep(f"{what} of a pointer and an integer") + return Select(c, t, f) # type: ignore[arg-type] + + def op_bitcast(self, op: Op) -> None: + # A pointer cast that keeps the element width keeps element offsets + # (atomic_max on floats casts f32 -> i32 pointers); any other + # bitcast reinterprets data. + src = parse_type(op.operand_types[0]) + dst = parse_type(op.result_types[0]) + x = self.val(op.operands[0]) + if src.pointee is not None and dst.pointee is not None: + if _pointee_bits(src.pointee) == _pointee_bits(dst.pointee): + self.bind(op, x) + else: + self.bind(op, DataDep("pointer bitcast changes the element width")) + return + self.bind(op, _data_dep((x,), f"unmodeled op {op.name} at line {op.line_no}")) + + # ── accesses ── + def access( + self, + op: Op, + kind: str, + mask_v: int | None, + ) -> None: + ptr_v = op.operands[0] + ptr = self.val(ptr_v) + if not isinstance(ptr, PtrValue): + self.refuse( + _address_kind(ptr), + f"{kind} of a non-pointer value{_why(ptr)}", + op, + ) + pointee = parse_type(parse_type(self.m.values[ptr_v].type).elem).pointee + elem_bits = _pointee_bits(pointee) + base = self.arg(ptr.base_param) # type: ignore[union-attr] + if base is None or base.elem_bits != elem_bits: + self.refuse( + TTIRKind.OTHER, + f"{kind} element width {elem_bits} differs from its base " + f"{ptr.base_param!r}", # type: ignore[union-attr] + op, + ) + mask: Term | None = None + mask_dropped = False + if mask_v is not None: + mv = self.val(mask_v) + if isinstance(mv, DataDep): + # Mask derived from loaded data: over-approximate it as free + # (any lane may be active) instead of failing the kernel. + # See AccessEvent.mask_dropped for the soundness discipline. + mask_dropped = True + elif isinstance(mv, PtrValue): + self.refuse(TTIRKind.OTHER, "pointer as mask", op) + else: + mask = mv # type: ignore[assignment] + guarded, path, in_loop = self.branch_state() + assert op.line_no is not None + self.accesses.append( + AccessEvent( + kind=kind, + base_param=ptr.base_param, # type: ignore[union-attr] + offset=ptr.offset, # type: ignore[union-attr] + mask=mask, + elem_bits=elem_bits, + loc=op.loc, + line_no=op.line_no, + guarded=guarded, + path=path, + in_loop=in_loop, + mask_dropped=mask_dropped, + ) + ) + + def mask_operand(self, op: Op, index: int) -> int | None: + if len(op.operands) <= index: + return None + v = op.operands[index] + if self.int_bits(v) != 1: + self.refuse(TTIRKind.OTHER, f"{op.name} operand {index} is not a mask", op) + return v + + def observed_binding(self, op: Op) -> object: + """The value of the just-recorded atomic's result: Observed for an + integer-typed result, DataDep otherwise (floats stay outside the Int + model).""" + if op.results and self.int_bits(op.results[0]) is not None: + index = len(self.accesses) - 1 + if parse_type(self.m.values[op.results[0]].type).shape: + self.lane_observed.add(index) + return Observed(index) + return DataDep("atomic result") + + def op_load(self, op: Op) -> None: + # ODS operands (ptr, mask?, other?). Both are optional, so the + # generic form spells which are present in operandSegmentSizes (the + # custom form prints a lone second operand as the mask). + segments = op.attrs.get("operandSegmentSizes") + if segments is None: + mask_v = self.mask_operand(op, 1) + elif ( + len(segments) != 3 or segments[0] != 1 or sum(segments) != len(op.operands) + ): + self.refuse(TTIRKind.OTHER, f"tt.load operand segments {segments}", op) + else: + mask_v = self.mask_operand(op, 1) if segments[1] else None + self.access(op, "load", mask_v) + self.bind(op, DataDep("loaded value")) + + def op_store(self, op: Op) -> None: + # ODS operands (ptr, value, mask?) + self.access(op, "store", self.mask_operand(op, 2)) + + def op_atomic_rmw(self, op: Op) -> None: + # ODS operands (ptr, val, mask?) + self.access( + op, + "atomic_rmw", + self.mask_operand(op, 2), + ) + self.bind(op, self.observed_binding(op)) + + def op_atomic_cas(self, op: Op) -> None: + # ODS operands (ptr, cmp, val); CAS has no mask: unconditional footprint + self.access( + op, + "atomic_cas", + None, + ) + self.bind(op, self.observed_binding(op)) + + # ── structured control flow ── + def op_for(self, op: Op) -> None: + # The single-path model has one induction variable: a second loop + # (sequential or nested) cannot be represented, and a loop under an + # scf.if runs a branch-dependent iteration count — a control-flow + # limitation, not one more induction variable. + if self.loops_opened or self.frames: + self.refuse( + TTIRKind.CONTROL_FLOW + if any(isinstance(f, _IfFrame) for f in self.frames) + else TTIRKind.NESTED_LOOP, + "multiple/nested loops", + op, + ) + self.loops_opened += 1 + bounds: dict[str, Term] = {} + for label, v in zip(("lower", "upper", "step"), op.operands[:3]): + bv = self.val(v) + if isinstance(bv, DataDep): + # The CSR shape: for k in range(loaded_start, loaded_end). + self.refuse( + TTIRKind.DATA_DEPENDENT_BOUND + if _from_memory(bv) + else TTIRKind.OTHER, + f"loop {label} bound: data-dependent ({bv.why})", + op, + ) + if any(isinstance(n, Observed) for n in _nodes((bv,), None)): + # A trip count driven by an atomic observation is a dynamic + # work-fetch loop. + self.refuse( + TTIRKind.DATA_DEPENDENT_BOUND, + f"loop {label} bound depends on an atomic observation", + op, + ) + bounds[label] = self.as_term(bv, f"loop {label}", op) + body = self.m.blocks[op.regions[0][0]] + iv, carried = body.args[0], body.args[1:] + self.env[iv] = LoopVar(_LOOP_SSA) + ptr_args: list[tuple[int, int]] = [] # (iter_arg position, arg_id) + for k, (arg_v, init_v) in enumerate(zip(carried, op.operands[3:])): + init = self.val(init_v) + if isinstance(init, PtrValue): + aid = len(self.iter_args) + self.iter_args.append( + IterArgInfo(aid, init.base_param, init.offset, Const(0), _LOOP_SSA) + ) + self.env[arg_v] = PtrValue(init.base_param, IterArgOffset(aid)) + ptr_args.append((k, aid)) + elif isinstance(init, DataDep): + self.env[arg_v] = DataDep(init.why) + else: + self.env[arg_v] = DataDep(_LOOP_CARRIED) + self.frames.append(_ForFrame()) + self.block(op.regions[0][0]) + self.frames.pop() + yield_op = self.m.ops[body.ops[-1]] + for k, aid in ptr_args: + delta = self.loop_delta(self.val(yield_op.operands[k]), aid, yield_op) + self.iter_args[aid] = replace(self.iter_args[aid], delta=delta) + # Expanded tiles advance by their source's delta, expanded alike (a + # source precedes the tiles expanded from it). + for (src, axis), aid in self.expanded.items(): + source = self.iter_args[src].delta + if any( + isinstance(n, Observed) and n.access_index in self.lane_observed + for n in _nodes((source,), None) + ): + self.refuse( + TTIRKind.OTHER, + f"loop-carried pointer {src} is expanded but advances by a " + "per-lane atomic result", + yield_op, + ) + expanded = _expand_dims(source, axis) + self.iter_args[aid] = replace(self.iter_args[aid], delta=expanded) # type: ignore[arg-type] + self.loop = LoopInfo( + loop_ssa=_LOOP_SSA, + lower=bounds["lower"], + upper=bounds["upper"], + step=bounds["step"], + bits=self.int_bits(iv), + unsigned=bool(op.attrs["unsignedCmp"]), + line_no=op.line_no, + loc=op.loc, + ) + self.bind(op, DataDep("loop result")) + + def loop_delta(self, y: object, aid: int, op: Op) -> Term: + """The loop-invariant advance of loop-carried pointer ``aid`` from + its yielded value, which must be ``IterArgOffset(aid) + delta`` on the + same base: the model reads iteration k's offset as + ``offset0 + k * delta``.""" + info = self.iter_args[aid] + if not isinstance(y, PtrValue) or y.base_param != info.base_param: + self.refuse( + TTIRKind.LOOP_VARIANT_ADVANCE, + f"loop-carried pointer {aid} (base {info.base_param!r}) is not " + "advanced from itself", + op, + ) + delta = _loop_delta(y.offset, aid) # type: ignore[union-attr] + if delta is None: + self.refuse( + TTIRKind.LOOP_VARIANT_ADVANCE, + f"loop-carried pointer {aid} is not advanced by addptr from " + "its own previous value", + op, + ) + for n in _nodes((delta,), None): + if ( + isinstance(n, (LoopVar, IterArgOffset)) + or isinstance(n, Observed) + and self.accesses[n.access_index].in_loop + ): + self.refuse( + TTIRKind.LOOP_VARIANT_ADVANCE, + f"loop-carried pointer {aid} advances by a loop-variant amount", + op, + ) + return delta # type: ignore[return-value] + + def op_if(self, op: Op) -> None: + cv = self.val(op.operands[0]) + # A pointer can't be a condition; loaded data (DataDep) can't be + # modeled -> the region stays pessimistically ``guarded``. + cond = None if isinstance(cv, (DataDep, PtrValue)) else cv + frame = _IfFrame(cond) # type: ignore[arg-type] + yields: list[list[object]] = [] + for ri, region in enumerate(op.regions): + if not region: + continue + if len(region) != 1: + self.op_control_flow(op) + frame.branch = "then" if ri == 0 else "else" + self.frames.append(frame) + self.block(region[0]) + self.frames.pop() + y = self.m.ops[self.m.blocks[region[0]].ops[-1]] + yields.append([self.val(v) for v in y.operands]) + if op.results and len(yields) != 2: + self.refuse(TTIRKind.OTHER, "scf.if with results but no else region", op) + for k, r in enumerate(op.results): + if cond is None: + self.env[r] = _data_dep((cv,), "select over loaded data") + else: + self.env[r] = self.merge(cond, yields[0][k], yields[1][k], "scf.if") + + # ── everything else ── + def op_control_flow(self, op: Op) -> None: + self.refuse(TTIRKind.CONTROL_FLOW, f"control flow {op.name} is unsupported", op) + + def op_call(self, op: Op) -> None: + self.refuse(TTIRKind.CALL, "calls are not modeled", op) + + def opaque_hazard(self, op: Op) -> str | None: + """Why an opaque op (inline asm, an extern function) may touch + memory the graph cannot see, or None: it is impure, or an operand + hands it an address, as a pointer or as an integer made from one.""" + if op.attrs.get("pure") is not True: + return "side effects" + if any(parse_type(t).pointee is not None for t in op.operand_types): + return "a pointer operand" + if self.from_pointer(op.operands): + return "an operand computed from a pointer (tt.ptr_to_int)" + return None + + def from_pointer(self, values: Iterable[int]) -> bool: + """True when a ``tt.ptr_to_int`` result flows into one of ``values`` + (through any op, region results and loop-carried values included; + loaded data carries no address).""" + m = self.m + stack = list(values) + seen: set[int] = set() + while stack: + v = stack.pop() + if v in seen: + continue + seen.add(v) + value = m.values[v] + if value.op is not None: + src = m.ops[value.op] + if src.name == "tt.ptr_to_int": + return True + if src.name in _ACCESS_OPS: + continue + else: + assert value.block is not None + src = m.ops[m.blocks[value.block].op] + if src.name == "tt.func": + continue + stack += src.operands + for region in src.regions: + for b in region: + ops = m.blocks[b].ops + if ops: + stack += m.ops[ops[-1]].operands + return False + + def op_inline_asm(self, op: Op) -> None: + hazard = self.opaque_hazard(op) + if hazard is not None: + self.refuse(TTIRKind.INLINE_ASM, f"inline asm with {hazard}", op) + self.bind(op, _data_dep(map(self.val, op.operands), "unmodeled inline asm")) + + def op_extern(self, op: Op) -> None: + hazard = self.opaque_hazard(op) + if hazard is not None: + self.refuse( + TTIRKind.OUT_OF_VOCABULARY, + f"{op.name} {op.attrs.get('symbol')!r} with {hazard}", + op, + ) + self.op_other(op) + + def op_other(self, op: Op) -> None: + name = op.name + if name.startswith(_MEMORY_PREFIXES): + self.refuse(TTIRKind.OUT_OF_VOCABULARY, f"unsupported memory op {name}", op) + if name.split(".", 1)[0] not in _DIALECTS: + self.refuse(TTIRKind.OUT_OF_VOCABULARY, f"op {name} is not TTIR", op) + if name.startswith(("scf.", "cf.")): + self.op_control_flow(op) + if op.regions: + self.pure_regions(op) + elif not op.results: + # a result-free op the reader does not know may write memory + self.refuse(TTIRKind.OUT_OF_VOCABULARY, f"unmodeled op {name}", op) + # Plain data (floats, dots, reductions, ...): memory-derived exactly + # when an operand is. + self.bind( + op, + _data_dep( + map(self.val, op.operands), f"unmodeled op {name} at line {op.line_no}" + ), + ) + + def pure_regions(self, op: Op) -> None: + """The regions of an op the reader does not follow (tt.reduce / + tt.scan combine bodies) compute values only: refuse any memory, + call or opaque op inside.""" + stack = [b for region in op.regions for b in region] + while stack: + for oi in self.m.blocks[stack.pop()].ops: + inner = self.m.ops[oi] + name = inner.name + if ( + name in _ACCESS_OPS + or name == "tt.call" + or name.startswith(_MEMORY_PREFIXES) + or (name.split(".", 1)[0] not in _DIALECTS and name not in _INERT) + or (name in _OPAQUE and self.opaque_hazard(inner) is not None) + or ( + not inner.results + and not inner.regions + and name not in _INERT + and name != "scf.condition" + ) + ): + self.refuse( + TTIRKind.OUT_OF_VOCABULARY, + f"{name} inside the region of {op.name}", + inner, + ) + stack.extend(b for region in inner.regions for b in region) + + +def _data_dep(operands: Iterable[object], why: str) -> DataDep: + """The DataDep of an op the reader cannot model over ``operands``: a + memory-derived reason when any operand derives from memory (``why`` + itself when it is one), else the first unmodelable operand's own reason + (a modeling gap stays a modeling gap), else ``why``.""" + first: DataDep | None = None + for x in operands: + if _from_memory(x): + return DataDep( + why + if why.startswith(_MEMORY_WHYS) + else f"arith over loaded data ({why})" + ) + if first is None and isinstance(x, DataDep): + first = x + return DataDep(first.why) if first is not None else DataDep(why) + + +def _why(v: object) -> str: + return f" ({v.why})" if isinstance(v, DataDep) else "" + + +def _address_kind(v: object) -> TTIRKind: + """The refusal kind of an unmodelable value in an address: loaded data + is indirection, a loop-carried integer a loop-variant address.""" + if _from_memory(v): + return TTIRKind.INDIRECT_ADDRESS + if isinstance(v, DataDep) and v.why == _LOOP_CARRIED: + return TTIRKind.LOOP_VARIANT_ADVANCE + return TTIRKind.OTHER + + +def _with_children(t: object, kids: tuple) -> object: + if isinstance(t, (Bin, Cmp, BoolBin)): + return replace(t, a=kids[0], b=kids[1]) + if isinstance(t, Select): + return replace(t, cond=kids[0], t=kids[1], f=kids[2]) + if isinstance(t, Not): + return replace(t, a=kids[0]) + if isinstance(t, IntCast): + return replace(t, x=kids[0]) + raise TypeError(type(t).__name__) + + +def _expand_dims( + v: object, + axis: int, + iter_arg: Callable[[IterArgOffset], object] | None = None, +) -> object: + """``tt.expand_dims`` inserts a size-1 dimension at ``axis``: every + Arange lane index moves to its position in the new shape (a 1D range + sits at 0; a dimension at or after ``axis`` shifts up by one). Pointer + tiles follow like integer tiles, so an address and a mask expanded the + same way keep sharing their lane variables; ``iter_arg`` maps a + loop-carried pointer's IterArgOffset to its expanded tile's. Iterative, + sharing-preserving (terms can be deeper than the recursion limit).""" + if isinstance(v, PtrValue): + return PtrValue(v.base_param, _expand_dims(v.offset, axis, iter_arg)) # type: ignore[arg-type] + memo: dict[int, object] = {} + stack: list[tuple[object, bool]] = [(v, False)] + while stack: + t, ready = stack.pop() + if id(t) in memo: + continue + kids = _children(t) + if isinstance(t, Arange): + pos = 0 if t.dim < 0 else t.dim + memo[id(t)] = replace(t, dim=pos + 1 if pos >= axis else pos) + elif isinstance(t, IterArgOffset) and iter_arg is not None: + memo[id(t)] = iter_arg(t) + elif not kids: + memo[id(t)] = t + elif not ready: + stack.append((t, True)) + stack.extend((k, False) for k in kids if id(k) not in memo) + else: + new = tuple(memo[id(k)] for k in kids) + memo[id(t)] = ( + t if all(n is k for n, k in zip(new, kids)) else _with_children(t, new) + ) + return memo[id(v)] + + +def _loop_delta(offset: Term, arg_id: int) -> Term | None: + """From a yielded pointer offset ``IterArgOffset(arg_id) + d1 + ... + dn`` + (a chain of addptr sums, any association), pull out the delta + ``d1 + ... + dn`` (``Const(0)`` for none); None for any other shape.""" + found = 0 + rest: list[Term] = [] + stack: list[Term] = [offset] + while stack: + t = stack.pop() + if isinstance(t, Bin) and t.op == "+" and t.bits is None: + stack += [t.b, t.a] + elif isinstance(t, IterArgOffset) and t.arg_id == arg_id: + found += 1 + else: + rest.append(t) + if found != 1: + return None + if not rest: + return Const(0) + delta = rest[0] + for t in rest[1:]: + delta = Bin("+", delta, t) + return delta + + +_HANDLERS = { + "tt.get_program_id": _Builder.op_program_id, + "tt.get_num_programs": _Builder.op_num_programs, + "tt.make_range": _Builder.op_make_range, + "arith.constant": _Builder.op_constant, + "tt.splat": _Builder.op_passthrough, + "tt.broadcast": _Builder.op_passthrough, + "tt.expand_dims": _Builder.op_expand_dims, + "arith.extsi": _Builder.op_cast, + "arith.extui": _Builder.op_cast, + "arith.trunci": _Builder.op_cast, + "tt.addptr": _Builder.op_addptr, + **{name: _Builder.op_bin for name in _BIN_OPS}, + "arith.cmpi": _Builder.op_cmpi, + "arith.andi": _Builder.op_boolbin, + "arith.ori": _Builder.op_boolbin, + "arith.select": _Builder.op_select, + "tt.bitcast": _Builder.op_bitcast, + "tt.load": _Builder.op_load, + "tt.store": _Builder.op_store, + "tt.atomic_rmw": _Builder.op_atomic_rmw, + "tt.atomic_cas": _Builder.op_atomic_cas, + "scf.for": _Builder.op_for, + "scf.if": _Builder.op_if, + "tt.call": _Builder.op_call, + "tt.elementwise_inline_asm": _Builder.op_inline_asm, + "tt.extern_elementwise": _Builder.op_extern, +}