# Copyright 2023 The JAX Authors.
#
# 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
#
#     https://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.

"""Contains registrations for pallas_call on TPU."""

from __future__ import annotations

from collections.abc import Hashable, Sequence
import dataclasses
import json
from typing import cast

import jax
from jax._src import core as jax_core
from jax._src import dtypes
from jax._src import flattree as ft
from jax._src import frozen_dict
from jax._src import jaxpr_util
from jax._src import sharding_impls
from jax._src import state
from jax._src import tpu_custom_call
from jax._src import xla_metadata
from jax._src.interpreters import mlir
from jax._src.interpreters import partial_eval as pe
from jax._src.lib.mlir import ir
from jax._src.lib.mlir import passmanager
from jax._src.pallas import core as pallas_core
from jax._src.pallas import mpmd
from jax._src.pallas.mosaic import core as tpu_core
from jax._src.pallas.mosaic import helpers
from jax._src.pallas.mosaic import lowering
from jax._src.pallas.mosaic import primitives as tpu_primitives
from jax._src.pallas.mosaic import sc_core
from jax._src.pallas.mosaic import sc_lowering  # noqa: F401
from jax._src.pallas.mosaic import tpu_info
from jax._src.state import types as state_types
from jax.experimental import mosaic
from jax.experimental.mosaic.dialects import tpu


def _get_memory_space_from_aval(
    out_aval: jax_core.AbstractValue, kernel_type: tpu_core.CoreType | None
) -> tpu_custom_call.MemorySpace | None:
  if not isinstance(out_aval, jax_core.ShapedArray):
    raise ValueError("Memory spaces not defined for non-ShapedArrays")
  if not isinstance(
      ms := getattr(out_aval, "memory_space", None),
      (
          tpu_core.MemorySpace,
          pallas_core.MemorySpace,
          pallas_core.CoreMemorySpace,
          jax_core.MemorySpace,
      ),
  ):
    return None  # If we are passed a non-TPU memory space, ignore it.
  # If we are passed an aval with an explicit memory space tag, we use it
  # to constrain the memory space.
  match ms:
    case tpu_core.MemorySpace.HBM:
      return tpu_custom_call.MemorySpace.HBM
    case tpu_core.MemorySpace.VMEM:
      return tpu_custom_call.MemorySpace.VMEM
    case tpu_core.MemorySpace.SMEM:
      match kernel_type:
        case tpu_core.CoreType.SC_SCALAR_SUBCORE:
          return tpu_custom_call.MemorySpace.SC_SCALAR_SMEM
        case tpu_core.CoreType.SC_VECTOR_SUBCORE:
          return tpu_custom_call.MemorySpace.SC_VECTOR_SMEM
        case tpu_core.CoreType.TC:
          return tpu_custom_call.MemorySpace.SMEM
        case _:
          raise ValueError(f"Invalid kernel type for SMEM: {kernel_type}")
    case tpu_core.MemorySpace.SEMAPHORE:
      match kernel_type:
        case tpu_core.CoreType.SC_SCALAR_SUBCORE:
          return tpu_custom_call.MemorySpace.SC_SCALAR_SEMAPHORE_MEM
        case tpu_core.CoreType.TC:
          return tpu_custom_call.MemorySpace.SEMAPHORE_MEM
        case _:
          raise ValueError(f"Invalid kernel type for semaphore: {kernel_type}")
    case jax_core.MemorySpace.Host:
      return tpu_custom_call.MemorySpace.HOST
    case pallas_core.CoreMemorySpace(tpu_core.MemorySpace.VMEM, mesh):
      match mesh.core_type:
        case tpu_core.CoreType.TC:
          return tpu_custom_call.MemorySpace.VMEM
        case _:
          raise ValueError(f"Invalid core type for VMEM: {mesh.core_type}")
    case pallas_core.CoreMemorySpace(tpu_core.MemorySpace.SMEM, mesh):
      match mesh.core_type:
        case tpu_core.CoreType.SC_SCALAR_SUBCORE:
          return tpu_custom_call.MemorySpace.SC_SCALAR_SMEM
        case tpu_core.CoreType.SC_VECTOR_SUBCORE:
          return tpu_custom_call.MemorySpace.SC_VECTOR_SMEM
        case tpu_core.CoreType.TC:
          return tpu_custom_call.MemorySpace.SMEM
        case _:
          raise ValueError(f"Invalid core type for SMEM: {mesh.core_type}")
    case pallas_core.CoreMemorySpace(tpu_core.MemorySpace.SEMAPHORE, mesh):
      match mesh.core_type:
        case tpu_core.CoreType.SC_SCALAR_SUBCORE:
          return tpu_custom_call.MemorySpace.SC_SCALAR_SEMAPHORE_MEM
        case tpu_core.CoreType.SC_VECTOR_SUBCORE:
          return tpu_custom_call.MemorySpace.SC_VECTOR_SEMAPHORE_MEM
        case tpu_core.CoreType.TC:
          return tpu_custom_call.MemorySpace.SEMAPHORE_MEM
        case _:
          raise ValueError(
              f"Invalid core type for semaphore: {mesh.core_type}"
          )
    case _:
      pass
  return None


