Support Triton repo's triton_kernels custom types via lazy import; add test for single launch and Tensor type (#101)

Summary:
- Add first-class support for custom data types provided by the Triton repository's `triton_kernels` module (e.g., `Tensor`, `Storage`, `StridedLayout`).
- Avoid top-level imports for optional `triton_kernels.tensor` by probing availability and importing lazily at use sites.
- Update CUDA test to emit a single launch event to logs and verify:
  - Exactly one `.ndjson` log file is produced in the temp logs directory
  - Exactly one `launch` event exists in that file
  - `extracted_args["X"].type == "triton_kernels.tensor.Tensor"`

## Motivation
- The custom types (`Tensor`, `Storage`, `StridedLayout`) are defined in the Triton repo under `triton_kernels` and may not be installed by default. Importing them at module import time causes hard failures for users without the extra.
- We still want to reliably validate parsing and logging by asserting that the launch event is recorded and carries the expected structured argument metadata.

## Changes
### tritonparse/reproducer/utils.py
- Add top-level probe (non-import) for optional dependency:
  - `TRITON_KERNELS_CUSTOM_TYPES = importlib.util.find_spec("triton_kernels.tensor") is not None`
- Add lazy loader with caching:
  - `_get_triton_tensor_types()` loads `Tensor`, `Storage`, `StridedLayout` on first use and is decorated with `functools.lru_cache(maxsize=1)`
- Update `_create_arg_from_info` branches to:
  - Guard with `if not TRITON_KERNELS_CUSTOM_TYPES: raise RuntimeError(...)`
  - Import via the cached lazy loader at the point of use
  - Preserve existing tensor dtype handling for random data generation

### tests/test_tritonparse.py
- Modify `test_triton_kernels_Tensor` to:
  - Import `from tritonparse.reproducer import utils as reproducer_utils`
  - Initialize `tritonparse.structured_logging` with a temp logs directory
  - Run `_topk_forward` once to generate a single `launch` event
  - Read the only `.ndjson` file under the temp logs directory and assert:
    - Exactly one file
    - Exactly one `launch` event
    - `extracted_args["X"].type == "triton_kernels.tensor.Tensor"`
  - Synchronize CUDA and clean up the temp directory (unless `TEST_KEEP_OUTPUT=1`)

## Backward compatibility
- When the optional package is absent, constructing the custom types now raises a clear `RuntimeError` instructing that the dependency is missing. Callers can catch this or ensure the extra is installed.
- Code paths that do not touch the custom types are unaffected.

## Testing
- Unit test updated: `TestTritonparseCUDA.test_triton_kernels_Tensor`
- How to run (CUDA required):
  - `python -m unittest tests.test_tritonparse -v -k test_triton_kernels_Tensor`
- Expected: test passes, asserting single launch event and correct `X` type in `extracted_args`.

Pull Request resolved: #101

Reviewed By: sfzhu93

Differential Revision: D82361016

Pulled By: FindHao

fbshipit-source-id: 83ccfa40a7f6a21aa3ba377ba18b3cb5390d9116

Author: FindHao

.ci/install-triton-kernels.sh

