[FEAT] Add cuTile interpreter and tracing support

.github/workflows/python-app.yml

@@ -66,18 +66,18 @@ jobs:
         uv run pytest tests -n auto --ignore=tests/end_to_end/test_triton_kernels.py --ignore=tests/unit/test_multithreading.py
         TRITON_INTERPRET=1 uv run pytest tests/unit/test_multithreading.py tests/end_to_end/test_core.py::test_autotune_interpreter_mode
 
-    - name: Install TileLens with NKI extras
+    - name: Install TileLens with NKI and cuTile extras
       run: |
         cd tilelens
-        uv sync --extra test --extra nki
+        uv sync --extra test --extra nki --extra cutile
 
     - name: Upgrade Triton to the latest 3.8 release (NKI)
       run: |
         cd tilelens
         uv pip install --pre -U "triton<3.9"
         uv pip install "git+https://github.com/triton-lang/triton.git@v3.8.0#subdirectory=python/triton_kernels"
 
-    - name: Run full (Triton + NKI) pytest suite
+    - name: Run full (Triton + NKI + cuTile) pytest suite
       run: |
         cd tilelens
         uv run pytest tests -m "" -n auto --ignore=tests/end_to_end/test_triton_kernels.py --ignore=tests/unit/test_multithreading.py
@@ -94,22 +94,27 @@ jobs:
         python-version: '3.12'
 
     - name: Install with pip
-      run: pip install ".[test]"
+      run: pip install ".[test,cutile]"
 
     - name: Verify installation
       working-directory: ${{ runner.temp }}
       run: |
         python - <<'PY'
         import tilelens
+        import cuda.tile
         import triton_viz
         from tilelens.clients import Sanitizer
         from triton_viz.clients import Sanitizer as LegacySanitizer
         from tilelens.core.trace import launches
+        from tilelens.core.frontend.cutile import frontend
+        from tilelens.core.simulation.cutile import CuTileInterpretedFunction
         from triton_viz.core.trace import launches as legacy_launches
 
         assert triton_viz.trace is tilelens.trace
         assert triton_viz.config is tilelens.config
         assert LegacySanitizer is Sanitizer
         assert legacy_launches is launches
+        assert frontend.name == "cutile"
+        assert CuTileInterpretedFunction is not None
         print(tilelens.__version__)
         PY

README.md

@@ -120,10 +120,13 @@ uv sync --extra nki # NKI support but no testing
 uv sync --extra test # tests but no NKI support
 ```
 
+cuTile support is optional too. Install it with `pip install "tilelens[cutile]"`, or from source with `uv sync --extra cutile` (add `--extra test` for its tests). cuda-tile publishes wheels only for Linux and Windows.
+
 ### Testing
 * To run core TileLens tests, run `pytest tests/`.
 * (if NKI installed) To run NKI-specific tests, run `pytest tests/ -m nki`.
-* To run all tests (Triton + NKI), run `pytest tests/ -m ""`.
+* To run cuTile-specific tests, run `pytest tests/unit/test_cutile.py tests/end_to_end/test_cutile.py`.
+* To run all tests (Triton + NKI + cuTile), install all extras and run `pytest tests/ -m ""`.
 * To run visualizer web UI tests, run `npm run test:frontend`.
 
 ## Working with Examples
@@ -171,18 +174,41 @@ tilelens.launch()
 
 ## DSL Frontends
 
-Triton is the default DSL frontend. NKI support is optional and selected with
-the `frontend` argument:
+Triton is the default DSL frontend. NKI and cuTile are also supported. Select
+them with the `frontend` argument:
 
 ```py
 tilelens.trace("tracer")  # Triton
 tilelens.trace("tracer", frontend="nki")  # NKI
 tilelens.trace("tracer", frontend="nki_beta2")  # NKI Beta 2
