# Copyright 2024 The JAX Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

from collections.abc import Callable, Iterable, Sequence
import contextlib
import ctypes
import dataclasses
import enum
import functools
import hashlib
import io
import itertools
import json
import math
import os
import pathlib
import time
from typing import Any
import weakref

import jax
from jax._src import core as jax_core
from jax._src import dtypes
from jax._src import lib
from jax._src import mesh as mesh_lib
from jax._src import sharding_impls
from jax._src import util as jax_util
from jax._src.interpreters import mlir
from jax._src.lib import mosaic_gpu_dialect as dialect
from jax._src.pallas.mosaic import error_handling as error
from jaxlib.mlir import ir
from jaxlib.mlir import passmanager
from jaxlib.mlir.dialects import _gpu_ops_gen
from jaxlib.mlir.dialects import arith
from jaxlib.mlir.dialects import builtin
from jaxlib.mlir.dialects import func
from jaxlib.mlir.dialects import gpu
from jaxlib.mlir.dialects import llvm
from jaxlib.mlir.dialects import memref
import numpy as np

from . import dialect_lowering
from . import launch_context
from . import layout_inference
from . import layouts
from . import profiler
from . import tcgen05
from . import utils

# Point Mosaic GPU tools (like nvdisasm) to the CUDA path
cuda_root = lib.cuda_path or "/usr/local/cuda"
os.environ["MOSAIC_GPU_CUDA_ROOT"] = cuda_root
PYTHON_RUNFILES = os.environ.get("PYTHON_RUNFILES")
BAZEL_TEST = os.environ.get("BAZEL_TEST", "0")

_SMEM_SIZE_BOUND = None  # For test purposes.

@contextlib.contextmanager
def artificial_shared_memory_limit(limit):
    global _SMEM_SIZE_BOUND
    old_limit = _SMEM_SIZE_BOUND
    _SMEM_SIZE_BOUND = limit
    try:
        yield
    finally:
        _SMEM_SIZE_BOUND = old_limit

# This tracks the latest Mosaic GPU IR version with a monthly delay.
FWD_COMPAT_IR_VERSION = 9

# jaxlib 0.12.0 changes the Mosaic GPU C API (MosaicGpuLaunch takes kernel and
# device_ordinal arguments).
# TODO(sohaibiftikhar): Remove conditioning once the minimum jaxlib version is
# >= 0.12.0.
_ABI_V2 = lib.version >= (0, 12, 0)

c = utils.c  # This is too common to fully qualify.


RUNTIME_PATH = None
try:
  from jax._src.lib import mosaic_gpu as mosaic_gpu_lib

  RUNTIME_PATH = (
      pathlib.Path(mosaic_gpu_lib._mosaic_gpu_ext.__file__).parent
      / "libmosaic_gpu_runtime.so"
  )
except ImportError:
  pass

if RUNTIME_PATH and RUNTIME_PATH.exists():
  # Set this so that the custom call can find it
  os.environ["MOSAIC_GPU_RUNTIME_LIB_PATH"] = str(RUNTIME_PATH)


try:
  from nvidia import nvshmem  # pyrefly: ignore[missing-import]
except ImportError:
  # Try to find the nvshmem library in Bazel test runfiles.
  if BAZEL_TEST == "1" and PYTHON_RUNFILES:
    libdevice_path = os.path.join(
        PYTHON_RUNFILES, "nvidia_nvshmem", "lib", "libnvshmem_device.bc"
    )
    if os.path.exists(libdevice_path):
      os.environ["MOSAIC_GPU_NVSHMEM_BC_PATH"] = libdevice_path
    for solib_path in ["_solib_x86_64", "_solib_aarch64"]:
      if not os.path.exists(
          os.path.join(PYTHON_RUNFILES, "__main__", solib_path)
      ):
        continue
      for root, _, files in os.walk(
          os.path.join(PYTHON_RUNFILES, "__main__", solib_path)
      ):
        if "libnvshmem_host.so.3" in files:
          os.environ["MOSAIC_GPU_NVSHMEM_SO_PATH"] = os.path.join(
              root, "libnvshmem_host.so.3"
          )
          break
      if "MOSAIC_GPU_NVSHMEM_SO_PATH" in os.environ:
        break
  else:
    pass
else:
  if os.environ.get("MOSAIC_GPU_NVSHMEM_BC_PATH") is None:
    os.environ["MOSAIC_GPU_NVSHMEM_BC_PATH"] = os.path.join(
        nvshmem.__path__[0], "lib/libnvshmem_device.bc"
    )
  if os.environ.get("MOSAIC_GPU_NVSHMEM_SO_PATH") is None:
    os.environ["MOSAIC_GPU_NVSHMEM_SO_PATH"] = os.path.join(
        nvshmem.__path__[0], "lib/libnvshmem_host.so.3"
    )


def is_nvshmem_available():
  try:
    nvshmem_bc_path = os.environ["MOSAIC_GPU_NVSHMEM_BC_PATH"]
  except KeyError:
    return False
  if nvshmem_so_path := os.environ.get("MOSAIC_GPU_NVSHMEM_SO_PATH", ""):
    try:
      # This both ensures that the file exists, and it populates the dlopen
      # cache, helping XLA find the library even if the RPATH is not right...
      ctypes.CDLL(nvshmem_so_path)
    except OSError:
      return False
  xla_flags = os.environ.get("XLA_FLAGS", "")
  return (
      os.path.exists(nvshmem_bc_path)
      and "--xla_gpu_experimental_enable_nvshmem" in xla_flags
  )


def is_single_process_multi_device_topology():
  return (jax.device_count() > 1
          and jax.device_count() == jax.local_device_count())


mosaic_gpu_p = jax_core.Primitive("mosaic_gpu_p")
mosaic_gpu_p.multiple_results = True


@mosaic_gpu_p.def_abstract_eval
def _mosaic_gpu_abstract_eval(*_, module, out_types, inout_types):
  del module # Unused.
  return [
      jax_core.ShapedArray(t.shape, t.dtype)
      for t in itertools.chain(out_types, inout_types)
  ]


def _has_communication(module, **_):
  if launch_context.uses_collective_metadata(module):
    return True
  empty_str_attr = ir.StringAttr.get("")
  for op in module.body:
    if "nvshmem" in getattr(op, "sym_name", empty_str_attr).value:
      return True
  return False


# TODO(apaszke): Implement a proper system for managing kernel lifetimes
# Maps kernel ID to the compiled kernel ASM.
KNOWN_KERNELS: dict[bytes, bytes] = {}


# Frontend attributes understood by the XLA:GPU compiler. They let a custom
# call request a specific memory space for its operands and results. See
# xla/service/gpu/gpu_memory_space_assignment.h for the definitions.
_OPERANDS_MEMORY_SPACES_ATTR = "operands_memory_spaces"
_RESULTS_MEMORY_SPACES_ATTR = "results_memory_spaces"

