import pytest
import torch
import transformer_engine as te
import transformer_engine_torch as tex
import warnings
try:
    import lightop
    enable_lightop = True
except ImportError:
    enable_lightop = False
from transformer_engine.pytorch import get_device_compute_capability
from transformer_engine.pytorch.constants import TE_DType
from transformer_engine.pytorch.fp8 import (FP8GlobalStateManager, blockwise_fp8_block_len)
from transformer_engine.pytorch.tensor.float8_blockwise_tensor import (
    Float8BlockQuantizer,
    Float8BlockwiseQTensor,
)
from transformer_engine.pytorch.triton.blockwise_int8_gemm_nt import w8a8_block_int8_matmul
from transformer_engine.pytorch.triton.blockwise_int8_gemm_nt_wgrad import w8a8_block_int8_matmul_wgrad

from references.blockwise_quantizer_reference import CuBLASScaleMunger
from references.blockwise_fp8_gemm_reference import CuBLASRefBlockwiseGemm


def fp8_blockwise_gemm_supported() -> bool:
    supported, _ = FP8GlobalStateManager.is_fp8_block_scaling_available()
    return supported


def cublas_gemm_fp8_blockwise_case_fw(
    x_dtype,
    w_dtype,
    out_dtype,
    M,
    K,
    N,
    noise_type,
    x_magnitude,
    w_magnitude,
    accumulate,
    use_split_accumulator,
    is_x_1d_scaled,
    is_w_1d_scaled,
    *,
    x_columnwise: bool = False,
    w_columnwise: bool = False,
    use_bias: bool = False,
    use_gelu: bool = False,
    use_grad: bool = False,
    atol: float = 5e-1,
    rtol: float = 5e-1
):
    if x_dtype == torch.float8_e5m2 and w_dtype == torch.float8_e5m2:
        pytest.skip("FP8 GEMM doesn't support both a and b types being torch.float8_e5m2")
    if not (is_x_1d_scaled or is_w_1d_scaled):
        pytest.skip("FP8 GEMM doesn't support 2dimensional qtile by 2dimensional qtile")
    if not fp8_blockwise_gemm_supported():
        pytest.skip("CUDA version does not support blockwise FP8 gemm.")
    # Setup device and random seed
    device = "cuda"
    seed = 0
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)

    x_shape = (K, M) if x_columnwise else (M, K)
    w_shape = (K, N) if w_columnwise else (N, K)
    # generate random input and weight
    if noise_type == "uniform":
        x = torch.rand(x_shape, dtype=torch.bfloat16, device=device) * x_magnitude * 2 - x_magnitude
        w = torch.rand(w_shape, dtype=torch.bfloat16, device=device) * w_magnitude * 2 - w_magnitude
    elif noise_type == "normal":
        x = torch.randn(x_shape, dtype=torch.bfloat16, device=device) * x_magnitude
        w = torch.randn(w_shape, dtype=torch.bfloat16, device=device) * w_magnitude
    else:
        assert False

    bf16_out = torch.matmul(x, w.t())

    # print(f"x.shape: {x.shape}, w.shape: {w.shape}")
    # print("bf16 gemm output: ", bf16_out)
    # print("bf16 gemm output shape: ", bf16_out.shape)

    # Setup out tensor if accumulate is True
    if accumulate:
        out = torch.randn((M, N), dtype=out_dtype, device=device) * x_magnitude
    else:
        out = None

    assert not (use_bias and use_grad), "Bias grad not supported by GEMM"
    # Set quantize_op and quantization parameters
    block_len = blockwise_fp8_block_len
    x_quant_tile_shape = (1, block_len) if is_x_1d_scaled else (block_len, block_len)
    w_quant_tile_shape = (1, block_len) if is_w_1d_scaled else (block_len, block_len)
    x_block_scaling_dim = 1 if is_x_1d_scaled else 2
    w_block_scaling_dim = 1 if is_w_1d_scaled else 2
    x_te_dtype = TE_DType[x_dtype]
    w_te_dtype = TE_DType[w_dtype]
    x_quantizer = Float8BlockQuantizer(
        fp8_dtype=x_te_dtype,
        rowwise=True,
        columnwise=True,
        amax_epsilon=0.0,
        force_pow_2_scales=True,
        block_scaling_dim=x_block_scaling_dim,
    )
    w_quantizer = Float8BlockQuantizer(
        fp8_dtype=w_te_dtype,
        rowwise=True,
        columnwise=True,
        amax_epsilon=0.0,
        force_pow_2_scales=True,
        block_scaling_dim=w_block_scaling_dim,
    )

    # Quantize x and w
    qx = x_quantizer.make_empty(x_shape, dtype=x_dtype, device=device, requires_grad=False)
    qx = x_quantizer.update_quantized(x, qx)
    qw = w_quantizer.make_empty(w_shape, dtype=w_dtype, device=device, requires_grad=False)
    qw = w_quantizer.update_quantized(w, qw)

    if not use_bias:
        bias = None
    else:
        bias = torch.randn((1, N), dtype=torch.bfloat16, device=device)

    # Reference GEMM
    ref_gemm = CuBLASRefBlockwiseGemm()
    scale_decoder = CuBLASScaleMunger()
    qx_data = (
        qx._columnwise_data.view(dtype=x_dtype)
        if x_columnwise
        else qx._rowwise_data.view(dtype=x_dtype)
    )
    qw_data = (
        qw._columnwise_data.view(dtype=w_dtype)
        if w_columnwise
        else qw._rowwise_data.view(dtype=w_dtype)
    )
    ref_scales_x = qx._columnwise_scale_inv if x_columnwise else qx._rowwise_scale_inv
    ref_scales_w = qw._columnwise_scale_inv if w_columnwise else qw._rowwise_scale_inv

    # print(f"qx_data.shape: {qx_data.shape}, qw_data.shape: {qw_data.shape}")
    # print(f"ref_scales_x.shape: {ref_scales_x.shape}, ref_scales_w.shape: {ref_scales_w.shape}")
    # print(f"ref_scales_x_t.shape: {ref_scales_x.t().shape}")
    # print(f"ref_scales_x_columnwise.shape: {qx._columnwise_scale_inv.shape}")

    y_ref = ref_gemm.qgemm(
        qx=qx_data,
        qw=qw_data,
        out_dtype=out_dtype,
        demunged_sx=CuBLASScaleMunger.demunge_scale_shape_from_backend(
            qtensor_shape=(M, K), scales=ref_scales_x, tile_shape=x_quant_tile_shape
        ),
        demunged_sw=CuBLASScaleMunger.demunge_scale_shape_from_backend(
            qtensor_shape=(N, K), scales=ref_scales_w, tile_shape=w_quant_tile_shape
        ),
        quant_tile_shape_x=x_quant_tile_shape,
        quant_tile_shape_w=w_quant_tile_shape,
        bias=bias,
        out=out.clone() if accumulate else None,
        accumulate=accumulate,
        use_split_accumulator=use_split_accumulator,
    )

    # print("fp8 gemm output: ", y_ref)
    # print("fp8 gemm output shape: ", y_ref.shape)

    x_te_dtype = TE_DType[torch.int8]
    w_te_dtype = TE_DType[torch.int8]
    x_quantizer = Float8BlockQuantizer(
        fp8_dtype=x_te_dtype,
        rowwise=True,
        columnwise=True,
        amax_epsilon=0.0,
        force_pow_2_scales=True,
        block_scaling_dim=x_block_scaling_dim,
    )
    w_quantizer = Float8BlockQuantizer(
        fp8_dtype=w_te_dtype,
        rowwise=True,
        columnwise=True,
        amax_epsilon=0.0,
        force_pow_2_scales=True,
        block_scaling_dim=w_block_scaling_dim,
    )
    # Quantize x and w
    qx = x_quantizer.make_empty(x_shape, dtype=torch.int8, device=device, requires_grad=False)
    qx = x_quantizer.update_quantized(x, qx)
    qw = w_quantizer.make_empty(w_shape, dtype=torch.int8, device=device, requires_grad=False)
    qw = w_quantizer.update_quantized(w, qw)

    qx_data = (
        qx._columnwise_data.view(dtype=torch.int8)
        if x_columnwise
        else qx._rowwise_data.view(dtype=torch.int8)
    )
    qw_data = (
        qw._columnwise_data.view(dtype=torch.int8)
        if w_columnwise
        else qw._rowwise_data.view(dtype=torch.int8)
    )
    ref_scales_x = qx._columnwise_scale_inv if x_columnwise else qx._rowwise_scale_inv
    ref_scales_w = qw._columnwise_scale_inv if w_columnwise else qw._rowwise_scale_inv
    
    if get_device_compute_capability() < (9, 3) or block_len != 128 or not enable_lightop:
        y, _ = w8a8_block_int8_matmul(
            qx_data, qw_data, ref_scales_x, ref_scales_w, [block_len, block_len],
            output_dtype=out_dtype
        )
    else:
        y = lightop.gemm_w8a8_asm(
            qx_data, qw_data, ref_scales_x, ref_scales_w, [block_len, block_len], out_dtype
        )

    # print("int8 gemm output: ", y)
    # print("int8 gemm output shape: ", y.shape)

    torch.testing.assert_close(y, bf16_out, atol=atol, rtol=rtol)
    torch.testing.assert_close(y_ref, bf16_out, atol=atol, rtol=rtol)


def cublas_gemm_fp8_blockwise_case_bw_xgrad(
    dout_dtype,
    w_dtype,
    dx_dtype,    
    M,
    K,
    N,
    noise_type,
    dout_magnitude,
    w_magnitude,
    accumulate,
    use_split_accumulator,
    is_dout_1d_scaled,
    is_w_1d_scaled,
    *,
    dout_columnwise: bool = False,
    w_columnwise: bool = True,
    use_bias: bool = False,
    use_gelu: bool = False,
    use_grad: bool = False,
    atol: float = 5e-1,
    rtol: float = 5e-1
):
    if dout_dtype == torch.float8_e5m2 and w_dtype == torch.float8_e5m2:
        pytest.skip("FP8 GEMM doesn't support both a and b types being torch.float8_e5m2")
    if not (is_dout_1d_scaled or is_w_1d_scaled):
        pytest.skip("FP8 GEMM doesn't support 2dimensional qtile by 2dimensional qtile")
    if not fp8_blockwise_gemm_supported():
        pytest.skip("CUDA version does not support blockwise FP8 gemm.")
    # Setup device and random seed
    device = "cuda"
    seed = 0
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)

    dout_shape = (M, N)
    w_shape = (N, K)
    # generate random input and weight
    if noise_type == "uniform":
        dout = torch.rand(dout_shape, dtype=torch.bfloat16, device=device) * dout_magnitude * 2 - dout_magnitude
        w = torch.rand(w_shape, dtype=torch.bfloat16, device=device) * w_magnitude * 2 - w_magnitude
    elif noise_type == "normal":
        dout = torch.randn(dout_shape, dtype=torch.bfloat16, device=device) * dout_magnitude
        w = torch.randn(w_shape, dtype=torch.bfloat16, device=device) * w_magnitude
    else:
        assert False

    bf16_dx = torch.matmul(dout, w)
    # print("bf16 gemm dx: ", bf16_dx)

    # Setup out tensor if accumulate is True
    if accumulate:
        dx = torch.randn((M, K), dtype=dx_dtype, device=device) * dout_magnitude
    else:
        dx = None

    assert not (use_bias and use_grad), "Bias grad not supported by GEMM"
    # Set quantize_op and quantization parameters
    block_len = blockwise_fp8_block_len
    dout_quant_tile_shape = (1, block_len) if is_dout_1d_scaled else (block_len, block_len)
    w_quant_tile_shape = (1, block_len) if is_w_1d_scaled else (block_len, block_len)
    dout_block_scaling_dim = 1 if is_dout_1d_scaled else 2
    w_block_scaling_dim = 1 if is_w_1d_scaled else 2
    dout_te_dtype = TE_DType[dout_dtype]
    w_te_dtype = TE_DType[w_dtype]
    dout_quantizer = Float8BlockQuantizer(
        fp8_dtype=dout_te_dtype,
        rowwise=True,
        columnwise=True,
        amax_epsilon=0.0,
        force_pow_2_scales=True,
        block_scaling_dim=dout_block_scaling_dim,
    )
    w_quantizer = Float8BlockQuantizer(
        fp8_dtype=w_te_dtype,
        rowwise=True,
        columnwise=True,
        amax_epsilon=0.0,
        force_pow_2_scales=True,
        block_scaling_dim=w_block_scaling_dim,
    )

    # Quantize dout and w
    qdout = dout_quantizer.make_empty(dout_shape, dtype=dout_dtype, device=device, requires_grad=False)
    qdout = dout_quantizer.update_quantized(dout, qdout)
    qw = w_quantizer.make_empty(w_shape, dtype=w_dtype, device=device, requires_grad=False)
    qw = w_quantizer.update_quantized(w, qw)

    if not use_bias:
        bias = None
    else:
        bias = torch.randn((1, N), dtype=torch.bfloat16, device=device)

    # Reference GEMM
    ref_gemm = CuBLASRefBlockwiseGemm()
    scale_decoder = CuBLASScaleMunger()
    qdout_data = (
        qdout._columnwise_data.view(dtype=dout_dtype)
        if dout_columnwise
        else qdout._rowwise_data.view(dtype=dout_dtype)
    )
    qw_data = (
        qw._columnwise_data.view(dtype=w_dtype)
        if w_columnwise
        else qw._rowwise_data.view(dtype=w_dtype)
    )
    ref_scales_dout = qdout._columnwise_scale_inv if dout_columnwise else qdout._rowwise_scale_inv
    ref_scales_w = qw._columnwise_scale_inv if w_columnwise else qw._rowwise_scale_inv

    y_ref = ref_gemm.qgemm(
        qx=qdout_data,
        qw=qw_data,
        out_dtype=dx_dtype,
        demunged_sx=CuBLASScaleMunger.demunge_scale_shape_from_backend(
            qtensor_shape=(M, N), scales=ref_scales_dout, tile_shape=dout_quant_tile_shape
        ),
        demunged_sw=CuBLASScaleMunger.demunge_scale_shape_from_backend(
            qtensor_shape=(K, N), scales=ref_scales_w, tile_shape=w_quant_tile_shape
        ),
        quant_tile_shape_x=dout_quant_tile_shape,
        quant_tile_shape_w=w_quant_tile_shape,
        bias=bias,
        out=dx.clone() if accumulate else None,
        accumulate=accumulate,
        use_split_accumulator=use_split_accumulator,
    )

    # print("fp8 gemm dx: ", y_ref)

    dout_te_dtype = TE_DType[torch.int8]
    w_te_dtype = TE_DType[torch.int8]
    dout_quantizer = Float8BlockQuantizer(
        fp8_dtype=dout_te_dtype,
        rowwise=True,
        columnwise=True,
        amax_epsilon=0.0,
        force_pow_2_scales=True,
        block_scaling_dim=dout_block_scaling_dim,
    )
    w_quantizer = Float8BlockQuantizer(
        fp8_dtype=w_te_dtype,
        rowwise=True,
        columnwise=True,
        amax_epsilon=0.0,
        force_pow_2_scales=True,
        block_scaling_dim=w_block_scaling_dim,
    )
    # Quantize x and w
    qdout = dout_quantizer.make_empty(dout_shape, dtype=torch.int8, device=device, requires_grad=False)
    qdout = dout_quantizer.update_quantized(dout, qdout)
    qw = w_quantizer.make_empty(w_shape, dtype=torch.int8, device=device, requires_grad=False)
    qw = w_quantizer.update_quantized(w, qw)

    qdout_data = (
        qdout._columnwise_data.view(dtype=torch.int8)
        if dout_columnwise
        else qdout._rowwise_data.view(dtype=torch.int8)
    )
    qw_data = (
        qw._columnwise_data.view(dtype=torch.int8)
        if w_columnwise
        else qw._rowwise_data.view(dtype=torch.int8)
    )
    ref_scales_dout = qdout._columnwise_scale_inv if dout_columnwise else qdout._rowwise_scale_inv
    ref_scales_w = qw._columnwise_scale_inv if w_columnwise else qw._rowwise_scale_inv

    if get_device_compute_capability() < (9, 3) or block_len != 128 or not enable_lightop:
        y, _ = w8a8_block_int8_matmul(
            qdout_data, qw_data, ref_scales_dout, ref_scales_w, [block_len, block_len],
            output_dtype=dx_dtype
        )
    else:
        y = lightop.gemm_w8a8_asm(
            qdout_data, qw_data, ref_scales_dout, ref_scales_w, [block_len, block_len], dx_dtype
        )

    # print("int8 gemm dx: ", y)

    torch.testing.assert_close(y_ref, bf16_dx, atol=atol, rtol=rtol)
    torch.testing.assert_close(y, bf16_dx, atol=atol, rtol=rtol)


