# 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 TPU-specific Pallas abstractions."""

from __future__ import annotations

import collections
from collections.abc import Mapping, Sequence
import dataclasses
import enum
from typing import Any, Literal

import jax
from jax._src import core as jax_core
from jax._src import linear_util as lu
from jax._src import state
from jax._src import util
from jax._src.frozen_dict import FrozenDict
from jax._src.interpreters import partial_eval as pe
from jax._src.pallas import core as pallas_core
import jax.numpy as jnp
import numpy as np


map, unsafe_map = util.safe_map, map
zip, unsafe_zip = util.safe_zip, zip

no_block_spec = pallas_core.no_block_spec
_out_shape_to_aval_mapping = pallas_core._out_shape_to_aval_mapping


class CoreType(enum.Enum):
  TC = "tc"
  SC_SCALAR_SUBCORE = "sc_scalar_subcore"
  SC_VECTOR_SUBCORE = "sc_vector_subcore"

  def __str__(self) -> str:
    return self.value

  def __repr__(self) -> str:
    return self.name


class GridDimensionSemantics(enum.Enum):
  PARALLEL = "parallel"
  CORE_PARALLEL = "core_parallel"
  SUBCORE_PARALLEL = "subcore_parallel"
  ARBITRARY = "arbitrary"


PARALLEL = GridDimensionSemantics.PARALLEL
CORE_PARALLEL = GridDimensionSemantics.CORE_PARALLEL
SUBCORE_PARALLEL = GridDimensionSemantics.SUBCORE_PARALLEL
ARBITRARY = GridDimensionSemantics.ARBITRARY


LiteralDimensionSemantics = Literal[
    "parallel", "core_parallel", "subcore_parallel", "arbitrary"
]
DimensionSemantics = LiteralDimensionSemantics | GridDimensionSemantics


class SideEffectType(enum.Enum):
  # No side effects, can be deduplicated / removed if unused.
  PURE = "pure"
  # Cannot be deduplicated, but can be removed if unused.
  DATAFLOW_SIDE_EFFECTING = "dataflow_side_effecting"
  # Cannot be deduplicated or removed.
  SIDE_EFFECTING = "side_effecting"


