tritonparse 0.2.1.dev20250915071616__py3-none-any.whl → 0.2.1.dev20250916071516__py3-none-any.whl

tritonparse/reproducer/utils.py

@@ -0,0 +1,147 @@
+import importlib
+import importlib.util
+import json
+import sys
+from functools import lru_cache
+
+import torch
+
+TRITON_KERNELS_CUSTOM_TYPES = (
+    importlib.util.find_spec("triton_kernels.tensor") is not None
+)
+
+
+def create_args_from_json(json_path):
+    """
+    Creates a list of arguments for a kernel launch from a JSON file.
+
+    Args:
+        json_path (str): The path to the JSON file containing the kernel
+                         launch information.
+
+    Returns:
+        tuple: A tuple containing the grid and a dictionary of arguments.
+    """
+    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
+
+
+@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.
+    """
+    arg_type = arg_info.get("type")
+
+    if arg_type in ["int", "bool"]:
+        return arg_info.get("value")
+
+    elif arg_type == "tensor":
+        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() or 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

tritonparse/structured_logging.py

@@ -183,13 +183,80 @@ def convert(obj):
             asdict(obj)
         )  # Convert dataclass to dict and then process that dict
 
+    if _is_triton_kernels_layout(obj):
+        layout_info = {"type": type(obj).__name__}
+        if hasattr(obj, "initial_shape"):
+            layout_info["initial_shape"] = convert(obj.initial_shape)
+        if hasattr(obj, "name"):
+            layout_info["name"] = convert(obj.name)
+        return layout_info
+
     # 4. Common Triton constexpr objects
     if isinstance(obj, dtype):
         return f"triton.language.core.dtype('{str(obj)}')"
+
+    if TORCH_INSTALLED and isinstance(obj, torch.dtype):
+        return str(obj)
+
     log.warning(f"Unknown type: {type(obj)}")
     return str(obj)  # Return primitive types as-is
 
 
+def _is_triton_kernels_layout(obj):
+    """
+    Check if an object is an instance of a Layout class from a
+    triton_kernels module by checking its MRO.
+    """
+    t = type(obj)
+    for base_class in t.__mro__:
+        module_name = getattr(base_class, "__module__", "")
+        type_name = getattr(base_class, "__name__", "")
+        if type_name == "Layout" and module_name.startswith("triton_kernels"):
+            return True
+    return False
+
+
+def _is_from_triton_kernels_module(obj):
+    """
+    Check if an object is an instance of Tensor or Storage from a
+    triton_kernels module.
+    """
+    t = type(obj)
+    module_name = getattr(t, "__module__", "")
+    type_name = getattr(t, "__name__", "")
+    return type_name in ("Tensor", "Storage") and module_name.startswith(
+        "triton_kernels"
+    )
+
+
+def _log_torch_tensor_info(tensor_value):
+    """
+    Extracts metadata from a torch.Tensor object.
+
+    Args:
+        tensor_value (torch.Tensor): The tensor to extract information from.
+
+    Returns:
+        dict: A dictionary containing tensor metadata.
+    """
+    arg_info = {}
+    arg_info["type"] = "tensor"
+    arg_info["shape"] = list(tensor_value.shape)
+    arg_info["dtype"] = str(tensor_value.dtype)
+    arg_info["device"] = str(tensor_value.device)
+    arg_info["stride"] = list(tensor_value.stride())
+    arg_info["numel"] = tensor_value.numel()
+    arg_info["is_contiguous"] = tensor_value.is_contiguous()
+    arg_info["element_size"] = tensor_value.element_size()
+    arg_info["storage_offset"] = tensor_value.storage_offset()
+    # Memory usage in bytes
+    arg_info["memory_usage"] = tensor_value.numel() * tensor_value.element_size()
+    # Add data_ptr for memory tracking (optional)
+    if hasattr(tensor_value, "data_ptr"):
+        arg_info["data_ptr"] = hex(tensor_value.data_ptr())
+    return arg_info
+
+
 def maybe_enable_debug_logging():
     """
     This logging is for logging module itself, not for logging the triton compilation.
