from torch.utils.cpp_extension import load_inline
import triton
import triton.language as tl

from task import input_t, output_t


CPP_WRAPPER = """
void fp8_mm(torch::Tensor a, torch::Tensor b, torch::Tensor as, torch::Tensor bs, torch::Tensor c);
"""

CUDA_SRC = """
#include "ck/ck.hpp"
#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
#include "ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3_ab_scale.hpp"
#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
#include "ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp"

#include "ck/library/utility/device_memory.hpp"
#include "ck/library/utility/host_tensor.hpp"
#include "ck/library/utility/host_tensor_generator.hpp"
#include "ck/library/utility/literals.hpp"
#include "ck/library/reference_tensor_operation/cpu/reference_gemm.hpp"
#include "ck/library/utility/check_err.hpp"

#include "ck/utility/blkgemmpipe_scheduler.hpp"


template <ck::index_t... Is>
using S = ck::Sequence<Is...>;

using BF16 = ck::bhalf_t;
using FP8  = ck::f8_t;
using F32  = float;

using Row = ck::tensor_layout::gemm::RowMajor;
using Col = ck::tensor_layout::gemm::ColumnMajor;

using A0DataType       = FP8;
using A1DataType       = F32;
using B0DataType       = FP8;
using B1DataType       = F32;
using AccDataType      = F32;
using CShuffleDataType = F32;
using DsDataType       = ck::Tuple<>;
using EDataType        = BF16;

using A0Layout = Row;
using B0Layout = Col;
using D0Layout = Row;
using D1Layout = Col;
using DsLayout = ck::Tuple<>;
using ELayout  = Row;

using PassThrough = ck::tensor_operation::element_wise::PassThrough;

using AElementOp   = PassThrough;
using BElementOp   = PassThrough;
using CDEElementOp = PassThrough;

static constexpr auto GemmSpec = ck::tensor_operation::device::GemmSpecialization::Default;

static constexpr ck::index_t Scale_Block_M = 1;
static constexpr ck::index_t Scale_Block_N = 128;
static constexpr ck::index_t Scale_Block_K = 128;

using DeviceOpInstance = ck::tensor_operation::device::DeviceGemmMultiD_ABScale_Xdl_CShuffle_V3
         <A0Layout, B0Layout, DsLayout, ELayout,
          A0DataType, A1DataType, B0DataType, B1DataType, DsDataType, EDataType, AccDataType, CShuffleDataType,
          AElementOp,  BElementOp, CDEElementOp, GemmSpec,
          256, Scale_Block_M, Scale_Block_N, Scale_Block_K,
          16, 128,
          256, 16, 16,
          16,   16,
          1,    2,
          S<16, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0,
          S<16, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0,
          1,    2,  S<1, 16, 1, 16>,  S<8>,
          ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v1, FP8>;


void fp8_mm(torch::Tensor a, torch::Tensor b, torch::Tensor as, torch::Tensor bs, torch::Tensor c) {

    // GEMM shape
    ck::index_t M = a.size(0);
    ck::index_t N = b.size(0);
    ck::index_t K = a.size(1);

    ck::index_t StrideA = K;
    ck::index_t StrideB = K;
    ck::index_t StrideE = N;

    auto a_element_op   = AElementOp{};
    auto b_element_op   = BElementOp{};
    auto cde_element_op = CDEElementOp{};

    constexpr ck::index_t NumDTensor = DsDataType::Size();

    // do GEMM
    auto device_op = DeviceOpInstance{};
    auto invoker   = device_op.MakeInvoker();
    auto argument  = device_op.MakeArgument(a.data_ptr(),
                                           b.data_ptr(),
                                           std::array<const void*, NumDTensor>{},
                                           c.data_ptr(),
                                           M,
                                           N,
                                           K,
                                           StrideA,
                                           StrideB,
                                           std::array<ck::index_t, NumDTensor>{},
                                           StrideE,
                                           as.data_ptr(),
                                           bs.data_ptr(),
                                           a_element_op,
                                           b_element_op,
                                           cde_element_op);
    
    invoker.Run(argument, StreamConfig{nullptr, false, 0, 50, 100});
}
"""

