gpu-mode/popcorn-cli/load-inline-native-code

import torch
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t
CUDA_SRC = """
template <typename scalar_t>
__global__ void my_kernel(const scalar_t* __restrict__ input,
                          scalar_t* __restrict__ output,
                          int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        output[idx] = input[idx];
    }
}
torch::Tensor my_op(torch::Tensor input, torch::Tensor output) {
    int N = input.numel();
    const int threads = 256;
    const int blocks = (N + threads - 1) / threads;
    AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.scalar_type(), "my_kernel", ([&] {
        my_kernel<scalar_t><<<blocks, threads>>>(
            input.data_ptr<scalar_t>(),
            output.data_ptr<scalar_t>(),
            N
        );
    }));
    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        throw std::runtime_error(cudaGetErrorString(err));
    }
    return output;
}
"""
CPP_SRC = """
torch::Tensor my_op(torch::Tensor input, torch::Tensor output);
"""
module = load_inline(
    name='my_module',
    cpp_sources=[CPP_SRC],
    cuda_sources=[CUDA_SRC],
    functions=['my_op'],
    verbose=True,
)
def custom_kernel(data: input_t) -> output_t:
    input, output = data
    return module.my_op(input, output)

import os
os.environ['PYTORCH_ROCM_ARCH'] = 'gfx942'
os.environ['CXX'] = 'clang++'
import torch
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t
CUDA_SRC = """
#include <hip/amd_detail/amd_hip_bf16.h>
__global__ void my_kernel(const float* input, float* output, int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        output[idx] = input[idx];
    }
}
void my_op(torch::Tensor input, torch::Tensor output) {
    int N = input.numel();
    const int threads = 256;
    const int blocks = (N + threads - 1) / threads;
    my_kernel<<<blocks, threads>>>(
        input.data_ptr<float>(),
        output.data_ptr<float>(),
        N
    );
}
"""
CPP_SRC = """
void my_op(torch::Tensor input, torch::Tensor output);
"""
module = load_inline(
    name='my_module',
    cpp_sources=[CPP_SRC],
    cuda_sources=[CUDA_SRC],
    functions=['my_op'],
    verbose=True,
    extra_cuda_cflags=["--offload-arch=gfx942", "-std=c++20"],
)
def custom_kernel(data: input_t) -> output_t:
    input, output = data
    module.my_op(input, output)
    return output