# SPDX-License-Identifier: Apache-2.0
# Written by Dr. Hicham Badri @Mobius Labs GmbH - 2025

import torch
import triton
import math
import triton.language as tl
from triton.runtime import driver
from ..dtypes import *


@triton.jit
def swizzle_tile_v1(
    pid,
    M,
    N,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
):
    grid_m = tl.cdiv(M, BLOCK_SIZE_M)
    grid_n = tl.cdiv(N, BLOCK_SIZE_N)
    width = GROUP_SIZE_M * grid_n
    group_id = pid // width
    group_size = tl.minimum(grid_m - group_id * GROUP_SIZE_M, GROUP_SIZE_M)
    pid_m = group_id * GROUP_SIZE_M + (pid % group_size)
    pid_n = (pid % width) // group_size
    return pid_m, pid_n


@triton.jit
def swizzle_tile_v2(
    pid,
    M,
    N,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
):
    grid_m = tl.cdiv(M, BLOCK_SIZE_M)
    grid_n = tl.cdiv(N, BLOCK_SIZE_N)
    width = GROUP_SIZE_M * grid_m
    group_id = pid // width
    group_size = tl.minimum(grid_n - group_id * GROUP_SIZE_M, GROUP_SIZE_M)
    pid_n = group_id * GROUP_SIZE_M + (pid % group_size)
    pid_m = (pid % width) // group_size
    return pid_m, pid_n


@triton.jit
def swizzle_tile_v3(
    pid,
    M,
    N,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
):
    pid_m = pid % tl.cdiv(M, BLOCK_SIZE_M)
    pid_n = pid // tl.cdiv(M, BLOCK_SIZE_M)
    grid_m = tl.cdiv(M, BLOCK_SIZE_M)
    grid_n = tl.cdiv(N, BLOCK_SIZE_N)
    return tl.swizzle2d(pid_m, pid_n, grid_m, grid_n, GROUP_SIZE_M)


@triton.jit
def swizzle_tile_persistent(
    tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M: tl.constexpr
):
    group_id = tile_id // num_pid_in_group
    first_pid_m = group_id * GROUP_SIZE_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
    pid_m = first_pid_m + (tile_id % group_size_m)
    pid_n = (tile_id % num_pid_in_group) // group_size_m
    return pid_m, pid_n


swizzle_tile = swizzle_tile_v1


@triton.jit
def linear_tile(
    pid,
    M,
    N,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
):
    pid_m = pid % tl.cdiv(M, BLOCK_SIZE_M)
    pid_n = pid // tl.cdiv(M, BLOCK_SIZE_M)
    return pid_m, pid_n


#################################################################################################################
@triton.jit
def dequantize(
    b,
    scales,
    zeros,
    q_shift,
    meta_dtype,
    unpack_mask,
    elements_per_sample: tl.constexpr,
    W_group_mode: tl.constexpr,
    zero_is_scalar: tl.constexpr,
):
    # Unpack
    if elements_per_sample > 1:
        b = (b >> q_shift) & unpack_mask  # int32 -> int32

    if W_group_mode == 1:  # Shift
        b = b.to(meta_dtype) - zeros

    if W_group_mode == 2:
        b = b.to(meta_dtype) * scales  # Symmetric no shift (Grouped)

    if (
        W_group_mode == 3
    ):  # Asymmetric / Symmetric with shift(Grouped - (b - zeros) * scales)
        # b = (b - zeros) * scales
        if zero_is_scalar:
            b = (b - zeros).to(meta_dtype) * scales
        else:
            b = (b.to(meta_dtype) - zeros) * scales

    if W_group_mode == 4:
        b = tl.fma(
            b.to(meta_dtype), scales, zeros
        )  # Asymmetric (Grouped - b*scales + zeros)

    return b


@triton.jit
def atomic_add_cas(ptr, value, Lock, mask=None, sem: tl.constexpr = "release"):
    while tl.atomic_cas(Lock, 0, 1, sem=sem) == 1:
        pass
    tl.store(ptr, tl.load(ptr, mask=mask) + value, mask=mask)
    tl.debug_barrier()
    tl.atomic_xchg(Lock, 0)


