from __future__ import annotations

from contextlib import suppress
import contextvars
import hashlib
import linecache
import sys
from typing import Any
from typing import cast

import torch

from .. import _compat as _compat  # ensure Triton compatibility patches run
from .. import exc
from .._utils import triton_is_available
from .config import Config as Config
from .kernel import Kernel as Kernel
from .kernel import kernel as kernel

_CUTLASS_SHUTDOWN_PATCHED = False


def _patch_cutlass_jit_shutdown_unload() -> None:
    """Avoid CUDA library unload hangs during interpreter shutdown.

    On current CUTLASS DSL builds, ``CudaDialectJitModule.__del__`` unconditionally
    calls ``cudaLibraryUnload``. On B200 this can hang during Python finalization
    after a CuTe kernel has already finished executing. Skipping that unload during
    interpreter teardown lets the process exit cleanly while preserving the normal
    unload path during regular runtime GC.
    """

    global _CUTLASS_SHUTDOWN_PATCHED
    if _CUTLASS_SHUTDOWN_PATCHED:
        return

    try:
        import cutlass.cutlass_dsl.cuda_jit_executor as cuda_jit_executor
    except ImportError:
        return

    module_type = cuda_jit_executor.CudaDialectJitModule
    if getattr(module_type, "_helion_shutdown_patch", False):
        _CUTLASS_SHUTDOWN_PATCHED = True
        return

    original_del = cast("Any", module_type.__del__)

    def _helion_del(self: object) -> None:
        module = cast("Any", self)
        if sys.is_finalizing():
            with suppress(Exception):
                module._unloaded = True
            return
        original_del(module)

    module_type.__del__ = _helion_del
    module_type._helion_shutdown_patch = True
    _CUTLASS_SHUTDOWN_PATCHED = True


if triton_is_available():
    import triton

    from .triton_helpers import triton_send_signal as triton_send_signal
    from .triton_helpers import (
        triton_wait_multiple_signal as triton_wait_multiple_signal,
    )
    from .triton_helpers import triton_wait_signal as triton_wait_signal

    def _alloc_fn(size: int, alignment: int, stream: int | None) -> torch.Tensor:
        # Dynamically get device from Triton backend
        current_target = triton.runtime.driver.active.get_current_target()
        if current_target is None:
            raise RuntimeError("No active Triton target available")
        backend = current_target.backend
        return torch.empty(size, device=backend, dtype=torch.int8)

    def set_triton_allocator() -> None:
        try:
            from triton import set_allocator
            from triton.runtime._allocation import NullAllocator
            from triton.runtime._allocation import _allocator
        except ImportError:
            return
        if isinstance(_allocator, contextvars.ContextVar):
            existing = _allocator.get()
        else:  # older versions of Triton
            existing = _allocator
        # if allocator isn't NullAllocator, we assume it is set by the user
        if isinstance(existing, NullAllocator):
            set_allocator(_alloc_fn)
else:

    def set_triton_allocator() -> None:  # type: ignore[misc]
        pass


def get_num_sm(device: torch.device, *, reserved_sms: int = 0) -> int:
    """
    Get the number of streaming multiprocessors (SMs) for the specified device.

    Args:
        device: Device to query.
        reserved_sms: Number of SMs to keep free for other work (e.g., communication
            kernels). Defaults to 0 meaning all device SMs are available to Helion.

    Returns:
        Grid size to use for a persistent kernel on the device after accounting
        for any reserved SMs. Always at least 1.
    """
    available_sms: int
    assert device.type in [
        "cuda",
        "xpu",
        "mtia",
        "mps",
    ], "TODO: implement for other devices"
    if device.type == "cuda":
        available_sms = torch.cuda.get_device_properties(
            device.index
        ).multi_processor_count
    # TODO(EikanWang): gpu_subslice_count is an out-of-date term. we change update it to XeCore number.
    elif device.type == "xpu":
        available_sms = torch.xpu.get_device_properties(device.index).gpu_subslice_count
    elif device.type == "mps":
        available_sms = torch.backends.mps.get_core_count()
    elif device.type == "mtia":
        device_props = torch.mtia.get_device_properties(device.index)
        if "max_grid_height" in device_props and "max_grid_width" in device_props:
            available_sms = (
                device_props["max_grid_height"] * device_props["max_grid_width"]
            )
        else:
            raise RuntimeError(
                f"Unable to determine SM count for MTIA device. "
                f"Available properties: {list(device_props.keys())}"
            )
    else:
        raise NotImplementedError(
            f"get_num_sm not implemented for device type: {device.type}"
        )

    if reserved_sms <= 0:
        return available_sms
    return max(available_sms - reserved_sms, 1)


