tritonparse 0.2.4.dev20250922071528__py3-none-any.whl → 0.2.4.dev20250924071525__py3-none-any.whl

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/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/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,
+    )

tritonparse/reproducer/templates/example.py

@@ -0,0 +1,239 @@
+import hashlib
+import importlib
+import json
+import sys
+from functools import lru_cache
+from pathlib import Path
+
+import torch
+
+# {{KERNEL_SYSPATH_PLACEHOLDER}}
+
+# {{KERNEL_IMPORT_PLACEHOLDER}}
+
+TRITON_KERNELS_CUSTOM_TYPES = (
+    importlib.util.find_spec("triton_kernels") is not None
+    and importlib.util.find_spec("triton_kernels.tensor") is not None
+)
+
+
+@lru_cache(maxsize=1)
+def _get_triton_tensor_types():
+    """
+    Import and cache Triton custom tensor types.
+
+    Returns:
+        tuple: (Tensor, Storage, StridedLayout) classes from triton_kernels.tensor.
+
+    Raises:
+        ImportError: If the optional module 'triton_kernels.tensor' is not available.
+    """
+    mod = importlib.import_module("triton_kernels.tensor")
+    return (
+        mod.Tensor,
+        mod.Storage,
+        mod.StridedLayout,
+    )
+
+
+def load_tensor(tensor_file_path: str, device: str = None) -> torch.Tensor:
+    """
+    Load a tensor from its file path and verify its integrity using the hash in the filename.
+
+    Args:
+        tensor_file_path (str): Direct path to the tensor .bin file. The filename should be
+                               the hash of the file contents followed by .bin extension.
+        device (str, optional): Device to load the tensor to (e.g., 'cuda:0', 'cpu').
+                               If None, keeps the tensor on its original device.
+
+    Returns:
+        torch.Tensor: The loaded tensor (moved to the specified device if provided)
+
+    Raises:
+        FileNotFoundError: If the tensor file doesn't exist
+        RuntimeError: If the tensor cannot be loaded
+        ValueError: If the computed hash doesn't match the filename hash
+    """
+    blob_path = Path(tensor_file_path)
+
+    if not blob_path.exists():
+        raise FileNotFoundError(f"Tensor blob not found: {blob_path}")
+
+    # Extract expected hash from filename (remove .bin extension)
+    expected_hash = blob_path.stem
+
+    # Compute actual hash of file contents
+    with open(blob_path, "rb") as f:
+        file_contents = f.read()
+        computed_hash = hashlib.blake2b(file_contents).hexdigest()
+
+    # Verify hash matches filename
+    if computed_hash != expected_hash:
+        raise ValueError(
+            f"Hash verification failed: expected '{expected_hash}' but computed '{computed_hash}'"
+        )
+
+    try:
+        # Load the tensor using torch.load (tensors are saved with torch.save)
+        # If device is None, keep tensor on its original device, otherwise move to specified device
+        tensor = torch.load(blob_path, map_location=device)
+        return tensor
+    except Exception as e:
+        raise RuntimeError(f"Failed to load tensor from {blob_path}: {str(e)}") from e
+
+
+def create_args_from_json(json_path):
+    """
+    Parse a reproducer JSON and build kernel grid and argument dictionary.
+
+    Args:
+        json_path (str): Path to the JSON file describing the kernel launch.
+
+    Returns:
+        tuple[list, dict]: Grid specification list and map of argument name to value.
+    """
+    with open(json_path, "r") as f:
+        data = json.load(f)
+    # Handle data format validation and extraction
+    if isinstance(data, list):
+        if len(data) != 1:
+            print(
+                f"Error: Expected single element list, got list with {len(data)} elements"
+            )
+            sys.exit(1)
+        data = data[0]
+    elif not isinstance(data, dict):
+        print(f"Error: Expected list or dict, got {type(data)}")
+        sys.exit(1)
+
+    grid = data.get("grid", [])
+    args_dict = {}
+    extracted_args = data.get("extracted_args", {})
+
+    for arg_name, arg_info in extracted_args.items():
+        args_dict[arg_name] = _create_arg_from_info(arg_info)
+
+    return grid, args_dict
+
+
+def _create_arg_from_info(arg_info):
+    """
+    Recursively construct a kernel argument from its JSON schema.
+
+    Args:
+        arg_info (dict): JSON object describing a single argument, including
+            fields like 'type', 'value', 'dtype', 'shape', 'device', etc.
+
+    Returns:
+        Any: The constructed Python object suitable for kernel invocation.
+
+    Raises:
+        RuntimeError: When required optional dependencies are missing.
+        NotImplementedError: When a dtype or type is not supported yet.
+    """
+    arg_type = arg_info.get("type")
+
+    if arg_type in ["int", "bool"]:
+        return arg_info.get("value")
+
+    elif arg_type == "tensor":
+        if arg_info.get("blob_path"):
+            return load_tensor(arg_info.get("blob_path"), arg_info.get("device"))
+        dtype_str = arg_info.get("dtype")
+        try:
+            torch_dtype = getattr(torch, dtype_str.split(".")[-1])
+        except AttributeError:
+            torch_dtype = torch.float32
+
+        shape = arg_info.get("shape", [])
+        device = arg_info.get("device", "cpu")
+
+        # 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_()
+        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)
+            else:
+                return torch.empty(shape, dtype=torch_dtype, device=device).random_()
+        elif torch_dtype in [
+            torch.int8,
+            torch.int16,
+            torch.int32,
+            torch.int64,
+            torch.uint8,
+            torch.bool,
+        ]:
+            return torch.empty(shape, dtype=torch_dtype, device=device).random_()
+        # Case 2: Complex numbers need special handling
+        elif tensor_props.is_complex():
+            float_dtype = (
+                torch.float32 if torch_dtype == torch.complex64 else torch.float64
+            )
+            real_part = torch.rand(shape, dtype=float_dtype, device=device)
+            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_()
+        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
+        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(
+                "Optional dependency 'triton_kernels.tensor' is not installed; cannot construct Tensor."
+            )
+        Tensor, Storage, StridedLayout = _get_triton_tensor_types()
+        storage = _create_arg_from_info(arg_info.get("storage"))
+        dtype_str = arg_info.get("dtype")
+        torch_dtype = getattr(torch, dtype_str.split(".")[-1])
+        return Tensor(
+            storage=storage,
+            shape=arg_info.get("shape"),
+            shape_max=arg_info.get("shape_max"),
+            dtype=torch_dtype,
+        )
+
+    elif arg_type == "triton_kernels.tensor.Storage":
+        if not TRITON_KERNELS_CUSTOM_TYPES:
+            raise RuntimeError(
+                "Optional dependency 'triton_kernels.tensor' is not installed; cannot construct Storage."
+            )
+        Tensor, Storage, StridedLayout = _get_triton_tensor_types()
+        data = _create_arg_from_info(arg_info.get("data"))
+        layout = _create_arg_from_info(arg_info.get("layout"))
+        return Storage(data=data, layout=layout)
+
+    elif arg_type == "StridedLayout":
+        if not TRITON_KERNELS_CUSTOM_TYPES:
+            raise RuntimeError(
+                "Optional dependency 'triton_kernels.tensor' is not installed; cannot construct StridedLayout."
+            )
+        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
+
+
+if __name__ == "__main__":
+    script_dir = Path(__file__).resolve().parent
+    json_file = script_dir / "{{JSON_FILE_NAME_PLACEHOLDER}}"
+    grid, args_dict = create_args_from_json(str(json_file))
+
+    print("Generated kernel arguments dictionary:")
+    for name, arg in args_dict.items():
+        print(f"  {name}: {arg}")
+    print(f"Grid: {grid}")
+
+    # {{KERNEL_INVOCATION_PLACEHOLDER}}
+
+    torch.cuda.synchronize()
+    print("Kernel execution finished.")