def _get_memory_spaces_from_avals(
    avals: Sequence[jax_core.AbstractValue],
    kernel_type: tpu_core.CoreType | None,
) -> tuple[tpu_custom_call.MemorySpace | None, ...] | None:
  memory_spaces = None
  if any(isinstance(aval, jax_core.ShapedArray)
         and (
             not isinstance(aval.memory_space, jax_core.MemorySpace)
             or aval.memory_space is jax_core.MemorySpace.Host
         )
         for aval in avals):
    memory_spaces = tuple(
        _get_memory_space_from_aval(aval, kernel_type=kernel_type)
        for aval in avals)
  return memory_spaces


def _resolve_memory_spaces(
    in_avals: Sequence[jax_core.AbstractValue],
    out_avals: Sequence[jax_core.AbstractValue],
    *,
    input_output_aliases: tuple[tuple[int, int], ...],
    kernel_type: tpu_core.CoreType | None,
) -> tuple[
    tuple[tpu_custom_call.MemorySpace | None, ...] | None,
    tuple[tpu_custom_call.MemorySpace | None, ...] | None,
]:
  output_memory_spaces = _get_memory_spaces_from_avals(
      out_avals, kernel_type=kernel_type
  )
  input_memory_spaces = None
  if any(isinstance(aval, jax_core.ShapedArray)
         and not isinstance(aval.memory_space, jax_core.MemorySpace)
         for aval in in_avals):
    input_memory_spaces = _get_memory_spaces_from_avals(
        in_avals, kernel_type=kernel_type
    )
    # ShapedArrayWithMemorySpace wasn't allowed to escape the Pallas kernel, so
    # it was stripped from the original out_avals. We need to resolve this for
    # outputs aliased to inputs with pltpu.with_memory_space_constraint.
    if input_memory_spaces is not None:
      output_memory_spaces_list: list[tpu_custom_call.MemorySpace | None]
      if output_memory_spaces is None:
        output_memory_spaces_list = [None] * len(out_avals)
      else:
        output_memory_spaces_list = list(output_memory_spaces)
      modified = False
      for in_idx, out_idx in input_output_aliases:
        if (ms := input_memory_spaces[in_idx]) is not None:
          if (out_ms := output_memory_spaces_list[out_idx]) not in (None, ms):
            raise ValueError(
                f"In/out memory space conflict for alias {in_idx}->{out_idx}:"
                f" {ms} != {out_ms}"
            )
          output_memory_spaces_list[out_idx] = ms
          modified = True
      if modified:
        output_memory_spaces = tuple(output_memory_spaces_list)
  if output_memory_spaces is not None:
    input_memory_spaces_list: list[tpu_custom_call.MemorySpace | None]
    if input_memory_spaces is None:
      input_memory_spaces_list = [None] * len(in_avals)
    else:
      input_memory_spaces_list = list(input_memory_spaces)
    modified = False
    for in_idx, out_idx in input_output_aliases:
      if (ms := output_memory_spaces[out_idx]) is not None:
        if (in_ms := input_memory_spaces_list[in_idx]) not in (None, ms):
          raise ValueError(
              f"In/out memory space conflict for alias {in_idx}->{out_idx}:"
              f"{in_ms} != {ms}"
          )
        input_memory_spaces_list[in_idx] = ms
        modified = True
    if modified:
      input_memory_spaces = tuple(input_memory_spaces_list)
  if input_memory_spaces is not None:
    input_memory_spaces = tuple(
        i
        if i
        in {
            tpu_custom_call.MemorySpace.HBM,
            tpu_custom_call.MemorySpace.VMEM,
            tpu_custom_call.MemorySpace.SMEM,
        }
        else None
        for i in input_memory_spaces
    )
  return input_memory_spaces, output_memory_spaces


