tritonparse 0.2.4.dev20251001071529__tar.gz → 0.2.4.dev20251002071459__tar.gz

{tritonparse-0.2.4.dev20251001071529/tritonparse.egg-info → tritonparse-0.2.4.dev20251002071459}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tritonparse
-Version: 0.2.4.dev20251001071529
+Version: 0.2.4.dev20251002071459
 Summary: TritonParse: A Compiler Tracer, Visualizer, and mini-Reproducer Generator for Triton Kernels
 Author-email: Yueming Hao <yhao@meta.com>
 License-Expression: BSD-3-Clause

{tritonparse-0.2.4.dev20251001071529 → tritonparse-0.2.4.dev20251002071459}/pyproject.toml

@@ -46,3 +46,6 @@ local_scheme = "no-local-version"
 
 [project.urls]
 "Homepage" = "https://github.com/meta-pytorch/tritonparse"
+
+[project.scripts]
+tritonparse = "tritonparse.run:main"

tritonparse-0.2.4.dev20251002071459/run.py

@@ -0,0 +1,80 @@
+#!/usr/bin/env python3
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
+import argparse
+from importlib.metadata import PackageNotFoundError, version
+
+from .reproducer.cli import _add_reproducer_args
+from .reproducer.orchestrator import reproduce
+from .utils import _add_parse_args, unified_parse
+
+
+def _get_package_version() -> str:
+    try:
+        return version("tritonparse")
+    except PackageNotFoundError:
+        return "0+unknown"
+
+
+def main():
+    pkg_version = _get_package_version()
+
+    parser = argparse.ArgumentParser(
+        prog="tritonparse",
+        description=(
+            "TritonParse: parse structured logs and generate minimal reproducers"
+        ),
+        epilog=(
+            "Examples:\n"
+            "  tritonparse parse /path/to/logs --out parsed_output\n"
+            "  tritonparse reproduce /path/to/trace.ndjson --line 1 --out-dir repro_output\n"
+        ),
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument(
+        "--version",
+        action="version",
+        version=f"%(prog)s {pkg_version}",
+        help="Show program's version number and exit",
+    )
+
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    # parse subcommand
+    parse_parser = subparsers.add_parser(
+        "parse",
+        help="Parse triton structured logs",
+        conflict_handler="resolve",
+    )
+    _add_parse_args(parse_parser)
+    parse_parser.set_defaults(func="parse")
+
+    # reproduce subcommand
+    repro_parser = subparsers.add_parser(
+        "reproduce",
+        help="Build reproducer from trace file",
+    )
+    _add_reproducer_args(repro_parser)
+    repro_parser.set_defaults(func="reproduce")
+
+    args = parser.parse_args()
+
+    if args.func == "parse":
+        parse_args = {
+            k: v for k, v in vars(args).items() if k not in ["command", "func"]
+        }
+        unified_parse(**parse_args)
+    elif args.func == "reproduce":
+        reproduce(
+            input_path=args.input,
+            line_index=args.line,
+            out_dir=args.out_dir,
+            template=args.template,
+        )
+    else:
+        raise RuntimeError(f"Unknown command: {args.func}")
+
+
+if __name__ == "__main__":
+    # Do not add code here, it won't be run. Add them to the function called below.
+    main()  # pragma: no cover

{tritonparse-0.2.4.dev20251001071529 → tritonparse-0.2.4.dev20251002071459}/tests/test_tritonparse.py

@@ -19,22 +19,20 @@ from typing import Any, Union
 
 import torch
 
-# @manual=//triton:triton
-import triton
+import triton  # @manual=//triton:triton
 
-# @manual=//triton:triton
-import triton.language as tl
+import triton.language as tl  # @manual=//triton:triton
 
 import tritonparse.structured_logging
+import tritonparse.tools.disasm
 import tritonparse.utils
 
-# @manual=//triton:triton
-from triton.compiler import ASTSource, IRSource
+from triton.compiler import ASTSource, IRSource  # @manual=//triton:triton
 
-# @manual=//triton:triton
-from triton.knobs import CompileTimes
+from triton.knobs import CompileTimes  # @manual=//triton:triton
 from tritonparse.common import is_fbcode
 from tritonparse.structured_logging import convert, extract_python_source_info
+from tritonparse.tools.disasm import is_nvdisasm_available
 
 HAS_TRITON_KERNELS = importlib.util.find_spec("triton_kernels") is not None
 
@@ -181,7 +179,7 @@ class TestTritonparseCUDA(unittest.TestCase):
 
     @unittest.skipUnless(torch.cuda.is_available(), "CUDA not available")
     def test_whole_workflow(self):
-        """Test unified_parse functionality"""
+        """Test unified_parse functionality including SASS extraction"""
 
         # Define a simple kernel directly in the test function
         @triton.jit
@@ -211,9 +209,18 @@ class TestTritonparseCUDA(unittest.TestCase):
         os.makedirs(temp_dir_logs, exist_ok=True)
         os.makedirs(temp_dir_parsed, exist_ok=True)
         print(f"Temporary directory: {temp_dir}")
+        nvdisasm_available = is_nvdisasm_available()
+        if nvdisasm_available:
+            print("✓ nvdisasm tool is available, enabling SASS dumping")
+        else:
+            print("⚠️  nvdisasm tool not available, SASS dumping will be disabled")
 
-        # Initialize logging
-        tritonparse.structured_logging.init(temp_dir_logs, enable_trace_launch=True)
+        # Initialize logging with conditional SASS dumping
+        tritonparse.structured_logging.init(
+            temp_dir_logs,
+            enable_trace_launch=True,
+            enable_sass_dump=nvdisasm_available,
+        )
 
         # Generate test data and run kernels
         torch.manual_seed(0)
@@ -237,8 +244,8 @@ class TestTritonparseCUDA(unittest.TestCase):
         print(f"Found {len(log_files)} log files in {temp_dir_logs}: {log_files}")
 
         def check_event_type_counts_in_logs(log_dir: str) -> dict:
-            """Count 'launch' and unique 'compilation' events in all log files"""
-            event_counts = {"launch": 0}
+            """Count 'launch' and unique 'compilation' events in all log files and verify SASS content"""
+            event_counts = {"launch": 0, "sass_found": False}
             # Track unique compilation hashes
             compilation_hashes = set()
 
@@ -267,6 +274,43 @@ class TestTritonparseCUDA(unittest.TestCase):
                                         print(
                                             f"  Line {line_num}: event_type = 'compilation' (unique hash: {compilation_hash[:8]}...)"
                                         )
+
+                                    # Check for SASS content in compilation events
+                                    file_content = event_data.get("payload", {}).get(
+                                        "file_content", {}
+                                    )
+                                    sass_files = [
+                                        key
+                                        for key in file_content.keys()
+                                        if key.endswith(".sass")
+                                    ]
+
+                                    if sass_files and not event_counts["sass_found"]:
+                                        event_counts["sass_found"] = True
+                                        sass_content = file_content[sass_files[0]]
+                                        print(f"✓ Found SASS file: {sass_files[0]}")
+                                        print(
+                                            f"  SASS content preview (first 200 chars): {sass_content[:200]}..."
+                                        )
+
+                                        # Verify SASS content looks like assembly
+                                        assert (
+                                            "Function:" in sass_content
+                                        ), "SASS content should contain function declaration"
+                                        # Basic check for NVIDIA GPU assembly patterns
+                                        assert any(
+                                            pattern in sass_content.lower()
+                                            for pattern in [
+                                                "mov",
+                                                "add",
+                                                "mul",
+                                                "ld",
+                                                "st",
+                                                "lop",
+                                                "s2r",
+                                            ]
+                                        ), "SASS content should contain GPU assembly instructions"
+
                             except (json.JSONDecodeError, KeyError, TypeError) as e:
                                 print(f"  Line {line_num}: Error processing line - {e}")
 
@@ -277,7 +321,7 @@ class TestTritonparseCUDA(unittest.TestCase):
             )
             return event_counts
 
