Add support for Triton dtype parameters in reproducer

Summary:
Add support for parsing and handling Triton dtype parameters (like `tl.bfloat16`, `tl.float16`) in TritonParse reproducer scripts.

**Problem**:
TritonParse reproducer scripts failed when Triton kernels had dtype parameters (e.g., `ab_dtype`, `c_dtype`), showing "Warning: Unhandled argument type 'dtype'. Returning None" and passing `None` instead of the actual dtype object like `tl.bfloat16`.

**Solution**:
1. Added `TRITON_DTYPE_MAP` dictionary mapping dtype string representations (e.g., 'bf16') to Triton dtype objects (e.g., `tl.bfloat16`)
2. Added dtype handling branch in `_create_arg_from_info()` to parse dtype arguments
3. Updated function extractor to include triton.language import and TRITON_DTYPE_MAP in generated reproducers

The mapping covers all Triton dtypes:
- Signed/unsigned integers: int1, int8-64, uint8-64
- Standard floating point: fp16, bf16, fp32, fp64
- FP8 variants: fp8e4b15, fp8e4nv, fp8e4b8, fp8e5, fp8e5b16

**Impact**:
Future generated reproducer scripts will correctly handle dtype parameters, allowing kernels with dtype arguments to run successfully.

Reviewed By: htyu

Differential Revision: D87575372

fbshipit-source-id: b68e1c14a4945fd73c73d434267aa95f9a7db333

Author: FindHao

tritonparse/reproducer/function_extractor.py

@@ -58,6 +58,11 @@ def extract_utility_functions() -> str:
     if constant:
         extracted_parts.append(constant)
 
+    # Extract TRITON_DTYPE_MAP constant
+    dtype_map = _extract_assignment(utils_tree, utils_lines, "TRITON_DTYPE_MAP")
+    if dtype_map:
+        extracted_parts.append(dtype_map)
+
     # Extract load_tensor functions
     extracted_parts.extend(
         _extract_functions(
@@ -218,5 +223,6 @@ def _generate_imports() -> str:
         "from typing import Union",
         "",
         "import torch",
+        "import triton.language as tl",
     ]
     return "\n".join(imports)

tritonparse/reproducer/utils.py

@@ -11,6 +11,8 @@
 from pathlib import Path
 
 import torch
+import triton.language as tl
+
 from tritonparse.tools.load_tensor import load_tensor
 from tritonparse.tp_logger import logger
 
@@ -19,6 +21,32 @@
     and importlib.util.find_spec("triton_kernels.tensor") is not None
 )
 
+# Mapping from dtype string representation to Triton dtype objects
+TRITON_DTYPE_MAP = {
+    # Signed integers
+    "int8": tl.int8,
+    "int16": tl.int16,
+    "int32": tl.int32,
+    "int64": tl.int64,
+    # Unsigned integers
+    "int1": tl.int1,
+    "uint8": tl.uint8,
+    "uint16": tl.uint16,
+    "uint32": tl.uint32,
+    "uint64": tl.uint64,
+    # Standard floating point types
+    "fp16": tl.float16,
+    "bf16": tl.bfloat16,
+    "fp32": tl.float32,
+    "fp64": tl.float64,
+    # FP8 variants
+    "fp8e4b15": tl.float8e4b15,
+    "fp8e4nv": tl.float8e4nv,
+    "fp8e4b8": tl.float8e4b8,
+    "fp8e5": tl.float8e5,
+    "fp8e5b16": tl.float8e5b16,
+}
+
 
 @lru_cache(maxsize=1)
 def _get_triton_tensor_types():
@@ -322,6 +350,14 @@ def _create_arg_from_info(arg_info):
             )
         Tensor, Storage, StridedLayout = _get_triton_tensor_types()
         return StridedLayout(shape=arg_info.get("initial_shape"))
+
+    elif arg_type == "dtype":
+        dtype_repr = arg_info.get("repr")
+        if dtype_repr in TRITON_DTYPE_MAP:
+            return TRITON_DTYPE_MAP[dtype_repr]
+        else:
+            raise NotImplementedError(f"Unsupported Triton dtype: {dtype_repr}")
+
     else:
         print(f"Warning: Unhandled argument type '{arg_type}'. Returning None.")
         return None