def _resolve_side_effect_type(
    has_side_effects: bool | tpu_core.SideEffectType,
) -> bool | tpu_custom_call.TpuSideEffectType:
  match has_side_effects:
    case bool():
      return has_side_effects
    case tpu_core.SideEffectType.PURE:
      return tpu_custom_call.TpuSideEffectType.PURE
    case tpu_core.SideEffectType.DATAFLOW_SIDE_EFFECTING:
      return tpu_custom_call.TpuSideEffectType.DATAFLOW_SIDE_EFFECTING
    case tpu_core.SideEffectType.SIDE_EFFECTING:
      return tpu_custom_call.TpuSideEffectType.SIDE_EFFECTING
    case _:
      raise ValueError(f"Invalid side effect type: {has_side_effects}")


def _resolve_tiling(
    mosaic_params: tpu_core.CompilerParams,
    kernel_type: tpu_core.CoreType | None,
) -> tpu_custom_call.Tiling | None:
  if kernel_type is tpu_core.CoreType.TC:
    if mosaic_params.use_tc_tiling_on_sc is None:
      return None
    raise ValueError(
        "use_tc_tiling_on_sc= is not supported for TC kernels"
    )

  return (
      tpu_custom_call.Tiling.COMPACT
      if mosaic_params.use_tc_tiling_on_sc
      or mosaic_params.use_tc_tiling_on_sc is None
      else tpu_custom_call.Tiling.SPARSE_CORE
  )


def _jaxpr_kernel_aval_to_mosaic(
    aval: jax_core.AbstractValue,
) -> jax_core.AbstractValue:
  match aval:
    case state_types.AbstractLinVal():
      if dtypes.issubdtype(aval.dtype, jax.numpy.bool_):
        raise NotImplementedError  # TODO(mattjj,sharadmv)
      return aval
    case jax_core.ShapedArray():
      if dtypes.issubdtype(aval.dtype, jax.numpy.bool_):
        return aval.update(dtype=lowering.BOOL_MEMREF_TYPE)
      return aval
    case _:
      raise ValueError(f"Unsupported JAX aval type: {type(aval)}")


def _jax_value_to_mosaic_value(x: jax.Array) -> jax.Array:
  if dtypes.issubdtype(x.dtype, dtypes.bool_):
    return x.astype(lowering.BOOL_MEMREF_TYPE)
  return x