-        # Verify event counts
+        # Verify event counts and conditional SASS extraction
         event_counts = check_event_type_counts_in_logs(temp_dir_logs)
         assert (
             event_counts["compilation"] == 1
@@ -285,6 +329,16 @@ class TestTritonparseCUDA(unittest.TestCase):
         assert (
             event_counts["launch"] == 2
         ), f"Expected 2 'launch' events, found {event_counts['launch']}"
+
+        # Conditionally verify SASS content based on nvdisasm availability
+        if nvdisasm_available:
+            assert event_counts[
+                "sass_found"
+            ], "SASS content was not found in compilation events"
+            print("✓ Successfully verified SASS extraction functionality")
+        else:
+            print("⚠️  SASS verification skipped: nvdisasm not available")
+
         print(
             "✓ Verified correct event type counts: 1 unique compilation hash, 2 launch events"
         )
@@ -297,6 +351,49 @@ class TestTritonparseCUDA(unittest.TestCase):
             # Verify parsing output
             parsed_files = os.listdir(temp_dir_parsed)
             assert len(parsed_files) > 0, "No files found in parsed output directory"
+
+            # Verify that SASS is preserved in parsed output
+            ndjson_gz_files = [f for f in parsed_files if f.endswith(".ndjson.gz")]
+            assert (
+                len(ndjson_gz_files) > 0
+            ), "No .ndjson.gz files found in parsed output"
+
+            sass_found_in_parsed = False
+            for ndjson_gz_file in ndjson_gz_files:
+                ndjson_gz_path = os.path.join(temp_dir_parsed, ndjson_gz_file)
+                with gzip.open(ndjson_gz_path, "rt", encoding="utf-8") as f:
+                    for line in f:
+                        try:
+                            event_data = json.loads(line.strip())
+                            if event_data.get("event_type") == "compilation":
+                                file_content = event_data.get("payload", {}).get(
+                                    "file_content", {}
+                                )
+                                sass_files = [
+                                    key
+                                    for key in file_content.keys()
+                                    if key.endswith(".sass")
+                                ]
+                                if sass_files:
+                                    sass_found_in_parsed = True
+                                    print("✓ SASS content preserved in parsed output")
+                                    break
+                        except json.JSONDecodeError:
+                            continue
+
+                if sass_found_in_parsed:
+                    break
+
+            # Conditionally verify SASS content is preserved in parsed output
+            if nvdisasm_available:
+                assert (
+                    sass_found_in_parsed
+                ), "SASS content was not preserved in parsed output"
+            else:
+                print(
+                    "⚠️  SASS preservation verification skipped: nvdisasm not available"
+                )
+
         finally:
             # Clean up
             if should_keep_output():
@@ -745,7 +842,9 @@ class TestTritonparseCUDA(unittest.TestCase):
             f.write(kernel_src)
 
         # 3) Generate logs by running the kernel once
-        tritonparse.structured_logging.init(logs_dir, enable_trace_launch=True)
+        tritonparse.structured_logging.init(
+            logs_dir, enable_trace_launch=True, enable_more_tensor_information=True
+        )
         try:
             if kernel_dir not in _sys.path:
                 _sys.path.insert(0, kernel_dir)

tritonparse-0.2.4.dev20251002071459/tritonparse/__main__.py

@@ -0,0 +1,5 @@
+from .run import main
+
+
+if __name__ == "__main__":
+    main()

{tritonparse-0.2.4.dev20251001071529 → tritonparse-0.2.4.dev20251002071459}/tritonparse/reproducer/cli.py

@@ -5,9 +5,13 @@ def _add_reproducer_args(parser: argparse.ArgumentParser) -> None:
     """Add common arguments for the reproducer to a parser."""
     parser.add_argument("input", help="Path to the ndjson/ndjson.gz log file")
     parser.add_argument(
-        "--line-index",
+        "--line",
         type=int,
-        help="The line number of the launch event in the input file to reproduce.",
+        default=1,
+        help=(
+            "The line number (1-based) of the launch event in the input file to reproduce. "
+            "Defaults to 1."
+        ),
     )
     parser.add_argument(
         "--out-dir",

{tritonparse-0.2.4.dev20251001071529 → tritonparse-0.2.4.dev20251002071459}/tritonparse/reproducer/templates/example.py

@@ -1,6 +1,12 @@
+"""
+This file is automatically generated by TritonParse reproducer.
+It contains a smallest testing example for a Triton kernel.
+"""
+
 import hashlib
 import importlib
 import json
+import logging
 import sys
 from functools import lru_cache
 from pathlib import Path
@@ -139,25 +145,62 @@ def _create_arg_from_info(arg_info):
     elif arg_type == "tensor":
         if arg_info.get("blob_path"):
             return load_tensor(arg_info.get("blob_path"), arg_info.get("device"))
+
+        # Extract basic tensor properties
         dtype_str = arg_info.get("dtype")
         try:
             torch_dtype = getattr(torch, dtype_str.split(".")[-1])
         except AttributeError:
+            logging.error(f"Unsupported dtype: {dtype_str}. Defaulting to float32.")
             torch_dtype = torch.float32
 
         shape = arg_info.get("shape", [])
         device = arg_info.get("device", "cpu")
 
+        # Extract statistical information if available
+        mean = arg_info.get("mean")
+        std = arg_info.get("std")
+        min_val = arg_info.get("min")
+        max_val = arg_info.get("max")
+        has_stats = (
+            mean is not None
+            and std is not None
+            and min_val is not None
+            and max_val is not None
+        )
+
+        if arg_info.get("tensor_capture_error", False):
+            logging.error(
+                f"Error: Tensor '{arg_info.get('name', '')}' had capture error. Generating random tensor instead."
+            )
+
         # Use a dummy tensor to check properties of the dtype
         tensor_props = torch.empty(0, dtype=torch_dtype)
 
-        # Case 1: Floating point, signed integers, uint8, and bool are supported by random_()
+        # Case 1: Floating point types
         if tensor_props.is_floating_point():
-            if torch_dtype in [torch.float8_e4m3fn, torch.float8_e5m2]:
-                tmp = torch.rand(shape, dtype=torch.float32, device=device)
-                return tmp.to(torch_dtype)
+            if has_stats:
+                # Generate tensor with statistical properties matching original data
+                if std == 0 or min_val == max_val:
+                    # Constant tensor
+                    return torch.full(shape, mean, dtype=torch_dtype, device=device)
+                # Generate normal distribution with mean and std, then clamp to [min, max]
+                tensor = (
+                    torch.randn(shape, dtype=torch.float32, device=device) * std + mean
+                )
+                tensor = torch.clamp(tensor, min=min_val, max=max_val)
+                return tensor.to(torch_dtype)
             else:
-                return torch.empty(shape, dtype=torch_dtype, device=device).random_()
+                # Fallback to original random generation
+                if torch_dtype in [torch.float8_e4m3fn, torch.float8_e5m2]:
+                    tmp = torch.rand(shape, dtype=torch.float32, device=device)
+                    return tmp.to(torch_dtype)
+                else:
+                    return torch.empty(
+                        shape, dtype=torch_dtype, device=device
+                    ).random_()
+
+        # Case 2: Integer types
         elif torch_dtype in [
             torch.int8,
             torch.int16,
@@ -166,9 +209,26 @@ def _create_arg_from_info(arg_info):
             torch.uint8,
             torch.bool,
         ]:
-            return torch.empty(shape, dtype=torch_dtype, device=device).random_()
-        # Case 2: Complex numbers need special handling
+            if has_stats and torch_dtype != torch.bool:
+                # Generate tensor with statistical properties, then round for integers
+                if std == 0 or min_val == max_val:
+                    # Constant tensor
+                    return torch.full(
+                        shape, int(mean), dtype=torch_dtype, device=device
+                    )
+                tensor = (
+                    torch.randn(shape, dtype=torch.float32, device=device) * std + mean
+                )
+                tensor = torch.clamp(tensor, min=min_val, max=max_val)
+                return torch.round(tensor).to(torch_dtype)
+            else:
+                # Fallback to original random generation
+                return torch.empty(shape, dtype=torch_dtype, device=device).random_()
+
+        # Case 3: Complex numbers need special handling
         elif tensor_props.is_complex():
+            # Complex types: fallback to original logic for now
+            # TODO: Could be improved to use statistical info if available
             float_dtype = (
                 torch.float32 if torch_dtype == torch.complex64 else torch.float64
             )
@@ -176,10 +236,24 @@ def _create_arg_from_info(arg_info):
             imag_part = torch.rand(shape, dtype=float_dtype, device=device)
             return torch.complex(real_part, imag_part)
 
-        # Case 3: Handle other unsigned integers (like uint32) which fail with random_()
+        # Case 4: Handle other unsigned integers (like uint32) which fail with random_()
         elif "uint" in str(torch_dtype):
-            return torch.randint(0, 1000, shape, dtype=torch_dtype, device=device)
-        # Case 4: If we don't know how to handle the type, raise an error
+            if has_stats:
+                # Generate tensor with statistical properties for unsigned integers
+                if std == 0 or min_val == max_val:
+                    return torch.full(
+                        shape, int(mean), dtype=torch_dtype, device=device
+                    )
+                tensor = (
+                    torch.randn(shape, dtype=torch.float32, device=device) * std + mean
+                )
+                tensor = torch.clamp(tensor, min=min_val, max=max_val)
+                return torch.round(tensor).to(torch_dtype)
+            else:
+                # Fallback to original random generation
+                return torch.randint(0, 1000, shape, dtype=torch_dtype, device=device)
+
+        # Case 5: If we don't know how to handle the type, raise an error
         else:
             raise NotImplementedError(
                 f"Random data generation not implemented for dtype: {torch_dtype}"

{tritonparse-0.2.4.dev20251001071529 → tritonparse-0.2.4.dev20251002071459}/tritonparse/reproducer/utils.py

@@ -84,25 +84,57 @@ def _create_arg_from_info(arg_info):
     elif arg_type == "tensor":
         if arg_info.get("blob_path"):
             return load_tensor(arg_info.get("blob_path"), arg_info.get("device"))
+
+        # Extract basic tensor properties
         dtype_str = arg_info.get("dtype")
         try:
             torch_dtype = getattr(torch, dtype_str.split(".")[-1])
         except AttributeError:
+            logger.error(f"Unsupported dtype: {dtype_str}. Defaulting to float32.")
             torch_dtype = torch.float32
 
         shape = arg_info.get("shape", [])
         device = arg_info.get("device", "cpu")
 
+        # Extract statistical information if available
+        mean = arg_info.get("mean")
+        std = arg_info.get("std")
+        min_val = arg_info.get("min")
+        max_val = arg_info.get("max")
+        has_stats = (
+            mean is not None
+            and std is not None
+            and min_val is not None
+            and max_val is not None
+        )
+
         # Use a dummy tensor to check properties of the dtype
         tensor_props = torch.empty(0, dtype=torch_dtype)
 
-        # Case 1: Floating point, signed integers, uint8, and bool are supported by random_()
+        # Case 1: Floating point types
         if tensor_props.is_floating_point():
-            if torch_dtype in [torch.float8_e4m3fn, torch.float8_e5m2]:
-                tmp = torch.rand(shape, dtype=torch.float32, device=device)
-                return tmp.to(torch_dtype)
+            if has_stats:
+                # Generate tensor with statistical properties matching original data
+                if std == 0 or min_val == max_val:
+                    # Constant tensor
+                    return torch.full(shape, mean, dtype=torch_dtype, device=device)
+                # Generate normal distribution with mean and std, then clamp to [min, max]
+                tensor = (
+                    torch.randn(shape, dtype=torch.float32, device=device) * std + mean
+                )
+                tensor = torch.clamp(tensor, min=min_val, max=max_val)
+                return tensor.to(torch_dtype)
             else:
-                return torch.empty(shape, dtype=torch_dtype, device=device).random_()
+                # Fallback to original random generation
+                if torch_dtype in [torch.float8_e4m3fn, torch.float8_e5m2]:
+                    tmp = torch.rand(shape, dtype=torch.float32, device=device)
+                    return tmp.to(torch_dtype)
+                else:
+                    return torch.empty(
+                        shape, dtype=torch_dtype, device=device
+                    ).random_()
+
+        # Case 2: Integer types
         elif torch_dtype in [
             torch.int8,
             torch.int16,
@@ -111,9 +143,26 @@ def _create_arg_from_info(arg_info):
             torch.uint8,
             torch.bool,
         ]:
-            return torch.empty(shape, dtype=torch_dtype, device=device).random_()
-        # Case 2: Complex numbers need special handling
+            if has_stats and torch_dtype != torch.bool:
+                # Generate tensor with statistical properties, then round for integers
+                if std == 0 or min_val == max_val:
+                    # Constant tensor
+                    return torch.full(
+                        shape, int(mean), dtype=torch_dtype, device=device
+                    )
+                tensor = (
+                    torch.randn(shape, dtype=torch.float32, device=device) * std + mean
+                )
+                tensor = torch.clamp(tensor, min=min_val, max=max_val)
+                return torch.round(tensor).to(torch_dtype)
+            else:
+                # Fallback to original random generation
+                return torch.empty(shape, dtype=torch_dtype, device=device).random_()
+
+        # Case 3: Complex numbers need special handling
         elif tensor_props.is_complex():
+            # Complex types: fallback to original logic for now
+            # TODO: Could be improved to use statistical info if available
             float_dtype = (
                 torch.float32 if torch_dtype == torch.complex64 else torch.float64
             )
@@ -121,14 +170,29 @@ def _create_arg_from_info(arg_info):
             imag_part = torch.rand(shape, dtype=float_dtype, device=device)
             return torch.complex(real_part, imag_part)
 
-        # Case 3: Handle other unsigned integers (like uint32) which fail with random_()
+        # Case 4: Handle other unsigned integers (like uint32) which fail with random_()
         elif "uint" in str(torch_dtype):
-            return torch.randint(0, 1000, shape, dtype=torch_dtype, device=device)
-            # Case 4: If we don't know how to handle the type, raise an error
+            if has_stats:
+                # Generate tensor with statistical properties for unsigned integers
+                if std == 0 or min_val == max_val:
+                    return torch.full(
+                        shape, int(mean), dtype=torch_dtype, device=device
+                    )
+                tensor = (
+                    torch.randn(shape, dtype=torch.float32, device=device) * std + mean
+                )
+                tensor = torch.clamp(tensor, min=min_val, max=max_val)
+                return torch.round(tensor).to(torch_dtype)
+            else:
+                # Fallback to original random generation
+                return torch.randint(0, 1000, shape, dtype=torch_dtype, device=device)
+
+        # Case 5: If we don't know how to handle the type, raise an error
         else:
             raise NotImplementedError(
                 f"Random data generation not implemented for dtype: {torch_dtype}"
             )
+
     elif arg_type == "triton_kernels.tensor.Tensor":
         if not TRITON_KERNELS_CUSTOM_TYPES:
             raise RuntimeError(
@@ -162,7 +226,6 @@ def _create_arg_from_info(arg_info):
             )
         Tensor, Storage, StridedLayout = _get_triton_tensor_types()
         return StridedLayout(shape=arg_info.get("initial_shape"))
-
     else:
         print(f"Warning: Unhandled argument type '{arg_type}'. Returning None.")
         return None