tritonparse/reproducer/templates/loader.py

@@ -0,0 +1,57 @@
+from importlib.resources import files as pkg_files
+from pathlib import Path
+from typing import List
+
+
+BUILTIN_TEMPLATES_PACKAGE = "tritonparse.reproducer.templates"
+
+
+def _is_path_like(template_arg: str) -> bool:
+    return "/" in template_arg or "\\" in template_arg or template_arg.endswith(".py")
+
+
+def _read_file_text(path: Path) -> str:
+    p = path.expanduser().resolve()
+    if not p.exists() or not p.is_file():
+        raise FileNotFoundError(f"Template not found: {p}")
+    return p.read_text(encoding="utf-8")
+
+
+def _read_builtin_template_text(name: str) -> str:
+    resource = pkg_files(BUILTIN_TEMPLATES_PACKAGE).joinpath(f"{name}.py")
+    # resource may not exist if an invalid name is provided
+    try:
+        with resource.open("r", encoding="utf-8") as f:
+            return f.read()
+    except FileNotFoundError as exc:
+        available = ", ".join(list_builtin_templates())
+        raise FileNotFoundError(
+            f"Builtin template '{name}' not found. Available: {available}"
+        ) from exc
+
+
+def list_builtin_templates() -> List[str]:
+    """
+    Return the list of available builtin template names (without .py suffix).
+    """
+    names: List[str] = []
+    for entry in pkg_files(BUILTIN_TEMPLATES_PACKAGE).iterdir():
+        try:
+            if entry.is_file():
+                filename = entry.name
+                if filename.endswith(".py") and not filename.startswith("__"):
+                    names.append(filename[:-3])
+        except (OSError, FileNotFoundError):
+            # Defensive: in case entry access fails in some environments
+            continue
+    names.sort()
+    return names
+
+
+def load_template_code(template_arg: str) -> str:
+    """
+    Load template code by name (builtin, without .py) or by filesystem path.
+    """
+    if _is_path_like(template_arg):
+        return _read_file_text(Path(template_arg))
+    return _read_builtin_template_text(template_arg)

