diff --git a/tests/end_to_end/test_sanitizer.py b/tests/end_to_end/test_sanitizer.py index a039a906..8bf903bc 100644 --- a/tests/end_to_end/test_sanitizer.py +++ b/tests/end_to_end/test_sanitizer.py @@ -1209,6 +1209,28 @@ def test_sanitizer_supports_data_dependent_cumsum_index(): assert len(cumsum_sanitizer.records) == 0 +rand_sanitizer = SymbolicSanitizer(abort_on_error=False) + + +@tilelens.trace(client=rand_sanitizer) +@triton.jit +def rand_store_kernel(out_ptr, seed, BLOCK: tl.constexpr): + offs = tl.arange(0, BLOCK) + tl.store(out_ptr + offs, tl.rand(seed, offs)) + + +def test_sanitizer_supports_tl_rand(): + """tl.rand's Philox rounds call tl.math.umulhi, which must reach the + interpreter builder's create_umulhi like any other builder op.""" + rand_sanitizer.records.clear() + + out = torch.empty((8,), dtype=torch.float32) + + rand_store_kernel[(1,)](out, 123, BLOCK=out.numel()) + + assert len(rand_sanitizer.records) == 0 + + ashr_index_sanitizer = SymbolicSanitizer(abort_on_error=False) diff --git a/tilelens/core/frontend/triton.py b/tilelens/core/frontend/triton.py index 30b4846a..13f1c5b7 100644 --- a/tilelens/core/frontend/triton.py +++ b/tilelens/core/frontend/triton.py @@ -173,9 +173,6 @@ "xor_sum": ReduceXor, "reduce_or": ReduceOr, }, - tl.math: { - "umulhi": Umulhi, - }, } @@ -260,10 +257,8 @@ def __init__(self): self._triton_extra_modules = self._triton_extra_builtin_modules() self._triton_builtin_attr_cache: dict[int, tuple[str, ...]] = {} # Most patched ops are interpreter-builder ops and can call overriders - # directly. Cache the smaller frontend-wrapped sets for quick routing. - self._math_op_types = frozenset(self.namespaces[tl.math].values()) - self._tl_op_types = frozenset(self.namespaces[tl].values()) - self._frontend_wrapped_op_types = self._math_op_types | self._tl_op_types + # directly. Cache the smaller frontend-wrapped set for quick routing. + self._frontend_wrapped_op_types = frozenset(self.namespaces[tl].values()) self._bind_interpreter_builder_thread_state() def _bind_interpreter_builder_thread_state(self) -> None: @@ -854,8 +849,6 @@ def run_op_overrider( # namespace ops need Triton tensor <-> handle conversion around clients. if op_type not in self._frontend_wrapped_op_types: return op_overrider(*args, **kwargs) - if op_type in self._math_op_types: - raise NotImplementedError("Patching math ops not yet supported") # see triton.runtime.interpreter:ReduceOps.sum # First, convert input from tl.tensor to TensorHandle. Here, input tensor is args[0]