[FIX] Route tl.math.umulhi through the interpreter builder

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)
 
 

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]