@dataclasses.dataclass(frozen=True)
class CompilerParams:
  """Mosaic TPU compiler parameters.

  Attributes:
    dimension_semantics: A list of dimension semantics for each grid dimension
      of the kernel. Either "parallel" for dimensions that can execute in any
      order, or "arbitrary" for dimensions that must be executed sequentially.
    allow_input_fusion: A list of booleans indicating whether input fusion is
      allowed for each argument.
    vmem_limit_bytes: Overrides the default VMEM limit for a kernel. Note that
      this must be used in conjunction with the
      --xla_tpu_scoped_vmem_limit_kib=N flag with N*1kib > vmem_limit_bytes.
    collective_id: Indicates which barrier semaphore to use for the kernel. Note
      that using the same collective_id does not guarantee that the same barrier
      semaphore will be allocated between kernels.
    has_side_effects: Set to True to prevent kernel being CSEd by XLA.
    flags: A dictionary of command line flags for the kernel.
    internal_scratch_in_bytes: The size of the internal scratch space used by
      Mosaic.
    serialization_format: The serialization format for the kernel body.
    kernel_type: Specify if the kernel is meant to run on TensorCore or one of
      the SparseCores
    disable_bounds_checks: Disable bounds checks in the kernel.
    disable_semaphore_checks: Disable semaphore checks in the kernel.
    skip_device_barrier: Skip the default device barrier for the kernel.
    allow_collective_id_without_custom_barrier: Allow the use of collective_id
      without a custom barrier.
    use_tc_tiling_on_sc: Use TensorCore tiling for SparseCore. This flag is only
      used for ``SC_*_SUBCORE`` kernels.
  """

  dimension_semantics: tuple[DimensionSemantics, ...] | None = None
  allow_input_fusion: tuple[bool, ...] | None = None
  vmem_limit_bytes: int | None = None
  collective_id: int | None = None
  has_side_effects: bool | SideEffectType = False
  flags: dict[str, Any] | None = None
  internal_scratch_in_bytes: int | None = None
  serialization_format: int = 1
  kernel_type: CoreType = CoreType.TC
  disable_bounds_checks: bool = False
  disable_semaphore_checks: bool = False
  skip_device_barrier: bool = False
  allow_collective_id_without_custom_barrier: bool = False
  shape_invariant_numerics: bool = True
  use_tc_tiling_on_sc: bool | None = None

  def __init__(
      self,
      dimension_semantics: Sequence[DimensionSemantics] | None = None,
      allow_input_fusion: Sequence[bool] | None = None,
      vmem_limit_bytes: int | None = None,
      collective_id: int | None = None,
      has_side_effects: bool | SideEffectType = False,
      flags: Mapping[str, Any] | None = None,
      internal_scratch_in_bytes: int | None = None,
      serialization_format: int = 1,
      kernel_type: CoreType = CoreType.TC,
      disable_bounds_checks: bool = False,
      disable_semaphore_checks: bool = False,
      skip_device_barrier: bool = False,
      allow_collective_id_without_custom_barrier: bool = False,
      shape_invariant_numerics: bool = True,
      use_tc_tiling_on_sc: bool | None = None,
  ):
    object.__setattr__(
        self,
        "dimension_semantics",
        None if dimension_semantics is None else tuple(dimension_semantics),
    )
    object.__setattr__(
        self,
        "allow_input_fusion",
        None if allow_input_fusion is None else tuple(allow_input_fusion),
    )
    object.__setattr__(self, "vmem_limit_bytes", vmem_limit_bytes)
    object.__setattr__(self, "collective_id", collective_id)
    object.__setattr__(self, "has_side_effects", has_side_effects)
    object.__setattr__(
        self, "flags", None if flags is None else FrozenDict(flags)
    )
    object.__setattr__(
        self, "internal_scratch_in_bytes", internal_scratch_in_bytes
    )
    object.__setattr__(self, "serialization_format", serialization_format)
    object.__setattr__(self, "kernel_type", kernel_type)
    object.__setattr__(self, "disable_bounds_checks", disable_bounds_checks)
    object.__setattr__(
        self, "disable_semaphore_checks", disable_semaphore_checks
    )
    object.__setattr__(self, "skip_device_barrier", skip_device_barrier)
    object.__setattr__(
        self,
        "allow_collective_id_without_custom_barrier",
        allow_collective_id_without_custom_barrier,
    )
    object.__setattr__(
        self, "shape_invariant_numerics", shape_invariant_numerics
    )
    object.__setattr__(self, "use_tc_tiling_on_sc", use_tc_tiling_on_sc)

  # Replace is a method, not a field.
  replace = dataclasses.replace


class MemoryRef(pallas_core.MemoryRef):

  def __matmul__(self, other, /):
    if not isinstance(other, CoreType):
      return NotImplemented
    return dataclasses.replace(self, memory_space=self.memory_space @ other)


class MemorySpace(enum.Enum):
  VMEM = "vmem"
  VMEM_SHARED = "vmem_shared"
  SMEM = "smem"
  CMEM = "cmem"
  SEMAPHORE = "semaphore_mem"
  HBM = "hbm"
  HOST = "host"

  def __str__(self) -> str:
    return self.value

  def __repr__(self) -> str:
    return self.name

  def from_type(self, ty):
    return MemoryRef(ty, memory_space=self)

  def __call__(self, shape: Sequence[int], dtype: jnp.dtype[Any]):
    # A convenience function for constructing MemoryRef types of ShapedArrays.
    return self.from_type(jax_core.ShapedArray(tuple(shape), dtype))

  def __matmul__(self, other, /):
    if not isinstance(other, CoreType):
      return NotImplemented
    return CoreMemorySpace(self, other)


@dataclasses.dataclass(frozen=True)
class CoreMemorySpace:
  """A memory space tied to a specific core type."""

  memory_space: MemorySpace
  core_type: CoreType

  def __post_init__(self):
    match self.memory_space, self.core_type:
      case MemorySpace.VMEM, CoreType.TC | CoreType.SC_VECTOR_SUBCORE:
        ...
      case MemorySpace.SMEM | MemorySpace.SEMAPHORE, _:
        ...
      case MemorySpace.CMEM, CoreType.TC:
        ...
      case _, _:
        raise ValueError(
            "Unsupported core memory space:"
            f" {self.memory_space, self.core_type}"
        )

  def __call__(self, shape: Sequence[int], dtype: jnp.dtype[Any]):
    return MemoryRef(jax_core.ShapedArray(tuple(shape), dtype), self)

  def __str__(self) -> str:
    return f"{self.memory_space}@{self.core_type}"

  def __repr__(self) -> str:
    return f"{self.memory_space!r}@{self.core_type!r}"