@@ -769,7 +836,8 @@ def maybe_trace_triton(
 
 def extract_arg_info(arg_dict):
     """
-    Extract detailed information from kernel arguments, especially for PyTorch tensors.
+    Extract detailed information from kernel arguments, especially for PyTorch
+    tensors.
 
     Args:
         arg_dict: Dictionary of kernel arguments
@@ -785,19 +853,40 @@ def extract_arg_info(arg_dict):
         # Check if it's a PyTorch tensor
         if TORCH_INSTALLED and isinstance(arg_value, torch.Tensor):
             arg_info["type"] = "tensor"
-            arg_info["shape"] = list(arg_value.shape)
-            arg_info["dtype"] = str(arg_value.dtype)
-            arg_info["device"] = str(arg_value.device)
-            arg_info["stride"] = list(arg_value.stride())
-            arg_info["numel"] = arg_value.numel()
-            arg_info["is_contiguous"] = arg_value.is_contiguous()
-            arg_info["element_size"] = arg_value.element_size()
-            arg_info["storage_offset"] = arg_value.storage_offset()
-            # Memory usage in bytes
-            arg_info["memory_usage"] = arg_value.numel() * arg_value.element_size()
-            # Add data_ptr for memory tracking (optional)
-            if hasattr(arg_value, "data_ptr"):
-                arg_info["data_ptr"] = hex(arg_value.data_ptr())
+            arg_info.update(_log_torch_tensor_info(arg_value))
+        # Handle custom Tensor/Storage types from triton_kernels
+        elif _is_from_triton_kernels_module(arg_value):
+            type_name = type(arg_value).__name__
+            arg_info["type"] = f"triton_kernels.tensor.{type_name}"
+
+            if type_name == "Tensor":
+                # Dump all attributes needed to reconstruct the Tensor wrapper
+                if hasattr(arg_value, "shape"):
+                    arg_info["shape"] = convert(arg_value.shape)
+                if hasattr(arg_value, "shape_max"):
+                    arg_info["shape_max"] = convert(arg_value.shape_max)
+                if hasattr(arg_value, "dtype"):
+                    arg_info["dtype"] = convert(arg_value.dtype)
+                if hasattr(arg_value, "storage"):
+                    # Recursively process the storage, which can be another
+                    # custom type or a torch.Tensor
+                    storage_arg = {"storage": arg_value.storage}
+                    arg_info["storage"] = extract_arg_info(storage_arg)["storage"]
+
+            elif type_name == "Storage":
+                # Dump all attributes needed to reconstruct the Storage object
+                if (
+                    hasattr(arg_value, "data")
+                    and TORCH_INSTALLED
+                    and isinstance(arg_value.data, torch.Tensor)
+                ):
+                    # The 'data' is a torch.Tensor, log its metadata fully
+                    arg_info["data"] = _log_torch_tensor_info(arg_value.data)
+                if hasattr(arg_value, "layout"):
+                    arg_info["layout"] = convert(arg_value.layout)
+            else:
+                log.warning(f"Unknown type: {type(arg_value)}")
+
         # Handle scalar values
         elif isinstance(arg_value, (int, float, bool)):
             arg_info["type"] = type(arg_value).__name__

{tritonparse-0.2.1.dev20250915071616.dist-info → tritonparse-0.2.1.dev20250916071516.dist-info}/METADATA

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

@@ -7,18 +7,20 @@ tritonparse/mapper.py,sha256=prrczfi13P7Aa042OrEBsmRF1HW3jDhwxicANgPkWIM,4150
 tritonparse/shared_vars.py,sha256=-c9CvXJSDm9spYhDOJPEQProeT_xl3PaNmqTEYi_u4s,505
 tritonparse/source_type.py,sha256=nmYEQS8rfkIN9BhNhQbkmEvKnvS-3zAxRGLY4TaZdi8,1676
 tritonparse/sourcemap_utils.py,sha256=qsQmTDuEe9yuUVyxSHRbjTR38gi0hvJEijnPkrJVAV4,2037
-tritonparse/structured_logging.py,sha256=mrApsVigIHZln6De2ElwNTUSZfO61OHXgd08o2X26sM,40430
+tritonparse/structured_logging.py,sha256=7wSKc9HnV40eLoOgFNosc7dXWvSxsJJo5mnRVoDZqIk,43678
 tritonparse/tp_logger.py,sha256=vXzY7hMDmVnRBGBhIjFZe3nHZzG5NKKPONGUszJhGgU,242
 tritonparse/trace_processor.py,sha256=QzUOKwnOkBbwTTKBsa5ZMUABPLMJIBFtTcG2SkhO0I8,12771
 tritonparse/utils.py,sha256=wt61tpbkqjGqHh0c7Nr2WlOv7PbQssmjULd6uA6aAko,4475
+tritonparse/reproducer/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tritonparse/reproducer/utils.py,sha256=VfMBwnTEZO8Ug9_ZRlZUVTMaMczDkviAykXpnK5dacU,5093
 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/readme.md,sha256=w6PWYfYnRgoPArLjxG9rVrpcLUkoVMGuRlbpF-o0IQM,110
-tritonparse-0.2.1.dev20250915071616.dist-info/licenses/LICENSE,sha256=4ZciugpyN7wcM4L-9pyDh_etvMUeIfBhDTyH1zeZlQM,1515
-tritonparse-0.2.1.dev20250915071616.dist-info/METADATA,sha256=CigDDEKt_dqH4EA7Le5lt9a-HQkeidokxdsYJBadVA0,6306
-tritonparse-0.2.1.dev20250915071616.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-tritonparse-0.2.1.dev20250915071616.dist-info/top_level.txt,sha256=ITcTKgp3vf_bXV9vixuQU9IrZa3L1EfDSZwvRzRaoJU,12
-tritonparse-0.2.1.dev20250915071616.dist-info/RECORD,,
+tritonparse-0.2.1.dev20250916071516.dist-info/licenses/LICENSE,sha256=4ZciugpyN7wcM4L-9pyDh_etvMUeIfBhDTyH1zeZlQM,1515
+tritonparse-0.2.1.dev20250916071516.dist-info/METADATA,sha256=5ic4t2gOHTBKvj4oZlYVOERcUOIaO93Ufkbyu6NG5aI,6306
+tritonparse-0.2.1.dev20250916071516.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tritonparse-0.2.1.dev20250916071516.dist-info/top_level.txt,sha256=ITcTKgp3vf_bXV9vixuQU9IrZa3L1EfDSZwvRzRaoJU,12
+tritonparse-0.2.1.dev20250916071516.dist-info/RECORD,,