# Fixed and Optimized FP8 Matrix Multiplication for AMD MI300X (gfx942)
import os 
os.environ['PYTORCH_ROCM_ARCH'] = 'gfx942'

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>
#include <hip/hip_runtime.h>

constexpr const int BLOCK_SIZE = 128;

// Basic optimized kernel following the original template structure
__global__ void custom_kernel_basic(const __hip_fp8_e4m3_fnuz* __restrict__ a, 
                                   const __hip_fp8_e4m3_fnuz* __restrict__ b, 
                                   const float* __restrict__ as, 
                                   const float* __restrict__ bs, 
                                   __hip_bfloat16* __restrict__ c, 
                                   int m, int n, int k) {
    
    const int cx = threadIdx.x + blockDim.x * blockIdx.x;
    const int cy = threadIdx.y + blockDim.y * blockIdx.y;
    
    if (cx >= m || cy >= n) return;
    
    const int scale_n = (n + BLOCK_SIZE - 1) / BLOCK_SIZE;
    
    float result = 0.0f;
    
    // Main computation loop - process K dimension in blocks of BLOCK_SIZE
    for (int k_block = 0; k_block < k; k_block += BLOCK_SIZE) {
        float block_result = 0.0f;
        
        // Inner loop within current block
        const int k_end = min(k_block + BLOCK_SIZE, k);
        #pragma unroll 8
        for (int kk = k_block; kk < k_end; kk++) {
            // Load FP8 values and convert to float
            float a_val = (float)a[cx + kk * m];
            float b_val = (float)b[cy + kk * n];
            block_result += a_val * b_val;
        }
        
        // Apply block-wise scaling - this is the key part from the reference
        float a_scale_val = as[cx + (k_block / BLOCK_SIZE) * m];
        float b_scale_val = bs[(cy / BLOCK_SIZE) + (k_block / BLOCK_SIZE) * scale_n]; 
        
        result += block_result * a_scale_val * b_scale_val;
    }
    
    // Write result as bfloat16
    c[cx * n + cy] = (__hip_bfloat16)result;
}

// Vectorized kernel for better performance
__global__ void custom_kernel_vectorized(const __hip_fp8_e4m3_fnuz* __restrict__ a, 
                                         const __hip_fp8_e4m3_fnuz* __restrict__ b, 
                                         const float* __restrict__ as, 
                                         const float* __restrict__ bs, 
                                         __hip_bfloat16* __restrict__ c, 
                                         int m, int n, int k) {
    
    const int tid = threadIdx.x + blockDim.x * blockIdx.x;
    const int stride = blockDim.x * gridDim.x;
    const int scale_n = (n + BLOCK_SIZE - 1) / BLOCK_SIZE;
    
    // Process multiple output elements per thread
    for (int idx = tid; idx < m * n; idx += stride) {
        const int cx = idx % m;  // row index
        const int cy = idx / m;  // col index
        
        float result = 0.0f;
        
        // Process K dimension in blocks
        for (int k_block = 0; k_block < k; k_block += BLOCK_SIZE) {
            float block_result = 0.0f;
            const int k_end = min(k_block + BLOCK_SIZE, k);
            
            // Vectorized inner loop with manual unrolling
            int kk = k_block;
            for (; kk + 3 < k_end; kk += 4) {
                // Process 4 elements at once for better ILP
                float a1 = (float)a[cx + kk * m];
                float b1 = (float)b[cy + kk * n];
                float a2 = (float)a[cx + (kk + 1) * m];
                float b2 = (float)b[cy + (kk + 1) * n];
                float a3 = (float)a[cx + (kk + 2) * m];
                float b3 = (float)b[cy + (kk + 2) * n];
                float a4 = (float)a[cx + (kk + 3) * m];
                float b4 = (float)b[cy + (kk + 3) * n];
                
                block_result += a1 * b1 + a2 * b2 + a3 * b3 + a4 * b4;
            }
            
            // Handle remaining elements
            for (; kk < k_end; kk++) {
                float a_val = (float)a[cx + kk * m];
                float b_val = (float)b[cy + kk * n];
                block_result += a_val * b_val;
            }
            
            // Apply scaling
            float a_scale_val = as[cx + (k_block / BLOCK_SIZE) * m];
            float b_scale_val = bs[(cy / BLOCK_SIZE) + (k_block / BLOCK_SIZE) * scale_n];
            result += block_result * a_scale_val * b_scale_val;
        }
        
        c[cx * n + cy] = (__hip_bfloat16)result;
    }
}