def default_launcher(
    triton_kernel: object,
    grid: tuple[int, ...],
    *args: object,
    num_warps: int,
    num_stages: int,
    ptx_options: str | None = None,
    launch_cooperative_grid: bool = False,
    **kwargs: dict,
) -> object:
    """Default launcher function that executes the kernel immediately."""
    # For both CUDA and MTIA, use the same kernel execution
    run_kwargs: dict = {
        "grid": grid,
        "warmup": False,
        "num_warps": num_warps,
        "num_stages": num_stages,
        "launch_cooperative_grid": launch_cooperative_grid,
        **kwargs,
    }
    if ptx_options is not None:
        run_kwargs["ptx_options"] = ptx_options
    return triton_kernel.run(  # type: ignore[union-attr]
        *args,
        **run_kwargs,
    )


def _pallas_make_block_spec(
    pl: object,
    jnp: object,
    tensor: torch.Tensor,
    entry: tuple[tuple[int | None, ...], tuple[int | tuple[int, int, int] | None, ...]]
    | None,
) -> object:
    """Build one ``pl.BlockSpec`` from compile-time ``(block_shape, grid_dims)``."""
    if entry is None:
        ndim = tensor.ndim
        full_shape = tuple(tensor.shape)

        def index_map_full(*grid_args: object, _nd: int = ndim) -> tuple[object, ...]:
            # pyrefly: ignore[missing-attribute]
            return tuple(jnp.int32(0) for _ in range(_nd))

        return pl.BlockSpec(full_shape, index_map_full)  # type: ignore[union-attr]

    block_shape_template, grid_dims = entry
    block_shape = tuple(
        min(bs, tensor.shape[d]) if bs is not None else tensor.shape[d]
        for d, bs in enumerate(block_shape_template)
    )

    def _index_for_dim(
        grid_args: tuple[object, ...],
        g: int | tuple[int, int, int] | None,
        jnp: object = jnp,
    ) -> object:
        if g is None:
            return jnp.int32(0)  # pyrefly: ignore[missing-attribute]
        if isinstance(g, tuple):
            # Flat grid decomposition: (grid_dim, stride, num_blocks)
            grid_dim, stride, num_blocks = g
            val = grid_args[grid_dim]
            if stride > 1:
                val = val // stride  # type: ignore[operator]
            val = val % num_blocks  # type: ignore[operator]
            return jnp.int32(val)  # pyrefly: ignore[missing-attribute]
        return jnp.int32(grid_args[g])  # pyrefly: ignore[missing-attribute]

    def index_map(
        *grid_args: object,
        _grid_dims: tuple[int | tuple[int, int, int] | None, ...] = grid_dims,
    ) -> tuple[object, ...]:
        return tuple(_index_for_dim(grid_args, g) for g in _grid_dims)

    return pl.BlockSpec(block_shape, index_map)  # type: ignore[union-attr]


# Per-tensor block spec info: see ``_pallas_make_block_spec``.
# grid_dims entries are int (direct grid dim), tuple (flat decomposition),
# or None (untiled dim).
_BlockSpecInfo = list[
    tuple[tuple[int | None, ...], tuple[int | tuple[int, int, int] | None, ...]] | None
]


def _pallas_build_block_specs(
    pl: object,
    jnp: object,
    grid: tuple[int, ...],
    args: tuple[object, ...],
    tensor_arg_indices: list[int],
    output_indices: list[int],
    block_spec_info: _BlockSpecInfo | None = None,
) -> tuple[list[object] | None, object | None]:
    """Build ``in_specs`` and ``out_specs`` for ``pl.pallas_call``."""
    if block_spec_info is None or len(grid) == 0:
        return None, None

    in_specs = []
    for tensor_pos, idx in enumerate(tensor_arg_indices):
        t = args[idx]
        assert isinstance(t, torch.Tensor)
        in_specs.append(
            _pallas_make_block_spec(pl, jnp, t, block_spec_info[tensor_pos])
        )

    arg_to_tensor_pos = {orig: tpos for tpos, orig in enumerate(tensor_arg_indices)}
    out_specs_list = []
    for idx in output_indices:
        t = args[idx]
        assert isinstance(t, torch.Tensor)
        out_specs_list.append(
            _pallas_make_block_spec(pl, jnp, t, block_spec_info[arg_to_tensor_pos[idx]])
        )

    out_specs = out_specs_list if len(out_specs_list) > 1 else out_specs_list[0]
    return in_specs, out_specs


