diff --git a/tests/unit/ir/reader_kernels.py b/tests/unit/ir/reader_kernels.py index 5aaa8e97..203310a4 100644 --- a/tests/unit/ir/reader_kernels.py +++ b/tests/unit/ir/reader_kernels.py @@ -361,6 +361,70 @@ def poll_while(flag_ptr, out_ptr, n): tl.store(out_ptr + pid, 1) +@triton.jit +def mp_loops(x_ptr, n, m): + # nested loops (the inner pointer advances by the outer induction + # variable), then a sequential loop under a branch + for i in range(0, n): + q = x_ptr + i + for j in range(0, m, 2): + tl.store(q + j, 1) + q += i + if tl.program_id(0) < n: + p = x_ptr + for k in range(0, n): + tl.store(p, 2) + p += m + + +@triton.jit +def mp_loaded(idx_ptr, x_ptr, f_ptr, n, B: tl.constexpr): + # integer loads in an address, a mask and a branch; a float load and a + # dropped mask; a loaded tile expanded; masks with a modelable conjunct + # (kept, also expanded) or disjunct (dropped); a float atomic operand + offs = tl.arange(0, B) + j = tl.load(idx_ptr + offs, mask=offs < n, other=0) + tl.store(x_ptr + j, 1.0) + v = tl.load(f_ptr + offs) + tl.store(x_ptr + offs, v, mask=j > 0) + t = tl.load(idx_ptr + offs, mask=offs < n, other=0) + tl.store(x_ptr + t[:, None] + offs[None, :], 3.0) + w = tl.load(idx_ptr + offs, mask=v > 0) + m = (offs < n) & (w > 0) + tl.store(x_ptr + offs, 4.0, mask=m) + if tl.load(idx_ptr) > 0: + tl.store(x_ptr, 5.0) + tl.store(x_ptr + offs, 6.0, mask=(offs < n) | (w > 0)) + tl.store(x_ptr + offs[:, None] * B + offs[None, :], 7.0, mask=m[:, None]) + tl.atomic_add(f_ptr + offs, v) + + +@triton.jit +def mp_early_return(x_ptr, f_ptr, n): + # an early return under a float condition: the rest is guarded + pid = tl.program_id(0) + if tl.load(f_ptr) > 0.0: + return + tl.store(x_ptr + pid, 1.0) + + +@triton.jit +def mp_loaded_advance(idx_ptr, x_ptr, n): + # an advance by a value loaded in the loop: variant + p = x_ptr + for i in range(0, n): + tl.store(p, 1) + tl.store(idx_ptr, i) + p += tl.load(idx_ptr) + + +@triton.jit +def mp_for_in_while(flag_ptr, x_ptr, n): + while tl.load(flag_ptr) == 0: + for i in range(0, n): + tl.store(x_ptr + i, 1) + + _F32 = "*fp32" # corpus name (reader_ttir/.ttir, the kernel of that name) -> # (signature, constexprs), compiled for sm80 with the default options @@ -425,4 +489,12 @@ def poll_while(flag_ptr, out_ptr, n): "fences": ({"x_ptr": "*i32", "n": "i32"}, {}), "pid_in_combine": ({"x_ptr": "*i32", "o_ptr": "*i32", "B": "constexpr"}, {"B": 8}), "poll_while": ({"flag_ptr": "*i32", "out_ptr": "*i32", "n": "i32"}, {}), + "mp_loops": ({"x_ptr": "*i32", "n": "i32", "m": "i32"}, {}), + "mp_loaded": ( + {"idx_ptr": "*i32", "x_ptr": _F32, "f_ptr": _F32, "n": "i32", "B": "constexpr"}, + {"B": 8}, + ), + "mp_early_return": ({"x_ptr": _F32, "f_ptr": _F32, "n": "i32"}, {}), + "mp_loaded_advance": ({"idx_ptr": "*i32", "x_ptr": "*i32", "n": "i32"}, {}), + "mp_for_in_while": ({"flag_ptr": "*i32", "x_ptr": "*i32", "n": "i32"}, {}), } diff --git a/tests/unit/ir/test_lowering.py b/tests/unit/ir/test_lowering.py index 8b2595cd..cbc0649c 100644 --- a/tests/unit/ir/test_lowering.py +++ b/tests/unit/ir/test_lowering.py @@ -27,6 +27,7 @@ IntCast, IterArgInfo, IterArgOffset, + Loaded, LoopInfo, LoopVar, Not, @@ -75,6 +76,9 @@ def iteration(self, loop_ssa): def observed(self, t): return self._var(f"observed_{t.access_index}") + def loaded(self, t): + return self._var(f"loaded_{t.access_index}") + def data_dep(self, t): raise _Refusal(t.why) @@ -201,7 +205,9 @@ def kids(t) -> list[int]: assert kids(LoopVar("%loop")) == [id(LOOP.lower), id(LOOP.step)] for leaf in (Const(1), Pid(0), NumPrograms(1), Arange("%r", 0, 4), X, Observed(0)): assert kids(leaf) == [] - assert kids(DataDep()) == [] + # a loaded value's terms are the client's; keep is no value dependency + assert kids(DataDep()) == kids(DataDep("x", keep=less)) == [] + assert kids(Loaded(0, X, less, Y)) == [] # ─────────────────────────── operator semantics ─────────────────────────── @@ -325,11 +331,30 @@ def test_loop_terms_read_the_clients_iteration(): assert _proved(lowerer.lower(IterArgOffset(0)) == _v("o0") + k * _v("d0")) assert _proved(lowerer.lower(IterArgOffset(1)) == _v("o1") + other * _v("d1")) assert "hi" not in leaves.made # the upper bound is not the variable's value + # with several loops, each variable reads its own loop by name + inner = LoopInfo( + "%loop1", lower=Param("lo1"), upper=Param("hi1"), step=Param("st1") + ) + graph = AccessGraph("k", (), (), None, loops=(LOOP, inner)) + lowerer = Lowerer(graph, _Leaves()) + assert _proved( + lowerer.lower(LoopVar("%loop1")) == _v("lo1") + _v("k%loop1") * _v("st1") + ) + assert _proved(lowerer.lower(LoopVar("%loop")) == _v("lo") + k * _v("st")) + with pytest.raises(ValueError, match="without a loop named '%loop2'"): + lowerer.lower(LoopVar("%loop2")) # ─────────────────────────── leaves and bugs ─────────────────────────── +def test_a_loaded_value_is_the_clients_leaf(): + leaves = _Leaves() + term = Bin("+", X, Loaded(3, Param("o"), None, None)) + assert _proved(Lowerer(_graph(), leaves).lower(term) == _v("x") + _v("loaded_3")) + assert "o" not in leaves.made + + def test_a_leaf_refusal_propagates_unchanged(): with pytest.raises(_Refusal, match="loaded value"): Lowerer(_graph(), _Leaves()).lower(BoolBin("and", X, DataDep("loaded value"))) diff --git a/tests/unit/ir/test_mlir_walk.py b/tests/unit/ir/test_mlir_walk.py index 681268a2..98677cda 100644 --- a/tests/unit/ir/test_mlir_walk.py +++ b/tests/unit/ir/test_mlir_walk.py @@ -226,6 +226,21 @@ def test_cf_blocks_align(): 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")) + # a branch's successors are its target blocks, read from the text + by_label = {m.blocks[b].label: b for b in func.regions[0][1:]} + succ = [op.attrs["successors"] for op in _ops(m, "cf.br") + _ops(m, "cf.cond_br")] + assert succ == [ + (by_label["^bb5"],), + (by_label["^bb1"], by_label["^bb2"]), + (by_label["^bb3"], by_label["^bb4"]), + (by_label["^bb1"], by_label["^bb5"]), + ] + same = W._walk(_text("crafted_same_dest.ttir")) + first = _ops(same, "cf.cond_br")[0] + assert first.attrs["successors"] == (first.block + 1,) * 2 # ^bb1(%a), ^bb1(%b) + bad = _text("crafted_same_dest.ttir").replace("^bb2(%y, %x : i32, i32)", "^bb2") + with pytest.raises(W.MisalignedModule, match="0 operands for its arguments"): + W._walk(_text("crafted_same_dest.ttir"), scan_text=bad) def test_parser_invented_locs_are_absent(): diff --git a/tests/unit/ir/test_ttir_reader.py b/tests/unit/ir/test_ttir_reader.py index fce2bff2..54b053f6 100644 --- a/tests/unit/ir/test_ttir_reader.py +++ b/tests/unit/ir/test_ttir_reader.py @@ -9,6 +9,7 @@ import itertools import pickle +from dataclasses import replace from pathlib import Path import pytest @@ -21,10 +22,14 @@ Arange, AtomicInfo, Bin, + BoolBin, Cmp, Const, IntCast, IterArgOffset, + Loaded, + LoopVar, + Not, Observed, Param, Pid, @@ -194,6 +199,11 @@ def test_refusals_carry_the_op_line_and_source_loc(): "reader_ttir/pure_asm_int_addr.ttir": "inline-asm", "reader_ttir/observed_lanes.ttir": "indirect-address", "reader_ttir/poll_while.ttir": "control-flow", + "reader_ttir/mp_loops.ttir": "nested-loop", + "reader_ttir/mp_loaded.ttir": "indirect-address", + "reader_ttir/mp_early_return.ttir": "control-flow", + "reader_ttir/mp_loaded_advance.ttir": "indirect-address", + "reader_ttir/mp_for_in_while.ttir": "control-flow", } @@ -1064,3 +1074,276 @@ def test_while_regions_are_read_only_on_request(): assert parse_ttir(_text(name), while_regions=True) == _graph(name) outcome = ParseCache().get(text, while_regions=True) assert len(outcome.graph.whiles) == 4 + # with multipath too, a poll's load stays an observation + for name in ("reader_ttir/poll_while.ttir", "ttir/spike_spin_while.ttir"): + both = parse_ttir(_text(name), while_regions=True, multipath=True) + assert pickle.dumps(both) == pickle.dumps( + parse_ttir(_text(name), while_regions=True) + ) + + +# ─────────────────────────── multipath ─────────────────────────── + + +def _mp(name: str) -> AccessGraph: + return parse_ttir(_text(name), multipath=True) + + +@pytest.mark.parametrize("name", FILES) +def test_multipath_reads_a_default_text_alike_but_for_loaded_values(name): + try: + g = _graph(name) + except UnsupportedTTIR: + return + # the default's loop facts: its one loop + assert g.loops == ((g.loop,) if g.loop else ()) + for x in g.accesses + g.fences: + assert x.loops == (("%loop",) if x.in_loop else ()) + m = _mp(name) + roots = [f.path for f in m.fences] + [ + t for a in m.accesses for t in (a.offset, a.mask, a.path, a.atomic_val) + ] + keep = any(a.mask_dropped and a.mask is not None for a in m.accesses) + if not keep and not any(R.loaded_indices(t, m) for t in roots): + # pickle: byte for byte, and iterative (kernel_deep_chain) + assert pickle.dumps(m) == pickle.dumps(g) + + +def test_several_loops_each_with_their_own_variable(): + assert _refusal("reader_ttir/mp_loops.ttir").kind is TTIRKind.NESTED_LOOP + g = _mp("reader_ttir/mp_loops.ttir") + outer, inner, under_if = g.loops + # a graph built with one of loop and loops gets the other + assert AccessGraph("k", (), (), outer).loops == (outer,) + assert AccessGraph("k", (), (), None, loops=(outer,)).loop is outer + assert g.loop is None and [lp.loop_ssa for lp in g.loops] == [ + "%loop", + "%loop1", + "%loop2", + ] + assert (inner.upper, inner.step, under_if.upper) == ( + Param("m"), + Const(2), + Param("n"), + ) + nested, seq = g.accesses + assert (nested.loops, seq.loops, seq.in_loop) == ( + ("%loop", "%loop1"), + ("%loop2",), + True, + ) + assert nested.offset == Bin("+", IterArgOffset(0), LoopVar("%loop1")) + assert seq.path == Cmp("slt", Pid(0), Param("n"), 32) + # the inner pointer advances by the outer variable: invariant inside + q, p = g.iter_args + assert (q.loop_ssa, q.delta, p.loop_ssa) == ("%loop1", LoopVar("%loop"), "%loop2") + # one increment obligation per enclosing loop, outermost first + latches = [ob.term for ob in width_obligations(g, nested) if ob.role == "loop"] + assert [t.b for t in latches] == [outer.step, inner.step] + assert [ob.term.b for ob in width_obligations(g, seq)] == [under_if.step] + # an advance by the loop's own variable, or by a value loaded in the + # loop, is still variant; a loop in an scf.while still refuses + for name in ("p1_variant_delta", "mp_loaded_advance"): + with pytest.raises(UnsupportedTTIR, match="loop-variant amount") as e: + _mp(f"reader_ttir/{name}.ttir") + assert e.value.kind is TTIRKind.LOOP_VARIANT_ADVANCE + text = _text("reader_ttir/mp_for_in_while.ttir") + with pytest.raises(UnsupportedTTIR) as e: + parse_ttir(text, while_regions=True, multipath=True) + assert e.value.kind is TTIRKind.NESTED_LOOP + + +def test_integer_loads_read_as_loaded_values(): + assert _refusal("reader_ttir/mp_loaded.ttir").kind is TTIRKind.INDIRECT_ADDRESS + g = _mp("reader_ttir/mp_loaded.ttir") + j, gather, _, by_j, _, tile, w, kept, _, under, either, kept2, add = g.accesses + lane = Arange(j.offset.b.ssa, 0, 8) + loaded = Loaded(0, j.offset, Cmp("slt", lane, Param("n"), 32), Const(0)) + # in an address, a mask and a branch condition + assert gather.offset == Bin("+", Const(0), loaded) + assert by_j.mask == Cmp("sgt", loaded, Const(0), 32) and not by_j.mask_dropped + assert under.path == Cmp("sgt", Loaded(8, Const(0), None, None), Const(0), 32) + assert R.loaded_indices(under.path, g) == {8} + # an expanded tile is re-placed along the new dimension, its mask too + (row,) = [n for n in R._nodes((tile.offset,), None) if isinstance(n, Loaded)] + row_lane = replace(lane, dim=0) + assert row == Loaded( + 4, + Bin("+", Const(0), row_lane), + Cmp("slt", row_lane, Param("n"), 32), + Const(0), + ) + # a load whose mask is dropped stays DataDep; the modelable conjunct of + # a dropped mask is kept (re-placed by an expand_dims), a disjunct not + assert (w.mask_dropped, kept.mask_dropped, kept2.mask_dropped) == (True,) * 3 + assert kept.mask == Cmp("slt", lane, Param("n"), 32) + assert kept2.mask == Cmp("slt", row_lane, Param("n"), 32) + assert (either.mask, either.mask_dropped) == (None, True) + # a float load never reaches an integer term (an atomic operand here) + assert add.atomic is not None and add.atomic_val is None + # loaded_indices follows a loop variable into its loop's bounds + body = """%c1 = arith.constant 1 : i32 + %lo = tt.load %q : !tt.ptr + scf.for %k = %lo to %n step %c1 : i32 { + %a = tt.addptr %q, %k : !tt.ptr, i32 + tt.store %a, %c1 : !tt.ptr + }""" + csr = parse_ttir(_module(body, "%q: !tt.ptr, %n: i32"), multipath=True) + assert csr.accesses[1].offset == Bin("+", Const(0), LoopVar("%loop")) + assert R.loaded_indices(csr.accesses[1].offset, csr) == {0} + + +def test_cf_blocks_read_with_block_predicates_and_merged_arguments(): + g = _mp("ttir/golden_nested_guard_merge_sm80.ttir") + pid, n, T = Pid(0), Param("n"), Param("T") + go_on = Not(Cmp("sge", pid, T, 32)) + first = BoolBin( + "and", + BoolBin("and", go_on, Cmp("eq", pid, Const(0), 32)), + Not(Cmp("slt", n, Const(0), 32)), + ) + other = BoolBin("and", go_on, Not(Cmp("eq", pid, Const(0), 32))) + load, store = g.accesses + # ^bb5's predicate, its two incoming edges, is the path; its argument + # (`base`) a Select over what they pass + assert load.path == store.path == BoolBin("or", first, other) + base = Select(first, Const(0), Bin("+", Bin("*", pid, Const(64), 32), n, 32)) + assert load.offset == Bin( + "+", Const(0), Bin("+", base, Arange(load.offset.b.b.ssa, 0, 64), 32) + ) + # both arms of one cond_br into one block + (store,) = _mp("ttir/crafted_same_dest.ttir").accesses + c, a, b = Cmp("slt", Param("arg1"), Param("arg2"), 32), Param("arg1"), Param("arg2") + assert store.path == BoolBin("and", BoolBin("or", c, Not(c)), c) + assert store.offset == Bin("+", Const(0), Bin("+", Select(c, a, b), a, 32)) + # an edge under an unmodeled condition leaves the block guarded + _, store = _mp("reader_ttir/mp_early_return.ttir").accesses + assert (store.guarded, store.path) == (True, None) + # an address handed to inline asm through a block argument + asm = ( + "%i = tt.ptr_to_int %p : !tt.ptr -> i64\ncf.br ^bb1(%i : i64)\n" + '^bb1(%x: i64):\n%y = tt.elementwise_inline_asm "mov.b64 $0, $1;" ' + '{constraints = "=l,l", packed_element = 1 : i32, pure = true} %x : i64 -> i64' + ) + with pytest.raises(UnsupportedTTIR, match="computed from a pointer"): + parse_ttir(_module(asm), multipath=True) + # but not one whose block argument only the branch condition's is + clean = asm.replace( + "cf.br ^bb1(%i : i64)", + "%z = arith.constant 0 : i64\n%c = arith.cmpi ne, %i, %z : i64\n" + "cf.cond_br %c, ^bb1(%z : i64), ^bb1(%z : i64)", + ) + parse_ttir(_module(clean), multipath=True) + + +@pytest.mark.parametrize( + "blocks", + [ + # a block entered before its predecessor holds an access + "cf.cond_br %c, ^bb2, ^bb1\n^bb1:\ntt.store %p, %n : !tt.ptr\n" + "tt.return\n^bb2:\ncf.cond_br %c, ^bb1, ^bb3\n^bb3:", + # a block branching back to itself + "cf.cond_br %c, ^bb1, ^bb2\n^bb1:\ncf.br ^bb1\n^bb2:", + # an address made from the argument of a block entered early + "cf.cond_br %c, ^bb2, ^bb1(%n : i32)\n^bb1(%x: i32):\n" + "%q = tt.addptr %p, %x : !tt.ptr, i32\ntt.return\n" + "^bb2:\ncf.br ^bb1(%n : i32)\n^bb3:", + ], +) +def test_cf_blocks_out_of_order_or_cyclic_refuse(blocks): + text = _module(blocks, "%p: !tt.ptr, %n: i32, %c: i1") + with pytest.raises(UnsupportedTTIR, match="entered") as e: + parse_ttir(text, multipath=True) + assert e.value.kind is TTIRKind.CONTROL_FLOW + + +_STORE_AT_X = "^bb1(%x: i32):\n%q = tt.addptr %p, %x : !tt.ptr, i32\n" +_STORE_AT_X += "tt.store %q, %n : !tt.ptr" + + +@pytest.mark.parametrize( + "blocks, match", + [ + # an edge under an unmodeled condition: a merged value is unknown + ( + "%v = tt.load %f : !tt.ptr\n%z = arith.constant 0.0 : f32\n" + "%c = arith.cmpf ogt, %v, %z : f32\ncf.cond_br %c, ^bb1(%n : i32), " + "^bb1(%n2 : i32)\n" + _STORE_AT_X, + "block argument", + ), + # an unmodelable incoming value + ( + "%v = tt.load %f : !tt.ptr\n%i = arith.fptosi %v : f32 to i32\n" + "%c = arith.cmpi slt, %n, %n2 : i32\ncf.cond_br %c, ^bb1(%i : i32), " + "^bb1(%n : i32)\n" + _STORE_AT_X, + "block argument", + ), + # pointers of two bases + ( + "%c = arith.cmpi slt, %n, %n2 : i32\ncf.cond_br %c, ^bb1(%p : " + "!tt.ptr), ^bb1(%p2 : !tt.ptr)\n^bb1(%x: !tt.ptr):\n" + "tt.store %x, %n : !tt.ptr", + "block argument", + ), + # a branch in generic form names no successors the reader reads + ('"cf.br"()[^bb1] : () -> ()\n^bb1:', "control flow cf.br"), + ], +) +def test_cf_block_arguments_merge_only_modeled_values(blocks, match): + args = "%p: !tt.ptr, %p2: !tt.ptr, %f: !tt.ptr, %n: i32, %n2: i32" + with pytest.raises(UnsupportedTTIR, match=match): + parse_ttir(_module(blocks, args), multipath=True) + + +@pytest.mark.parametrize( + "body, kind", + [ + # a bitwise op the reader does not model over an integer load + ( + "%j = tt.load %q : !tt.ptr\n%c1 = arith.constant 1 : i32\n" + "%x = arith.xori %j, %c1 : i32\n" + "%a = tt.addptr %q, %x : !tt.ptr, i32\n" + "tt.store %a, %c1 : !tt.ptr", + TTIRKind.INDIRECT_ADDRESS, + ), + # casts of an integer load, in a loop bound + ( + "%j = tt.load %q : !tt.ptr\n%f = arith.sitofp %j : i32 to f32\n" + "%u = arith.fptosi %f : f32 to i32\n%c0 = arith.constant 0 : i32\n" + "%c1 = arith.constant 1 : i32\n" + "scf.for %i = %c0 to %u step %c1 : i32 {\n" + "tt.store %q, %c0 : !tt.ptr\n}", + TTIRKind.DATA_DEPENDENT_BOUND, + ), + # a float load cast to an integer, in an address + ( + "%v = tt.load %p : !tt.ptr\n%i = arith.fptosi %v : f32 to i32\n" + "%a = tt.addptr %q, %i : !tt.ptr, i32\n" + "tt.store %a, %i : !tt.ptr", + TTIRKind.INDIRECT_ADDRESS, + ), + # i1 arithmetic over a load + ( + "%b = tt.load %r : !tt.ptr\n%s = arith.addi %b, %b : i1\n" + "%e = arith.extui %s : i1 to i32\n" + "%a = tt.addptr %q, %e : !tt.ptr, i32\n" + "tt.store %a, %e : !tt.ptr", + TTIRKind.INDIRECT_ADDRESS, + ), + # a loop-carried integer that starts as a load + ( + "%j = tt.load %q : !tt.ptr\n%c0 = arith.constant 0 : i32\n" + "%c1 = arith.constant 1 : i32\n%o = scf.for %i = %c0 to %c1 step %c1 " + "iter_args(%x = %j) -> (i32) : i32 {\n" + "%a = tt.addptr %q, %x : !tt.ptr, i32\n" + "tt.store %a, %c0 : !tt.ptr\nscf.yield %x : i32\n}", + TTIRKind.INDIRECT_ADDRESS, + ), + ], +) +def test_a_value_made_from_a_load_refuses_alike_in_both_modes(body, kind): + text = _module(body, "%p: !tt.ptr, %q: !tt.ptr, %r: !tt.ptr") + for options in ({}, {"multipath": True}): + with pytest.raises(UnsupportedTTIR) as e: + parse_ttir(text, **options) + assert e.value.kind is kind diff --git a/tests/unit/sanitizer_compiled/test_oob.py b/tests/unit/sanitizer_compiled/test_oob.py index fe2265dc..255620de 100644 --- a/tests/unit/sanitizer_compiled/test_oob.py +++ b/tests/unit/sanitizer_compiled/test_oob.py @@ -43,6 +43,7 @@ Const, DataDep, FuncArg, + Loaded, LoopInfo, LoopVar, Observed, @@ -513,10 +514,11 @@ def test_an_unbound_loop_bound_refuses_the_loop(): def test_unmodeled_values_and_unusable_facts_refuse(): """Defensive: the reader never lets a DataDep into an address, and bind_launch never builds such facts.""" - g = _synthetic(_access(Bin("+", Const(0), DataDep("loaded value")))) - r = check_graph(g, _bind(p=_facts(4))) - assert r.abstained == ((0, K.UNMODELED_VALUE),) - assert "an unmodeled value (loaded value)" in r.refusal.message + for leaf in (DataDep("loaded value"), Loaded(0, Const(0), None, None)): + g = _synthetic(_access(Bin("+", Const(0), leaf))) + r = check_graph(g, _bind(p=_facts(4))) + assert r.abstained == ((0, K.UNMODELED_VALUE),) + assert "an unmodeled value (loaded value)" in r.refusal.message bad = TensorFacts(PTR, 4, 4, (4,), (), "torch.float32", False) r = check_graph(_synthetic(_access(Const(0))), _bind(p=bad)) assert r.abstained == ((0, K.MISSING_BINDING),) diff --git a/tilelens/clients/sanitizer/compiled/oob.py b/tilelens/clients/sanitizer/compiled/oob.py index b401eff3..51c70507 100644 --- a/tilelens/clients/sanitizer/compiled/oob.py +++ b/tilelens/clients/sanitizer/compiled/oob.py @@ -107,6 +107,7 @@ Arange, Bin, DataDep, + Loaded, LoopInfo, LoopVar, NumPrograms, @@ -452,6 +453,10 @@ def observed(self, t: Observed) -> ArithRef: def data_dep(self, t: DataDep) -> NoReturn: raise _Refused(SanitizerKind.UNMODELED_VALUE, f"an unmodeled value ({t.why})") + def loaded(self, t: Loaded) -> NoReturn: + # what a default parse reads as DataDep("loaded value") + self.data_dep(DataDep("loaded value")) + # ── the loop ── def bounds(self) -> tuple[ArithRef, ArithRef, ArithRef]: diff --git a/tilelens/ir/_mlir_walk.py b/tilelens/ir/_mlir_walk.py index 3acdb11f..037e58dc 100644 --- a/tilelens/ir/_mlir_walk.py +++ b/tilelens/ir/_mlir_walk.py @@ -15,8 +15,9 @@ 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. + ``scf.for``'s ``unsigned`` keyword (a unit attribute no getter reads), a + generic-form ``tt.load``'s operand segments and a branch's successors + (the bindings have no successor getter). 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 @@ -143,7 +144,8 @@ class Op: # their printed keyword ("slt"), program-id axes 0-2, a barrier's # address spaces their bits; # arith.constant "value" is an int (sign-extended) / bool for i1, else - # ("float", None) / ("dense", None); scf.for "unsignedCmp" is a bool. + # ("float", None) / ("dense", None); scf.for "unsignedCmp" is a bool; + # cf.br / cf.cond_br "successors" are the target block indices. 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 @@ -706,9 +708,29 @@ def _text_attrs(t: _TOp, name: str) -> dict[str, Any]: if m is None: raise ValueError("generic tt.load without operandSegmentSizes") return {"operandSegmentSizes": tuple(int(x) for x in m.group(1).split(","))} + if name in ("cf.br", "cf.cond_br") and not t.generic: + return {"successors": _successor_labels(t)} return {} +# A branch's successors: Triton 3.8's bindings have no successor getter, so +# they are the labels the branch prints, each with the number of operands it +# passes (`^bb5(%a, %b : i32, i32)`); the aligner resolves them to blocks. +_RE_SUCCESSOR = re.compile(r"\^[-\w.$]+") +_RE_USE = re.compile(r"%[-\w.$]+(?:#\d+)?") + + +def _successor_labels(t: _TOp) -> tuple[tuple[str, int], ...]: + s = t.masked + out = [] + for m in _RE_SUCCESSOR.finditer(s, t.name_end): + n, j = 0, m.end() + if s.startswith("(", j): + n = len(_RE_USE.findall(s[j : _match_close(s, j)].partition(":")[0])) + out.append((m.group(0), n)) + return tuple(out) + + # A block-pointer type anywhere in a text: `ptr>`, `!tt>>`, whitespace, # newlines or `//` comments in between), and any type alias definition, which @@ -1302,6 +1324,37 @@ def text_loc(self, t: _TOp) -> tuple | None: return None return self.locp.parse(raw[span[0] : span[1]]) + def successors( + self, + labels: tuple[tuple[str, int], ...], + block: int, + n_operands: int, + cond: bool, + ) -> tuple[int, ...]: + """The block indices a branch in ``block`` names (its region's + blocks are all built before its ops are visited), each passed as + many operands as it has arguments; a cf.cond_br also takes its + condition. Raises ValueError otherwise.""" + parent = self.blocks[block] + by_label = { + b.label: i + for i, b in enumerate(self.blocks) + if b.op == parent.op and b.region == parent.region + } + if len(labels) != (2 if cond else 1): + raise ValueError(f"{len(labels)} successors printed") + out = [] + for label, n in labels: + dest = by_label.get(label) + if dest is None: + raise ValueError(f"unknown successor {label}") + if n != len(self.blocks[dest].args): + raise ValueError(f"{label}: {n} operands for its arguments") + out.append(dest) + if n_operands != cond + sum(n for _, n in labels): + raise ValueError(f"{n_operands} operands for its successors") + return tuple(out) + # ── traversal ── def run(self) -> Module: root_w = [i for i, b in enumerate(self.bw.ops) if b.block_id is None] @@ -1413,6 +1466,15 @@ def visit(self, t: _TOp, wi: int, block: int | None, pos: int) -> list: self.bad(line, f"{b.name}: {p}") try: attrs.update(_text_attrs(t, b.name)) + if "successors" in attrs: + if block is None: + raise ValueError("a branch outside a block") + attrs["successors"] = self.successors( + attrs["successors"], + block, + len(b.operands), + b.name == "cf.cond_br", + ) except ValueError as e: self.bad(line, f"{b.name}: {e}") self.stats["attrs"] += len(attrs) diff --git a/tilelens/ir/lowering.py b/tilelens/ir/lowering.py index 20e97af5..067b76f8 100644 --- a/tilelens/ir/lowering.py +++ b/tilelens/ir/lowering.py @@ -3,13 +3,14 @@ Mechanism only: the operator semantics live here, every leaf's meaning comes from the client. A client's :class:`TermLeaves` says what a scalar argument, a program id, the grid, a lane, a loop's -iteration, an atomic observation and an unmodeled value are, and a client -that cannot model a leaf raises its own exception from it, which propagates -unchanged. The lowering itself raises only on a malformed graph, a bug and -never a limit of the model: ``TypeError`` or ``ValueError`` for a term or op -outside the algebra, or for a ``LoopVar`` (or an ``IterArgOffset`` whose -IterArgInfo names no loop) in a graph without a loop, and ``IndexError`` for -an ``arg_id`` outside ``iter_args``. +iteration, an atomic observation, a loaded value and an unmodeled value +are, and a client that cannot model a leaf raises its own exception from +it, which propagates unchanged. The lowering itself raises only on a +malformed graph, a bug and never a limit of the model: ``TypeError`` or +``ValueError`` for a term or op outside the algebra, or for a ``LoopVar`` +whose loop the graph does not have (or an ``IterArgOffset`` whose +IterArgInfo names no loop in a graph without a loop), and ``IndexError`` +for an ``arg_id`` outside ``iter_args``. Operator semantics, the reader's integer model (see ``ttir_reader``): @@ -23,8 +24,9 @@ * an i1 value is a Bool in a boolean position (a mask or path, a Select's condition, ``and`` / ``or`` / negation) and 0/1 in an integer one, and a Select whose arms differ in sort is an Int; -* ``LoopVar`` is ``lower + k * step`` and ``IterArgOffset`` is ``offset0 + - k * delta``, with ``k`` the client's iteration of that loop. +* ``LoopVar`` is ``lower + k * step`` of its loop (``graph.loop_info``) and + ``IterArgOffset`` is ``offset0 + k * delta``, with ``k`` the client's + iteration of that loop. Z3 context: the lowering never creates one. Constants are made in ``leaves.ctx`` (None: Z3's main context) and every other term in its @@ -54,6 +56,7 @@ DataDep, IntCast, IterArgOffset, + Loaded, LoopInfo, LoopVar, Not, @@ -93,6 +96,10 @@ def iteration(self, loop_ssa: str) -> ArithRef: def observed(self, t: Observed) -> ArithRef: """The value the access at ``t.access_index`` observed.""" + def loaded(self, t: Loaded) -> ArithRef: + """The value the load at ``t.access_index`` read at the lane (its + ``offset``, ``mask`` and ``other`` are for the client to lower).""" + def data_dep(self, t: DataDep) -> ArithRef: """A value the reader could not model (``t.why`` says which).""" @@ -167,20 +174,23 @@ def _cmp(pred: str, a: Any, b: Any) -> BoolRef: raise ValueError(f"unknown cmpi predicate {pred!r}") -def _loop(graph: AccessGraph, t: object) -> LoopInfo: - if graph.loop is None: +def _loop(graph: AccessGraph, t: LoopVar | IterArgOffset, name: str) -> LoopInfo: + """The loop named ``name`` (none: the graph's only loop).""" + loop = graph.loop_info(name) if name else graph.loop + if loop is None: raise ValueError( f"kernel {graph.kernel_name!r}: a loop term ({type(t).__name__}) " - "without a loop" + "without a loop" + (f" named {name!r}" if name else "") ) - return graph.loop + return loop def children(t: object, graph: AccessGraph) -> tuple: """The terms ``t``'s value is computed from: its operands; for a loop-carried pointer's offset, its IterArgInfo's ``offset0`` and - ``delta``; for the induction variable, the loop's ``lower`` and - ``step``. A leaf (a DataDep included) has none.""" + ``delta``; for the induction variable, its loop's ``lower`` and + ``step``. A leaf has none: a Loaded's terms are its client's to lower, + and a DataDep's ``keep`` is not its value.""" if isinstance(t, (Bin, Cmp, BoolBin)): return (t.a, t.b) if isinstance(t, Select): @@ -193,7 +203,7 @@ def children(t: object, graph: AccessGraph) -> tuple: info = graph.iter_args[t.arg_id] return (info.offset0, info.delta) if isinstance(t, LoopVar): - loop = _loop(graph, t) + loop = _loop(graph, t, t.loop_ssa) return (loop.lower, loop.step) return () @@ -258,7 +268,7 @@ def _apply(self, t: object, kids: Sequence[Any]) -> Any: return as_int(kids[0]) + k * as_int(kids[1]) if isinstance(t, IterArgOffset): loop_ssa = self.graph.iter_args[t.arg_id].loop_ssa - k = leaves.iteration(loop_ssa or _loop(self.graph, t).loop_ssa) + k = leaves.iteration(loop_ssa or _loop(self.graph, t, "").loop_ssa) return as_int(kids[0]) + k * as_int(kids[1]) if isinstance(t, Bin): return _bin(t.op, as_int(kids[0]), as_int(kids[1])) @@ -283,6 +293,8 @@ def _apply(self, t: object, kids: Sequence[Any]) -> Any: return as_int(kids[0]) if isinstance(t, Observed): return leaves.observed(t) + if isinstance(t, Loaded): + return leaves.loaded(t) if isinstance(t, DataDep): return leaves.data_dep(t) raise TypeError(f"unknown term {type(t).__name__}") diff --git a/tilelens/ir/ttir_reader.py b/tilelens/ir/ttir_reader.py index 8e708ded..f631e5e9 100644 --- a/tilelens/ir/ttir_reader.py +++ b/tilelens/ir/ttir_reader.py @@ -11,7 +11,8 @@ 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``. +kernel itself gathers — the data-dependent case, marked with ``DataDep`` +(with ``parse_ttir(multipath=True)``, an integer load is a ``Loaded``). This module is mechanism-only: it reads and flags (``DataDep`` markers, ``guarded`` accesses, width obligations, ``UnsupportedTTIR``); what to do @@ -34,7 +35,7 @@ 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 +the loops' own increments); 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 @@ -282,6 +283,11 @@ class IntCast: @dataclass(frozen=True) class DataDep: why: str = "value derived from loaded data" + # Under parse_ttir(multipath=True), for an i1 ``and`` with an operand + # that cannot be modeled: the modelable conjunct(s). The value implies + # ``keep``, so an access's mask is ``keep`` instead of free (still + # ``mask_dropped``: wider than the real mask). Not a value dependency. + keep: "Term | None" = None @dataclass(frozen=True) @@ -307,6 +313,29 @@ class Observed: access_index: int +@dataclass(frozen=True) +class Loaded: + """The value the ``tt.load`` at ``graph.accesses[access_index]`` read, + bound under ``parse_ttir(multipath=True)`` to an integer-typed load + whose mask is modeled (a float or pointer load, or one with a dropped + mask, keeps DataDep): at the lane the surrounding term is read at, the + element at ``offset`` of the load's base where ``mask`` holds, and + ``other`` (None: any value) where it does not. The element is read in + the load's own integer type: through a same-width pointer bitcast (an + ``f32`` base read as ``i32``) it is the base's bits reinterpreted, not + its value converted. ``offset``, ``mask`` and ``other`` are the load's + own terms with their lanes placed where the value is read (an + ``expand_dims`` of the loaded tile re-places them like any lane term), + so unlike an Observed, lanes of one load may meet. What a loaded value + means (memory contents, a free value, a refusal) is each consumer's + policy; find them with :func:`loaded_indices`.""" + + access_index: int + offset: "Term" + mask: "Term | None" + other: "Term | None" + + Term = ( Const | Pid @@ -323,6 +352,7 @@ class Observed: | IntCast | DataDep | Observed + | Loaded ) @@ -335,6 +365,10 @@ def _children(t: object) -> tuple: return (t.a,) if isinstance(t, IntCast): return (t.x,) + if isinstance(t, Loaded): + return tuple(x for x in (t.offset, t.mask, t.other) if x is not None) + if isinstance(t, DataDep) and t.keep is not None: + return (t.keep,) return () @@ -342,7 +376,8 @@ 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 + identity), a Loaded's offset, mask and other and a DataDep's keep + included; 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() @@ -378,6 +413,35 @@ def observed_indices(term: object, graph: "AccessGraph") -> frozenset[int]: ) +def loaded_indices(term: object, graph: "AccessGraph") -> frozenset[int]: + """The access indices of every :class:`Loaded` leaf ``term`` reaches, + through loop-carried pointers (see :func:`mentions_observed`) and + through the ``lower``, ``upper`` and ``step`` of the loop of each + LoopVar and loop-carried pointer it reaches (the bounds give the + iterations, and so the values, there are). Whether an access runs at + all also depends on the bounds of its enclosing loops: scan + ``graph.loop_info(name)`` for each name in ``access.loops`` too.""" + found: set[int] = set() + seen: set[str] = set() + roots: list[object] = [term] + while roots: + names: list[str | None] = [] + for n in _nodes(roots, graph.iter_args): + if isinstance(n, Loaded): + found.add(n.access_index) + elif isinstance(n, LoopVar): + names.append(n.loop_ssa) + elif isinstance(n, IterArgOffset): + names.append(graph.iter_args[n.arg_id].loop_ssa) + roots = [] + for name in names: + lp = graph.loop_info(name) if name else graph.loop + if lp is not None and lp.loop_ssa not in seen: + seen.add(lp.loop_ssa) + roots += [lp.lower, lp.upper, lp.step] + return frozenset(found) + + # 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 @@ -438,16 +502,18 @@ class AccessEvent: 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 + # True when some enclosing scf.if condition (or, under multipath, a + # branch into its cf block) 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. + # else-regions negated (Not), after its cf block's predicate under + # multipath. The access executes iff path ∧ mask. path: Term | None = None - # True when the access sits inside the scf.for body: it executes once + # True when the access sits inside an 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 @@ -478,13 +544,17 @@ class AccessEvent: # to the access's own block (one block per region the access sits in). # () only in a hand-built graph, where it reads as the function body. block_path: tuple[int, ...] = () + # The enclosing scf.for loops (LoopInfo.loop_ssa), outermost first; + # ``in_loop == bool(loops)`` (a hand-built graph may leave it empty). + loops: tuple[str, ...] = () @dataclass(frozen=True) class Fence: """A ``ttg.barrier`` (what ``tl.debug_barrier()`` prints), at its place among the accesses: after ``position + 0.5`` of them. ``path``, - ``guarded``, ``in_loop`` and ``block_path`` read as an access's do.""" + ``guarded``, ``in_loop``, ``block_path`` and ``loops`` read as an + access's do.""" position: float op: str @@ -497,6 +567,7 @@ class Fence: block_path: tuple[int, ...] line_no: int loc: SourceLoc | None + loops: tuple[str, ...] = () @dataclass(frozen=True) @@ -539,6 +610,8 @@ class IterArgInfo: @dataclass(frozen=True) class LoopInfo: + # "%loop" for the first scf.for to open, "%loop" for the n-th after + # it (parse_ttir(multipath=True) reads several) loop_ssa: str lower: Term upper: Term @@ -557,7 +630,7 @@ class AccessGraph: kernel_name: str func_args: tuple[FuncArg, ...] accesses: tuple[AccessEvent, ...] - loop: LoopInfo | None + loop: LoopInfo | None # the scf.for when there is exactly one # Loop-carried pointers, indexed by arg_id (``iter_args[k].arg_id == k``), # expanded tiles of them included (see IterArgInfo). iter_args: tuple[IterArgInfo, ...] = () @@ -567,11 +640,21 @@ class AccessGraph: pid_axes: frozenset[int] = frozenset() fences: tuple[Fence, ...] = () whiles: tuple[WhileRegion, ...] = () # only under while_regions=True + # Every scf.for in opening order (outer before inner); more than one + # only under multipath=True. A graph built with only ``loop`` gets + # ``loops == (loop,)``, and one built with a single entry in ``loops`` + # gets it as ``loop``. + loops: tuple[LoopInfo, ...] = () def __post_init__(self) -> None: object.__setattr__(self, "pid_axes", frozenset(self.pid_axes)) - for name in ("func_args", "accesses", "iter_args", "fences", "whiles"): + names = ("func_args", "accesses", "iter_args", "fences", "whiles", "loops") + for name in names: object.__setattr__(self, name, tuple(getattr(self, name))) + if not self.loops and self.loop is not None: + object.__setattr__(self, "loops", (self.loop,)) + elif self.loop is None and len(self.loops) == 1: + object.__setattr__(self, "loop", self.loops[0]) def arg(self, name: str) -> FuncArg | None: for a in self.func_args: @@ -579,6 +662,10 @@ def arg(self, name: str) -> FuncArg | None: return a return None + def loop_info(self, loop_ssa: str) -> LoopInfo | None: + """The LoopInfo named ``loop_ssa``, or None.""" + return next((lp for lp in self.loops if lp.loop_ssa == loop_ssa), None) + # ─────────────────────────── width obligations ─────────────────────────── @@ -603,7 +690,7 @@ def width_obligations( ) -> 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 + and each loop enclosing the access) equals the IR's fixed-width arithmetic: * every integer ``Bin`` result fits its width, signed; @@ -613,24 +700,28 @@ def width_obligations( * 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 + * each 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. + ``role`` says where an obligation comes from: ``"loop"`` (the loops' + bounds and increments, outermost loop first), ``"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, ``path`` ones hold unconditionally and ``loop`` ones wherever the + loop is reached (the increment's only when it runs at least once), which + is everywhere unless a multipath graph nests the loop under a branch; + the access's own path implies that its loops are reached, so it guards + their obligations soundly. 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. Likewise a :class:`Loaded`'s + ``offset``, ``mask`` and ``other`` are listed under the role of the term + that reads it, without the load's mask. 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() @@ -650,7 +741,7 @@ def need(term: object, bits: int, signed: bool, site: object, role: str) -> None ) groups: list[tuple[str, tuple[object, ...]]] = [] - if loop is not None: + for loop in _enclosing_loops(graph, access): if loop.bits is not None: if loop.unsigned: for bound in (loop.lower, loop.upper, loop.step): @@ -691,17 +782,54 @@ def need(term: object, bits: int, signed: bool, site: object, role: str) -> None return tuple(out) +def _enclosing_loops(graph: AccessGraph, access: AccessEvent) -> list[LoopInfo]: + """The LoopInfos of ``access.loops``, outermost first (``graph.loop`` for + an in-loop access of a hand-built single-loop graph that names none).""" + if not access.loops: + return [graph.loop] if access.in_loop and graph.loop is not None else [] + found = [graph.loop_info(name) for name in access.loops] + if None in found: + raise ValueError( + f"kernel {graph.kernel_name!r}: unknown loop in {access.loops}" + ) + return found # type: ignore[return-value] + + # ─────────────────────────── the reader ─────────────────────────── -def parse_ttir(text: str, *, while_regions: bool = False) -> AccessGraph: - """Read one TTIR module into an AccessGraph (single-path model: at most - one scf.for, structured scf.if only). +def parse_ttir( + text: str, *, while_regions: bool = False, multipath: bool = False +) -> AccessGraph: + """Read one TTIR module into an AccessGraph (by default the single-path + model: at most one scf.for, structured scf.if only). ``while_regions=True`` reads each ``scf.while`` too (a ``WhileRegion`` in ``graph.whiles``) instead of refusing it as control flow; a kernel without one reads the same either way. + ``multipath=True`` reads four more facts: + + * several ``scf.for`` loops, sequential, nested or under a branch, each + with its own LoopInfo (``graph.loops``) and induction variable, and + each access's enclosing loops (``AccessEvent.loops``); ``graph.loop`` + is set only for a single loop; + * the unstructured blocks Triton prints for an ``if`` that holds a + ``return`` (``cf.br`` / ``cf.cond_br`` / ``tt.return``): a block's + predicate, the disjunction of its incoming edges (the predecessor's + predicate and the branch condition, negated on the false edge), is + the outermost conjunct of the path of every access in it, and a block + argument is a Select over its incoming edges' operands; an edge whose + condition is not modelable keeps only the predecessor's predicate and + leaves the block ``guarded`` with its arguments DataDep; + * an integer ``tt.load`` result whose mask is modeled as + :class:`Loaded` instead of DataDep; + * a mask that is an ``and`` of modelable and unmodelable values as its + modelable conjuncts (``DataDep.keep``), still ``mask_dropped``. + + Each starts at a refusal of the default model (a second loop, a cf + block) or a DataDep, so a kernel without them reads the same either way. + 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 @@ -722,7 +850,7 @@ def parse_ttir(text: str, *, while_regions: bool = False) -> AccessGraph: raise UnsupportedTTIR( TTIRKind.UNPARSABLE, e.diagnostic, line_no=e.line_no ) from e - return _Builder(module, while_regions).build() + return _Builder(module, while_regions, multipath).build() # Memory ops outside the modeled vocabulary (TMA descriptors, ...: targets @@ -771,10 +899,11 @@ def parse_ttir(text: str, *, while_regions: bool = False) -> AccessGraph: # 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 +_LOOP_SSA = "%loop" # the first loop's name (LoopInfo.loop_ssa) # 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" +_EARLY_ARG = "argument of a block entered early" _RE_ADDR_SPACE = re.compile(r", (\d+)>$") @@ -793,9 +922,31 @@ class _IfFrame: branch: str = "then" +@dataclass class _ForFrame: """Marks the scf.for body being read.""" + loop_ssa: str + + +@dataclass +class _Block: + """Walker state for one block of a function body with several blocks + (multipath=True). ``edges`` are its incoming edges recorded at the + predecessors' branches, (path, exact, operand values): ``exact`` is + False when the branch condition could not be modeled. A block entered + before all its predecessors' branches were read is not ``resolved``: its + predicate is unknown, so an access, fence or branch in it refuses, and + so does an address made from one of its arguments (MLIR's block merging + can print a shared tail, a return-only one or not, before one of its + predecessors).""" + + n_preds: int + edges: list[tuple[Term | None, bool, list[object]]] = field(default_factory=list) + pred: Term | None = None + guarded: bool = False + resolved: bool = True + @dataclass class _WhileFrame: @@ -825,16 +976,18 @@ def _pointee_bits(pointee: str | None) -> int: class _Builder: - def __init__(self, module: Module, while_regions: bool = False) -> None: + def __init__( + self, module: Module, while_regions: bool = False, multipath: bool = False + ) -> None: self.m = module self.while_regions = while_regions + self.multipath = multipath # 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.loops: list[LoopInfo | None] = [] # in opening order, set on closing self.fences: list[Fence] = [] self.whiles: list[WhileRegion | None] = [] self.frames: list[_IfFrame | _ForFrame | _WhileFrame] = [] @@ -848,6 +1001,12 @@ def __init__(self, module: Module, while_regions: bool = False) -> None: # it reads lane for lane (see positional) self.access_op: dict[int, int] = {} self.prov: dict[int, frozenset[int]] = {} + # the block being read of a function body with several blocks, and + # all of them by block index + self.cf: _Block | None = None + self.cf_table: dict[int, _Block] = {} + # id(term) -> (term, whether it reaches a Loaded), see reads_loaded + self.loaded_memo: dict[int, tuple[object, bool]] = {} # ── helpers ── def refuse(self, kind: TTIRKind, message: str, op: Op | None) -> NoReturn: @@ -870,6 +1029,21 @@ def bind(self, op: Op, value: object) -> None: for r in op.results: self.env[r] = value + def reads_loaded(self, x: object) -> bool: + """True when the term ``x`` reaches a Loaded (multipath only): such a + value derives from memory, as its DataDep would by default.""" + if not self.multipath or isinstance(x, (DataDep, PtrValue)): + return False + hit = self.loaded_memo.get(id(x)) + if hit is None or hit[0] is not x: + reach = any(isinstance(n, Loaded) for n in _nodes((x,), None)) + hit = self.loaded_memo[id(x)] = (x, reach) + return hit[1] + + def data_dep(self, operands: Iterable[object], why: str) -> DataDep: + """_data_dep, a term that reads a Loaded counted as memory-derived.""" + return _data_dep(operands, why, self.reads_loaded) + 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) @@ -880,17 +1054,19 @@ def as_term(self, v: object, ctx: str, op: Op) -> Term: 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.""" + def branch_state(self, op: Op) -> tuple[bool, Term | None, tuple[str, ...]]: + """(guarded, path, loops) for an access (or fence) ``op`` under the + open frames: ``guarded`` if any enclosing condition is unmodeled; + ``path`` is the conjunction of the modeled ones (else-regions + negated), outermost first, after the predicate of the cf block + holding them; ``loops`` the enclosing scf.for loops, outermost + first.""" guarded = False path: Term | None = None - in_loop = False + loops: list[str] = [] for f in self.frames: if isinstance(f, _ForFrame): - in_loop = True + loops.append(f.loop_ssa) continue if isinstance(f, _WhileFrame): continue @@ -898,8 +1074,19 @@ def branch_state(self) -> tuple[bool, Term | None, bool]: 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 + path = _conj(path, c) + blk = self.cf + if blk is not None: + if not blk.resolved: + self.refuse( + TTIRKind.CONTROL_FLOW, + f"{op.name} in a block entered before its predecessors", + op, + ) + guarded = guarded or blk.guarded + if blk.pred is not None: + path = blk.pred if path is None else BoolBin("and", blk.pred, path) + return guarded, path, tuple(loops) def arg(self, name: str) -> FuncArg | None: return next((a for a in self.func_args if a.name == name), None) @@ -953,18 +1140,22 @@ def build(self) -> AccessGraph: for fa in func.args: self.bind_arg(fa, fop) body = fop.regions[0] - self.block(body[0]) - if len(body) > 1: + if self.multipath and len(body) > 1: + self.cf_blocks(body) + else: + self.block(body[0]) + if len(body) > 1 and not self.multipath: self.refuse( TTIRKind.CONTROL_FLOW, f"tt.func {func.sym_name!r} has {len(body)} blocks", fop, ) + loops = tuple(self.loops) return AccessGraph( kernel_name=func.sym_name, func_args=tuple(self.func_args), accesses=tuple(self.accesses), - loop=self.loop, + loop=loops[0] if len(loops) == 1 else None, # type: ignore[arg-type] iter_args=tuple(self.iter_args), pid_axes=frozenset( op.attrs["axis"] @@ -973,6 +1164,7 @@ def build(self) -> AccessGraph: ), fences=tuple(self.fences), whiles=tuple(self.whiles), # type: ignore[arg-type] + loops=loops, # type: ignore[arg-type] ) def bind_arg(self, fa, fop: Op) -> None: @@ -1072,7 +1264,7 @@ def expanded_iter_arg(self, t: IterArgOffset, axis: int) -> IterArgOffset: 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] + IterArgInfo(aid, src.base_param, offset0, Const(0), src.loop_ssa) # type: ignore[arg-type] ) self.expanded[key] = aid return IterArgOffset(aid) @@ -1080,6 +1272,9 @@ def expanded_iter_arg(self, t: IterArgOffset, axis: int) -> IterArgOffset: def op_cast(self, op: Op) -> None: x = self.val(op.operands[0]) if isinstance(x, (DataDep, PtrValue)): + # a cast of an i1 widens it: its keep (an i1 fact) stays behind + if isinstance(x, DataDep) and x.keep is not None: + x = DataDep(x.why) self.bind(op, x) return kind = op.name.split(".", 1)[1] @@ -1126,9 +1321,9 @@ 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")) + self.bind(op, self.data_dep((a, b), "arith over loaded data")) elif bits == 1: - self.bind(op, DataDep(f"{op.name} on i1")) + self.bind(op, self.data_dep((a, b), f"{op.name} on i1")) else: self.bind( op, @@ -1147,10 +1342,10 @@ def op_cmpi(self, op: Op) -> None: 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")) + self.bind(op, self.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")) + self.bind(op, self.data_dep((a, b), "signed comparison of i1 values")) else: self.bind( op, @@ -1177,8 +1372,15 @@ def op_boolbin(self, op: Op) -> None: 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)) + why = self.data_dep((a, b), "bool op over loaded data").why + keep: Term | None = None + if is_and and self.multipath: + # ``modelable and unmodelable`` implies ``modelable`` + for x in (a, b): + part = x.keep if isinstance(x, DataDep) else x + if part is not None: + keep = _conj(keep, part) # type: ignore[arg-type] + self.bind(op, DataDep(why, keep)) return self.bind( op, @@ -1201,7 +1403,7 @@ def merge(self, c: object, t: object, f: object, what: str) -> object: or isinstance(t, DataDep) or (isinstance(f, DataDep)) ): - return _data_dep((c, t, f), "select over loaded data") + return self.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") @@ -1223,7 +1425,9 @@ def op_bitcast(self, op: Op) -> None: 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}")) + self.bind( + op, self.data_dep((x,), f"unmodeled op {op.name} at line {op.line_no}") + ) # ── accesses ── def access( @@ -1259,14 +1463,16 @@ def access( 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. + # (any lane may be active), or as its modelable conjuncts, + # instead of failing the kernel. See AccessEvent.mask_dropped + # for the soundness discipline. mask_dropped = True + mask = mv.keep 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() + guarded, path, loops = self.branch_state(op) assert op.line_no is not None # the value operands: a store's value and mask, an atomic_rmw's val # and mask, an atomic_cas's cmp and val (ODS order, after the pointer) @@ -1283,7 +1489,7 @@ def access( line_no=op.line_no, guarded=guarded, path=path, - in_loop=in_loop, + in_loop=bool(loops), mask_dropped=mask_dropped, atomic=atomic, atomic_val=atomic_val, @@ -1293,6 +1499,7 @@ def access( and parse_type(pointee).pointee is None, deps=deps, block_path=op.path, + loops=loops, ) ) @@ -1366,9 +1573,23 @@ def op_load(self, op: Op) -> None: if any(isinstance(f, _WhileFrame) and f.stage == "before" for f in self.frames): # a poll: what the loop's condition reads self.bind(op, self.observed_binding(op, "loaded value")) + elif self.multipath: + self.bind(op, self.loaded_binding(op, mask_v)) else: self.bind(op, DataDep("loaded value")) + def loaded_binding(self, op: Op, mask_v: int | None) -> object: + """The value of the just-recorded load: Loaded for an integer-typed + result whose mask is modeled, DataDep otherwise (see Loaded).""" + index = len(self.accesses) - 1 + acc = self.accesses[index] + if acc.mask_dropped or self.int_bits(op.results[0]) is None: + return DataDep("loaded value") + other = None # ODS (ptr, mask?, other?): other follows a mask + if mask_v is not None and len(op.operands) == 3: + other = _operand_term(self.val(op.operands[2])) + return Loaded(index, acc.offset, acc.mask, other) + def op_store(self, op: Op) -> None: # ODS operands (ptr, value, mask?) self.access(op, "store", self.mask_operand(op, 2)) @@ -1401,8 +1622,10 @@ 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: + # limitation, not one more induction variable. Multipath reads each + # loop with its own, except inside an scf.while. + in_while = any(isinstance(f, _WhileFrame) for f in self.frames) + if (self.loops or self.frames) and (not self.multipath or in_while): self.refuse( TTIRKind.CONTROL_FLOW if any(isinstance(f, _IfFrame) for f in self.frames) @@ -1410,7 +1633,9 @@ def op_for(self, op: Op) -> None: "multiple/nested loops", op, ) - self.loops_opened += 1 + index = len(self.loops) + loop_ssa = f"{_LOOP_SSA}{index}" if index else _LOOP_SSA + self.loops.append(None) # set on closing, after the nested ones bounds: dict[str, Term] = {} for label, v in zip(("lower", "upper", "step"), op.operands[:3]): bv = self.val(v) @@ -1434,22 +1659,22 @@ def op_for(self, op: Op) -> None: 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) + 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) + 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.env[arg_v] = self.data_dep((init,), _LOOP_CARRIED) + self.frames.append(_ForFrame(loop_ssa)) self.block(op.regions[0][0]) self.frames.pop() yield_op = self.m.ops[body.ops[-1]] @@ -1459,6 +1684,8 @@ def op_for(self, op: Op) -> None: # 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(): + if self.iter_args[src].loop_ssa != loop_ssa: + continue source = self.iter_args[src].delta if any( isinstance(n, Observed) and n.access_index in self.lane_observed @@ -1472,8 +1699,8 @@ def op_for(self, op: Op) -> None: ) 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, + self.loops[index] = LoopInfo( + loop_ssa=loop_ssa, lower=bounds["lower"], upper=bounds["upper"], step=bounds["step"], @@ -1490,6 +1717,7 @@ def loop_delta(self, y: object, aid: int, op: Op) -> Term: same base: the model reads iteration k's offset as ``offset0 + k * delta``.""" info = self.iter_args[aid] + loop = info.loop_ssa if not isinstance(y, PtrValue) or y.base_param != info.base_param: self.refuse( TTIRKind.LOOP_VARIANT_ADVANCE, @@ -1505,11 +1733,17 @@ def loop_delta(self, y: object, aid: int, op: Op) -> Term: "its own previous value", op, ) + # variant in this loop: its induction variable, one of its carried + # pointers, or a value an access inside it read (an enclosing + # loop's are invariant here) for n in _nodes((delta,), None): if ( - isinstance(n, (LoopVar, IterArgOffset)) - or isinstance(n, Observed) - and self.accesses[n.access_index].in_loop + isinstance(n, LoopVar) + and n.loop_ssa == loop + or isinstance(n, IterArgOffset) + and self.iter_args[n.arg_id].loop_ssa == loop + or isinstance(n, (Observed, Loaded)) + and loop in self.accesses[n.access_index].loops ): self.refuse( TTIRKind.LOOP_VARIANT_ADVANCE, @@ -1540,7 +1774,7 @@ def op_if(self, op: Op) -> None: 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") + self.env[r] = self.data_dep((cv,), "select over loaded data") else: self.env[r] = self.merge(cond, yields[0][k], yields[1][k], "scf.if") @@ -1589,7 +1823,7 @@ def op_condition(self, op: Op) -> None: frame.stage, frame.split = "after", len(self.accesses) def op_barrier(self, op: Op) -> None: - guarded, path, in_loop = self.branch_state() + guarded, path, loops = self.branch_state(op) assert op.line_no is not None self.fences.append( Fence( @@ -1598,13 +1832,85 @@ def op_barrier(self, op: Op) -> None: addr_space=op.attrs["addr_space"], path=path, guarded=guarded, - in_loop=in_loop, + in_loop=bool(loops), block_path=op.path, line_no=op.line_no, loc=op.loc, + loops=loops, ) ) + # ── the blocks of an if that holds a return (multipath) ── + def cf_blocks(self, body: Sequence[int]) -> None: + """Read a function body with several blocks in printed order (see + parse_ttir). Triton places a block after its predecessors, except a + shared tail that MLIR's block merging may print before one of them, + and never branches back: the first block of a cycle is entered + before a predecessor, so its branch refuses. A terminator other than + a custom-form cf.br / cf.cond_br or tt.return refuses before any + block is read, so it is reported ahead of an earlier refusal.""" + n_preds = dict.fromkeys(body, 0) + for b in body: + term = self.m.ops[self.m.blocks[b].ops[-1]] + if term.name in ("cf.br", "cf.cond_br") and "successors" in term.attrs: + for d in term.attrs["successors"]: + n_preds[d] += 1 + elif term.name != "tt.return": + self.op_control_flow(term) # cf.switch, a generic-form branch + self.cf_table = {b: _Block(n_preds[b]) for b in body} + for b in body: + blk = self.cf = self.cf_table[b] + args = self.m.blocks[b].args if b != body[0] else () + if len(blk.edges) < blk.n_preds: + blk.resolved = False + for v in args: + self.env[v] = DataDep(_EARLY_ARG) + elif b != body[0]: + pred: Term | None = Cmp("ne", Const(0), Const(0)) # unreachable + for i, (path, _, _) in enumerate(blk.edges): + pred = path if i == 0 else _disj(pred, path) + blk.pred = pred + blk.guarded = not all(exact for _, exact, _ in blk.edges) + for k, v in enumerate(args): + self.env[v] = ( + DataDep("block argument") + if blk.guarded + else _merge_block_arg(blk.edges, k) + ) + self.block(b) + self.cf = None + + def op_branch(self, op: Op) -> None: + """cf.br / cf.cond_br at the end of a block being read: its targets' + incoming edges.""" + blk = self.cf + if blk is None or self.frames: + self.op_control_flow(op) + if not blk.resolved: + self.refuse( + TTIRKind.CONTROL_FLOW, + f"{op.name} in a block entered before its predecessors", + op, + ) + operands = list(op.operands) + edges: list[tuple[Term | None, bool]] = [(blk.pred, not blk.guarded)] + if op.name == "cf.cond_br": + c = self.val(operands.pop(0)) + if isinstance(c, (DataDep, PtrValue)): + # both targets stay reachable under the predecessor's + # predicate alone (a widening) + edges = [(blk.pred, False)] * 2 + else: + ct: Term = c # type: ignore[assignment] + edges = [ + (_conj(blk.pred, ct), not blk.guarded), + (_conj(blk.pred, Not(ct)), not blk.guarded), + ] + for d, (path, exact) in zip(op.attrs["successors"], edges): + n = len(self.m.blocks[d].args) + passed, operands = operands[:n], operands[n:] + self.cf_table[d].edges.append((path, exact, [self.val(v) for v in passed])) + # ── everything else ── def op_control_flow(self, op: Op) -> NoReturn: self.refuse(TTIRKind.CONTROL_FLOW, f"control flow {op.name} is unsupported", op) @@ -1647,6 +1953,16 @@ def from_pointer(self, values: Iterable[int]) -> bool: assert value.block is not None src = m.ops[m.blocks[value.block].op] if src.name == "tt.func": + # an argument of a later block (multipath): what each + # branch into it passes at its position + k = m.blocks[value.block].args.index(v) + for b in src.regions[0]: + term = m.ops[m.blocks[b].ops[-1]] + rest = term.operands[term.name == "cf.cond_br" :] + for d in term.attrs.get("successors", ()): + if d == value.block: + stack.append(rest[k]) + rest = rest[len(m.blocks[d].args) :] continue stack += src.operands for region in src.regions: @@ -1660,7 +1976,7 @@ 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")) + self.bind(op, self.data_dep(map(self.val, op.operands), "unmodeled inline asm")) def op_extern(self, op: Op) -> None: hazard = self.opaque_hazard(op) @@ -1689,7 +2005,7 @@ def op_other(self, op: Op) -> None: # when an operand is. self.bind( op, - _data_dep( + self.data_dep( map(self.val, op.operands), f"unmodeled op {name} at line {op.line_no}" ), ) @@ -1724,14 +2040,19 @@ def pure_regions(self, op: Op) -> None: stack.extend(b for region in inner.regions for b in region) -def _data_dep(operands: Iterable[object], why: str) -> DataDep: +def _data_dep( + operands: Iterable[object], + why: str, + reads_loaded: Callable[[object], bool] | None = None, +) -> DataDep: """The DataDep of an op the reader cannot model over ``operands``: a - memory-derived reason when any operand derives from memory (``why`` + memory-derived reason when any operand derives from memory (a + memory-derived DataDep, or a term ``reads_loaded`` holds for; ``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): + if _from_memory(x) or (reads_loaded is not None and reads_loaded(x)): return DataDep( why if why.startswith(_MEMORY_WHYS) @@ -1742,6 +2063,39 @@ def _data_dep(operands: Iterable[object], why: str) -> DataDep: return DataDep(first.why) if first is not None else DataDep(why) +def _conj(a: Term | None, b: Term) -> Term: + return b if a is None else BoolBin("and", a, b) + + +def _disj(a: Term | None, b: Term | None) -> Term | None: + return None if a is None or b is None else BoolBin("or", a, b) + + +def _merge_block_arg(edges: Sequence[tuple[Term | None, bool, list]], k: int) -> object: + """Block argument ``k`` as a Select over the incoming edges (all exact): + the last edge's operand is the fallback, each earlier one is selected + under its own path. Pointers of one base select their offsets; anything + else unmodelable is DataDep.""" + if not edges: + return DataDep("argument of an unreachable block") + vals = [e[2][k] for e in edges] + if all(isinstance(v, PtrValue) for v in vals): + if len({v.base_param for v in vals}) != 1: # type: ignore[union-attr] + return DataDep("block argument merging different bases") + offsets = [v.offset for v in vals] # type: ignore[union-attr] + return PtrValue(vals[0].base_param, _select_chain(edges, offsets)) # type: ignore[union-attr] + if any(isinstance(v, (DataDep, PtrValue)) for v in vals): + return DataDep("block argument") + return _select_chain(edges, vals) + + +def _select_chain(edges: Sequence[tuple[Term | None, bool, list]], vals: list) -> Term: + term = vals[-1] + for (path, _, _), v in reversed(list(zip(edges[:-1], vals[:-1]))): + term = v if path is None else Select(path, v, term) + return term + + def _why(v: object) -> str: return f" ({v.why})" if isinstance(v, DataDep) else "" @@ -1753,6 +2107,8 @@ def _address_kind(v: object) -> TTIRKind: return TTIRKind.INDIRECT_ADDRESS if isinstance(v, DataDep) and v.why == _LOOP_CARRIED: return TTIRKind.LOOP_VARIANT_ADVANCE + if isinstance(v, DataDep) and v.why == _EARLY_ARG: + return TTIRKind.CONTROL_FLOW return TTIRKind.OTHER @@ -1770,6 +2126,16 @@ def _with_children(t: object, kids: tuple) -> object: return replace(t, a=kids[0]) if isinstance(t, IntCast): return replace(t, x=kids[0]) + if isinstance(t, Loaded): + it = iter(kids) + return Loaded( + t.access_index, + next(it), + None if t.mask is None else next(it), + None if t.other is None else next(it), + ) + if isinstance(t, DataDep): + return replace(t, keep=kids[0]) raise TypeError(type(t).__name__) @@ -1859,6 +2225,8 @@ def _loop_delta(offset: Term, arg_id: int) -> Term | None: "tt.store": _Builder.op_store, "tt.atomic_rmw": _Builder.op_atomic_rmw, "tt.atomic_cas": _Builder.op_atomic_cas, + "cf.br": _Builder.op_branch, + "cf.cond_br": _Builder.op_branch, "scf.for": _Builder.op_for, "scf.if": _Builder.op_if, "scf.while": _Builder.op_while,