tritonparse/reproducer/utils.py

@@ -2,25 +2,40 @@ import importlib
 import importlib.util
 import json
 import sys
+from datetime import datetime
 from functools import lru_cache
+from pathlib import Path
 
 import torch
 
+from tritonparse.tools.load_tensor import load_tensor
+from tritonparse.tp_logger import logger
+
 TRITON_KERNELS_CUSTOM_TYPES = (
-    importlib.util.find_spec("triton_kernels.tensor") is not None
+    importlib.util.find_spec("triton_kernels") is not None
+    and importlib.util.find_spec("triton_kernels.tensor") is not None
 )
 
 
+@lru_cache(maxsize=1)
+def _get_triton_tensor_types():
+    mod = importlib.import_module("triton_kernels.tensor")
+    return (
+        mod.Tensor,
+        mod.Storage,
+        mod.StridedLayout,
+    )
+
+
 def create_args_from_json(json_path):
     """
-    Creates a list of arguments for a kernel launch from a JSON file.
+    Parse a reproducer JSON and build kernel grid and argument dictionary.
 
     Args:
-        json_path (str): The path to the JSON file containing the kernel
-                         launch information.
+        json_path (str): Path to the JSON file describing the kernel launch.
 
     Returns:
-        tuple: A tuple containing the grid and a dictionary of arguments.
+        tuple[list, dict]: Grid specification list and map of argument name to value.
     """
     with open(json_path, "r") as f:
         data = json.load(f)
@@ -46,19 +61,20 @@ def create_args_from_json(json_path):
     return grid, args_dict
 
 