# Must be kept in sync with xla::gpu::MemorySpaceColor::kCollective. Buffers
# in this memory space are allocated as symmetric memory, which is a
# prerequisite for peer-to-peer and multimem accesses.
_COLLECTIVE_MEMORY_SPACE = 1


def _memory_spaces_attr_value(indices: Sequence[int]) -> str:
  """Formats ``indices`` as an XLA ``{index:memory_space,...}`` attribute."""
  return "{%s}" % ",".join(
      f"{index}:{_COLLECTIVE_MEMORY_SPACE}" for index in indices
  )


def _mosaic_gpu_lowering_rule(
    ctx,
    *args,
    module,
    out_types,
    inout_types,
    input_output_aliases: tuple[tuple[int, int], ...] = (),
    cost_estimate: Any = None,
    use_custom_barrier: bool = False,
    skip_device_barrier: bool = False,
):
  axis_context = ctx.module_context.axis_context
  replica_ids = []
  if is_multi_device_module := _has_communication(module):
    # Those checks are trying to ensure that the logical device ids are
    # consistent with the NVSHMEM PE ids that Mosaic will be using for
    # communication. Any divergence here would require us to implement a logical
    # to physical translation, which is currently not implemented.
    if isinstance(axis_context, sharding_impls.SPMDAxisContext):
      mesh = axis_context.mesh
      if isinstance(mesh, mesh_lib.Mesh):
        replica_ids = mesh.device_ids.ravel()
        if not np.array_equal(mesh.device_ids.ravel(), np.arange(mesh.size)):
          raise NotImplementedError(
              "Mosaic GPU only supports meshes with device ordering that follows"
              f" row-major device ids. Got: {mesh.device_ids.ravel()} device ids."
          )
      elif isinstance(mesh, mesh_lib.AbstractMesh):
        replica_ids = np.arange(mesh.size)
    elif isinstance(axis_context, sharding_impls.ShardingContext):
      if axis_context.num_devices != 1:
        raise NotImplementedError(
            "Mosaic GPU only supports single-device meshes in ShardingContext."
            f" Got: {axis_context.num_devices} devices."
        )
    else:
      raise NotImplementedError(f"Unsupported sharding context: {axis_context}")

  if inout_types:
    if input_output_aliases:
      raise ValueError(
          "input_output_aliases and inout_types are mutually exclusive"
      )
    num_inputs = len(ctx.avals_in)
    num_outputs = len(ctx.avals_out)
    input_output_aliases = tuple(
        (num_inputs - 1 - i, num_outputs - 1 - i)
        for i in range(len(inout_types))
    )
  assert len(ctx.avals_in) == len(args)
  assert len(ctx.avals_out) == len(out_types) + len(inout_types)
  module = _run_serde_pass(
      module,
      serialize=True,
      ir_version=FWD_COMPAT_IR_VERSION if ctx.is_forward_compat() else None,
  )

  # Prior to computing the kernel hash, we temporarily remove the dump basename
  # attribute from the module. This ensures that modules with different dump
  # basenames but otherwise identical will have the same kernel hash, and thus
  # will share the same compiled kernel.
  dump_basename_attr_name = "mosaic_gpu.dump_basename"
  if dump_basename_attr_name in module.operation.attributes:
    dump_basename_attr = module.operation.attributes[dump_basename_attr_name]
  else:
    dump_basename_attr = None
  module.operation.attributes[dump_basename_attr_name] = ir.UnitAttr.get()

  bytecode_buffer = io.BytesIO()
  module.operation.write_bytecode(bytecode_buffer, desired_version=0)
  module_asm = bytecode_buffer.getvalue()
  kernel_id = hashlib.sha256(module_asm).digest()
  # Note that this is technically only a half measure. Someone might load a
  # compiled module with a hash collision from disk. But that's so unlikely with
  # SHA256 that it shouldn't be a problem.
  if (kernel_text := KNOWN_KERNELS.get(kernel_id, None)) is not None:
    if kernel_text != module_asm:
      raise RuntimeError("Kernel hash collision!")
  else:
    KNOWN_KERNELS[kernel_id] = module_asm

  if dump_basename_attr is not None:
    module.operation.attributes[dump_basename_attr_name] = dump_basename_attr

  backend_config: dict[str, ir.Attribute] = dict(
      kernel_hash=ir.StringAttr.get(kernel_id),
      module=ir.StringAttr.get(module_asm),
      use_custom_barrier=ir.BoolAttr.get(use_custom_barrier),
      skip_device_barrier=ir.BoolAttr.get(skip_device_barrier),
      uses_xla_collective_metadata=ir.BoolAttr.get(
          launch_context.uses_collective_metadata(module)
      ),
  )
  if cost_estimate is not None:
    backend_config["cost_estimate_json"] = ir.StringAttr.get(
        json.dumps(dataclasses.asdict(cost_estimate))
    )

  frontend_attributes: dict[str, ir.Attribute] = {}

  # If NVSHMEM is available it will be used by default, otherwise we will use
  # collective metadata.
  if is_multi_device_module and (
      is_single_process_multi_device_topology() or not is_nvshmem_available()
  ):
    backend_config["xla_replica_ids"] = ir.StringAttr.get(
        ",".join(map(str, replica_ids))
    )

    if launch_context.SYMMETRIC_MEMORY_ARGS_ATTR in module.operation.attributes:
      symmetric_memory_args = np.array(
          ir.DenseIntElementsAttr(
              module.operation.attributes[
                  launch_context.SYMMETRIC_MEMORY_ARGS_ATTR
              ]
          ),
          dtype=bool,
      )
      # The Mosaic runtime uses this to decide which buffers to exchange with
      # the peers through the collective metadata.
      backend_config["symmetric_memory_parameters"] = ir.StringAttr.get(
          ",".join(map(str, map(int, symmetric_memory_args)))
      )
      # Mosaic kernel parameters are the flat list of the custom call buffers:
      # all the operands first, followed by all the results. Ask the XLA
      # compiler to allocate the symmetric ones in the collective memory space.
      num_operands = len(ctx.avals_in)
      operand_indices = [
          i for i, s in enumerate(symmetric_memory_args[:num_operands]) if s
      ]
      result_indices = [
          i for i, s in enumerate(symmetric_memory_args[num_operands:]) if s
      ]
      if operand_indices:
        frontend_attributes[_OPERANDS_MEMORY_SPACES_ATTR] = ir.StringAttr.get(
            _memory_spaces_attr_value(operand_indices)
        )
      if result_indices:
        frontend_attributes[_RESULTS_MEMORY_SPACES_ATTR] = ir.StringAttr.get(
            _memory_spaces_attr_value(result_indices)
        )

  extra_attributes: dict[str, ir.Attribute] | None = None
  if frontend_attributes:
    extra_attributes = {
        "mhlo.frontend_attributes": ir.DictAttr.get(frontend_attributes)
    }

  result_types, _ = mlir.ir_tree_registry.flatten([
      mlir.aval_to_ir_type(ctx.module_context, aval) for aval in ctx.avals_out
  ])
  return mlir.custom_call(
      call_target_name="mosaic_gpu_v2",
      result_types=result_types,
      operands=args,
      operand_layouts=[list(reversed(range(a.ndim))) for a in ctx.avals_in],
      result_layouts=[list(reversed(range(a.ndim))) for a in ctx.avals_out],
      backend_config=backend_config,
      operand_output_aliases=dict(input_output_aliases),
      api_version=4,
      extra_attributes=extra_attributes,
  ).results

