import torch
import triton

# For now, there is 1-to-1 correspondence between arch and device
_ARCH_TO_DEVICE = {
    "gfx928": "K100_AI",
    "gfx936": "BW200",
    "gfx938": "BW200B",
    "gfx92a": "K200_AI",
}


def get_arch():
    return triton.runtime.driver.active.get_current_target().arch


def get_device():
    return _ARCH_TO_DEVICE.get(get_arch(), "Unknown")


def is_fp4_avail():
    return get_arch() in ("gfx946")


def is_fp8_avail():
    return get_arch() in ("gfx938", "gfx92a")


def is_mls_avail():
    return get_arch() in ("gfx938", "gfx92a")


def get_fp8_dtypes():
    e5m2_dtype = torch.float8_e5m2
    e4m3_dtype = torch.float8_e4m3fn
    return e5m2_dtype, e4m3_dtype


def get_fp8_e4m3_dtype():
    e4m3_dtype = torch.float8_e4m3fn
    return e4m3_dtype


def get_num_sms():
    # Returns the Compute Unit count of the current device
    current_device_index = torch.cuda.current_device()
    current_device = torch.cuda.get_device_properties(current_device_index)
    num_sms = current_device.multi_processor_count
    return num_sms
