submission 66663

# 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_runtime.h>
#include <cuda_bf16.h>
#include <cuda_fp16.h>
#include <type_traits>

__constant__ float kLuma[3] = {0.2989f, 0.5870f, 0.1140f};

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_scalar(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 = static_cast<acc_t>(kLuma[0]);
    const acc_t wg = static_cast<acc_t>(kLuma[1]);
    const acc_t wb = static_cast<acc_t>(kLuma[2]);

    // Process 4 pixels per thread for better memory throughput
    const size_t stride = static_cast<size_t>(gridDim.x) * blockDim.x * 4;
    size_t idx = (blockIdx.x * blockDim.x + threadIdx.x) * 4;

    // main unrolled loop
    for (; idx + 3 < n_pixels; idx += stride) {
        #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);
    }
}

__global__ void grayscale_kernel_float(const float* __restrict__ in,
                                       float* __restrict__ out,
                                       size_t n_pixels)
{
    const float wr = kLuma[0];
    const float wg = kLuma[1];
    const float wb = kLuma[2];

    const size_t stride = static_cast<size_t>(gridDim.x) * blockDim.x * 4;
    size_t idx = (blockIdx.x * blockDim.x + threadIdx.x) * 4;

    for (; idx + 3 < n_pixels; idx += stride) {
        const size_t base = idx * 3;
        const float4* vec = reinterpret_cast<const float4*>(in + base);
        const float4 v0 = vec[0];
        const float4 v1 = vec[1];
        const float4 v2 = vec[2];

        const float gray0 = fmaf(v0.y, wg, fmaf(v0.x, wr, v0.z * wb));
        const float gray1 = fmaf(v1.x, wg, fmaf(v0.w, wr, v1.y * wb));
        const float gray2 = fmaf(v1.w, wg, fmaf(v1.z, wr, v2.x * wb));
        const float gray3 = fmaf(v2.z, wg, fmaf(v2.y, wr, v2.w * wb));

        float4* out_vec = reinterpret_cast<float4*>(out + idx);
        out_vec[0] = make_float4(gray0, gray1, gray2, gray3);
    }

    for (; idx < n_pixels; ++idx) {
        const size_t base = idx * 3;
        const float r = in[base + 0];
        const float g = in[base + 1];
        const float b = in[base + 2];
        out[idx] = r * wr + g * wg + b * wb;
    }
}

__global__ void grayscale_kernel_half(const at::Half* __restrict__ in,
                                      at::Half* __restrict__ out,
                                      size_t n_pixels)
{
    const float wr = kLuma[0];
    const float wg = kLuma[1];
    const float wb = kLuma[2];

    const size_t stride = static_cast<size_t>(gridDim.x) * blockDim.x * 4;
    size_t idx = (blockIdx.x * blockDim.x + threadIdx.x) * 4;

    const __half* in_half = reinterpret_cast<const __half*>(in);
    __half* out_half = reinterpret_cast<__half*>(out);

    for (; idx + 3 < n_pixels; idx += stride) {
        const size_t base = idx * 3;
        const __half2* vec = reinterpret_cast<const __half2*>(in_half + base);

        const float2 rg0 = __half22float2(vec[0]);
        const float2 br1 = __half22float2(vec[1]);
        const float2 gb1 = __half22float2(vec[2]);
        const float2 rg2 = __half22float2(vec[3]);
        const float2 br3 = __half22float2(vec[4]);
        const float2 gb3 = __half22float2(vec[5]);

        const float gray0 = rg0.x * wr + rg0.y * wg + br1.x * wb;
        const float gray1 = br1.y * wr + gb1.x * wg + gb1.y * wb;
        const float gray2 = rg2.x * wr + rg2.y * wg + br3.x * wb;
        const float gray3 = br3.y * wr + gb3.x * wg + gb3.y * wb;

        __half2* out_vec = reinterpret_cast<__half2*>(out_half + idx);
        out_vec[0] = __floats2half2_rn(gray0, gray1);
        out_vec[1] = __floats2half2_rn(gray2, gray3);
    }

    for (; idx < n_pixels; ++idx) {
        const size_t base = idx * 3;
        const float r = __half2float(in_half[base + 0]);
        const float g = __half2float(in_half[base + 1]);
        const float b = __half2float(in_half[base + 2]);
        out_half[idx] = __float2half(r * wr + g * wg + b * wb);
    }
}