mlir.register_lowering(mosaic_gpu_p, _mosaic_gpu_lowering_rule, "cuda")


# ShapeTrees currently can not contain unions.
ShapeTree = Any
RefTree = Any


@dataclasses.dataclass(frozen=True)
class Union[T]:
  members: Sequence[T]

  def __iter__(self):
    return iter(self.members)

@dataclasses.dataclass(frozen=True)
class Barrier:
  arrival_count: int
  num_barriers: int = 1
  orders_tensor_core: bool = False

  def __post_init__(self):
    if self.arrival_count < 1:
      raise ValueError(
          f"Arrival count must be at least 1, but got {self.arrival_count}"
      )

@dataclasses.dataclass(frozen=True)
class ClusterBarrier:
  collective_dims: Sequence[gpu.Dimension | Sequence[gpu.Dimension]]
  arrival_count: int = 1
  num_barriers: int = 1
  _: dataclasses.KW_ONLY
  orders_tensor_core: bool = False
  leader_tracked: bool = False


@dataclasses.dataclass(frozen=True)
class CtaBarrier:
  arrival_count: int
  num_barriers: int = 1

  def __post_init__(self):
    if (self.arrival_count < 1) or (self.arrival_count % utils.WARP_SIZE != 0):
      raise ValueError(
          f"Arrival count must be a multiple of {utils.WARP_SIZE}, but got"
          f" {self.arrival_count}"
      )
    if self.num_barriers < 1:
      raise ValueError(
          f"Num barriers must be at least 1, but got {self.num_barriers}"
      )


@dataclasses.dataclass(frozen=True)
class TMEM:
  shape: tuple[int, int]
  dtype: Any
  _: dataclasses.KW_ONLY
  layout: tcgen05.TMEMLayout | None = None
  collective: bool = False
  packing: int | None = None

  def __post_init__(self):
    if self.layout is not None:
      self.layout.check_type(
          self.shape, utils.bitwidth(utils.dtype_to_ir_type(self.dtype))
      )
      if self.packing is not None:
        raise ValueError("Cannot specify both layout and packing")


def _count_buffer_bytes(shape_dtype: jax.ShapeDtypeStruct) -> int:
  return math.prod(shape_dtype.shape) * dtypes.itemsize_bits(dtypes.dtype(shape_dtype.dtype)) // 8


class LoweringSemantics(enum.Enum):
  """Semantics for the kernel's instruction stream."""

  Lane = enum.auto()
  Warpgroup = enum.auto()


@dataclasses.dataclass(frozen=True)
class _TMEMAlloc:
  addr_ref: ir.Value
  num_cols: int
  collective: bool

  def alloc(self) -> int:
    """Allocates TMEM and returns the number of columns allocated."""
    return tcgen05.tmem_alloc(
        self.addr_ref, self.num_cols, collective=self.collective, exact=False
    )

  def dealloc(self):
    addr = memref.load(self.addr_ref, [])
    tcgen05.tmem_dealloc(
        addr, self.num_cols, collective=self.collective, exact=False
    )


@dataclasses.dataclass()
class _TMEMDialectAlloc:
  addr_ref: ir.Value
  shape: tuple[int, int]
  dtype: ir.Type
  packing: int
  collective: bool
  tmem_ref: ir.Value | None = dataclasses.field(init=False, default=None)

  def alloc(self) -> int:
    """Allocates TMEM and returns the number of columns allocated."""
    result_type = ir.MemRefType.get(
        self.shape,
        self.dtype,
        memory_space=utils.tmem(),
    )
    self.tmem_ref = dialect.tmem_alloc(
        result_type,
        self.addr_ref,
        collective=self.collective,
        packing=self.packing,
    )
    ncols = self.shape[1] // self.packing
    return tcgen05.tmem_alloc_exact_ncols(ncols, exact=False)

  def dealloc(self):
    assert self.tmem_ref is not None
    dialect.tmem_dealloc(self.tmem_ref)


def _slice_smem(
    result: ir.Type,
    smem_base: ir.Value,
    offset: int,
    lowering_semantics: LoweringSemantics,
) -> ir.Value:
  if lowering_semantics == LoweringSemantics.Warpgroup:
    return dialect.slice_smem(result, offset)
  else:
    ir_offset = arith.constant(ir.IndexType.get(), offset)
    return memref.view(result, smem_base, ir_offset, [])


