Get Inductor kernel params (#89)

Summary:

X-link: pytorch/pytorch#161953

Save the config args that Inductor burns into `inductor_metadata` so we can optionally pass them to any Jit Hooks that are set. This allows us to pass them to Tritonparse.

Reviewed By: davidberard98, FindHao

Differential Revision: D80994791

Author: NikhilAPatel

tritonparse/structured_logging.py

@@ -15,6 +15,7 @@
 from dataclasses import asdict, is_dataclass
 from datetime import date, datetime
 from enum import Enum
+from functools import partial
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Optional, Union
 
@@ -819,10 +820,18 @@ def extract_arg_info(arg_dict):
     return extracted_args
 
 
-def add_launch_metadata(grid, metadata, arg_dict):
+def add_launch_metadata(grid, metadata, arg_dict, inductor_args=None):
     # Extract detailed argument information
     extracted_args = extract_arg_info(arg_dict)
-    return {"launch_metadata_tritonparse": (grid, metadata._asdict(), extracted_args)}
+    extracted_inductor_args = extract_arg_info(inductor_args) if inductor_args else {}
+    return {
+        "launch_metadata_tritonparse": (
+            grid,
+            metadata._asdict(),
+            extracted_args,
+            extracted_inductor_args,
+        )
+    }
 
 
 class JITHookImpl(JITHook):
@@ -848,6 +857,7 @@ def __call__(
         compile,
         is_manual_warmup: bool,
         already_compiled: bool,
+        inductor_args: Optional[Dict[str, Any]] = None,
     ) -> Optional[bool]:
         """
         Override or set the launch_metadata function for the JIT-compiled kernel.
@@ -882,7 +892,9 @@ def __call__(
             log.warning(
                 f"fn {fn} launch_metadata is not None: {current_launch_metadata}. It will be overridden by tritonparse."
             )
-        function.launch_metadata = add_launch_metadata
+        function.launch_metadata = partial(
+            add_launch_metadata, inductor_args=inductor_args
+        )
         return True
 
 
@@ -946,6 +958,7 @@ def __call__(self, metadata):
             trace_data["extracted_args"] = launch_metadata_tritonparse[
                 2
             ]  # Now contains detailed arg info
+            trace_data["extracted_inductor_args"] = launch_metadata_tritonparse[3]
         trace_structured_triton("launch", metadata_fn=lambda: convert(trace_data))