def _pallas_prepare_args(
    args: tuple[object, ...],
    _output_indices: list[int],
) -> tuple[
    set[int],
    list[int],
    dict[int, object],
    int,
    dict[int, int],
    list[object],
    set[int],
    tuple[object, ...],
]:
    """Extract and organize tensor/non-tensor args for Pallas launchers.

    Returns a tuple of:
    - output_set: set of output arg positions
    - tensor_arg_indices: positions of tensor args
    - non_tensor_args: mapping of non-tensor arg positions to values
    - n_tensor_inputs: count of tensor args
    - arg_to_tensor_pos: mapping from original position to tensor-only position
    - outputs: list of output tensors
    - inplace_positions: positions that are both input and output
    - out_shapes: JAX placeholders for output shapes
    """
    from torch_tpu._internal.pallas.pallas import (  # pyrefly: ignore[missing-import]
        jax_placeholder,
    )

    output_set = set(_output_indices)
    tensor_arg_indices = [
        i for i in range(len(args)) if isinstance(args[i], torch.Tensor)
    ]
    non_tensor_args: dict[int, object] = {
        i: args[i] for i in range(len(args)) if not isinstance(args[i], torch.Tensor)
    }
    n_tensor_inputs = len(tensor_arg_indices)
    arg_to_tensor_pos = {orig: tpos for tpos, orig in enumerate(tensor_arg_indices)}
    outputs = [args[i] for i in _output_indices]
    inplace_positions = output_set & set(tensor_arg_indices)
    out_shapes = tuple(jax_placeholder(out) for out in outputs)  # type: ignore[arg-type]

    return (
        output_set,
        tensor_arg_indices,
        non_tensor_args,
        n_tensor_inputs,
        arg_to_tensor_pos,
        outputs,
        inplace_positions,
        out_shapes,
    )


def _pallas_make_reordered_kernel(
    pallas_kernel: object,
    args: tuple[object, ...],
    tensor_arg_indices: list[int],
    non_tensor_args: dict[int, object],
    n_tensor_inputs: int,
    _output_indices: list[int],
    inplace_positions: set[int],
    arg_to_tensor_pos: dict[int, int],
    n_extra_refs: int = 0,
    skip_inplace_copy: bool = False,
) -> object:
    """Create a wrapper kernel that reorders pallas_call refs to the original arg order.

    ``pallas_call`` provides refs as ``[inputs..., outputs...]``, but Helion
    kernels expect the original parameter order.  When *n_extra_refs* > 0
    (e.g. scratch buffers), those trailing refs are appended after the
    reordered args.

    When *skip_inplace_copy* is True, the initial ``out_ref[...] = in_ref[...]``
    copy for inplace positions is skipped.  This is needed for the pipeline
    launcher where refs are in HBM (``pl.ANY``) and direct load/store is not
    allowed — ``input_output_aliases`` already handles the aliasing.
    """

    def reordered_kernel(*refs: object) -> None:
        n_kernel_params = len(args)
        original_order: list[object] = [None] * n_kernel_params
        for tensor_pos, orig_pos in enumerate(tensor_arg_indices):
            original_order[orig_pos] = refs[tensor_pos]
        for orig_pos, value in non_tensor_args.items():
            original_order[orig_pos] = value
        for out_idx, orig_pos in enumerate(_output_indices):
            out_ref = refs[n_tensor_inputs + out_idx]
            if orig_pos in inplace_positions and not skip_inplace_copy:
                in_ref = refs[arg_to_tensor_pos[orig_pos]]
                out_ref[...] = in_ref[...]  # type: ignore[index]
            original_order[orig_pos] = out_ref
        extra_refs = refs[n_tensor_inputs + len(_output_indices) :]
        pallas_kernel(*original_order, *extra_refs)  # type: ignore[operator]

    return reordered_kernel


