# SPDX-License-Identifier: MIT
import os
import functools
import subprocess


@functools.lru_cache(maxsize=1)
def get_gfx():
    gfx = os.getenv("GPU_ARCHS", "native")
    if gfx == "native":
        try:
            result = subprocess.run(
                ["rocminfo"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
            )
            output = result.stdout
            for line in output.split("\n"):
                if "gfx" in line.lower():
                    return line.split(":")[-1].strip()
        except Exception as e:
            raise RuntimeError(f"Get GPU arch from rcominfo failed {str(e)}")
    return gfx


@functools.lru_cache(maxsize=1)
def get_cu_num():
    import torch

    device = torch.cuda.current_device()
    cu_num = torch.cuda.get_device_properties(device).multi_processor_count
    return cu_num