def cublas_gemm_fp8_blockwise_case_bw_wgrad(
    dout_dtype,
    x_dtype,
    dw_dtype,
    M,
    K,
    N,
    noise_type,
    dout_magnitude,
    x_magnitude,
    accumulate,
    use_split_accumulator,
    is_dout_1d_scaled,
    is_x_1d_scaled,
    *,
    dout_columnwise: bool = True,
    x_columnwise: bool = True,
    use_bias: bool = False,
    use_gelu: bool = False,
    use_grad: bool = False,
    atol: float = 5e-1,
    rtol: float = 5e-1
):
    if dout_dtype == torch.float8_e5m2 and x_dtype == torch.float8_e5m2:
        pytest.skip("FP8 GEMM doesn't support both a and b types being torch.float8_e5m2")
    if not (is_dout_1d_scaled or is_x_1d_scaled):
        pytest.skip("FP8 GEMM doesn't support 2dimensional qtile by 2dimensional qtile")
    if not fp8_blockwise_gemm_supported():
        pytest.skip("CUDA version does not support blockwise FP8 gemm.")
    # Setup device and random seed
    device = "cuda"
    seed = 0
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)

    dout_shape = (M, N)
    x_shape = (M, K)
    # generate random input and weight
    if noise_type == "uniform":
        dout = torch.rand(dout_shape, dtype=torch.bfloat16, device=device) * dout_magnitude * 2 - dout_magnitude
        x = torch.rand(x_shape, dtype=torch.bfloat16, device=device) * x_magnitude * 2 - x_magnitude
    elif noise_type == "normal":
        dout = torch.randn(dout_shape, dtype=torch.bfloat16, device=device) * dout_magnitude
        x = torch.randn(x_shape, dtype=torch.bfloat16, device=device) * x_magnitude
    else:
        assert False

    bf16_dw = torch.matmul(dout.t(), x)
    # print("bf16 gemm dw: ", bf16_dw)

    # Setup out tensor if accumulate is True
    if accumulate:
        dw = torch.randn((N, K), dtype=dw_dtype, device=device) * dout_magnitude
    else:
        dw = None

    assert not (use_bias and use_grad), "Bias grad not supported by GEMM"
    # Set quantize_op and quantization parameters
    block_len = blockwise_fp8_block_len
    dout_quant_tile_shape = (1, block_len) if is_dout_1d_scaled else (block_len, block_len)
    x_quant_tile_shape = (1, block_len) if is_x_1d_scaled else (block_len, block_len)
    dout_block_scaling_dim = 1 if is_dout_1d_scaled else 2
    x_block_scaling_dim = 1 if is_x_1d_scaled else 2
    dout_te_dtype = TE_DType[dout_dtype]
    x_te_dtype = TE_DType[x_dtype]
    dout_quantizer = Float8BlockQuantizer(
        fp8_dtype=dout_te_dtype,
        rowwise=True,
        columnwise=True,
        amax_epsilon=0.0,
        force_pow_2_scales=True,
        block_scaling_dim=dout_block_scaling_dim,
    )
    x_quantizer = Float8BlockQuantizer(
        fp8_dtype=x_te_dtype,
        rowwise=True,
        columnwise=True,
        amax_epsilon=0.0,
        force_pow_2_scales=True,
        block_scaling_dim=x_block_scaling_dim,
    )

    # Quantize dout and w
    qdout = dout_quantizer.make_empty(dout_shape, dtype=dout_dtype, device=device, requires_grad=False)
    qdout = dout_quantizer.update_quantized(dout, qdout)
    qx = x_quantizer.make_empty(x_shape, dtype=x_dtype, device=device, requires_grad=False)
    qx = x_quantizer.update_quantized(x, qx)

    if not use_bias:
        bias = None
    else:
        bias = torch.randn((1, N), dtype=torch.bfloat16, device=device)

    # Reference GEMM
    ref_gemm = CuBLASRefBlockwiseGemm()
    scale_decoder = CuBLASScaleMunger()
    qdout_data = (
        qdout._columnwise_data.view(dtype=dout_dtype)
        if dout_columnwise
        else qdout._rowwise_data.view(dtype=dout_dtype)
    )
    qx_data = (
        qx._columnwise_data.view(dtype=x_dtype)
        if x_columnwise
        else qx._rowwise_data.view(dtype=x_dtype)
    )
    ref_scales_dout = qdout._columnwise_scale_inv if dout_columnwise else qdout._rowwise_scale_inv
    ref_scales_x = qx._columnwise_scale_inv if x_columnwise else qx._rowwise_scale_inv

    y_ref = ref_gemm.qgemm(
        qx=qdout_data,
        qw=qx_data,
        out_dtype=dw_dtype,
        demunged_sx=CuBLASScaleMunger.demunge_scale_shape_from_backend(
            qtensor_shape=(N, M), scales=ref_scales_dout, tile_shape=dout_quant_tile_shape
        ),
        demunged_sw=CuBLASScaleMunger.demunge_scale_shape_from_backend(
            qtensor_shape=(K, M), scales=ref_scales_x, tile_shape=x_quant_tile_shape
        ),
        quant_tile_shape_x=dout_quant_tile_shape,
        quant_tile_shape_w=x_quant_tile_shape,
        bias=bias,
        out=dw.clone() if accumulate else None,
        accumulate=accumulate,
        use_split_accumulator=use_split_accumulator,
    )

    # print("fp8 gemm dw: ",y_ref)

    dout_te_dtype = TE_DType[torch.int8]
    x_te_dtype = TE_DType[torch.int8]
    dout_quantizer = Float8BlockQuantizer(
        fp8_dtype=dout_te_dtype,
        rowwise=True,
        columnwise=True,
        amax_epsilon=0.0,
        force_pow_2_scales=True,
        block_scaling_dim=dout_block_scaling_dim,
    )
    x_quantizer = Float8BlockQuantizer(
        fp8_dtype=x_te_dtype,
        rowwise=True,
        columnwise=True,
        amax_epsilon=0.0,
        force_pow_2_scales=True,
        block_scaling_dim=x_block_scaling_dim,
    )
    # Quantize x and w
    qdout = dout_quantizer.make_empty(dout_shape, dtype=torch.int8, device=device, requires_grad=False)
    qdout = dout_quantizer.update_quantized(dout, qdout)
    qx = x_quantizer.make_empty(x_shape, dtype=torch.int8, device=device, requires_grad=False)
    qx = x_quantizer.update_quantized(x, qx)

    qdout_data = (
        qdout._columnwise_data.view(dtype=torch.int8)
        if dout_columnwise
        else qdout._rowwise_data.view(dtype=torch.int8)
    )
    qx_data = (
        qx._columnwise_data.view(dtype=torch.int8)
        if x_columnwise
        else qx._rowwise_data.view(dtype=torch.int8)
    )
    ref_scales_dout = qdout._columnwise_scale_inv if dout_columnwise else qdout._rowwise_scale_inv
    ref_scales_x = qx._columnwise_scale_inv if x_columnwise else qx._rowwise_scale_inv

    # print(f"qdout_data.shape: {qdout_data.shape}, qx_data.shape: {qx_data.shape}")
    # print(f"ref_scales_dout.shape: {ref_scales_dout.shape}, ref_scales_x.shape: {ref_scales_x.shape}")
    if get_device_compute_capability() < (9, 3) or block_len != 128  or not enable_lightop:
        y, _ = w8a8_block_int8_matmul_wgrad(
            qdout_data, qx_data, ref_scales_dout, ref_scales_x, dw.clone() if accumulate else None,
            accumulate, [block_len, block_len],
            output_dtype=dw_dtype
        )
    else:
        y = lightop.gemm_w8a8_wgrad_asm(
            qdout_data, qx_data, ref_scales_dout, ref_scales_x, dw.clone() if accumulate else None,
            accumulate, [block_len, block_len], dw_dtype
        )

    # print("int8 gemm dw: ",y)

    torch.testing.assert_close(y_ref, bf16_dw, atol=atol, rtol=rtol)
    torch.testing.assert_close(y, bf16_dw, atol=atol, rtol=rtol)