def _construct_smem_reftree(
    cluster_shape: tuple[int, int, int],
    dynamic_smem: ir.Value,
    smem_buffers: ShapeTree,
    tmem_allocs: list[
        _TMEMAlloc | _TMEMDialectAlloc
    ],  # Mutated by this function!
    lowering_semantics: LoweringSemantics,
    dynamic_smem_offset: int = 0,
    cta_barrier_offset: int = 0,
) -> Callable[[], RefTree]:
  i32 = ir.IntegerType.get_signless(32)
  i64 = ir.IntegerType.get_signless(64)
  flat_ref_tys, smem_buffer_tree = jax.tree.flatten(
      smem_buffers, is_leaf=lambda x: isinstance(x, Union)
  )
  smem_refs = []

  for ref_ty in flat_ref_tys:
    def barrier_memref(num_barriers: int) -> ir.Value:
      nonlocal dynamic_smem_offset
      barrier_ty = ir.MemRefType.get(
          (num_barriers,),
          dialect.BarrierType.get()
          if lowering_semantics == LoweringSemantics.Warpgroup
          else i64,
          memory_space=utils.smem(),
      )
      barrier_memref = _slice_smem(
            barrier_ty,
            dynamic_smem,
            dynamic_smem_offset,
            lowering_semantics,
        )
      dynamic_smem_offset += num_barriers * utils.MBARRIER_BYTES
      return barrier_memref
    ref: Any
    match ref_ty:
      case Union(members):
        member_thunks = [
            _construct_smem_reftree(
                cluster_shape,
                dynamic_smem,
                m,
                tmem_allocs,
                lowering_semantics,
                dynamic_smem_offset,
                cta_barrier_offset,
            )
            for m in members
        ]
        # TODO(apaszke): This is quadratic, but it shouldn't matter for now...
        dynamic_smem_offset += _smem_tree_size(ref_ty)
        cta_barrier_offset += _cta_barrier_tree_count(ref_ty)

        def ref(member_thunks=member_thunks):
          return Union([t() for t in member_thunks])

      case Barrier(arrival_count, num_barriers, orders_tensor_core):
        init_fn = (
            functools.partial(
                utils.DialectBarrierRef.initialize,
                orders_tensor_core=orders_tensor_core,
            )
            if lowering_semantics == LoweringSemantics.Warpgroup
            else utils.BarrierRef.initialize
        )
        ref = init_fn(barrier_memref(num_barriers), arrival_count=arrival_count)
      case ClusterBarrier(collective_dims, arrival_count, num_barriers):
        ref = utils.CollectiveBarrierRef.initialize(
            barrier_memref(num_barriers), arrival_count, collective_dims,
            cluster_shape, leader_tracked=ref_ty.leader_tracked
        )
      case CtaBarrier(arrival_count, num_barriers):
        ref = utils.CtaBarrierRef.initialize(
            cta_barrier_offset,
            arrival_count=arrival_count,
            num_barriers=num_barriers,
        )
        cta_barrier_offset += num_barriers
      case TMEM(shape, dtype, layout=layout, collective=collective, packing=packing):
        addr_ref = _slice_smem(
            ir.MemRefType.get([], i32, memory_space=utils.smem()),
            dynamic_smem,
            dynamic_smem_offset,
            lowering_semantics,
        )
        packing = 1 if packing is None else packing
        ir_dtype = utils.dtype_to_ir_type(dtype)
        if lowering_semantics == LoweringSemantics.Warpgroup:
          if layout is not None:
            packing = layout.vector_length

          alloc = _TMEMDialectAlloc(
              addr_ref, shape, ir_dtype, packing, collective
          )
          tmem_allocs.append(alloc)
          def ref(alloc=alloc, layout=layout):
            assert alloc.tmem_ref is not None
            if layout is not None:
              layout_attr = layouts.to_layout_attr(layout)
              return dialect.tmem_layout_cast(alloc.tmem_ref, layout_attr)
            else:
              return alloc.tmem_ref

        else:
          if layout is None:
            layout = tcgen05._infer_tmem_layout(shape, collective, packing)
          num_cols = layout.cols_in_shape(shape, utils.bitwidth(ir_dtype))
          tmem_allocs.append(_TMEMAlloc(addr_ref, num_cols, collective))
          def ref(addr_ref=addr_ref, shape=shape, ir_dtype=ir_dtype, layout=layout):
            addr = memref.load(addr_ref, [])
            return tcgen05.TMEMRef(addr, shape, ir_dtype, layout)

        dynamic_smem_offset += 4  # i32 takes up 4 bytes
      case _:
        mlir_dtype = utils.dtype_to_ir_type(ref_ty.dtype)
        tile_smem = _slice_smem(
            ir.MemRefType.get(ref_ty.shape, mlir_dtype, memory_space=utils.smem()),
            dynamic_smem,
            dynamic_smem_offset,
            lowering_semantics,
        )
        dynamic_smem_offset += _count_buffer_bytes(ref_ty)
        ref = tile_smem
    smem_refs.append(ref)
  def ref_tree_thunk():
    refs = []
    for ref in smem_refs:
      if callable(ref):
        ref = ref()
      refs.append(ref)
    return jax.tree.unflatten(smem_buffer_tree, refs)
  return ref_tree_thunk


def _smem_tree_size(smem_buffers: ShapeTree) -> int:
  leaves = jax.tree.leaves(
      smem_buffers, is_leaf=lambda x: isinstance(x, Union)
  )
  size = 0
  for l in leaves:
    match l:
      case Union(members):
        size += max(_smem_tree_size(s) for s in members)
      case (
          Barrier(num_barriers=num_barriers)
          | ClusterBarrier(num_barriers=num_barriers)
      ):
        if size % utils.MBARRIER_BYTES:
          raise NotImplementedError(
              "Misaligned barrier allocation. Expected smem size"
              f" ({size} bytes) to be divisible by the size of the barrier:"
              f" {utils.MBARRIER_BYTES} bytes."
          )
        size += num_barriers * utils.MBARRIER_BYTES
      case CtaBarrier():
        pass
      case TMEM(_):
        # TODO(justinfu): This can trigger misaligned barrier allocations
        # if TMEM is requested before barriers b/c it's not divisible by 8.
        size += 4  # i32 takes up 4 bytes
      case _:
        size += _count_buffer_bytes(l)
  return size


def _cta_barrier_tree_count(smem_buffers: ShapeTree) -> int:
  leaves = jax.tree.leaves(smem_buffers, is_leaf=lambda x: isinstance(x, Union))
  count = 0
  for l in leaves:
    match l:
      case Union(members):
        count += max(_cta_barrier_tree_count(s) for s in members)
      case CtaBarrier(num_barriers=num_barriers):
        count += num_barriers
      case _:
        pass
  return count


def _is_known_multihost_mesh(
    mesh: mesh_lib.Mesh | mesh_lib.AbstractMesh | None,
) -> bool:
  """Returns whether the given mesh is a known multihost mesh.

  Mesh devices should have either slice_index, logical_task_id, or process_index
  attributes. Without these attributes we can't reliably determine if the mesh
  is multihost.

  Args:
    mesh: The mesh to check.

  Returns:
    True if the mesh is a known multihost mesh, False otherwise.
  """
  if mesh is None or mesh.empty or not isinstance(mesh, mesh_lib.Mesh):
    return False
  tasks = {
      (
          getattr(d, "slice_index", 0),
          getattr(d, "logical_task_id", 0),
          getattr(d, "process_index", 0),
      )
      for d in mesh.devices.flat
  }
  return len(tasks) > 1


def _cluster_arrive_relaxed_aligned():
  llvm.inline_asm(
      ir.Type.parse("!llvm.void"),
      [],
      "barrier.cluster.arrive.relaxed.aligned;",
      "",
      has_side_effects=True,
  )


def _cluster_wait_aligned():
  llvm.inline_asm(
      ir.Type.parse("!llvm.void"),
      [],
      "barrier.cluster.wait.aligned;",
      "",
      has_side_effects=True,
  )


