submission 66676

# 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>
#include <cstdint>

__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);
   }
};

__device__ inline float4 ld_float4_cs(const float* ptr) {
    float4 v;
    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"(ptr));
    return v;
}

__device__ inline void st_float4_cs(float* ptr, const float4& v) {
    asm volatile("st.global.cs.v4.f32 [%0], {%1, %2, %3, %4};"
                 :
                 : "l"(ptr), "f"(v.x), "f"(v.y), "f"(v.z), "f"(v.w));
}

// 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;

#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
    extern __shared__ float4 shared_tiles[];
    float4* tile = shared_tiles + threadIdx.x * 3;

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

        const unsigned smem_addr0 = __cvta_generic_to_shared(tile + 0);
        const unsigned smem_addr1 = __cvta_generic_to_shared(tile + 1);
        const unsigned smem_addr2 = __cvta_generic_to_shared(tile + 2);
        const unsigned long long gmem_addr0 = reinterpret_cast<const unsigned long long>(vec + 0);
        const unsigned long long gmem_addr1 = reinterpret_cast<const unsigned long long>(vec + 1);
        const unsigned long long gmem_addr2 = reinterpret_cast<const unsigned long long>(vec + 2);

        asm volatile("cp.async.ca.shared.global [%0], [%1], 16;\n" :: "r"(smem_addr0), "l"(gmem_addr0));
        asm volatile("cp.async.ca.shared.global [%0], [%1], 16;\n" :: "r"(smem_addr1), "l"(gmem_addr1));
        asm volatile("cp.async.ca.shared.global [%0], [%1], 16;\n" :: "r"(smem_addr2), "l"(gmem_addr2));
        asm volatile("cp.async.commit_group;\n");
        asm volatile("cp.async.wait_group 0;\n");

        const float4 v0 = tile[0];
        const float4 v1 = tile[1];
        const float4 v2 = tile[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));

        const float4 packed = make_float4(gray0, gray1, gray2, gray3);
        st_float4_cs(out + idx, packed);
    }
#else
    for (; idx + 3 < n_pixels; idx += stride) {
        const size_t base = idx * 3;
        const float4 v0 = ld_float4_cs(in + base + 0);
        const float4 v1 = ld_float4_cs(in + base + 4);
        const float4 v2 = ld_float4_cs(in + base + 8);

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

        const float4 packed = make_float4(gray0, gray1, gray2, gray3);
        st_float4_cs(out + idx, packed);
    }
#endif

    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) {
    const size_t shared_bytes = static_cast<size_t>(threads) * 3 * sizeof(float4);
    grayscale_kernel_float<<<blocks, threads, shared_bytes, 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