def _lower_to_custom_call(
    ctx: mlir.LoweringRuleContext,
    *in_nodes,
    mosaic_module: ir.Module,
    mosaic_params: tpu_core.CompilerParams,
    kernel_type: tpu_core.CoreType | None,
    num_dynamic_grid_bounds: int,
    input_output_aliases: tuple[tuple[int, int], ...],
    cost_estimate: pallas_core.CostEstimate | None,
    out_avals: tuple[jax_core.AbstractValue, ...],
    effects: jax_core.Effects,
    metadata: frozen_dict.FrozenDict[str, str] | None,
    name: str,
    jax_mesh,
    collective_id_tag: Hashable | None = None,
):
  input_output_aliases = tuple(
      (a[0] + num_dynamic_grid_bounds, a[1]) for a in input_output_aliases
  )

  # Replace in_avals to physical avals.
  # This step is required for mapping logical types to physical types.
  # (e.g. PRNG key -> uint32[2])
  physical_avals = [jax_core.physical_aval(aval) for aval in ctx.avals_in]
  physical_out_avals = [jax_core.physical_aval(aval) for aval in ctx.avals_out]
  ctx = ctx.replace(avals_in=physical_avals, avals_out=physical_out_avals)

  # Booleans are loaded into the kernel as integers.
  def _maybe_cast_inputs(*args):
    args = [_jax_value_to_mosaic_value(x) for x in args]
    return args

  kernel_in_avals = [_jaxpr_kernel_aval_to_mosaic(x) for x in physical_avals]
  kernel_out_avals = [
      _jaxpr_kernel_aval_to_mosaic(x) for x in physical_out_avals
  ]
  cast_ctx = ctx.replace(avals_out=kernel_in_avals)
  in_nodes = mlir.lower_fun(_maybe_cast_inputs)(cast_ctx, *in_nodes)

  # Dynamic grid bounds have to go at the front.
  dynamic_grid_args, args = (
      in_nodes[:num_dynamic_grid_bounds],
      in_nodes[num_dynamic_grid_bounds:],
  )
  kernel_ctx = ctx.replace(avals_in=kernel_in_avals, avals_out=kernel_out_avals)
  input_memory_spaces, output_memory_spaces = _resolve_memory_spaces(
      ctx.avals_in,
      out_avals,
      input_output_aliases=input_output_aliases,
      kernel_type=kernel_type,
  )

  if cost_estimate is not None:
    mosaic_cost_estimate = cast(
        tpu_custom_call.CostEstimate, dataclasses.asdict(cost_estimate)
    )
  else:
    mosaic_cost_estimate = None

  dict_metadata = dict(metadata) if metadata is not None else {}
  del metadata
  if jax_mesh is not None:
    mesh_axes = {
        e.name
        for e in effects
        if isinstance(e, jax_core.NamedAxisEffect)
        # Filter for only device mesh axis name effects
        and e.name in jax_mesh.axis_names
    }
    # Only put mesh axes in metadata if there are any.
    if mesh_axes:
      if "mesh_axes" in dict_metadata:
        raise ValueError("Metadata already contains mesh axes.")
      mesh_axes_list = list(mesh_axes)
      if all(isinstance(a, str) for a in mesh_axes):
        mesh_axes_list = sorted(mesh_axes)  # pyrefly: ignore[bad-specialization]
      dict_metadata["mesh_axes"] = json.dumps(mesh_axes_list)
  out_nodes = mosaic.lower_module_to_custom_call(
      kernel_ctx,
      *dynamic_grid_args,
      *args,
      module=mosaic_module,
      out_type=kernel_out_avals,
      kernel_name=mlir.sanitize_name(name),
      cost_estimate=mosaic_cost_estimate,
      vmem_limit_bytes=mosaic_params.vmem_limit_bytes,
      flags=mosaic_params.flags,
      allow_input_fusion=mosaic_params.allow_input_fusion,
      input_output_aliases=input_output_aliases,
      serialization_format=mosaic_params.serialization_format,
      internal_scratch_in_bytes=mosaic_params.internal_scratch_in_bytes,
      collective_id=mosaic_params.collective_id,
      collective_id_tag=collective_id_tag,
      has_side_effects=_resolve_side_effect_type(
          mosaic_params.has_side_effects
      ),
      output_memory_spaces=output_memory_spaces,
      disable_bounds_checks=mosaic_params.disable_bounds_checks,
      disable_semaphore_checks=mosaic_params.disable_semaphore_checks,
      input_memory_spaces=input_memory_spaces,
      metadata=dict_metadata,
      skip_device_barrier=mosaic_params.skip_device_barrier,
      allow_collective_id_without_custom_barrier=mosaic_params.allow_collective_id_without_custom_barrier,
      shape_invariant_numerics=mosaic_params.shape_invariant_numerics,
      needs_layout_passes=mosaic_params.needs_layout_passes,
      tiling=_resolve_tiling(mosaic_params, kernel_type),
      opt_level=mosaic_params.opt_level,
  )
  _maybe_cast_to_bool = (
      lambda x, aval: x.astype(jax.numpy.bool_)
      if aval.dtype == jax.numpy.bool_
      else x
  )

  def _maybe_cast_outputs(*args):
    args = [_maybe_cast_to_bool(x, aval) for x, aval in zip(args, out_avals)]
    return args

  cast_ctx = ctx.replace(avals_in=kernel_out_avals)
  return mlir.lower_fun(_maybe_cast_outputs)(cast_ctx, *out_nodes)