def _pallas_build_callable(
    pallas_kernel: object,
    grid: tuple[int, ...],
    jit_fn: object,
    _output_indices: list[int],
    arg_to_tensor_pos: dict[int, int],
    tensor_arg_indices: list[int],
    cache_attr: str,
    trace_key_suffix: str = "",
) -> object:
    """Build a ``JaxCallable``, cache it on the kernel, and return it."""
    import jax
    from torch_tpu._internal.pallas.pallas import (  # pyrefly: ignore[missing-import]
        JaxCallable,
    )

    kernel_name = getattr(pallas_kernel, "__name__", "pallas_kernel")

    call_aliases: dict[int, int] = {}
    for out_idx, orig_pos in enumerate(_output_indices):
        call_aliases[arg_to_tensor_pos[orig_pos]] = out_idx

    jax_callable = JaxCallable(
        name=kernel_name,
        jit_fn=jax.jit(jit_fn),  # pyrefly: ignore[no-matching-overload]
        trace_key=f"{kernel_name}_{id(pallas_kernel)}_{grid}{trace_key_suffix}",
        input_output_aliases=call_aliases,
    )
    setattr(pallas_kernel, cache_attr, (grid, jax_callable, tensor_arg_indices))
    return jax_callable


def default_pallas_launcher(
    pallas_kernel: object,
    grid: tuple[int, ...],
    *args: object,
    _output_indices: list[int] | None = None,
    _block_spec_info: _BlockSpecInfo | None = None,
    **kwargs: object,
) -> None:
    """Default launcher for Pallas kernels on TPU.

    Uses ``JaxCallable`` from ``torch_tpu`` to compile and run the Pallas
    kernel on TPU.  Output tensors are donated via ``input_output_aliases``
    so the kernel writes directly into their buffers (zero-copy).
    """
    if _output_indices is None:
        _output_indices = []

    cache = getattr(pallas_kernel, "_pallas_cache", None)
    if cache is not None and cache[0] == grid:
        _, jax_callable, tensor_arg_indices = cache
    else:
        from jax.experimental import pallas as pl
        import jax.numpy as jnp

        (
            output_set,
            tensor_arg_indices,
            non_tensor_args,
            n_tensor_inputs,
            arg_to_tensor_pos,
            outputs,
            inplace_positions,
            out_shapes,
        ) = _pallas_prepare_args(args, _output_indices)

        in_specs, out_specs = _pallas_build_block_specs(
            pl,
            jnp,
            grid,
            args,
            tensor_arg_indices,
            _output_indices,
            _block_spec_info,
        )

        reordered_kernel = _pallas_make_reordered_kernel(
            pallas_kernel,
            args,
            tensor_arg_indices,
            non_tensor_args,
            n_tensor_inputs,
            _output_indices,
            inplace_positions,
            arg_to_tensor_pos,
        )

        out_shape_arg = out_shapes if len(out_shapes) > 1 else out_shapes[0]

        pallas_aliases = {
            arg_to_tensor_pos[orig_pos]: out_idx
            for out_idx, orig_pos in enumerate(_output_indices)
        }

        pallas_call_kwargs: dict[str, object] = {
            "out_shape": out_shape_arg,
            "input_output_aliases": pallas_aliases,
            "grid": grid,
        }
        if in_specs is not None:
            pallas_call_kwargs["in_specs"] = in_specs
            pallas_call_kwargs["out_specs"] = out_specs

        jit_fn = pl.pallas_call(
            reordered_kernel,  # pyrefly: ignore[bad-argument-type]
            **pallas_call_kwargs,  # type: ignore[arg-type]
        )

        jax_callable = _pallas_build_callable(
            pallas_kernel,
            grid,
            jit_fn,
            _output_indices,
            arg_to_tensor_pos,
            tensor_arg_indices,
            cache_attr="_pallas_cache",
        )

    input_tensors = [
        cast("torch.Tensor", args[i]).contiguous() for i in tensor_arg_indices
    ]
    jax_callable(*input_tensors)  # type: ignore[operator]