// Tiled kernel with shared memory for large matrices
__global__ void custom_kernel_tiled(const __hip_fp8_e4m3_fnuz* __restrict__ a, 
                                   const __hip_fp8_e4m3_fnuz* __restrict__ b, 
                                   const float* __restrict__ as, 
                                   const float* __restrict__ bs, 
                                   __hip_bfloat16* __restrict__ c, 
                                   int m, int n, int k) {
    
    const int TILE_M = 16;
    const int TILE_N = 16;
    const int TILE_K = 16;
    
    __shared__ float a_shared[TILE_M][TILE_K];
    __shared__ float b_shared[TILE_N][TILE_K];
    
    const int bx = blockIdx.x;
    const int by = blockIdx.y;
    const int tx = threadIdx.x;
    const int ty = threadIdx.y;
    
    const int global_m = by * TILE_M + ty;
    const int global_n = bx * TILE_N + tx;
    
    if (global_m >= m || global_n >= n) return;
    
    const int scale_n = (n + BLOCK_SIZE - 1) / BLOCK_SIZE;
    float result = 0.0f;
    
    // Process K dimension in blocks
    for (int k_block = 0; k_block < k; k_block += BLOCK_SIZE) {
        float block_result = 0.0f;
        
        // Process tiles within current K block
        for (int tile_k = k_block; tile_k < min(k_block + BLOCK_SIZE, k); tile_k += TILE_K) {
            // Cooperative loading into shared memory
            if (ty < TILE_M && tx < TILE_K && tile_k + tx < k) {
                a_shared[ty][tx] = (float)a[global_m + (tile_k + tx) * m];
            }
            if (ty < TILE_N && tx < TILE_K && tile_k + tx < k) {
                b_shared[ty][tx] = (float)b[global_n + (tile_k + tx) * n];
            }
            
            __syncthreads();
            
            // Compute using shared data
            const int k_tile_end = min(TILE_K, k - tile_k);
            #pragma unroll
            for (int i = 0; i < k_tile_end; i++) {
                if (ty < TILE_M && tx < TILE_N) {
                    block_result += a_shared[ty][i] * b_shared[tx][i];
                }
            }
            
            __syncthreads();
        }
        
        // Apply block scaling
        float a_scale_val = as[global_m + (k_block / BLOCK_SIZE) * m];
        float b_scale_val = bs[(global_n / BLOCK_SIZE) + (k_block / BLOCK_SIZE) * scale_n];
        result += block_result * a_scale_val * b_scale_val;
    }
    
    c[global_m * n + global_n] = (__hip_bfloat16)result;
}

void fp8_mm(torch::Tensor a, torch::Tensor b, torch::Tensor as, torch::Tensor bs, torch::Tensor c) {
    const int m = a.size(0);
    const int n = b.size(0);  
    const int k = a.size(1);
    
    // Adaptive kernel selection based on problem size
    if (m >= 2048 && n >= 2048) {
        // Large matrices: use tiled approach
        const int TILE_SIZE = 16;
        dim3 grid((n + TILE_SIZE - 1) / TILE_SIZE, (m + TILE_SIZE - 1) / TILE_SIZE);
        dim3 block(TILE_SIZE, TILE_SIZE);
        
        custom_kernel_tiled<<<grid, block>>>(
            (__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
        );
    } else if (m * n > 100000) {
        // Medium matrices: use vectorized approach
        const int threads = 256;
        const int blocks = min(65535, (m * n + threads - 1) / threads);
        
        custom_kernel_vectorized<<<blocks, threads>>>(
            (__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
        );
    } else {
        // Small matrices: use basic 2D grid
        dim3 grid((m + 15) / 16, (n + 15) / 16);
        dim3 block(16, 16);
        
        custom_kernel_basic<<<grid, block>>>(
            (__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
        );
    }
    
    // Error checking
    hipError_t err = hipGetLastError();
    if (err != hipSuccess) {
        printf("HIP kernel error: %s\\n", hipGetErrorString(err));
    }
}
"""

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

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

def custom_kernel(data: input_t) -> output_t:
    """
    Highly optimized FP8 matrix multiplication kernel for AMD MI300X
    """
    a, b, a_scale, b_scale, c = data
    module.fp8_mm(a, b, a_scale, b_scale, c)
    return c
python submission.py
[92mSuccessfully preprocessed all matching files.[0m

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


Total number of replaced kernel launches: 3
[1/3] clang++ -MMD -MF main.o.d -DTORCH_EXTENSION_NAME=fp8_mm_optimized -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_optimized/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_optimized -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 --offload-arch=gfx942 -std=c++20 -O3 -ffast-math -c /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm_optimized/hip.hip -o hip.cuda.o 
[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_optimized.so

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

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


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