diff --git a/tests/test_core.py b/tests/test_core.py index 2748b621c..7b19013f3 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -52,6 +52,36 @@ def my_kernel(x_ptr, y_ptr, out_ptr, BLOCK_SIZE: tl.constexpr): assert sum(c == "tracer" for c in clients) == 1 +# ======== Nested JIT Call Tests ========= +@triton_viz.trace(clients=Sanitizer(abort_on_error=True)) +@triton.jit +def trace_nested_inner_kernel(x): + return x * 2 + + +def test_trace_nested_jit_calls(): + """ + Test that Trace class properly handles nested JIT function calls via __call__ method. + + When a traced JIT function is called from within another JIT function, + the Trace wrapper needs to properly delegate to the underlying function. + This test ensures compatibility with nested Triton JIT function calls. + """ + + @triton_viz.trace(clients=Sanitizer(abort_on_error=True)) + @triton.jit + def trace_nested_call_kernel(ptr, n: tl.constexpr): + x = tl.load(ptr + tl.arange(0, n)) + y = trace_nested_inner_kernel( + x + ) # This nested call requires __call__ method when wrapped by trace + tl.store(ptr + tl.arange(0, n), y) + + # Test execution + data = torch.ones(8) + trace_nested_call_kernel[(1,)](data, 8) + + # ======== Wrapper Tests ========= # Test that triton_viz.wrapper works correctly: # 1. It should patch triton.jit / triton.language.jit / diff --git a/triton_viz/core/trace.py b/triton_viz/core/trace.py index 2e5f52923..6cba7815a 100644 --- a/triton_viz/core/trace.py +++ b/triton_viz/core/trace.py @@ -69,6 +69,11 @@ def run(self, *args, **kwargs): self.finalize() return ret + def __call__(self, *args, **kwargs): + # When a traced JIT function is called from within another JIT function, + # we need to execute the underlying function directly + return self.base_fn(*args, **kwargs) + def warmup(self, *args, **kwargs): raise NotImplementedError