def default_pallas_pipeline_launcher(
    pallas_kernel: object,
    grid: tuple[int, ...],
    *args: object,
    _output_indices: list[int] | None = None,
    _block_spec_info: _BlockSpecInfo | None = None,
    _scratch_shapes: list[tuple[tuple[int, ...], str]] | None = None,
    **kwargs: object,
) -> None:
    """Launcher for Pallas kernels using PrefetchScalarGridSpec with scratch memory.

    Used when ``pallas_loop_type='emit_pipeline'``.  Passes all tensors as
    ``memory_space=pl.ANY`` (HBM refs) and adds scratch buffers as
    ``pltpu.VMEM`` shapes.  The kernel uses ``pltpu.emit_pipeline``
    internally for DMA pipelining.
    """
    if _output_indices is None:
        _output_indices = []
    if _scratch_shapes is None:
        _scratch_shapes = []

    cache = getattr(pallas_kernel, "_pallas_pipeline_cache", None)
    if cache is not None and cache[0] == grid:
        _, jax_callable, tensor_arg_indices = cache
    else:
        from jax.experimental import pallas as pl
        from jax.experimental.pallas import tpu as pltpu
        import jax.numpy as jnp

        (
            output_set,
            tensor_arg_indices,
            non_tensor_args,
            n_tensor_inputs,
            arg_to_tensor_pos,
            outputs,
            inplace_positions,
            out_shapes,
        ) = _pallas_prepare_args(args, _output_indices)

        # Build scratch shapes for VMEM
        _jnp_dtype_map: dict[str, object] = {
            "jnp.float32": jnp.float32,
            "jnp.float16": jnp.float16,
            "jnp.bfloat16": jnp.bfloat16,
            "jnp.int32": jnp.int32,
            "jnp.int16": jnp.int16,
            "jnp.int8": jnp.int8,
            "jnp.uint8": jnp.uint8,
            "jnp.bool_": jnp.bool_,
        }
        scratch_shapes = []
        for scratch_entry in _scratch_shapes:
            if len(scratch_entry) == 3:
                shape, dtype_str, scratch_type = scratch_entry
            else:
                shape, dtype_str = scratch_entry  # type: ignore[misc]
                scratch_type = "vmem"
            if scratch_type == "dma_semaphore":
                scratch_shapes.append(pltpu.SemaphoreType.DMA(()))
            else:
                jnp_dtype = _jnp_dtype_map.get(dtype_str, jnp.float32)
                scratch_shapes.append(
                    pltpu.VMEM(shape, jnp_dtype)  # pyrefly: ignore[bad-argument-type]
                )

        # Build in_specs/out_specs with memory_space=pl.ANY (HBM refs)
        in_specs_list = [pl.BlockSpec(memory_space=pl.ANY) for _ in tensor_arg_indices]
        out_specs_list = [pl.BlockSpec(memory_space=pl.ANY) for _ in _output_indices]
        out_specs = out_specs_list if len(out_specs_list) > 1 else out_specs_list[0]

        reordered_kernel = _pallas_make_reordered_kernel(
            pallas_kernel,
            args,
            tensor_arg_indices,
            non_tensor_args,
            n_tensor_inputs,
            _output_indices,
            inplace_positions,
            arg_to_tensor_pos,
            n_extra_refs=len(scratch_shapes),
            skip_inplace_copy=True,
        )

        out_shape_arg = out_shapes if len(out_shapes) > 1 else out_shapes[0]

        pallas_aliases = {
            arg_to_tensor_pos[orig_pos]: out_idx
            for out_idx, orig_pos in enumerate(_output_indices)
        }

        grid_spec = pltpu.PrefetchScalarGridSpec(
            num_scalar_prefetch=0,
            in_specs=in_specs_list,
            out_specs=out_specs,
            scratch_shapes=scratch_shapes,
            grid=grid,
        )

        jit_fn = pl.pallas_call(
            reordered_kernel,  # pyrefly: ignore[bad-argument-type]
            out_shape=out_shape_arg,
            input_output_aliases=pallas_aliases,
            grid_spec=grid_spec,
            compiler_params=pltpu.CompilerParams(  # pyrefly: ignore[bad-instantiation]
                dimension_semantics=tuple("parallel" for _ in grid),
            ),
        )

        jax_callable = _pallas_build_callable(
            pallas_kernel,
            grid,
            jit_fn,
            _output_indices,
            arg_to_tensor_pos,
            tensor_arg_indices,
            cache_attr="_pallas_pipeline_cache",
            trace_key_suffix="_pipeline",
        )

    input_tensors = [
        cast("torch.Tensor", args[i]).contiguous() for i in tensor_arg_indices
    ]
    jax_callable(*input_tensors)  # type: ignore[operator]