+tilelens.trace("tracer", frontend="cutile")  # cuTile
 ```
 
-The runtime integration code lives under `tilelens/core/frontend/`. NKI
+The runtime integration code lives under `tilelens/core/frontend/`. NKI and cuTile
 simulation runtimes live under `tilelens/core/simulation/`.
 
+### cuTile tracing
+
+cuTile kernels run on the CPU with NumPy arrays and support tracing and visualization:
+
+```sh
+uv sync --extra test --extra cutile
+uv run --extra test --extra cutile python -m examples.cutile.matmul --visualize
+```
+
+```py
+@tilelens.trace("tracer", frontend="cutile")
+@ct.kernel
+def kernel(lhs, rhs, result):
+    ...
+
+
+kernel[grid](lhs, rhs, result)
+```
+
+Use contiguous NumPy arrays for visualization. cuTile analysis currently supports
+only the tracer client; numerical behavior may differ from GPU execution.
+
 ## Analysis Clients
 
 Analyze kernels across visualization, profiling, and sanitization with a single line of code.

docs/index.html

@@ -571,7 +571,7 @@
         </div>
         <h1>Analyze and debug your tile-based DSL kernels</h1>
         <p style="display: block; font-size: 0.95rem; color: var(--text-muted); margin-top: -1rem; margin-bottom: 2.5rem;">
-            Note: Triton is the default DSL frontend; Amazon NKI support is optional and in active development.
+            Note: Triton is the default DSL frontend; NVIDIA cuTile is also supported, and Amazon NKI support is optional and in active development.
         </p>
         <div class="install-container">
             <div class="install-box">
@@ -593,6 +593,15 @@ <h1>Analyze and debug your tile-based DSL kernels</h1>
                         </button>
                     </div>
                 </div>
+                <div class="install-step">
+                    <div class="install-label">Install (cuTile support)</div>
+                    <div class="install-command">
+                        <code style="white-space: pre-wrap;">pip install "tilelens[cutile]"</code>
+                        <button class="copy-btn" onclick="navigator.clipboard.writeText('pip install &quot;tilelens[cutile]&quot;')" title="Copy to clipboard">
+                            <svg width="18" height="18" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><rect x="9" y="9" width="13" height="13" rx="2" ry="2"></rect><path d="M5 15H4a2 2 0 0 1-2-2V4a2 2 0 0 1 2-2h9a2 2 0 0 1 2 2v1"></path></svg>
+                        </button>
+                    </div>
+                </div>
             </div>
             <p style="font-size: 0.9rem; color: var(--text-muted); text-align: center; margin-top: 1rem;">
                 Note: A browser with <strong>WebGL/OpenGL enabled</strong> is required (standard in modern browsers). <br>

examples/cutile/matmul.py

@@ -0,0 +1,81 @@
+"""Trace a tiled cuTile matmul with optional visualization and trace export."""
+
+import argparse
+import math
+
+import cuda.tile as ct
+import numpy as np
+
+import tilelens
+
+
+TILE_M = 2
+TILE_K = 4
+TILE_N = 8
+
+
+@ct.kernel
+def matmul_kernel(lhs, rhs, result):
+    """Compute matrix multiplication one output tile at a time.
+
+    Args:
+        lhs: Left-hand input with shape [M, K].
+        rhs: Right-hand input with shape [K, N].
+        result: Output with shape [M, N].
+    """
+    _, K = lhs.shape
+
+    # Compute one output tile per block.
+    m, n = ct.bid(0), ct.bid(1)
+    result_tile = ct.full((TILE_M, TILE_N), 0, ct.float32)
+
+    # Accumulate partial sums.
+    for k in range(ct.cdiv(K, TILE_K)):
+        lhs_tile = ct.load(
+            lhs,
+            (m, k),
+            (TILE_M, TILE_K),
+            padding_mode=ct.PaddingMode.ZERO,
+        )
+        rhs_tile = ct.load(
+            rhs,
+            (k, n),
+            (TILE_K, TILE_N),
+            padding_mode=ct.PaddingMode.ZERO,
+        )
+        result_tile = result_tile + ct.matmul(lhs_tile, rhs_tile)
+
+    ct.store(result, (m, n), result_tile)
+
+
+def _run_demo():
+    parser = argparse.ArgumentParser(description="Trace a cuTile matmul on the CPU")
+    parser.add_argument(
+        "--visualize", action="store_true", help="Open the local visualizer"
+    )
+    parser.add_argument("--save", metavar="PATH", help="Save the trace as a .tvz file")
+    args = parser.parse_args()
+    M = 4
+    K = 8
+    N = 16
+
+    lhs = np.arange(M * K, dtype=np.float32).reshape(M, K)
+    rhs = np.arange(K * N, dtype=np.float32).reshape(K, N)
+    result = np.empty((M, N), dtype=np.float32)
+    kernel_grid = (math.ceil(M / TILE_M), math.ceil(N / TILE_N))
+
+    print("Executing matmul_kernel with the cuTile NumPy interpreter...")
+    traced_kernel = tilelens.trace("tracer", frontend="cutile")(matmul_kernel)
+    traced_kernel[kernel_grid](lhs, rhs, result)
+
+    expected = lhs @ rhs
+    print(np.max(np.abs(expected - result)))
+    assert np.allclose(expected, result)
+    if args.save:
+        tilelens.save(args.save)
+    if args.visualize:
+        tilelens.launch(share=False)
+
+
+if __name__ == "__main__":
+    _run_demo()

pyproject.toml

@@ -65,6 +65,9 @@ test = [
   "numpy",
   "torch",
 ]
+cutile = [
+  "cuda-tile>=1.0.0,<2",
+]
 nki = [
   "neuronx-cc==2.22.12471.0+b4a00d10; python_version < '3.13'",
   "nki==0.2.0+g82fdb402; python_version < '3.13'"