# TODO(apaszke): Inline this
@contextlib.contextmanager
def _launch(
    token,
    grid: tuple[int, int, int],
    cluster: tuple[int, int, int],
    block: tuple[int, int, int],
    smem_buffers: ShapeTree | Union[ShapeTree],
    lowering_semantics: LoweringSemantics,
    module: ir.Module,
    inout_buffers_ptr: ir.Value,
    profiler_spec: profiler.ProfilerSpec | None = None,
    maybe_prof_buffer: ir.Value | None = None,
    num_peers: int = 0,
    num_params: int = 0,
    jax_mesh: mesh_lib.Mesh | None = None,
    is_multi_process: bool = False,
):
  if (profiler_spec is None) != (maybe_prof_buffer is None):
    raise ValueError(
        "Both profiler_spec and maybe_prof_buffer must be specified or"
        " left unspecified."
    )
  index = ir.IndexType.get()
  i32 = ir.IntegerType.get_signless(32)
  i8 = ir.IntegerType.get_signless(8)
  grid_vals = [c(i, index) for i in grid]
  block_vals = [c(i, index) for i in block]

  user_smem_bytes = _smem_tree_size(smem_buffers)

  smem_bytes = user_smem_bytes
  if profiler_spec is not None:
    # Profiler array stores values in 64 bit chunks (vectors of size 2
    # of 32-bit elements), and so the starting address needs to be 64
    # bit = 8 byte aligned.
    # https://docs.nvidia.com/cuda/parallel-thread-execution/#addresses-as-operands:~:text=The%20address%20must%20be%20naturally%20aligned%20to%20a%20multiple%20of%20the%20access%20size.
    align = 8
    profiler_start = (smem_bytes + align - 1) & ~(align - 1)
    smem_bytes = profiler_start + profiler_spec.smem_bytes(block=block)

  device = jax.local_devices()[0]
  # For ahead-of-time compilation purposes, that is when a CUDA device
  # isn't available to query directly, we default to 227 KB, the
  # maximum amount of shared memory per thread block available in
  # compute capabilities 9.0 and 10.x:
  # https://docs.nvidia.com/cuda/cuda-c-programming-guide/#features-and-technical-specifications-technical-specifications-per-compute-capability
  # Note in either case we assume all devices have the same amount of
  # shared memory.
  max_smem_bytes = getattr(device, "shared_memory_per_block_optin", 227 * 1024)
  if _SMEM_SIZE_BOUND is not None:
    max_smem_bytes = min(max_smem_bytes, _SMEM_SIZE_BOUND)
  if smem_bytes > max_smem_bytes:
    raise ValueError("Mosaic GPU kernel exceeds available shared memory: "
                     f"{smem_bytes=} > {max_smem_bytes=}")
  if math.prod(cluster) != 1:
    if len(cluster) != 3:
      raise ValueError(f"Clusters must be 3D. Got: {cluster}")
    cluster_kwargs = {
        "clusterSize" + d: c(s, index) for s, d in zip(cluster, "XYZ")
    }
    for d, grid_size, cluster_size in zip("xyz", grid, cluster):
      if grid_size % cluster_size != 0:
        raise ValueError(
            f"Grid dimension {d} must be divisible by cluster dimension:"
            f" {grid_size} % {cluster_size} != 0"
        )
  else:
    cluster_kwargs = {}
  # `gpu.LaunchOp` is missing the clusterSize{X,Y,Z} arguments.
  launch_op = _gpu_ops_gen.LaunchOp(
      token.type,
      [token],
      *grid_vals,
      *block_vals,
      dynamicSharedMemorySize=c(smem_bytes, i32),
      **cluster_kwargs,
  )
  launch_op.body.blocks.append(*([index] * (12 + 2 * len(cluster_kwargs))))  # Append an empty block
  with ir.InsertionPoint(launch_op.body.blocks[0]):
    dynamic_smem = gpu.dynamic_shared_memory(
        ir.MemRefType.get((utils.DYNAMIC,), i8, memory_space=utils.smem())
    )

    if profiler_spec:
      prof_smem = _slice_smem(
          ir.MemRefType.get(
              (profiler_spec.smem_i32_elements(block=block),),
              i32,
              memory_space=utils.smem(),
          ),
          dynamic_smem,
          profiler_start,  # pyrefly: ignore[unbound-name]
          lowering_semantics,
      )
      if lowering_semantics == LoweringSemantics.Warpgroup:
        prof_smem = dialect.with_transforms(prof_smem, ir.ArrayAttr.get([]))
        wrap_in_custom_primitive = True
      else:
        wrap_in_custom_primitive = False
      assert maybe_prof_buffer is not None
      prof = profiler.OnDeviceProfiler(
          profiler_spec,
          prof_smem,
          maybe_prof_buffer,
          wrap_in_custom_primitive,
      )
    else:
      prof = None

    ctx = launch_context.LaunchContext(
        module,
        launch_context.Scratch(launch_op),
        cluster,
        inout_buffers_ptr,
        prof,
        num_peers=num_peers,
        num_params=num_params,
        multihost_kernel=is_multi_process or _is_known_multihost_mesh(
            jax_mesh
        ),
    )
    with ctx.named_region("Init"):
      num_warpgroups = math.prod(block) // utils.WARPGROUP_SIZE
      reserved_cta_barriers = 1 + num_warpgroups
      total_cta_barriers = _cta_barrier_tree_count(smem_buffers)
      if reserved_cta_barriers + total_cta_barriers > 16:
        raise ValueError(
            f"Too many CTA barriers requested: {total_cta_barriers} requested,"
            f" but only {16 - reserved_cta_barriers} available"
            f" ({num_warpgroups} warpgroup barriers and 1 CTA-wide barrier"
            " reserved)."
        )
      tmem_allocs: list[_TMEMAlloc | _TMEMDialectAlloc] = []
      smem_ref_tree_thunk = _construct_smem_reftree(
          cluster,
          dynamic_smem,
          smem_buffers,
          tmem_allocs,
          lowering_semantics,
          cta_barrier_offset=reserved_cta_barriers,
      )
      # TODO(apaszke): Skip fences if no barriers or TMEM is initialized.
      # TODO(apaszke): Only initialize cluster barriers before the cluster wait.
      if utils.get_arch().major >= 9:
        llvm.inline_asm(
            ir.Type.parse("!llvm.void"),
            [],
            "fence.mbarrier_init.release.cluster;",
            "",
            has_side_effects=True,
        )
      if math.prod(cluster) != 1:
        _cluster_arrive_relaxed_aligned()
        _cluster_wait_aligned()
      if tmem_allocs:
        init_warp_ctx: contextlib.AbstractContextManager
        if lowering_semantics == LoweringSemantics.Warpgroup:
          init_warp_ctx = contextlib.nullcontext()
        else:
          eq = arith.CmpIPredicate.eq
          is_init_warp = arith.cmpi(eq, utils.warp_idx(sync=False), c(0, i32))
          init_warp_ctx = utils.when(is_init_warp)
        with init_warp_ctx:
          cols_used = 0
          for alloc in tmem_allocs:
            cols_used += alloc.alloc()
          if cols_used > tcgen05.TMEM_MAX_COLS:
            raise ValueError(
                "Total TMEM allocation exceeds memory limit. "
                f"Requested {cols_used} columns which exceeds limit of "
                f"{tcgen05.TMEM_MAX_COLS}."
            )
          collective_types = {alloc.collective for alloc in tmem_allocs}
          if len(collective_types) > 1:
            raise ValueError(
                "Can't mix collective and non-collective TMEM allocations"
                " within the same kernel."
            )
          collective = True in collective_types
          if collective and math.prod(cluster) % 2:
            raise ValueError(
                "Collective TMEM allocations are only supported for clusters"
                " with an even number of blocks in them. Got cluster:"
                f" {cluster}"
            )
          if lowering_semantics == LoweringSemantics.Warpgroup:
            dialect.tmem_relinquish_alloc_permit(collective=collective)
          else:
            tcgen05.tmem_relinquish_alloc_permit(collective=collective)
      gpu.barrier()  # Make sure the init is visible to all threads.
      smem_ref_tree = smem_ref_tree_thunk()

    yield ctx, smem_ref_tree

    if tmem_allocs:
      gpu.barrier()  # Make sure everyone is done before we release TMEM.
      if any(alloc.collective for alloc in tmem_allocs):
        _cluster_arrive_relaxed_aligned()
        _cluster_wait_aligned()
      if lowering_semantics == LoweringSemantics.Warpgroup:
        init_warp_ctx = contextlib.nullcontext()
      else:
        init_warp_ctx = utils.when(is_init_warp)  # pyrefly: ignore[unbound-name]
      with init_warp_ctx:
        for alloc in tmem_allocs:
          alloc.dealloc()
    if prof is not None:
      prof.finalize(grid=grid, block=block)
    gpu.terminator()