def default_pallas_fori_launcher(
    pallas_kernel: object,
    grid: tuple[int, ...],
    *args: object,
    _output_indices: list[int] | None = None,
    _block_spec_info: _BlockSpecInfo | None = None,
    _scratch_shapes: list[tuple[tuple[int, ...], str | None, str]] | None = None,
    **kwargs: object,
) -> None:
    """Launcher for Pallas kernels using fori_loop with manual DMA.

    Used when ``pallas_loop_type="fori_loop"``.  Passes all tensors as
    ``memory_space=pl.ANY`` (HBM refs) and adds scratch buffers as
    ``pltpu.VMEM`` shapes plus ``pltpu.SemaphoreType.DMA`` for async copies.
    The kernel uses ``jax.lax.fori_loop`` with ``pltpu.make_async_copy``
    internally for DMA control.
    """
    if _output_indices is None:
        _output_indices = []
    if _scratch_shapes is None:
        _scratch_shapes = []

    cache = getattr(pallas_kernel, "_pallas_fori_cache", None)
    if cache is not None and cache[0] == grid:
        _, jax_callable, tensor_arg_indices = cache
    else:
        from jax.experimental import pallas as pl
        from jax.experimental.pallas import tpu as pltpu
        import jax.numpy as jnp

        (
            output_set,
            tensor_arg_indices,
            non_tensor_args,
            n_tensor_inputs,
            arg_to_tensor_pos,
            outputs,
            inplace_positions,
            out_shapes,
        ) = _pallas_prepare_args(args, _output_indices)

        # Build scratch shapes: VMEM buffers + DMA semaphores
        _jnp_dtype_map: dict[str, object] = {
            "jnp.float32": jnp.float32,
            "jnp.float16": jnp.float16,
            "jnp.bfloat16": jnp.bfloat16,
            "jnp.int32": jnp.int32,
            "jnp.int16": jnp.int16,
            "jnp.int8": jnp.int8,
            "jnp.uint8": jnp.uint8,
            "jnp.bool_": jnp.bool_,
        }
        scratch_shapes = []
        for shape, dtype_str, scratch_type in _scratch_shapes:
            if scratch_type == "dma_semaphore":
                scratch_shapes.append(pltpu.SemaphoreType.DMA(()))
            else:  # "vmem"
                assert dtype_str is not None
                jnp_dtype = _jnp_dtype_map.get(dtype_str, jnp.float32)
                scratch_shapes.append(
                    pltpu.VMEM(shape, jnp_dtype)  # pyrefly: ignore[bad-argument-type]
                )

        # Build in_specs/out_specs with memory_space=pl.ANY (HBM refs)
        in_specs_list = [pl.BlockSpec(memory_space=pl.ANY) for _ in tensor_arg_indices]
        out_specs_list = [pl.BlockSpec(memory_space=pl.ANY) for _ in _output_indices]
        out_specs = out_specs_list if len(out_specs_list) > 1 else out_specs_list[0]

        reordered_kernel = _pallas_make_reordered_kernel(
            pallas_kernel,
            args,
            tensor_arg_indices,
            non_tensor_args,
            n_tensor_inputs,
            _output_indices,
            inplace_positions,
            arg_to_tensor_pos,
            n_extra_refs=len(scratch_shapes),
            skip_inplace_copy=True,
        )

        out_shape_arg = out_shapes if len(out_shapes) > 1 else out_shapes[0]

        pallas_aliases = {
            arg_to_tensor_pos[orig_pos]: out_idx
            for out_idx, orig_pos in enumerate(_output_indices)
        }

        grid_spec = pltpu.PrefetchScalarGridSpec(
            num_scalar_prefetch=0,
            in_specs=in_specs_list,
            out_specs=out_specs,
            scratch_shapes=scratch_shapes,
            grid=grid,
        )

        jit_fn = pl.pallas_call(
            reordered_kernel,  # pyrefly: ignore[bad-argument-type]
            out_shape=out_shape_arg,
            input_output_aliases=pallas_aliases,
            grid_spec=grid_spec,
            compiler_params=pltpu.CompilerParams(  # pyrefly: ignore[bad-instantiation]
                dimension_semantics=tuple("parallel" for _ in grid),
            ),
        )

        jax_callable = _pallas_build_callable(
            pallas_kernel,
            grid,
            jit_fn,
            _output_indices,
            arg_to_tensor_pos,
            tensor_arg_indices,
            cache_attr="_pallas_fori_cache",
            trace_key_suffix="_fori",
        )

    input_tensors = [
        cast("torch.Tensor", args[i]).contiguous() for i in tensor_arg_indices
    ]
    jax_callable(*input_tensors)  # type: ignore[operator]


def _torch_dtype_to_cutlass(dtype: torch.dtype) -> object:
    _patch_cutlass_jit_shutdown_unload()
    import cutlass

    mapping: dict[torch.dtype, object] = {
        torch.float16: cutlass.Float16,
        torch.float32: cutlass.Float32,
        torch.float64: cutlass.Float64,
        torch.bfloat16: cutlass.BFloat16,
        torch.int8: cutlass.Int8,
        torch.int16: cutlass.Int16,
        torch.int32: cutlass.Int32,
        torch.int64: cutlass.Int64,
        torch.uint8: cutlass.Uint8,
    }
    if dtype not in mapping:
        raise exc.BackendUnsupported("cute", f"dtype: {dtype}")
    return mapping[dtype]


