submission 66881

import torch
from torch.utils.cpp_extension import load_inline

from task import input_t, output_t


CUDA_SRC = r"""
#include <ATen/cuda/CUDAContext.h>
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cuda.h>

#include <algorithm>
#include <cstdint>

namespace {

constexpr int kWarpSize = 32;
constexpr int kAsyncStages = 2;
constexpr int kCopiesPerThread = 4;

__device__ inline double warp_reduce_sum(double value) {
#pragma unroll
    for (int offset = kWarpSize / 2; offset > 0; offset >>= 1) {
        value += __shfl_down_sync(0xffffffff, value, offset);
    }
    return value;
}

#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
__device__ inline void cp_async_commit_group() { asm volatile("cp.async.commit_group;\n" ::); }
__device__ inline void cp_async_wait_group_0() { asm volatile("cp.async.wait_group 0;\n" ::); }
__device__ inline void cp_async_wait_all() { asm volatile("cp.async.wait_all;\n" ::); }
#else
__device__ inline void cp_async_commit_group() {}
__device__ inline void cp_async_wait_group_0() {}
__device__ inline void cp_async_wait_all() {}
#endif

__device__ inline bool async_copy_float4(float4* dst,
                                         const float4* base,
                                         const long long index,
                                         const long long vec_numel) {
    const bool valid = index < vec_numel;
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
    if (valid) {
        const unsigned int smem_addr = static_cast<unsigned int>(__cvta_generic_to_shared(dst));
        const unsigned long long gmem_addr = reinterpret_cast<unsigned long long>(base + index);
        asm volatile(
            "cp.async.cg.shared.global.L2::128B [%0], [%1], 16;\n"
            :: "r"(smem_addr), "l"(gmem_addr));
    } else {
        dst->x = 0.f;
        dst->y = 0.f;
        dst->z = 0.f;
        dst->w = 0.f;
    }
#else
    if (valid) {
        *dst = base[index];
    } else {
        dst->x = 0.f;
        dst->y = 0.f;
        dst->z = 0.f;
        dst->w = 0.f;
    }
#endif
    return valid;
}

template <int Threads>
__global__ void vector_sum_kernel_fallback(const float* __restrict__ input,
                                           double* __restrict__ accumulator,
                                           const long long numel,
                                           const int use_vectorized) {
    extern __shared__ double warp_sums[];
    float thread_sum_f = 0.0f;

    const long long thread_linear =
        static_cast<long long>(blockIdx.x) * blockDim.x + threadIdx.x;
    const long long stride =
        static_cast<long long>(blockDim.x) * gridDim.x;

    if (use_vectorized) {
        const float4* __restrict__ input4 = reinterpret_cast<const float4*>(input);
        const long long vec_numel = numel >> 2;
        constexpr int kVectorUnroll = 8;
        float partials[kVectorUnroll] = {0.0f};

        long long index = thread_linear;
        const long long vec_stride = stride;
        const long long unroll_stride = vec_stride * kVectorUnroll;
        for (; index + static_cast<long long>(kVectorUnroll - 1) * vec_stride < vec_numel;
             index += unroll_stride) {
#pragma unroll
            for (int unroll = 0; unroll < kVectorUnroll; ++unroll) {
                const long long load_index =
                    index + static_cast<long long>(unroll) * vec_stride;
                float4 v;
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
                asm volatile("ld.global.cs.v4.f32 {%0, %1, %2, %3}, [%4];"
                             : "=f"(v.x), "=f"(v.y), "=f"(v.z), "=f"(v.w)
                             : "l"(input4 + load_index));
#elif defined(__CUDA_ARCH__)
                v = __ldg(input4 + load_index);
#else
                v = input4[load_index];
#endif
                const float pair0 = v.x + v.y;
                const float pair1 = v.z + v.w;
                partials[unroll] += pair0 + pair1;
            }
        }
        for (; index < vec_numel; index += vec_stride) {
            float4 v;
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
            asm volatile("ld.global.cs.v4.f32 {%0, %1, %2, %3}, [%4];"
                         : "=f"(v.x), "=f"(v.y), "=f"(v.z), "=f"(v.w)
                         : "l"(input4 + index));
#elif defined(__CUDA_ARCH__)
            v = __ldg(input4 + index);
#else
            v = input4[index];
#endif
            const float pair0 = v.x + v.y;
            const float pair1 = v.z + v.w;
            partials[0] += pair0 + pair1;
        }

        const long long processed = vec_numel << 2;
        float tail_sum = 0.f;
        for (long long idx = processed + thread_linear; idx < numel; idx += stride) {
            tail_sum += input[idx];
        }
        float vector_sum = tail_sum;
#pragma unroll
        for (int unroll = 0; unroll < kVectorUnroll; ++unroll) {
            vector_sum += partials[unroll];
        }
        thread_sum_f += vector_sum;
    } else {
        float accum0 = 0.0f;
        float accum1 = 0.0f;
        bool toggle = false;
        for (long long idx = thread_linear; idx < numel; idx += stride) {
            const float value = input[idx];
            if (!toggle) {
                accum0 += value;
            } else {
                accum1 += value;
            }
            toggle = !toggle;
        }
        thread_sum_f += accum0 + accum1;
    }

    const int lane = threadIdx.x & (kWarpSize - 1);
    const int warp_id = threadIdx.x >> 5;
    const int warp_count = (Threads + kWarpSize - 1) / kWarpSize;

    double thread_sum = static_cast<double>(thread_sum_f);
    double warp_sum = warp_reduce_sum(thread_sum);
    if (lane == 0) {
        warp_sums[warp_id] = warp_sum;
    }
    __syncthreads();

    if (warp_id == 0) {
        double block_sum = lane < warp_count ? warp_sums[lane] : 0.0;
        block_sum = warp_reduce_sum(block_sum);
        if (lane == 0) {
            atomicAdd(accumulator, block_sum);
        }
    }
}

__global__ void vector_sum_kernel_async(const float* __restrict__ input,
                                        double* __restrict__ accumulator,
                                        const long long numel,
                                        const long long vec_numel) {
    extern __shared__ unsigned char shared_storage[];
    float4* tile_storage = reinterpret_cast<float4*>(shared_storage);
    const int warp_count = (blockDim.x + kWarpSize - 1) / kWarpSize;
    const int warp_stride = warp_count * kWarpSize;
    const int stage_area = warp_stride * kCopiesPerThread;
    double* warp_sums =
        reinterpret_cast<double*>(tile_storage + kAsyncStages * stage_area);

    const long long thread_linear =
        static_cast<long long>(blockIdx.x) * blockDim.x + threadIdx.x;
    const long long stride =
        static_cast<long long>(blockDim.x) * gridDim.x;
    const float4* __restrict__ input4 = reinterpret_cast<const float4*>(input);

    const long long iteration_stride = stride * kCopiesPerThread;
    const long long total_iters =
        (vec_numel + iteration_stride - 1) / iteration_stride;
    float thread_acc[kCopiesPerThread] = {};

    int stage = 0;
    const int lane = threadIdx.x & (kWarpSize - 1);
    const int warp_id = threadIdx.x >> 5;

    if (total_iters > 0) {
        const long long initial_base = thread_linear;
#pragma unroll
        for (int copy = 0; copy < kCopiesPerThread; ++copy) {
            const long long index = initial_base + static_cast<long long>(copy) * stride;
            float4* dst =
                tile_storage + stage * stage_area + copy * warp_stride +
                warp_id * kWarpSize + lane;
            async_copy_float4(dst, input4, index, vec_numel);
        }
        cp_async_commit_group();

        for (long long iteration = 0; iteration < total_iters; ++iteration) {
            cp_async_wait_group_0();

            const float4* stage_base =
                tile_storage + stage * stage_area + warp_id * kWarpSize;
            const long long iteration_base =
                thread_linear + iteration * iteration_stride;
#pragma unroll
            for (int copy = 0; copy < kCopiesPerThread; ++copy) {
                const long long vec_index =
                    iteration_base + static_cast<long long>(copy) * stride;
                if (vec_index < vec_numel) {
                    const float4 value =
                        stage_base[copy * warp_stride + lane];
                    // Keep per-copy accumulators to increase ILP for the reduction.
                    const float pair0 = value.x + value.y;
                    const float pair1 = value.z + value.w;
                    const float partial = pair0 + pair1;
                    thread_acc[copy] += partial;
                }
            }

            const long long next_base =
                thread_linear + (iteration + 1) * iteration_stride;
            const int next_stage = stage ^ 1;

            if (iteration + 1 < total_iters) {
#pragma unroll
                for (int copy = 0; copy < kCopiesPerThread; ++copy) {
                    float4* next_dst =
                        tile_storage + next_stage * stage_area +
                        copy * warp_stride + warp_id * kWarpSize + lane;
                    const long long next_index =
                        next_base + static_cast<long long>(copy) * stride;
                    async_copy_float4(next_dst, input4, next_index, vec_numel);
                }
                cp_async_commit_group();
            }

            stage = next_stage;
        }
        cp_async_wait_all();
    }

    const long long processed = vec_numel << 2;
    float tail_sum = 0.f;
    for (long long idx = processed + thread_linear; idx < numel; idx += stride) {
        tail_sum += input[idx];
    }
    float thread_sum_f = tail_sum;
#pragma unroll
    for (int copy = 0; copy < kCopiesPerThread; ++copy) {
        thread_sum_f += thread_acc[copy];
    }

    double thread_sum = static_cast<double>(thread_sum_f);
    double warp_sum = warp_reduce_sum(thread_sum);
    if (lane == 0) {
        warp_sums[warp_id] = warp_sum;
    }
    __syncthreads();

    if (warp_id == 0) {
        const int warp_count_shared = (blockDim.x + kWarpSize - 1) / kWarpSize;
        double block_sum = lane < warp_count_shared ? warp_sums[lane] : 0.0;
        block_sum = warp_reduce_sum(block_sum);
        if (lane == 0) {
            atomicAdd(accumulator, block_sum);
        }
    }
}

template <int Threads>
void launch_fallback(const float* input,
                     double* accumulator,
                     const long long numel,
                     const int use_vectorized,
                     const int blocks,
                     cudaStream_t stream) {
    const int warp_count = (Threads + kWarpSize - 1) / kWarpSize;
    const size_t shared_bytes = static_cast<size_t>(warp_count) * sizeof(double);
    vector_sum_kernel_fallback<Threads>
        <<<blocks, Threads, shared_bytes, stream>>>(
            input, accumulator, numel, use_vectorized);
}

template <int Threads, int Unroll>
__global__ void vector_sum_kernel_stream(const float* __restrict__ input,
                                         double* __restrict__ accumulator,
                                         const long long numel) {
    extern __shared__ double warp_sums[];

    const long long thread_linear =
        static_cast<long long>(blockIdx.x) * blockDim.x + threadIdx.x;
    const long long stride =
        static_cast<long long>(blockDim.x) * gridDim.x;
    const float4* __restrict__ input4 = reinterpret_cast<const float4*>(input);
    const long long vec_numel = numel >> 2;

    float thread_acc_even = 0.0f;
    float thread_acc_odd = 0.0f;

    long long index = thread_linear;
    const long long vec_stride = stride;
    const long long unroll_step =
        vec_stride * static_cast<long long>(Unroll);
    for (; index + static_cast<long long>(Unroll - 1) * vec_stride < vec_numel;
         index += unroll_step) {
#pragma unroll
        for (int i = 0; i < Unroll; ++i) {
            const long long load_index =
                index + static_cast<long long>(i) * vec_stride;
            float4 v;
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
            asm volatile("ld.global.cs.v4.f32 {%0, %1, %2, %3}, [%4];"
                         : "=f"(v.x), "=f"(v.y), "=f"(v.z), "=f"(v.w)
                         : "l"(input4 + load_index));
#elif defined(__CUDA_ARCH__)
            v = __ldg(input4 + load_index);
#else
            v = input4[load_index];
#endif
            const float pair0 = v.x + v.y;
            const float pair1 = v.z + v.w;
            const float partial = pair0 + pair1;
            if ((i & 1) == 0) {
                thread_acc_even += partial;
            } else {
                thread_acc_odd += partial;
            }
        }
    }

    for (; index < vec_numel; index += vec_stride) {
        float4 v;
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
        asm volatile("ld.global.cs.v4.f32 {%0, %1, %2, %3}, [%4];"
                     : "=f"(v.x), "=f"(v.y), "=f"(v.z), "=f"(v.w)
                     : "l"(input4 + index));
#elif defined(__CUDA_ARCH__)
        v = __ldg(input4 + index);
#else
        v = input4[index];
#endif
        const float pair0 = v.x + v.y;
        const float pair1 = v.z + v.w;
        thread_acc_even += pair0 + pair1;
    }

    const long long processed = vec_numel << 2;
    float tail_sum = 0.f;
    for (long long idx = processed + thread_linear; idx < numel; idx += stride) {
        tail_sum += input[idx];
    }

    float thread_sum_f = tail_sum;
    thread_sum_f += thread_acc_even + thread_acc_odd;

    double thread_sum = static_cast<double>(thread_sum_f);
    const int lane = threadIdx.x & (kWarpSize - 1);
    const int warp_id = threadIdx.x >> 5;
    const int warp_count = (Threads + kWarpSize - 1) / kWarpSize;

    double warp_sum = warp_reduce_sum(thread_sum);
    if (lane == 0) {
        warp_sums[warp_id] = warp_sum;
    }
    __syncthreads();

    if (warp_id == 0) {
        double block_sum = lane < warp_count ? warp_sums[lane] : 0.0;
        block_sum = warp_reduce_sum(block_sum);
        if (lane == 0) {
            atomicAdd(accumulator, block_sum);
        }
    }
}

template <int Threads, int Unroll>
void launch_stream(const float* input,
                   double* accumulator,
                   const long long numel,
                   const int blocks,
                   cudaStream_t stream) {
    const int warp_count = (Threads + kWarpSize - 1) / kWarpSize;
    const size_t shared_bytes =
        static_cast<size_t>(warp_count) * sizeof(double);
    vector_sum_kernel_stream<Threads, Unroll>
        <<<blocks, Threads, shared_bytes, stream>>>(input, accumulator, numel);
}

}  // namespace

static inline int select_block_dim(const long long numel) {
    if (numel >= 1024) return 1024;
    if (numel >= 512) return 512;
    if (numel >= 256) return 256;
    if (numel >= 128) return 128;
    if (numel >= 64) return 64;
    if (numel >= 32) return 32;
    if (numel >= 16) return 16;
    if (numel >= 8) return 8;
    if (numel >= 4) return 4;
    return std::max<int>(static_cast<int>(numel), 1);
}

void vector_sum_launcher(const at::Tensor& input,
                         const at::Tensor& accumulator) {
    TORCH_CHECK(input.is_cuda(), "input must be on CUDA");
    TORCH_CHECK(accumulator.is_cuda(), "accumulator must be on CUDA");
    TORCH_CHECK(input.dtype() == torch::kFloat32, "input must be float32");
    TORCH_CHECK(accumulator.dtype() == torch::kFloat64, "accumulator must be float64");
    TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
    TORCH_CHECK(accumulator.is_contiguous(), "accumulator must be contiguous");
    TORCH_CHECK(accumulator.numel() == 1, "accumulator must have exactly one element");

    const long long numel = input.numel();
    if (numel == 0) {
        return;
    }

    const bool vector_aligned =
        (numel >= 4) &&
        ((reinterpret_cast<std::uintptr_t>(input.data_ptr<float>()) & 0xF) == 0);

    const auto* device_props = at::cuda::getCurrentDeviceProperties();
    const int supports_cp_async = device_props->major >= 8;
    const int sm_count = device_props->multiProcessorCount;
    const long long vec_numel = numel >> 2;

    int threads = select_block_dim(numel);
    const long long elements_per_thread = vector_aligned ? 4LL : 1LL;
    long long tentative_blocks =
        (numel + threads * elements_per_thread - 1) /
        (threads * elements_per_thread);
    tentative_blocks = std::max<long long>(tentative_blocks, 1LL);
    const int max_grid = 65535;
    int blocks = static_cast<int>(std::min<long long>(tentative_blocks, max_grid));

    const auto stream = at::cuda::getCurrentCUDAStream();

    constexpr int kStreamThreads = 640;
    constexpr int kStreamUnroll = 5;
    const bool use_stream = vector_aligned && numel >= 4096;
    if (use_stream) {
        threads = kStreamThreads;
        const long long stream_elements_per_thread =
            4LL * static_cast<long long>(kStreamUnroll);
        long long stream_blocks =
            (numel + threads * stream_elements_per_thread - 1) /
            (threads * stream_elements_per_thread);
        stream_blocks = std::max<long long>(stream_blocks, 1LL);
        stream_blocks = std::min<long long>(stream_blocks, max_grid);

        const int warp_count_stream = (threads + kWarpSize - 1) / kWarpSize;
        const size_t shared_bytes_stream =
            static_cast<size_t>(warp_count_stream) * sizeof(double);
        int occupancy_blocks_per_sm = 0;
        if (cudaOccupancyMaxActiveBlocksPerMultiprocessor(
                &occupancy_blocks_per_sm,
                vector_sum_kernel_stream<kStreamThreads, kStreamUnroll>,
                threads,
                shared_bytes_stream) == cudaSuccess &&
            occupancy_blocks_per_sm > 0) {
            const long long occupancy_blocks =
                static_cast<long long>(sm_count) *
                static_cast<long long>(occupancy_blocks_per_sm);
            stream_blocks = std::min<long long>(
                stream_blocks,
                std::min<long long>(occupancy_blocks, max_grid));
            stream_blocks = std::max(stream_blocks, 1LL);
        }

        launch_stream<kStreamThreads, kStreamUnroll>(
            input.data_ptr<float>(),
            accumulator.data_ptr<double>(),
            numel,
            static_cast<int>(stream_blocks),
            stream);
        C10_CUDA_KERNEL_LAUNCH_CHECK();
        return;
    }

    const bool use_async =
        supports_cp_async && vector_aligned && vec_numel > 0 && numel >= 1024;

    if (use_async) {
        threads = 768;
        const long long async_elements_per_thread =
            4LL * static_cast<long long>(kCopiesPerThread);
        long long async_blocks =
            (numel + threads * async_elements_per_thread - 1) /
            (threads * async_elements_per_thread);
        async_blocks = std::max<long long>(async_blocks, 1LL);
        blocks = static_cast<int>(std::min<long long>(async_blocks, max_grid));

        const int warp_count = (threads + kWarpSize - 1) / kWarpSize;
        const size_t warp_stride = static_cast<size_t>(warp_count) * kWarpSize;
        const size_t stage_bytes =
            static_cast<size_t>(kAsyncStages) * warp_stride *
            static_cast<size_t>(kCopiesPerThread) * sizeof(float4);
        const size_t warp_bytes =
            static_cast<size_t>(warp_count) * sizeof(double);
        const size_t shared_bytes = stage_bytes + warp_bytes;

        int max_optin_smem = 0;
        C10_CUDA_CHECK(cudaDeviceGetAttribute(
            &max_optin_smem,
            cudaDevAttrMaxSharedMemoryPerBlockOptin,
            at::cuda::current_device()));
        TORCH_CHECK(
            shared_bytes <= static_cast<size_t>(max_optin_smem),
            "vector_sum_kernel_async requires ",
            shared_bytes,
            " bytes of shared memory, exceeding device limit of ",
            max_optin_smem);
        C10_CUDA_CHECK(cudaFuncSetAttribute(
            vector_sum_kernel_async,
            cudaFuncAttributeMaxDynamicSharedMemorySize,
            static_cast<int>(shared_bytes)));

        int occupancy_blocks_per_sm = 0;
        if (cudaOccupancyMaxActiveBlocksPerMultiprocessor(
                &occupancy_blocks_per_sm,
                vector_sum_kernel_async,
                threads,
                static_cast<size_t>(shared_bytes)) == cudaSuccess &&
            occupancy_blocks_per_sm > 0) {
            const long long occupancy_blocks =
                static_cast<long long>(sm_count) *
                static_cast<long long>(occupancy_blocks_per_sm);
            blocks = static_cast<int>(std::min<long long>(
                blocks,
                std::min<long long>(occupancy_blocks, max_grid)));
            blocks = std::max(blocks, 1);
        }
        vector_sum_kernel_async<<<blocks, threads, shared_bytes, stream>>>(
            input.data_ptr<float>(),
            accumulator.data_ptr<double>(),
            numel,
            vec_numel);
    } else {
        switch (threads) {
            case 1024:
                launch_fallback<1024>(
                    input.data_ptr<float>(),
                    accumulator.data_ptr<double>(),
                    numel,
                    static_cast<int>(vector_aligned),
                    blocks,
                    stream);
                break;
            case 512:
                launch_fallback<512>(
                    input.data_ptr<float>(),
                    accumulator.data_ptr<double>(),
                    numel,
                    static_cast<int>(vector_aligned),
                    blocks,
                    stream);
                break;
            case 256:
                launch_fallback<256>(
                    input.data_ptr<float>(),
                    accumulator.data_ptr<double>(),
                    numel,
                    static_cast<int>(vector_aligned),
                    blocks,
                    stream);
                break;
            case 128:
                launch_fallback<128>(
                    input.data_ptr<float>(),
                    accumulator.data_ptr<double>(),
                    numel,
                    static_cast<int>(vector_aligned),
                    blocks,
                    stream);
                break;
            case 64:
                launch_fallback<64>(
                    input.data_ptr<float>(),
                    accumulator.data_ptr<double>(),
                    numel,
                    static_cast<int>(vector_aligned),
                    blocks,
                    stream);
                break;
            case 32:
                launch_fallback<32>(
                    input.data_ptr<float>(),
                    accumulator.data_ptr<double>(),
                    numel,
                    static_cast<int>(vector_aligned),
                    blocks,
                    stream);
                break;
            case 16:
                launch_fallback<16>(
                    input.data_ptr<float>(),
                    accumulator.data_ptr<double>(),
                    numel,
                    static_cast<int>(vector_aligned),
                    blocks,
                    stream);
                break;
            case 8:
                launch_fallback<8>(
                    input.data_ptr<float>(),
                    accumulator.data_ptr<double>(),
                    numel,
                    static_cast<int>(vector_aligned),
                    blocks,
                    stream);
                break;
            case 4:
                launch_fallback<4>(
                    input.data_ptr<float>(),
                    accumulator.data_ptr<double>(),
                    numel,
                    static_cast<int>(vector_aligned),
                    blocks,
                    stream);
                break;
            case 3:
                launch_fallback<3>(
                    input.data_ptr<float>(),
                    accumulator.data_ptr<double>(),
                    numel,
                    static_cast<int>(vector_aligned),
                    blocks,
                    stream);
                break;
            case 2:
                launch_fallback<2>(
                    input.data_ptr<float>(),
                    accumulator.data_ptr<double>(),
                    numel,
                    static_cast<int>(vector_aligned),
                    blocks,
                    stream);
                break;
            case 1:
                launch_fallback<1>(
                    input.data_ptr<float>(),
                    accumulator.data_ptr<double>(),
                    numel,
                    static_cast<int>(vector_aligned),
                    blocks,
                    stream);
                break;
            default:
                TORCH_CHECK(false, "Unsupported block dimension");
        }
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("vector_sum_launcher", &vector_sum_launcher, "vector sum reduction launcher");
}
"""


_module = load_inline(
    name="vectorsum_cuda_kernel",
    cpp_sources="",
    cuda_sources=CUDA_SRC,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


def custom_kernel(data: input_t) -> output_t:
    tensor, output = data  # type: ignore[assignment]

    if tensor.device.type != "cuda":
        raise RuntimeError("custom_kernel expects CUDA tensors")

    if not tensor.is_contiguous():
        tensor = tensor.contiguous()

    accumulator = torch.zeros(
        (),
        device=tensor.device,
        dtype=torch.float64,
    )

    _module.vector_sum_launcher(tensor, accumulator)

    result = accumulator.to(dtype=torch.float32)
    output.view(-1)[0] = result
    return result