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

import sys

from absl.testing import absltest
from absl.testing import parameterized
import jax
from jax._src import test_util as jtu
from jax._src.pallas.triton import gpu_info

if sys.platform != "win32":
  from jax.experimental.pallas import triton as plgpu

  GpuVersion = plgpu.GpuVersion
else:
  plgpu = None
  GpuVersion = gpu_info.GpuVersion


class GpuInfoTest(jtu.JaxTestCase):

  def setUp(self):
    super().setUp()
    if not plgpu:
      self.skipTest("GPU jaxlib is not available")
    if not jtu.is_device_cuda() and not jtu.is_device_rocm():
      self.skipTest("Needs a GPU device")
    try:
      plgpu.get_gpu_info()
    except ValueError:
      self.skipTest("Unsupported GPU device")

  def test_get_gpu_info(self):
    device = jax.devices()[0]
    info = plgpu.get_gpu_info()
    if info.gpu_version is None:
      self.skipTest("device has a free-form arch with no GpuVersion")
    self.assertIsInstance(info, plgpu.GpuInfo)
    self.assertEqual(info.arch_name, device.compute_capability)
    expected_version = gpu_info.gpu_version_from_device_kind(device.device_kind)
    self.assertEqual(info.gpu_version, expected_version)

    cc = info.compute_capability
    match info.gpu_version:
      case GpuVersion.T4:
        self.assertEqual(cc, 75)
      case GpuVersion.A100 | GpuVersion.A30:
        self.assertEqual(cc, 80)
      case GpuVersion.A10:
        self.assertEqual(cc, 86)
      case GpuVersion.L4 | GpuVersion.L40 | GpuVersion.RTX_4090:
        self.assertEqual(cc, 89)
      case GpuVersion.H100 | GpuVersion.H200 | GpuVersion.GH200:
        self.assertEqual(cc, 90)
      case GpuVersion.B200 | GpuVersion.GB200:
        self.assertEqual(cc, 100)
      case GpuVersion.B300 | GpuVersion.GB300:
        self.assertEqual(cc, 103)
      case GpuVersion.VR200:
        self.assertEqual(cc, 107)
      case GpuVersion.THOR:
        self.assertEqual(cc, 110)
      case (
          GpuVersion.RTX_PRO_4500
          | GpuVersion.RTX_PRO_5000
          | GpuVersion.RTX_PRO_6000
      ):
        self.assertEqual(cc, 120)
      case GpuVersion.GB10:
        self.assertEqual(cc, 121)
      case _:
        self.assertEqual(cc, 0)

  def test_get_gpu_info_given_gpu_version(self):
    info = plgpu.get_gpu_info()
    if info.gpu_version is None:
      self.skipTest("device has a free-form arch with no GpuVersion")
    info_version = plgpu.get_gpu_info_from_version(info.gpu_version)
    self.assertEqual(info, info_version)

  @parameterized.parameters([
      ("NVIDIA A100-SXM4-40GB", GpuVersion.A100),
      ("NVIDIA A100-SXM4-80GB", GpuVersion.A100),
      ("NVIDIA A100-PCIE-40GB", GpuVersion.A100),
      ("NVIDIA A100 80GB PCIe", GpuVersion.A100),
      ("NVIDIA A10 WHATEVER", GpuVersion.A10),
      ("NVIDIA H100 80GB HBM3", GpuVersion.H100),
      ("NVIDIA H100 PCIe", GpuVersion.H100),
      ("NVIDIA H100 NVL", GpuVersion.H100),
      ("NVIDIA VR200", GpuVersion.VR200),
      ("NVIDIA RTX 123", None),
      ("UNKNOWN", None),
      ("gfx908", None),
      ("gfx90a:sramecc+:xnack-", None),
      ("gfx942", None),
      ("gfx950", None),
      ("AMD Instinct MI355X", None),
  ])
  def test_gpu_version_from_device_kind(self, device_kind, expected):
    info = gpu_info.gpu_version_from_device_kind(device_kind)
    self.assertEqual(info, expected)

  @parameterized.parameters([
      "gfx942",
      "gfx950",
      "gfx950:sramecc+:xnack-",
      "gfx1201",
  ])
  def test_gfx_device_kind_used_as_arch(self, device_kind):
    abstract_device = jax.sharding.AbstractDevice(
        device_kind=device_kind, num_cores=None, platform="gpu"
    )
    abstract_mesh = jax.sharding.AbstractMesh(
        (1,),
        ("x",),
        (jax.sharding.AxisType.Explicit,),
        abstract_device=abstract_device,
    )
    with jax.sharding.use_abstract_mesh(abstract_mesh):
      info = plgpu.get_gpu_info()
      self.assertIsNone(info.gpu_version)
      self.assertEqual(info.arch_name, device_kind)
      self.assertTrue(plgpu.is_gpu_device())


if __name__ == "__main__":
  jax.config.parse_flags_with_absl()
  absltest.main(testLoader=jtu.JaxTestLoader())