class dma_semaphore(pallas_core.semaphore_dtype):
  pass


class DMASemaphore(pallas_core.AbstractSemaphoreTy):
  type = dma_semaphore
  name = "dma_sem"


class SemaphoreType(enum.Enum):
  REGULAR = "regular"
  DMA = "dma"
  BARRIER = "barrier"

  @property
  def dtype(self) -> Any:
    if self == SemaphoreType.DMA:
      return DMASemaphore()
    elif self == SemaphoreType.BARRIER:
      return pallas_core.BarrierSemaphore()
    else:
      return pallas_core.Semaphore()

  def __call__(self, shape: tuple[int, ...]):
    return MemoryRef(jax_core.ShapedArray(shape, self.dtype), MemorySpace.SEMAPHORE)

  def __matmul__(self, other, /):
    if not isinstance(other, CoreType):
      return NotImplemented
    return CoreMemorySpace(MemorySpace.SEMAPHORE, other)((), self.dtype)

  def get_array_aval(self) -> jax_core.ShapedArray:
    return self(()).get_array_aval()

  def get_ref_aval(self) -> state.AbstractRef:
    return self(()).get_ref_aval()


@dataclasses.dataclass(frozen=True)
class AbstractSemaphore(jax_core.AbstractValue):
  sem_type: SemaphoreType


@dataclasses.dataclass(init=False, kw_only=True, unsafe_hash=True)
class PrefetchScalarGridSpec(pallas_core.GridSpec):
  num_scalar_prefetch: int

  def __init__(
      self,
      num_scalar_prefetch: int,
      grid: pallas_core.Grid = (),
      in_specs: pallas_core.BlockSpecTree = no_block_spec,
      out_specs: pallas_core.BlockSpecTree = no_block_spec,
      scratch_shapes: pallas_core.ScratchShapeTree = (),
  ):
    super().__init__(grid, in_specs, out_specs, scratch_shapes)
    self.num_scalar_prefetch = num_scalar_prefetch
    self.scratch_shapes = tuple(scratch_shapes)

  def _make_scalar_ref_aval(self, aval):
    return state.AbstractRef(
        jax_core.ShapedArray(aval.shape, aval.dtype), MemorySpace.SMEM
    )


@dataclasses.dataclass(frozen=True)
class TensorCore:
  id: int


@dataclasses.dataclass(frozen=True)
class TensorCoreMesh(pallas_core.Mesh):
  """A mesh of TensorCores."""

  devices: np.ndarray
  axis_names: Sequence[str]

  def __init__(self, devices: np.ndarray, axis_names: Sequence[str]):
    devices = np.copy(devices)
    devices.setflags(write=False)
    object.__setattr__(self, "devices", devices)
    object.__setattr__(self, "axis_names", tuple(axis_names))

  def __hash__(self) -> int:
    return hash(
        (self.devices.shape, tuple(np.ravel(self.devices)), self.axis_names)
    )

  @property
  def kernel_type(self) -> CoreType:
    return CoreType.TC

  @property
  def default_memory_space(self) -> pallas_core.MemorySpace:
    return pallas_core.MemorySpace.ANY

  @property
  def shape(self):
    return collections.OrderedDict(zip(self.axis_names, self.devices.shape))

  @property
  def dimension_semantics(self) -> Sequence[str]:
    return ["parallel"]

  def discharges_effect(self, effect: jax_core.Effect) -> Literal[False]:
    del effect
    return False

  def check_is_compatible_with(self, other_mesh):
    if isinstance(other_mesh, TensorCoreMesh) and self != other_mesh:
      raise ValueError("You can't use two different TensorCoreMeshes.")
    # TODO: Add support for mpmd with SparseCore meshes.
    return super().check_is_compatible_with(other_mesh)