def _extract_single_barrier_semaphore_tag(
    jaxpr: jax_core.Jaxpr,
) -> Hashable | None:
  tags = {
      eqn.params["tag"]
      for _, eqn in jaxpr_util.all_eqns(jaxpr)
      if eqn.primitive is tpu_primitives.get_barrier_semaphore_p
      and eqn.params.get("tag") is not None
  }
  if len(tags) > 1:
    raise ValueError(
        "Cannot specify multiple distinct barrier semaphore tags within a"
        f" single kernel: {tags}. XLA only supports one collective_id per"
        " kernel."
    )
  return next(iter(tags)) if tags else None


def pallas_call_tpu_lowering_rule(
    ctx: mlir.LoweringRuleContext,
    *in_nodes,
    jaxpr: jax_core.Jaxpr,
    grid_mapping: pallas_core.GridMapping,
    mesh: pallas_core.Mesh | None,
    input_output_aliases: tuple[tuple[int, int], ...],
    debug: bool,
    interpret: bool,
    compiler_params: pallas_core.CompilerParams | None,
    cost_estimate: pallas_core.CostEstimate | None,
    out_avals: tuple[jax_core.AbstractValue, ...],
    metadata: frozen_dict.FrozenDict[str, str] | None,
    name: str | None,
):
  """Lowers a pallas_call to a Mosaic TPU custom call."""
  del interpret  # Unused.

  debug_info = jaxpr.debug_info
  if debug:
    print(f"\nThe kernel jaxpr for pallas_call {debug_info.func_src_info}:")
    print(jaxpr)

  if compiler_params is None:
    mosaic_params = tpu_core.CompilerParams()
  else:
    assert isinstance(compiler_params, tpu_core.CompilerParams)
    mosaic_params = compiler_params

  tag = _extract_single_barrier_semaphore_tag(jaxpr)

  jax_mesh = None
  axis_context = ctx.module_context.axis_context
  if axis_context is not None:
    if isinstance(axis_context, sharding_impls.SPMDAxisContext):
      jax_mesh = axis_context.mesh
  mlir_ctx = ctx.module_context.context
  tpu.register_dialect(mlir_ctx)

  with (
      mlir_ctx,
      ir.Location.unknown(mlir_ctx),
      xla_metadata.clear_xla_metadata(),
  ):
    mosaic_module = lowering.lower_jaxpr_to_pipelined_module(
        ctx,
        grid_mapping,
        jaxpr,
        dimension_semantics=mosaic_params.dimension_semantics,
        kernel_type=tpu_core.CoreType.TC,
        mesh=jax_mesh,
        dynamic_shape_replacement_enabled=pallas_core.dynamic_shapes_export_enabled(),
        fuse_transposed_lhs_in_matmul=mosaic_params.fuse_transposed_lhs_in_matmul,
    )

  if debug:
    pm = passmanager.PassManager.parse("builtin.module(canonicalize)", mlir_ctx)
    pm.run(mosaic_module.operation)
    print(f"\nThe Mosaic module for pallas_call {debug_info.func_src_info}:")
    print(mosaic_module.operation.get_asm(use_name_loc_as_prefix=True))

  return _lower_to_custom_call(
      ctx,
      *in_nodes,
      mosaic_module=mosaic_module,
      mosaic_params=mosaic_params,
      kernel_type=tpu_core.CoreType.TC,
      num_dynamic_grid_bounds=grid_mapping.num_dynamic_grid_bounds,
      input_output_aliases=input_output_aliases,
      cost_estimate=cost_estimate,
      out_avals=out_avals,
      effects=jaxpr.effects,
      metadata=metadata,
      name=name or debug_info.func_name,
      jax_mesh=jax_mesh,
      collective_id_tag=tag,
  )


