tritonparse 0.3.2.dev20251210071601__py3-none-any.whl

tritonparse/reproducer/orchestrator.py

@@ -0,0 +1,110 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
+from pathlib import Path
+from typing import Optional
+
+from tritonparse.info.kernel_query import find_launch_index_by_kernel
+from tritonparse.reproducer.ingestion.ndjson import build_context_bundle
+from tritonparse.reproducer.placeholder_replacer import (
+    DefaultPlaceholderReplacer,
+    PlaceholderReplacer,
+)
+from tritonparse.reproducer.templates.loader import load_template_code
+from tritonparse.reproducer.types import KernelImportMode
+from tritonparse.reproducer.utils import determine_output_paths, format_python_code
+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,
+    kernel_name: Optional[str] = None,
+    launch_id: int = 0,
+    replacer: Optional[PlaceholderReplacer] = None,
+    kernel_import: KernelImportMode = KernelImportMode.DEFAULT,
+) -> dict[str, str]:
+    """
+    Generate a reproducer script from NDJSON trace file.
+
+    Must provide either line_index OR (kernel_name + launch_id), not both.
+    If kernel_name is provided, the line_index parameter will be ignored and
+    recalculated from the kernel lookup.
+
+    Args:
+        input_path: Path to ndjson file. Supports uncompressed (.ndjson),
+            gzip compressed (.ndjson.gz), and gzip member concatenation (.bin.ndjson) formats.
+        line_index: 0-based index in events list. Ignored if kernel_name is provided.
+        out_dir: Output directory for reproducer files.
+        template: Template name to use for the reproducer.
+        kernel_name: Exact kernel name to match (case-sensitive). If provided, line_index will be recalculated.
+        launch_id: 0-based launch index for the kernel (default: 0, first launch).
+        replacer: Optional custom PlaceholderReplacer instance. If None, uses DefaultPlaceholderReplacer.
+        kernel_import: Kernel import mode (DEFAULT or COPY).
+    """
+    events = load_ndjson(Path(input_path))
+    logger.debug(f"Loaded {len(events)} events")
+
+    # If kernel_name is provided, lookup the actual line_index (overrides the parameter)
+    if kernel_name is not None:
+        logger.debug(
+            f"Looking up kernel '{kernel_name}' launch_id={launch_id} in {input_path}"
+        )
+        line_index = find_launch_index_by_kernel(events, kernel_name, launch_id)
+        logger.debug(
+            f"Found kernel '{kernel_name}' launch_id={launch_id} at line {line_index}"
+        )
+
+    logger.debug(f"Building bundle from {input_path} at line {line_index}")
+
+    # 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, template
+    )
+    save_prettified_json(context_bundle.raw_launch_event, temp_json_path)
+
+    # Save compilation event JSON if using OVERRIDE_TTIR mode
+    comp_json_path = None
+    if kernel_import == KernelImportMode.OVERRIDE_TTIR:
+        comp_json_path = (
+            temp_json_path.parent / f"{temp_json_path.stem}_compilation.json"
+        )
+        save_prettified_json(context_bundle.raw_comp_event, comp_json_path)
+
+    logger.debug("Loading reproducer template.")
+    template_code = load_template_code(template)
+
+    # Use PlaceholderReplacer to replace all placeholders
+    # If no custom replacer provided, use the default one
+    if replacer is None:
+        replacer = DefaultPlaceholderReplacer()
+    final_code = replacer.replace(
+        template_code,
+        context_bundle,
+        temp_json_path=temp_json_path,
+        kernel_import=kernel_import,
+        comp_json_filename=comp_json_path.name if comp_json_path else None,
+    )
+
+    # Format the generated code
+    final_code = format_python_code(final_code)
+
+    out_py_path.write_text(final_code, encoding="utf-8")
+
+    filepath = context_bundle.kernel_info.file_path
+    filepath = "/".join(filepath.split("/")[5:])
+    ret = {
+        "kernel_src_path": filepath,
+        "kernel": context_bundle.kernel_info.function_name,
+        "repro_script": str(out_py_path.resolve()),
+        "repro_context": str(temp_json_path.resolve()),
+    }
+    logger.info("REPRODUCER_OUTPUT\n%s", ret)
+
+    return ret

tritonparse/reproducer/placeholder_replacer.py