import os
os.environ["CXX"] = "clang++"
os.environ["PYTORCH_ROCM_ARCH"] = "gfx942"

module = load_inline(
    name='fp8_mm',
    cpp_sources=[CPP_WRAPPER],
    cuda_sources=[CUDA_SRC],
    functions=['fp8_mm'],
    verbose=True,
    extra_cuda_cflags=[
        "-DCK_ENABLE_BF16 -DCK_ENABLE_BF8 -DCK_ENABLE_FP16 -DCK_ENABLE_FP32 -DCK_ENABLE_FP64 -DCK_ENABLE_FP8",
        "-DCK_TIME_KERNEL=1 -DCK_USE_FNUZ_FP8 -DCK_USE_GFX94 -DCK_USE_XDL -DDL_KERNELS -DDPP_KERNELS -DUSE_PROF_API=1",
        "-D__HIP_PLATFORM_AMD__=1 -D__HIP_PLATFORM_HCC__=1 -D__HIP_ROCclr__=1 -U__HIP_NO_HALF_CONVERSIONS__ -U__HIP_NO_HALF_OPERATORS__",
        "-DDEBUG_LOG",
        "--offload-arch=gfx942",
        # "-Wall -Wextra -Wcomment -Wendif-labels -Wformat -Winit-self -Wreturn-type -Wsequence-point -Wswitch -Wtrigraphs -Wundef -Wuninitialized -Wunreachable-code -Wunused -Wno-reserved-identifier -Wno-option-ignored -Wsign-compare -Wno-extra-semi-stmt -Wno-unused-template -Wno-missing-field-initializers -Wno-deprecated-declarations -Weverything -Wno-c++98-compat -Wno-c++98-compat-pedantic -Wno-conversion -Wno-double-promotion -Wno-exit-time-destructors -Wno-extra-semi -Wno-float-conversion -Wno-gnu-anonymous-struct -Wno-gnu-zero-variadic-macro-arguments -Wno-missing-prototypes -Wno-nested-anon-types -Wno-padded -Wno-return-std-move-in-c++11 -Wno-shorten-64-to-32 -Wno-sign-conversion -Wno-unknown-warning-option -Wno-unused-command-line-argument -Wno-weak-vtables -Wno-covered-switch-default -Wno-unsafe-buffer-usage -Wno-unused-lambda-capture -Wno-nvcc-compat -Wno-bit-int-extension -Wno-pass-failed -Wno-switch-default -Wno-unique-object-duplication -fno-offload-uniform-block -mllvm --lsr-drop-solution=1 -mllvm -enable-post-misched=0 -mllvm -amdgpu-coerce-illegal-types=1 -mllvm -amdgpu-early-inline-all=true -mllvm -amdgpu-function-calls=false -fcolor-diagnostics -std=gnu++17 -x hip"
        "-Wno-unused-result",
        "-Wno-switch-bool",
        "-Wno-vla-cxx-extension",
        "-Wno-undefined-func-template",
        "-Wno-macro-redefined",
        "-fgpu-flush-denormals-to-zero",
        "-fno-offload-uniform-block",
        "-mllvm --amdgpu-kernarg-preload-count=16",
        "-mllvm -enable-post-misched=0",
        "-mllvm -amdgpu-early-inline-all=true -mllvm -amdgpu-function-calls=false",
        "-mllvm -amdgpu-coerce-illegal-types=1"
    ],
)


def custom_kernel_hip(data: input_t) -> output_t:
    a, b, a_scale, b_scale, c = data
    a = a.contiguous()
    b = b.contiguous()
    a_scale = a_scale.contiguous()
    b_scale = b_scale.contiguous()
    module.fp8_mm(a, b, a_scale, b_scale, c)
    return c