def _lower_as_gpu_kernel(
    body,
    grid: tuple[int, int, int],
    cluster: tuple[int, int, int],
    block: tuple[int, int, int],
    in_shapes: tuple[Any, ...],
    out_shape,
    inout_shape,
    smem_scratch_shape: ShapeTree | Union[ShapeTree],
    lowering_semantics: LoweringSemantics,
    module_name: str,
    kernel_name: str,
    prof_spec: profiler.ProfilerSpec | None = None,
    jax_mesh: mesh_lib.Mesh | None = None,
    base_loc: ir.Location | None = None,
    uses_pdl: bool = False,
    is_multi_process: bool = False,
):
  ptr_ty = llvm.PointerType.get()
  token_ty = gpu.AsyncTokenType.get()
  i8 = ir.IntegerType.get_signless(8)
  i32 = ir.IntegerType.get_signless(32)

  def _shape_to_ref_ty(shape: jax.ShapeDtypeStruct) -> ir.MemRefType:
    return ir.MemRefType.get(shape.shape, utils.dtype_to_ir_type(shape.dtype))

  in_ref_tys = [_shape_to_ref_ty(t) for t in in_shapes]
  inout_ref_tys = [_shape_to_ref_ty(t) for t in inout_shape]

  unwrap_output_tuple = False
  if isinstance(out_shape, list):
    out_shape = tuple(out_shape)
  elif not isinstance(out_shape, tuple):
    out_shape = (out_shape,)
    unwrap_output_tuple = not inout_shape
  out_ref_tys = [_shape_to_ref_ty(t) for t in out_shape]
  if prof_spec is not None:
    out_shape = (*out_shape, prof_spec.jax_buffer_type(grid, block))
    out_ref_tys.append(prof_spec.mlir_buffer_type(grid, block))

  module = ir.Module.create(loc=base_loc)
  dialect.register_dialect(module.context)
  attrs = module.operation.attributes
  attrs["sym_name"] = ir.StringAttr.get(module_name)
  arch_major, arch_minor = utils._infer_arch()
  attrs["mosaic_gpu.arch_major"] = ir.IntegerAttr.get(i32, arch_major)
  attrs["mosaic_gpu.arch_minor"] = ir.IntegerAttr.get(i32, arch_minor)
  if uses_pdl:
    attrs["mosaic_gpu.uses_pdl"] = ir.UnitAttr.get()

  # These are needed as nonlocal below.
  launch_ctx = None
  with ir.InsertionPoint(module.body):
    _declare_runtime_functions()
    global_scratch = llvm.GlobalOp(
        llvm.ArrayType.get(i8, 0),  # We don't know the shape yet.
        "global_scratch",
        ir.Attribute.parse("#llvm.linkage<external>"),
        addr_space=ir.IntegerAttr.get(i32, 4),  # GPU constant memory.
    )
    @func.FuncOp.from_py_func(ptr_ty, ptr_ty, name=f"{kernel_name}_mosaic_gpu")
    def main(token_ptr, buffers):
      nonlocal launch_ctx
      token = builtin.unrealized_conversion_cast([token_ty], [token_ptr])
      arg_refs = []
      # XLA will pass in inout refs again as outputs, but we ignore them.
      for i, ref_ty in enumerate([*in_ref_tys, *inout_ref_tys, *out_ref_tys]):
        ptr = llvm.load(ptr_ty, utils.getelementptr(buffers, [i], ptr_ty))
        arg_memref = utils.ptr_as_memref(ptr, ir.MemRefType(ref_ty))
        # Annotate so we can find the corresponding kernel argument during the
        # lowering.
        arg_memref.owner.attributes[launch_context.KERNEL_ARG_ID_ATTR] = (
            ir.IntegerAttr.get(i32, i)
        )
        arg_memref.owner.attributes[launch_context.ORIGINAL_KERNEL_ARG_ATTR] = (
            ir.UnitAttr.get()
        )
        arg_refs.append(arg_memref)

      num_peers = 0
      num_params = 0

      if (
          jax_mesh is not None
          and jax_mesh.size > 1
          and (
              is_single_process_multi_device_topology()
              or not is_nvshmem_available()
          )
      ):
        num_params = len(arg_refs) + len(inout_ref_tys)
        num_peers = jax_mesh.size

      prof_buffer = arg_refs.pop() if prof_spec is not None else None

      with _launch(
          token,
          grid,
          cluster,
          block,
          smem_scratch_shape,
          lowering_semantics,
          module,
          buffers,
          prof_spec,
          prof_buffer,
          num_peers,
          num_params,
          jax_mesh,
          is_multi_process=is_multi_process,
      ) as (_launch_ctx, smem_refs):
        launch_ctx = _launch_ctx
        body(launch_ctx, *arg_refs, smem_refs)
    main.func_op.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get()
  sym_tab = ir.SymbolTable(module.operation)
  sym_tab.insert(main.func_op)
  sym_tab.insert(global_scratch)
  try:
    module.operation.verify()
  except ir.MLIRError as e:
    try:
      new_error = error.mlir_error_to_verification_error(e)
    except:
      new_error = None
    if new_error is None:
      raise
    raise error.mlir_error_to_verification_error(e) from e

  assert launch_ctx is not None
  return module, out_shape, unwrap_output_tuple, launch_ctx


def _run_serde_pass(
    module: ir.Module, *, serialize: bool, ir_version: int | None = None
) -> ir.Module:
  bytecode_buffer = io.BytesIO()
  module.operation.write_bytecode(bytecode_buffer)
  module = ir.Module.parse(bytecode_buffer.getvalue(), context=module.context)
  pipeline = passmanager.PassManager.parse(
      "builtin.module(mosaic_gpu-serde{serialize="
      + str(serialize).lower()
      + (f" target-version={ir_version}" if ir_version is not None else "")
      + "})",
      module.context,
  )
  allow_unregistered_dialects = module.context.allow_unregistered_dialects
  module.context.allow_unregistered_dialects = True
  try:
    pipeline.run(module.operation)
  except ir.MLIRError as e:
    raise error.mlir_error_to_verification_error(e) from e
  finally:
    module.context.allow_unregistered_dialects = allow_unregistered_dialects
  return module