def test_cublas_gemm_fp8_blockwise_fw(
    x_dtype,
    w_dtype,
    out_dtype,
    M,
    K,
    N,
    noise_type,
    x_magnitude,
    w_magnitude,
    accumulate,
    use_split_accumulator,
    is_x_1d_scaled,
    is_w_1d_scaled,
):
    cublas_gemm_fp8_blockwise_case_fw(
        x_dtype,
        w_dtype,
        out_dtype,
        M,
        K,
        N,
        noise_type,
        x_magnitude,
        w_magnitude,
        accumulate,
        use_split_accumulator,
        is_x_1d_scaled,
        is_w_1d_scaled,
    )


def test_cublas_gemm_fp8_blockwise_bw_xgrad(
    dout_dtype,
    w_dtype,
    dx_dtype,
    M,
    K,
    N,
    noise_type,
    dout_magnitude,
    w_magnitude,
    accumulate,
    use_split_accumulator,
    is_dout_1d_scaled,
    is_w_1d_scaled,
):
    cublas_gemm_fp8_blockwise_case_bw_xgrad(
        dout_dtype,
        w_dtype,
        dx_dtype,
        M,
        K,
        N,
        noise_type,
        dout_magnitude,
        w_magnitude,
        accumulate,
        use_split_accumulator,
        is_dout_1d_scaled,
        is_w_1d_scaled,
    )
        