def create_tensorcore_mesh(
    axis_name: str,
    devices: Sequence[jax.Device] | None = None,
    num_cores: int | None = None,
) -> TensorCoreMesh:
  if devices is not None and num_cores is not None:
    raise ValueError("cannot specify both devices and num_cores")
  if num_cores is None:
    if devices is None:
      abstract_device = jax.sharding.get_abstract_mesh().abstract_device
      if abstract_device is None:
        devices = [jax.devices()[0]]
      else:
        devices = [abstract_device]
    num_cores = devices[0].num_cores
  return TensorCoreMesh(
      np.array([TensorCore(i) for i in range(num_cores)]),
      [axis_name],
  )


def pass_scalars_as_refs(
    jaxpr: jax_core.Jaxpr,
    args: Sequence[Any],
    in_avals: Sequence[jax_core.AbstractValue],
    out_avals: Sequence[jax_core.AbstractValue],
    mesh,
    copy_to_smem: bool = False,
) -> tuple[
    jax_core.Jaxpr,
    tuple[Any, ...],
    tuple[jax_core.AbstractValue, ...],
    tuple[jax_core.AbstractValue, ...],
    tuple[bool, ...],
]:
  """Rewrites a jaxpr to pass scalars as refs instead of values."""
  def allowed_aval(aval):
    if isinstance(aval, state.AbstractRef):
      return True
    if isinstance(aval, jax_core.ShapedArray):
      # Only scalars are allowed.
      return not aval.shape
    return False

  assert all(allowed_aval(v.aval) for v in jaxpr.constvars + jaxpr.invars)

  is_scalar_const = [
      isinstance(v.aval, jax_core.ShapedArray) and not v.aval.shape
      for v in jaxpr.constvars
  ]
  if not any(is_scalar_const):
    return (
        jaxpr,
        tuple(in_avals),
        tuple(out_avals),
        tuple(args),
        tuple(is_scalar_const),
    )
  non_scalar_const_avals, scalar_const_avals = util.partition_list(
      is_scalar_const,
      [v.aval for v in jaxpr.constvars],
  )
  non_scalar_consts, scalar_consts = util.partition_list(
      is_scalar_const, args
  )
  if copy_to_smem:
    smem_alloc = [
        state.AbstractRef(
            jax_core.ShapedArray((1,), aval.dtype),  # pyrefly: ignore[missing-attribute]
            memory_space=MemorySpace.SMEM,
        )
        for aval in scalar_const_avals
    ]
  else:
    smem_alloc = []

  # Rewrite body jaxpr to take in scalar values as Refs.
  def new_body(*args):
    scalar_const_refs, non_scalar_const_refs, args = util.split_list(
        args, [len(scalar_consts), len(non_scalar_consts)]
    )
    if copy_to_smem:
      smem, args = util.split_list(args, [len(smem_alloc)])
      assert len(smem) == len(scalar_const_refs)
      from jax._src.pallas.mosaic.helpers import sync_copy

      sync_copy(scalar_const_refs, smem)
    else:
      smem = scalar_const_refs
    scalar_const_values = [s[0] for s in smem]
    new_consts = util.merge_lists(
        is_scalar_const, non_scalar_const_refs, scalar_const_values
    )
    return jax_core.eval_jaxpr(jaxpr, new_consts, *args)

  # TODO(sharadmv): Remove this once Mosaic support passing scalars as values.
  scalar_const_trace_avals = [
      state.AbstractRef(
          jax_core.ShapedArray((1,), aval.dtype),  # pyrefly: ignore[missing-attribute]
          memory_space=MemorySpace.HBM if copy_to_smem else MemorySpace.SMEM,
      )
      for aval in scalar_const_avals
  ]
  new_trace_avals = [
      *scalar_const_trace_avals,
      *non_scalar_const_avals,
      *smem_alloc,
      *[v.aval for v in jaxpr.invars],
  ]
  with (
      pallas_core.tracing_grid_env(
          tuple(mesh.shape.values()), mapped_dims=()
      ),
      jax_core.extend_axis_env_nd(mesh.shape.items()),
  ):
    new_jaxpr, _, _ = pe.trace_to_jaxpr_dynamic(
        lu.wrap_init(
            new_body, debug_info=jaxpr.debug_info.with_unknown_names()
        ),
        new_trace_avals,
    )
  jaxpr = new_jaxpr.replace(
      constvars=new_jaxpr.invars[: len(jaxpr.constvars)],
      invars=new_jaxpr.invars[len(jaxpr.constvars) :],
  )
  args = [
      *[a[None] for a in scalar_consts],
      *non_scalar_consts,
  ]
  in_avals, out_avals, _ = util.split_list(
      new_trace_avals, [len(in_avals), len(out_avals)]
  )
  return jaxpr, tuple(in_avals), tuple(out_avals), tuple(args), tuple(is_scalar_const)