-@lru_cache(maxsize=1)
-def _get_triton_tensor_types():
-    mod = importlib.import_module("triton_kernels.tensor")
-    return (
-        getattr(mod, "Tensor"),
-        getattr(mod, "Storage"),
-        getattr(mod, "StridedLayout"),
-    )
-
-
 def _create_arg_from_info(arg_info):
     """
-    Recursively creates a kernel argument from its JSON info dictionary.
+    Recursively construct a kernel argument from its JSON schema.
+
+    Args:
+        arg_info (dict): JSON object describing a single argument, including
+            fields like 'type', 'value', 'dtype', 'shape', 'device', etc.
+
+    Returns:
+        Any: The constructed Python object suitable for kernel invocation.
+
+    Raises:
+        RuntimeError: When required optional dependencies are missing.
+        NotImplementedError: When a dtype or type is not supported yet.
     """
     arg_type = arg_info.get("type")
 
@@ -66,6 +82,8 @@ def _create_arg_from_info(arg_info):
         return arg_info.get("value")
 
     elif arg_type == "tensor":
+        if arg_info.get("blob_path"):
+            return load_tensor(arg_info.get("blob_path"), arg_info.get("device"))
         dtype_str = arg_info.get("dtype")
         try:
             torch_dtype = getattr(torch, dtype_str.split(".")[-1])
@@ -79,7 +97,13 @@ def _create_arg_from_info(arg_info):
         tensor_props = torch.empty(0, dtype=torch_dtype)
 
         # Case 1: Floating point, signed integers, uint8, and bool are supported by random_()
-        if tensor_props.is_floating_point() or torch_dtype in [
+        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)
+            else:
+                return torch.empty(shape, dtype=torch_dtype, device=device).random_()
+        elif torch_dtype in [
             torch.int8,
             torch.int16,
             torch.int32,
@@ -88,7 +112,6 @@ def _create_arg_from_info(arg_info):
             torch.bool,
         ]:
             return torch.empty(shape, dtype=torch_dtype, device=device).random_()