@@ -0,0 +1,335 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
+from abc import ABC
+from typing import Any, Dict, Optional, Protocol
+
+from tritonparse.reproducer.function_extractor import extract_utility_functions
+from tritonparse.reproducer.ingestion.ndjson import ContextBundle
+from tritonparse.reproducer.templates.utils import (
+    _disable_triton_autotune,
+    get_function_source,
+)
+from tritonparse.reproducer.types import KernelImportMode
+from tritonparse.reproducer.utils import (
+    _generate_import_statements,
+    _generate_invocation_snippet,
+    _parse_kernel_signature,
+)
+from tritonparse.tp_logger import logger
+
+
+class HandlerProtocol(Protocol):
+    def __call__(
+        self, code: str, context_bundle: ContextBundle, **kwargs: Any
+    ) -> str: ...
+
+
+class PlaceholderReplacer(ABC):
+    """
+    Abstract base class for template placeholder replacement.
+
+    Subclasses should register replacement handlers in their __init__ method
+    by calling self.register(placeholder, handler_function).
+
+    Each handler function should have the signature:
+        handler(code: str, context_bundle: ContextBundle, **kwargs) -> str
+    """
+
+    def __init__(self):
+        # Dictionary mapping placeholder strings to handler functions
+        self.handlers: Dict[str, HandlerProtocol] = {}
+
+    def register(self, placeholder: str, handler: HandlerProtocol):
+        """
+        Register a handler function for a specific placeholder.
+
+        Args:
+            placeholder: The placeholder string to replace (e.g., "{{JSON_FILE_NAME_PLACEHOLDER}}")
+            handler: A callable that takes (code, context_bundle, **kwargs) and returns modified code
+        """
+        self.handlers[placeholder] = handler
+
+    def replace(
+        self, template_code: str, context_bundle: ContextBundle, **kwargs: Any
+    ) -> str:
+        """
+        Replace all registered placeholders in the template code.
+
+        Args:
+            template_code: The template code containing placeholders
+            context_bundle: Context information about the kernel
+            **kwargs: Additional keyword arguments passed to handler functions
+
+        Returns:
+            The code with all placeholders replaced
+        """
+        code = template_code
+        for handler in self.handlers.values():
+            code = handler(code, context_bundle, **kwargs)
+        return code
+
+
+class DefaultPlaceholderReplacer(PlaceholderReplacer):
+    """
+    Default implementation of PlaceholderReplacer.
+
+    Handles the following placeholders:
+    - {{JSON_FILE_NAME_PLACEHOLDER}}: Replaced with the JSON file name
+    - # {{KERNEL_SYSPATH_PLACEHOLDER}}: Replaced with sys.path setup code
+    - # {{KERNEL_IMPORT_PLACEHOLDER}}: Replaced with kernel import statement
+    - # {{KERNEL_INVOCATION_PLACEHOLDER}}: Replaced with kernel invocation code
+    """
+
+    KERNEL_NAME_PLACEHOLDER = "{{KERNEL_NAME_PLACEHOLDER}}"
+    JSON_FILE_NAME_PLACEHOLDER = "{{JSON_FILE_NAME_PLACEHOLDER}}"
+    IR_OVERRIDE_SETUP_PLACEHOLDER = "# {{IR_OVERRIDE_SETUP_PLACEHOLDER}}"
+    KERNEL_SYSPATH_PLACEHOLDER = "# {{KERNEL_SYSPATH_PLACEHOLDER}}"
+    KERNEL_IMPORT_PLACEHOLDER = "# {{KERNEL_IMPORT_PLACEHOLDER}}"
+    UTILITY_FUNCTIONS_PLACEHOLDER = "# {{UTILITY_FUNCTIONS_PLACEHOLDER}}"
+    KERNEL_INVOCATION_PLACEHOLDER = "# {{KERNEL_INVOCATION_PLACEHOLDER}}"
+
+    def __init__(self):
+        super().__init__()
+        # Register all default handlers
+        self.register(self.JSON_FILE_NAME_PLACEHOLDER, self._replace_json_filename)
+        self.register(
+            self.IR_OVERRIDE_SETUP_PLACEHOLDER, self._replace_ir_override_setup
+        )
+        self.register(self.KERNEL_SYSPATH_PLACEHOLDER, self._replace_kernel_syspath)
+        self.register(self.KERNEL_IMPORT_PLACEHOLDER, self._replace_kernel_import)
+        self.register(
+            self.UTILITY_FUNCTIONS_PLACEHOLDER, self._replace_utility_functions
+        )
+        self.register(
+            self.KERNEL_INVOCATION_PLACEHOLDER, self._replace_kernel_invocation
+        )
+        self.register(self.KERNEL_NAME_PLACEHOLDER, self._replace_kernel_name)
+
+    def _replace_kernel_name(
+        self, code: str, context_bundle: ContextBundle, **kwargs
+    ) -> str:
+        """Replace the kernel name placeholder."""
+        kernel_name = context_bundle.kernel_info.function_name
+        if not kernel_name:
+            raise ValueError("Kernel function name is not available")
+        return code.replace(self.KERNEL_NAME_PLACEHOLDER, kernel_name)
+
+    def _replace_json_filename(
+        self, code: str, context_bundle: ContextBundle, **kwargs
+    ) -> str:
+        """Replace the JSON file name placeholder."""
+        temp_json_path = kwargs.get("temp_json_path")
+        if temp_json_path is None:
+            raise ValueError("temp_json_path is required for JSON filename replacement")
+        return code.replace(self.JSON_FILE_NAME_PLACEHOLDER, temp_json_path.name)
+
+    def _replace_ir_override_setup(
+        self, code: str, context_bundle: ContextBundle, **kwargs
+    ) -> str:
+        """Replace the IR override setup placeholder."""
+        kernel_import = kwargs.get("kernel_import", KernelImportMode.DEFAULT)
+
+        if kernel_import != KernelImportMode.OVERRIDE_TTIR:
+            return code.replace(self.IR_OVERRIDE_SETUP_PLACEHOLDER, "")
+
+        comp_json_filename = kwargs.get("comp_json_filename")
+        if not comp_json_filename:
+            raise ValueError("comp_json_filename is required for OVERRIDE_TTIR mode")
+
+        setup_code = f'''
+def create_ttir_tempfile():
+    """Extract TTIR from compilation event and create temporary file."""
+    script_dir = Path(__file__).resolve().parent
+    comp_json_file = script_dir / "{comp_json_filename}"
+
+    with open(comp_json_file, 'r') as f:
+        comp_data = json.load(f)
+
+    # Extract TTIR content
+    kernel_name = comp_data['payload']['metadata']['name']
+    ttir_key = f"{{kernel_name}}.ttir"
+    ttir_content = comp_data['payload']['file_content'][ttir_key]
+
+    # Create temporary file
+    temp_file = tempfile.NamedTemporaryFile(
+        mode='w',
+        suffix='.ttir',
+        delete=False,
+        prefix=f'{{kernel_name}}_'
+    )
+    temp_file.write(ttir_content)
+    temp_file.close()
+    return temp_file.name
+
+
+# Monkeypatch triton.autotune to use our TTIR
+_ttir_file = create_ttir_tempfile()
+_original_autotune = None
+
+def _patched_autotune(configs, key=None, **kwargs):
+    """Patched autotune that uses our TTIR file."""
+    import triton
+    # Replace configs with our single config using ir_override
+    new_configs = [triton.Config(kwargs={{}}, ir_override=_ttir_file)]
+    # Call original autotune with our config
+    return _original_autotune(new_configs, key=[], **kwargs)
+
+# Apply the monkeypatch before importing the kernel
+import triton
+_original_autotune = triton.autotune
+triton.autotune = _patched_autotune
+'''
+
+        return code.replace(self.IR_OVERRIDE_SETUP_PLACEHOLDER, setup_code)
+
+    def _replace_kernel_syspath(
+        self, code: str, context_bundle: ContextBundle, **kwargs
+    ) -> str:
+        """Replace the kernel sys.path placeholder."""
+        kernel_import = kwargs.get("kernel_import", KernelImportMode.DEFAULT)
+
+        if kernel_import == KernelImportMode.DEFAULT:
+            sys_stmt, _ = _generate_import_statements(context_bundle.kernel_info)
+            return code.replace(self.KERNEL_SYSPATH_PLACEHOLDER, sys_stmt)
+        elif kernel_import == KernelImportMode.COPY:
+            comment = (
+                "# Kernel sys.path setup skipped - kernel source code embedded below"
+            )
+            return code.replace(self.KERNEL_SYSPATH_PLACEHOLDER, comment)
+        elif kernel_import == KernelImportMode.OVERRIDE_TTIR:
+            comment = "# Kernel sys.path setup skipped - using IR override mode"
+            return code.replace(self.KERNEL_SYSPATH_PLACEHOLDER, comment)
+        else:
+            raise ValueError(f"Unknown kernel_import mode: {kernel_import}")
+
+    def _replace_kernel_import(
+        self, code: str, context_bundle: ContextBundle, **kwargs
+    ) -> str:
+        """Replace the kernel import placeholder."""
+        kernel_import = kwargs.get("kernel_import", KernelImportMode.DEFAULT)
+
+        if kernel_import == KernelImportMode.DEFAULT:
+            _, import_statement = _generate_import_statements(
+                context_bundle.kernel_info
+            )
+
+            final_stmt = "\n".join(
+                [import_statement, ""] + get_function_source(_disable_triton_autotune)
+            )
+            return code.replace(self.KERNEL_IMPORT_PLACEHOLDER, final_stmt)
+        elif kernel_import == KernelImportMode.COPY:
+            source_code = context_bundle.kernel_info.source_code
+            func_name = context_bundle.kernel_info.function_name
+
+            if not source_code or not source_code.strip():
+                raise ValueError("Kernel source code is empty, cannot use 'copy' mode")
+            if not func_name:
+                raise ValueError(
+                    "Cannot determine kernel function name for 'copy' mode"
+                )
+
+            if kernel_import == KernelImportMode.COPY:
+                dependent_source_map = get_dependent_source_map(
+                    context_bundle.kernel_info.function_name,
+                    context_bundle.kernel_info.file_path,
+                )
+                # Only add dependent functions if extraction was successful
+                if dependent_source_map:
+                    # Add separator, import statements, and dependent functions
+                    dependent_code = (
+                        "\n\n# Dependent functions extracted from source file\n\n"
+                    )
+                    dependent_code += "\n\n".join(dependent_source_map.values())
+                    source_code += "\n\n" + dependent_code
+                    logger.debug("Appended dependent functions to kernel source code")
+
+            # Add common imports needed for most Triton kernels
+            import_lines = [
+                "import torch",
+                "import numpy as np",
+                "import triton",
+                "import triton.language as tl",
+                "from typing import List, Tuple",
+                "",
+            ] + get_function_source(_disable_triton_autotune)
+
+            # Combine: imports + kernel source code + alias
+            embedded_code = "\n".join(import_lines)
+            embedded_code += "\n" + source_code
+            embedded_code += f"\n\n# Use kernel function directly\nimported_kernel_function = {func_name}"
+
+            return code.replace(self.KERNEL_IMPORT_PLACEHOLDER, embedded_code)
+        elif kernel_import == KernelImportMode.OVERRIDE_TTIR:
+            comment = "# Kernel import skipped - using IR override mode with TTIR"
+            return code.replace(self.KERNEL_IMPORT_PLACEHOLDER, comment)
+        else:
+            raise ValueError(f"Unknown kernel_import mode: {kernel_import}")
+
+    def _replace_utility_functions(
+        self, code: str, context_bundle: ContextBundle, **kwargs
+    ) -> str:
+        """Replace the utility functions placeholder with extracted functions."""
+        utility_code = extract_utility_functions()
+        return code.replace(self.UTILITY_FUNCTIONS_PLACEHOLDER, utility_code)
+
+    def _replace_kernel_invocation(
+        self, code: str, context_bundle: ContextBundle, **kwargs
+    ) -> str:
+        """Replace the kernel invocation placeholder."""
+        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)
+        return code.replace(self.KERNEL_INVOCATION_PLACEHOLDER, invocation_snippet)
+
+
+def get_dependent_source_map(
+    function_name: str, file_path: str
+) -> Optional[dict[str, str]]:
+    """
+    Extract dependent functions and their required imports.
+
+    Returns:
+        A tuple of (functions_dict, import_statements_list) or None if extraction fails.
+        - functions_dict: Maps qualified function names to their source code
+        - import_statements_list: List of formatted import statements needed by dependent functions
+    """
+    from pathlib import Path
+
+    from tritonparse.tp_logger import logger
+
+    source_path = Path(file_path)
+    if not source_path.exists():
+        return None
+
+    try:
+        # Use MultiFileCallGraphAnalyzer for multi-file analysis
+        from tritonparse.reproducer.multi_file_analyzer import (
+            MultiFileCallGraphAnalyzer,
+        )
+
+        analyzer = MultiFileCallGraphAnalyzer(
+            entry_file=file_path,
+            entry_function=function_name,
+        )
+        result = analyzer.analyze()
+
+        logger.info(
+            f"Extracted {result.stats.total_functions_found} dependent functions "
+            f"from {result.stats.total_files_analyzed} files with "
+            f"{result.stats.total_imports} imports"
+        )
+
+        # Print dependent functions' short names
+        logger.info("\nDependent functions (short names):")
+        for func_name in sorted(result.function_short_names.keys()):
+            short_name = result.function_short_names[func_name]
+            logger.info(
+                "  - %s. %s", short_name, result.functions[func_name].splitlines()[0]
+            )
+
+        return result.functions
+
+    except Exception as e:
+        # If AST analysis fails, continue without dependent functions
+        logger.warning(f"Failed to extract dependent functions: {e}", exc_info=True)
+        return None

tritonparse/reproducer/templates/example.py

@@ -0,0 +1,38 @@
+"""
+This file is automatically generated by TritonParse reproducer.
+It contains a smallest testing example for a Triton kernel.
+"""
+
+import logging
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+# {{IR_OVERRIDE_SETUP_PLACEHOLDER}}
+
+# {{KERNEL_SYSPATH_PLACEHOLDER}}
+
+# {{KERNEL_IMPORT_PLACEHOLDER}}
+
+# {{UTILITY_FUNCTIONS_PLACEHOLDER}}
+
+
+def launch_kernel():
+    script_dir = Path(__file__).resolve().parent  # noqa: F821
+    json_file = script_dir / "{{JSON_FILE_NAME_PLACEHOLDER}}"
+    grid, args_dict = create_args_from_json_file(str(json_file))  # noqa: F821
+
+    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.")
+
+
+if __name__ == "__main__":
+    launch_kernel()

tritonparse/reproducer/templates/loader.py

@@ -0,0 +1,59 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
+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/templates/tritonbench.py

@@ -0,0 +1,106 @@
+# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
+
+import logging
+from pathlib import Path
+from typing import Any, Callable, Dict, Optional, Tuple
+
+import torch
+from tritonbench.utils.triton_op import (
+    BenchmarkOperator,
+    register_benchmark,
+    REGISTERED_X_VALS,
+)
+
+logger = logging.getLogger(__name__)
+
+
+imported_kernel_function: Optional[Callable[[Tuple[int], Dict[str, Any]], None]] = None
+
+# {{IR_OVERRIDE_SETUP_PLACEHOLDER}}
+
+# {{KERNEL_SYSPATH_PLACEHOLDER}}
+
+# {{KERNEL_IMPORT_PLACEHOLDER}}
+
+# {{UTILITY_FUNCTIONS_PLACEHOLDER}}
+
+assert imported_kernel_function is not None, "imported_kernel_function is missing"
+
+KERNEL_NAME = "{{KERNEL_NAME_PLACEHOLDER}}"
+REPRO_CONTEXT_FILE_NAME = "{{JSON_FILE_NAME_PLACEHOLDER}}"
+
+
+def _get_launch_kernel_args() -> Tuple[Tuple[int], Dict[str, Any]]:
+    script_dir = Path(__file__).resolve().parent  # noqa: F821
+    json_file = script_dir / REPRO_CONTEXT_FILE_NAME
+
+    grid, args_dict = create_args_from_json_file(json_file)  # noqa: F821, F841
+
+    print("Recorded kernel arguments dictionary:")
+    for name, arg in args_dict.items():
+        if isinstance(arg, torch.Tensor):
+            print(
+                f"  {name}: Tensor:  {arg.shape} {arg.dtype} stride: {arg.stride()}, is_contiguous: {arg.is_contiguous()}"
+            )
+        else:
+            print(f"  {name}: {arg}")
+    print(f"Grid: {grid}")
+
+    return tuple(grid), args_dict
+
+
+grid, args_dict = _get_launch_kernel_args()
+
+
+def _launch_kernel(grid: tuple[int], args_dict: dict[str, Any]):
+    try:
+        assert grid is not None
+        assert args_dict is not None
+
+        # {{KERNEL_INVOCATION_PLACEHOLDER}}
+
+    except Exception as e:
+        print(f"Error: {e}")
+        print("Failed to launch kernel!")
+
+
+# HACK: @register_x_val doesn't allow us to pass `operator_name`` as a parameter
+tensor_args = {k: v for k, v in args_dict.items() if isinstance(v, torch.Tensor)}
+x_vals_label = ", ".join(tensor_args.keys())
+REGISTERED_X_VALS[KERNEL_NAME] = x_vals_label
+
+
+class Operator(BenchmarkOperator):
+    @register_benchmark(operator_name=KERNEL_NAME)
+    def run_kernel(self, grid, args_dict):
+        return lambda: _launch_kernel(grid, args_dict)
+
+    def get_input_iter(self):
+        yield {"grid": grid, "args_dict": args_dict}
+
+    def get_x_val(self, example_inputs):
+        tensors_shapes = [
+            tuple(v.shape)
+            for v in example_inputs["args_dict"].values()
+            if isinstance(v, torch.Tensor)
+        ]
+        return tuple(tensors_shapes)
+
+
+if __name__ == "__main__":
+    print("do_benchmark...")
+
+    args = [
+        "--benchmark-name",
+        KERNEL_NAME,
+    ]
+
+    from tritonbench.utils.parser import get_parser
+
+    parser = get_parser(args)
+    tb_args, extra_args = parser.parse_known_args(args)
+    bench = Operator(tb_args, extra_args)
+    bench.run()
+
+    print(bench.output)
+    print("Benchmark completed successfully!")