tritonparse 0.2.4.dev20250923071613__tar.gz → 0.2.4.dev20250925071528__tar.gz

{tritonparse-0.2.4.dev20250923071613 → tritonparse-0.2.4.dev20250925071528}/.gitignore

@@ -69,5 +69,5 @@ ENV/
 env.bak/
 venv.bak/
 *.mdc
-
+repro_output/
 # end

{tritonparse-0.2.4.dev20250923071613/tritonparse.egg-info → tritonparse-0.2.4.dev20250925071528}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tritonparse
-Version: 0.2.4.dev20250923071613
+Version: 0.2.4.dev20250925071528
 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.dev20250925071528/run.py

@@ -0,0 +1,48 @@
+#!/usr/bin/env python3
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
+import argparse
+
+from tritonparse.reproducer.cli import _add_reproducer_args
+from tritonparse.reproducer.orchestrator import reproduce
+
+from tritonparse.utils import _add_parse_args, unified_parse
+
+
+# We need this as an entrace for fbpkg
+def main():
+    parser = argparse.ArgumentParser(description="tritonparse CLI")
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    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")
+
+    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":
+        # Filter out routing-specific arguments before passing to unified_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_index,
+            out_dir=args.out_dir,
+            template=args.template,
+        )
+
+
+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.dev20250923071613 → tritonparse-0.2.4.dev20250925071528}/tests/test_tritonparse.py

@@ -33,6 +33,7 @@ from triton.compiler import ASTSource, IRSource
 
 # @manual=//triton:triton
 from triton.knobs import CompileTimes
+from tritonparse.common import is_fbcode
 from tritonparse.structured_logging import convert, extract_python_source_info
 
 HAS_TRITON_KERNELS = importlib.util.find_spec("triton_kernels") is not None
@@ -706,6 +707,122 @@ class TestTritonparseCUDA(unittest.TestCase):
                 print("✓ Cleaned up temporary directory")
             tritonparse.structured_logging.clear_logging_config()
 
