# mypy: allow-untyped-defs
from __future__ import annotations

import builtins
import copy
import dataclasses
import enum
import functools
import hashlib
import inspect
import itertools
import logging
import math
import operator
import os
import os.path
import re
import sys
import threading
import time
from collections import namedtuple
from typing import (
    Any,
    cast,
    Final,
    Generic,
    get_args,
    Literal,
    TYPE_CHECKING,
    TypeAlias,
    TypeVar,
)

import torch
from torch._dynamo.utils import counters, set_feature_use
from torch._inductor import metrics
from torch._inductor.config import triton as inductor_triton_config
from torch._prims_common import compute_required_storage_length
from torch.utils._debug_mode import get_active_debug_mode
from torch.utils._ordered_set import OrderedSet
from torch.utils._triton import get_triton_version, has_triton_stable_tma_api

from ..triton_bundler import TritonBundler
from ..utils import (
    GPU_KERNEL_BIN_EXTS,
    prefix_is_reduction,
    tlx_only_cuda_options,
    TMA_ALIGNMENT,
    triton_version_uses_attrs_dict,
    XPU_KERNEL_FORMAT,
)
from . import triton_helpers
from .autotune_cache import AutotuneCache
from .benchmarking import benchmarker
from .coordinate_descent_tuner import CoordescTuner
from .hints import (
    AutotuneHint,
    DeviceProperties,
    HeuristicType,
    InductorMeta,
    native_matmul_block_numel,
    ReductionHint,
    TileHint,
    TRITON_MAX_BLOCK,
    TRITON_MAX_TENSOR_NUMEL,
    TritonMeta,
)
from .runtime_utils import (
    cache_dir,
    ceildiv,
    conditional_product,
    create_bandwidth_info_str,
    dynamo_timed,
    get_first_attr,
    get_max_y_grid,
    get_num_bytes,
    next_power_of_2,
    triton_cache_dir,
    triton_config_to_hashable,
    triton_hash_to_path_key,
    validate_triton_config,
)
from .static_triton_launcher import (
    statically_launched_kernel_by_device,
    StaticallyLaunchedCudaKernel,
    StaticallyLaunchedXpuKernel,
)
from .triton_compat import (
    ASTSource,
    autograd_profiler,
    CompiledKernel,
    Config,
    GPUTarget,
    HAS_WARP_SPEC,
    IntelGPUError,
    KernelInterface,
    knobs,
    OutOfResources,
    PTXASError,
    triton,
)
from .triton_helpers import get_constexprs


class BenchmarkFailureReason(enum.Enum):
    """Reasons why a triton config benchmark may return float('inf')."""

    REGISTER_SPILLING = "register_spilling"
    INVALID_CONFIG = "invalid_config"


class InductorConfig(Config):
    """Inductor-specific Triton config with additional control flags"""

    def __init__(self, *args, dynamic_scale_rblock=True, **kwargs):
        super().__init__(*args, **kwargs)
        self.dynamic_scale_rblock = dynamic_scale_rblock


class NoTritonConfigsError(RuntimeError):
    pass


def _should_enable_triton_debug_asserts(inductor_meta: InductorMeta) -> bool:
    """
    Enable Triton debug asserts whenever indirect indexing asserts are on,
    except on HIP where older Triton releases lack the required support.

    Triton 3.7 is the first release that includes the upstream debug-assert
    support needed by ROCm, so HIP kernels must keep this disabled below that
    version to remain compatible.
    """
    if not inductor_meta.get("assert_indirect_indexing", True):
        return False
    if not inductor_meta.get("is_hip", False):
        return True
    return get_triton_version() >= (3, 7)


if TYPE_CHECKING:
    from collections.abc import Callable, Container, Hashable

    from torch._C._profiler import _RecordFunctionFast
    from torch._guards import CompileId
    from torch.utils._debug_mode import _TritonKernelCall

    LauncherType = Any

_KernelType: TypeAlias = (
    CompiledKernel | StaticallyLaunchedCudaKernel | StaticallyLaunchedXpuKernel
)
_T = TypeVar(
    "_T",
    CompiledKernel,
    StaticallyLaunchedCudaKernel,
    StaticallyLaunchedXpuKernel,
)
if get_args(_KernelType) != _T.__constraints__:
    raise AssertionError("_KernelType args must match _T type constraints")

log = logging.getLogger(__name__)
autotuning_inputs_log = torch._logging.getArtifactLogger(__name__, "autotuning_inputs")

triton_name_sub = re.compile(r"^def [^(]+\(")


def generate_lookup_hash_from_source_code(size_hints_str: str, source_code: str) -> str:
    # Name agnostic + strip white space
    fn_strip_name = re.sub(triton_name_sub, "(", source_code.strip(), count=1)
    hash_str = size_hints_str + fn_strip_name
    fn_hash = hashlib.sha256(hash_str.encode("utf-8")).hexdigest()

    return fn_hash


def lookup_autotune_config(size_hints, fn) -> Config | None:
    lookup_table = torch._inductor.config.autotune_lookup_table
    cached_config = None
    if len(lookup_table) > 0 and "_fused_" in fn.src:
        fn_hash = generate_lookup_hash_from_source_code(str(size_hints), fn.src)
        if fn_hash in lookup_table:
            config_dict = lookup_table[fn_hash]
            block_configs = {k: v for k, v in config_dict.items() if "BLOCK" in k}
            cached_config = Config(
                block_configs,
                num_warps=config_dict["num_warps"],
                num_stages=config_dict["num_stages"],
            )

    return cached_config


def get_total_reduction_numel(numels: dict[str, int]) -> int:
    return conditional_product(
        *[numel for prefix, numel in numels.items() if prefix_is_reduction(prefix)]
    )


def _resolve_dims(dims, cfg_kwargs, constants):
    """Resolve a list of block/shape/stride dims to concrete ints."""
    result = []
    for s in dims:
        if isinstance(s, int):
            result.append(s)
        elif isinstance(s, str) and s in constants:
            result.append(int(constants[s]))
        elif isinstance(s, str) and s in cfg_kwargs:
            result.append(int(cfg_kwargs[s]))
        else:
            log.debug("host-side TMA: unresolved descriptor dim %r; skipping", s)
            return None
    return result


@functools.lru_cache(None)
def _warn_host_tma_clone(name: str) -> None:
    log.warning(
        "host-side TMA: input %s is not %d-byte aligned; cloning it (an extra "
        "copy per launch). Pass aligned inputs to avoid this.",
        name,
        TMA_ALIGNMENT,
    )


def _host_tma_aligned(tensor, name):
    """Return a TMA_ALIGNMENT-aligned view of `tensor`, cloning (with a one-time
    warning) only if its base address is not aligned. Used by the host-side TMA
    launcher to build TensorDescriptors from aligned storage."""
    if tensor.data_ptr() % TMA_ALIGNMENT == 0:
        return tensor
    _warn_host_tma_clone(name)
    return tensor.clone()


