[FIX] Stop the sanitizer from skipping repeated launches

benchmarks/bench_sanitizer.py

@@ -36,10 +36,7 @@
 import triton.language as tl
 
 import tilelens
-from tilelens.clients.sanitizer.sanitizer import (
-    SymbolicSanitizer,
-    _fn_symbolic_cache_set,
-)
+from tilelens.clients.sanitizer.sanitizer import SymbolicSanitizer
 
 # ---------------------------------------------------------------------------
 # Benchmark timeout helper
@@ -694,14 +691,12 @@ def flaggems_ln_wb_backward_kernel(
 
 
 def _reset_sanitizer(san: SymbolicSanitizer):
-    """Clear sanitizer state and the global cache so re-runs are not skipped."""
-    _fn_symbolic_cache_set.clear()
+    """Clear sanitizer records between benchmark runs."""
     san.records.clear()
 
 
 # ---------------------------------------------------------------------------
 # GEMM benchmarks: multiple (M, N, K, TILE_SIZE) constexpr combos
-# Each unique combo is a separate symbolic cache entry.
 # ---------------------------------------------------------------------------
 
 _GEMM_CONFIGS = [

tests/end_to_end/test_sanitizer.py

@@ -1615,3 +1615,43 @@ def test_strided_view_warns_on_large_numel(_isolate_per_element_threshold):
 
     # Warning path must not change correctness -- no OOB with correct mask.
     assert len(strided_view_sanitizer.records) == 0
+
+
+# ======== Repeated Launch Tests ===========
+
+
+relaunch_sanitizer = SymbolicSanitizer(abort_on_error=False)
+
+
+@tilelens.trace(client=relaunch_sanitizer)
+@triton.jit
+def relaunch_gather_kernel(x_ptr, idx_ptr, out_ptr, BLOCK: tl.constexpr):
+    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    idx = tl.load(idx_ptr + offs)
+    tl.store(out_ptr + offs, tl.load(x_ptr + idx))
+
+
+@pytest.mark.parametrize("num_blocks", [1, 2])
+def test_relaunch_with_new_loaded_indices_is_checked(num_blocks):
+    """
+    A launch that repeats an earlier launch's kernel, grid, and argument
+    shapes is checked against its own data.
+
+    Setup:
+      - x holds 8 float32 elements.
+      - The first launch gathers with in-bounds indices.
+      - The second launch keeps every shape but sets idx[0] = 100, so
+        block 0 loads x[100].
+    """
+    relaunch_sanitizer.records.clear()
+
+    x = torch.arange(8, dtype=torch.float32)
+    idx = torch.zeros((4 * num_blocks,), dtype=torch.int64)
+    out = torch.empty((4 * num_blocks,), dtype=torch.float32)
+
+    relaunch_gather_kernel[(num_blocks,)](x, idx, out, BLOCK=4)
+    assert relaunch_sanitizer.records == []
+
+    idx[0] = 100
+    relaunch_gather_kernel[(num_blocks,)](x, idx, out, BLOCK=4)
+    assert len(relaunch_sanitizer.records) > 0, "Expected OOB to be detected"

tests/unit/test_sanitizer.py

@@ -19,7 +19,6 @@
 from tilelens.clients.sanitizer.sanitizer import (
     NullSanitizer,
     SymbolicSanitizer,
-    _fn_symbolic_cache_set,
 )
 from tilelens.clients.sanitizer.range_summary import IntRange, access_interval_summary
 from tilelens.clients.symbolic_engine import LoopContext, PendingCheck, SymbolicExpr
@@ -50,69 +49,28 @@ def test_sanitizer_init(_isolate_sanitizer_cfg):
 #
 # These guard the two spots that break most easily when SymbolicClient's
 # launch-lifecycle machinery is reshuffled:
-#   1. the cross-launch fn cache (sanitizer's pre_run_callback short-circuit)
+#   1. sanitizer's pre_run_callback must not skip a launch that repeats an
+#      earlier launch's kernel, grid, and argument shapes.
 #   2. the post_run_callback guard that must NOT wipe launch state mid-
 #      full-grid enumeration.
 
 
-@pytest.fixture
-def _isolate_fn_symbolic_cache():
-    """Reset the cross-launch sanitizer fn-cache around each test."""
-    saved = set(_fn_symbolic_cache_set)
-    _fn_symbolic_cache_set.clear()
-    yield _fn_symbolic_cache_set
-    _fn_symbolic_cache_set.clear()
-    _fn_symbolic_cache_set.update(saved)
-
-
-def _populate_for_launch(sanitizer: SymbolicSanitizer, grid: tuple[int, ...]) -> None:
-    """Prime a sanitizer as if ``arg_callback`` + ``grid_callback`` had fired.
-
-    Uses a fake cache_args payload instead of real tensors so the test stays
-    focused on pre_run_callback's cache-hit logic.
-    """
-    sanitizer.cache_args = [("fake_shape",)]
-    sanitizer.cache_grid = grid
-
-
-def test_pre_run_callback_cache_hit_on_replay(_isolate_fn_symbolic_cache):
-    """Identical (fn, grid, args) launched twice: 1st runs, 2nd is cached-skip."""
+def test_pre_run_callback_runs_replayed_launch():
+    """Identical (fn, grid) launched twice: both run, because loaded values
+    may differ between the two launches."""
 
     def kernel():
         pass
 
-    first = SymbolicSanitizer(abort_on_error=False)
-    _populate_for_launch(first, (2, 1, 1))
-    assert first.pre_run_callback(kernel) is True
-    # pre_run_callback clears the cache after hashing, matching production flow.
-    assert first.cache_grid is None and first.cache_args == []
-    assert len(_isolate_fn_symbolic_cache) == 1
-
-    second = SymbolicSanitizer(abort_on_error=False)
-    _populate_for_launch(second, (2, 1, 1))
-    assert second.pre_run_callback(kernel) is False
-    assert len(_isolate_fn_symbolic_cache) == 1
-
-
-def test_pre_run_callback_cache_miss_on_grid_change(_isolate_fn_symbolic_cache):
-    """Changing the grid must NOT trigger a cache hit — it's part of the key."""
-
-    def kernel():
-        pass
-
-    a = SymbolicSanitizer(abort_on_error=False)
-    _populate_for_launch(a, (2, 1, 1))
-    assert a.pre_run_callback(kernel) is True
-
-    b = SymbolicSanitizer(abort_on_error=False)
-    _populate_for_launch(b, (4, 1, 1))
-    assert b.pre_run_callback(kernel) is True
-    assert len(_isolate_fn_symbolic_cache) == 2
+    for _ in range(2):
+        sanitizer = SymbolicSanitizer(abort_on_error=False)
+        sanitizer.grid_callback((2, 1, 1))
+        assert sanitizer.pre_run_callback(kernel) is True
 
 
 def test_post_run_callback_preserves_state_mid_full_grid():
     """When need_full_grid=True and we're NOT on the last block, launch state
-    (tensors / tensor_addrs / tensor_names / cache_*) must survive
+    (tensors / tensor_addrs / tensor_names) must survive
     post_run_callback so later blocks still see registered tensors."""
     sanitizer = SymbolicSanitizer(abort_on_error=False)
 
@@ -121,8 +79,6 @@ def test_post_run_callback_preserves_state_mid_full_grid():
     sanitizer.tensors.append(cast(object, sentinel_tensor))  # type: ignore[arg-type]
     sanitizer.tensor_addrs.append((0, 0, cast(object, sentinel_tensor)))  # type: ignore[arg-type]
     sanitizer.tensor_names[id(sentinel_tensor)] = {"X"}
-    sanitizer.cache_args = [("fake",)]
-    sanitizer.cache_grid = (4, 1, 1)
 
     sanitizer.last_grid = (3, 0, 0)  # last block is (3,0,0) in a 4-block grid
     sanitizer.grid_idx = (1, 0, 0)  # we're mid-enumeration
@@ -134,8 +90,6 @@ def test_post_run_callback_preserves_state_mid_full_grid():
     assert sanitizer.tensors == [sentinel_tensor]
     assert sanitizer.tensor_addrs == [(0, 0, sentinel_tensor)]
     assert sanitizer.tensor_names == {id(sentinel_tensor): {"X"}}
-    assert sanitizer.cache_args == [("fake",)]
-    assert sanitizer.cache_grid == (4, 1, 1)
 
 
 def test_post_run_callback_clears_state_on_last_block():
@@ -146,8 +100,6 @@ def test_post_run_callback_clears_state_on_last_block():
     sanitizer.tensors.append(cast(object, sentinel_tensor))  # type: ignore[arg-type]
     sanitizer.tensor_addrs.append((0, 0, cast(object, sentinel_tensor)))  # type: ignore[arg-type]
     sanitizer.tensor_names[id(sentinel_tensor)] = {"X"}
-    sanitizer.cache_args = [("fake",)]
-    sanitizer.cache_grid = (4, 1, 1)
 
     sanitizer.last_grid = (3, 0, 0)
     sanitizer.grid_idx = (3, 0, 0)  # final block
@@ -158,8 +110,6 @@ def test_post_run_callback_clears_state_on_last_block():
     assert sanitizer.tensors == []
     assert sanitizer.tensor_addrs == []
     assert sanitizer.tensor_names == {}
-    assert sanitizer.cache_args == []
-    assert sanitizer.cache_grid is None
 
 
 def _make_concrete_ptr_and_mask(

tilelens/clients/sanitizer/sanitizer.py

@@ -1,6 +1,4 @@
 from collections.abc import Callable
-from dataclasses import dataclass
-from functools import cached_property
 from typing import (
     Any,
     ClassVar,
@@ -11,7 +9,6 @@
 
 import numpy as np
 from torch import Tensor
-from triton.tools.tensor_descriptor import TensorDescriptor
 from z3 import And, BoolVal, Not, Or, sat
 from z3.z3 import BoolRef, IntNumRef
 
@@ -136,34 +133,18 @@ def _make_signature(
     return hash((addr_hash, constr_hash))
 
 
-@dataclass(frozen=True)
-class _FnSymbolicCache:
-    fn: Callable
-    grid: tuple[int, ...]
-    args: tuple
-
-    @cached_property
-    def hash_value(self) -> int:
-        return hash((self.fn, self.grid, self.args))
-
-    def __hash__(self) -> int:
-        return self.hash_value
-
-
-_fn_symbolic_cache_set: set[_FnSymbolicCache] = set()
-
-
 class SymbolicSanitizer(Sanitizer, SymbolicClient):
     def __init__(self, abort_on_error: bool = True):
         super().__init__(abort_on_error=abort_on_error)
         self.records: list[OutOfBoundsRecordZ3] = []
-        self.cache_args: list[Any] = []
-        self.cache_grid: tuple[int, ...] | None = None
 
     # Explicit forwarders to SymbolicClient: the Sanitizer factory
     # carries concrete stubs (NotImplementedError or ``return True``) to
     # satisfy Client's @abstractmethod contract, and those stubs would
     # otherwise shadow SymbolicClient's impls in the subclass MRO.
+    def grid_callback(self, grid: tuple[int, ...]) -> None:
+        SymbolicClient.grid_callback(self, grid)
+
     def grid_idx_callback(self, grid_idx: tuple[int, ...]) -> None:
         SymbolicClient.grid_idx_callback(self, grid_idx)
 
@@ -187,6 +168,12 @@ def register_op_callback(self, op_type: type[Op]) -> OpCallbacks:
             return cached
         return SymbolicClient.register_op_callback(self, op_type)
 
+    # Every launch is analyzed, including one whose kernel, grid, and
+    # argument shapes match an earlier launch: the earlier result does not
+    # carry over when loaded values feed addresses, masks, or branches.
+    def pre_run_callback(self, fn: Callable) -> bool:
+        return SymbolicClient.pre_run_callback(self, fn)
+
     def post_run_callback(self, fn: Callable) -> bool:
         return SymbolicClient.post_run_callback(self, fn)
 
@@ -214,37 +201,6 @@ def _addr_ok_premise(self) -> BoolRef:
         # ranges instead of the global union of all registered tensors.
         return BoolVal(True)
 
-    def _cache_non_tensor_arg(self, name: str, arg: Any) -> None:
-        # TODO: init a reserved_args field per frontend to filter out these args
-        if name not in ["num_warps", "num_stages", "maxnreg", "num_ctas"]:
-            if isinstance(arg, TensorDescriptor):
-                self.cache_args.append(
-                    (
-                        "TensorDescriptor",
-                        repr(arg.base),
-                        tuple(arg.shape),
-                        tuple(arg.strides),
-                        tuple(arg.block_shape),
-                        arg.padding,
-                    )
-                )
-                return
-            self.cache_args.append(arg)
-
-    def _cache_tensor_arg(self, arg: Tensor) -> None:
-        self.cache_args.append((arg.shape, arg.stride(), arg.dtype))
-
-    def _clear_cache(self) -> None:
-        self.cache_args.clear()
-        self.cache_grid = None
-
-    def grid_callback(self, grid: tuple[int, ...]) -> None:
-        # Sanitizer needs ``cache_grid`` tracked alongside the shared setup so
-        # ``pre_run_callback`` can consult it before deciding whether to skip
-        # the launch.
-        self.cache_grid = tuple(int(g) for g in grid)
-        SymbolicClient.grid_callback(self, grid)
-
     def _op_load_overrider(self, ptr, mask, other, *_args):
         ptr = self._materialize_memory_operand(ptr)
         mask = self._materialize_memory_operand(mask)
@@ -351,16 +307,6 @@ def _concrete_addr_range_info(
             return None
         return ConcreteRangeInfo(tensor, ranges)
 
-    def pre_run_callback(self, fn: Callable) -> bool:
-        if self.cache_grid:
-            fn_cache = _FnSymbolicCache(fn, self.cache_grid, tuple(self.cache_args))
-            self._clear_cache()
-            if fn_cache not in _fn_symbolic_cache_set:
-                _fn_symbolic_cache_set.add(fn_cache)
-                return SymbolicClient.pre_run_callback(self, fn)
-            return False
-        return SymbolicClient.pre_run_callback(self, fn)
-
     def _find_tensor_for_expr(
         self, symbolic_expr: SymbolicExpr, violation_addr: int
     ) -> Tensor: