submission 66304

import sys

import torch
from torch.utils.cpp_extension import load_inline

import triton
import triton.language as tl

from typing import Callable
from types import ModuleType

from task import input_t, output_t
from utils import DeterministicContext

cuda_source = """
#include <stdio.h>
#include <stdint.h>
#include <assert.h>

#define ArrayCount(array) (sizeof(array)/sizeof(array[0]))

typedef uint8_t u8;
typedef uint16_t u16;
typedef uint32_t u32;
typedef uint64_t u64;

typedef u8 b8;
typedef u32 b32;

typedef int16_t s16;
typedef int32_t s32;
typedef int64_t s64;

typedef float f32;
typedef double f64;

#define CUDACheck(fn_call) (fn_call);

template <typename T, u32 block_dim, u32 coarse_factor>
__global__ void ReduceKernel(u64 count, T *input, f64 *output)
{
    __shared__ f64 segment[block_dim];

    u64 segment_start = blockIdx.x*coarse_factor*block_dim;
    u64 index = segment_start + threadIdx.x;

    f64 sum = f64(0);
    for (u8 i = 0; i < coarse_factor; ++i)
    {
        if (index + i*block_dim < count)
            sum += f64(input[index + i*block_dim]);
    }

    segment[threadIdx.x] = sum;
    
    for (int stride = block_dim/2; stride >= 1; stride /= 2)
    {
        __syncthreads();
        if (threadIdx.x < stride)
        {
            f64 temp = f64(segment[threadIdx.x + stride]) + f64(segment[threadIdx.x]);
            segment[threadIdx.x] = temp;
        }
    }

    if (threadIdx.x == 0)
        atomicAdd(output, segment[0]);
}

template <typename T, u32 block_dim, u32 coarse_factor>
void Reduce_(u64 count, T *input, f64 *output, cudaStream_t stream)
{
    int elements_per_block = coarse_factor*block_dim;
    int grid_dim = (count + elements_per_block - 1)/(elements_per_block);
    CUDACheck((ReduceKernel<T, block_dim, coarse_factor><<<grid_dim, block_dim, 0, stream>>>(count, input, output)));
}

void Reduce(torch::Tensor input, torch::Tensor output)
{
    u64 array_count = input.numel();
    using T = f32;
    T *d_input = input.data_ptr<T>();

    f64 *d_output = output.data_ptr<f64>();
    // T *d_output = output.data_ptr<T>();

    Reduce_<T, 1024, 32>(array_count, d_input, d_output, 0);
}
"""

def load_cuda_vectorsum() -> ModuleType:
    vectorsum_module: ModuleType = None
    cpp_source = """
    #include <torch/extension.h>

    void Reduce(torch::Tensor input, torch::Tensor output);
    """
    vectorsum_module = load_inline(
        name="Reduce",
        cpp_sources=cpp_source,
        cuda_sources=cuda_source,
        extra_cuda_cflags=["-DCOMPILING_FROM_PYTORCH"],
        functions=["Reduce"], verbose=True)        

    return vectorsum_module

vectorsum_module: ModuleType = load_cuda_vectorsum()

def ref_kernel(data: input_t) -> output_t:
    """
    Reference implementation of vector sum reduction using PyTorch.
    Args:
        data: Input tensor to be reduced
    Returns:
        Tensor containing the sum of all elements
    """
    data, output = data
    # Let's be on the safe side here, and do the reduction in 64 bit
    output = data.to(torch.float64).sum().to(torch.float32)
    return output

def triton_vectorsum_(input: torch.Tensor, output: torch.Tensor) -> None:
    @triton.jit
    def triton_vectorsum_kernel(input_ptr: tl.tensor, output_ptr: tl.tensor, N: int, BLOCK_DIM: tl.constexpr):
        segment_start = tl.program_id(0)*BLOCK_DIM
        offsets = segment_start + tl.arange(0, BLOCK_DIM)
        mask = offsets < N

        a: tl.tensor = tl.load(input_ptr + offsets, mask)
        result: float = tl.sum(a)
        
        # How does Triton know that only one thread should do this operation?
        tl.atomic_add(output_ptr, result)
    
    N = input.numel()
    grid = lambda metaparams: ((N + metaparams["BLOCK_DIM"] - 1)//metaparams["BLOCK_DIM"], 1, 1)
    triton_vectorsum_kernel[grid](input, output, input.numel(), BLOCK_DIM=1024)

triton_vectorsum: Callable[[torch.Tensor, torch.Tensor], None] = torch.compile(triton_vectorsum_, mode="reduce-overhead")

def custom_kernel_(data: input_t) -> output_t:
    with DeterministicContext():
        input, output = data
        output_f64 = output.to(torch.float64).zero_()
        # output.zero_()
        # vectorsum_module.Reduce(input.to(torch.float64), output_f64)
        vectorsum_module.Reduce(input, output_f64)
        return output_f64.to(torch.float32)[0]
        # return output[0]

custom_kernel = custom_kernel_

# custom_kernel = ref_kernel
# custom_kernel = torch.compile(ref_kernel, mode="reduce-overhead")