+    @unittest.skipIf(is_fbcode(), "Skip in internal FB environment")
+    @unittest.skipUnless(torch.cuda.is_available(), "CUDA not available")
+    def test_reproducer_end_to_end(self):
+        """End-to-end test for reproducer: generate logs, build script, run it."""
+        import subprocess as _subprocess
+        import sys as _sys
+        from pathlib import Path as _Path
+
+        # 1) Prepare temp dirs
+        temp_dir = tempfile.mkdtemp()
+        logs_dir = os.path.join(temp_dir, "logs")
+        out_dir = os.path.join(temp_dir, "repro_output")
+        os.makedirs(logs_dir, exist_ok=True)
+        os.makedirs(out_dir, exist_ok=True)
+
+        # 2) Write a simple module-level Triton kernel to a temp file
+        kernel_dir = os.path.join(temp_dir, "kernels")
+        os.makedirs(kernel_dir, exist_ok=True)
+        kernel_file = os.path.join(kernel_dir, "simple_kernel.py")
+        kernel_src = (
+            "import triton\n"
+            "import triton.language as tl\n"
+            "import torch\n"
+            "\n"
+            "@triton.jit\n"
+            "def add_kernel(x_ptr, y_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr):\n"
+            "    pid = tl.program_id(axis=0)\n"
+            "    block_start = pid * BLOCK_SIZE\n"
+            "    offsets = block_start + tl.arange(0, BLOCK_SIZE)\n"
+            "    mask = offsets < n_elements\n"
+            "    x = tl.load(x_ptr + offsets, mask=mask)\n"
+            "    y = tl.load(y_ptr + offsets, mask=mask)\n"
+            "    tl.store(out_ptr + offsets, x + y, mask=mask)\n"
+        )
+        with open(kernel_file, "w", encoding="utf-8") as f:
+            f.write(kernel_src)
+
+        # 3) Generate logs by running the kernel once
+        tritonparse.structured_logging.init(logs_dir, enable_trace_launch=True)
+        try:
+            if kernel_dir not in _sys.path:
+                _sys.path.insert(0, kernel_dir)
+            import importlib as _importlib
+
+            mod = _importlib.import_module("simple_kernel")
+            device = torch.device("cuda:0")
+            torch.manual_seed(0)
+            n = 256
+            x = torch.randn((n,), device=device, dtype=torch.float32)
+            y = torch.randn((n,), device=device, dtype=torch.float32)
+            out = torch.empty_like(x)
+            BLOCK_SIZE = 64
+            grid = (triton.cdiv(n, BLOCK_SIZE),)
+            mod.add_kernel[grid](x, y, out, n, BLOCK_SIZE)
+            torch.cuda.synchronize()
+        finally:
+            tritonparse.structured_logging.clear_logging_config()
+
+        # 4) Find the NDJSON and compute launch event index
+        ndjson_files = [
+            os.path.join(logs_dir, f)
+            for f in os.listdir(logs_dir)
+            if f.endswith(".ndjson")
+        ]
+        assert ndjson_files, f"No ndjson found in {logs_dir}"
+        ndjson_path = max(ndjson_files, key=os.path.getmtime)
+
+        from tritonparse.tools.prettify_ndjson import load_ndjson as _load_ndjson
+
+        events = _load_ndjson(_Path(ndjson_path), save_irs=True)
+        launch_indices = [
+            i for i, ev in enumerate(events) if ev.get("event_type") == "launch"
+        ]
+        assert launch_indices, "No launch event found in ndjson"
+        line_index = launch_indices[0]
+
+        # 5) Build reproducer
+        from tritonparse.reproducer.orchestrator import reproduce
+
+        reproduce(
+            input_path=ndjson_path,
+            line_index=line_index,
+            out_dir=out_dir,
+            template="example",
+        )
+
+        # 6) Locate generated script and context under out_dir/add_kernel/
+        kernel_out_dir = os.path.join(out_dir, "add_kernel")
+        assert os.path.isdir(
+            kernel_out_dir
+        ), f"Kernel output dir not found: {kernel_out_dir}"
+        gen_scripts = [f for f in os.listdir(kernel_out_dir) if f.endswith(".py")]
+        gen_jsons = [f for f in os.listdir(kernel_out_dir) if f.endswith(".json")]
+        assert gen_scripts, f"No generated script in {kernel_out_dir}"
+        assert gen_jsons, f"No generated context json in {kernel_out_dir}"
+        script_path = os.path.join(kernel_out_dir, sorted(gen_scripts)[-1])
+
+        # 7) Execute generated script and assert success output
+        proc = _subprocess.run(
+            [
+                _sys.executable,
+                script_path,
+            ],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+        self.assertIn("Kernel execution finished.", proc.stdout)
+
+        # Cleanup
+        if should_keep_output():
+            print(f"✓ Preserving temporary directory (TEST_KEEP_OUTPUT=1): {temp_dir}")
+        else:
+            shutil.rmtree(temp_dir)
+            print("✓ Cleaned up temporary directory")
+
 
 if __name__ == "__main__":
     unittest.main()

tritonparse-0.2.4.dev20250925071528/tritonparse/reproducer/cli.py

@@ -0,0 +1,27 @@
+import argparse
+
+
+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",
+        type=int,
+        help="The line number of the launch event in the input file to reproduce.",
+    )
+    parser.add_argument(
+        "--out-dir",
+        default="repro_output",
+        help=(
+            "Directory to save the reproducer script and context JSON. Defaults to "
+            "'repro_output/<kernel_name>/' if not provided."
+        ),
+    )
+    parser.add_argument(
+        "--template",
+        default="example",
+        help=(
+            "Template name (builtin, without .py) or a filesystem path to a .py file. "
+            "Defaults to 'example'."
+        ),
+    )

tritonparse-0.2.4.dev20250925071528/tritonparse/reproducer/ingestion/ndjson.py

