diff --git a/tests/end_to_end/test_sanitizer.py b/tests/end_to_end/test_sanitizer.py index a039a906..ecfc761b 100644 --- a/tests/end_to_end/test_sanitizer.py +++ b/tests/end_to_end/test_sanitizer.py @@ -1209,6 +1209,117 @@ def test_sanitizer_supports_data_dependent_cumsum_index(): assert len(cumsum_sanitizer.records) == 0 +def _offsets_into_b(a, b, indices): + return [ + (b.data_ptr() + i * b.element_size() - a.data_ptr()) // a.element_size() + for i in indices + ] + + +@pytest.mark.parametrize("op", ["load", "store"]) +@pytest.mark.parametrize( + "lanes, reported", + [ + ("a a a a", False), + ("b b b b", True), + ("a b b b", True), + ("b a a a", True), + ("b b past past", True), + ], +) +def test_sanitizer_checks_loaded_offsets_against_their_base_tensor(op, lanes, reported): + """An address computed from loaded offsets is concretized; it must still be + checked against the tensor of the pointer it is computed from (a_ptr), not + the tensor that its first address falls in.""" + sanitizer = SymbolicSanitizer(abort_on_error=False) + + # Defined per test, so that no launch repeats an earlier one. + @tilelens.trace(client=sanitizer) + @triton.jit + def load_kernel(a_ptr, b_ptr, idx_ptr, out_ptr, N: tl.constexpr): + offs = tl.arange(0, N) + idx = tl.load(idx_ptr + offs) # offsets from a_ptr, loaded from memory + y = tl.load(b_ptr + offs) # in-bounds use of b, so b is registered + tl.store(out_ptr + offs, tl.load(a_ptr + idx) + y) + + @tilelens.trace(client=sanitizer) + @triton.jit + def store_kernel(a_ptr, b_ptr, idx_ptr, N: tl.constexpr): + offs = tl.arange(0, N) + idx = tl.load(idx_ptr + offs) + tl.store(a_ptr + idx, tl.load(b_ptr + offs)) + + a = torch.zeros(16, dtype=torch.float32) + b = torch.ones(16, dtype=torch.float32) + into_b = _offsets_into_b(a, b, range(4)) + where = { + "a": list(range(4)), + "b": into_b, + "past": [into_b[0] + 1000, into_b[0] + 2000], + } + used = {name: 0 for name in where} + idx = [] + for lane in lanes.split(): + idx.append(where[lane][used[lane]]) + used[lane] += 1 + idx = torch.tensor(idx, dtype=torch.int64) + + if op == "load": + load_kernel[(1,)](a, b, idx, torch.empty(4), N=4) + else: + store_kernel[(1,)](a, b, idx, N=4) + + names = {record.tensor_name for record in sanitizer.records} + assert names == ({"a_ptr"} if reported else set()) + + +def test_sanitizer_checks_loaded_offsets_in_a_loop_against_their_base_tensor(): + """The pointer tile is built before the loop and shared by its iterations; + concretizing the first iteration's address must not lose its base.""" + sanitizer = SymbolicSanitizer(abort_on_error=False) + + @tilelens.trace(client=sanitizer) + @triton.jit + def loop_kernel(a_ptr, b_ptr, idx_ptr, N: tl.constexpr, ITERS: tl.constexpr): + offs = tl.arange(0, N) + a_ptrs = a_ptr + offs * 0 + for i in range(ITERS): + idx = tl.load(idx_ptr + i * N + offs) # loaded offsets from a_ptrs + tl.store(a_ptrs + idx, tl.load(b_ptr + offs)) + + a = torch.zeros(16, dtype=torch.float32) + b = torch.ones(16, dtype=torch.float32) + # Iteration 0 stays in a; iteration 1 lands entirely in b. + idx = torch.tensor( + list(range(4)) + _offsets_into_b(a, b, range(4)), dtype=torch.int64 + ) + + loop_kernel[(1,)](a, b, idx, N=4, ITERS=2) + + assert {record.tensor_name for record in sanitizer.records} == {"a_ptr"} + + +def test_sanitizer_resolves_a_loaded_pointer_from_its_value(): + """A pointer loaded from memory has no pointer argument to start from; + its value decides the tensor, here src, which is in bounds.""" + sanitizer = SymbolicSanitizer(abort_on_error=False) + + @tilelens.trace(client=sanitizer) + @triton.jit + def pointer_table_kernel(table_ptr, src_ptr, out_ptr, N: tl.constexpr): + offs = tl.arange(0, N) + src = tl.load(table_ptr).to(tl.pointer_type(tl.float32)) + tl.store(out_ptr + offs, tl.load(src + offs + tl.load(src_ptr).to(tl.int32))) + + src = torch.zeros(8, dtype=torch.float32) + table = torch.tensor([src.data_ptr()], dtype=torch.int64) + out = torch.empty(4, dtype=torch.float32) + + pointer_table_kernel[(1,)](table, src, out, N=4) + + assert len(sanitizer.records) == 0 + + ashr_index_sanitizer = SymbolicSanitizer(abort_on_error=False) diff --git a/tilelens/clients/symbolic_engine.py b/tilelens/clients/symbolic_engine.py index af435c88..0126f1b3 100644 --- a/tilelens/clients/symbolic_engine.py +++ b/tilelens/clients/symbolic_engine.py @@ -527,7 +527,10 @@ def signature(self) -> int: value_part: Any = None if self.op == "const": - value_part = self._value_signature(getattr(self, "value", None)) + value_part = ( + self._value_signature(getattr(self, "value", None)), + getattr(self, "origin_base", None), + ) self.expr_signature = hash( ( @@ -804,11 +807,14 @@ def replace_subtree(self, anchor_op: str | None = None) -> SymbolicExpr: if not isinstance(concrete, SymbolicTensorValue): raise TypeError(f"Unexpected dtype: {type(concrete)}!") + origin = getattr(self, "origin_base", None) self = SymbolicExpr.create( "const", concrete, type_spec(concrete.dtype, concrete.shape), ) + if origin is not None: + self.origin_base = origin return self @@ -2920,6 +2926,15 @@ def _materialize_memory_operand(self, expr: Any) -> Any: if materialized: self._on_data_dependent_value() if expr.has_vector_const(): + # The full replacement turns a pointer into its concrete + # addresses; keep the pointer argument it is computed from, + # which decides the tensor whose bounds the access must keep. + # It replaces subtrees in place, and they may be shared with + # later accesses (a pointer tile built before a loop), so every + # pointer node on the way keeps it. + base = self._pointer_base(expr) + if base is not None: + self._mark_pointer_base(expr, base) expr = expr.replace_subtree() return expr @@ -3124,6 +3139,62 @@ def register_for_loop_callback(self) -> ForLoopCallbacks: # ── Tensor resolution helpers ───────────────────────────────────── + @classmethod + def _pointer_base(cls, expr: SymbolicExpr) -> Any | None: + """Return the pointer argument that ``expr`` is computed from, or None. + + Follow only pointer-valued operands (pointer arithmetic, broadcasts, + selects) down to a scalar pointer constant. A pointer loaded from + memory or cast from an integer has no such base: only its concrete + value tells where it points. + """ + if not is_pointer_dtype(expr.dtype) or expr.op == "load": + return None + if expr.op == "const": + origin = getattr(expr, "origin_base", None) + if origin is not None: + return origin + value = getattr(expr, "value", None) + if isinstance(value, SymbolicTensorValue) and value.data.size != 1: + return None + return expr.to_py() + bases = set() + for child in expr.children.values(): + for item in child if isinstance(child, tuple) else (child,): + if ( + item is None + or not is_pointer_dtype(item.dtype) + or isinstance( + getattr(item, "value", None), + (SymbolicScalarDType, SymbolicPointerDType, SymbolicTypeSpec), + ) + ): + continue + base = cls._pointer_base(item) + if base is None: + return None + bases.add(int(base)) + return bases.pop() if len(bases) == 1 else None + + @classmethod + def _mark_pointer_base(cls, expr: SymbolicExpr, base: Any) -> None: + """Set ``origin_base`` on ``expr`` and on the pointer-valued operands + that ``_pointer_base`` followed to find ``base``.""" + expr.origin_base = base + if expr.op == "const": + return + for child in expr.children.values(): + for item in child if isinstance(child, tuple) else (child,): + if ( + item is not None + and is_pointer_dtype(item.dtype) + and not isinstance( + getattr(item, "value", None), + (SymbolicScalarDType, SymbolicPointerDType, SymbolicTypeSpec), + ) + ): + cls._mark_pointer_base(item, base) + def _collect_tensor_base(self, expr: SymbolicExpr) -> Any | None: def walk(node: SymbolicExpr) -> Any | None: if ( @@ -3134,7 +3205,8 @@ def walk(node: SymbolicExpr) -> Any | None: (SymbolicScalarDType, SymbolicPointerDType, SymbolicTypeSpec), ) ): - return node.to_py() + origin = getattr(node, "origin_base", None) + return origin if origin is not None else node.to_py() for child in node.children.values(): if child is None: continue