submission 66641

+ # fast_grayscale.py
  import torch
- from task import input_t, output_t
+ from torch.utils.cpp_extension import load_inline
   
- # Standard luminance coefficients expressed as Python floats.
- _WEIGHT_R = 0.2989
- _WEIGHT_G = 0.5870
- _WEIGHT_B = 0.1140
+ _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;
   
- def custom_kernel(data: input_t) -> output_t:
-     data, output = data
+ 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; }
+ };
   
-     # Avoid allocating temporary tensors by writing directly into the provided output buffer.
-     # Using in-place arithmetic keeps the computation bandwidth-bound and reuses the output memory.
-     r = data[..., 0]
-     g = data[..., 1]
-     b = data[..., 2]
+ 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; }
+ };
   
-     torch.mul(r, _WEIGHT_R, out=output)
-     output.add_(g, alpha=_WEIGHT_G)
-     output.add_(b, alpha=_WEIGHT_B)
+ 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);
+     }
+ };
   
-     return output
+ 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