from torch.utils.cpp_extension import load_inline
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 <hip/amd_detail/amd_hip_fp8.h>
#include <hip/amd_detail/amd_hip_bf16.h>

#define DIV_UP(a,b) (((a) + (b) - 1)/(b))

constexpr const int BLOCK = 128;
constexpr const int TILE_SIZE = 16;

__global__ void custom_kernel(const __hip_fp8_e4m3_fnuz* a, const __hip_fp8_e4m3_fnuz* b, const float* as, const float* bs, 
                   __hip_bfloat16* c, int m, int n, int k) {
                   
    // Your implementation here
    int cx = threadIdx.x + blockDim.x * blockIdx.x;
    int cy = threadIdx.y + blockDim.y * blockIdx.y;
    if(cx >= m || cy >= n) return;

    int sn = DIV_UP(n, BLOCK);
    __shared__ __hip_fp8_e4m3_fnuz As[TILE_SIZE][TILE_SIZE];
    __shared__ __hip_fp8_e4m3_fnuz Bs[TILE_SIZE][TILE_SIZE];

    
    float result = 0;
    // split loop into an outer loop over different blocks, and an inner loop within one block.
    // we can assume k % BLOCK == 0.
    for(int i = 0; i < k; i += BLOCK) {
   	float block_result = 0;
        for(int ii = 0; ii < BLOCK; ii += TILE_SIZE) {
		As[threadIdx.y][threadIdx.x] = a[cx + (i + ii + threadIdx.y) * m];
		Bs[threadIdx.y][threadIdx.x] = b[cy + (i + ii + threadIdx.x) * n];
		__syncthreads();
		for(int iii = 0; iii < TILE_SIZE; ++iii) {
			float av = (float) As[iii][threadIdx.x]; // we need a[m * (i + ii + iii) + 
			float bv = (float) Bs[threadIdx.y][iii];
			block_result += av * bv;
		}
		__syncthreads();
        }
        
        // before we can go to the next block, scale the result of the current block
        // and accumulate to final result
        // note the different indexing into as and bs
        result += block_result * as[m * i/BLOCK + cx] * bs[cy/BLOCK + i/BLOCK * sn];
    }
    
    // finally, write the result as bf16
    c[cx * n + cy] = (__hip_bfloat16)result;
}

void fp8_mm(torch::Tensor a, torch::Tensor b, torch::Tensor as, torch::Tensor bs, torch::Tensor c) {
    int m = a.size(0);
    int n = b.size(0);
    int k = a.size(1);
    custom_kernel<<<dim3(DIV_UP(m, TILE_SIZE), 
		    DIV_UP(n, TILE_SIZE)), 
	    	dim3(TILE_SIZE,
			TILE_SIZE), 0, 0>>> 
		((__hip_fp8_e4m3_fnuz*)a.data_ptr(), 
		 (__hip_fp8_e4m3_fnuz*)b.data_ptr(), 
    as.data_ptr<float>(), 
    bs.data_ptr<float>(), 
    (__hip_bfloat16*)c.data_ptr(), 
    m, 
    n, 
    k);
}
"""

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

module = load_inline(
    name='fp8_mm',
    cpp_sources=[CPP_WRAPPER],
    cuda_sources=[CUDA_SRC],
    functions=['fp8_mm'],
    verbose=True,
    extra_cuda_cflags=["--offload-arch=gfx942", "-std=c++20", "-O3", "-use_fast_math"],
)


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

python submission.py
Using /home/runner/.cache/torch_extensions/py310_cpu as PyTorch extensions root...
Creating extension directory /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm...
[92mSuccessfully preprocessed all matching files.[0m
Detected CUDA files, patching ldflags
Emitting ninja build file /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/build.ninja...
Building extension module fp8_mm...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
Loading extension module fp8_mm...

/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/cuda.cu -> /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip [ok]
/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]
Total number of unsupported CUDA function calls: 0


Total number of replaced kernel launches: 1
[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 /home/runner/.local/lib/python3.10/site-packages/torch/include -isystem /home/runner/.local/lib/python3.10/site-packages/torch/include/torch/csrc/api/include -isystem /home/runner/.local/lib/python3.10/site-packages/torch/include/THH -isystem /opt/rocm-6.3.1/include -isystem /home/runner/_work/_tool/Python/3.10.17/x64/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 /home/runner/.local/lib/python3.10/site-packages/torch/include -isystem /home/runner/.local/lib/python3.10/site-packages/torch/include/torch/csrc/api/include -isystem /home/runner/.local/lib/python3.10/site-packages/torch/include/THH -isystem /opt/rocm-6.3.1/include -isystem /home/runner/_work/_tool/Python/3.10.17/x64/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 --offload-arch=gfx942 -std=c++20 -O3 -use_fast_math -fno-gpu-rdc -c /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip -o hip.cuda.o 
clang++: warning: argument unused during compilation: '-u se_fast_math' [-Wunused-command-line-argument]
[3/3] clang++ main.o hip.cuda.o -shared -L/home/runner/.local/lib/python3.10/site-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 benchmark /tmp/tmprwc88gxd
Using /home/runner/.cache/torch_extensions/py310_cpu as PyTorch extensions root...
[92mSuccessfully preprocessed all matching files.[0m
Detected CUDA files, patching ldflags
Emitting ninja build file /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/build.ninja...
Building extension module fp8_mm...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
Loading extension module fp8_mm...

/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]
/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]
Total number of unsupported CUDA function calls: 0


Total number of replaced kernel launches: 1
[1/1] clang++ hip.cuda.o main.o -shared -L/home/runner/.local/lib/python3.10/site-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
