submission 773732

#!POPCORN leaderboard vectoradd_v2
#!POPCORN gpu A100

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 vectoradd_kernel_basic(
    const scalar_t* __restrict__ A,
    const scalar_t* __restrict__ B,
    scalar_t* __restrict__ output,
    int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        output[idx] = A[idx] + B[idx];
    }
}

// 优化版本1：向量化内存访问（使用float4一次处理4个float）
// 适用于float类型，提高内存带宽利用率
__global__ void vectoradd_kernel_vectorized(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ output,
    int N) {
    int idx = (blockIdx.x * blockDim.x + threadIdx.x) * 4;
    
    if (idx + 3 < N) {
        // 使用float4进行向量化加载/存储，一次处理4个元素
        float4 a_val = reinterpret_cast<const float4*>(A)[idx / 4];
        float4 b_val = reinterpret_cast<const float4*>(B)[idx / 4];
        float4 c_val;
        c_val.x = a_val.x + b_val.x;
        c_val.y = a_val.y + b_val.y;
        c_val.z = a_val.z + b_val.z;
        c_val.w = a_val.w + b_val.w;
        reinterpret_cast<float4*>(output)[idx / 4] = c_val;
    } else {
        // 处理剩余元素
        for (int i = 0; i < 4 && idx + i < N; i++) {
            output[idx + i] = A[idx + i] + B[idx + i];
        }
    }
}

// 优化版本2：循环展开 - 每个线程处理多个元素
// 减少线程块数量，提高每个线程的工作量
template <typename scalar_t, int UNROLL_FACTOR>
__global__ void vectoradd_kernel_unrolled(
    const scalar_t* __restrict__ A,
    const scalar_t* __restrict__ B,
    scalar_t* __restrict__ output,
    int N) {
    int base_idx = blockIdx.x * blockDim.x * UNROLL_FACTOR + threadIdx.x;
    int stride = blockDim.x;
    
    #pragma unroll
    for (int i = 0; i < UNROLL_FACTOR; i++) {
        int idx = base_idx + i * stride;
        if (idx < N) {
            output[idx] = A[idx] + B[idx];
        }
    }
}

// 优化版本3：结合向量化和循环展开
template <int VEC_SIZE, int UNROLL_FACTOR>
__global__ void vectoradd_kernel_optimized(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ output,
    int N) {
    
    const int ELEMENTS_PER_THREAD = VEC_SIZE * UNROLL_FACTOR;
    int thread_id = blockIdx.x * blockDim.x + threadIdx.x;
    int base_idx = thread_id * ELEMENTS_PER_THREAD;
    
    #pragma unroll
    for (int u = 0; u < UNROLL_FACTOR; u++) {
        int vec_idx = base_idx + u * VEC_SIZE;
        
        if (vec_idx + VEC_SIZE - 1 < N) {
            // 向量化加载
            float4 a_val = reinterpret_cast<const float4*>(A)[vec_idx / 4];
            float4 b_val = reinterpret_cast<const float4*>(B)[vec_idx / 4];
            float4 c_val;
            c_val.x = a_val.x + b_val.x;
            c_val.y = a_val.y + b_val.y;
            c_val.z = a_val.z + b_val.z;
            c_val.w = a_val.w + b_val.w;
            reinterpret_cast<float4*>(output)[vec_idx / 4] = c_val;
        } else {
            // 处理剩余元素
            for (int v = 0; v < VEC_SIZE && vec_idx + v < N; v++) {
                output[vec_idx + v] = A[vec_idx + v] + B[vec_idx + v];
            }
        }
    }
}

torch::Tensor vectoradd_op(torch::Tensor A, torch::Tensor B, torch::Tensor output) {
    int N = A.numel();
    const int threads = 256;
    
    // 根据数据类型选择不同的kernel
    if (A.scalar_type() == torch::kFloat32) {
        // 对于float类型，使用优化版本（向量化+循环展开）
        // 每个线程处理16个元素（4个float4）
        const int ELEMENTS_PER_THREAD = 16;
        const int blocks = (N + threads * ELEMENTS_PER_THREAD - 1) / (threads * ELEMENTS_PER_THREAD);
        
        vectoradd_kernel_optimized<4, 4><<<blocks, threads>>>(
            A.data_ptr<float>(),
            B.data_ptr<float>(),
            output.data_ptr<float>(),
            N
        );
    } else {
        // 其他类型使用循环展开版本
        const int UNROLL_FACTOR = 4;
        const int blocks = (N + threads * UNROLL_FACTOR - 1) / (threads * UNROLL_FACTOR);
        
        AT_DISPATCH_FLOATING_TYPES_AND_HALF(A.scalar_type(), "vectoradd_kernel_unrolled", ([&] {
            vectoradd_kernel_unrolled<scalar_t, UNROLL_FACTOR><<<blocks, threads>>>(
                A.data_ptr<scalar_t>(),
                B.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 vectoradd_op(torch::Tensor A, torch::Tensor B, torch::Tensor output);
"""

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

def custom_kernel(data: input_t) -> output_t:
    A, B, output = data
    return module.vectoradd_op(A, B, output)