submission 782505

#!POPCORN leaderboard matmul_v2
#!POPCORN gpu B200

import torch
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t

CUDA_SRC = r"""
#include <cuda_fp16.h>
#include <mma.h>

using namespace nvcuda;

#define WMMA_M 16
#define WMMA_N 16
#define WMMA_K 16

__global__ void wmma_matmul_kernel(const __half* __restrict__ A,
                                   const __half* __restrict__ B,
                                   __half* __restrict__ C,
                                   int M, int N, int K) {
    int warpRow = blockIdx.y * WMMA_M;
    int warpCol = blockIdx.x * WMMA_N;

    wmma::fragment<wmma::matrix_a, WMMA_M, WMMA_N, WMMA_K, __half, wmma::row_major> a_frag;
    wmma::fragment<wmma::matrix_b, WMMA_M, WMMA_N, WMMA_K, __half, wmma::row_major> b_frag;
    wmma::fragment<wmma::accumulator, WMMA_M, WMMA_N, WMMA_K, float> acc_frag;
    wmma::fragment<wmma::accumulator, WMMA_M, WMMA_N, WMMA_K, __half> c_frag;

    wmma::fill_fragment(acc_frag, 0.0f);

    for (int k = 0; k < K; k += WMMA_K) {
        wmma::load_matrix_sync(a_frag, A + warpRow * K + k, K);
        wmma::load_matrix_sync(b_frag, B + k * N + warpCol, N);
        wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag);
    }

    for (int i = 0; i < acc_frag.num_elements; i++) {
        c_frag.x[i] = __float2half(acc_frag.x[i]);
    }

    wmma::store_matrix_sync(C + warpRow * N + warpCol, c_frag, N, wmma::mem_row_major);
}

torch::Tensor naive_matmul(torch::Tensor A, torch::Tensor B, torch::Tensor C) {
    const int M = A.size(0);
    const int K = A.size(1);
    const int N = B.size(1);

    TORCH_CHECK(A.scalar_type() == at::kHalf, "Only fp16 (Half) is supported");
    TORCH_CHECK(M % WMMA_M == 0 && N % WMMA_N == 0 && K % WMMA_K == 0,
                "M, N, K must be multiples of 16");

    dim3 threads(32, 1, 1);
    dim3 blocks(N / WMMA_N, M / WMMA_M, 1);

    wmma_matmul_kernel<<<blocks, threads>>>(
        reinterpret_cast<const __half*>(A.data_ptr<at::Half>()),
        reinterpret_cast<const __half*>(B.data_ptr<at::Half>()),
        reinterpret_cast<__half*>(C.data_ptr<at::Half>()),
        M, N, K
    );

    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        throw std::runtime_error(cudaGetErrorString(err));
    }
    return C;
}
"""

CPP_SRC = r"""
torch::Tensor naive_matmul(torch::Tensor A, torch::Tensor B, torch::Tensor C);
"""

module = load_inline(
    name='wmma_matmul_module',
    cpp_sources=[CPP_SRC],
    cuda_sources=[CUDA_SRC],
    functions=['naive_matmul'],
    verbose=True,
)

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