def _declare_runtime_functions():
  """Declares the runtime functions that can be used by the generated code."""
  ptr_ty = llvm.PointerType.get()
  i64 = ir.IntegerType.get_signless(64)
  arg_tys = [ptr_ty, ptr_ty, i64, i64, ptr_ty, ptr_ty, i64, ptr_ty]
  init_tma_desc_type = ir.FunctionType.get(arg_tys, [])
  func.FuncOp(
      "mosaic_gpu_init_tma_desc", init_tma_desc_type, visibility="private"
  )
  im2col_arg_tys = [
      ptr_ty, ptr_ty, i64, i64, ptr_ty, ptr_ty, i64, ptr_ty, ptr_ty, i64, i64, ptr_ty,
  ]
  init_tma_im2col_desc_type = ir.FunctionType.get(im2col_arg_tys, [])
  func.FuncOp(
      "mosaic_gpu_init_tma_im2col_desc",
      init_tma_im2col_desc_type,
      visibility="private",
  )


def lower_mgpu_module(
    module: ir.Module,
    launch_ctx: launch_context.LaunchContext,
    lowering_semantics: LoweringSemantics,
    *,
    auto_barriers: bool = True,
) -> None:
  if lowering_semantics == LoweringSemantics.Warpgroup:
    dump_options = dialect.get_or_set_dump_options(module)

    # We need to run a pass that removes dead-code for which layout inference
    # does not work.
    pm = mlir.passmanager.PassManager.parse("builtin.module(canonicalize,cse)", module.context)
    pm.enable_verifier(False)
    pm.run(module.operation)

    # Run Python lowering passes. The remaining passes will be run in C++ in
    # jax/jaxlib/mosaic/gpu/custom_call.cc
    if dump_options.mlir_passes:
      utils.dump_to_file_or_stdout(
          str(module),
          f"{dump_options.module_basename}.before_layout_inference.txt",
          dump_options.dump_path
      )

    layout_inference.infer_layout(module, arch=utils._infer_arch())

    if dump_options.mlir_passes:
      utils.dump_to_file_or_stdout(
          str(module),
          f"{dump_options.module_basename}.after_layout_inference.txt",
          dump_options.dump_path
      )

    dialect_lowering.lower_mgpu_dialect(
        module, launch_ctx, auto_barriers=auto_barriers
    )

  launch_ctx.scratch.finalize_size()
  try:
    module.operation.verify()
  except ir.MLIRError as e:
    raise error.mlir_error_to_verification_error(e) from e


def _kernel_to_module(
    body,
    grid: tuple[int, int, int],
    block: tuple[int, int, int],
    in_shape,
    out_shape,
    smem_scratch_shape: ShapeTree | Union[ShapeTree],
    prof_spec: profiler.ProfilerSpec | None = None,
    cluster: tuple[int, int, int] = (1, 1, 1),
    module_name: str = "unknown",
    kernel_name: str | None = None,
    thread_semantics: LoweringSemantics = LoweringSemantics.Lane,
    inout_shape = (),
    is_multi_process: bool = False,
):
  if isinstance(in_shape, list):
    in_shape = tuple(in_shape)
  elif not isinstance(in_shape, tuple):
    in_shape = (in_shape,)
  if isinstance(inout_shape, list):
    inout_shape = tuple(inout_shape)
  elif not isinstance(inout_shape, tuple):
    inout_shape = (inout_shape,)
  if kernel_name is None:
    kernel_name = jax_util.fun_name(body, "anonymous")

  inout_shape = jax.tree.map(jax.ShapeDtypeStruct.like, inout_shape)
  out_shape = jax.tree.map(jax.ShapeDtypeStruct.like, out_shape)
  jax_mesh = mesh_lib.get_concrete_mesh()
  if jax_mesh.empty:
    jax_mesh = mesh_lib.thread_resources.env.physical_mesh
  module, out_shape, unwrap_output_tuple, launch_ctx = (
      _lower_as_gpu_kernel(
          body, grid, cluster, block, in_shape, out_shape, inout_shape,
          smem_scratch_shape, thread_semantics, module_name, kernel_name,
          prof_spec, jax_mesh=jax_mesh, is_multi_process=is_multi_process,
      )
  )

  lower_mgpu_module(module, launch_ctx, thread_semantics)

  return (
      module,
      in_shape,
      inout_shape,
      out_shape,
      unwrap_output_tuple,
      launch_ctx.is_device_collective,
  )


def as_gpu_kernel(
    body,
    grid: tuple[int, int, int],
    block: tuple[int, int, int],
    in_shape,
    out_shape,
    smem_scratch_shape: ShapeTree | Union[ShapeTree],
    prof_spec: profiler.ProfilerSpec | None = None,
    cluster: tuple[int, int, int] = (1, 1, 1),
    module_name: str = "unknown",
    kernel_name: str | None = None,
    ir_version: int | None = None,
    thread_semantics: LoweringSemantics = LoweringSemantics.Lane,
    inout_shape = (),
    is_multi_process: bool = False,
):
  module, in_shape, inout_shape, out_shape, unwrap_output_tuple, is_device_collective = _kernel_to_module(
      body, grid, block, in_shape, out_shape, smem_scratch_shape, prof_spec,
      cluster, module_name, kernel_name, thread_semantics, inout_shape,
      is_multi_process=is_multi_process,
  )

  expected_arg_tys, expected_arg_treedef = jax.tree.flatten((*in_shape, *inout_shape))
  def _check_args(*args):
    arg_treedef = jax.tree.structure(args)
    if arg_treedef != expected_arg_treedef:
      raise ValueError(
          f"Invalid argument structure: expected {expected_arg_treedef}, got"
          f" {arg_treedef}, ({args=})"
      )
    for arg, expected_ty in zip(args, expected_arg_tys):
      if arg.shape != expected_ty.shape:
        raise ValueError(
            f"Argument shape mismatch: expected {expected_ty.shape}, got"
            f" {arg.shape}"
        )
      if arg.dtype != expected_ty.dtype:
        hint = ""
        if not arg.shape:
          hint = f". Hint: cast the scalar to {expected_ty.dtype} explicitly."
        raise ValueError(
            f"Argument dtype mismatch: expected {expected_ty.dtype}, got"
            f" {arg.dtype}{hint}"
        )

  def bind(*args) -> Any:
    return mosaic_gpu_p.bind(
        *args,
        module=module,
        out_types=out_shape,
        inout_types=inout_shape,
    )

  if prof_spec is not None:
    @jax.jit
    def prof_kernel(*args):
      _check_args(*args)
      *results, prof_buffer = bind(*args)
      def dump_profile(prof_buffer):
        out_file = os.path.join(prof_spec.dump_path, f"{time.time_ns()}-trace.json")
        try:
          with open(out_file, "x") as f:
            prof_spec.dump(prof_buffer, f, grid=grid, block=block)
        except FileExistsError:
          pass  # TODO: Retry
      jax.debug.callback(dump_profile, prof_buffer)
      return results[0] if unwrap_output_tuple else results
    return prof_kernel
  else:
    @jax.jit
    def kernel(*args):
      _check_args(*args)
      results = bind(*args)
      return results[0] if unwrap_output_tuple else results
    return kernel