__global__ void grayscale_kernel_bfloat(const at::BFloat16* __restrict__ in,
                                        at::BFloat16* __restrict__ out,
                                        size_t n_pixels)
{
    const float wr = kLuma[0];
    const float wg = kLuma[1];
    const float wb = kLuma[2];

    const size_t stride = static_cast<size_t>(gridDim.x) * blockDim.x * 4;
    size_t idx = (blockIdx.x * blockDim.x + threadIdx.x) * 4;

    const __nv_bfloat16* in_bf16 = reinterpret_cast<const __nv_bfloat16*>(in);
    __nv_bfloat16* out_bf16 = reinterpret_cast<__nv_bfloat16*>(out);

    for (; idx + 3 < n_pixels; idx += stride) {
        const size_t base = idx * 3;
        const __nv_bfloat162* vec = reinterpret_cast<const __nv_bfloat162*>(in_bf16 + base);

        const float2 rg0 = __bfloat1622float2(vec[0]);
        const float2 br1 = __bfloat1622float2(vec[1]);
        const float2 gb1 = __bfloat1622float2(vec[2]);
        const float2 rg2 = __bfloat1622float2(vec[3]);
        const float2 br3 = __bfloat1622float2(vec[4]);
        const float2 gb3 = __bfloat1622float2(vec[5]);

        const float gray0 = rg0.x * wr + rg0.y * wg + br1.x * wb;
        const float gray1 = br1.y * wr + gb1.x * wg + gb1.y * wb;
        const float gray2 = rg2.x * wr + rg2.y * wg + br3.x * wb;
        const float gray3 = br3.y * wr + gb3.x * wg + gb3.y * wb;

        __nv_bfloat162* out_vec = reinterpret_cast<__nv_bfloat162*>(out_bf16 + idx);
        out_vec[0] = __floats2bfloat162_rn(gray0, gray1);
        out_vec[1] = __floats2bfloat162_rn(gray2, gray3);
    }

    for (; idx < n_pixels; ++idx) {
        const size_t base = idx * 3;
        const float r = __bfloat162float(in_bf16[base + 0]);
        const float g = __bfloat162float(in_bf16[base + 1]);
        const float b = __bfloat162float(in_bf16[base + 2]);
        out_bf16[idx] = __float2bfloat16(r * wr + g * wg + b * wb);
    }
}

template <typename scalar_t>
inline void launch_grayscale_kernel(const scalar_t* in,
                                    scalar_t* out,
                                    size_t n_pixels,
                                    int blocks,
                                    int threads,
                                    cudaStream_t stream) {
    grayscale_kernel_scalar<scalar_t><<<blocks, threads, 0, stream>>>(in, out, n_pixels);
}

template <>
inline void launch_grayscale_kernel<float>(const float* in,
                                           float* out,
                                           size_t n_pixels,
                                           int blocks,
                                           int threads,
                                           cudaStream_t stream) {
    grayscale_kernel_float<<<blocks, threads, 0, stream>>>(in, out, n_pixels);
}

template <>
inline void launch_grayscale_kernel<at::Half>(const at::Half* in,
                                              at::Half* out,
                                              size_t n_pixels,
                                              int blocks,
                                              int threads,
                                              cudaStream_t stream) {
    grayscale_kernel_half<<<blocks, threads, 0, stream>>>(in, out, n_pixels);
}

template <>
inline void launch_grayscale_kernel<at::BFloat16>(const at::BFloat16* in,
                                                  at::BFloat16* out,
                                                  size_t n_pixels,
                                                  int blocks,
                                                  int threads,
                                                  cudaStream_t stream) {
    grayscale_kernel_bfloat<<<blocks, threads, 0, stream>>>(in, out, n_pixels);
}

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");

    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_>();
        constexpr int pixels_per_thread = 4;
        const int threads = std::is_same_v<scalar_t_, float> ? 1024 : 256;
        const auto n_pixels = static_cast<size_t>(output.numel());
        const int blocks = std::min<int>((int)((n_pixels + (threads*pixels_per_thread - 1)) / (threads*pixels_per_thread)), 32768);
        launch_grayscale_kernel<scalar_t_>(in_ptr, out_ptr, n_pixels, blocks, threads, stream);
    });

    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