def test_cublas_gemm_fp8_blockwise_bw_wgrad(
    dout_dtype,
    x_dtype,
    dw_dtype,
    M,
    K,
    N,
    noise_type,
    dout_magnitude,
    x_magnitude,
    accumulate,
    use_split_accumulator,
    is_dout_1d_scaled,
    is_x_1d_scaled,
):
    cublas_gemm_fp8_blockwise_case_bw_wgrad(
        dout_dtype,
        x_dtype,
        dw_dtype,
        M,
        K,
        N,
        noise_type,
        dout_magnitude,
        x_magnitude,
        accumulate,
        use_split_accumulator,
        is_dout_1d_scaled,
        is_x_1d_scaled,
    )


if __name__ == "__main__":
    test_cublas_gemm_fp8_blockwise_fw(
        x_dtype=torch.float8_e4m3fn, # torch.float8_e4m3fnuz if te.e4m3 use funz
        w_dtype=torch.float8_e4m3fn, # torch.float8_e4m3fnuz if te.e4m3 use funz
        out_dtype=torch.bfloat16, 
        M=128, # batch_size * seq_len
        K=512, # in_feature
        N=256, # out_feature
        noise_type="normal", 
        x_magnitude=1e-1,
        w_magnitude=1, 
        accumulate=False, 
        use_split_accumulator=True, 
        is_x_1d_scaled=True, 
        is_w_1d_scaled=False
    )

    test_cublas_gemm_fp8_blockwise_bw_xgrad(
        dout_dtype=torch.float8_e4m3fn, # torch.float8_e4m3fnuz if te.e4m3 use funz
        w_dtype=torch.float8_e4m3fn, # torch.float8_e4m3fnuz if te.e4m3 use funz 
        dx_dtype=torch.bfloat16,
        M=128, # batch_size * seq_len
        K=512, # in_feature
        N=256, # out_feature
        noise_type="normal", 
        dout_magnitude=1e-1,
        w_magnitude=1, 
        accumulate=False, 
        use_split_accumulator=True, 
        is_dout_1d_scaled=True, 
        is_w_1d_scaled=False,
    )

    test_cublas_gemm_fp8_blockwise_bw_wgrad(
        dout_dtype=torch.float8_e4m3fn, # torch.float8_e4m3fnuz if te.e4m3 use funz
        x_dtype=torch.float8_e4m3fn, # torch.float8_e4m3fnuz if te.e4m3 use funz 
        dw_dtype=torch.bfloat16,
        M=128, # batch_size * seq_len
        K=512, # in_feature
        N=256, # out_feature
        noise_type="normal", 
        dout_magnitude=1e-1,
        x_magnitude=1, 
        accumulate=False, 
        use_split_accumulator=True, 
        is_dout_1d_scaled=True, 
        is_x_1d_scaled=True,
    )