def _rewrite_jaxpr_for_lowering(
    jaxpr: jax_core.Jaxpr,
    mesh: pallas_core.Mesh,
    all_meshes: tuple[pallas_core.Mesh, ...],
) -> jax_core.Jaxpr:
  # If the jaxpr has any scalar shaped arrays as inputs, they must have come
  # from closed over scalars. We need to rewrite the jaxpr to actually take in
  # a (1,) shaped Ref as an input. On TC and SC Vector Subcore, this Ref will be
  # in SMEM and on the SC scalar subcore, it will be in SMEM.
  core_type = mesh.core_type
  is_scalar_input = [
      isinstance(v.aval, jax_core.ShapedArray) and not v.aval.shape
      for v in jaxpr.invars
  ]
  if not any(is_scalar_input):
    return jaxpr

  new_in_avals = []
  for is_scalar, v in zip(is_scalar_input, jaxpr.invars):
    if not is_scalar:
      new_in_avals.append(v.aval)
      continue
    if core_type == tpu_core.CoreType.TC:
      # TC compiler supports passing Refs in SMEM directly.
      mem_space = tpu_core.MemorySpace.SMEM
    elif core_type in (
        tpu_core.CoreType.SC_SCALAR_SUBCORE,
        tpu_core.CoreType.SC_VECTOR_SUBCORE,
    ):
      # SC compiler doesn't support pass-via-SMEM so we need to pass the Refs
      # in HBM and then copy them into SMEM inside the kernel.
      mem_space = tpu_core.MemorySpace.HBM
    else:
      raise ValueError(f"Unsupported core type: {core_type}")
    new_in_avals.append(
        state.AbstractRef(
            jax_core.ShapedArray((1,), v.aval.dtype), memory_space=mem_space  # pyrefly: ignore[missing-attribute]
        )
    )

  def new_body(*args):
    sync_copy_srcs = []
    sync_copy_dsts = []
    refs = []
    # First pass: collect the destinations and sources for all scalar inputs
    # so we can perform a single, grouped sync_copy for performance.
    for is_scalar, arg in zip(is_scalar_input, args):
      if not is_scalar:
        refs.append(None)
        continue
      if core_type == tpu_core.CoreType.TC:
        refs.append(arg)
      elif core_type == tpu_core.CoreType.SC_SCALAR_SUBCORE:
        smem_ref = jax.empty_ref(
            jax_core.ShapedArray((1,), arg.dtype),
            memory_space=tpu_core.MemorySpace.SMEM,
        )
        sync_copy_srcs.append(arg)
        sync_copy_dsts.append(smem_ref)
        refs.append(smem_ref)
      elif core_type == tpu_core.CoreType.SC_VECTOR_SUBCORE:
        num_lanes = sc_core.get_sparse_core_info().num_lanes
        vmem_ref = jax.empty_ref(
            jax_core.ShapedArray((num_lanes,), arg.dtype),
            memory_space=tpu_core.MemorySpace.VMEM,
        )
        sync_copy_srcs.append(arg)
        sync_copy_dsts.append(vmem_ref.at[:1])
        refs.append(vmem_ref)

    # Perform bulk sync_copy for all scalars before any reads are emitted.
    if sync_copy_srcs:
      helpers.sync_copy(tuple(sync_copy_srcs), tuple(sync_copy_dsts))

    processed_args = []
    # Second pass: load values from the populated SMEM/VMEM references.
    # Doing this in a second pass ensures the read operations appear after
    # the sync_copy operation in the traced jaxpr.
    for is_scalar, arg, ref in zip(is_scalar_input, args, refs):
      if is_scalar:
        assert ref is not None
        if core_type == tpu_core.CoreType.SC_VECTOR_SUBCORE:
          processed_args.append(ref[...][0])
        else:
          processed_args.append(ref[0])
      else:
        processed_args.append(arg)

    return jax_core.eval_jaxpr(jaxpr, jaxpr.constvars, *processed_args)

  with mpmd.mpmd_map_tracing_context(mesh, all_meshes):
    # TODO(necula): use trace_to_jaxpr_nocache instead
    new_jaxpr, _ = pe.trace_to_jaxpr(
        new_body, ft.flatten_args(*new_in_avals),
        jaxpr.debug_info.with_unknown_names(),
    )
  assert not new_jaxpr.consts
  return new_jaxpr