@triton.jit
def fp8_gemm_kernel(a_ptr, b_ptr, c_ptr,
                    a_s_ptr, b_s_ptr,
                    M, N: tl.constexpr, K: tl.constexpr,
                    BLOCK_SIZE_M: tl.constexpr,
                    BLOCK_SIZE_N: tl.constexpr,
                    BLOCK_SIZE_K: tl.constexpr,
                    GROUP_SIZE_M: tl.constexpr):
    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)

    # pid_m = pid // num_pid_n
    # pid_n = pid % num_pid_n

    NUM_XCDS: tl.constexpr = 8
    GRID_MN = tl.cdiv(M, BLOCK_SIZE_M) * tl.cdiv(N, BLOCK_SIZE_N)
    ## pid remapping on xcds
    # Number of pids per XCD in the new arrangement
    pids_per_xcd = (GRID_MN + NUM_XCDS - 1) // NUM_XCDS
    # When GRID_MN cannot divide NUM_XCDS, some xcds will have
    # pids_per_xcd pids, the other will have pids_per_xcd - 1 pids.
    # We calculate the number of xcds that have pids_per_xcd pids as
    # tall_xcds
    tall_xcds = GRID_MN % NUM_XCDS
    tall_xcds = NUM_XCDS if tall_xcds == 0 else tall_xcds
    # Compute current XCD and local pid within the XCD
    xcd = pid % NUM_XCDS
    local_pid = pid // NUM_XCDS
    # Calculate new pid based on the new grouping
    # Note that we need to consider the following two cases:
    # 1. the current pid is on a tall xcd
    # 2. the current pid is on a short xcd
    if xcd < tall_xcds:
        pid = xcd * pids_per_xcd + local_pid
    else:
        pid = (
            tall_xcds * pids_per_xcd
            + (xcd - tall_xcds) * (pids_per_xcd - 1)
            + local_pid
        )

    num_pid_in_group = GROUP_SIZE_M * num_pid_n
    group_id = pid // 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 + ((pid % num_pid_in_group) % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    k = tl.cdiv(K, BLOCK_SIZE_K)
    offs_m = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
    offs_n = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
    offs_k = tl.arange(0, BLOCK_SIZE_K)
    a_ptrs = a_ptr + offs_m[:, None] * K + offs_k[None, :]
    b_ptrs = b_ptr + offs_n[None, :] * K + offs_k[:, None]
    a_s_ptrs = a_s_ptr + offs_m
    b_s_ptrs = b_s_ptr + (offs_n // BLOCK_SIZE_K) * k

    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    for i in range(k):
        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - i * BLOCK_SIZE_K, other=0.0)
        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - i * BLOCK_SIZE_K, other=0.0)
        a_s = tl.load(a_s_ptrs)
        b_s = tl.load(b_s_ptrs)
        accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
        a_ptrs += BLOCK_SIZE_K
        b_ptrs += BLOCK_SIZE_K
        a_s_ptrs += M
        b_s_ptrs += 1
    c = accumulator.to(c_ptr.dtype.element_ty)
    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    c_ptrs = c_ptr + offs_m[:, None] * N + offs_n[None, :]
    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
    tl.store(c_ptrs, c, mask=mask)


def custom_kernel_triton(data: input_t) -> output_t:
    a, b, a_scale, b_scale, c = data
    a = a.contiguous()
    b = b.contiguous()
    a_scale = a_scale.t().contiguous()
    b_scale = b_scale.contiguous()
    K = a.size(-1)
    M = a.numel() // K
    N = b.size(0)

    if M < 2048:
        BLOCK_SIZE_M = 32
    else:
        BLOCK_SIZE_M = 64
    BLOCK_SIZE_N = 64
    BLOCK_SIZE_K = 128
    GROUP_SIZE_M = 16
    num_stages = 2
    num_warps = 8
    waves_per_eu = 2

    grid = (triton.cdiv(M, BLOCK_SIZE_M) * triton.cdiv(N, BLOCK_SIZE_N),)
    fp8_gemm_kernel[grid](a, b, c, a_scale, b_scale, M, N, K,
                        BLOCK_SIZE_M=BLOCK_SIZE_M,
                        BLOCK_SIZE_N=BLOCK_SIZE_N,
                        BLOCK_SIZE_K=BLOCK_SIZE_K,
                        GROUP_SIZE_M=GROUP_SIZE_M,
                        num_stages=num_stages,
                        num_warps=num_warps,
                        waves_per_eu=waves_per_eu,
                        kpack=2,
                        matrix_instr_nonkdim=16,)
    # bm = fp8_gemm_kernel.best_config.kwargs["BLOCK_SIZE_M"]
    # bn = fp8_gemm_kernel.best_config.kwargs["BLOCK_SIZE_N"]
    # ns = fp8_gemm_kernel.best_config.num_stages
    # m = M / 64
    # n = N / 64
    # k = K / 64
    # if (m != 16):
    #     print(f"{N},{K},{bm},{bn},{ns}")
    return c