@@ -0,0 +1,235 @@
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple
+
+from tritonparse.tp_logger import logger
+
+# Sentinel object to mark arguments that should be skipped during processing
+_SKIP = object()
+
+
+@dataclass
+class KernelInfo:
+    """Information about a Triton kernel extracted from compilation events."""
+
+    file_path: str
+    function_name: str
+    source_code: str
+    call_stack: List[Dict[str, Any]]
+
+
+@dataclass
+class ContextBundle:
+    """Bundle of all context information needed to reproduce a kernel launch."""
+
+    kernel_info: KernelInfo
+    compile: Dict[str, Any]
+    launch: Dict[str, Any]
+    args: Dict[str, Any]
+    tensor_args: Dict[str, Any]
+    raw_launch_event: Dict[str, Any]
+    raw_comp_event: Dict[str, Any]
+
+
+def get_launch_and_compilation_events(
+    events: List[Dict[str, Any]], line_index: Optional[int] = None
+) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+    """
+    Extract launch and compilation events from the event list.
+
+    Args:
+        events: List of parsed event dictionaries.
+        line_index: Index of the launch event to process.
+
+    Returns:
+        Tuple of (launch_event, compilation_event).
+
+    Raises:
+        ValueError: If the event at line_index is not a launch event.
+        RuntimeError: If compilation event cannot be found or is ambiguous.
+    """
+    if line_index is None or line_index >= len(events):
+        raise ValueError(f"Invalid line_index: {line_index}")
+
+    launch_event = events[line_index]
+    if launch_event["event_type"] != "launch":
+        raise ValueError(f"Event at index {line_index} is not a launch event")
+
+    comp_meta = launch_event.get("compilation_metadata", {})
+    comp_hash = comp_meta.get("hash")
+    if not comp_hash:
+        raise RuntimeError("Could not find compilation hash in launch event.")
+
+    comp_event = None
+    for event in events:
+        if (
+            event["event_type"] == "compilation"
+            and event.get("payload", {}).get("metadata", {}).get("hash") == comp_hash
+        ):
+            comp_event = event
+            break
+    if not comp_event:
+        raise RuntimeError(f"Could not find compilation event for hash {comp_hash}.")
+    return launch_event, comp_event
+
+
+def get_kernel_info(comp_event: Dict[str, Any]) -> KernelInfo:
+    """
+    Extract kernel information from a compilation event.
+
+    Args:
+        comp_event: Compilation event dictionary containing kernel metadata.
+
+    Returns:
+        KernelInfo object with extracted kernel details.
+
+    Raises:
+        RuntimeError: If file path or function name cannot be resolved.
+    """
+    payload = comp_event.get("payload") or {}
+    py_source = payload.get("python_source") or {}
+    code = py_source.get("code", "")
+
+    # Extract file path and function name
+    file_path = py_source.get("file_path")
+    # The function name is in the compilation metadata payload
+    func_name = (comp_event.get("payload", {}).get("metadata") or {}).get("name")
+
+    # Find '@triton.jit' decorator and slice the string from there
+    jit_marker = "@triton.jit"
+    jit_pos = code.find(jit_marker)
+    if jit_pos != -1:
+        code = code[jit_pos:]
+        logger.debug("Extracted kernel source starting from '@triton.jit'.")
+
+    if not file_path or not func_name:
+        raise RuntimeError(
+            "Could not resolve kernel file path or function name from compilation event."
+            " The import-based strategy cannot proceed."
+        )
+    return KernelInfo(file_path, func_name, code, comp_event.get("stack", []))
+
+
+def _decode_arg(raw: Any) -> Any:
+    """
+    Decode a raw argument value from event data.
+
+    Args:
+        raw: Raw argument value from event data.
+
+    Returns:
+        Decoded argument value, or _SKIP sentinel for tensors.
+    """
+    if not isinstance(raw, dict):
+        return raw
+    t = raw.get("type")
+    if t == "tensor":
+        return _SKIP
+    if t == "NoneType":
+        return None
+    return raw.get("value", raw.get("repr"))
+
+
+def _pack_args(args: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Pack argument values into a standardized format.
+
+    Args:
+        args: Dictionary of argument names to values.
+
+    Returns:
+        Dictionary with packed argument information including type and metadata.
+    """
+    packed = {}
+    for k, v in args.items():
+        t = v.get("type") if isinstance(v, dict) else None
+        if t == "tensor":
+            packed[k] = {
+                "type": "tensor",
+                "shape": v.get("shape") if isinstance(v, dict) else None,
+                "dtype": v.get("dtype") if isinstance(v, dict) else None,
+                "device": v.get("device") if isinstance(v, dict) else None,
+                "stride": v.get("stride") if isinstance(v, dict) else None,
+                "is_contiguous": (
+                    v.get("is_contiguous") if isinstance(v, dict) else None
+                ),
+                "numel": v.get("numel") if isinstance(v, dict) else None,
+            }
+        else:
+            # scalar / NoneType etc
+            if isinstance(v, dict):
+                packed[k] = {
+                    "type": v.get("type"),
+                    "value": v.get("value", v.get("repr")),
+                }
+            else:
+                packed[k] = {
+                    "type": None,
+                    "value": v,
+                }
+    return packed
+
+
+def build_context_bundle(
+    events: List[Dict[str, Any]], line_index: Optional[int] = None
+):
+    """
+    Build a complete context bundle from events and line index.
+
+    Args:
+        events: List of parsed event dictionaries.
+        line_index: Index of the launch event to process.
+
+    Returns:
+        ContextBundle containing all information needed to reproduce the kernel launch.
+
+    Raises:
+        ValueError: If line_index is invalid or event is not a launch event.
+        RuntimeError: If compilation event cannot be found.
+    """
+    launch_event, comp_event = get_launch_and_compilation_events(events, line_index)
+    kernel_info = get_kernel_info(comp_event)
+    grid = launch_event.get("grid")
+    extracted_args = launch_event.get("extracted_args", {})
+    comp_meta = launch_event.get("compilation_metadata", {})
+
+    # Compile metadata subset we care about
+    compile_block = {
+        "num_warps": comp_meta.get("num_warps"),
+        "num_stages": comp_meta.get("num_stages"),
+        "arch": comp_meta.get("arch"),
+        "backend": comp_meta.get("backend_name") or comp_meta.get("backend"),
+        "triton_version": comp_meta.get("triton_version"),
+        "hash": comp_meta.get("hash"),
+    }
+
+    # kwargs: include constexpr + explicit scalars used for launch (skip tensor args)
+    kwargs = {}
+    for k, v in extracted_args.items():
+        val = _decode_arg(v)
+        if val is _SKIP:
+            continue
+        kwargs[k] = val
+
+    # tensor args: only tensors
+    raw_tensor_args = {
+        k: v
+        for k, v in extracted_args.items()
+        if isinstance(v, dict) and v.get("type") == "tensor"
+    }
+
+    primitive_args = _pack_args(extracted_args)
+    tensor_args = _pack_args(raw_tensor_args)
+    launch_block = {
+        "grid": grid,
+        "kwargs": kwargs,
+    }
+
+    return ContextBundle(
+        kernel_info,
+        compile_block,
+        launch_block,
+        primitive_args,
+        tensor_args,
+        launch_event,
+        comp_event,
+    )

tritonparse-0.2.4.dev20250925071528/tritonparse/reproducer/orchestrator.py

@@ -0,0 +1,63 @@
+from pathlib import Path
+
+from tritonparse.reproducer.ingestion.ndjson import build_context_bundle
+from tritonparse.reproducer.templates.loader import load_template_code
+from tritonparse.reproducer.utils import (
+    _generate_import_statements,
+    _generate_invocation_snippet,
+    _parse_kernel_signature,
+    determine_output_paths,
+)
+
+from tritonparse.tools.prettify_ndjson import load_ndjson, save_prettified_json
+from tritonparse.tp_logger import logger
+
+
+def reproduce(
+    input_path: str,
+    line_index: int,
+    out_dir: str,
+    template: str,
+):
+    """
+    Generate a reproducer script from NDJSON trace file.
+
+    Args:
+        input_path: Path to the NDJSON trace file.
+        line_index: Line index of the launch event to reproduce.
+        out_dir: Output directory for reproducer files.
+    """
+    logger.debug(f"Building bundle from {input_path} at line {line_index}")
+    events = load_ndjson(Path(input_path), save_irs=True)
+    logger.debug(f"Loaded {len(events)} events")
+
+    # Build context bundle from the specified launch event
+    context_bundle = build_context_bundle(events, line_index)
+    logger.debug(
+        f"Built context bundle for kernel: {context_bundle.kernel_info.function_name}"
+    )
+    out_py_path, temp_json_path = determine_output_paths(
+        out_dir, context_bundle.kernel_info.function_name
+    )
+    save_prettified_json(context_bundle.raw_launch_event, temp_json_path)
+    logger.debug("Loading reproducer template.")
+    template_code = load_template_code(template)
+    final_code = template_code.replace(
+        "{{JSON_FILE_NAME_PLACEHOLDER}}", temp_json_path.name
+    )
+    sys_stmt, import_statement = _generate_import_statements(context_bundle.kernel_info)
+    final_code = final_code.replace("# {{KERNEL_SYSPATH_PLACEHOLDER}}", sys_stmt)
+    final_code = final_code.replace("# {{KERNEL_IMPORT_PLACEHOLDER}}", import_statement)
+    source_code = context_bundle.kernel_info.source_code
+    pos_args, kw_args = _parse_kernel_signature(source_code)
+    invocation_snippet = _generate_invocation_snippet(pos_args, kw_args)
+    final_code = final_code.replace(
+        "# {{KERNEL_INVOCATION_PLACEHOLDER}}", invocation_snippet
+    )
+    out_py_path.write_text(final_code, encoding="utf-8")
+    logger.info(
+        "REPRODUCER_OUTPUT script=%s json=%s kernel=%s",
+        str(out_py_path.resolve()),
+        str(temp_json_path.resolve()),
+        context_bundle.kernel_info.function_name,
+    )