def mpmd_map_tpu_lowering_rule(
    ctx: mlir.LoweringRuleContext,
    *in_nodes,
    meshes,
    jaxprs,
    out_avals,
    input_output_aliases,
    compiler_params,
    interpret,
    debug,
    cost_estimate,
    metadata,
    name,
    external_meshes,
    num_scratch,
):
  del interpret  # Unused.

  if debug:
    for idx, jaxpr in enumerate(jaxprs):
      print(
          f"\nThe kernel jaxpr {idx=} for mpmd_map"
          f" {jaxpr.debug_info.func_src_info}:"
      )
      print(jaxpr)

  if compiler_params is None:
    mosaic_params = tpu_core.CompilerParams()
  else:
    assert isinstance(compiler_params, tpu_core.CompilerParams)
    mosaic_params = compiler_params

  all_tags: set[Hashable] = set()
  for j in jaxprs:
    if (tag := _extract_single_barrier_semaphore_tag(j)) is not None:
      all_tags.add(tag)
  if len(all_tags) > 1:
    raise ValueError(
        "Cannot specify multiple distinct barrier semaphore tags across"
        f" kernels in mpmd_map: {all_tags}. XLA only supports one"
        " collective_id per custom call."
    )
  tag = next(iter(all_tags)) if all_tags else None

  # TODO(slebedev): Check kernel type and raise if it is set.
  if mosaic_params.dimension_semantics is not None:
    raise ValueError(
        "mpmd_map does not support dimension_semantics= in compiler_params="
    )

  mpmd_meshes_map = {
      mesh.core_type: mesh for mesh in [*meshes, *external_meshes]
  }
  jax_mesh = None
  axis_context = ctx.module_context.axis_context
  if axis_context is not None:
    if isinstance(axis_context, sharding_impls.SPMDAxisContext):
      jax_mesh = axis_context.mesh
  mlir_ctx = ctx.module_context.context
  tpu.register_dialect(mlir_ctx)

  some_jaxpr = jaxprs[0]
  is_scalar_input = [
      isinstance(v.aval, jax_core.ShapedArray) and not v.aval.shape
      for v in some_jaxpr.invars
  ]

  with mlir_ctx, ir.Location.unknown(mlir_ctx):
    mosaic_module = ir.Module.create()
    for mesh, jaxpr in zip(meshes, jaxprs, strict=True):
      if not hasattr(mesh, "core_type") or not hasattr(
          mesh, "dimension_semantics"
      ):
        raise ValueError(
            "mpmd_map requires the mesh to define its ``core_type`` and"
            " ``dimension_semantics``"
        )

      match kernel_type := mesh.core_type:
        case tpu_core.CoreType.TC:
          if mpmd_meshes_map is not None and mpmd_meshes_map.keys() != {
              tpu_core.CoreType.TC
          }:
            raise NotImplementedError(
                "mpmd_map does not support TC kernels yet."
            )
          if mesh.num_cores > 1:
            actual_invars = (
                jaxpr.invars[:-num_scratch]
                if num_scratch > 0
                else jaxpr.invars
            )
            if any(
                pallas_core.get_memory_space_aval(v.aval)
                == tpu_core.MemorySpace.VMEM
                for v in actual_invars
            ):
              raise NotImplementedError(
                  "TensorCoreMesh does not support VMEM inputs/outputs when there"
                  " are >1 cores. Use HBM or ANY instead."
              )
        case (
            tpu_core.CoreType.SC_SCALAR_SUBCORE
            | tpu_core.CoreType.SC_VECTOR_SUBCORE
        ):
          if not tpu_info.is_tpu_device():
            raise ValueError(
                "SparseCore kernels are only supported on TPU, but the current"
                f" device is {tpu_info.get_device_kind()}."
            )
          info = tpu_info.get_tpu_info()
          if not info.sparse_core:
            raise ValueError(
                "SparseCore is not available on the current device"
                f" ({info.chip_version}), but the kernel type is set to"
                " SparseCore."
            )
        case _:
          raise ValueError(f"Unsupported kernel type: {kernel_type}")

      if any(is_scalar_input):
        jaxpr = _rewrite_jaxpr_for_lowering(
            jaxpr, mesh, (*meshes, *external_meshes)
        )

      lowering.lower_jaxpr_into_unpipelined_module(
          ctx,
          mosaic_module,
          jaxpr,
          jax_mesh=jax_mesh,
          pallas_mesh=mesh,
          name=mlir.sanitize_name(jaxpr.debug_info.func_name),
          dynamic_shape_replacement_enabled=pallas_core.dynamic_shapes_export_enabled(),
          num_scratch=num_scratch,
          needs_layout_passes=mosaic_params.needs_layout_passes,
      )

  if debug:
    pm = passmanager.PassManager.parse("builtin.module(canonicalize)", mlir_ctx)
    pm.run(mosaic_module.operation)
    print("\nThe Mosaic module for mpmd_map:")
    print(mosaic_module.operation.get_asm(use_name_loc_as_prefix=True))

  if name is None:
    name = "_".join(jaxpr.debug_info.func_name for jaxpr in jaxprs)

  match [*{mesh.core_type for mesh in meshes}]:
    case [kernel_type]:
      pass
    case _:
      # Use a stub ``kernel_type`` if we are lowering multiple kernels.
      # This will hopefully cause a runtime error if ``kernel_type`` is ever
      # accessed.
      kernel_type = None

  def _maybe_expand_scalar_input(is_scalar, in_node, aval):
    expand_ctx = ctx.replace(
        avals_in=[aval],
        avals_out=[aval.update(shape=(1,))]
    )
    if is_scalar:
      return mlir.lower_fun(lambda x: x[None], multiple_results=False)(
          expand_ctx, in_node
      )[0]
    return in_node
  in_nodes = [
      _maybe_expand_scalar_input(is_scalar, in_node, aval)
      for is_scalar, in_node, aval in zip(is_scalar_input, in_nodes, ctx.avals_in)
  ]
  ctx = ctx.replace(
      avals_in=[
          aval.update(shape=(1,))
          if is_scalar
          else aval
          for is_scalar, aval in zip(is_scalar_input, ctx.avals_in)
      ]
  )

  return _lower_to_custom_call(
      ctx,
      *in_nodes,
      mosaic_module=mosaic_module,
      mosaic_params=mosaic_params,
      kernel_type=kernel_type,
      num_dynamic_grid_bounds=0,
      input_output_aliases=tuple(input_output_aliases.items()),
      cost_estimate=cost_estimate,
      out_avals=out_avals,
      effects=set().union(*(jaxpr.effects for jaxpr in jaxprs)),
      metadata=metadata,
      name=name,
      jax_mesh=jax_mesh,
      collective_id_tag=tag,
  )


pallas_core.register_lowering_rule(tpu_core.CompilerParams, pallas_call_tpu_lowering_rule, "tpu")