def custom_kernel(data: input_t) -> output_t:
    K = data[0].size(-1)
    M = data[0].numel() // K
    N = data[1].size(0)

    test_m = M // 64
    test_n = N // 64
    test_k = K // 64
    if test_m == 96 and test_n == 72 and test_k == 112:
        return custom_kernel_hip(data)
    return custom_kernel_triton(data)
python submission.py
[92mSuccessfully preprocessed all matching files.[0m

/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/main.cpp -> /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/main.cpp [skipped, no changes]
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/cuda.cu -> /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip [ok]
Total number of unsupported CUDA function calls: 0


Total number of replaced kernel launches: 0
[1/3] clang++ -MMD -MF main.o.d -DTORCH_EXTENSION_NAME=fp8_mm -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1016\" -isystem /usr/local/lib/python3.10/dist-packages/torch/include -isystem /usr/local/lib/python3.10/dist-packages/torch/include/torch/csrc/api/include -isystem /usr/local/lib/python3.10/dist-packages/torch/include/THH -isystem /opt/rocm-6.3.1/include -isystem /usr/include/python3.10 -fPIC -std=c++17 -c /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/main.cpp -o main.o -D__HIP_PLATFORM_AMD__=1 -DUSE_ROCM=1 -DHIPBLAS_V2 -fPIC
[2/3] /opt/rocm-6.3.1/bin/hipcc  -DWITH_HIP -DTORCH_EXTENSION_NAME=fp8_mm -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1016\" -isystem /usr/local/lib/python3.10/dist-packages/torch/include -isystem /usr/local/lib/python3.10/dist-packages/torch/include/torch/csrc/api/include -isystem /usr/local/lib/python3.10/dist-packages/torch/include/THH -isystem /opt/rocm-6.3.1/include -isystem /usr/include/python3.10 -fPIC -std=c++17 -D__HIP_PLATFORM_AMD__=1 -DUSE_ROCM=1 -DHIPBLAS_V2 -fPIC -DCUDA_HAS_FP16=1 -D__HIP_NO_HALF_OPERATORS__=1 -D__HIP_NO_HALF_CONVERSIONS__=1 -DHIP_ENABLE_WARP_SYNC_BUILTINS=1 --offload-arch=gfx942 -fno-gpu-rdc -DCK_ENABLE_BF16 -DCK_ENABLE_BF8 -DCK_ENABLE_FP16 -DCK_ENABLE_FP32 -DCK_ENABLE_FP64 -DCK_ENABLE_FP8 -DCK_TIME_KERNEL=1 -DCK_USE_FNUZ_FP8 -DCK_USE_GFX94 -DCK_USE_XDL -DDL_KERNELS -DDPP_KERNELS -DUSE_PROF_API=1 -D__HIP_PLATFORM_AMD__=1 -D__HIP_PLATFORM_HCC__=1 -D__HIP_ROCclr__=1 -U__HIP_NO_HALF_CONVERSIONS__ -U__HIP_NO_HALF_OPERATORS__ -DDEBUG_LOG --offload-arch=gfx942 -Wno-unused-result -Wno-switch-bool -Wno-vla-cxx-extension -Wno-undefined-func-template -Wno-macro-redefined -fgpu-flush-denormals-to-zero -fno-offload-uniform-block -mllvm --amdgpu-kernarg-preload-count=16 -mllvm -enable-post-misched=0 -mllvm -amdgpu-early-inline-all=true -mllvm -amdgpu-function-calls=false -mllvm -amdgpu-coerce-illegal-types=1 -c /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip -o hip.cuda.o 
In file included from /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip:8:
In file included from /opt/rocm-6.3.1/include/ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3_ab_scale.hpp:15:
/opt/rocm-6.3.1/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_multi_d_ab_scale.hpp:39:5: warning: failed to meet occupancy target given by 'amdgpu-waves-per-eu' in '_ZN2ck27kernel_gemm_xdl_cshuffle_v3INS_42GridwiseGemmMultiD_ABScale_xdl_cshuffle_v3INS_13tensor_layout4gemm8RowMajorENS3_11ColumnMajorENS_5TupleIJEEES4_DB8_S8_ffS7_tNS_16tensor_operation12element_wise11PassThroughESB_SB_LNS9_6device18GemmSpecializationE0ELi256ELi1ELi128ELi128ELi16ELi128ELi256ELi16ELi16ELi16ELi16ELi1ELi2ENS_8SequenceIJLi16ELi16ELi1EEEENSE_IJLi1ELi0ELi2EEEESG_Li2ELi16ELi16ELb0ELi0ESF_SG_SG_Li2ELi16ELi16ELb0ELi0ELi1ELi2ENSE_IJLi1ELi16ELi1ELi16EEEENSE_IJLi8EEEELNS_26BlockGemmPipelineSchedulerE0ELNS_24BlockGemmPipelineVersionE0ES8_S8_S8_S8_EELb1ELNS_25InMemoryDataOperationEnumE0ELi2ELNS_10TailNumberE10EEEvNT_8ArgumentE': desired occupancy was 2, final occupancy is 1 [-Wpass-failed]
   39 |     kernel_gemm_xdl_cshuffle_v3(typename GridwiseGemm::Argument karg)
      |     ^
/opt/rocm-6.3.1/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_multi_d_ab_scale.hpp:39:5: warning: failed to meet occupancy target given by 'amdgpu-waves-per-eu' in '_ZN2ck27kernel_gemm_xdl_cshuffle_v3INS_42GridwiseGemmMultiD_ABScale_xdl_cshuffle_v3INS_13tensor_layout4gemm8RowMajorENS3_11ColumnMajorENS_5TupleIJEEES4_DB8_S8_ffS7_tNS_16tensor_operation12element_wise11PassThroughESB_SB_LNS9_6device18GemmSpecializationE0ELi256ELi1ELi128ELi128ELi16ELi128ELi256ELi16ELi16ELi16ELi16ELi1ELi2ENS_8SequenceIJLi16ELi16ELi1EEEENSE_IJLi1ELi0ELi2EEEESG_Li2ELi16ELi16ELb0ELi0ESF_SG_SG_Li2ELi16ELi16ELb0ELi0ELi1ELi2ENSE_IJLi1ELi16ELi1ELi16EEEENSE_IJLi8EEEELNS_26BlockGemmPipelineSchedulerE0ELNS_24BlockGemmPipelineVersionE0ES8_S8_S8_S8_EELb0ELNS_25InMemoryDataOperationEnumE0ELi2ELNS_10TailNumberE10EEEvNT_8ArgumentE': desired occupancy was 2, final occupancy is 1 [-Wpass-failed]
2 warnings generated when compiling for gfx942.
[3/3] clang++ hip.cuda.o main.o -shared -L/usr/local/lib/python3.10/dist-packages/torch/lib -lc10 -lc10_hip -ltorch_cpu -ltorch_hip -ltorch -ltorch_python -L/opt/rocm-6.3.1/lib -lamdhip64 -o fp8_mm.so

python eval.py test /tmp/tmp6io2v2u1
[92mSuccessfully preprocessed all matching files.[0m

/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/main.cpp -> /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/main.cpp [skipped, no changes]
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/cuda.cu -> /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip [skipped, already hipified]
Total number of unsupported CUDA function calls: 0


Total number of replaced kernel launches: 0
ninja: no work to do.