-
         # Case 2: Complex numbers need special handling
         elif tensor_props.is_complex():
             float_dtype = (
@@ -101,13 +124,11 @@ def _create_arg_from_info(arg_info):
         # Case 3: 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
+            # Case 4: 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(
@@ -145,3 +166,137 @@ def _create_arg_from_info(arg_info):
     else:
         print(f"Warning: Unhandled argument type '{arg_type}'. Returning None.")
         return None
+
+
+def determine_output_paths(out_dir: str, kernel_name: str):
+    """
+    Determine output file paths for reproducer script and context data.
+
+    Args:
+        out_dir: Output directory path. If empty, uses default location.
+        kernel_name: Name of the kernel for default directory naming.
+
+    Returns:
+        Tuple of (python_script_path, json_context_path) as Path objects.
+    """
+    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
+    output_directory = Path(out_dir) / kernel_name
+    output_directory.mkdir(parents=True, exist_ok=True)
+
+    out_py_path = output_directory / f"repro_{timestamp}.py"
+    temp_json_path = output_directory / f"repro_context_{timestamp}.json"
+
+    return out_py_path, temp_json_path
+
+
+def _generate_import_statements(kernel_info) -> tuple[str, str]:
+    """
+    Generate (sys.path insertion statement, import statement) for the kernel.
+
+    Strategy:
+    - Always add the kernel file's parent directory to sys.path.
+    - If the filename (without .py) is a valid identifier, import using that
+      module name: `from <stem> import <func> as imported_kernel_function`.
+    - Otherwise, fall back to dynamic import via importlib.util and bind
+      `imported_kernel_function` from the loaded module.
+    """
+    file_path = Path(kernel_info.file_path)
+    function_name = kernel_info.function_name
+
+    if not file_path or not function_name:
+        raise ValueError("Kernel file path or function name missing from context.")
+
+    # Always add the file's parent directory to sys.path
+    sys_stmt = (
+        "import sys; p = r'" + str(file_path.parent) + "';\n"
+        "if p not in sys.path: sys.path.insert(0, p)"
+    )
+
+    module_name = file_path.with_suffix("").name
+    if module_name.isidentifier():
+        import_stmt = (
+            f"from {module_name} import {function_name} as imported_kernel_function"
+        )
+        logger.debug("Generated direct import statement: %s", import_stmt)
+        return sys_stmt, import_stmt
+
+    # Fallback: dynamic import when filename is not a valid identifier
+    import_stmt = (
+        "import importlib.util\n"
+        f"_spec = importlib.util.spec_from_file_location('kernel_mod', r'{str(file_path)}')\n"
+        "_mod = importlib.util.module_from_spec(_spec)\n"
+        "_spec.loader.exec_module(_mod)\n"
+        f"imported_kernel_function = getattr(_mod, '{function_name}')"
+    )
+    logger.debug("Generated dynamic import for file: %s", file_path)
+    return sys_stmt, import_stmt
+
+
+def _parse_kernel_signature(kernel_source_code: str) -> tuple[list[str], list[str]]:
+    """
+    Parses a Triton kernel's source code to distinguish positional args
+    from keyword args (those with default values).
+    """
+    signature_lines = []
+    in_signature = False
+    for line in kernel_source_code.splitlines():
+        # Mark beginning of signature when function definition is found
+        if line.strip().startswith("def "):
+            in_signature = True
+        if in_signature:
+            # Strip comments and leading/trailing whitespace
+            clean_line = line.split("#")[0].strip()
+            signature_lines.append(clean_line)
+            # Stop capturing after the signature ends
+            if "):" in line:
+                break
+
+    full_signature = "".join(signature_lines)
+    # Extract content between the first '(' and the last '):'
+    try:
+        params_str = full_signature[
+            full_signature.find("(") + 1 : full_signature.rfind("):")
+        ]
+    except IndexError as exc:
+        raise ValueError("Could not parse kernel signature.") from exc
+
+    # Clean up and split the parameters string
+    params = [p.strip() for p in params_str.replace("\n", "").split(",") if p.strip()]
+
+    positional_args = []
+    keyword_args = []
+
+    for param in params:
+        if "=" in param:
+            # Keyword arguments have a default value
+            arg_name = param.split("=")[0].strip()
+            keyword_args.append(arg_name)
+        else:
+            # Positional arguments do not have a default value
+            arg_name = param.split(":")[0].strip()
+            positional_args.append(arg_name)
+
+    logger.debug("Parsed positional args: %s", positional_args)
+    logger.debug("Parsed keyword args: %s", keyword_args)
+    return positional_args, keyword_args
+
+
+def _generate_invocation_snippet(
+    positional_args: list[str], keyword_args: list[str]
+) -> str:
+    """Generates a single-line Python code snippet for kernel invocation."""
+    # Prepare positional args for direct injection into the call
+    pos_args_str = ", ".join([f'args_dict["{arg}"]' for arg in positional_args])
+
+    # Prepare keyword args for direct injection
+    kw_args_str = ", ".join([f'{arg}=args_dict["{arg}"]' for arg in keyword_args])
+
+    # Combine them, ensuring proper comma separation
+    all_args = []
+    if pos_args_str:
+        all_args.append(pos_args_str)
+    if kw_args_str:
+        all_args.append(kw_args_str)
+
+    # Create the single-line call
+    return f"imported_kernel_function[tuple(grid)]({', '.join(all_args)})"

tritonparse/tools/prettify_ndjson.py

@@ -40,7 +40,7 @@ import argparse
 import json
 import sys
 from pathlib import Path
-from typing import Any, List
+from typing import Any, List, Union
 
 
 def parse_line_ranges(lines_arg: str) -> set[int]:
@@ -174,7 +174,7 @@ def load_ndjson(
     except FileNotFoundError:
         print(f"Error: File '{file_path}' not found.", file=sys.stderr)
         raise
-    except Exception as e:
+    except (OSError, UnicodeDecodeError) as e:
         print(f"Error reading file '{file_path}': {e}", file=sys.stderr)
         raise
 
@@ -201,19 +201,21 @@ def load_ndjson(
     return json_objects
 
 
-def save_prettified_json(json_objects: List[Any], output_path: Path) -> None:
+def save_prettified_json(
+    json_objects: Union[List[Any], Any], output_path: Path
+) -> None:
     """
-    Save list of JSON objects to a prettified JSON file.
+    Save JSON data to a prettified JSON file.
 
     Args:
-        json_objects: List of JSON objects to save
+        json_objects: Either a list of JSON objects or a single JSON-serializable object
         output_path: Path where to save the prettified JSON file
     """
     try:
         with open(output_path, "w", encoding="utf-8") as f:
             json.dump(json_objects, f, indent=2, ensure_ascii=False, sort_keys=True)
         print(f"Successfully converted to prettified JSON: {output_path}")
-    except Exception as e:
+    except OSError as e:
         print(f"Error writing to file '{output_path}': {e}", file=sys.stderr)
         raise
 

tritonparse/utils.py

@@ -16,21 +16,15 @@ from .common import (
 )
 from .source_type import Source, SourceType
 
-# argument parser for OSS
-parser = None
 
-
-def init_parser():
-    global parser
-
-    parser = argparse.ArgumentParser(
-        description="analyze triton structured logs", conflict_handler="resolve"
-    )
-
-    # Add arguments for the parse command
+def _add_parse_args(parser: argparse.ArgumentParser) -> None:
+    """Add common 'parse' subcommand arguments to a parser."""
     parser.add_argument(
         "source",
-        help="Source of torch logs to be analyzed. It is expected to path to a local directory or log",
+        help=(
+            "Source of torch logs to be analyzed. It is expected to path to a local "
+            "directory or log"
+        ),
     )
     parser.add_argument(
         "-o",
@@ -40,7 +34,9 @@ def init_parser():
     )
     parser.add_argument(
         "--overwrite",
-        help="Delete out directory if it already exists. Only does something if --out is set",
+        help=(
+            "Delete out directory if it already exists. Only does something if --out is set"
+        ),
         action="store_true",
     )
     parser.add_argument("-r", "--rank", help="Rank of logs to be analyzed", type=int)
@@ -54,7 +50,6 @@ def init_parser():
         from tritonparse.fb.utils import append_parser
 
         append_parser(parser)
-    return parser
 
 
 def oss_run(
@@ -113,12 +108,6 @@ def oss_run(
     print_parsed_files_summary(out_dir)
 
 
-def unified_parse_from_cli():
-    parser = init_parser()
-    args = parser.parse_args()
-    return unified_parse(**vars(args))
-
-
 def unified_parse(
     source: str,
     out: Optional[str] = None,

{tritonparse-0.2.4.dev20250922071528.dist-info → tritonparse-0.2.4.dev20250924071525.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tritonparse
-Version: 0.2.4.dev20250922071528
+Version: 0.2.4.dev20250924071525
 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.dev20250922071528.dist-info → tritonparse-0.2.4.dev20250924071525.dist-info}/RECORD

@@ -10,17 +10,22 @@ tritonparse/sourcemap_utils.py,sha256=qsQmTDuEe9yuUVyxSHRbjTR38gi0hvJEijnPkrJVAV
 tritonparse/structured_logging.py,sha256=7r9pv6miUdb8-CCZfj8SkD3XItzwPeONmszEL7TZak4,43949
 tritonparse/tp_logger.py,sha256=vXzY7hMDmVnRBGBhIjFZe3nHZzG5NKKPONGUszJhGgU,242
 tritonparse/trace_processor.py,sha256=QzUOKwnOkBbwTTKBsa5ZMUABPLMJIBFtTcG2SkhO0I8,12771
-tritonparse/utils.py,sha256=wt61tpbkqjGqHh0c7Nr2WlOv7PbQssmjULd6uA6aAko,4475
+tritonparse/utils.py,sha256=ujx9iUrpOthJ5vWzaNs6RXtqX0dp_GeozOaQLqlUDxg,4269
 tritonparse/reproducer/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tritonparse/reproducer/utils.py,sha256=VfMBwnTEZO8Ug9_ZRlZUVTMaMczDkviAykXpnK5dacU,5093
+tritonparse/reproducer/cli.py,sha256=JvnAi1FKSpNa6eHEapRn9jdXxsj1vvyrrEEnfxTJYa8,871
+tritonparse/reproducer/orchestrator.py,sha256=9SQ_rATY-s4r3BZQZdKLw7WYGz8IQJ1StPMgRbKAs5s,2456
+tritonparse/reproducer/utils.py,sha256=qi4XTKk0pWV4hgYg_GPBISEfVXlrI6tZR0A5ZZbwVyo,11132
+tritonparse/reproducer/ingestion/ndjson.py,sha256=pEujTl5xXW2E2DEW8ngxXQ8qP9oawb90wBVTWHDs1jk,7372
+tritonparse/reproducer/templates/example.py,sha256=XWfXD4tDOiE213YlWWK1l1ZgXbK3BX61NnvuVTkO-S0,8595
+tritonparse/reproducer/templates/loader.py,sha256=HqjfThdDVg7q2bYWry78sIaVRkUpkcA8KQDt83YrlVE,1920
 tritonparse/tools/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tritonparse/tools/decompress_bin_ndjson.py,sha256=kpt7DM_sSA334F1X45xdkP2OR9LuB27Pc50EkGr6CPM,4144
 tritonparse/tools/format_fix.py,sha256=Ol0Sjui8D7OzHwbamAfGnq8V5Y63uwNaFTKSORN5HkQ,3867
 tritonparse/tools/load_tensor.py,sha256=tfdmNVd9gsZqO6msQBhbXIhOvUzgc83yF64k2GDWPNk,2122
-tritonparse/tools/prettify_ndjson.py,sha256=VOzVWoXpCbaAXYA4i_wBcQIHfh-JhAx7xR4cF_L8yDs,10928
+tritonparse/tools/prettify_ndjson.py,sha256=YpJ7SFXTkZPZEXQeN1w5wkOf9pFrGqaqhhfHV7eobWA,10998
 tritonparse/tools/readme.md,sha256=w6PWYfYnRgoPArLjxG9rVrpcLUkoVMGuRlbpF-o0IQM,110
-tritonparse-0.2.4.dev20250922071528.dist-info/licenses/LICENSE,sha256=4ZciugpyN7wcM4L-9pyDh_etvMUeIfBhDTyH1zeZlQM,1515
-tritonparse-0.2.4.dev20250922071528.dist-info/METADATA,sha256=pm3r6Z1nR3gOJ35Ztyen1MhOvxKfPqz18_06ASPNYlc,6580
-tritonparse-0.2.4.dev20250922071528.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-tritonparse-0.2.4.dev20250922071528.dist-info/top_level.txt,sha256=ITcTKgp3vf_bXV9vixuQU9IrZa3L1EfDSZwvRzRaoJU,12
-tritonparse-0.2.4.dev20250922071528.dist-info/RECORD,,
+tritonparse-0.2.4.dev20250924071525.dist-info/licenses/LICENSE,sha256=4ZciugpyN7wcM4L-9pyDh_etvMUeIfBhDTyH1zeZlQM,1515
+tritonparse-0.2.4.dev20250924071525.dist-info/METADATA,sha256=cxNEHWh9EoRq332ybQO3nsai4gC5eASCWvyPloW0gko,6580
+tritonparse-0.2.4.dev20250924071525.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tritonparse-0.2.4.dev20250924071525.dist-info/top_level.txt,sha256=ITcTKgp3vf_bXV9vixuQU9IrZa3L1EfDSZwvRzRaoJU,12
+tritonparse-0.2.4.dev20250924071525.dist-info/RECORD,,