def init_to_zero(name):
    return lambda nargs: nargs[name].zero_()


def next_power_of_2(v):
    return 2 ** int(math.ceil(math.log2(v)))


def is_divisible(dividend, divisor):
    return dividend % divisor == 0


def is_hip():
    return triton.runtime.driver.active.get_current_target().backend == "hip"


def gpu_has_more_shared_memory(
    ref_gpus=["a100", "h100", "h200", "h20", "h800", "b100", "b200"],
):
    gpu_name = torch.cuda.get_device_properties(0).name.lower()
    return True in [g in gpu_name for g in ref_gpus]


def gpu_supports_float16_acc(
    ref_gpus=[
        "5090",
        "5080",
        "5070",
        "5060",
        "4090",
        "4080",
        "4070",
        "4060",
        "3090",
        "3080",
        "3070",
        "3060",
        "2080",
        "2070",
    ],
):
    gpu_name = torch.cuda.get_device_properties(0).name.lower()
    return True in [g in gpu_name for g in ref_gpus]


def gpu_supports_bfloat16_atomicadd():
    # Triton tl.atomic_add doens't support bfloat16 even for Hopper and above.
    # return torch.cuda.get_device_capability()[0] >= 9 #Hopper and above
    return False


NUM_SMS = torch.cuda.get_device_properties("cuda").multi_processor_count


def get_num_SMs(device):
    # return torch.cuda.get_device_properties(device).multi_processor_count
    return NUM_SMS  # cache it to avoid driver ping - should be ok as multi-gpu systems tend to have the same devices


# Only powers of 2
def generate_autotune_lookup_v1(max_m=16384):
    return [
        min(2 ** int(math.ceil(math.log2(M))), max_m) if (M > 0) else 0
        for M in range(max_m + 1)
    ]


# Powers of 2 but also (power of 2 + next power of 2) / divisor,
def generate_autotune_lookup_v2(
    max_m=16384, min_split=32, divisors=[2, 4], mode="next", include_vllm_config=False
):
    lookup = [0] * (max_m + 1)
    autotune_vals = set()

    i = 0
    while (val := 2**i) <= max_m:
        autotune_vals.add(val)
        next_val = 2 ** (i + 1)
        if val >= min_split and next_val <= max_m:
            for d in divisors:
                interpolated = (val + next_val) // d
                autotune_vals.add(interpolated)
        i += 1

    if include_vllm_config:
        autotune_vals.update(
            [
                1,
                2,
                4,
                8,
                16,
                24,
                32,
                40,
                48,
                56,
                64,
                72,
                80,
                88,
                96,
                104,
                112,
                120,
                128,
                136,
                144,
                152,
                160,
                168,
                176,
                184,
                192,
                200,
                208,
                216,
                224,
                232,
                240,
                248,
                256,
                384,
                512,
            ]
        )

    sorted_vals = sorted(autotune_vals)

    for m in range(max_m + 1):
        if mode == "next":
            lookup[m] = min(
                (x for x in sorted_vals if x >= m), default=None
            )  # Next-value
        elif mode == "closest":
            lookup[m] = min(
                sorted_vals, key=lambda x: (abs(x - m), x < m)
            )  # Closest-Value
        else:
            raise Exception("Invalid mode.")
    return lookup


M_MAXVAL = 4096  # 1024, 4096, 16384
M_MAPPING = generate_autotune_lookup_v2(M_MAXVAL, mode="next")


def get_closest_m(M):
    return M_MAPPING[M] if M <= M_MAXVAL else M_MAXVAL


def get_gpu_shared_memory():
    return driver.active.utils.get_device_properties(0).get("max_shared_mem", 0)


###################################################################################
# Cached results to avoid runtime driver calls

IS_HIP = is_hip()
NATIVE_ATOMIC = gpu_supports_bfloat16_atomicadd()