def autotune_hints_to_configs(
    hints: OrderedSet[AutotuneHint],
    size_hints,
    block_size: int,
    device_props: DeviceProperties,
) -> list[Config]:
    """
    AutotuneHints can be attached to the metadata of triton kernels for providing
    suggestions about what to try for autotuning. One reason to do this is if there are
    some configs that are only useful in specific scenarios, in which case we can avoid
    wasting compile time on autotuning unless we know we are in one of those scenarios.

    Based on those hints, this function will generate a list of additional autotuning
    configs to try.
    """
    xyz_options: tuple[tuple[int, int | None, int | None], ...]
    configs: list[Config] = []
    for hint in hints:
        if hint == AutotuneHint.ONE_ELEMENT_PER_THREAD:
            if device_props.warp_size is None:
                log.debug(
                    "Skipping %s autotune hint because device %s does not report warp_size",
                    AutotuneHint.ONE_ELEMENT_PER_THREAD,
                    device_props.type,
                )
                continue
            if len(size_hints) == 1:
                xyz_options = ((block_size // 4, None, None),)
            elif len(size_hints) == 2:
                xyz_options = ((block_size // 4, 1, None), (1, block_size // 4, None))
            elif len(size_hints) == 3:
                xyz_options = (
                    (block_size // 4, 1, 1),
                    (1, block_size // 4, 1),
                    (1, 1, block_size // 4),
                )
            warp_size = device_props.warp_size_or_default
            configs.extend(
                triton_config(
                    size_hints,
                    *xyz,
                    num_elements_per_warp=warp_size,
                    warp_size=warp_size,
                )
                for xyz in xyz_options
            )

    return configs


def _dump_launch_params(args, kwargs, launcher, kernel_name, grid):
    call_args = []
    call_kwargs = dict(kwargs)
    for arg in args:
        if isinstance(arg, (int, bool)):
            call_args.append(str(arg))
        else:
            call_args.append("T")
    call_kwargs.update(launcher.config.kwargs)
    call_kwargs["num_warps"] = launcher.config.num_warps
    call_kwargs["num_stages"] = launcher.config.num_stages
    if HAS_WARP_SPEC:
        call_kwargs["num_consumer_groups"] = getattr(
            launcher.config, "num_consumer_groups", 0
        )
        call_kwargs["num_buffers_warp_spec"] = getattr(
            launcher.config, "num_buffers_warp_spec", 0
        )
    args_str = [*call_args]
    args_str.extend(f"{k}={v}" for k, v in call_kwargs.items())
    args_str = ", ".join(args_str)
    abs_path = os.path.abspath(sys.argv[0])
    with open(f"{abs_path}.launch_params", "a") as f:
        f.write(f"{kernel_name} | {args_str} | {grid!r}\n")


def _dump_launch_tensors(args, kernel_path, kernel_hash, kernel_name):
    tensor_list = [arg for arg in args if isinstance(arg, torch.Tensor)]

    run_index = 0

    # Some kernels don't have path and hash stored
    # Using only the name to differentiate between those
    if not kernel_path:
        kernel_hash = kernel_name

    # Saving only the last N runs of the kernels to avoid bloating the folder
    if kernel_hash in inductor_triton_config.debug_dump_kernel_inputs:
        run_index = inductor_triton_config.debug_dump_kernel_inputs[kernel_hash] + 1

        if run_index >= inductor_triton_config.max_kernel_dump_occurrences:
            run_index = 0

    inductor_triton_config.debug_dump_kernel_inputs[kernel_hash] = run_index

    # Default path for kernels with no hash
    if not kernel_path:
        directory_path = os.path.join(cache_dir(), "unhashed_kernel_inputs")
    else:
        directory_path = os.path.dirname(kernel_path)
    directory_path = f"{directory_path}/{kernel_name}_run_{run_index}"
    os.makedirs(directory_path, exist_ok=True)

    log.info(
        "Dumping %d tensor(s) for kernel %s to %s",
        len(tensor_list),
        kernel_name,
        directory_path,
    )

    for index, tensor in enumerate(tensor_list):
        torch.save(tensor, f"{directory_path}/tensor_{index}.pt")


def _combo_has_reduction_subkernel(inductor_meta: InductorMeta) -> bool:
    combo_meta = inductor_meta.get("combo_grid_meta")
    if combo_meta is None or "heuristic_0" not in combo_meta:
        return False
    # Stitched combos fix the reduction blocks at codegen time: no-bench uses a
    # single config, compile-time autotune passes the chosen blocks via
    # default_config and only autotunes launch candidates. Neither runtime-scales
    # rblock, so don't add scaling candidates for them.
    if "stitched_num_warps" in combo_meta or "stitched_launch_candidates" in combo_meta:
        return False
    return any(
        combo_meta.get(f"heuristic_{i}") == "reduction"
        for i in range(combo_meta["num_kernels"])
    )


def _could_dynamic_scale_rblock(
    *,
    size_hints: list[int] | None,
    heuristic_type: HeuristicType,
    device_prop: DeviceProperties | None,
    inductor_meta: InductorMeta,
) -> bool:
    return (
        device_prop is not None
        and not inductor_meta.get("deterministic", False)
        and inductor_meta.get("dynamic_scale_rblock", True)
        and not inductor_meta.get("persistent_reduction")
        and heuristic_type == HeuristicType.REDUCTION
        # Combo kernels with per-subkernel blocks set size_hints=None but
        # carry per-subkernel size hints in combo_grid_meta.
        and (size_hints is not None or _combo_has_reduction_subkernel(inductor_meta))
        # Disable for Intel as Triton is not ready to return n_regs for a compiled_binary.
        and device_prop.type in ["cuda", "hip"]
        and bool(device_prop.major)
        and (device_prop.major >= 8 or torch.version.hip)
        and device_prop.regs_per_multiprocessor is not None
        and device_prop.warp_size is not None
    )


def check_autotune_cache(
    configs: list[Config],
    filename: str | None,
    inductor_meta: InductorMeta,
    dynamic_scale_rblock_eligible: bool = False,
) -> tuple[list[Config], AutotuneCache | None, dict[str, Any]]:
    """
    Given a list of configs, checks autotune cache and return metadata
    """
    autotune_cache = None
    autotune_cache_info = {}
    disabled = inductor_meta.get("force_disable_caches", False)
    if (
        not disabled
        and filename is not None
        and (
            len(configs) > 1
            or inductor_meta.get("coordinate_descent_tuning")
            or dynamic_scale_rblock_eligible
        )
        and os.environ.get("TRITON_INTERPRET", "0") != "1"
    ):
        configs_hash = hash_configs(configs)

        from torch._inductor.compile_worker import watchdog

        watchdog.report_phase(watchdog.Phase.QUERYING_CACHE)
        autotune_cache = AutotuneCache.create(inductor_meta, filename, configs_hash)
        if autotune_cache:
            if best_config := autotune_cache.read_best(inductor_meta, configs):
                configs = [best_config]
                autotune_cache_info["best_config"] = triton_config_to_hashable(
                    best_config
                )
                autotune_cache_info["autotune_cache_state"] = "hit"

            else:
                autotune_cache_info["autotune_cache_state"] = "miss"
                autotune_cache_info["num_configs"] = len(configs)
                if inductor_meta.get("coordinate_descent_tuning"):
                    autotune_cache_info["coordesc_tuning"] = True
                    if len(configs) == 1:
                        # This is the config that coordinate descent tuning started at, which
                        # is not the same as the final config chosen (i.e. only_config, best_config)
                        autotune_cache_info["coordesc_tuning_start_config"] = (
                            triton_config_to_hashable(configs[0])
                        )
    else:
        if len(configs) == 1:
            autotune_cache_info["autotune_cache_state"] = "only 1 config"
            autotune_cache_info["only_config"] = triton_config_to_hashable(configs[0])

        if disabled:
            autotune_cache_info["autotune_cache_state"] = "force_disabled"
            log.debug("autotune caching is disabled by config.force_disable_caches")

    return configs, autotune_cache, autotune_cache_info


# Sentinel returned by plugin hooks to defer to the next plugin or to the
# default behavior. Tested with ``is DEFER``.
DEFER: Final[object] = object()


class CachingAutotunerPlugin:
    """Base class for ``CachingAutotuner`` plugins.

    Each hook returns ``DEFER`` to fall through to the next plugin / default
    behavior, or any other value to short-circuit ``run()`` with that value.
    """

    def pre_compile(self, autotuner: CachingAutotuner) -> object:
        """Fires at the top of ``CachingAutotuner.precompile`` (the
        parent-side, non-warm-cache-only path).

        Returning ``DEFER`` lets the standard precompile flow run
        (``_precompile_worker`` → ``_make_launchers`` →
        ``_dynamic_scale_rblock``). Returning anything else short-
        circuits ``precompile`` entirely — the plugin owns compile and
        launcher creation from this point on, including any
        ``TritonBundler`` interaction (the standard
        ``static_triton_bundle_key`` bundling path is skipped).
        """
        return DEFER

    def pre_dispatch(
        self,
        autotuner: CachingAutotuner,
        *args: object,
        stream: object,
        **kwargs: object,
    ) -> object:
        """Fires before kernel dispatch, ahead of precompile or autotune.

        Note: ``run()``'s steady-state fast path returns ``self._cached_launcher``
        directly once it is populated (after autotune converges to one
        launcher), bypassing this hook entirely on subsequent calls.
        Plugins that need every-call interception cannot rely on
        ``pre_dispatch`` alone.
        """
        return DEFER

    def pre_autotune(
        self,
        autotuner: CachingAutotuner,
        *args: object,
        stream: object,
        **kwargs: object,
    ) -> object:
        """Fires after precompile when more than one launcher remains, in
        place of ``autotune_to_one_config()``.

        Returning non-``DEFER`` short-circuits the *entire* remainder of
        ``run()`` — the plugin must perform ``_pre_launch`` / kernel
        launch / ``_post_launch`` / ``TritonBundler.put_winner`` /
        ``save_gpu_kernel`` and any post-autotune transforms (combo
        tuning, coordinate descent) itself, and return the launch
        result. Plugins that only want to influence config selection
        should mutate ``autotuner.launchers`` and return ``DEFER``
        instead.
        """
        return DEFER


def get_caching_autotuner_plugins(
    autotuner: CachingAutotuner,
) -> list[CachingAutotunerPlugin]:
    """Build the list of plugins active for ``autotuner``.

    Each plugin adds an entry here, gated on its own config flag, with
    imports kept inside the relevant branch.
    """
    plugins: list[CachingAutotunerPlugin] = []
    if autotuner.inductor_meta.get("incremental_autotune", False):
        try:
            from .fb.incremental import IncrementalAutotunePlugin

            plugins.append(IncrementalAutotunePlugin())
        except ImportError:
            pass
    return plugins


def _resolve_load_device(device: int | None, device_type: str) -> int | None:
    # compile-on-one-rank: a None device is the rank-agnostic marker; resolve it to the
    # current device at load time so a shared kernel's cubin loads on the running rank's
    # GPU rather than a baked compile-time index. CPU has no device index and its
    # DeviceGuard is a no-op -- resolving its None to 0 makes the guard call
    # exchange_device, which CPU does not implement -- so leave CPU as-is.
    if device is not None or device_type == "cpu":
        return device
    from torch._dynamo.device_interface import get_interface_for_device

    return get_interface_for_device(device_type.replace("hip", "cuda")).current_device()


class CachingAutotuner(KernelInterface):
    """
    Simplified version of Triton autotuner that has no invalidation
    key and caches the best config to disk to improve cold start times.
    Unlike the main triton Autotuner, this version can precompile all
    configs, and does not rely on the Triton JIT.
    """

    def __init__(
        self,
        fn,
        triton_meta: TritonMeta,  # passed directly to triton
        configs,
        save_cache_hook,
        mutated_arg_names: list[str],  # see [Note: clone mutated buffers]
        optimize_mem,
        heuristic_type,
        size_hints=None,
        inductor_meta: InductorMeta | None = None,  # metadata not relevant to triton
        custom_kernel=False,  # whether the kernel is inductor-generated or custom
        filename: str | None = None,
        reset_to_zero_arg_names: list[str] | None = None,
        autotune_cache_info: dict[str, Any] | None = None,
    ):
        super().__init__()

        if len(configs) == 0:
            raise AssertionError("Non-empty TritonConfig list required for compiling")
        # makes sure there are no pre-hooks on any of the triton configs
        for cfg in configs:
            validate_triton_config(cfg)

        self.fn = fn
        self.device_props: DeviceProperties = triton_meta["device"]
        # device may be None under compile-on-one-rank (rank-agnostic); it is resolved to
        # the current device at load time (see _resolve_load_device), not baked here.
        self.triton_meta: TritonMeta = cast(
            TritonMeta,
            {
                **triton_meta,
                "device": self.device_props.index,
                "device_type": self.device_props.type,
            },
        )
        self.inductor_meta: InductorMeta = (
            {} if inductor_meta is None else inductor_meta
        )
        # Add device properties to inductor_meta for use by coordinate descent tuner
        self.inductor_meta["warp_size"] = self.device_props.warp_size
        self.inductor_meta["max_threads_per_block"] = (
            self.device_props.max_threads_per_block
        )
        self.deterministic_mode = self.inductor_meta.get("deterministic", False)

        self.save_cache_hook = save_cache_hook
        self.mutated_arg_names = mutated_arg_names
        self.reset_to_zero_arg_names = (
            [] if reset_to_zero_arg_names is None else reset_to_zero_arg_names
        )
        self.optimize_mem = optimize_mem
        cached_config = lookup_autotune_config(size_hints, fn)
        self.configs = [cached_config] if cached_config else configs

        self.heuristic_type = heuristic_type
        self.custom_kernel = custom_kernel
        self.cuda_kernel_saved = False
        self.cpu_kernel_saved = False
        self.autotune_cache_info = autotune_cache_info
        if log.isEnabledFor(logging.DEBUG):
            log.debug(
                "CachingAutotuner gets %d configs for %s",
                len(self.configs),
                self.fn.__name__,
            )
            for c in self.configs:
                log.debug(c)

        self.compile_results: list[_KernelCompileResult] = []
        self.launchers: list[LauncherType] = []
        self.lock = threading.Lock()
        self.benchmark_failure_reasons: dict[Any, BenchmarkFailureReason] = {}
        if os.getenv("TRITON_CACHE_DIR") is None:
            os.environ["TRITON_CACHE_DIR"] = triton_cache_dir(
                cast(int, self.triton_meta.get("device", 0))
            )
        log.debug("Triton cache dir: %s", os.environ["TRITON_CACHE_DIR"])

        self.size_hints = size_hints
        self.is_mix_order_reduction = self.inductor_meta.get("RSPLIT_SIZE") is not None
        self.coordesc_tuner = CoordescTuner(
            is_mm=False,
            is_native_matmul=triton_meta.get("native_matmul", False),
            is_mix_order_reduction=self.is_mix_order_reduction,
            name=self.fn.__name__,
            size_hints=size_hints,
            inductor_meta=self.inductor_meta,
        )
        self.filename = filename

        # used for profiling
        self.kernel_hash: str = ""

        # Kernels are stored in the codecache with the filename as a hash of the code.
        # We rely on this to obtain the kernel hash
        if self.filename is not None:
            base_name = os.path.basename(self.filename)
            if ".py" in base_name:
                self.kernel_hash = os.path.splitext(base_name)[0]

        self.precompile_time_taken_ns = 0
        self.autotune_time_taken_ns = 0
        # Dumps the launch configs after autotuning.
        self.dump_launch_params = (
            os.environ.get("TORCHINDUCTOR_DUMP_LAUNCH_PARAMS", "0") == "1"
        )
        self.dump_launch_tensors = (
            os.environ.get("TORCHINDUCTOR_DUMP_LAUNCH_TENSORS", "0") == "1"
        )
        self.kernels_to_dump = os.environ.get(
            "TORCHINDUCTOR_KERNELS_TO_DUMP", ""
        ).split(",")

        self.triton_interpret = os.environ.get("TRITON_INTERPRET", "0") == "1"

        self._debug_call: _TritonKernelCall | None = None
        self._profiler_ctx: _RecordFunctionFast | None = None

        # Cached launcher for fast path — bypasses all preamble after first
        # successful steady-state launch.  Set to None until populated.
        self._cached_launcher: LauncherType | None = None
        # Pre-compute static eligibility for launcher caching.  These flags
        # are set once in __init__ and never change, so we avoid re-checking
        # them on every kernel launch.
        self._cache_eligible = (
            not self.triton_interpret
            and not self.dump_launch_params
            and not self.dump_launch_tensors
        )

        self._plugins = get_caching_autotuner_plugins(self)

        # Compile-time info included in runtime logginging
        self.compile_id: CompileId | None = None
        self.is_backward = False

        # Mode for launch grid calculation
        self.grid_mode: Literal["python", "cpp"] = "python"

    @staticmethod
    def _close_compiled_kernel(kernel: Any) -> None:
        if kernel is None:
            return
        close = getattr(kernel, "close", None)
        if close is not None:
            close()
            return
        module = getattr(kernel, "module", None)
        if module is not None:
            # Transitional fallback for Triton versions before CompiledKernel.close().
            # Their __del__ unloads the module and clears kernel.module, so repeated
            # calls from our idempotent benchmark cleanup paths are harmless.
            delete = getattr(kernel, "__del__", None)
            if delete is not None:
                delete()

    def release_benchmark_artifacts(self) -> None:
        for launcher in self.launchers:
            kernel = getattr(launcher, "__self__", None)
            self._close_compiled_kernel(kernel)

        for result in self.compile_results:
            self._close_compiled_kernel(getattr(result, "kernel", None))

        self.launchers = []
        self.compile_results = []
        self.benchmark_failure_reasons.clear()
        self._cached_launcher = None
        self._debug_call = None

    def is_statically_launchable(self):
        """
        Checks if every compiled kernel is statically launchable, which
        allows us to efficiently cache it in FXGraphCache
        """
        if not self.compile_results:
            return False
        return all(
            isinstance(x, StaticTritonCompileResult) for x in self.compile_results
        )

    def recheck_autotune_cache(
        self, reload_kernel_from_src: Callable[[], CachingAutotuner]
    ) -> None:
        """
        On cache load on static autotuner, we need to recheck the autotune cache, since
        a best config could have been found from a previous run
        """
        if not self.is_statically_launchable():
            raise AssertionError("Expected statically launchable kernel")

        configs = [result.config for result in self.compile_results]

        (cached_configs, _, autotune_cache_info) = check_autotune_cache(
            configs,
            self.filename,
            self.inductor_meta,
            dynamic_scale_rblock_eligible=self._could_rblock_scale,
        )
        self.autotune_cache_info = autotune_cache_info
        # I.e. there was an autotune cache hit
        if len(cached_configs) == 1:
            best_config = cached_configs[0]
            found_by_coordesc = getattr(best_config, "found_by_coordesc", False)
            # Grab the best compiled config, if it's in the list of available ones
            best_config_hash = triton_config_to_hashable(best_config)

            for compile_result in self.compile_results:
                if triton_config_to_hashable(compile_result.config) == best_config_hash:
                    compile_result.config.found_by_coordesc = found_by_coordesc
                    self.compile_results = [compile_result]
                    return

            # The best config isn't in our compile results — it was
            # found dynamically (coordesc tuning or _dynamic_scale_rblock)
            # after the static autotuner was saved. Compile it now.
            with dynamo_timed("CachingAutotuner.slow_precompile_config"):
                if self.fn.fn is None:
                    self.fn = reload_kernel_from_src().fn
                self.compile_results = [self._precompile_config(best_config)]

    def set_compile_info(self, compile_id: CompileId | None, is_backward: bool) -> None:
        self.compile_id = compile_id
        self.is_backward = is_backward

    def precompile(
        self,
        warm_cache_only=False,
        reload_kernel: Callable[[], CachingAutotuner] | None = None,
        static_triton_bundle_key: str | None = None,
    ):
        if warm_cache_only:
            self._precompile_worker()
            return
        with self.lock:
            # Helper function for reloading a kernel generated in a worker
            # in the parent class. Normally we don't need to reload the kernel
            # in the parent process, but in certain cases (coordesc tuning, dynamic_scale_rblock),
            # we need to actually run compilation on the parent process
            if reload_kernel is not None:
                self._reload_kernel = reload_kernel
            # Plugin opt-out: ``pre_compile`` returning anything other
            # than ``DEFER`` means the plugin owns compile + launcher
            # creation entirely. We return without running
            # ``_precompile_worker`` / ``_make_launchers`` /
            # ``_dynamic_scale_rblock``.
            for plugin in self._plugins:
                if plugin.pre_compile(self) is not DEFER:
                    return
            self._precompile_worker()
            if static_triton_bundle_key is not None and self.is_statically_launchable():
                TritonBundler.put_static_autotuner(static_triton_bundle_key, self)
            self._make_launchers()
            self._dynamic_scale_rblock()

    def _precompile_worker(self):
        if self.compile_results:
            for result in self.compile_results:
                TritonBundler.put(
                    triton_hash_to_path_key(result.kernel.hash),  # type: ignore[attr-defined]
                    cast(int, self.triton_meta.get("device", 0)),
                )
            return
        if self.launchers:
            raise AssertionError("launchers already populated before precompile")
        if not self.configs:
            raise NoTritonConfigsError("No triton configs are available")

        compile_results = []
        exc = None
        for c in self.configs:
            try:
                compile_results.append(self._precompile_config(c))
            except (OutOfResources, PTXASError, IntelGPUError) as e:
                exc = e
        if len(compile_results) == 0:
            raise NoTritonConfigsError(
                f"No valid triton configs. {type(exc).__name__}: {exc}"
            )
        self.compile_results = compile_results
        self.configs = None

    @functools.cached_property
    def _could_rblock_scale(self) -> bool:
        """Whether ``_dynamic_scale_rblock`` should attempt occupancy-
        driven rblock halving for this autotuner.
        """
        if "strict_sum_rblock" in self.inductor_meta:
            # Strict numerics: don't scale/tune R0_BLOCK for reductions (it shifts the order).
            return False
        return _could_dynamic_scale_rblock(
            size_hints=self.size_hints,
            heuristic_type=self.heuristic_type,
            device_prop=self.device_props,
            inductor_meta=self.inductor_meta,
        )

    @functools.cached_property
    def _combo_has_reduction_subkernel(self) -> bool:
        """True for a combo kernel (per-subkernel blocks) with a non-persistent
        reduction sub-kernel; these carry size_hints=None at the autotuner level."""
        return _combo_has_reduction_subkernel(self.inductor_meta)

    def _iter_rblock_scale_candidates(self):
        """Yield new configs with halved rblock for occupancy improvement.

        Caller is responsible for gating on ``self._could_rblock_scale``.
        """
        # TODO(jansel): we should find a way to move this extra compile into the worker process
        # Currently it relies on _make_launchers(), which requires a cuda context, to populate nreg.
        device_prop = self.device_props
        if not device_prop.regs_per_multiprocessor:
            raise AssertionError("device_prop.regs_per_multiprocessor is not set")
        if not device_prop.max_threads_per_multi_processor:
            raise AssertionError(
                "device_prop.max_threads_per_multi_processor is not set"
            )
        if not device_prop.multi_processor_count:
            raise AssertionError("device_prop.multi_processor_count is not set")
        if device_prop.warp_size is None:
            raise AssertionError("device_prop.warp_size is not set")
        seen_config_hashes: OrderedSet[Hashable] | None = None
        warp_size = device_prop.warp_size
        # Combo kernels with per-subkernel blocks pass size_hints=None and carry
        # per-subkernel xnumel_i / XBLOCK_i in combo_grid_meta. Only total_block
        # is derived differently; reduction-block selection and the occupancy
        # logic below are shared with the single-kernel path.
        combo_meta = (
            self.inductor_meta.get("combo_grid_meta")
            if self.size_hints is None
            else None
        )
        for result in self.compile_results:
            triton_config = result.config
            compiled_binary = result.kernel
            if combo_meta is not None:
                # Combo grid sums each sub-kernel's blocks (SequentialFlatten);
                # xnumel_i is None for dynamic shapes, so those dims are skipped.
                total_block = 0
                for i in range(combo_meta["num_kernels"]):
                    xnumel = combo_meta.get(f"xnumel_{i}")
                    if isinstance(xnumel, int):
                        xblock = triton_config.kwargs.get(f"XBLOCK_{i}", 1)
                        total_block += (xnumel + xblock - 1) // xblock
            else:
                if len(self.size_hints) < 2:
                    raise AssertionError(
                        f"Expected at least 2 size_hints, got {len(self.size_hints)}"
                    )
                xblock = triton_config.kwargs.get("XBLOCK", 1)
                total_block = (self.size_hints["x"] + xblock - 1) // xblock
            # Tunable reduction blocks. Combo's per-subkernel kwargs are suffixed
            # (R0_BLOCK_0, R1_BLOCK_0, R0_BLOCK_1, ...); persistent/pointwise
            # sub-kernels carry no R*_BLOCK kwarg so they drop out naturally. Same
            # rule for both paths.
            reduction_kwargs = [
                kwarg for kwarg in triton_config.kwargs if kwarg.startswith("R")
            ]
            if not reduction_kwargs:
                continue
            rblocks = [triton_config.kwargs[kwarg] for kwarg in reduction_kwargs]
            nreg = getattr(compiled_binary, "n_regs", None)
            if nreg is None:
                continue

            # make sure rblocks are not too small
            if conditional_product(*rblocks) <= 64:
                continue

            # each SM of A100 has 65536 32-bit registers. To maximize
            # the theoretical occupancy, we need run 2048 threads on each
            # SM. So each thread should use no more than 65536 / 2048
            # = 32 registers. In cases where occupancy matters, and each
            # thread uses too many registers, reduce R0_BLOCK to reduce
            # the register usage.
            # For kernel https://gist.github.com/shunting314/e4cccc031fe30d378b9b23c08c238cbd
            # from PLBartForCausalLM, latency improve from
            # 7.795ms to 4.883ms.
            #
            if (
                nreg
                <= device_prop.regs_per_multiprocessor
                // device_prop.max_threads_per_multi_processor
            ):
                continue

            nreg_per_warp = nreg * warp_size
            nreg_per_block = nreg_per_warp * triton_config.num_warps

            # Previously we set max_blocks_per_sm to 'max_threads_per_multi_processor / (warp_size * num_warps)'
            # The formula below is a tighter upper bound since we have the assumption that
            #   nreg > device_prop.regs_per_multiprocessor // device_prop.max_threads_per_multi_processor
            # due to the if condition above and:
            #   regs_per_multiprocessor / nreg_per_block
            #   = regs_per_multiprocessor / (nreg * warp_size * num_warps)
            #   < regs_per_multiprocessor / ((regs_per_multiprocessor / max_threads_per_multi_processor) * warp_size * num_warps)
            #   = max_threads_per_multi_processor / (warp_size * num_warps)
            # Using a tighter upper bound can reveal more optimization opportunities.
            max_blocks_per_sm = max(
                device_prop.regs_per_multiprocessor // nreg_per_block, 1
            )

            if total_block <= max_blocks_per_sm * device_prop.multi_processor_count:
                # no need to improve occupancy
                continue

            # Reduce the largest Rn_BLOCK by a factor of 2.
            largest_rkwarg: str = max(
                reduction_kwargs, key=triton_config.kwargs.__getitem__
            )
            new_rblock = triton_config.kwargs[largest_rkwarg] // 2
            min_rblock = self.inductor_meta.get("min_rblock")
            if (
                min_rblock is not None
                and largest_rkwarg.startswith("R0_BLOCK")
                and new_rblock < min_rblock
            ):
                continue
            new_config = copy.deepcopy(triton_config)
            new_config.kwargs[largest_rkwarg] = new_rblock

            if seen_config_hashes is None:
                seen_config_hashes = OrderedSet(
                    [triton_config_to_hashable(x.config) for x in self.compile_results]
                )
            new_config_hash = triton_config_to_hashable(new_config)
            if new_config_hash in seen_config_hashes:
                continue
            seen_config_hashes.add(new_config_hash)
            log.debug(
                "Dynamically scale down %s from TritonConfig(%s) and get a new TritonConfig(%s)",
                largest_rkwarg,
                triton_config,
                new_config,
            )
            self._ensure_kernel_loaded()
            yield new_config

            # Combo kernels additionally offer an all-reduction-blocks-halved
            # candidate. A combo's register count is the max across sub-kernel
            # branches, so a balanced combo (sub-kernels with similar register
            # use) only drops below the occupancy limit when every branch
            # shrinks, not just the largest.
            if combo_meta is not None and len(reduction_kwargs) > 1:
                all_halved = copy.deepcopy(triton_config)
                too_small = False
                for kwarg in reduction_kwargs:
                    halved = triton_config.kwargs[kwarg] // 2
                    # A block that is already 1 would halve to 0 (invalid -> Triton
                    # compile error). Skip the candidate rather than emit a 0 block.
                    if halved < 1 or (
                        min_rblock is not None
                        and kwarg.startswith("R0_BLOCK")
                        and halved < min_rblock
                    ):
                        too_small = True
                        break
                    all_halved.kwargs[kwarg] = halved
                if not too_small:
                    all_hash = triton_config_to_hashable(all_halved)
                    if all_hash not in seen_config_hashes:
                        seen_config_hashes.add(all_hash)
                        self._ensure_kernel_loaded()
                        yield all_halved

    def _dynamic_scale_rblock(self):
        if (
            self.autotune_cache_info
            and self.autotune_cache_info.get("autotune_cache_state") == "hit"
        ):
            return
        if not self._could_rblock_scale:
            return
        for new_config in self._iter_rblock_scale_candidates():
            self.compile_results.append(self._precompile_config(new_config))  # noqa: B909
        self._make_launchers()

    def compile_by_disabling_pipelining(self, config):
        self._ensure_kernel_loaded()
        cfg = copy.deepcopy(config)
        cfg.num_stages = 1
        if "NUM_STAGES" in cfg.kwargs:
            cfg.kwargs["NUM_STAGES"] = 1
        result = self._precompile_config(cfg)
        self.compile_results = [result]
        return result.make_launcher()

    def _make_launcher(
        self, compile_result: _KernelCompileResult
    ) -> tuple[LauncherType, None] | tuple[None, Exception]:
        """Create a launcher from a compile result.

        Caller must hold a DeviceGuard for the target device.
        Returns (launcher, None) on success, or (None, exception) on failure.
        """
        try:
            return compile_result.make_launcher(), None
        except (
            OutOfResources,
            PTXASError,
            torch.cuda.OutOfMemoryError,
            IntelGPUError,
        ) as e:
            return None, e

    def _make_launchers(self):
        if len(self.launchers) == len(self.compile_results):
            return

        from torch._dynamo.device_interface import DeviceGuard

        device_interface = self.get_device_interface()
        launchers = []
        exc = None
        try:
            load_device = _resolve_load_device(
                self.triton_meta["device"], self.device_props.type
            )
            # DeviceGuard ensures each launcher's binary loads onto the right device.
            with DeviceGuard(device_interface, cast(int, load_device)):
                for result in self.compile_results:
                    launcher, exc = self._make_launcher(result)
                    if launcher is not None:
                        launchers.append(launcher)
                if len(launchers) == 0:
                    result = self.compile_results[-1]
                    config = result.config
                    if (
                        isinstance(exc, (OutOfResources, torch.cuda.OutOfMemoryError))
                        and (
                            config.num_stages > 1
                            or config.kwargs.get("NUM_STAGES", 1) > 1
                        )
                        and self.inductor_meta.get("dynamic_disable_pipelining", True)
                    ):
                        self.launchers = [self.compile_by_disabling_pipelining(config)]
                        return
                    raise RuntimeError(
                        f"No valid triton configs. {type(exc).__name__}: {exc}"
                    )
            self.launchers = launchers
        finally:
            # Drop the retained failed-config exception. Holding it keeps its traceback
            # (and thus the whole benchmarking frame chain up to do_bench via f_back)
            # alive; do_bench's 256MB L2-flush buffer would then leak one-per-autotuned
            # kernel until cyclic GC runs -- which never happens under gc.disable(). We
            # only needed exc's type/message above.
            exc = None

    def _prune_compile_results_to_launcher(self, launcher: LauncherType) -> None:
        if not self.compile_results:
            return

        launcher_config = launcher.config  # type: ignore[attr-defined]
        for result in self.compile_results:
            if result.config is launcher_config:
                self.compile_results = [result]
                return

        launcher_config_hash = triton_config_to_hashable(launcher_config)
        for result in self.compile_results:
            if triton_config_to_hashable(result.config) == launcher_config_hash:
                self.compile_results = [result]
                return

        raise AssertionError(
            f"Autotuned launcher config does not match any compile result: {launcher_config}"
        )

    def _ensure_kernel_loaded(self) -> None:
        """Reload the kernel in the parent process if needed.

        When this autotuner was compiled in a worker subprocess and
        unpickled into the parent, ``prepare_for_pickle`` cleared
        ``self.fn.fn`` so the worker's JITFunction wouldn't follow
        across the pickle. Any path that needs the live function in
        the parent (coordesc, dynamic_scale_rblock, kernel_autotune
        metrics) must reload first via the stashed ``_reload_kernel``
        callback. No-op when the function is already present.
        """
        if self.fn.fn is None:
            if not hasattr(self, "_reload_kernel"):
                raise AssertionError("_reload_kernel attribute not set")
            if not callable(self._reload_kernel):
                raise AssertionError("_reload_kernel must be callable")
            self.fn = self._reload_kernel().fn

    def prepare_for_pickle(self) -> tuple[Any, ...]:
        """Drop stuff from triton.JITFunction that does not pickle.
        This must be called after precompile so that these things are no longer needed.
        Returns a tuple of old values
        """
        old_values = (
            self.fn.fn,
            self.fn.__globals__,
            self.fn.used_global_vals,
            self.fn.repr,
            self.launchers,
            getattr(self.fn, "_hash_lock", None),
            self.benchmark_failure_reasons,
        )
        self.fn.fn = None
        self.fn.__globals__ = None
        self.fn.used_global_vals = None
        self.fn.repr = _ConstRepr(self.fn.repr(self.fn))
        self.launchers = []
        self._cached_launcher = None
        self.benchmark_failure_reasons = {}
        self.fn._hash_lock = None
        return old_values

    def restore_after_unpickle(self, old_values: tuple[Any, ...] | None) -> None:
        self._cached_launcher = None
        if old_values:
            (
                self.fn.fn,
                self.fn.__globals__,
                self.fn.used_global_vals,
                self.fn.repr,
                self.launchers,
                self.fn._hash_lock,
                self.benchmark_failure_reasons,
            ) = old_values
        else:
            # even if we don't need/have specific values, we do need the
            # _hash_lock to be a valid RLock
            self.fn._hash_lock = threading.RLock()

    def prepare_for_caching(self) -> None:
        """
        Statically Launched CUDA Kernels have a raw cubin on them
        that we don't need to store in the cache(since TritonBundler handles the collection for us),
        this behavior is gated by keep_static_cubin_raw config.
        """
        # Only cubin_raw must be retained: __getstate__ already nulls cubin_path
        # on every serialize, so a cold-container load rehydrates the cubin from
        # cubin_raw (reload_cubin_path calls reload_cubin_from_raw) instead of
        # pointing at a missing file.
        if torch._inductor.config.keep_static_cubin_raw:
            return
        for result in self.compile_results:
            if isinstance(result, StaticTritonCompileResult):
                # Don't save this in the inductor cache, as it is very large
                result.kernel.cubin_raw = None

    def __getstate__(self) -> dict[str, Any]:
        if self.launchers:
            raise AssertionError("pickle should not be called after make_launchers()")
        return {
            **self.__dict__,
            "lock": None,
            "_plugins": [],
        }

    def __setstate__(self, state: dict[str, Any]) -> None:
        self.__dict__.update(state)
        self.lock = threading.Lock()
        self._plugins = get_caching_autotuner_plugins(self)

    def get_device_interface(self):
        # this code cannot run in compile workers, because it imports from torch
        from torch._dynamo.device_interface import get_interface_for_device

        return get_interface_for_device(self.device_props.type.replace("hip", "cuda"))

    def _create_compile_meta(self, cfg: Config) -> dict[str, Any]:
        """
        Create compilation metadata for a given autotuner config. This involves
        processing the Config kwargs so that the kwargs that are not part
        of the triton signature are passed in as options to triton.compile
        instead
        """
        compile_meta: dict[str, Any] = cast(
            dict[str, Any], copy.deepcopy(self.triton_meta)
        )
        compile_meta["num_warps"] = cfg.num_warps
        compile_meta["num_stages"] = cfg.num_stages

        cfg_kwargs = {**cfg.kwargs}
        if self.device_props.type == "hip":
            kernel_arg_names = OrderedSet(compile_meta["signature"])
            combo_meta = self.inductor_meta.get("combo_grid_meta") or {}
            kernel_arg_names.update(combo_meta.get("block_arg_names", ()))
            # AttrsDescriptor signatures omit constexprs, so combo block argument
            # names are carried separately in combo_grid_meta.
            # Any HIP config kwarg that is *not* in that set is not a kernel
            # argument at all; it is a backend compile option that should be forwarded
            # to triton.compile via `options`, not materialized as a constexpr.
            backend_options = {
                key: value
                for key, value in cfg_kwargs.items()
                if key not in kernel_arg_names
            }
            cfg_kwargs = {
                key: value
                for key, value in cfg_kwargs.items()
                if key in kernel_arg_names
            }
            if backend_options:
                # Stash backend-only options separately so they do not get mixed into
                # `constants`, which are interpreted as signature-bound constexpr args.
                compile_meta["backend_options"] = {
                    **compile_meta.get("backend_options", {}),
                    **backend_options,
                }
        compile_meta["constants"].update(cfg_kwargs)

        for i in get_constexprs(self.fn):
            arg_name = self.fn.arg_names[i]
            if arg_name not in compile_meta["constants"] and (
                arg_name == "num_warps" or arg_name == "num_stages"
            ):
                compile_meta["constants"][arg_name] = getattr(cfg, arg_name)
        if HAS_WARP_SPEC:
            compile_meta["num_consumer_groups"] = getattr(cfg, "num_consumer_groups", 0)
            compile_meta["num_buffers_warp_spec"] = getattr(
                cfg, "num_buffers_warp_spec", 0
            )

        host_tma_args = self.inductor_meta.get("host_tma_descriptor_args")
        if host_tma_args:
            all_constants = compile_meta["constants"]
            for key in list(compile_meta["signature"]):
                desc_info = host_tma_args.get(key)
                if desc_info is None or not isinstance(desc_info, dict):
                    continue
                block_shape_vals = _resolve_dims(
                    desc_info["block_shape"], cfg_kwargs, all_constants
                )
                shape_vals = _resolve_dims(
                    desc_info["shape"], cfg_kwargs, all_constants
                )
                stride_vals = _resolve_dims(
                    desc_info["strides"], cfg_kwargs, all_constants
                )
                if (
                    block_shape_vals is None
                    or shape_vals is None
                    or stride_vals is None
                    or any(v <= 0 for v in block_shape_vals)
                ):
                    continue
                ty = compile_meta["signature"][key]
                if isinstance(ty, str) and ty.startswith("*"):
                    dtype_str = ty[1:]
                elif isinstance(ty, str) and ty.startswith("tensordesc<"):
                    dtype_str = ty.split("<")[1].split("[")[0]
                else:
                    continue
                compile_meta["signature"][key] = (
                    f"tensordesc<{dtype_str}{list(block_shape_vals)}>"
                )

        compile_meta["debug"] = _should_enable_triton_debug_asserts(self.inductor_meta)

        # device type will be "hip" rather than "cuda" here
        compile_meta["device_type"] = self.device_props.type
        compile_meta["cc"] = self.device_props.cc

        for k in tlx_only_cuda_options():
            if v := getattr(cfg, k, None):
                compile_meta[k] = v

        return compile_meta

    def _create_compile_options(
        self, cfg: Config, compile_meta: dict[str, Any]
    ) -> dict[str, Any]:
        """
        Create options to pass to triton.compile based on the compile metadata
        and the given config.
        """
        options = {
            "num_warps": compile_meta["num_warps"],
            "num_stages": compile_meta["num_stages"],
            "debug": compile_meta["debug"],
            "sanitize_overflow": False,  # turn off additional asserts added for overflow checks
        }
        if "enable_fp_fusion" in compile_meta:
            options["enable_fp_fusion"] = compile_meta["enable_fp_fusion"]
        if HAS_WARP_SPEC:
            options.update(
                {
                    "num_consumer_groups": compile_meta.get("num_consumer_groups", 0),
                    "num_buffers_warp_spec": compile_meta.get(
                        "num_buffers_warp_spec", 0
                    ),
                }
            )
        if self.device_props.type == "cuda":
            options.update(
                {
                    "launch_cooperative_grid": compile_meta.get(
                        "launch_cooperative_grid", False
                    ),
                    "launch_pdl": compile_meta.get("launch_pdl", False),  # True
                }
            )
            if compile_meta.get("disable_ftz", False):
                options["enable_reflect_ftz"] = False
            for k in tlx_only_cuda_options():
                if v := getattr(cfg, k, None):
                    options[k] = v
        # Backend options are consumed by Triton out-of-band from the kernel
        # signature. They are intentionally *not* present in `constants`.
        options.update(compile_meta.get("backend_options", {}))

        if self.device_props.type == "xpu" and XPU_KERNEL_FORMAT == "zebin":
            options["generate_native_code"] = True

        return options

    def _precompile_config(
        self, cfg: Config, *, cc_override: str | int | None = None
    ) -> _KernelCompileResult:
        """Ahead of time compile a given autotuner config."""
        compile_meta = self._create_compile_meta(cfg)
        if cc_override is not None:
            compile_meta["cc"] = cc_override

        if self.device_props.type == "cpu":
            triton_helpers.set_driver_to_cpu()
        else:
            triton_helpers.set_driver_to_gpu()

        if not ASTSource:
            raise RuntimeError("Installed triton version too old, please upgrade")

        compile_args = (
            ASTSource(
                self.fn,
                compile_meta["signature"],
                compile_meta["constants"],
                compile_meta["configs"][0],
            ),
        )

        if self.device_props.type == "mtia":
            from mtia.host_runtime.torch_mtia.acc_flags import (  # type: ignore[import-not-found]
                build_codename,
            )

            arch = build_codename()
        else:
            arch = compile_meta["cc"]

        target = GPUTarget(
            compile_meta["device_type"],
            arch,
            self.device_props.warp_size_or_default,
        )

        options = self._create_compile_options(cfg, compile_meta)

        compile_kwargs = {
            "target": target,
            "options": options,
        }

        try:
            binary = triton.compile(*compile_args, **compile_kwargs)
        except Exception:
            log.exception(
                "Triton compilation failed: %s\n%s\nmetadata: %s",
                self.inductor_meta.get("kernel_name", "triton_"),
                self.fn.src,
                compile_meta,
            )
            raise

        # Simulate JIT Hook call
        if (
            torch._inductor.config.run_jit_post_compile_hook
            and knobs
            and getattr(knobs.runtime, "jit_post_compile_hook", None)
        ):
            try:
                hook = knobs.runtime.jit_post_compile_hook

                # base args everyone should get
                call_kwargs = dict(
                    key=getattr(self.fn, "cache_key", self.kernel_hash or str(self.fn)),
                    repr=getattr(self.fn, "src", None),
                    fn=self.fn,
                    compile=binary,
                    is_manual_warmup=False,
                    already_compiled=True,
                )

                # only add inductor_args if the hook takes it
                sig = inspect.signature(hook)
                params = sig.parameters
                if "inductor_args" in params and "config_args" in self.inductor_meta:
                    call_kwargs["inductor_args"] = self.inductor_meta["config_args"]

                hook(**call_kwargs)
            except Exception:
                log.exception("jit_post_compile_hook failed")

        TritonBundler.put(
            triton_hash_to_path_key(binary.hash),
            cast(int, self.triton_meta.get("device", 0)),
        )
        # If the binary has a cubin file to directly launch, save it on the binary
        static_launcher = StaticTritonCompileResult.can_statically_launch(
            binary, self.inductor_meta, self.triton_meta, self.heuristic_type
        )

        if static_launcher is not None:
            result = StaticTritonCompileResult(
                static_launcher, cfg, compile_meta, self.inductor_meta
            )
            return result

        return TritonCompileResult(binary, cfg, compile_meta, self.inductor_meta)

    def bench(self, launcher, *args, with_profiler=False, **kwargs):
        """Measure the performance of a given launcher."""
        # we don't skip configs with spilled registers when auto-tuning custom
        # (user-written) Triton kernels, as (i) we don't have any knowledge or
        # control over the kernel code; (ii) there is empirical evidence that
        # for some (complicated) custom Triton kernels, a register-spilling
        # config may yield the best latency.
        if (
            not self.custom_kernel
            and launcher.n_spills is not None
            and launcher.n_spills
            > self.inductor_meta.get("spill_threshold", 32 if torch.version.hip else 16)
        ):
            log.debug(
                "Skip config %s because of register spilling: %d",
                launcher.config,
                launcher.n_spills,
            )
            self.benchmark_failure_reasons[launcher] = (
                BenchmarkFailureReason.REGISTER_SPILLING
            )
            return float("inf")

        device_interface = self.get_device_interface()

        cpu_copies = self.copy_args_to_cpu_if_needed(*args, **kwargs)

        def kernel_call():
            # Resolve the raw stream at call time rather than at closure-creation
            # time so that CUDA graph capture on a different stream works correctly.
            stream = device_interface.get_raw_stream(device_interface.current_device())
            cloned_args, cloned_kwargs = self.maybe_clone_args(
                cpu_copies, *args, **kwargs
            )
            kernel_name = self.inductor_meta.get("kernel_name", "triton kernel")
            # reset to zero before evaluating any config
            self.reset_to_zero_args(*args, **kwargs)
            if autograd_profiler._is_profiler_enabled:
                profiler_kwargs = self.get_profiler_kwargs(stream, launcher)
                with torch._C._profiler._RecordFunctionFast(
                    kernel_name,
                    cloned_args,
                    profiler_kwargs,
                ):
                    try:
                        launcher(
                            *cloned_args,
                            **cloned_kwargs,
                            stream=stream,
                        )
                    except Exception as e:
                        if isinstance(e, TypeError):
                            self._check_launcher_call_args(launcher, cloned_args)
                        log.error(
                            "Failed during launch %s with config: %s (num_warps=%s, num_stages=%s, kwargs=%s)",
                            kernel_name,
                            launcher.config,
                            launcher.config.num_warps,
                            launcher.config.num_stages,
                            launcher.config.kwargs,
                        )
                        raise

            else:
                try:
                    launcher(
                        *cloned_args,
                        **cloned_kwargs,
                        stream=stream,
                    )
                except Exception as e:
                    if isinstance(e, TypeError):
                        self._check_launcher_call_args(launcher, cloned_args)
                    log.error(
                        "Failed during launch %s with config: %s (num_warps=%s, num_stages=%s, kwargs=%s)",
                        kernel_name,
                        launcher.config,
                        launcher.config.num_warps,
                        launcher.config.num_stages,
                        launcher.config.kwargs,
                    )
                    raise
            self.restore_args_from_cpu(cpu_copies)

        # only use profiler when not already in a profiler instance
        if with_profiler and not autograd_profiler._is_profiler_enabled:
            from torch._inductor.utils import do_bench_using_profiling

            return do_bench_using_profiling(kernel_call, warmup=10, rep=40)

        benchmark_kwargs = (
            {}
            if self.device_props.type == "cpu"
            else {"rep": 40, "is_vetted_benchmarking": True}
        )
        result = benchmarker.benchmark(
            fn=kernel_call,
            device=self.device_props.type,
            **benchmark_kwargs,  # type: ignore[arg-type]
        )
        # benchmarker.benchmark() only returns float("inf") when catching an
        # "invalid configuration" exception - all other exceptions are re-raised.
        # Therefore, if result is inf here, it must be due to invalid config.
        if result == float("inf"):
            self.benchmark_failure_reasons[launcher] = (
                BenchmarkFailureReason.INVALID_CONFIG
            )
        return result

    def copy_args_to_cpu_if_needed(self, *args, **kwargs):
        """
        To support benchmarking in the presence of mutated args, we need to avoid
        autotuning contaminating them. We try to pass cloned args to the kernel.
        If those clones would increase the peak memory usage, however, we instead
        copy to cpu and restore them after each iteration. Figure out the args
        to be copied and do the copying.
        """
        if not self.optimize_mem:
            return {}

        copies = {}
        try:
            if torch.accelerator.current_accelerator() is None:
                # No initialized accelerator; skip memory-optimized path
                return {}
            budget = (
                torch.accelerator.max_memory_allocated()
                - torch.accelerator.memory_allocated()
            )
        except RuntimeError:
            # Possibly a custom CUDA allocator, see https://github.com/pytorch/pytorch/issues/163257
            return {}

        def maybe_copy(name, arg):
            if name in self.mutated_arg_names and arg.device.type in (
                "cuda",
                "xpu",
            ):
                nonlocal budget
                if not isinstance(arg, torch.Tensor):
                    raise AssertionError(
                        f"Expected torch.Tensor for mutated arg, got {type(arg)}"
                    )
                required_storage_length = compute_required_storage_length(
                    arg.size(),
                    arg.stride(),
                    0,
                )
                size = required_storage_length * arg.element_size()
                if size > budget:
                    cpu_arg = torch.empty_strided(
                        (required_storage_length,),
                        (1,),
                        dtype=arg.dtype,
                        device="cpu",
                        pin_memory=True,
                    )
                    cpu_arg.copy_(
                        arg.as_strided((required_storage_length,), (1,)),
                        non_blocking=True,
                    )
                    copies[name] = (arg, cpu_arg)
                else:
                    budget -= size

        for name, arg in zip(self.fn.arg_names, args):
            maybe_copy(name, arg)

        for name, arg in kwargs.items():
            maybe_copy(name, arg)

        return copies

    def restore_args_from_cpu(self, cpu_copies):
        for pair in cpu_copies.values():
            arg, cpu_arg = pair
            required_storage_length = compute_required_storage_length(
                arg.size(),
                arg.stride(),
                0,
            )
            arg.as_strided((required_storage_length,), (1,)).copy_(
                cpu_arg, non_blocking=True
            )

    def reset_to_zero_args(self, *args, **kwargs):
        if not self.reset_to_zero_arg_names:
            return
        for i, arg in enumerate(args):
            if self.fn.arg_names[i] in self.reset_to_zero_arg_names:
                if not isinstance(arg, torch.Tensor):
                    raise AssertionError(
                        "self.reset_to_zero_arg_names should only contain valid argument names"
                    )
                arg.zero_()

        for name, arg in kwargs.items():
            if name in self.reset_to_zero_arg_names:
                if not isinstance(arg, torch.Tensor):
                    raise AssertionError(
                        "self.reset_to_zero_arg_names should only contain valid argument names"
                    )
                arg.zero_()

    def maybe_clone_args(
        self, exclude: Container[str], *args, **kwargs
    ) -> tuple[list[Any], dict[str, Any]]:
        """
        Prepare new args and kwargs by cloning any in-place buffers
        (that are not in the provided exclusion list), to avoid autotune
        contaminating them. Avoid cloning the other buffers because it
        leads to increased memory usage.
        """
        from ..compile_fx import clone_preserve_strides

        def prepare_arg(name, arg):
            if name in self.mutated_arg_names and name not in exclude:
                if not isinstance(arg, torch.Tensor):
                    raise AssertionError(
                        f"Expected torch.Tensor for mutated arg '{name}', got {type(arg)}"
                    )
                return clone_preserve_strides(arg)
            else:
                return arg

        cloned_args = [
            prepare_arg(name, arg)
            for name, arg in itertools.zip_longest(self.fn.arg_names[: len(args)], args)
        ]
        cloned_kwargs = {name: prepare_arg(name, arg) for name, arg in kwargs.items()}
        return cloned_args, cloned_kwargs

    def clone_args(self, *args, **kwargs) -> tuple[list[Any], dict[str, Any]]:
        return self.maybe_clone_args(OrderedSet(), *args, **kwargs)

    @staticmethod
    def _close_static_launcher(launcher: LauncherType) -> None:
        if not getattr(launcher, "_is_static", False):
            return
        runner = getattr(launcher, "__globals__", {}).get("runner")
        kernel = getattr(runner, "__self__", None)
        close = getattr(kernel, "close", None)
        if close is not None:
            close()

    def _release_static_launchers_except(self, keep_launcher: LauncherType) -> None:
        for launcher in self.launchers:
            if launcher is not keep_launcher:
                self._close_static_launcher(launcher)
        if not getattr(keep_launcher, "_is_static", False):
            return
        keep_hash = getattr(keep_launcher, "cache_hash", None)
        if keep_hash is None:
            return
        keep_results = [
            result
            for result in self.compile_results
            if isinstance(result, StaticTritonCompileResult)
            and triton_hash_to_path_key(result.kernel.hash) == keep_hash
        ]
        if len(keep_results) == 1:
            for result in self.compile_results:
                if result is not keep_results[0] and isinstance(
                    result, StaticTritonCompileResult
                ):
                    close = getattr(result.kernel, "close", None)
                    if close is not None:
                        close()
            self.compile_results = keep_results

    def _log_autotune_inputs(self, args, kwargs) -> None:
        """Log input tensor shapes/dtypes and scalar values for autotuning."""
        kernel_name = self.inductor_meta.get("kernel_name", self.fn.__name__)
        signature = self.triton_meta.get("signature", {})
        arg_names = list(signature.keys())

        autotuning_inputs_log.debug("=" * 60)
        autotuning_inputs_log.debug("Autotuning inputs for kernel: %s", kernel_name)
        autotuning_inputs_log.debug("=" * 60)
        autotuning_inputs_log.debug("  Heuristic type: %s", self.heuristic_type)
        autotuning_inputs_log.debug("  Size hints: %s", self.size_hints)
        autotuning_inputs_log.debug(
            "  Num configs to benchmark: %d", len(self.launchers)
        )
        autotuning_inputs_log.debug(
            "  Device: %s (index=%s)", self.device_props.type, self.device_props.index
        )
        autotuning_inputs_log.debug("-" * 60)
        autotuning_inputs_log.debug("Arguments:")

        for i, arg in enumerate(args):
            arg_name = arg_names[i] if i < len(arg_names) else f"arg_{i}"
            arg_signature = signature.get(arg_name, "unknown")
            if isinstance(arg, torch.Tensor):
                autotuning_inputs_log.debug(
                    "  [%d] %s (%s): Tensor(shape=%s, dtype=%s, device=%s, stride=%s, contiguous=%s)",
                    i,
                    arg_name,
                    arg_signature,
                    tuple(arg.shape),
                    arg.dtype,
                    arg.device,
                    arg.stride(),
                    arg.is_contiguous(),
                )
            elif isinstance(arg, (int, float, bool)):
                autotuning_inputs_log.debug(
                    "  [%d] %s (%s): %s (type=%s)",
                    i,
                    arg_name,
                    arg_signature,
                    arg,
                    type(arg).__name__,
                )
            else:
                autotuning_inputs_log.debug(
                    "  [%d] %s (%s): %s (type=%s)",
                    i,
                    arg_name,
                    arg_signature,
                    repr(arg)[:100],
                    type(arg).__name__,
                )

        if kwargs:
            autotuning_inputs_log.debug("-" * 60)
            autotuning_inputs_log.debug("Keyword arguments:")
            for k, v in kwargs.items():
                if isinstance(v, torch.Tensor):
                    autotuning_inputs_log.debug(
                        "  %s: Tensor(shape=%s, dtype=%s, device=%s)",
                        k,
                        tuple(v.shape),
                        v.dtype,
                        v.device,
                    )
                else:
                    autotuning_inputs_log.debug(
                        "  %s: %s (type=%s)", k, v, type(v).__name__
                    )

        autotuning_inputs_log.debug("=" * 60)

    def benchmark_all_configs(self, *args, **kwargs):
        with (
            dynamo_timed(
                "CachingAutotuner.benchmark_all_configs",
                log_pt2_compile_event=True,
                metadata={"kernel_name": self.inductor_meta.get("kernel_name")},
                dynamo_compile_column_us="runtime_triton_autotune_time_us",
                compile_id=self.compile_id,
                is_backward=self.is_backward,
                log_waitcounter=True,
                waitcounter_name_override="triton_autotuner",
            ),
            # Temporarily disable due to spam
            # compilation_callback.callback_handler.install_callbacks(
            #     compilation_callback.CallbackTrigger.TRITON_AUTOTUNING,
            #     str(self.compile_id),
            # ),
        ):
            timings = {}
            best_launcher = None
            best_timing = float("inf")
            for launcher in self.launchers:
                timing = self.bench(launcher, *args, **kwargs)
                timings[launcher] = timing
                # Close losing static launchers eagerly so exhaustive autotuning
                # keeps only the current winner and candidate modules loaded.
                if best_launcher is None or timing < best_timing:
                    if best_launcher is not None:
                        self._close_static_launcher(best_launcher)
                    best_launcher = launcher
                    best_timing = timing
                else:
                    self._close_static_launcher(launcher)

            for k, v in timings.items():
                self.coordesc_tuner.cache_benchmark_result(k.config, v)

            if log.isEnabledFor(logging.DEBUG):
                log.debug("Benchmark all input configs for %s, get:", self.fn.__name__)
                for k, v in timings.items():
                    log.debug(
                        "%s: %f, nreg %d, nspill %d, #shared-mem %s",
                        k.config,
                        v,
                        k.n_regs,
                        k.n_spills,
                        k.shared,
                    )

            if metrics.is_metric_table_enabled("kernel_autotune"):
                self._ensure_kernel_loaded()

                kernel_path = self.fn.fn.__code__.co_filename
                kernel_name = self.fn.__name__

                for k, v in timings.items():
                    metrics.log_kernel_autotune_result(
                        kernel_path, kernel_name, k.config, v
                    )

            self.reset_to_zero_args(*args, **kwargs)
            return timings

    def autotune_to_one_config(self, *args, **kwargs):
        """Execute autotuning to select the optimal kernel configuration."""
        if autotuning_inputs_log.isEnabledFor(logging.DEBUG):
            self._log_autotune_inputs(args, kwargs)

        start_time = time.time_ns()
        timings = self.benchmark_all_configs(*args, **kwargs)
        benchmark_time_taken_ns = time.time_ns() - start_time

        # Check if any configs failed (have inf timing) and log which one was selected
        failed_launchers = [
            launcher for launcher, timing in timings.items() if timing == float("inf")
        ]
        if failed_launchers:
            valid_timings = [(k, v) for k, v in timings.items() if v != float("inf")]
            if valid_timings:
                best_launcher, best_time = min(valid_timings, key=lambda x: x[1])

                # Count failures by reason
                spill_count = sum(
                    1
                    for launcher in failed_launchers
                    if self.benchmark_failure_reasons.get(launcher)
                    == BenchmarkFailureReason.REGISTER_SPILLING
                )
                invalid_config_count = sum(
                    1
                    for launcher in failed_launchers
                    if self.benchmark_failure_reasons.get(launcher)
                    == BenchmarkFailureReason.INVALID_CONFIG
                )

                reason_parts = []
                if spill_count > 0:
                    reason_parts.append(f"{spill_count} register spilling")
                if invalid_config_count > 0:
                    reason_parts.append(f"{invalid_config_count} invalid config")
                reason_str = ", ".join(reason_parts) if reason_parts else "unknown"

                log.info(
                    "Skipped %d/%d configs for %s (%s). Selected: %s (%.4f ms)",
                    len(failed_launchers),
                    len(timings),
                    self.fn.__name__,
                    reason_str,
                    best_launcher.config,
                    best_time,
                )

        best_launcher = builtins.min(timings, key=timings.get)
        self._release_static_launchers_except(best_launcher)
        self.launchers = [best_launcher]
        self._prune_compile_results_to_launcher(best_launcher)
        self.autotune_time_taken_ns = (
            self.precompile_time_taken_ns + benchmark_time_taken_ns
        )

        # log the best config
        launcher = self.launchers[0]
        log.debug(
            "Best config for %s: %s: %f, nreg %d, nspill %d, #shared-mem %s",
            self.fn.__name__,
            launcher.config,
            timings[launcher],
            launcher.n_regs,
            launcher.n_spills,
            launcher.shared,
        )

        TritonBundler.put_winner(launcher.cache_hash)

        if self.save_cache_hook:
            self.save_cache_hook(
                launcher.config,
                self.autotune_time_taken_ns,
                found_by_coordesc=self.inductor_meta.get(
                    "coordinate_descent_tuning", False
                ),
                triton_cache_hash=launcher.cache_hash,
            )

    def _combo_sequential_autotune(self, launcher, *args, **kwargs):
        """
        Chain block-size decisions for combo kernels: tune one group at a time,
        each step building on the previous winner.

        Phase 1: Tune block sizes with warps/stages fixed from the base config.
        Phase 2: Re-tune warps/stages with finalized block sizes.
        """
        combo_tuning_groups = self.inductor_meta.get("combo_tuning_groups")
        if not combo_tuning_groups:
            return launcher

        self._ensure_kernel_loaded()

        combo_meta = self.inductor_meta.get("combo_grid_meta") or {}
        block_arg_names = OrderedSet(combo_meta.get("block_arg_names", ()))
        best_config = launcher.config
        current_kwargs = dict(best_config.kwargs)
        base_num_warps = best_config.num_warps
        base_num_stages = best_config.num_stages

        start_time = time.time_ns()
        best_time = self.bench(launcher, *args, **kwargs)
        counters["inductor"]["combo_autotune_bench"] += 1
        self.coordesc_tuner.cache_benchmark_result(launcher.config, best_time)
        log.debug(
            "  Phase 1 baseline: %s warps=%d time=%f",
            dict(current_kwargs),
            base_num_warps,
            best_time,
        )

        # Phase 1: Tune block sizes per sub-kernel (largest first).
        # warps/stages stay fixed at base config values.
        for gi, group in enumerate(combo_tuning_groups):
            member_indices = group["member_indices"]
            cfgs = group["configs"]
            skip_rblock = group["skip_rblock"]

            if len(cfgs) <= 1:
                log.debug("  Phase 1 group %d SK%s: 1 config, skip", gi, member_indices)
                continue

            log.debug(
                "  Phase 1 group %d SK%s: trying %d configs, current_kwargs=%s",
                gi,
                member_indices,
                len(cfgs),
                dict(current_kwargs),
            )
            for ci, cfg in enumerate(cfgs):
                trial_kwargs = dict(current_kwargs)
                for idx in member_indices:
                    _update_combo_kernel_kwargs(
                        trial_kwargs, cfg.kwargs, idx, skip_rblock, block_arg_names
                    )

                if trial_kwargs == current_kwargs:
                    log.debug("    cfg[%d] skip (same as current)", ci)
                    continue

                trial_config = triton.Config(
                    trial_kwargs,
                    num_warps=base_num_warps,
                    num_stages=base_num_stages,
                )

                with self.lock:
                    trial_launcher = self._precompile_config(
                        trial_config
                    ).make_launcher()
                trial_time = self.bench(trial_launcher, *args, **kwargs)
                counters["inductor"]["combo_autotune_bench"] += 1
                self.coordesc_tuner.cache_benchmark_result(trial_config, trial_time)

                improved = trial_time < best_time
                log.debug(
                    "    cfg[%d] trial=%s time=%f%s",
                    ci,
                    dict(trial_kwargs),
                    trial_time,
                    " (BETTER)" if improved else "",
                )
                if improved:
                    best_time = trial_time
                    launcher = trial_launcher
                    current_kwargs = trial_kwargs

            log.debug(
                "  Phase 1 group %d winner: current_kwargs=%s",
                gi,
                dict(current_kwargs),
            )

        # Phase 2: Re-tune num_warps/num_stages with finalized block sizes.
        # Block sizes are now optimal — find the best warp/stage pair for them.
        warp_stage_candidates = self.inductor_meta.get("combo_warp_stage_candidates")
        log.debug(
            "  Phase 2: blocks=%s, trying %d warp/stage pairs",
            dict(current_kwargs),
            len(warp_stage_candidates),
        )
        best_warps = launcher.config.num_warps
        best_stages = launcher.config.num_stages
        for num_warps, num_stages in warp_stage_candidates:
            if num_warps == best_warps and num_stages == best_stages:
                log.debug(
                    "    warps=%d stages=%d skip (same as current)",
                    num_warps,
                    num_stages,
                )
                continue

            trial_config = triton.Config(
                dict(current_kwargs),
                num_warps=num_warps,
                num_stages=num_stages,
            )
            with self.lock:
                trial_launcher = self._precompile_config(trial_config).make_launcher()
            trial_time = self.bench(trial_launcher, *args, **kwargs)
            counters["inductor"]["combo_autotune_bench"] += 1
            self.coordesc_tuner.cache_benchmark_result(trial_config, trial_time)

            improved = trial_time < best_time
            log.debug(
                "    warps=%d stages=%d time=%f%s",
                num_warps,
                num_stages,
                trial_time,
                " (BETTER)" if improved else "",
            )
            if improved:
                best_time = trial_time
                launcher = trial_launcher
                best_warps = num_warps
                best_stages = num_stages

        log.debug(
            "Combo sequential autotune for %s: best config %s, time %f",
            self.fn.__name__,
            launcher.config,
            best_time,
        )
        launcher.config.found_by_combo_autotune = True
        self.autotune_time_taken_ns += time.time_ns() - start_time
        if self.save_cache_hook:
            self.save_cache_hook(launcher.config, self.autotune_time_taken_ns)
        return launcher

    def save_gpu_kernel(self, stream, launcher):
        """Save the selected GPU kernel binary and assembly for AOTI packaging."""
        key = self.inductor_meta.get("kernel_name", None)  # unique kernel name
        if key is None:
            raise AssertionError("kernel_name can not be None")
        binary = launcher.bin

        from torch._inductor import config as inductor_config

        target_cc = None
        cuda_arch = None
        if (
            self.device_props.type == "cuda"
            and torch.version.hip is None
            and inductor_config.aot_inductor.emit_multi_arch_kernel
        ):
            from torch._inductor.codegen.cuda import compile_utils as cuda_compile_utils

            cuda_arch = cuda_compile_utils._aoti_cuda_target_arch()
            if inductor_config.cuda.arch is not None:
                target_cc = cuda_compile_utils._cuda_arch_number(cuda_arch)
                device_cc = cuda_compile_utils._cuda_arch_number(
                    str(self.device_props.cc)
                )
                if target_cc == device_cc:
                    target_cc = None

        if target_cc is not None:
            # Autotuning must run kernels compiled for the actual GPU, but AOTI
            # needs to package PTX/cubin for the requested deployment target.
            # Recompile the selected config for the target arch without loading
            # or benchmarking that binary on the current device.
            target_result = self._precompile_config(
                launcher.config, cc_override=target_cc
            )
            if isinstance(target_result, TritonCompileResult):
                binary = target_result.kernel
            else:
                raise RuntimeError(
                    "AOTI CUDA target-arch packaging requires a Triton binary"
                )

        # Prefer Level 0 launch metadata schema (versioned, stable contract)
        # over hasattr probing of CompiledKernel internals.
        # TODO: When the AOTI C++ launch path gains cuLaunchKernelEx support for
        # CTA clusters, add num_ctas/cluster_dims here from the schema.
        # Currently num_ctas is already captured via config_to_dict(launcher.config)
        # for scratch space scaling, but is not used in the actual kernel launch.
        schema = getattr(binary, "launch_metadata_schema", None)
        if schema is not None and inductor_config.use_launch_metadata_schema:
            params: dict[str, Any] = {
                "mangled_name": schema["entry_name"],
                "num_warps": schema["num_warps"],
                "shared_mem": schema["shared_mem"],
                "stream": stream,
                "config": config_to_dict(launcher.config),
                "inductor_meta": self.inductor_meta,
                "triton_meta": self.triton_meta,
                "def_args": launcher.def_args,
                "call_args": launcher.call_args,
                "global_scratch": launcher.global_scratch,
                "profile_scratch": launcher.profile_scratch,
                "cuda_arch": cuda_arch,
            }
        else:
            # Fallback: hasattr probing for older Triton versions
            params: dict[str, Any] = {
                "mangled_name": (
                    binary.metadata.name
                    if hasattr(binary.metadata, "name")
                    else binary.metadata["name"]
                ),
                "num_warps": (
                    binary.num_warps
                    if hasattr(binary, "num_warps")
                    else binary.metadata.num_warps
                ),
                "shared_mem": (
                    binary.shared
                    if hasattr(binary, "shared")
                    else binary.metadata.shared
                ),
                "stream": stream,
                "config": config_to_dict(launcher.config),
                "inductor_meta": self.inductor_meta,
                "triton_meta": self.triton_meta,
                "def_args": launcher.def_args,
                "call_args": launcher.call_args,
                "global_scratch": launcher.global_scratch,
                "profile_scratch": launcher.profile_scratch,
                "cuda_arch": cuda_arch,
            }

        from torch._inductor.codecache import CudaKernelParamCache

        bin_type = {"hip": "hsaco", "xpu": XPU_KERNEL_FORMAT}.get(
            self.device_props.type, "cubin"
        )
        kernel_binary = binary.asm[bin_type]

        # ROCm multi-arch: capture LLVM IR
        if torch.version.hip and inductor_config.aot_inductor.emit_multi_arch_kernel:
            # Multi-arch ROCm: Capture LLVM IR for cross-architecture compilation
            asm_type = "ll"

            # llir is the key to obtain LLVM IR from triton
            asm = binary.asm.get("llir", None)

            # CRITICAL: Multi-arch compilation cannot proceed without LLVM IR
            # Fail fast with clear error message pointing to the issue
            if not asm:
                available_keys = list(binary.asm.keys())
                raise RuntimeError(
                    f"ROCm multi-arch requires LLVM IR, but none found. "
                    f"Available keys: {available_keys}. "
                    f"Triton may need to be patched to emit LLVM IR."
                )

        # Everything else: capture architecture-specific assembly
        else:
            asm_type = {"hip": "amdgcn", "cuda": "ptx", "xpu": "spv"}.get(
                self.device_props.type
            )
            asm = binary.asm.get(asm_type, None)

        CudaKernelParamCache.set(key, params, kernel_binary, bin_type, asm, asm_type)
        self.cuda_kernel_saved = True

    def save_cpu_kernel(self, launcher):
        """AOTI counterpart of save_gpu_kernel for CPU Triton kernels.

        Captures the kernel and launcher `.so` files into
        `CpuTritonKernelCache` for `CppWrapperCpu` to dlopen at runtime.
        Triton CPU backend needs to emit `run_from_nativert`.
        """
        from torch._inductor.codecache import CpuTritonKernelCache

        key = self.inductor_meta.get("kernel_name")
        if key is None:
            raise AssertionError("kernel_name can not be None")

        compiled = launcher.bin
        kernel_bytes = compiled.asm.get("so")
        launcher_bytes = compiled.asm.get("launcher.so")
        if kernel_bytes is None or launcher_bytes is None:
            raise RuntimeError(
                f"CPU AOTI requires a Triton CPU backend that emits a launcher "
                f"`.so` exporting `run_from_nativert`; kernel '{key}' has "
                f"compiled.asm keys {list(compiled.asm.keys())}."
            )
        kernel_symbol = (
            compiled.metadata.name
            if hasattr(compiled.metadata, "name")
            else compiled.metadata["name"]
        )
        signature = compiled.src.signature

        CpuTritonKernelCache.set(
            key,
            kernel_bytes=kernel_bytes,
            launcher_bytes=launcher_bytes,
            kernel_symbol=kernel_symbol,
            signature=signature,
        )
        self.cpu_kernel_saved = True

    @functools.cached_property
    def _should_coordesc_tune(self) -> bool:
        """Whether this autotuner is eligible for coordinate descent tuning."""
        if self.heuristic_type in (
            HeuristicType.TEMPLATE,
            HeuristicType.USER_AUTOTUNE,
            HeuristicType.FIXED,
        ):
            return False
        # Deterministic mode (and strict numerics) forbid tuning RBLOCK / num_warps for
        # reductions because those knobs shift numerics.
        if (
            self.deterministic_mode or "strict_sum_rblock" in self.inductor_meta
        ) and self.heuristic_type in (
            HeuristicType.REDUCTION,
            HeuristicType.PERSISTENT_REDUCTION,
            HeuristicType.SPLIT_SCAN,
        ):
            return False
        return True

    def coordinate_descent_tuning(self, launcher, *args, **kwargs):
        """
        Coordinate descent tuning can be run with or without max-autotune.

        The only difference between these two is the starting config for coordinate_descent tuning.
        E.g., assuming regular autotune only get one config C1; while max-autotune get 4 configs C1, C2, C3, C4
        and max-autotune figure out C3 is the best.

        Then if coordinate descent tuning is run with max-autotune disabled, it will start from C1;
        while if coordinate descent tuning is run with max-autotune enabled, it will start from C3.
        """
        if not self._should_coordesc_tune:
            return launcher

        with dynamo_timed(
            "CachingAutotuner.coordinate_descent_tuning",
            # These generate too many pt2_compile_event logs:
            log_pt2_compile_event=False,
            metadata={"kernel_name": self.inductor_meta.get("kernel_name")},
            dynamo_compile_column_us="runtime_triton_autotune_time_us",
            compile_id=self.compile_id,
            is_backward=self.is_backward,
            log_waitcounter=True,
            waitcounter_name_override="triton_autotuner",
        ):
            return self._coordinate_descent_tuning(launcher, *args, **kwargs)

    def _coordinate_descent_tuning(self, launcher, *args, **kwargs):
        config2launcher = {launcher.config: launcher}

        self._ensure_kernel_loaded()

        def benchmark_one_config(config):
            with self.lock:
                launcher = self._precompile_config(config).make_launcher()
            config2launcher[config] = launcher

            out = self.bench(launcher, *args, **kwargs)
            counters["inductor"]["coordesc_tuning_bench"] += 1
            log.debug(
                "COORDESC: %s: %f, nreg %d, nspill %d, #shared-mem %d",
                launcher.config,
                out,
                launcher.n_regs,
                launcher.n_spills,
                launcher.shared,
            )
            return out

        if (
            self.heuristic_type == HeuristicType.PERSISTENT_REDUCTION
            and "R0_BLOCK" in launcher.config.kwargs
        ):
            raise AssertionError(
                "Coordinate descent tuner relies on the assumption that persistent reduction's triton config does not have R0_BLOCK"
            )
        start_time = time.time_ns()
        best_config = self.coordesc_tuner.autotune(
            benchmark_one_config, launcher.config, None
        )
        coordesc_time_taken_ns = time.time_ns() - start_time
        best_config.found_by_coordesc = True

        if self.save_cache_hook:
            self.save_cache_hook(
                best_config,
                self.autotune_time_taken_ns + coordesc_time_taken_ns,
                found_by_coordesc=True,
            )

        if best_config not in config2launcher:
            # On a Coordesc cache hit, we might not have loaded the launcher
            # This can happen because PyCodeCache saves CachingAutotuners in memory,
            # even for separate compile IDs (which can have different inputs without changing output code)
            config2launcher[best_config] = self._precompile_config(
                best_config
            ).make_launcher()

        winner = config2launcher[best_config]
        TritonBundler.put_winner(winner.cache_hash)

        fn_hash = generate_lookup_hash_from_source_code(
            str(self.size_hints), self.fn.src
        )
        log.debug("Function hash %s has best config %s", fn_hash, best_config)
        return winner

    def get_profiler_kwargs(self, stream, launcher):
        kernel_kwargs_str = ",".join(
            f"{k}={v}" for (k, v) in launcher.config.kwargs.items()
        )

        ret = {
            "kernel_file": (self.filename or ""),
            "kernel_hash": self.kernel_hash,
            "kernel_backend": "triton",
            "stream": stream,
            "num_warps": launcher.config.num_warps,
            "num_stages": launcher.config.num_stages,
            "kernel_kwargs": kernel_kwargs_str,
        }
        if "kernel_name" in self.inductor_meta:
            ret["kernel_name"] = self.inductor_meta["kernel_name"]
        if "kernel_flop" in self.inductor_meta:
            ret["kernel_flop"] = self.inductor_meta["kernel_flop"]
        if "kernel_num_gb" in self.inductor_meta:
            ret["kernel_num_gb"] = self.inductor_meta["kernel_num_gb"]
        return ret

    def _pre_launch(self, launcher, *args, stream, **kwargs):
        """Pre-launch instrumentation: param/tensor dumping and profiler context entry."""
        if self.dump_launch_params:
            new_args, grid = self._interpret_args_grid(args, launcher.config)
            _dump_launch_params(new_args, kwargs, launcher, self.fn.__name__, grid)

        if self.dump_launch_tensors:
            if not self.kernels_to_dump or any(
                kernel_name in self.fn.__name__ for kernel_name in self.kernels_to_dump
            ):
                _dump_launch_tensors(
                    args, self.filename, self.kernel_hash, self.fn.__name__
                )

        if autograd_profiler._is_profiler_enabled:
            profiler_kwargs = self.get_profiler_kwargs(stream, launcher)
            profiler_ctx = torch._C._profiler._RecordFunctionFast(
                self.inductor_meta.get("kernel_name", "triton kernel"),
                tuple(args),
                profiler_kwargs,
            )
            profiler_ctx.__enter__()
            # set ctx after enter succeeds
            self._profiler_ctx = profiler_ctx
        else:
            self._profiler_ctx = None

    def _post_launch(self) -> None:
        """Post-launch instrumentation: profiler context exit and debug mode finalization."""
        if (profiler_ctx := self._profiler_ctx) is not None:
            self._profiler_ctx = None
            profiler_ctx.__exit__(None, None, None)
        if (debug_call := self._debug_call) is not None:
            self._debug_call = None
            debug_call.finalize(self.get_device_interface())

    def run(
        self,
        *args,
        stream,
        benchmark_run=False,
        **kwargs,
    ):  # type:ignore[override]
        """Launch triton kernel call and return result."""
        # --- FAST PATH ---
        # After the first successful launch in steady state, cache the launcher
        # and skip all preamble on subsequent calls (~2µs savings).
        # Conditions here deliberately differ from the cache-population block
        # below: we re-check dynamic conditions (profiler, debug mode) every
        # call so enabling them at runtime falls back to the slow path.
        # Static conditions (interpret, dump flags, launcher count) were
        # already validated when the cache was populated.
        # `not kwargs` is checked here but not at population time because
        # inductor steady-state never passes kwargs; if it somehow does, the
        # slow path handles it correctly.
        fast = self._cached_launcher
        if (
            fast is not None
            and not benchmark_run
            and not kwargs
            and not autograd_profiler._is_profiler_enabled
            and not get_active_debug_mode()
        ):
            return fast(*args, stream=stream)

        debug_mode = get_active_debug_mode()
        if debug_mode:
            arg_names = list(self.triton_meta.get("signature", {}).keys())
            kernel_kwargs = dict(zip(arg_names, args))
            kernel_kwargs.update(kwargs)
            self._debug_call = debug_mode.record_triton_kernel(
                kernel_name=self.fn.__name__, kwargs=kernel_kwargs
            )

        if hasattr(triton, "set_allocator"):

            def alloc_fn(size: int, align: int, stream: int | None):
                return torch.empty(
                    size, dtype=torch.int8, device=self.device_props.type
                )

            triton.set_allocator(alloc_fn)

        if self.triton_interpret:
            args, grid = self._interpret_args_grid(args, self.configs[0])
            return self.fn[grid](
                *args,
                **kwargs,
                **self.configs[0].kwargs,
            )

        for plugin in self._plugins:
            if (
                result := plugin.pre_dispatch(self, *args, stream=stream, **kwargs)
            ) is not DEFER:
                return result

        if len(self.launchers) != 1:
            if len(self.launchers) == 0:
                start_time = time.time_ns()
                self.precompile()
                self.precompile_time_taken_ns = time.time_ns() - start_time
            if len(self.launchers) > 1:
                for plugin in self._plugins:
                    if (
                        result := plugin.pre_autotune(
                            self, *args, stream=stream, **kwargs
                        )
                    ) is not DEFER:
                        return result
                # Re-check: a plugin may have mutated launchers down to one.
                if len(self.launchers) > 1:
                    self.autotune_to_one_config(*args, **kwargs)

        if self.inductor_meta.get("combo_tuning_groups") and not getattr(
            self.launchers[0].config, "found_by_combo_autotune", False
        ):
            with dynamo_timed(
                "CachingAutotuner.combo_sequential_autotune",
                log_pt2_compile_event=False,
                metadata={"kernel_name": self.inductor_meta.get("kernel_name")},
                dynamo_compile_column_us="runtime_triton_autotune_time_us",
                compile_id=self.compile_id,
                is_backward=self.is_backward,
                log_waitcounter=True,
                waitcounter_name_override="triton_autotuner",
            ):
                self.launchers = [
                    self._combo_sequential_autotune(self.launchers[0], *args, **kwargs)
                ]

        if not getattr(
            self.launchers[0].config, "found_by_coordesc", False
        ) and self.inductor_meta.get("coordinate_descent_tuning", False):
            self.launchers = [
                self.coordinate_descent_tuning(self.launchers[0], *args, **kwargs)
            ]

        (launcher,) = self.launchers
        # Ensure the final launcher is marked as a winner for bundle filtering.
        # For multi-config autotuning and coordesc, put_winner was already called
        # (this is an idempotent set-add). For single-config kernels that skip
        # autotuning entirely, this is the only call site that records the winner.
        TritonBundler.put_winner(launcher.cache_hash)
        if launcher.store_cubin and (not benchmark_run or not self.cuda_kernel_saved):
            if self.device_props.type == "cpu":
                if not self.cpu_kernel_saved:
                    self.save_cpu_kernel(launcher)
            else:
                self.save_gpu_kernel(stream, launcher)

        try:
            self._pre_launch(launcher, *args, stream=stream, **kwargs)
            try:
                result = launcher(*args, **kwargs, stream=stream)
            except Exception as e:
                if isinstance(e, TypeError):
                    self._check_launcher_call_args(launcher, args)
                raise
        finally:
            self._post_launch()

        # Populate fast path: cache the launcher for future calls.  Static
        # conditions (interpret, dump flags) are pre-computed in _cache_eligible;
        # only dynamic conditions are checked here.
        if (
            self._cached_launcher is None
            and self._cache_eligible
            and not benchmark_run
            and not debug_mode
            and not autograd_profiler._is_profiler_enabled
            and len(self.launchers) == 1
        ):
            self._cached_launcher = self._build_fast_launcher(launcher) or launcher
        return result

    def _check_launcher_call_args(
        self,
        launcher: LauncherType,
        args: tuple[Any, ...],
    ) -> None:
        """Raise TypeError with a helpful message when stream is passed positionally."""
        expected = getattr(launcher, "_expected_positional_count", None)
        if expected is None:
            return

        if len(args) > expected:
            kernel_name = self.inductor_meta.get("kernel_name", "triton kernel")
            raise TypeError(
                f"{kernel_name}: too many positional arguments - "
                f"expected {expected}, got {len(args)}. "
                "'stream' must be passed as a keyword argument."
            ) from None

    def _build_fast_launcher(self, launcher: LauncherType) -> LauncherType | None:
        """Try to build a _FastCudaLauncher-backed version of the launcher.

        Returns a new launcher function with the runner replaced by a
        _FastCudaLauncher instance (vectorcall C extension), or None if
        conditions are not met.  Falls back silently on expected errors,
        logs a warning on unexpected ones.
        """
        import types

        if not self.inductor_meta.get(
            "use_fast_triton_launcher",
            torch._inductor.config.use_fast_triton_launcher,
        ):
            return None

        # The _FastCudaLauncher vectorcall bypasses run(), so it can't do the
        # variable-length TMA descriptor expansion; fall back to the regular
        # (still static) launcher.
        if self.inductor_meta.get("host_tma_descriptor_args"):
            return None

        # _FastCudaLauncher binds CUDA/HIP function pointers; XPU static
        # launchers use SYCL kernels stored in PyCapsules.
        if self.device_props.type not in ("cuda", "hip"):
            return None

        try:
            from torch._C import _FastCudaLauncher
        except ImportError:
            return None

        try:
            # Only works for the static triton launcher path.
            if not getattr(launcher, "_is_static", False):
                return None

            # Resolve the bound kernel behind the launcher function.
            runner = launcher.__globals__.get("runner")
            if not callable(runner):
                return None
            kernel = runner.__self__
            # compile-on-one-rank kernels keep their loaded handles per device, so
            # kernel.function is None; the fast launcher binds a single device's
            # function pointer, so fall back to the per-device static launcher.
            if getattr(kernel, "device_agnostic", False):
                return None
            cu_function = kernel.function
            num_warps = kernel.num_warps
            shared = kernel.shared
            arg_tys = kernel.arg_tys
            if cu_function is None or num_warps is None:
                return None

            n_scratch = sum(
                [
                    getattr(kernel, "has_global_scratch", False),
                    getattr(kernel, "has_profile_scratch", False),
                ]
            )
            if torch.version.hip:
                n_scratch = max(n_scratch, 2)

            fast_runner = _FastCudaLauncher(
                cu_function, num_warps, shared, arg_tys, n_scratch
            )

            new_globals = {**launcher.__globals__, "runner": fast_runner}
            new_launcher = types.FunctionType(
                launcher.__code__,
                new_globals,
                launcher.__name__,
            )
            # Copy launcher attributes to the new function object.
            # NOTE: If new attributes are added to launchers in the future,
            # they must be added here too — otherwise the fast launcher will
            # silently drop them.
            for attr in (
                "config",
                "n_regs",
                "n_spills",
                "shared",
                "cache_hash",
                "store_cubin",
                "_is_static",
                "_expected_positional_count",
            ):
                val = getattr(launcher, attr, None)
                if val is not None:
                    setattr(new_launcher, attr, val)
            # _FastCudaLauncher bakes kernel.function (a raw CUfunction pointer)
            # into a C object and never re-reads it, and replacing the "runner"
            # global drops this launcher's only reference to the owning static
            # kernel. Without an explicit reference the kernel can be collected or
            # closed while this launcher is still cached and callable; its
            # close()/__del__ then unloads the module and leaves the baked pointer
            # dangling, producing a CUDA "misaligned address" error on the next
            # launch. Keep the owner alive for as long as the fast launcher is.
            new_launcher._static_kernel_owner = kernel  # type: ignore[attr-defined]
            return new_launcher
        except (AttributeError, TypeError, KeyError, ValueError):
            # Expected failures - silent fallback is OK.
            # ValueError is raised when nArgs > MAX_ARGS in the C extension.
            return None
        except Exception:
            # Unexpected failures - log for debugging
            log.warning("Unexpected error building fast launcher", exc_info=True)
            return None

    def _interpret_args_grid(
        self, args: tuple[Any, ...], cfg: Config
    ) -> tuple[tuple[Any, ...], tuple[int, int, int]]:
        if triton_version_uses_attrs_dict():

            def filtered_signature() -> list[str]:
                # constexprs are not passed in as args
                new_signature: list[str] = []
                from triton.runtime.interpreter import InterpretedFunction

                for i, x in enumerate(self.triton_meta["signature"].keys()):
                    if isinstance(self.fn, InterpretedFunction):
                        # These are torch compiled triton kernels that definitely
                        # have block size configs. Dynamo does not currently
                        # trace user defined triton kernels when TRITON_INTERPRET=1
                        if x not in cfg.kwargs:
                            new_signature.append(x)
                    elif i not in get_constexprs(self.fn):
                        # use constexprs rather than just configs since user
                        # defined triton kernels may not have any configs
                        new_signature.append(x)

                return new_signature

        else:

            def filtered_signature() -> list[str]:
                return list(self.triton_meta["signature"].keys())

        grid = GridExpr.from_meta(
            self.inductor_meta, cfg, mode=self.grid_mode
        ).eval_slow(
            dict(
                zip(
                    [
                        *filtered_signature(),
                        *self.inductor_meta.get("extra_launcher_args", ()),
                    ],
                    args,
                )
            )
        )
        if self.inductor_meta.get("extra_launcher_args"):
            args = args[: -len(self.inductor_meta["extra_launcher_args"])]
        return args, grid


class _ConstRepr:
    def __init__(self, value: str):
        self.value = value

    def __call__(self, _=None) -> str:
        return self.value


class CompileResult(Generic[_T]):
    """
    Base class representing compiled result.
    """

    def __init__(
        self,
        kernel: _T,
        config: Config,
        compile_meta: dict[str, Any],
        inductor_meta: InductorMeta,
    ):
        self.kernel = kernel
        self.config = config
        self.compile_meta = compile_meta
        self.inductor_meta = inductor_meta

    def make_launcher(self) -> LauncherType: ...

    def _host_tma_pre_runner_lines(
        self, runner_args: list[str], call_args: list[str]
    ) -> tuple[list[str], list[str]]:
        """Build host-side TMA descriptors in the launcher and swap them in for
        the raw tensor args. Shared by the dynamic and static launchers."""
        host_tma_args = self.inductor_meta.get("host_tma_descriptor_args")
        pre_runner_lines: list[str] = []
        if not host_tma_args:
            return pre_runner_lines, runner_args
        if not has_triton_stable_tma_api():
            raise RuntimeError(
                "host-side TMA requires a Triton with the stable TMA API "
                "(triton.tools.tensor_descriptor.TensorDescriptor)"
            )
        cfg_kwargs = self.config.kwargs
        all_constants = self.compile_meta["constants"]
        for inner_name, desc_info in host_tma_args.items():
            if inner_name not in call_args or not isinstance(desc_info, dict):
                continue
            block_shape_vals = _resolve_dims(
                desc_info["block_shape"], cfg_kwargs, all_constants
            )
            shape_vals = _resolve_dims(desc_info["shape"], cfg_kwargs, all_constants)
            stride_vals = _resolve_dims(desc_info["strides"], cfg_kwargs, all_constants)
            if block_shape_vals is None or shape_vals is None or stride_vals is None:
                continue
            desc_var = f"{inner_name}_host_tma_desc"
            aligned_var = f"{inner_name}_aligned"
            # _host_tma_aligned clones (warning once) only if misaligned.
            pre_runner_lines.append(
                f'{aligned_var} = _host_tma_aligned({inner_name}, "{inner_name}")'
            )
            pre_runner_lines.append(
                f"{desc_var} = TensorDescriptor({aligned_var}, {shape_vals},"
                f" {stride_vals}, {block_shape_vals})"
            )
            runner_args = [desc_var if a == inner_name else a for a in runner_args]
        return pre_runner_lines, runner_args

    def _gen_launcher_code(
        self, scope, def_args, runner_args, pre_runner_lines=None
    ) -> LauncherType:
        grid = GridExpr.from_meta(self.inductor_meta, self.config)
        # grid.prefix is usually empty, grid.x_grid is something like `-(xnumel//-1024)`
        lines = [
            f"def launcher({', '.join(def_args)}, stream):",
            *[f"    {line}" for line in grid.prefix],
            f"    grid_0 = {grid.x_grid}",
            f"    grid_1 = {grid.y_grid}",
            f"    grid_2 = {grid.z_grid}",
            *(f"    {l}" for l in (pre_runner_lines or [])),
            f"    runner({', '.join(runner_args)})",
        ]
        launcher_code = "\n".join(lines)
        exec(launcher_code, scope)
        launcher = scope["launcher"]
        # Stash expected positional arg count at codegen time so
        # _check_launcher_call_args can validate without inspect.signature().
        launcher._expected_positional_count = len(def_args)
        return launcher

    def _get_arg_lists(
        self, arg_names, constexprs
    ) -> tuple[list[str], list[str], OrderedSet[str]]:
        """
        Return a bunch of intermediate lists of args needed for generating
        launcher code.
        """
        compile_meta = self.compile_meta
        cfg = self.config
        known_constants = OrderedSet(
            arg for i, arg in enumerate(arg_names) if i in constexprs
        )

        """
        https://github.com/pytorch/pytorch/issues/115344

        self.fn.constexprs doesn't properly deal with None args, so when we filter out
        an arg in UserDefinedTritonKernel.codegen, we need to filter it here as well.
        We also don't want to modify self.fn.

        We know that we removed something from the signature if:
            1. It's in compile_meta["constants"]
            2. It isn't a constant we already know about
                Note: The value of interest has already been added to compile_meta['constants'],
                    so we use self.fn.constexprs instead.
            3. It isn't in the compile_meta signature
        """
        none_args = OrderedSet(
            k
            for k, v in compile_meta["constants"].items()
            if v is None and k not in known_constants
        )
        none_args = none_args.difference(OrderedSet(compile_meta["signature"].keys()))

        def _convert_constant(constant):
            if isinstance(constant, str):
                return "r'" + constant + "'"
            else:
                return repr(constant)

        if triton_version_uses_attrs_dict():
            call_args = arg_names
            def_args = arg_names
            implicit_constants = OrderedSet(
                (
                    "num_warps",
                    "num_stages",
                )
            ).union(OrderedSet(k for k in known_constants))
            if implicit_constants := implicit_constants & OrderedSet(
                compile_meta["constants"].keys()
            ):
                # num_warps/num_stages are special implicit args that are not in the signature
                # see test_triton_kernel_special_params
                def_args = [arg for arg in def_args if arg not in implicit_constants]
                repl = {
                    k: _convert_constant(compile_meta["constants"].get(k))
                    for k in implicit_constants
                }
                call_args = [repl.get(arg, arg) for arg in call_args]
        else:
            call_args = [
                arg
                for i, arg in enumerate(arg_names)
                if i not in constexprs and arg not in none_args
            ]
            cfg_dict = config_to_dict(cfg)
            def_args = [
                name
                for name in arg_names
                if name not in cfg_dict and name not in none_args
            ]

        if "extra_launcher_args" in self.inductor_meta:
            def_args = [*def_args, *self.inductor_meta["extra_launcher_args"]]

        return call_args, def_args, none_args


_KernelCompileResult: TypeAlias = (
    CompileResult[CompiledKernel]
    | CompileResult[StaticallyLaunchedCudaKernel]
    | CompileResult[StaticallyLaunchedXpuKernel]
)


class CannotStaticallyLaunchKernel(Exception):
    pass


class StaticTritonCompileResult(CompileResult[_T]):
    """
    TritonCompileResult that uses StaticCudaLauncher,
    which vastly simplifies the setup and metadata needed to be kept.
    """

    @staticmethod
    def can_statically_launch(
        kernel: CompiledKernel,
        inductor_meta: InductorMeta,
        triton_meta: TritonMeta,
        heuristic_type: HeuristicType,
    ) -> _KernelType | None:
        if not torch._inductor.config.use_static_triton_launcher:
            return None

        def check_can_launch() -> _KernelType:
            if triton_meta.get("device_type") not in ("cuda", "xpu", "hip"):
                raise CannotStaticallyLaunchKernel("Non-cuda/XPU/ROCm device")

            if triton_meta.get("device_type") == "xpu" and XPU_KERNEL_FORMAT == "spv":
                raise CannotStaticallyLaunchKernel(
                    "Static XPU Triton kernel launch does not support SPIR-V kernel."
                )

            if torch._inductor.config.cpp_wrapper:
                # If we're running with cpp wrapper, it doesn't
                # make sense to statically compile since everything
                # is codegenned anyway
                raise CannotStaticallyLaunchKernel("Cpp wrapper enabled")

            if (
                heuristic_type == HeuristicType.USER_AUTOTUNE
                and not torch._inductor.config.static_launch_user_defined_triton_kernels
            ):
                # Don't support user defined triton kernels yet
                raise CannotStaticallyLaunchKernel("User defined triton kernel")

            if inductor_meta.get("store_cubin"):
                # Requires storing the entire binary
                raise CannotStaticallyLaunchKernel("store_cubin is enabled")

            if getattr(kernel.metadata, "launch_pdl", False) or getattr(
                kernel.metadata, "launch_cooperative_grid", False
            ):
                raise CannotStaticallyLaunchKernel(
                    "static launch does not support launch attributes"
                )

            binary_ext = GPU_KERNEL_BIN_EXTS.get(
                triton_meta.get("device_type"), ".cubin"
            )
            cubin_location = os.path.join(
                triton_cache_dir(cast(int, triton_meta.get("device", 0))),
                triton_hash_to_path_key(kernel.hash),
                f"{kernel.src.fn.__name__}{binary_ext}",
            )

            if not os.path.exists(cubin_location):
                raise CannotStaticallyLaunchKernel(
                    f"Cubin path not found: {cubin_location}"
                )

            else:
                kernel._cubin_path = cubin_location

            try:
                static_kernel = statically_launched_kernel_by_device(
                    kernel, triton_meta.get("device_type")
                )
            except NotImplementedError as e:
                raise CannotStaticallyLaunchKernel(f"NotImplemented: {str(e)}") from e

            return static_kernel

        try:
            result = check_can_launch()
            return result
        except CannotStaticallyLaunchKernel as e:
            log.info("Bypassing StaticallyLaunchedCudaKernel due to %s", e)
            if torch._inductor.config.strict_static_triton_launcher:
                raise e
            return None

    def reload_cubin_path(self):
        """
        When loading from cache on disk, we want to reload cubin
        files from their appropriate location on disc.
        """
        device_type = (
            "hip" if torch.version.hip else self.compile_meta.get("device_type", "cuda")
        )
        binary_ext = GPU_KERNEL_BIN_EXTS.get(device_type, "cubin")
        cubin_location = os.path.join(
            triton_cache_dir(
                _resolve_load_device(self.compile_meta.get("device"), device_type)
            ),
            triton_hash_to_path_key(self.kernel.hash),
            f"{self.kernel.name}{binary_ext}",
        )
        if not os.path.exists(cubin_location):
            if self.kernel.cubin_raw is not None:
                # We saved the raw cubin, so write it to he appropriate location
                self.kernel.reload_cubin_from_raw(cubin_location)
            else:
                raise RuntimeError(
                    "Cubin file saved by TritonBundler not found at %s", cubin_location
                )
        self.kernel.cubin_path = cubin_location

    def make_launcher(self) -> LauncherType:
        # If at least one static make_launcher call occurs,
        # we're sure static cuda launcher was used for this compile
        set_feature_use("static_triton_launcher", True)
        # Load the binary on the parent
        if not self.kernel.cubin_path:
            self.reload_cubin_path()
        # compile-on-one-rank: a None device in compile_meta marks a rank/device-agnostic
        # kernel, so the launcher must keep its loaded handles per device.
        self.kernel.device_agnostic = self.compile_meta.get("device") is None
        device = _resolve_load_device(
            self.compile_meta.get("device"),
            self.compile_meta.get("device_type", "cuda"),
        )
        self.kernel.load_kernel(device)
        scope = {
            "runner": self.kernel.run,
        }

        # NOTE: Constexpr handling for triton and static cuda launcher

        # Triton kernels have two types of constexprs: *declared* ones, which are ones the user
        # has explicitly declared as tl.constexpr, and *implied* ones, which are expressions triton
        # deems constant while compiling/analyzing the code (i.e. unused parameters, for example)

        # Triton kernels handle constexprs slightly differently depending on which version of triton
        # we care about (we support 3.2.0 and 3.3.0).

        # In 3.2.0, triton kernels do not require passing any declared constexprs into the kernel
        # In 3.3.0, triton kernels require all declared constexprs be passed into the kernel, where
        # they are subsequently ignored.
        # When statically launching, since we're launching from the triton generated cubin, we actually want to
        # always get rid of all const exprs, declared or implied, since the underlying cubin file has all
        # of the constants stripped away anyway.

        # But CachingAutotuner.run will pass us a different number of arguments depending on
        # whether or not we're in triton 3.2.0 or later, so we grab def_args with the same logic
        # as the (non static) TritonCompileResult. We then generate call_args ourselves, since we
        # want only a subset of the arguments passed to triton.
        # Here, arg_names is exactly fn.src.arg_names and declared_constexprs is exactly fn.src.constexprs,
        # which matches behavior with regular TritonCompileResult
        _, def_args, none_args = self._get_arg_lists(
            self.kernel.arg_names, self.kernel.declared_constexprs
        )

        call_args = [
            arg
            for i, arg in enumerate(self.kernel.arg_names)
            if i not in self.kernel.full_constexprs and arg not in none_args
        ]

        # StaticallyLaunchedCudaKernel.run takes in order grid_0, grid_1, grid_2, stream, and call_args
        runner_args = ["grid_0", "grid_1", "grid_2", "stream", *call_args]
        pre_runner_lines, runner_args = self._host_tma_pre_runner_lines(
            runner_args, call_args
        )
        if self.inductor_meta.get("host_tma_descriptor_args"):
            # _host_tma_pre_runner_lines already validated the stable TMA API.
            from triton.tools.tensor_descriptor import TensorDescriptor

            scope["_host_tma_aligned"] = _host_tma_aligned
            scope["TensorDescriptor"] = TensorDescriptor
        launcher = self._gen_launcher_code(
            scope, def_args, runner_args, pre_runner_lines=pre_runner_lines
        )
        launcher.config = self.config  # type: ignore[attr-defined]
        launcher.n_regs = self.kernel.n_regs  # type: ignore[attr-defined]
        launcher.n_spills = self.kernel.n_spills  # type: ignore[attr-defined]
        launcher.shared = self.kernel.shared  # type: ignore[attr-defined]
        launcher.cache_hash = triton_hash_to_path_key(self.kernel.hash)  # type: ignore[attr-defined]
        launcher.store_cubin = False  # type: ignore[attr-defined]
        launcher._is_static = True  # type: ignore[attr-defined]
        return launcher


class TritonCompileResult(CompileResult[CompiledKernel]):
    """
    Upstream Triton CompileKernel can not be pickled.  This is a wrapper
    to support serialization and generate the launcher function.
    """

    @staticmethod
    @functools.lru_cache(32)
    def _kernel_metadata_cls(fields: tuple[str, ...]) -> Any:
        return namedtuple("KernelMetadata", sorted(fields))

    @staticmethod
    def _serialize_metadata(metadata):
        """
        Triton uses a nested class called KernelMetadata to store metadata information.
        Pickle does not work well with nested namedtuples, as the namedtuple doesn't appear
        in the toplevel namespace of the module. So these serialization/deser functions
        are used to convert the namedtuples to a dict and back.

        As for packed_metadata, depending on the triton backend, KernelMetadata can be
        a namedtuple, or a regular tuple! So the serialization function branches on whether
        the metadata to be serialized is a namedtuple or regular, serializable one.
        """

        def is_namedtuple(obj) -> bool:
            return (
                isinstance(obj, tuple)
                and hasattr(obj, "_asdict")
                and hasattr(obj, "_fields")
            )

        if is_namedtuple(metadata):
            return metadata._asdict()
        else:
            return metadata

    @staticmethod
    def _deserialize_metadata(metadata):
        if isinstance(metadata, dict):
            return TritonCompileResult._kernel_metadata_cls(tuple(metadata.keys()))(
                **metadata
            )
        else:
            return metadata

    def __getstate__(self) -> dict[str, Any]:
        kernel = self.kernel
        # replace the fields that don't pickle nicely
        kernel_state = {
            **kernel.__dict__,
            # See doc about serializing metadata above
            "metadata": self._serialize_metadata(kernel.metadata),
            "packed_metadata": self._serialize_metadata(
                getattr(kernel, "packed_metadata", None)
            ),
            "module": None,  # regenerated by kernel._init_handles()
            "function": None,  # regenerated by kernel._init_handles()
            "run": None,  # regenerated by kernel._init_handles()
        }
        return {**self.__dict__, "kernel": kernel_state}  # type: ignore[dict-item]

    def __setstate__(self, state: dict[str, Any]) -> None:
        # src = ASTSource.__new__(ASTSource)
        # src.__setstate__(state["kernel"]["src"])
        # TODO(jansel): need to fixup src.fn which is now None
        kernel = CompiledKernel.__new__(CompiledKernel)
        metadata = state["kernel"]["metadata"]
        packed_metadata = state["kernel"]["packed_metadata"]
        kernel.__dict__.update(
            {
                **state["kernel"],
                # "src": src,
                "metadata": self._deserialize_metadata(metadata),
                "packed_metadata": self._deserialize_metadata(packed_metadata),
            }
        )
        self.__dict__.update(state)
        self.kernel = kernel

    def make_launcher(self) -> LauncherType:
        """
        Launching triton kernels is performance sensitive, we compile
        a custom Python function get the grid() and reorder the args to
        the underlying wrapper.
        """
        cfg = self.config
        compile_meta = self.compile_meta
        binary = self.kernel
        fn = binary.src.fn
        binary._init_handles()
        (call_args, def_args, none_args) = self._get_arg_lists(
            fn.arg_names, get_constexprs(fn)
        )
        binary_shared = (
            binary.shared if hasattr(binary, "shared") else binary.metadata.shared
        )

        if knobs is None:
            launch_enter = binary.__class__.launch_enter_hook
            launch_exit = binary.__class__.launch_exit_hook
        else:
            launch_enter = knobs.runtime.launch_enter_hook
            launch_exit = knobs.runtime.launch_exit_hook

        import math as math_lib

        import triton as triton_lib

        import torch as torch_lib

        scope = {
            "grid_meta": cfg.kwargs,
            "bin": binary,
            "launch_enter_hook": launch_enter,
            "launch_exit_hook": launch_exit,
            "metadata": (
                binary.packed_metadata
                if hasattr(binary, "packed_metadata")
                else binary.metadata
            ),
            "shared": binary_shared,
            "num_warps": (
                binary.num_warps
                if hasattr(binary, "num_warps")
                else binary.metadata.num_warps
            ),
            "cta_args": (
                (
                    binary.num_ctas,
                    *get_first_attr(binary, "cluster_dims", "clusterDims"),
                )
                if hasattr(binary, "num_ctas")
                else (
                    (binary.metadata.num_ctas, *binary.metadata.cluster_dims)
                    if hasattr(binary, "metadata")
                    and hasattr(binary.metadata, "num_ctas")
                    and hasattr(binary.metadata, "cluster_dims")
                    else ()
                )
            ),
            "function": get_first_attr(binary, "function", "cu_function"),
            "runner": get_first_attr(binary, "run", "c_wrapper"),
            "math": math_lib,
            "torch": torch_lib,
            "triton": triton_lib,
        }

        if not hasattr(binary, "launch_metadata"):
            # launch args before CompiledKernel.launch_metadata is added.
            # TODO(jansel): delete this branch in mid-2025
            runner_args = [
                "grid_0",
                "grid_1",
                "grid_2",
                "num_warps",
                "*cta_args",
                "shared",
                "stream",
                "function",
                "launch_enter_hook",
                "launch_exit_hook",
                "metadata",
                *call_args,
            ]
        else:  # args after CompiledKernel.launch_metadata: https://github.com/triton-lang/triton/pull/3492
            # Getting the kernel launch args is extremely perf-sensitive.  Evaluating
            # `bin.launch_metadata` is relatively expensive, and returns None unless a
            # `launch_enter_hook` is installed.  So if we don't have that hook installed,
            # we want to burn None in to the launch args with zero overhead.
            # See https://github.com/pytorch/pytorch/issues/123597
            if launch_enter:
                launch_metadata = f"bin.launch_metadata((grid_0, grid_1, grid_2), stream, {', '.join(call_args)})"
            else:
                launch_metadata = "None"
            runner_args = [
                "grid_0",
                "grid_1",
                "grid_2",
                "stream",
                "function",
                "metadata",
                launch_metadata,
                "launch_enter_hook",
                "launch_exit_hook",
                *call_args,
            ]

        pre_runner_lines, runner_args = self._host_tma_pre_runner_lines(
            runner_args, call_args
        )
        if self.inductor_meta.get("host_tma_descriptor_args"):
            # _host_tma_pre_runner_lines already validated the stable TMA API.
            from triton.tools.tensor_descriptor import TensorDescriptor

            scope["_host_tma_aligned"] = _host_tma_aligned
            scope["TensorDescriptor"] = TensorDescriptor

        launcher = self._gen_launcher_code(
            scope, def_args, runner_args, pre_runner_lines=pre_runner_lines
        )

        launcher = scope["launcher"]
        launcher.config = cfg
        launcher.n_regs = getattr(binary, "n_regs", None)
        launcher.n_spills = getattr(binary, "n_spills", None)
        launcher.shared = binary_shared
        launcher.cache_hash = triton_hash_to_path_key(binary.hash)
        launcher.store_cubin = self.inductor_meta.get("store_cubin", False)
        # store this global variable to avoid the high overhead of reading it when calling run
        if launcher.store_cubin:
            launcher.fn = fn
            launcher.bin = binary
            if triton_version_uses_attrs_dict():
                # arg filtering wasn't done above
                cfg_dict = config_to_dict(cfg)
                def_args = [x for x in def_args if x not in cfg_dict]
                call_args = [
                    x
                    for x in call_args
                    if compile_meta["signature"].get(x, "constexpr") != "constexpr"
                    and x not in none_args
                ]
            launcher.def_args = def_args
            launcher.call_args = call_args
            kernel_metadata = getattr(self.kernel, "metadata", None)

            # for the scratch arguments: None indicates that the kernel doesn't
            # take any scratch argument; otherwise a number indicates the number
            # of bytes of scratch that need to be provided.

            # in AMD's Triton backend, the global scratch size is never provided
            # (but for AMD it's safe to pass an extra null arg, so always include it)
            global_scratch: int | None = getattr(
                kernel_metadata,
                "global_scratch_size",
                (0 if torch.version.hip else None),
            )
            profile_scratch: int | None = getattr(
                kernel_metadata, "profile_scratch_size", None
            )
            launcher.global_scratch = global_scratch
            launcher.profile_scratch = profile_scratch
        return launcher


def _find_names(obj):
    import gc
    import inspect

    frame = inspect.currentframe()
    while frame is not None:
        # On CPython <= 3.12 this access materializes the frame's locals
        # dict so gc.get_referrers below can find obj inside it. On 3.13+
        # f_locals is a fresh write-through proxy (PEP 667) and this loop
        # is a no-op, so function-local names are not discoverable there.
        _ = frame.f_locals
        frame = frame.f_back
    obj_names = []
    for referrer in gc.get_referrers(obj):
        if isinstance(referrer, dict):
            for k, v in referrer.items():
                if v is obj:
                    obj_names.append(k)
    return obj_names


collected_calls: list[Any] = []


def start_graph():
    collected_calls.clear()


def end_graph(output_file):
    if len(collected_calls) == 0:
        return
    overall_time = sum(call[0] for call in collected_calls)
    overall_gb = sum(call[1] for call in collected_calls)
    cur_file = inspect.stack()[1].filename
    summary_str = (
        f"SUMMARY ({cur_file})\n"
        f"{overall_time:.2f}ms   \t {overall_gb:.2f} GB\t {overall_gb / (overall_time / 1e3):.2f}GB/s"
    )
    log.info(
        "%s",
        summary_str,
    )
    if output_file is not None:
        # sort perf numbers in descending order, i.e. placing the
        # most runtime-heavy kernels at the top of the list
        sorted_calls = sorted(collected_calls, key=lambda c: float(c[0]), reverse=True)
        try:
            with open(output_file, "a") as file:
                log.info(
                    "Save profile bandwidth results to %s",
                    output_file,
                )
                file.write("====================\n")
                file.write(f"TRITON KERNELS BANDWIDTH INFO ({cur_file})\n")
                for ms, num_gb, gb_per_s, kernel_name in sorted_calls:
                    # also display the runtime percentage for each kernel
                    percentage = f"{ms / overall_time * 100:.2f}%"
                    suffix = f" \t {percentage} \t {kernel_name}"
                    bw_info_str = create_bandwidth_info_str(
                        ms,
                        num_gb,
                        gb_per_s,
                        suffix=suffix,
                        color=False,
                    )
                    file.write(bw_info_str + "\n")
                file.write(f"{summary_str}\n\n")
        except Exception:
            log.warning(
                "failed to write profile bandwidth result into %s",
                output_file,
                exc_info=True,
            )


class DebugAutotuner(CachingAutotuner):
    def __init__(
        self,
        *args,
        regex_filter="",
        with_profiler=False,
        with_bandwidth_info=True,
        **kwargs,
    ):
        self.regex_filter = regex_filter
        self.with_profiler = with_profiler
        self.with_bandwidth_info = with_bandwidth_info
        super().__init__(*args, **kwargs)
        self.cached = None

    def run(self, *args, stream, **kwargs):
        if not self.with_bandwidth_info:
            super().run(*args, stream=stream, **kwargs, benchmark_run=True)
            return
        else:
            possible_names = _find_names(self)
            if possible_names:
                kernel_name = f"{max(possible_names, key=len)}"
            else:
                # In the AOTI lazy compile path, the CachingAutotuner is not
                # bound to a module-level name, so _find_names returns empty.
                # Fall back to the kernel name recorded in inductor_meta.
                kernel_name = self.inductor_meta.get("kernel_name") or self.fn.__name__
            if not re.match(self.regex_filter, kernel_name):
                return
            if len(self.launchers) != 1:
                if len(self.launchers) == 0:
                    start_time = time.time_ns()
                    self.precompile()
                    self.precompile_time_taken_ns = time.time_ns() - start_time
                if len(self.launchers) > 1:
                    self.autotune_to_one_config(*args, **kwargs)
            (launcher,) = self.launchers

            if launcher.store_cubin:
                if self.device_props.type == "cpu":
                    self.save_cpu_kernel(launcher)
                else:
                    self.save_gpu_kernel(stream, launcher)

            if self.cached is None:
                ms = self.bench(launcher, *args, with_profiler=self.with_profiler)
                num_in_out_ptrs = len(
                    [
                        arg_name
                        for arg_name in self.fn.arg_names
                        if arg_name.startswith("in_out_ptr")
                    ]
                )
                num_gb = self.inductor_meta.get("kernel_num_gb", None)
                if num_gb is None:
                    num_gb = get_num_bytes(*args, num_in_out_args=num_in_out_ptrs) / 1e9
                gb_per_s = num_gb / (ms / 1e3)
                self.cached = ms, num_gb, gb_per_s, kernel_name
                collected_calls.append((ms, num_gb, gb_per_s, kernel_name))
                log.info(
                    "%s",
                    create_bandwidth_info_str(
                        ms, num_gb, gb_per_s, suffix=f" \t {kernel_name}"
                    ),
                )
            else:
                # in AOTI, we will call the kernel and its timing info has been cached already
                collected_calls.append(self.cached)


def hash_configs(configs: list[Config]):
    """
    Hash used to check for changes in configurations
    """
    hasher = hashlib.sha256()
    for cfg in configs:
        hasher.update(
            f"{sorted(cfg.kwargs.items())} {cfg.num_warps} {cfg.num_stages}\n".encode()
        )
    return hasher.hexdigest()


def cached_autotune(
    size_hints: list[int] | None,
    configs: list[Config],
    triton_meta: TritonMeta,
    heuristic_type,
    filename=None,
    inductor_meta: InductorMeta | None = None,
    custom_kernel=False,
    caching_autotuner_cls: type[CachingAutotuner] = CachingAutotuner,
    debug_autotuner_cls: type[DebugAutotuner] = DebugAutotuner,
):
    """
    A copy of triton.autotune that calls our subclass.  Our subclass
    has additional debugging, error handling, and on-disk caching.
    """
    inductor_meta = {} if inductor_meta is None else inductor_meta
    if size_hints is not None and heuristic_type in (
        HeuristicType.REDUCTION,
        HeuristicType.PERSISTENT_REDUCTION,
    ):
        configs = _enforce_reduction_config_block_minimums(
            configs, size_hints, inductor_meta
        )
    configs = unique_configs(configs)
    if len(configs) != 1 and not filename:
        raise AssertionError("filename required when multiple configs are provided")

    device_prop = triton_meta.get("device")
    if not isinstance(device_prop, DeviceProperties):
        device_prop = None
    dynamic_scale_rblock_eligible = _could_dynamic_scale_rblock(
        size_hints=size_hints,
        heuristic_type=heuristic_type,
        device_prop=device_prop,
        inductor_meta=inductor_meta,
    )
    configs, autotune_cache, autotune_cache_info = check_autotune_cache(
        configs,
        filename,
        inductor_meta,
        dynamic_scale_rblock_eligible=dynamic_scale_rblock_eligible,
    )
    mutated_arg_names = cast("list[str]", inductor_meta.pop("mutated_arg_names", ()))
    optimize_mem = inductor_meta.pop("optimize_mem", True)

    if "restore_value" in triton_meta:
        mutated_arg_names += triton_meta.pop("restore_value")

    reset_to_zero_arg_names: list[str] = []
    if "reset_to_zero" in triton_meta:
        reset_to_zero_arg_names.extend(triton_meta.pop("reset_to_zero"))

    def decorator(fn):
        # Remove XBLOCK from config if it's not a function argument.
        # This way, coordinate descent tuning will not try to tune it.
        #
        # Context: When TritonKernel.no_x_dim is True, we hardcode XBLOCK to 1.
        import inspect

        if "XBLOCK" not in inspect.signature(fn.fn).parameters:
            for tconfig in configs:
                if "XBLOCK" in tconfig.kwargs:
                    if tconfig.kwargs["XBLOCK"] != 1:
                        raise AssertionError(
                            f"Expected XBLOCK == 1 when not in fn params, got {tconfig.kwargs['XBLOCK']}"
                        )
                    tconfig.kwargs.pop("XBLOCK")

        if inductor_meta.get("profile_bandwidth"):
            return debug_autotuner_cls(
                fn,
                triton_meta=triton_meta,
                inductor_meta=inductor_meta,
                regex_filter=inductor_meta["profile_bandwidth_regex"],
                with_profiler=inductor_meta[
                    "profile_bandwidth_with_do_bench_using_profiling"
                ],
                configs=configs,
                save_cache_hook=autotune_cache and autotune_cache.save,
                mutated_arg_names=mutated_arg_names,
                reset_to_zero_arg_names=reset_to_zero_arg_names,
                optimize_mem=optimize_mem,
                heuristic_type=heuristic_type,
                size_hints=size_hints,
                custom_kernel=custom_kernel,
                filename=filename,
                with_bandwidth_info=True,
            )
        return caching_autotuner_cls(
            fn,
            triton_meta=triton_meta,
            inductor_meta=inductor_meta,
            configs=configs,
            save_cache_hook=autotune_cache and autotune_cache.save,
            mutated_arg_names=mutated_arg_names,
            reset_to_zero_arg_names=reset_to_zero_arg_names,
            optimize_mem=optimize_mem,
            heuristic_type=heuristic_type,
            size_hints=size_hints,
            custom_kernel=custom_kernel,
            filename=filename,
            autotune_cache_info=autotune_cache_info,
        )

    return decorator


def unique_configs(configs: list[Config]):
    """Remove duplicate configurations"""
    seen: OrderedSet[Hashable] = OrderedSet()
    pruned_configs = []

    for cfg in configs:
        key = triton_config_to_hashable(cfg)
        if key not in seen:
            seen.add(key)
            pruned_configs.append(cfg)
    return pruned_configs


def check_config(cfg, *, xnumel=None, ynumel=None, znumel=None):
    for numel, label in zip((xnumel, ynumel, znumel), "XYZ"):
        if numel is None:
            continue
        block = cfg[f"{label}BLOCK"]
        if numel == 1:
            if block != 1:
                raise AssertionError(
                    f"TritonKernel.indexing assumes numel == 1 => BLOCK == 1"
                    f" but {label.lower()}numel=={numel} and {label}BLOCK={block} (cfg={cfg})."
                )
        max_block = TRITON_MAX_BLOCK[label]
        max_block_str = f'config.triton.max_block["{label}"]'
        if max_block % block != 0:
            raise AssertionError(
                f"TritonKernel.indexing assumes {label}BLOCK divides {max_block_str}"
                f" but {label}BLOCK={block} and {max_block_str}={max_block} (cfg={cfg})."
            )


def check_max_block(cfg: dict[str, int]):
    """
    Check that block sizes are within the maximum allowed.
    """
    for var, val in cfg.items():
        block_suffix = "BLOCK"
        if block_suffix in var:
            prefix = var.removesuffix(block_suffix)
            max_block = TRITON_MAX_BLOCK[prefix]
            if val > max_block:
                raise AssertionError(
                    f"'{var}' too large. Maximum: {max_block}. Actual: {val}."
                )


def _check_native_matmul_block_numel(
    kwargs: dict[str, int], r0_block: int | None = None
) -> None:
    block_numel = native_matmul_block_numel(kwargs, r0_block=r0_block)
    if block_numel > TRITON_MAX_TENSOR_NUMEL:
        raise AssertionError(
            f"Block numel {block_numel} exceeds Triton maximum "
            f"{TRITON_MAX_TENSOR_NUMEL}"
        )


def _native_matmul_config_under_numel_limit(
    cfg: Config, r0_block: int | None = None
) -> bool:
    return (
        native_matmul_block_numel(cfg.kwargs, r0_block=r0_block)
        <= TRITON_MAX_TENSOR_NUMEL
    )


def _cap_native_matmul_configs(configs: list[Config], r0_block: int) -> list[Config]:
    capped_configs: list[Config] = []
    for cfg in configs:
        cfg = copy.deepcopy(cfg)
        while not _native_matmul_config_under_numel_limit(cfg, r0_block=r0_block):
            shrinkable_fields = [
                field for field in ("XBLOCK", "YBLOCK") if cfg.kwargs.get(field, 1) > 16
            ]
            if not shrinkable_fields:
                break
            field = max(shrinkable_fields, key=lambda field: cfg.kwargs[field])
            cfg.kwargs[field] //= 2

        if _native_matmul_config_under_numel_limit(cfg, r0_block=r0_block):
            capped_configs.append(cfg)

    return unique_configs(capped_configs)


def _enforce_reduction_config_block_minimums(
    configs: list[Config],
    size_hints: dict[str, int],
    inductor_meta: InductorMeta,
) -> list[Config]:
    min_xblock = inductor_meta.get("min_xblock")
    min_rblock = inductor_meta.get("min_rblock")
    if min_xblock is None and min_rblock is None:
        return configs

    for cfg in configs:
        if frozenset(("YBLOCK", "ZBLOCK", "R1_BLOCK")) & cfg.kwargs.keys():
            raise AssertionError(
                f"min_xblock/min_rblock only support 2D X/R0 configs: {cfg}"
            )
        has_xblock = "XBLOCK" in cfg.kwargs
        has_rblock = "R0_BLOCK" in cfg.kwargs
        if not (has_xblock or has_rblock):
            continue

        x_floor = min_xblock if min_xblock is not None else 1
        r_floor = min_rblock if min_rblock is not None else 1
        target_tile_product = (cfg.kwargs["XBLOCK"] if has_xblock else 1) * (
            cfg.kwargs["R0_BLOCK"] if has_rblock else 1
        )

        if has_xblock:
            cfg.kwargs["XBLOCK"] = max(cfg.kwargs["XBLOCK"], x_floor)
        if has_rblock:
            cfg.kwargs["R0_BLOCK"] = max(cfg.kwargs["R0_BLOCK"], r_floor)

        def current_tile_product() -> int:
            return (cfg.kwargs["XBLOCK"] if has_xblock else 1) * (
                cfg.kwargs["R0_BLOCK"] if has_rblock else 1
            )

        def shrink_to_budget(name: str, floor: int) -> None:
            while (
                name in cfg.kwargs
                and current_tile_product() > target_tile_product
                and cfg.kwargs[name] > floor
            ):
                cfg.kwargs[name] //= 2

        # Preserve the autotuner's original tile-size budget where possible:
        # raising one block to satisfy a floor should shrink the other block.
        shrink_to_budget("R0_BLOCK", r_floor)
        shrink_to_budget("XBLOCK", x_floor)

        check_max_block(cfg.kwargs)
        check_config(
            cfg.kwargs,
            xnumel=size_hints.get("x") if "XBLOCK" in cfg.kwargs else None,
            ynumel=size_hints.get("y") if "YBLOCK" in cfg.kwargs else None,
            znumel=size_hints.get("z") if "ZBLOCK" in cfg.kwargs else None,
        )

    return configs


def _num_warps(
    num_warps,
    max_num_warps=8,
    min_num_warps=2,
    register_intensive=False,
    *,
    warp_size: int = 32,
):
    # On wave64 AMD GPUs (CDNA / gfx9) each warp has 64 lanes, double NVIDIA
    # and RDNA, so use half the number of warps to keep total thread count
    # comparable. RDNA (wave32) follows the NVIDIA path.
    if warp_size == 64:
        max_num_warps = (max_num_warps + 1) // 2
        min_num_warps = (min_num_warps + 1) // 2
    # persistent reduction is register intensive
    if register_intensive:
        max_num_warps = max_num_warps // 2
    return next_power_of_2(min(max(num_warps, min_num_warps), max_num_warps))


def _check_max_grid_x(size_hints, x, num_warps, *, warp_size: int = 32):
    # Check if maxGridSize is exceeded - if so then must scale XBLOCK further
    max_grid_x = 2147483647
    max_block_x = TRITON_MAX_BLOCK["X"]
    num_blocks = (size_hints["x"] + x - 1) // x

    if torch.version.hip:
        # HIP has a 2^31-1 limit on total threads (num_blocks * num_warps * warp_size)
        while (
            (num_blocks * num_warps * warp_size) > max_grid_x
            and x < size_hints["x"]
            and x < max_block_x
        ):
            x *= 2
            num_blocks = num_blocks // 2
    else:
        # NVIDIA has a 2^31-1 limit on number of blocks in grid (not total threads)
        while num_blocks > max_grid_x and x < size_hints["x"] and x < max_block_x:
            x *= 2
            num_blocks = num_blocks // 2

    if num_blocks > max_grid_x:
        raise AssertionError(
            "Reduction config exceeds cudaDeviceProp maxGridSize. Please raise a pytorch issue"
        )
    return x, num_blocks


def triton_config(
    size_hints,
    x,
    y=None,
    z=None,
    num_stages=1,
    num_elements_per_warp=256,
    min_elem_per_thread=0,
    num_warps=None,
    matrix_instr=None,
    waves_per_eu=None,
    kpack=None,
    *,
    warp_size: int = 32,
) -> Config:
    """
    Construct a pointwise triton config with some adjustment heuristics
    based on size_hints. Size_hints is a tuple of numels in each tile
    dimension and will be rounded up to the nearest power of 2.

    num_elements_per_warp is a suggestion for controlling how many warps
    the triton config should contain. e.g.: if x=16, y=8, z=4 then
    num_elements = 16*8*4 = 512. Then if we set num_elements_per_warp=128,
    we'll launch 512 (elem) / 128 (elem/warp) = 4 warps. Note that it's
    just a suggestion, and sometimes other adjustment heuristics will
    override the num_elements_per_warp.

    min_elem_per_thread controls the minimum number of elements
    processed by each thread. It's always enforced.
    """
    # Ideally we want to read this from some device config

    maxGridSize = [2147483647, 65535, 65535]

    target = conditional_product(x, y, z)
    if conditional_product(*size_hints.values()) < target:
        target //= 8

    # shrink sizes to size hints
    x = min(x, size_hints["x"])
    if y:
        y = min(y, size_hints["y"])
    if z:
        z = min(z, size_hints["z"])

    # if we are below original block size, scale up where we can;
    # or if the calculated grid size is larger than the limit, we bump up the corresponding dimension
    while x < min(size_hints["x"], TRITON_MAX_BLOCK["X"]) and (
        x * maxGridSize[0] < size_hints["x"] or conditional_product(x, y, z) < target
    ):
        x *= 2
    while (
        y
        and y < min(size_hints["y"], TRITON_MAX_BLOCK["Y"])
        and (
            y * maxGridSize[1] < size_hints["y"]
            or conditional_product(x, y, z) < target
        )
    ):
        y *= 2
    while (
        z
        and z < min(size_hints["z"], TRITON_MAX_BLOCK["Z"])
        and (
            z * maxGridSize[2] < size_hints["z"]
            or conditional_product(x, y, z) < target
        )
    ):
        z *= 2

    # Calculate num_warps if they are not hard passed to config
    if num_warps is None:
        num_warps = _num_warps(
            conditional_product(x, y, z) // num_elements_per_warp,
            min_num_warps=1,
            warp_size=warp_size,
        )
    # we are going to arrive at 2 warps only if bs was too small due to
    # numel being too small. However to workaround some ptx bugs we still
    # want at least 4 warps if there's enough elements per thread
    # given that this is a rare situation, don't expect this to affect perf
    # in general
    # see https://github.com/pytorch/pytorch/pull/97950
    if conditional_product(x, y, z) >= 128 and not torch.version.hip:
        num_warps = max(num_warps, 4)
    xnumel = size_hints["x"]
    ynumel = size_hints.get("y")
    znumel = size_hints.get("z")

    # Increase x to satisfy min_elem_per_thread requirements.
    block_size = max(
        conditional_product(x, y, z),
        min_elem_per_thread * warp_size * num_warps,
    )
    x *= math.ceil(block_size / conditional_product(x, y, z))

    x, _num_blocks = _check_max_grid_x(size_hints, x, num_warps, warp_size=warp_size)
    x = min(x, size_hints["x"])

    cfg = {"XBLOCK": x}
    if y:
        cfg["YBLOCK"] = y
    if z:
        cfg["ZBLOCK"] = z
    check_max_block(cfg)
    check_config(cfg, xnumel=xnumel, ynumel=ynumel, znumel=znumel)
    config = Config(cfg, num_warps=num_warps, num_stages=num_stages)

    if torch.version.hip:
        if matrix_instr is not None:
            config.kwargs["matrix_instr_nonkdim"] = matrix_instr
        if waves_per_eu is not None:
            config.kwargs["waves_per_eu"] = waves_per_eu
        if kpack is not None:
            config.kwargs["kpack"] = kpack

    return config


def _get_nd_reduction_numels(r: int, size_hints: dict[str, int]) -> dict[str, int]:
    """
    Converts a linear reduction numel to ND, in row major order.
    This order is often desirable as it presents opportunities to coalesce memory
    accesses.
    For example, if r = 64 and size_hints = [32,32], this function returns [32, 2].
    This unraveling works because both r and size_hints are powers of 2.
    """
    # Shrink r to size_hints.
    r = min(r, get_total_reduction_numel(size_hints))
    num_reduction_dims = len(
        [prefix for prefix in size_hints if prefix_is_reduction(prefix)]
    )

    remaining = r
    rnumels = {}
    for idx in range(num_reduction_dims - 1, -1, -1):
        prefix = f"r{idx}_"
        max_size = min(size_hints[prefix], TRITON_MAX_BLOCK[prefix.upper()])
        dim = min(max_size, remaining)
        if remaining % dim != 0:
            raise AssertionError(
                f"Expected dimension '{dim}' to divide remaining size '{remaining}'"
            )
        rnumels[prefix] = dim
        remaining //= dim

    # Sanity check the results.
    final_numel = conditional_product(*rnumels.values())
    if r != final_numel:
        raise AssertionError(
            f"Expected ND reduction size ({rnumels}) to have {r} elements."
        )
    if not all(rnumels[prefix] <= size_hints[prefix] for prefix in rnumels):
        raise AssertionError(f"rnumels exceed size_hints. {rnumels} > {size_hints}")

    return rnumels


def triton_config_reduction(
    size_hints,
    x: int,
    r: int,
    num_stages=1,
    num_warps=None,
    register_intensive=False,
    waves_per_eu=None,
    dynamic_scale_rblock=True,
    reduction_hint=None,
    min_num_warps=None,
    *,
    warp_size: int = 32,
) -> Config:
    """
    Construct a reduction triton config with some adjustment heuristics
    based on size_hints. Size_hints is a tuple of numels in each tile
    dimension and will be rounded up to the nearest power of 2.
    """
    # Convert the linear reduction numel into a multi-dimensional block.
    rnumels = _get_nd_reduction_numels(r, size_hints)

    # shrink sizes to size hints
    x = min(x, size_hints["x"])
    target = conditional_product(x, *rnumels.values())
    if conditional_product(*size_hints.values()) < target:
        target //= 8

    def total_numel() -> int:
        return conditional_product(x, *rnumels.values())

    # if we are below original block size, scale up where we can
    while x < size_hints["x"] and total_numel() < target:
        x *= 2
    for prefix in sorted(rnumels):
        while rnumels[prefix] < size_hints[prefix] and total_numel() < target:
            rnumels[prefix] *= 2

    if num_warps is None:
        if reduction_hint == ReductionHint.INNER:
            # r is contiguous, ensure at least 8 elements per thread
            # xblock is usually 1-2, default to giving each thread more work
            num_warps = r // 128
        else:
            num_warps = total_numel() // 128

    max_num_warps = 16 if r <= 8192 else 32
    if min_num_warps is not None:
        _num_warps_func = functools.partial(_num_warps, min_num_warps=min_num_warps)
    else:
        _num_warps_func = _num_warps

    num_warps = _num_warps_func(
        num_warps,
        max_num_warps=max_num_warps,
        register_intensive=register_intensive,
        warp_size=warp_size,
    )

    x, _num_blocks = _check_max_grid_x(size_hints, x, num_warps, warp_size=warp_size)

    for prefix in sorted(rnumels):
        while total_numel() > target:
            if rnumels[prefix] == 1:
                break
            rnumels[prefix] //= 2

    cfg = _get_config({"x": x, **rnumels})
    check_max_block(cfg)
    check_config(cfg, xnumel=size_hints["x"])
    config = InductorConfig(
        cfg,
        num_warps=num_warps,
        num_stages=num_stages,
        dynamic_scale_rblock=dynamic_scale_rblock,
    )

    if torch.version.hip:
        if waves_per_eu is not None:
            config.kwargs["waves_per_eu"] = waves_per_eu

    return config


def _get_config(numels: dict[str, int]) -> dict[str, int]:
    """
    Convert numels ("x", "r0_", etc.) to block sizes ("XBLOCK", "R0_BLOCK"), etc.
    """

    return {prefix.upper() + "BLOCK": numel for prefix, numel in numels.items()}


def _subkernel_fingerprint(combo_meta: dict[str, Any], i: int) -> tuple[Any, ...]:
    """Per-sub-kernel heuristic inputs as a hashable tuple. Identical
    fingerprints imply identical heuristic output.

    Per-kernel fields (num_load, autotune_hints, tiling_scores, etc.) live
    inside combo_meta[f"inductor_meta_{i}"] (single source of truth — see
    TritonKernel.inductor_meta_per_kernel). Combo-level fields (heuristic,
    size_hints, tile_hint, reduction_hint) remain top-level in combo_meta.
    """
    sub_meta = combo_meta.get(f"inductor_meta_{i}", {})
    tma = sub_meta.get("tma_min_block_sizes") or {}
    tiling_scores = sub_meta.get("tiling_scores") or {}
    return (
        combo_meta[f"heuristic_{i}"],
        tuple(sorted(combo_meta[f"size_hints_{i}"].items())),
        sub_meta.get("num_load"),
        sub_meta.get("num_store"),
        sub_meta.get("num_reduction"),
        tuple(sorted(sub_meta.get("autotune_hints") or [], key=str)),
        sub_meta.get("atomic_add_found"),
        sub_meta.get("no_x_dim"),
        combo_meta.get(f"reduction_hint_{i}"),
        combo_meta.get(f"tile_hint_{i}"),
        sub_meta.get("add_persistent_rblock", False),
        sub_meta.get("has_loadstore_with_contiguous_rdim"),
        sub_meta.get("uses_device_tma", False),
        tuple(sorted(tma.items())),
        tuple(sorted(tiling_scores.items())),
    )


def _update_combo_kernel_kwargs(
    kwargs: dict[str, Any],
    cfg_kwargs: dict[str, Any],
    subkernel_idx: int,
    skip_rblock: bool,
    block_arg_names: OrderedSet[str],
) -> None:
    for key, value in cfg_kwargs.items():
        if skip_rblock and key.startswith("R") and "BLOCK" in key:
            continue
        suffixed_key = f"{key}_{subkernel_idx}"
        # Only suffix keys emitted as combo kernel block arguments.
        # Everything else must stay unsuffixed so HIP-specific compile options like
        # waves_per_eu continue to flow through the backend-options path above.
        kwargs[suffixed_key if suffixed_key in block_arg_names else key] = value


def _handle_combo_kernel_per_subkernel_blocks(
    size_hints: dict[str, int],
    inductor_meta: InductorMeta,
    triton_meta: TritonMeta,
    filename: str | None = None,
    reduction_hint: bool = False,
    tile_hint: Any = None,
    min_elem_per_thread: int = 0,
) -> list[Config] | None:
    """
    Handle per-subkernel config generation for combo kernels.

    Each sub-kernel gets its own block sizes (XBLOCK_0, XBLOCK_1, etc.) generated
    using the same heuristics as standalone Triton kernels.

    Returns base configs that vary (num_warps, num_stages) with all blocks at
    heuristic defaults. Stores per-subkernel candidate configs in
    inductor_meta["combo_tuning_groups"] for sequential chained autotuning
    in CachingAutotuner._combo_sequential_autotune().

    Returns:
        List of configs if combo kernel with combo_grid_meta and per-subkernel
        blocks enabled, None otherwise.
    """
    combo_meta = inductor_meta.get("combo_grid_meta")
    if combo_meta is None or "heuristic_0" not in combo_meta:
        return None

    # CAP no-bench: combo_grid_meta carries a stitched config; skip the
    # combo_tuning_groups computation and return a single stitched config.
    # default_config holds BLOCK keys for the grid lambda; backend kwargs
    # (HIP options like waves_per_eu) come from stitched_backend_kwargs.
    stitched_warps = combo_meta.get("stitched_num_warps")
    if "stitched_launch_candidates" in combo_meta or stitched_warps is not None:
        # Compile-time autotune emits the distinct winner launch configs (kwargs, num_warps,
        # num_stages) -> combo autotunes kernel-level knobs over them; the chosen block sizes
        # are passed as args via default_config. No-bench mode has no candidates and reuses
        # default_config for the explicitly recorded combo block arguments.
        if "stitched_launch_candidates" in combo_meta:
            launch_candidates = combo_meta["stitched_launch_candidates"]
            block_config = combo_meta.get("default_config") or {}
            return [
                triton.Config({**block_config, **kwargs}, num_warps=nw, num_stages=ns)
                for kwargs, nw, ns in launch_candidates
            ]
        block_arg_names = OrderedSet(combo_meta.get("block_arg_names", ()))
        block_config = {
            k: v
            for k, v in (combo_meta.get("default_config") or {}).items()
            if k in block_arg_names
        }
        return [
            triton.Config(
                {**block_config, **combo_meta["stitched_backend_kwargs"]},
                num_warps=stitched_warps,
                num_stages=combo_meta["stitched_num_stages"],
            )
        ]

    num_kernels = combo_meta["num_kernels"]
    inductor_meta_clean = {
        k: v for k, v in inductor_meta.items() if k != "combo_grid_meta"
    }

    combined_kwargs: dict[str, int] = {}
    all_num_warps: list[int] = []
    all_num_stages: list[int] = []
    unique_warp_stage_pairs: OrderedSet[tuple[int, int]] = OrderedSet()
    combo_coordesc_field_limits: dict[str, int] = {}
    block_arg_names = OrderedSet(combo_meta.get("block_arg_names", ()))

    # Group sub-kernels with identical config kwargs to skip redundant tuning.
    group_map: dict[tuple[Any, ...], dict[str, Any]] = {}

    for i in range(num_kernels):
        subkernel_heuristic = combo_meta[f"heuristic_{i}"]
        size_hints_i = combo_meta[f"size_hints_{i}"]
        # Per-sub-kernel inductor_meta passthrough packed by combo_grid_meta()
        # via TritonKernel.inductor_meta_per_kernel(). Forward into
        # inductor_meta_i so pointwise()/_reduction_configs()/_persistent_reduction_configs()
        # pick configs based on the actual sub-kernel .
        inductor_meta_i = cast(
            "InductorMeta",
            {
                **inductor_meta_clean,
                **combo_meta.get(f"inductor_meta_{i}", {}),
            },
        )

        if subkernel_heuristic == "pointwise":
            cfgs = pointwise(
                size_hints_i,
                triton_meta=triton_meta,
                tile_hint=(
                    TileHint.SQUARE
                    if combo_meta[f"tile_hint_{i}"] == "TileHint.SQUARE"
                    else TileHint.DEFAULT
                ),
                filename=filename,
                min_elem_per_thread=min_elem_per_thread,
                inductor_meta=inductor_meta_i,
                return_configs=True,
            )
            skip_rblock = False
        elif subkernel_heuristic == "reduction":
            cfgs = reduction(
                size_hints_i,
                reduction_hint=ReductionHint[combo_meta[f"reduction_hint_{i}"]],
                triton_meta=triton_meta,
                filename=filename,
                inductor_meta=inductor_meta_i,
                return_configs=True,
            )
            skip_rblock = False
        elif subkernel_heuristic == "persistent_reduction":
            cfgs = persistent_reduction(
                size_hints_i,
                reduction_hint=ReductionHint[combo_meta[f"reduction_hint_{i}"]],
                triton_meta=triton_meta,
                filename=filename,
                inductor_meta=inductor_meta_i,
                return_configs=True,
            )
            skip_rblock = True  # persistent reduction embeds RBLOCK in kernel body
        else:
            raise ValueError(f"Unknown heuristic: {subkernel_heuristic}")

        group_coordesc_fields: OrderedSet[str] = OrderedSet()
        cfg = cfgs[0]
        _update_combo_kernel_kwargs(
            combined_kwargs, cfg.kwargs, i, skip_rblock, block_arg_names
        )
        for key in cfg.kwargs:
            if skip_rblock and key.startswith("R") and "BLOCK" in key:
                continue
            if not key.endswith("BLOCK"):
                continue
            combined_key = f"{key}_{i}"
            group_coordesc_fields.add(combined_key)
            prefix = key.removesuffix("BLOCK").lower()
            if prefix in size_hints_i:
                combo_coordesc_field_limits[combined_key] = min(
                    TRITON_MAX_BLOCK[prefix.upper()],
                    size_hints_i[prefix],
                )

        all_num_warps.append(cfg.num_warps)
        all_num_stages.append(cfg.num_stages)
        for c in cfgs:
            unique_warp_stage_pairs.add((c.num_warps, c.num_stages))

        group_key = (
            _subkernel_fingerprint(combo_meta, i)
            if combo_meta.get("autotune_grouping")
            else (i,)
        )
        if group_key in group_map:
            group_map[group_key]["member_indices"].append(i)
        else:
            group_map[group_key] = {
                "member_indices": [i],
                "configs": cfgs,
                "skip_rblock": skip_rblock,
                "size_hints": size_hints_i,
                "coordesc_fields": list(group_coordesc_fields),
            }

    unique_warp_stage_pairs.add((max(all_num_warps), max(all_num_stages)))

    combo_tuning_groups = list(group_map.values())
    # Largest sub-kernels tuned first — they dominate runtime and get most freedom
    combo_tuning_groups.sort(
        key=lambda g: -functools.reduce(operator.mul, g["size_hints"].values())
    )
    inductor_meta["combo_tuning_groups"] = combo_tuning_groups
    inductor_meta["combo_coordesc_field_order"] = [
        field for group in combo_tuning_groups for field in group["coordesc_fields"]
    ]
    inductor_meta["combo_coordesc_field_limits"] = combo_coordesc_field_limits
    # Candidates for num_warps/num_stages re-tuning after block sizes are finalized
    inductor_meta["combo_warp_stage_candidates"] = list(unique_warp_stage_pairs)

    # Single base config: max warps/stages, all blocks at heuristic defaults.
    # Block sizes are tuned in _combo_sequential_autotune, then num_warps/num_stages
    # are re-tuned at the end with finalized block sizes.
    base_num_warps = max(all_num_warps)
    base_num_stages = max(all_num_stages)
    return [
        triton.Config(
            combined_kwargs,
            num_warps=base_num_warps,
            num_stages=base_num_stages,
        )
    ]


def triton_config_tiled_reduction(
    size_hints,
    x,
    y,
    r,
    num_stages=1,
    register_intensive=False,
    waves_per_eu=None,
    *,
    warp_size: int = 32,
):
    """
    Construct a tile reduction triton config with some adjustment
    heuristics based on size_hints. Size_hints is a tuple of numels in
    each tile dimension and will be rounded up to the nearest power of 2.
    """
    # Convert the linear reduction numel into a multi-dimensional block.
    rnumels = _get_nd_reduction_numels(r, size_hints)

    # shrink sizes to size hints
    x = min(x, size_hints["x"])
    y = min(y, size_hints["y"])

    def total_numel() -> int:
        return conditional_product(x, y, *rnumels.values())

    target = total_numel()
    if conditional_product(*size_hints.values()) < target:
        target //= 8

    # if we are below original block size, scale up where we can
    while x < size_hints["x"] and total_numel() < target:
        x *= 2
    for prefix in sorted(rnumels):
        while rnumels[prefix] < size_hints[prefix] and total_numel() < target:
            rnumels[prefix] *= 2
    while y < size_hints["y"] and total_numel() < target:
        y *= 2

    cfg = _get_config({"x": x, "y": y, **rnumels})
    num_warps = _num_warps(total_numel() // 256, min_num_warps=1, warp_size=warp_size)
    num_warps = _num_warps(
        num_warps,
        max_num_warps=16,
        register_intensive=register_intensive,
        warp_size=warp_size,
    )
    check_config(cfg, xnumel=size_hints["x"], ynumel=size_hints["y"])
    check_max_block(cfg)
    config = Config(cfg, num_warps=num_warps, num_stages=num_stages)
    if torch.version.hip:
        if waves_per_eu is not None:
            config.kwargs["waves_per_eu"] = waves_per_eu
    return config


def _maybe_filter_configs_for_tma_restrictions(
    inductor_meta: InductorMeta, configs: list[Config]
):
    tma_min_block_sizes: dict[str, int] | None
    if (tma_min_block_sizes := inductor_meta.get("tma_min_block_sizes")) and configs:
        # Rn blocks are not provided to the kernel for persistent reductions
        if inductor_meta.get("persistent_reduction"):
            tma_min_block_sizes = {
                block_type: block_size
                for block_type, block_size in tma_min_block_sizes.items()
                if not prefix_is_reduction(block_type.lower())
            }

        if not all(
            block_type in configs[0].kwargs for block_type in tma_min_block_sizes
        ):
            raise AssertionError(
                f"Not all TMA block types found in config kwargs: "
                f"missing {OrderedSet(tma_min_block_sizes) - OrderedSet(configs[0].kwargs)}"
            )

        # Add a config that is guaranteed to compile
        example_config = configs[0]
        config_block_sizes = {**example_config.kwargs}
        for block_type, min_block_value in tma_min_block_sizes.items():
            existing = config_block_sizes.get(block_type, 1)
            config_block_sizes[block_type] = max(existing, min_block_value)
        new_configs = [
            Config(
                config_block_sizes,
                num_warps=example_config.num_warps,
                num_stages=example_config.num_stages,
                maxnreg=example_config.maxnreg,
                pre_hook=example_config.pre_hook,
            )
        ]
        # Remove configs that will not compile
        for c in configs:
            if all(
                c.kwargs.get(block_type) >= min_block_value
                for block_type, min_block_value in tma_min_block_sizes.items()
            ):
                new_configs.append(c)

        log.debug(
            "Filtering configs for TMA API restrictions. Input configs size: %d. Output configs size: %d",
            len(configs),
            len(new_configs),
        )
        return new_configs
    return configs


def pointwise(
    size_hints,
    triton_meta: TritonMeta,
    tile_hint=None,
    filename=None,
    min_elem_per_thread=0,
    inductor_meta: InductorMeta | None = None,
    return_configs=False,
):
    """
    Construct @triton.heuristics() based on size_hints.
    """
    inductor_meta = {} if inductor_meta is None else inductor_meta

    configs = _handle_combo_kernel_per_subkernel_blocks(
        size_hints,
        inductor_meta,
        triton_meta,
        filename=filename,
        tile_hint=tile_hint,
        min_elem_per_thread=min_elem_per_thread,
    )
    if configs is not None:
        return cached_autotune(
            None,
            configs,
            triton_meta=triton_meta,
            inductor_meta=inductor_meta,
            heuristic_type=HeuristicType.POINTWISE,
            filename=filename,
        )

    if inductor_meta.get("no_x_dim"):
        raise AssertionError("no_x_dim should not be set for this heuristic")

    numel = functools.reduce(operator.mul, size_hints.values())
    bs = max(256, min(numel // 128, 1024))

    hinted_configs = autotune_hints_to_configs(
        inductor_meta.get("autotune_hints", OrderedSet()),
        size_hints,
        bs,
        triton_meta["device"],
    )

    warp_size = triton_meta["device"].warp_size_or_default
    triton_config_with_settings = functools.partial(
        triton_config, min_elem_per_thread=min_elem_per_thread, warp_size=warp_size
    )

    from torch._inductor.heuristics.registry import get_codegen_heuristic

    pointwise_heuristic = get_codegen_heuristic("pointwise", triton_meta["device"].type)
    configs = pointwise_heuristic.get_configs(
        size_hints,
        bs,
        triton_config_with_settings,
        hinted_configs,
        tile_hint=tile_hint,
        inductor_meta=inductor_meta,
    )

    configs = _maybe_filter_configs_for_tma_restrictions(inductor_meta, configs)
    if return_configs:
        return configs

    return cached_autotune(
        size_hints,
        configs,
        triton_meta=triton_meta,
        inductor_meta=inductor_meta,
        heuristic_type=HeuristicType.POINTWISE,
        filename=filename,
    )


def make_matmul_triton_config(sizes: dict[str, int], num_warps: int, num_stages: int):
    config = {
        "XBLOCK": sizes.get("x"),
        "YBLOCK": sizes.get("y"),
        "ZBLOCK": sizes.get("z"),
        "R0_BLOCK": sizes.get("r"),
    }
    # Remove keys with None values (i.e., missing in sizes)
    config = {k: v for k, v in config.items() if v is not None}
    _check_native_matmul_block_numel(config)
    return Config(config, num_warps=num_warps, num_stages=num_stages)


def _config_helper(bmm=False, persistent=False):
    # Each entry is: (sizes_dict, num_warps, num_stages)
    _base_mm_configs = [
        ({"x": 32, "y": 32, "r": 16}, 2, 1),
        ({"x": 32, "y": 32, "r": 128}, 4, 2),
        ({"x": 32, "y": 64, "r": 32}, 8, 5),
        ({"x": 64, "y": 32, "r": 32}, 8, 5),
        ({"x": 64, "y": 32, "r": 128}, 4, 5),
        ({"x": 64, "y": 64, "r": 16}, 4, 2),
        ({"x": 64, "y": 64, "r": 32}, 4, 2),
        ({"x": 64, "y": 64, "r": 64}, 8, 3),
        ({"x": 64, "y": 64, "r": 128}, 4, 5),
        ({"x": 64, "y": 128, "r": 32}, 4, 3),
        ({"x": 64, "y": 128, "r": 32}, 8, 4),
        ({"x": 64, "y": 128, "r": 64}, 4, 3),
        ({"x": 64, "y": 128, "r": 128}, 4, 4),
        ({"x": 128, "y": 64, "r": 32}, 4, 3),
        ({"x": 128, "y": 64, "r": 32}, 8, 4),
        ({"x": 128, "y": 128, "r": 32}, 8, 2),
        ({"x": 128, "y": 128, "r": 32}, 4, 3),
        ({"x": 128, "y": 128, "r": 64}, 4, 3),
        ({"x": 128, "y": 128, "r": 64}, 8, 5),
    ]
    out = []
    for sizes, w, s in _base_mm_configs:
        d = dict(sizes)
        if persistent:
            d.pop("r", None)
        if bmm:
            d["z"] = 1
        out.append((d, w, s))

    # Deduplicate by converting dicts to immutable frozensets
    deduped = {(frozenset(d.items()), w, s): (d, w, s) for d, w, s in out}

    return list(deduped.values())


triton_native_mm_configs = _config_helper(bmm=False, persistent=False)
triton_native_persistent_mm_configs = _config_helper(bmm=False, persistent=True)
triton_native_bmm_configs = _config_helper(bmm=True, persistent=False)
triton_native_persistent_bmm_configs = _config_helper(bmm=True, persistent=True)


def _reduction_configs(
    *,
    size_hints: dict[str, int],
    inductor_meta: InductorMeta,
    triton_meta: TritonMeta,
    num_dynamic=0,
) -> list[Config]:
    from torch._inductor.heuristics.registry import get_codegen_heuristic

    reduction_heuristic = get_codegen_heuristic("reduction", triton_meta["device"].type)
    configs = reduction_heuristic.get_configs(
        size_hints=size_hints,
        inductor_meta=inductor_meta,
        triton_meta=triton_meta,
        num_dynamic=num_dynamic,
    )
    r0 = inductor_meta.get("strict_sum_rblock")
    if r0 is not None:
        configs = copy.deepcopy(configs)
        for triton_config in configs:
            if "R0_BLOCK" in triton_config.kwargs:
                triton_config.kwargs["R0_BLOCK"] = r0
        configs = unique_configs(configs)
    return configs


def filter_reduction_configs_for_determinism(
    inductor_meta: InductorMeta, configs: list[Config]
) -> list[Config]:
    """
    Filter configs for reduction so the numerics can be deterministic.

    Heuristics:
    - skip reduction configs with too small RBLOCK
    - skip reduction configs with XBLOCK==1 if we are confident it will not perform well
    - if there is a tie, pick the config with second largest RBLOCK
    - if there is still a tie, pick the config with second largest num_warps
    - if there is still a tie, pick the config with second largest XBLOCK
    """
    configs = unique_configs(configs)
    if len(configs) == 0:
        raise AssertionError("No configs remaining after deduplication")

    def _do_filter_due_to_inductor_config():
        return (
            inductor_meta.get("deterministic", False)
            or inductor_meta.get("force_filter_reduction_configs", False)
        ) or inductor_meta.get("are_deterministic_algorithms_enabled")

    if not _do_filter_due_to_inductor_config() or len(configs) == 1:
        # no filtering happening if NOT in deterministic mode
        return configs

    if log.isEnabledFor(logging.DEBUG):
        log.debug("reduction configs before filtering:")
        for c in configs:
            log.debug("%s", c)
            log.debug("")

    def _has_too_small_rblock(config):
        rblock = config.kwargs.get("R0_BLOCK")
        # too small RBLOCK is likely to be bad
        return rblock is not None and rblock <= 4

    def _nonpromising_xblock_1(config):
        # kernel like https://gist.github.com/shunting314/0b3281c087e79bc915fe45985ff9d7d5
        # without a load/store having contiguous rdim is unlikely to perform well with XBLOCK==1
        return config.kwargs["XBLOCK"] == 1 and not inductor_meta.get(
            "has_loadstore_with_contiguous_rdim", True
        )

    newconfigs = [*filter(lambda x: not _has_too_small_rblock(x), configs)]
    # accept the filtering only if there are configs left
    if len(newconfigs) > 0:
        configs = newconfigs

    newconfigs = [*filter(lambda x: not _nonpromising_xblock_1(x), configs)]
    if len(newconfigs) > 0:
        configs = newconfigs

    if len(configs) == 0:
        raise AssertionError("No configs remaining after filtering")

    def _r0_block(c):
        return c.kwargs.get("R0_BLOCK", -1)

    def _xblock(c):
        return c.kwargs.get("XBLOCK", -1)

    def _num_warps(c):
        return c.num_warps

    def _pick_second_largest(accessor):
        nonlocal configs
        configs = sorted(configs, key=lambda x: accessor(x))
        if accessor(configs[0]) != accessor(configs[-1]):
            max_val = accessor(configs[-1])
            configs = [*filter(lambda x: accessor(x) != max_val, configs)]
            second_max_val = accessor(configs[-1])
            configs = [*filter(lambda x: accessor(x) == second_max_val, configs)]
        return configs

    def _pick_config():
        nonlocal configs
        if len(configs) == 0:
            raise AssertionError("No configs available for selection")
        if len(configs) == 1:
            return configs[0]

        # break tie by R0_BLOCK
        configs = _pick_second_largest(_r0_block)
        if len(configs) == 1:
            return configs[0]

        # break tie by num_warps
        configs = _pick_second_largest(_num_warps)
        if len(configs) == 1:
            return configs[0]

        # break tie by XBLOCK
        configs = _pick_second_largest(_xblock)

        # there is still a tie, pick the first one
        return configs[0]

    configs = [_pick_config()]

    if log.isEnabledFor(logging.DEBUG):
        log.debug("reduction configs after filtering:")
        for c in configs:
            log.debug("%s", c)
            log.debug("")
    return configs


def reduction(
    size_hints,
    reduction_hint=False,
    triton_meta: TritonMeta | None = None,
    filename=None,
    inductor_meta: InductorMeta | None = None,
    return_configs=False,
):
    """args to @triton.heuristics()"""
    inductor_meta = {} if inductor_meta is None else inductor_meta
    inductor_meta["reduction_hint"] = reduction_hint
    if inductor_meta.get("no_x_dim"):
        size_hints["x"] = 1

    if triton_meta is None:
        raise AssertionError("triton_meta must not be None")

    configs = _handle_combo_kernel_per_subkernel_blocks(
        size_hints,
        inductor_meta,
        triton_meta,
        filename=filename,
        reduction_hint=reduction_hint,
    )
    if configs is not None:
        return cached_autotune(
            None,
            configs,
            triton_meta=triton_meta,
            inductor_meta=inductor_meta,
            heuristic_type=HeuristicType.REDUCTION,
            filename=filename,
        )

    if triton_meta is None:
        raise AssertionError("triton_meta must not be None")

    num_dynamic = 0
    for k in triton_meta["signature"]:
        if "ks" in k:
            num_dynamic += 1

    configs = _reduction_configs(
        size_hints=size_hints,
        inductor_meta=inductor_meta,
        triton_meta=triton_meta,
        num_dynamic=num_dynamic,
    )

    configs = _maybe_filter_configs_for_tma_restrictions(inductor_meta, configs)
    configs = filter_reduction_configs_for_determinism(inductor_meta, configs)
    strict_rblock = inductor_meta.get("strict_sum_rblock")
    if strict_rblock is not None and any(
        triton_config.kwargs.get("R0_BLOCK", strict_rblock) != strict_rblock
        for triton_config in configs
    ):
        raise AssertionError("strict sum requires its planned R0_BLOCK")

    if return_configs:
        return configs

    return cached_autotune(
        size_hints,
        configs=configs,
        triton_meta=triton_meta,
        inductor_meta=inductor_meta,
        heuristic_type=HeuristicType.REDUCTION,
        filename=filename,
    )


def cooperative_reduction(
    size_hints,
    reduction_hint,
    triton_meta: TritonMeta,
    filename,
    inductor_meta: InductorMeta | None = None,
):
    inductor_meta = {} if inductor_meta is None else inductor_meta
    inductor_meta["reduction_hint"] = reduction_hint
    if inductor_meta.get("no_x_dim"):
        size_hints["x"] = 1

    from torch._inductor.heuristics.registry import get_codegen_heuristic

    reduction_heuristic = get_codegen_heuristic("reduction", triton_meta["device"].type)
    configs = reduction_heuristic.get_cooperative_configs(
        size_hints=size_hints,
        reduction_hint=reduction_hint,
        inductor_meta=inductor_meta,
        triton_meta=triton_meta,
    )
    # TODO(jansel): add more configs in max_autotune

    configs = _maybe_filter_configs_for_tma_restrictions(inductor_meta, configs)
    configs = filter_reduction_configs_for_determinism(inductor_meta, configs)
    return cached_autotune(
        size_hints,
        configs=configs,
        triton_meta=triton_meta,
        inductor_meta=inductor_meta,
        heuristic_type=HeuristicType.REDUCTION,
        filename=filename,
    )


def _persistent_reduction_configs(
    size_hints,
    reduction_hint=False,
    inductor_meta: InductorMeta | None = None,
    triton_meta: TritonMeta | None = None,
):
    from torch._inductor.heuristics.registry import get_codegen_heuristic

    reduction_heuristic = get_codegen_heuristic("reduction", triton_meta["device"].type)
    return reduction_heuristic.get_persistent_configs(
        size_hints=size_hints,
        reduction_hint=reduction_hint,
        inductor_meta=inductor_meta,
        triton_meta=triton_meta,
    )


def persistent_reduction(
    size_hints,
    reduction_hint=False,
    triton_meta: TritonMeta | None = None,
    filename=None,
    inductor_meta: InductorMeta | None = None,
    return_configs=False,
):
    """Generate persistent reductions + mix-order if available"""
    inductor_meta = {} if inductor_meta is None else inductor_meta
    inductor_meta["reduction_hint"] = reduction_hint
    if inductor_meta.get("no_x_dim"):
        size_hints["x"] = 1

    if triton_meta is None:
        raise AssertionError("triton_meta must not be None")

    configs = _handle_combo_kernel_per_subkernel_blocks(
        size_hints,
        inductor_meta,
        triton_meta,
        filename=filename,
        reduction_hint=reduction_hint,
    )
    if configs is not None:
        return cached_autotune(
            None,
            configs,
            triton_meta=triton_meta,
            inductor_meta=inductor_meta,
            heuristic_type=HeuristicType.PERSISTENT_REDUCTION,
            filename=filename,
        )

    configs = _persistent_reduction_configs(
        size_hints, reduction_hint, inductor_meta, triton_meta
    )

    # This key is not added to the inductor meta as its clear from the heuristic
    # choice that it is persistent. Add it and remove it below so that persistent
    # configs can be filtered appropriately by _maybe_filter_configs_for_tma_restrictions
    persistent_reduction_key = "persistent_reduction"
    inductor_meta[persistent_reduction_key] = True
    configs = _maybe_filter_configs_for_tma_restrictions(inductor_meta, configs)
    inductor_meta.pop(persistent_reduction_key)

    if inductor_meta.get("RSPLIT_SIZE"):
        from torch._inductor.heuristics.registry import get_codegen_heuristic

        reduction_heuristic = get_codegen_heuristic(
            "reduction", triton_meta["device"].type
        )
        configs = reduction_heuristic.apply_rsplit_size(
            configs,
            size_hints=size_hints,
            inductor_meta=inductor_meta,
        )

    configs = filter_reduction_configs_for_determinism(inductor_meta, configs)

    if return_configs:
        return configs

    return cached_autotune(
        size_hints,
        configs,
        triton_meta=triton_meta,
        inductor_meta=inductor_meta,
        filename=filename,
        heuristic_type=HeuristicType.PERSISTENT_REDUCTION,
    )


def split_scan(
    size_hints,
    reduction_hint=False,
    triton_meta: TritonMeta | None = None,
    filename=None,
    inductor_meta: InductorMeta | None = None,
):
    """Heuristic for TritonSplitScanKernel"""
    inductor_meta = {} if inductor_meta is None else inductor_meta
    inductor_meta["reduction_hint"] = reduction_hint
    if inductor_meta.get("no_x_dim"):
        size_hints["x"] = 1

    if triton_meta is None:
        raise AssertionError("triton_meta must not be None")
    if len(size_hints) != 2:
        raise NotImplementedError(f"size_hints: {size_hints}")

    from torch._inductor.heuristics.registry import get_codegen_heuristic

    reduction_heuristic = get_codegen_heuristic("reduction", triton_meta["device"].type)
    configs = reduction_heuristic.get_split_scan_configs(
        size_hints=size_hints,
        inductor_meta=inductor_meta,
        triton_meta=triton_meta,
    )

    configs = _maybe_filter_configs_for_tma_restrictions(inductor_meta, configs)
    configs = filter_reduction_configs_for_determinism(inductor_meta, configs)
    return cached_autotune(
        size_hints,
        configs=configs,
        triton_meta=triton_meta,
        inductor_meta=inductor_meta,
        heuristic_type=HeuristicType.SPLIT_SCAN,
        filename=filename,
    )


def template(
    num_stages,
    num_warps,
    triton_meta: TritonMeta,
    num_consumer_groups=0,
    num_buffers_warp_spec=0,
    filename=None,
    inductor_meta: InductorMeta | None = None,
    **kwargs,
):
    """
    Compile a triton template
    """
    # Prepare the base configuration
    config_args = {
        "num_stages": num_stages,
        "num_warps": num_warps,
    }

    # Conditionally add arguments based on HAS_WARP_SPEC
    if HAS_WARP_SPEC:
        config_args.update(
            {
                "num_consumer_groups": num_consumer_groups,
                "num_buffers_warp_spec": num_buffers_warp_spec,
            }
        )

    for k in tlx_only_cuda_options():
        if v := triton_meta.get(k, None):
            config_args[k] = v

    return cached_autotune(
        None,
        [triton.Config({}, **config_args)],
        triton_meta=triton_meta,
        inductor_meta=inductor_meta,
        heuristic_type=HeuristicType.TEMPLATE,
        filename=filename,
    )


def _pop_config_kwargs(config: dict[str, Any]) -> dict[str, Any]:
    """Extract triton.Config options that should become kwargs"""
    popped = {}
    for key in (
        "num_warps",
        "num_stages",
        "num_ctas",
        "maxnreg",
        "num_consumer_groups",
        "num_buffers_warp_spec",
    ):
        val = config.pop(key, None)
        if val is not None:
            popped[key] = val
    return popped


def config_to_dict(config: Config) -> dict[str, Any]:
    config_dict = {
        **config.kwargs,
        "num_warps": config.num_warps,
        "num_stages": config.num_stages,
    }
    if HAS_WARP_SPEC:
        config_dict.update(
            {
                "num_consumer_groups": getattr(config, "num_consumer_groups", 0),
                "num_buffers_warp_spec": getattr(config, "num_buffers_warp_spec", 0),
            }
        )
    return config_dict


def config_from_dict(config: dict[str, Any]) -> Config:
    config = {**config}
    return Config(config, **_pop_config_kwargs(config))


def fixed_config(
    config, filename, triton_meta: TritonMeta, inductor_meta: InductorMeta
):
    """
    Used when the configuration is already decided at compile time
    """
    config = {**config}
    return cached_autotune(
        None,
        [triton.Config(config, **_pop_config_kwargs(config))],
        triton_meta=triton_meta,
        inductor_meta=inductor_meta,
        heuristic_type=HeuristicType.FIXED,
        filename=filename,
    )


def user_autotune(
    configs,
    triton_meta: TritonMeta,
    filename=None,
    inductor_meta: InductorMeta | None = None,
    custom_kernel=False,
):
    """
    Compile a user defined triton kernel
    """
    if len(configs) == 0:
        configs = [triton.Config({})]
    else:
        configs = [*map(config_from_dict, configs)]
    return cached_autotune(
        None,
        configs,
        triton_meta=triton_meta,
        heuristic_type=HeuristicType.USER_AUTOTUNE,
        filename=filename,
        inductor_meta=inductor_meta,
        custom_kernel=custom_kernel,
    )


def foreach(
    triton_meta: TritonMeta, filename=None, inductor_meta: InductorMeta | None = None
):
    """
    Compile a triton foreach kernel
    """
    inductor_meta = {} if inductor_meta is None else inductor_meta
    configs = []

    # Naive autotuning path for num_warps
    if not (
        inductor_meta.get("max_autotune") or inductor_meta.get("max_autotune_pointwise")
    ):
        configs.append(triton.Config({}, num_stages=1, num_warps=8))
    else:
        for warps in [1, 2, 4, 8]:
            configs.append(triton.Config({}, num_stages=1, num_warps=warps))

    return cached_autotune(
        None,
        configs,
        triton_meta=triton_meta,
        inductor_meta=inductor_meta,
        heuristic_type=HeuristicType.TEMPLATE,
        filename=filename,
    )


@dataclasses.dataclass
class GridExpr:
    """Generate code for grid size expressions in launcher"""

    inductor_meta: InductorMeta
    mode: Literal["python", "cpp"] = "python"
    prefix: list[str] = dataclasses.field(default_factory=list)
    x_grid: str | int = 1
    y_grid: str | int = 1
    z_grid: str | int = 1

    def __post_init__(self) -> None:
        if self.mode not in ("python", "cpp"):
            raise AssertionError(f"mode must be 'python' or 'cpp', got {self.mode!r}")

    def generate(self, meta: dict[str, int], is_lazy: bool = False) -> None:
        raise NotImplementedError

    def ceildiv(self, numel: str | int, block: int | str | None) -> str | int:
        if block is None or block == 1:
            return numel
        if isinstance(numel, int) and isinstance(block, int):
            return ceildiv(numel, block)  # constant fold
        # This trick only works in python, where
        # negative integer division is floored
        if self.mode == "python":
            return f"-(({numel}) // -({block}))"
        # For cpp code gen
        return f"(({numel} + ({block} - 1)) / ({block}))"

    def maximum(self, seq: list[int | str]) -> int | str:
        """Codegen for max function with constant folding, constants are represented as int"""
        items = self._constant_fold(max, seq)
        if len(items) <= 1:
            return items[0]
        if self.mode == "python":
            return f"max({', '.join(map(str, items))})"
        # Cast int constants to (long) to avoid type deduction errors with std::max
        # when mixing long variables with int literals
        cpp_items = [f"(long){x}" if isinstance(x, int) else str(x) for x in items]
        return functools.reduce(lambda x, y: f"std::max({x}, {y})", cpp_items)

    def summation(self, seq: list[int | str]) -> int | str:
        """Codegen for sum function with constant folding, constants are represented as int"""
        items = self._constant_fold(sum, seq)
        if len(items) <= 1:
            return items[0]
        return " + ".join(map(str, items))

    def product(self, seq: list[int | str]) -> int | str:
        items = self._constant_fold(math.prod, seq)
        if len(items) <= 1:
            return items[0]
        return " * ".join(map(str, items))

    def _constant_fold(
        self, fn: Callable[[list[int]], int], seq: list[int | str]
    ) -> list[int | str]:
        """Constant fold through a commutative fn where ints are constants"""
        items: list[int | str] = [x for x in seq if not isinstance(x, int)]
        const_items = [x for x in seq if isinstance(x, int)]
        if const_items:
            items.append(fn(const_items))
        return items

    def assign_tmp(self, name: str, expr: str | int) -> str:
        # Grid functions are one per kernel, so name collisions are fine
        if self.mode == "python":
            return f"{name} = {expr}"
        if self.mode == "cpp":
            return f"uint32_t {name} = {expr};"
        raise AssertionError(f"invalid mode {self.mode}")

    @staticmethod
    def from_meta(
        inductor_meta: InductorMeta,
        cfg: Config | dict[str, int],
        mode: Literal["python", "cpp"] = "python",
    ) -> GridExpr:
        grid_cls = globals()[inductor_meta["grid_type"]]
        if not issubclass(grid_cls, GridExpr):
            raise AssertionError(f"Expected GridExpr subclass, got {grid_cls}")
        grid = grid_cls(inductor_meta=inductor_meta, mode=mode)
        if isinstance(cfg, Config):
            cfg = config_to_dict(cfg)
        grid.generate(cfg)
        return grid

    def eval_slow(self, meta: dict[str, int]) -> tuple[int, int, int]:
        scope = {**meta}
        for line in self.prefix:
            exec(line, scope)
        exec(f"grid_0 = {self.x_grid}", scope)
        exec(f"grid_1 = {self.y_grid}", scope)
        exec(f"grid_2 = {self.z_grid}", scope)
        return scope["grid_0"], scope["grid_1"], scope["grid_2"]

    def generate_lazy(self, kernel_name: str) -> None:
        """
        Creates a GridExpr for lazy compile, where config values are not known
        at codegen time and are instead referenced by variable names.
        """
        meta: dict[str, Any] = {
            "XBLOCK": f"{kernel_name}_result.xblocks[0]",
            "YBLOCK": f"{kernel_name}_result.yblocks[0]",
            "ZBLOCK": f"{kernel_name}_result.zblocks[0]",
            "R0_BLOCK": f"{kernel_name}_result.r0blocks[0]",
            "RSPLIT": f"{kernel_name}_result.rsplit",
            "RSPLIT_SIZE": f"{kernel_name}_result.rsplit_size",
        }
        # assertions are done based on real values, so we can skip here
        self.generate(meta, is_lazy=True)

    @classmethod
    def from_meta_lazy(
        cls,
        inductor_meta: InductorMeta | None,
        kernel_name: str,
    ) -> GridExpr:
        """Factory method for lazy compile mode."""
        if inductor_meta is None:
            raise AssertionError("inductor_meta must be specified for lazy compile")
        grid_type = inductor_meta.get("grid_type", None)
        if grid_type is None:
            raise AssertionError("grid_type must be specified for lazy compile")
        grid_cls = globals()[grid_type]
        if not issubclass(grid_cls, GridExpr):
            raise AssertionError(f"Expected GridExpr subclass, got {grid_cls}")
        grid = grid_cls(inductor_meta=inductor_meta, mode="cpp")
        grid.generate_lazy(kernel_name)
        return grid


class Grid1D(GridExpr):
    def generate(self, meta: dict[str, int], is_lazy: bool = False) -> None:
        self.x_grid = self.ceildiv("xnumel", meta.get("XBLOCK"))


class Grid2D(GridExpr):
    def generate(self, meta: dict[str, int], is_lazy: bool = False) -> None:
        self.x_grid = self.ceildiv("xnumel", meta.get("XBLOCK"))
        self.y_grid = self.ceildiv("ynumel", meta.get("YBLOCK"))


class Grid3D(GridExpr):
    def generate(self, meta: dict[str, int], is_lazy: bool = False) -> None:
        self.x_grid = self.ceildiv("xnumel", meta.get("XBLOCK"))
        self.y_grid = self.ceildiv("ynumel", meta.get("YBLOCK"))
        self.z_grid = self.ceildiv("znumel", meta.get("ZBLOCK"))


class BatchMatmulGrid3D(GridExpr):
    def generate(self, meta: dict[str, int], is_lazy: bool = False) -> None:
        self.z_grid = self.ceildiv("xnumel", meta.get("XBLOCK"))
        self.y_grid = self.ceildiv("ynumel", meta.get("YBLOCK"))
        self.x_grid = self.ceildiv("znumel", meta.get("ZBLOCK"))


class Grid2DWithYZOverflow(GridExpr):
    def generate(self, meta: dict[str, int], is_lazy: bool = False) -> None:
        self.x_grid = self.ceildiv("xnumel", meta.get("XBLOCK"))
        self.prefix.extend(
            [
                self.assign_tmp(
                    "y_grid_raw_", self.ceildiv("ynumel", meta.get("YBLOCK"))
                ),
                self.assign_tmp(
                    "y_grid_div_", self.ceildiv("y_grid_raw_", get_max_y_grid())
                ),
            ]
        )
        ceildiv_expr = self.ceildiv("y_grid_raw_", "y_grid_div_")
        if self.mode == "python":
            self.y_grid = f"(0 if y_grid_div_ == 0 else {ceildiv_expr})"
        else:
            self.y_grid = f"(y_grid_div_ == 0 ? 0 : {ceildiv_expr})"
        self.z_grid = "y_grid_div_"


class MixOrderReductionGrid(GridExpr):
    def generate(self, meta: dict[str, int], is_lazy: bool = False) -> None:
        split_size = meta.get("RSPLIT_SIZE")
        xblock = meta.get("XBLOCK")
        if not is_lazy:
            if not split_size:
                raise AssertionError("Missing RSPLIT_SIZE")
            if not xblock:
                raise AssertionError("Missing XBLOCK")
            if split_size % xblock != 0:
                raise AssertionError(f"{split_size=}, {xblock=}")
        self.x_grid = self.ceildiv("xnumel", split_size)


class CooperativeReductionGrid(GridExpr):
    def generate(self, meta: dict[str, int], is_lazy: bool = False) -> None:
        self.x_grid = str(meta["RSPLIT"])
        self.y_grid = self.ceildiv("xnumel", meta.get("XBLOCK"))


class SplitScanGrid(GridExpr):
    def generate(self, meta: dict[str, int], is_lazy: bool = False) -> None:
        if not is_lazy:
            if meta.get("XBLOCK", 1) != 1:
                raise AssertionError(
                    f"Expected XBLOCK == 1 for SplitScanGrid, got {meta.get('XBLOCK', 1)}"
                )
        self.x_grid = self.ceildiv("r0_numel", meta.get("R0_BLOCK"))
        self.y_grid = "xnumel"


class FixedGrid(GridExpr):
    @staticmethod
    def setup_grid_as_args() -> dict[str, Any]:
        """Inductor meta so the launcher takes three extra grid arguments"""
        return {
            "grid_type": FixedGrid.__name__,
            "fixed_grid": ["_grid_0", "_grid_1", "_grid_2"],
            "extra_launcher_args": ["_grid_0", "_grid_1", "_grid_2"],
        }

    def generate(self, meta: dict[str, int], is_lazy: bool = False) -> None:
        self.x_grid, self.y_grid, self.z_grid = self.inductor_meta["fixed_grid"]


class PrecomputedGrid(GridExpr):
    def generate(self, meta: dict[str, int], is_lazy: bool = False) -> None:
        for candidate in self.inductor_meta["precomputed_grids"]:
            if all(meta.get(k) == v for k, v in candidate["config"].items()):
                self.x_grid, self.y_grid, self.z_grid = candidate[self.mode]
                return
        raise AssertionError(
            f"Precomputed grid not found for {meta} in {self.inductor_meta['precomputed_grids']}"
        )


class ComboKernelGrid(GridExpr):
    def generate(self, meta: dict[str, int], is_lazy: bool = False) -> None:
        combo_meta = self.inductor_meta["combo_grid_meta"]
        if combo_meta["default_config"]:
            meta = {**combo_meta["default_config"], **meta}
        no_x_dims = []
        xnumels = []
        ynumels = []

        for num in range(combo_meta["num_kernels"]):
            if (
                combo_meta[f"xnumel_{num}"] is not None
                and combo_meta[f"xnumel_{num}"] <= 0
            ):
                raise AssertionError(
                    f"xnumel_{num} must be None or positive, got {combo_meta[f'xnumel_{num}']}"
                )
            no_x_dims.append(combo_meta[f"no_x_dim_{num}"])
            xnumels.append(combo_meta[f"xnumel_{num}"] or f"xnumel_{num}")
            if f"ynumel_{num}" in combo_meta:
                ynumels.append(combo_meta[f"ynumel_{num}"] or f"ynumel_{num}")

        self.x_grid = self.combo_x_grid(xnumels, no_x_dims, meta)
        if combo_meta["min_blocks"]:
            self.x_grid = self.maximum([self.x_grid, combo_meta["min_blocks"]])
        if ynumels:
            self.prefix.extend(
                [
                    self.assign_tmp(
                        "y_grid_raw_",
                        self.ceildiv(self.maximum(ynumels), meta.get("YBLOCK")),
                    ),
                    self.assign_tmp(
                        "y_grid_div_", self.ceildiv("y_grid_raw_", get_max_y_grid())
                    ),
                ]
            )
            ceildiv_expr = self.ceildiv("y_grid_raw_", "y_grid_div_")
            if self.mode == "python":
                self.y_grid = f"(0 if y_grid_div_ == 0 else {ceildiv_expr})"
            else:
                self.y_grid = f"(y_grid_div_ == 0 ? 0 : {ceildiv_expr})"
            self.z_grid = "y_grid_div_"

    def combo_x_grid(
        self,
        xnumels: list[int | str],
        no_x_dims: list[bool],
        meta: dict[str, int],
    ) -> str | int:
        raise NotImplementedError


class SequentialComboKernelGrid(ComboKernelGrid):
    def combo_x_grid(
        self,
        xnumels: list[int | str],
        no_x_dims: list[bool],
        meta: dict[str, int],
    ) -> str | int:
        if len(xnumels) != len(no_x_dims):
            raise AssertionError(
                f"xnumels and no_x_dims length mismatch: {len(xnumels)} != {len(no_x_dims)}"
            )
        return self.summation(
            [
                self.ceildiv(x, 1 if no_x_dim else meta.get("XBLOCK"))
                for x, no_x_dim in zip(xnumels, no_x_dims)
            ]
        )


class SequentialFlattenComboKernelGrid(GridExpr):
    """Flattened grid: (sum of x*y blocks, 1, 1) for per-subkernel with flattened dispatch."""

    def generate_lazy(self, kernel_name: str) -> None:
        combo_meta = self.inductor_meta["combo_grid_meta"]
        num_kernels = combo_meta["num_kernels"]
        meta: dict[str, Any] = {}
        for i in range(num_kernels):
            meta[f"XBLOCK_{i}"] = f"{kernel_name}_result.xblocks[{i}]"
            meta[f"YBLOCK_{i}"] = f"{kernel_name}_result.yblocks[{i}]"
        self.generate(meta, is_lazy=True)

    def generate(self, meta: dict[str, int], is_lazy: bool = False) -> None:
        combo_meta = self.inductor_meta["combo_grid_meta"]
        if combo_meta["default_config"]:
            meta = {**combo_meta["default_config"], **meta}

        total_blocks_list = []
        for num in range(combo_meta["num_kernels"]):
            xnumel = combo_meta[f"xnumel_{num}"]
            if xnumel is not None and xnumel <= 0:
                raise AssertionError(
                    f"xnumel_{num} must be None or positive, got {xnumel}"
                )
            xnumel = xnumel or f"xnumel_{num}"
            x_blocks = self.ceildiv(
                xnumel,
                1 if combo_meta[f"no_x_dim_{num}"] else meta.get(f"XBLOCK_{num}"),
            )
            y_blocks = (
                self.ceildiv(
                    combo_meta[f"ynumel_{num}"] or f"ynumel_{num}",
                    meta.get(f"YBLOCK_{num}"),
                )
                if f"ynumel_{num}" in combo_meta
                else 1
            )
            total_blocks_list.append(self.product([x_blocks, y_blocks]))

        self.x_grid = self.summation(total_blocks_list)
        if combo_meta["min_blocks"]:
            self.x_grid = self.maximum([self.x_grid, combo_meta["min_blocks"]])
        self.y_grid = 1
        self.z_grid = 1


class RoundRobinComboKernelGrid(ComboKernelGrid):
    def combo_x_grid(
        self,
        xnumels: list[int | str],
        no_x_dims: list[bool],
        meta: dict[str, int],
    ) -> str:
        if len(xnumels) != len(no_x_dims):
            raise AssertionError(
                f"xnumels and no_x_dims length mismatch: {len(xnumels)} != {len(no_x_dims)}"
            )
        num_kernels = self.inductor_meta["combo_grid_meta"]["num_kernels"]
        exprs = [x for x, no_x_dim in zip(xnumels, no_x_dims) if no_x_dim]
        xnumels_x_dim = [x for x, no_x_dim in zip(xnumels, no_x_dims) if not no_x_dim]
        if xnumels_x_dim:
            exprs.append(self.ceildiv(self.maximum(xnumels_x_dim), meta.get("XBLOCK")))
        return f"({self.maximum(exprs)}) * {num_kernels}"