[DEV] Support TRITON_INTERPRET=1 (#73)

Author: mark14wu

.github/workflows/python-app.yml

@@ -5,11 +5,13 @@ name: Python application
 
 on:
   push:
-    branches-ignore:
-      - '**'
+    branches:
+      - main
+      - keren/v2.0
   pull_request:
-    branches-ignore:
-      - '**'
+    branches:
+      - main
+      - keren/v2.0
 
 permissions:
   contents: read
@@ -21,6 +23,8 @@ concurrency:
 jobs:
   build:
     runs-on: ubuntu-latest
+    env:
+      TRITON_INTERPRET: "1"
 
     steps:
     - uses: actions/checkout@v3
@@ -44,7 +48,7 @@ jobs:
         pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121
         pip uninstall pytorch-triton -y
 
-    - name: Clone Triton and Install
+    - name: Install Triton
       run: |
         pip install triton==3.1.0
 

tests/test_autotune_add.py

@@ -1,3 +1,4 @@
+import pytest
 import torch
 import triton
 import triton.language as tl
@@ -7,8 +8,12 @@
 from triton_viz import config as cfg
 
 
-cfg.sanitizer_backend = "symexec"
+try:
+    torch.cuda.current_device()
+except:
+    pytest.skip("This test requires a CUDA-enabled environment.", allow_module_level=True)
 
+cfg.sanitizer_backend = "symexec"
 
 @triton.autotune(
     configs=[
@@ -39,8 +44,8 @@ def test_autotune_add_inrange():
     This test uses n_elements = 128, matching the size of the input tensors.
     It should NOT cause any out-of-bound access.
     """
-    x = torch.randn(128, device="cuda")
-    y = torch.randn(128, device="cuda")
+    x = torch.randn(128)
+    y = torch.randn(128)
     out = torch.empty_like(x)
 
     # The kernel launch uses n_elements=128, aligned with the tensor size.
@@ -55,8 +60,8 @@ def test_autotune_add_out_of_bound():
     This test deliberately sets n_elements = 256, exceeding the actual buffer size (128).
     It will likely cause out-of-bound reads/writes, which may trigger errors or warnings.
     """
-    x = torch.randn(128, device="cuda")
-    y = torch.randn(128, device="cuda")
+    x = torch.randn(128)
+    y = torch.randn(128)
     out = torch.empty_like(x)
 
     # The kernel launch uses n_elements=256, exceeding the valid tensor size.

tests/test_config.py

@@ -1,4 +1,5 @@
-import pytest
+import pytest, os
+os.environ["TRITON_SANITIZER_BACKEND"] = "off"
 import triton_viz.core.config as cfg
 
 

tests/test_print_traceback.py

@@ -27,7 +27,7 @@ def kernel_A(ptr, n):
 
 
 def test_print_nested_functions():
-    x = torch.arange(4, device="cuda", dtype=torch.float32)
+    x = torch.arange(4, dtype=torch.float32)
     print("Input:", x)
 
     # We'll launch a grid bigger than x.numel() to force a out-of-bounds error

tests/test_wrapper.py

@@ -59,6 +59,7 @@ def _decorator(fn):
         env = os.environ.copy()
         env["PYTHONPATH"] = str(tmp_path) + os.pathsep + env.get("PYTHONPATH", "")
         env["TRITON_SANITIZER_BACKEND"] = "symexec"
+        env["TRITON_INTERPRET"] = "1"
 
         # run the dummy program using triton-sanitizer
         cmd = ["triton-sanitizer", str(tmp_path / "dummy_program.py")]

triton_viz/core/trace.py

@@ -34,10 +34,20 @@ def add_client(self, new_client: Union[Client, str]) -> None:
         new_client_instance = self._normalize_client(new_client)
         self.client_manager.add_clients([new_client_instance])
 
-    def __init__(self, kernel: JITFunction, client: Union[str, Client]) -> None:
-        assert isinstance(kernel, JITFunction), "Kernel must be a JITFunction"
-        self.interpreter_fn = InterpretedFunction(kernel.fn)
+    def __init__(
+        self,
+        kernel: Union[JITFunction, InterpretedFunction],
+        client: Union[str, Client],
+    ) -> None:
         self.fn = kernel
+        if isinstance(kernel, InterpretedFunction):
+            self.interpreter_fn = kernel
+        elif isinstance(kernel, JITFunction):
+            self.interpreter_fn = InterpretedFunction(kernel.fn)
+        else:
+            raise TypeError(
+                f"Kernel must be JITFunction or InterpretedFunction, got {type(kernel)}"
+            )
         self.arg_names = kernel.arg_names
         self.client_manager = ClientManager()
         self.add_client(client)
@@ -76,7 +86,7 @@ def decorator(kernel) -> Trace:
             return kernel
 
         # First-time wrapping
-        if isinstance(kernel, JITFunction):
+        if isinstance(kernel, (JITFunction, InterpretedFunction)):
             return Trace(kernel, clients)
 
         # If the object is already a Trace, just append the new client(s)
@@ -85,8 +95,9 @@ def decorator(kernel) -> Trace:
             trace.add_client(clients)
             return trace
 
-        # If the object is neither a JITFunction nor Trace, raise an error
-        raise TypeError(f"Expected JITFunction, got {type(kernel)}")
+        raise TypeError(
+            f"Expected JITFunction, InterpretedFunction or Trace, got {type(kernel)}"
+        )
 
     return decorator