def _normalize_cute_scalar(arg: object) -> tuple[str, object]:
    if isinstance(arg, (bool, torch.SymBool)):
        return ("bool", bool(arg))
    if isinstance(arg, (int, torch.SymInt)):
        return ("int", int(arg))
    if isinstance(arg, (float, torch.SymFloat)):
        return ("float", float(arg))
    raise exc.BackendUnsupported("cute", f"launcher scalar argument type: {type(arg)}")


def _cute_scalar_annotation(kind: str) -> str:
    mapping = {
        "bool": "cutlass.Boolean",
        "int": "cutlass.Int64",
        "float": "cutlass.Float64",
    }
    return mapping[kind]


def _create_cute_wrapper(
    cute_kernel: object,
    schema_key: tuple[tuple[object, ...], ...],
) -> object:
    _patch_cutlass_jit_shutdown_unload()
    import cutlass
    import cutlass.cute as cute

    kernel_name = getattr(cast("Any", cute_kernel), "__name__", "cute_kernel")
    kernel_tag = f"{kernel_name}_{id(cute_kernel):x}"
    func_name = f"_helion_cute_launch_{kernel_tag}"
    params: list[str] = []
    body: list[str] = []
    call_args: list[str] = []

    for i, entry in enumerate(schema_key):
        kind = entry[0]
        if kind == "tensor":
            (_, _dtype, rank) = entry
            assert isinstance(rank, int)
            ptr_name = f"arg{i}_ptr"
            params.append(f"{ptr_name}: cute.Pointer")
            shape_names = [f"arg{i}_shape{d}" for d in range(rank)]
            stride_names = [f"arg{i}_stride{d}" for d in range(rank)]
            params.extend(f"{name}: cutlass.Int64" for name in shape_names)
            params.extend(f"{name}: cutlass.Int64" for name in stride_names)
            shape_tuple = (
                f"({shape_names[0]},)" if rank == 1 else f"({', '.join(shape_names)})"
            )
            stride_tuple = (
                f"({stride_names[0]},)" if rank == 1 else f"({', '.join(stride_names)})"
            )
            body.append(
                f"    arg{i} = cute.make_tensor({ptr_name}, layout=cute.make_layout({shape_tuple}, stride={stride_tuple}))"
            )
            call_args.append(f"arg{i}")
            continue

        assert kind == "scalar"
        (_, scalar_kind) = entry
        assert isinstance(scalar_kind, str)
        scalar_name = f"arg{i}"
        params.append(f"{scalar_name}: {_cute_scalar_annotation(scalar_kind)}")
        call_args.append(scalar_name)

    params.extend(
        (
            "grid_x: cutlass.Int32",
            "grid_y: cutlass.Int32",
            "grid_z: cutlass.Int32",
            "block_x: cutlass.Int32",
            "block_y: cutlass.Int32",
            "block_z: cutlass.Int32",
        )
    )
    body.extend(
        (
            f"    _helion_cute_kernel_tag = {kernel_tag!r}",
            "    _kernel("
            + ", ".join(call_args)
            + ").launch(grid=(grid_x, grid_y, grid_z), block=(block_x, block_y, block_z))",
        )
    )

    source = "\n".join(
        [
            "@cute.jit",
            f"def {func_name}({', '.join(params)}) -> None:",
            *body,
        ]
    )

    namespace: dict[str, Any] = {
        "cutlass": cutlass,
        "cute": cute,
        "_kernel": cute_kernel,
    }
    filename = f"<helion_cute_launcher:{kernel_tag}:{schema_key!r}>"
    linecache.cache[filename] = (
        len(source),
        None,
        [line + "\n" for line in source.splitlines()],
        filename,
    )
    exec(compile(source, filename, "exec"), namespace)
    return namespace[func_name]


def _get_compiled_cute_launcher(
    cute_kernel: object,
    schema_key: tuple[tuple[object, ...], ...],
    launch_args: tuple[object, ...],
) -> object:
    try:
        # pyrefly: ignore [missing-attribute]
        cache = cute_kernel._helion_cute_compiled_launchers
    except AttributeError:
        cache = {}
        # pyrefly: ignore [missing-attribute]
        cute_kernel._helion_cute_compiled_launchers = cache
    cached = cache.get(schema_key)
    if cached is not None:
        return cached

    wrapper = _create_cute_wrapper(cute_kernel, schema_key)
    cache[schema_key] = wrapper
    return wrapper