@@ -0,0 +1,87 @@
+#!/bin/bash
+
+# Install Triton kernels from triton-lang/triton/python/triton_kernels
+
+set -e
+
+echo "🚀 Installing Triton kernels from triton-lang/triton/python/triton_kernels..."
+START_TIME=$(date +%s)
+
+# Function to show elapsed time
+show_elapsed() {
+    CURRENT_TIME=$(date +%s)
+    ELAPSED=$((CURRENT_TIME - START_TIME))
+    echo "⏱️ Elapsed time: ${ELAPSED}s"
+}
+
+# Set Triton version/commit for cache consistency
+TRITON_COMMIT=${TRITON_COMMIT:-"main"}
+echo "🎯 Target Triton commit/branch: $TRITON_COMMIT"
+TRITON_SOURCE_DIR="/tmp/triton"
+
+# Ensure we're in the conda environment
+if [ -z "$CONDA_ENV" ]; then
+    echo "ERROR: CONDA_ENV is not set"
+    exit 1
+fi
+
+# Activate conda environment
+source /opt/miniconda3/etc/profile.d/conda.sh
+conda activate "$CONDA_ENV"
+
+# Ensure TRITON_SOURCE_DIR contains Triton source; otherwise, clone it
+echo "🔧 Ensuring Triton source exists at $TRITON_SOURCE_DIR..."
+
+if [ -d "$TRITON_SOURCE_DIR/.git" ]; then
+    REMOTE_URL=$(git -C "$TRITON_SOURCE_DIR" remote get-url origin 2>/dev/null || echo "")
+    if [[ "$REMOTE_URL" == *"triton-lang/triton"* ]]; then
+        echo "✅ Found existing Triton repository: $REMOTE_URL"
+    else
+        echo "⚠️ Existing directory is not triton-lang/triton (origin: $REMOTE_URL). Re-cloning..."
+        rm -rf "$TRITON_SOURCE_DIR"
+    fi
+fi
+
+if [ ! -d "$TRITON_SOURCE_DIR/.git" ]; then
+    echo "Cloning Triton repository..."
+    if ! git clone https://github.com/triton-lang/triton.git "$TRITON_SOURCE_DIR"; then
+        echo "❌ ERROR: Failed to clone Triton repository"
+        echo "This might be due to network issues or GitHub rate limiting"
+        exit 1
+    fi
+fi
+
+echo "Checking out Triton commit/branch: $TRITON_COMMIT"
+if ! git -C "$TRITON_SOURCE_DIR" checkout "$TRITON_COMMIT"; then
+    echo "❌ ERROR: Failed to checkout $TRITON_COMMIT"
+    exit 1
+fi
+
+# Install triton_kernels in editable mode
+KERNELS_DIR="$TRITON_SOURCE_DIR/python/triton_kernels"
+if [ ! -d "$KERNELS_DIR" ]; then
+    echo "❌ ERROR: triton_kernels directory not found at $KERNELS_DIR"
+    exit 1
+fi
+
+echo "📦 Installing triton_kernels from $KERNELS_DIR (editable)..."
+pip install -e "$KERNELS_DIR"
+show_elapsed
+
+# Verify installation with a simple import
+echo "🔎 Verifying triton_kernels installation..."
+set +e
+KERNELS_IMPORT_OUTPUT=$(python -c "import triton_kernels; import os; print('triton_kernels OK'); print(getattr(triton_kernels, '__file__', 'no_file'))" 2>&1)
+KERNELS_IMPORT_EXITCODE=$?
+set -e
+
+echo "Import exit code: $KERNELS_IMPORT_EXITCODE"
+echo "Import output: $KERNELS_IMPORT_OUTPUT"
+
+if [ $KERNELS_IMPORT_EXITCODE -ne 0 ]; then
+    echo "❌ ERROR: Failed to import triton_kernels"
+    exit 1
+fi
+
+echo "✅ triton_kernels installation verified"
+show_elapsed

.github/workflows/test.yml

@@ -142,7 +142,7 @@ jobs:
                   TRITON_COMMIT: ${{ steps.triton-commit.outputs.commit }}
               run: |
                   bash .ci/install-triton.sh
-
+                  bash .ci/install-triton-kernels.sh
             - name: Install project dependencies
               env:
                   CONDA_ENV: tritonparse
@@ -192,6 +192,12 @@ jobs:
               run: |
                   bash .ci/setup.sh
 
+            - name: Install Triton kernels
+              env:
+                  CONDA_ENV: tritonparse-pip
+              run: |
+                  bash .ci/install-triton-kernels.sh
+
             - name: Install project dependencies
               env:
                   CONDA_ENV: tritonparse-pip

.gitignore

@@ -54,6 +54,7 @@ parsed_output/
 *.ndjson.gz
 *.json
 !tests/example_output/
+!tests/example_output/repro/**
 !tests/example_output/logs/**
 !tests/example_output/parsed_output/**
 !tests/example_output/parsed_output_complex/**