def as_torch_gpu_kernel(
    body,
    grid: tuple[int, int, int],
    block: tuple[int, int, int],
    in_shape,
    out_shape,
    smem_scratch_shape: ShapeTree | Union[ShapeTree],
    prof_spec: profiler.ProfilerSpec | None = None,
    cluster: tuple[int, int, int] = (1, 1, 1),
    module_name: str = "unknown",
    kernel_name: str | None = None,
    thread_semantics: LoweringSemantics = LoweringSemantics.Lane,
    inout_shape=(),
):
  (
      module,
      in_shape,
      inout_shape,
      out_shape,
      unwrap_output_tuple,
      is_device_collective,
  ) = _kernel_to_module(
      body,
      grid,
      block,
      in_shape,
      out_shape,
      smem_scratch_shape,
      prof_spec,
      cluster,
      module_name,
      kernel_name,
      thread_semantics,
      inout_shape,
  )
  module = _run_serde_pass(module, serialize=True, ir_version=None)
  bytecode_buffer = io.BytesIO()
  module.operation.write_bytecode(bytecode_buffer)
  return _as_torch_gpu_kernel(
      bytecode_buffer.getvalue(),
      in_shape,
      out_shape,
      inout_shape,
      unwrap_output_tuple=unwrap_output_tuple,
  )


def _compile_as_torch_gpu_kernel(module_asm: bytes):
  try:
    import torch  # pyrefly: ignore[missing-import]
  except ImportError:
    raise RuntimeError("Can't compile for PyTorch: import torch failed") from None

  torch.cuda.init()  # Make sure CUDA context is set up.

  # Get our hands on the compilation and unload functions
  try:
    try:
      import jax_plugins.xla_cuda13 as cuda_plugin  # pyrefly: ignore[missing-import]
    except ImportError:
      import jax_plugins.xla_cuda12 as cuda_plugin  # pyrefly: ignore[missing-import]
  except ImportError:
    dll = ctypes.CDLL(None)
  else:
    dll = ctypes.CDLL(cuda_plugin._get_library_path())
  compile_func = dll.MosaicGpuCompile
  compile_func.argtypes = [ctypes.c_char_p, ctypes.c_int]
  compile_func.restype = ctypes.POINTER(ctypes.c_void_p)
  unload_func = dll.MosaicGpuUnload
  unload_func.argtypes = [compile_func.restype]
  unload_func.restype = None

  compiled = compile_func(ctypes.c_char_p(module_asm), ctypes.c_int(len(module_asm)))
  if not compiled:
    raise RuntimeError("Failed to compile the module")
  function, launch_ptr, kernel = compiled[0], compiled[1], compiled[2]
  launch_arg_types = [
      ctypes.c_void_p, # kernel
      ctypes.c_void_p, # function
      ctypes.c_void_p, # stream
      ctypes.POINTER(ctypes.c_void_p), # arg_ptrs
      ctypes.c_int, # device_ordinal
  ]
  if not _ABI_V2:
    # Drop kernel and device_ordinal for ABI v1.
    launch_arg_types = launch_arg_types[1:-1]
  launch_c = ctypes.CFUNCTYPE(None, *launch_arg_types)(launch_ptr)

  def launch(arg_ptrs, device):
    device = torch.device(device)
    args = [
        kernel,
        function,
        torch.cuda.default_stream(device)._as_parameter_,
        arg_ptrs,
        device.index
        if device.index is not None
        else torch.cuda.current_device(),
    ]
    if not _ABI_V2:
      args = args[1:-1]
    launch_c(*args)

  return launch, functools.partial(unload_func, compiled)


def _as_torch_gpu_kernel(
    module_asm: bytes,
    in_shape: Iterable[object],
    out_shape: Iterable[object],
    inout_shape: Iterable[object] = (),
    *,
    unwrap_output_tuple: bool = False,
    _prepare_args = None,
    _prepare_results = None,
):
  flat_arg_types, expected_arg_treedef = jax.tree.flatten((*in_shape, *inout_shape))
  flat_out_types, _ = jax.tree.flatten(out_shape)
  out_treedef = jax.tree.structure((*out_shape, *inout_shape))

  launch, unload = _compile_as_torch_gpu_kernel(module_asm)
  # _compile_as_torch_gpu_kernel checks that this succeeds
  import torch  # pyrefly: ignore[missing-import]

  def as_torch_dtype(dtype):
    # torch contains NumPy-compatible dtypes in its top namespace
    return getattr(torch, np.dtype(dtype).name)

  def apply(*args):
    flat_args, arg_treedef = jax.tree.flatten(args)
    if arg_treedef != expected_arg_treedef:
      raise ValueError(
          f"Invalid argument structure: expected {expected_arg_treedef}, got"
          f" {arg_treedef}, ({args=})"
      )
    for arg, expected_ty in zip(flat_args, flat_arg_types):
      if arg.shape != expected_ty.shape:
        raise ValueError(
            f"Argument shape mismatch: expected {expected_ty.shape}, got"
            f" {arg.shape}"
        )
      if arg.dtype != as_torch_dtype(expected_ty.dtype):
        raise ValueError(
            "Argument dtype mismatch: expected"
            f" {as_torch_dtype(expected_ty.dtype)}, got {arg.dtype}"
        )

    # Construct a device pointer list like in the XLA calling convention
    buffers = (ctypes.c_void_p * (arg_treedef.num_leaves + out_treedef.num_leaves))()
    i = -1  # Define i in case there are no args
    device = 'cuda'
    for i, arg in enumerate(flat_args):
      buffers[i] = arg.data_ptr()
      device = arg.device
    flat_outs = []
    for i, t in enumerate(flat_out_types, i + 1):
      out = torch.empty(t.shape, dtype=as_torch_dtype(t.dtype), device=device)
      flat_outs.append(out)
      buffers[i] = out.data_ptr()
    if num_inout_args := jax.tree.structure(inout_shape).num_leaves:
      flat_outs += flat_args[-num_inout_args:]
    launch(buffers, device)
    out = jax.tree.unflatten(out_treedef, flat_outs)
    return out[0] if unwrap_output_tuple else out

  # Unload the compiled code when the Python function is destroyed.
  apply.destructor = weakref.ref(apply, lambda _weak_ref: unload)  # pyrefly: ignore[missing-attribute]

  return apply