# SPDX-License-Identifier: MIT
 
import torch
import triton
import triton.language as tl
import aiter.ops.triton.utils.arch_info as arch_info


@triton.jit
def _compute_fp8_scaling_factors(x, fp8_max: tl.constexpr):
    # compute fp8 scaling and descaling factor for a block
    x_amax = tl.max(tl.abs(x))  # NOTE: abs deals with negative values
    x_amax = tl.where(x_amax <= 1e-9, 1e-9, x_amax)
    scale_x = fp8_max / x_amax
    descale_x = x_amax / fp8_max
    return scale_x, descale_x


def _is_fp8(x):
    if x.dtype in {
        torch.float8_e4m3fnuz,
        torch.float8_e4m3fn,
        torch.float8_e5m2,
        torch.float8_e5m2fnuz,
    }:
        if arch_info.is_fp8_avail():
            return True
        else:
            raise RuntimeError("This device does not support fp8")
    else:
        return False
