# Copyright 2026 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.

"""Exposes GPU hardware information."""

from collections.abc import Callable
import dataclasses
import enum
import re

from jax._src import mesh as mesh_lib
from jax._src import util as jax_util
from jax._src.interpreters import pxla


class GpuVersion(enum.Enum):
  """GPU version"""

  # NVIDIA GPUs
  A10 = "NVIDIA A10"
  A30 = "NVIDIA A30"
  A100 = "NVIDIA A100"
  H100 = "NVIDIA H100"
  H200 = "NVIDIA H200"
  GH200 = "NVIDIA GH200"
  B200 = "NVIDIA B200"
  GB200 = "NVIDIA GB200"
  B300 = "NVIDIA B300"
  GB300 = "NVIDIA GB300"
  GB10 = "NVIDIA GB10"
  L4 = "NVIDIA L4"
  L40 = "NVIDIA L40"
  T4 = "Tesla T4"
  RTX_4090 = "NVIDIA GeForce RTX 4090"
  RTX_PRO_4500 = "NVIDIA RTX PRO 4500"
  RTX_PRO_5000 = "NVIDIA RTX PRO 5000"
  RTX_PRO_6000 = "NVIDIA RTX PRO 6000"
  THOR = "NVIDIA Thor"
  VR200 = "NVIDIA VR200"

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


# Longer names first so e.g. "NVIDIA A100" wins over "NVIDIA A10".
_GPU_VERSION_RE = re.compile(
    r"\b("
    + "|".join(
        re.escape(e.value)
        for e in sorted(GpuVersion, key=lambda e: len(e.value), reverse=True)
    )
    + r")\b"
)


def gpu_version_from_device_kind(device_kind: str) -> GpuVersion | None:
  if m := _GPU_VERSION_RE.search(device_kind):
    return GpuVersion(m.group(1))
  return None


@dataclasses.dataclass(frozen=True, kw_only=True)
class GpuInfo:
  """GPU hardware information"""

  # None when device_kind does not map to a known GpuVersion but arch_name is
  # still a valid backend target (e.g. a ROCm GPU with a marketing device_kind).
  gpu_version: GpuVersion | None
  arch_name: str
  compute_capability: int


def is_gpu_device() -> bool:
  return get_device_platform() == "gpu"


registry: dict[str, Callable[[], GpuInfo]] = {}


def _get_gpu_info_impl(gpu_version: GpuVersion) -> GpuInfo:
  """Returns the GPU hardware info for the given its version.

  Args:
    gpu_version: The GPU version.
  """
  # https://developer.nvidia.com/cuda/gpus
  match gpu_version:
    case GpuVersion.T4:
      return GpuInfo(
          gpu_version=gpu_version,
          arch_name="7.5",
          compute_capability=75,
      )
    case GpuVersion.A100 | GpuVersion.A30:
      return GpuInfo(
          gpu_version=gpu_version,
          arch_name="8.0",
          compute_capability=80,
      )
    case GpuVersion.A10:
      return GpuInfo(
          gpu_version=gpu_version,
          arch_name="8.6",
          compute_capability=86,
      )
    case GpuVersion.L4 | GpuVersion.L40 | GpuVersion.RTX_4090:
      return GpuInfo(
          gpu_version=gpu_version,
          arch_name="8.9",
          compute_capability=89,
      )
    case GpuVersion.H100 | GpuVersion.H200 | GpuVersion.GH200:
      return GpuInfo(
          gpu_version=gpu_version,
          arch_name="9.0",
          compute_capability=90,
      )
    case GpuVersion.B200 | GpuVersion.GB200:
      return GpuInfo(
          gpu_version=gpu_version,
          arch_name="10.0",
          compute_capability=100,
      )
    case GpuVersion.B300 | GpuVersion.GB300:
      return GpuInfo(
          gpu_version=gpu_version,
          arch_name="10.3",
          compute_capability=103,
      )
    case GpuVersion.VR200:
      return GpuInfo(
          gpu_version=gpu_version,
          arch_name="10.7",
          compute_capability=107,
      )
    case (
        GpuVersion.RTX_PRO_4500
        | GpuVersion.RTX_PRO_5000
        | GpuVersion.RTX_PRO_6000
    ):
      return GpuInfo(
          gpu_version=gpu_version,
          arch_name="12.0",
          compute_capability=120,
      )
    case GpuVersion.GB10:
      return GpuInfo(
          gpu_version=gpu_version,
          arch_name="12.1",
          compute_capability=121,
      )
    case GpuVersion.THOR:
      return GpuInfo(
          gpu_version=gpu_version,
          arch_name="11.0",
          compute_capability=110,
      )
    case _:
      raise ValueError(f"Unsupported GPU version: {gpu_version}")


@jax_util.cache(trace_context_in_key=True)
def get_gpu_info() -> GpuInfo:
  """Returns the GPU hardware info for the current device."""
  device_kind = get_device_kind()
  gpu_version = gpu_version_from_device_kind(device_kind)
  if gpu_version is not None:
    return _get_gpu_info_impl(gpu_version)

  if get_device_platform() == "gpu" and device_kind.startswith("gfx"):
    return GpuInfo(
        gpu_version=None, arch_name=device_kind, compute_capability=0
    )

  # generic NVIDIA GPU device without GpuVersion entry
  if get_device_platform() == "gpu" and device_kind.startswith("NVIDIA"):
    gpu_arch_name = _get_device_arch_name(device_kind)
    gpu_compute_capability = int(gpu_arch_name.replace(".", ""))
    return GpuInfo(
        gpu_version=None,
        arch_name=gpu_arch_name,
        compute_capability=gpu_compute_capability,
    )

  if device_kind in registry:
    return registry[device_kind]()
  raise ValueError(f"Unsupported GPU device kind: {device_kind}")


@jax_util.cache(trace_context_in_key=True)
def get_gpu_info_from_version(gpu_version: GpuVersion) -> GpuInfo:
  """Returns the GPU hardware info for the given GPU version.

  Args:
    gpu_version: The GPU version.
  """
  return _get_gpu_info_impl(gpu_version)


def _get_device_arch_name(device_kind: str) -> str:
  concrete_device = pxla.get_default_device()

  if device_kind != concrete_device.device_kind:
    raise ValueError(
        f"Cannot infer architecture of an unknown abstract GPU {device_kind!r}"
    )

  return concrete_device.compute_capability


def get_device_kind() -> str:
  device = pxla.get_default_device()
  abstract_device = mesh_lib.get_abstract_mesh().abstract_device
  if abstract_device is not None:
    if abstract_device.device_kind.startswith("gfx"):
      return abstract_device.device_kind
    # A kind that differs from the concrete device is a deliberate AOT target;
    # if it matches, the abstract mesh just mirrors the local device.
    if abstract_device.device_kind != device.device_kind:
      return abstract_device.device_kind
  cc = getattr(device, "compute_capability", None)
  if isinstance(cc, str) and cc.startswith("gfx"):
    return cc
  return device.device_kind


_GPU_PLATFORMS = ("gpu", "rocm", "cuda")


def get_device_platform() -> str:
  if abstract_device := mesh_lib.get_abstract_mesh().abstract_device:
    platform = abstract_device.platform
  else:
    platform = pxla.get_default_device().platform
  return "gpu" if platform in _GPU_PLATFORMS else platform