def _tensorcore_mesh_discharge_rule(
    in_avals,
    out_avals,
    *args,
    mesh,
    jaxpr,
    compiler_params: Any | None,
    interpret: Any,
    debug: bool,
    cost_estimate: pallas_core.CostEstimate | None,
    name: str,
    metadata: FrozenDict[str, str] | None,
):
  assert isinstance(mesh, TensorCoreMesh)
  if compiler_params and not isinstance(compiler_params, CompilerParams):
    raise ValueError("compiler_params must be a pltpu.CompilerParams")
  if not compiler_params:
    compiler_params = CompilerParams()
  if len(mesh.shape) > 1:
    raise NotImplementedError("Mesh must be 1D")
  if compiler_params.dimension_semantics is not None:
    raise ValueError("dimension_semantics must be None for TensorCoreMesh")
  num_cores = len(mesh.devices)
  if num_cores > 1:
    # Since each core will have its own VMEM, we currently disallow VMEM inputs
    # and outputs since other ops might not agree on how they are sharded across
    # cores by the (core-mapped) kernel.
    if any(
        pallas_core.get_memory_space_aval(aval) == MemorySpace.VMEM
        for aval in in_avals
    ):
      raise NotImplementedError(
          "TensorCoreMesh does not support VMEM inputs/outputs when there are"
          " >1 cores. Use HBM or ANY instead."
      )
  jaxpr, in_avals, out_avals, args, is_scalar_const = pass_scalars_as_refs(
      jaxpr, args, in_avals, out_avals, mesh
  )
  refs_out, out = pallas_core.default_mesh_discharge_rule(
      in_avals,
      out_avals,
      *args,
      jaxpr=jaxpr,
      mesh=mesh,
      compiler_params=compiler_params,
      debug=debug,
      interpret=interpret,
      cost_estimate=cost_estimate,
      name=name,
      metadata=metadata,
  )
  refs_out = [
      a if not is_scalar else None
      for is_scalar, a in zip(is_scalar_const, refs_out)
  ]
  return refs_out, out


pallas_core._core_map_mesh_rules[TensorCoreMesh] = (
    _tensorcore_mesh_discharge_rule
)


def _convert_semaphore_type_to_aval(
    out_shape: SemaphoreType,
) -> jax_core.AbstractValue:
  return out_shape.get_array_aval()


pallas_core._out_shape_to_aval_mapping[SemaphoreType] = (
    _convert_semaphore_type_to_aval
)


def memory_space_to_tpu_memory_space(
    memory_space: (
        MemorySpace | pallas_core.MemorySpace | CoreMemorySpace | None
    ),
    kernel_type: CoreType,
) -> MemorySpace | pallas_core.MemorySpace | CoreMemorySpace:
  match memory_space:
    case None:
      match kernel_type:
        case CoreType.TC | CoreType.SC_VECTOR_SUBCORE:
          return MemorySpace.VMEM
        case CoreType.SC_SCALAR_SUBCORE:
          return MemorySpace.SMEM
        case _:
          raise ValueError(f"Unsupported kernel type: {kernel_type}")
    case pallas_core.MemorySpace.ANY:
      return pallas_core.MemorySpace.ANY
    case pallas_core.MemorySpace.HOST:
      return MemorySpace.HOST
    case (
        pallas_core.MemorySpace.ERROR
        | pallas_core.MemorySpace.INDEX
        | pallas_core.MemorySpace.KEY
    ):
      return MemorySpace.SMEM
    case CoreMemorySpace():
      return (
          memory_space.memory_space
          if memory_space.core_type is kernel_type
          else memory_space
      )
    case MemorySpace():
      return memory_space
    case _:
      raise ValueError(f"Invalid memory space: {memory_space!r}")
