submission 66642

# fast_grayscale.py
import torch
from torch.utils.cpp_extension import load_inline

_src = r"""
#include <ATen/ATen.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_bf16.h>
#include <cuda_fp16.h>

template <typename scalar_t>
struct Traits;

template <>
struct Traits<float> {
    using scalar_t = float;
    using acc_t    = float;
    __device__ static inline acc_t to_acc(scalar_t x){ return x; }
    __device__ static inline scalar_t from_acc(acc_t x){ return x; }
};

template <>
struct Traits<double> {
    using scalar_t = double;
    using acc_t    = double;
    __device__ static inline acc_t to_acc(scalar_t x){ return x; }
    __device__ static inline scalar_t from_acc(acc_t x){ return x; }
};

template <>
struct Traits<at::Half> {
    using scalar_t = at::Half;
    using acc_t    = float;
    __device__ static inline acc_t to_acc(scalar_t x){ return __half2float(*reinterpret_cast<const __half*>(&x)); }
    __device__ static inline scalar_t from_acc(acc_t x){
        __half h = __float2half(x);
        return *reinterpret_cast<scalar_t*>(&h);
    }
};

template <>
struct Traits<at::BFloat16> {
    using scalar_t = at::BFloat16;
    using acc_t    = float;
    __device__ static inline acc_t to_acc(scalar_t x){ return __bfloat162float(*reinterpret_cast<const __nv_bfloat16*>(&x)); }
    __device__ static inline scalar_t from_acc(acc_t x){
        __nv_bfloat16 h = __float2bfloat16(x);
        return *reinterpret_cast<scalar_t*>(&h);
    }
};

// grid-stride loop; each thread handles multiple pixels.
// We assume the last dimension is 3 (RGB) and contiguous memory.
template <typename scalar_t>
__global__ void grayscale_kernel(const scalar_t* __restrict__ in,
                                 scalar_t* __restrict__ out,
                                 size_t n_pixels)
{
    using T = Traits<scalar_t>;
    using acc_t = typename T::acc_t;

    const acc_t wr = (acc_t)0.2989f;
    const acc_t wg = (acc_t)0.5870f;
    const acc_t wb = (acc_t)0.1140f;

    // Process 4 pixels per thread for better memory throughput
    size_t idx = (blockIdx.x * blockDim.x + threadIdx.x) * 4;

    // main unrolled loop
    for (; idx + 3 < n_pixels; idx += gridDim.x * blockDim.x * 4) {
        #pragma unroll
        for (int k = 0; k < 4; ++k) {
            size_t p = idx + k;
            size_t base = p * 3;
            acc_t r = T::to_acc(in[base + 0]);
            acc_t g = T::to_acc(in[base + 1]);
            acc_t b = T::to_acc(in[base + 2]);
            out[p] = T::from_acc(r * wr + g * wg + b * wb);
        }
    }

    // tail
    for (; idx < n_pixels; ++idx) {
        size_t base = idx * 3;
        acc_t r = T::to_acc(in[base + 0]);
        acc_t g = T::to_acc(in[base + 1]);
        acc_t b = T::to_acc(in[base + 2]);
        out[idx] = T::from_acc(r * wr + g * wg + b * wb);
    }
}

at::Tensor grayscale_inline_cuda(const at::Tensor& input, at::Tensor output) {
    TORCH_CHECK(input.is_cuda(), "input must be CUDA");
    TORCH_CHECK(output.is_cuda(), "output must be CUDA");
    TORCH_CHECK(input.scalar_type() == output.scalar_type(), "dtype mismatch");
    TORCH_CHECK(input.is_contiguous(), "input must be contiguous (NHWC with C=3)");
    TORCH_CHECK(output.is_contiguous(), "output must be contiguous");
    TORCH_CHECK(input.size(-1) == 3, "last dimension must be 3 (RGB)");
    TORCH_CHECK(output.numel() == input.numel() / 3, "output must have N*H*W elements");

    const auto n_pixels = static_cast<size_t>(output.numel());
    const int threads = 256;
    const int blocks = std::min<int>((int)((n_pixels + (threads*4 - 1)) / (threads*4)), 32768);

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

    AT_DISPATCH_FLOATING_TYPES_AND2(at::kHalf, at::kBFloat16, input.scalar_type(), "grayscale_inline_cuda", [&](){
        using scalar_t_ = scalar_t;
        const scalar_t_* in_ptr  = input.data_ptr<scalar_t_>();
        scalar_t_* out_ptr       = output.data_ptr<scalar_t_>();
        grayscale_kernel<scalar_t_><<<blocks, threads, 0, stream>>>(in_ptr, out_ptr, n_pixels);
    });

    return output;
}
"""

_cpp = r"""
at::Tensor grayscale_inline_cuda(const at::Tensor& input, at::Tensor output);
"""

_mod = load_inline(
    name="grayscale_inline_cuda",
    cpp_sources=_cpp,
    cuda_sources=_src,
    functions=["grayscale_inline_cuda"],
    verbose=False,
)

def custom_kernel(data):
    x, y = data  # x: [*, 3], y: [*]
    # Make sure we have NHWC-with-3 contiguous; copy-as-needed for speed guarantees.
    if not x.is_contiguous():
        x = x.contiguous()
    if not y.is_contiguous():
        y = y.contiguous()
    # Dtype support: float16/float32/bfloat16 on CUDA
    if x.dtype not in (torch.float16, torch.float32, torch.bfloat16):
        raise TypeError(f"Unsupported dtype {x.dtype}. Use float16/float32/bfloat16.")
    _mod.grayscale_inline_cuda(x, y)
    return y