def _build_cute_schema_and_args(
    args: tuple[object, ...],
    grid: tuple[int, int, int],
    block: tuple[int, int, int],
) -> tuple[tuple[tuple[object, ...], ...], tuple[object, ...]]:
    _patch_cutlass_jit_shutdown_unload()
    import cutlass.cute as cute
    from cutlass.cute.runtime import make_ptr

    schema: list[tuple[object, ...]] = []
    launch_args: list[object] = []
    for arg in args:
        if isinstance(arg, torch.Tensor):
            if arg.device.type != "cuda":
                raise exc.BackendUnsupported("cute", "launcher requires CUDA tensors")
            if arg.ndim <= 0:
                raise exc.BackendUnsupported(
                    "cute", "launcher requires tensor rank >= 1"
                )
            schema.append(("tensor", str(arg.dtype), arg.ndim))
            launch_args.append(
                make_ptr(
                    cast("Any", _torch_dtype_to_cutlass(arg.dtype)),
                    arg.data_ptr(),
                    cute.AddressSpace.gmem,
                    assumed_align=16,
                )
            )
            launch_args.extend(int(arg.size(d)) for d in range(arg.ndim))
            launch_args.extend(int(arg.stride(d)) for d in range(arg.ndim))
            continue

        scalar_kind, scalar_value = _normalize_cute_scalar(arg)
        schema.append(("scalar", scalar_kind))
        launch_args.append(scalar_value)

    launch_args.extend((*grid, *block))
    return tuple(schema), tuple(launch_args)


def default_cute_launcher(
    cute_kernel: object,
    grid: tuple[int, ...],
    *args: object,
    **kwargs: object,
) -> object:
    block = kwargs.pop("block", (256, 1, 1))
    if not isinstance(block, tuple) or len(block) < 1:
        raise ValueError(f"Invalid block specification: {block}")
    if not isinstance(grid, tuple) or len(grid) < 1:
        raise ValueError(f"Invalid grid specification: {grid}")
    if kwargs:
        raise exc.BackendUnsupported("cute", f"launcher kwargs: {sorted(kwargs)}")

    grid_xyz = (
        int(grid[0]),
        int(grid[1]) if len(grid) > 1 else 1,
        int(grid[2]) if len(grid) > 2 else 1,
    )
    block_xyz = (
        int(block[0]),
        int(block[1]) if len(block) > 1 else 1,
        int(block[2]) if len(block) > 2 else 1,
    )

    schema_key, launch_args = _build_cute_schema_and_args(
        tuple(args), grid_xyz, block_xyz
    )
    compiled = _get_compiled_cute_launcher(cute_kernel, schema_key, launch_args)
    return cast("Any", compiled)(*launch_args)


def default_metal_launcher(
    metal_kernel: object,
    grid: tuple[int, ...],
    *args: object,
    _block_size: int = 256,
    **kwargs: object,
) -> None:
    """Default launcher for Metal kernels on Apple MPS devices.

    Compiles MSL source via ``torch.mps.compile_shader()`` and dispatches
    using the compiled library.  Caches the compiled library on the kernel
    object to avoid recompilation on subsequent calls.

    Only 1D grids are currently supported.
    """
    kwargs.pop("num_warps", None)
    kwargs.pop("num_stages", None)
    if kwargs:
        raise exc.BackendUnsupported(
            "metal", f"unexpected launcher kwargs: {sorted(kwargs)}"
        )

    assert len(grid) == 1, (
        f"Metal launcher only supports 1D grids, got {len(grid)}D: {grid}"
    )

    msl_source, kernel_name = metal_kernel()  # type: ignore[operator]
    source_hash = hashlib.sha256(msl_source.encode()).digest()
    cache = getattr(metal_kernel, "_metal_cache", None)
    if cache is not None and cache[0] == source_hash:
        lib = cache[1]
    else:
        lib = torch.mps.compile_shader(msl_source)  # type: ignore[attr-defined]
        metal_kernel._metal_cache = (source_hash, lib)  # type: ignore[attr-defined]

    tensor_args = [a for a in args if isinstance(a, torch.Tensor)]
    dispatch_fn = getattr(lib, kernel_name)
    total_threads = grid[0] * _block_size
    dispatch_fn(*tensor_args, threads=total_threads, group_size=_block_size)
