#!POPCORN leaderboard amd-identity
from task import input_t, output_t
import torch
from torch.utils.cpp_extension import load_inline
import time
import os
import sys

if "PYTORCH_ROCM_ARCH" not in os.environ:
    os.environ["PYTORCH_ROCM_ARCH"] = "gfx942:xnack-"

kernel_cpp = r"""
#include <hip/hip_fp16.h>
#include <hip/hip_fp8.h>
#include <pybind11/pybind11.h>

using fp16 = __hip_bfloat16;

// HIP-style kernel (same syntax as CUDA)
__global__ void  identity_kernel(const fp16* inp, fp16* out, int size) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < size) {
        out[i] = inp[i];
    }
}

// Wrapper using torch::Tensor
void identity_hip(uintptr_t inp, uintptr_t out, int size) {
    const int threads = 1024;
    const int blocks = (size + threads - 1) / threads;

    fp16* input = (fp16*)inp;
    fp16* output = (fp16*) out;

    identity_kernel<<<blocks, threads>>>(
        input,
        output,
        size
    );
}

PYBIND11_MODULE(identity, m) {
    m.def("identity", &identity_hip, "Identity (HIP)");
}
"""


hip_module = load_inline(
    name="identity",
    cpp_sources="",
    cuda_sources=kernel_cpp.replace('@', chr(92)),
    with_cuda=True,
    verbose=False,
    extra_cuda_cflags=["-std=c++20"],
    no_implicit_headers=True,
)

first = True

def custom_kernel(data: input_t) -> output_t:
    """
    Reference implementation of block-scale fp8 gemm
    Args:
        data: Tuple that expands to:
            a: torch.Tensor[float8_e4m3fnuz] of shape [m, k],
            b: torch.Tensor[float8_e4m3fnuz] of shape [n, k],
            a_scale: torch.Tensor[float32] of shape [m, k // BLOCK],
            b_scale: torch.Tensor[float32] of shape [n // BLOCK, k // BLOCK],
            c: torch.Tensor[bfloat16] of shape [m, n]
    Returns:
        Tensor containing output in bf16
    """

    # global first
    # if first:
    #     print("executing on", torch.cuda.get_device_name(), file=sys.stderr)
    #     first = False

    # a, b, a_scale, b_scale, c = data

    # m, n = c.shape
    # k = a.shape[1]

    # hip_module.fp8(
    #     a.data_ptr(),
    #     b.data_ptr(),
    #     a_scale.data_ptr(),
    #     b_scale.data_ptr(),
    #     c.data_ptr(),
    #     m,
    #     n,
    #     k,
    # )
    
    inp, out = data

    size = inp.numel()
    hip_module.identity(inp.data_ptr(), out.data_ptr(), size)

    return out
python eval.py benchmark /tmp/tmp2c77evih
[92mSuccessfully preprocessed all matching files.[0m
