submission 66514

  # submission.py
+ import os
  from task import input_t, output_t
  import torch
  from torch.utils.cpp_extension import load_inline
  
  _cuda = r"""
  #include <torch/extension.h>
- #include <cuda.h>
- #include <cuda_runtime.h>
  #include <ATen/cuda/CUDAContext.h>
+ #include <cuda.h>
  
- template <typename scalar_t>
- __global__ void rgb2gray_kernel(const scalar_t* __restrict__ in,
-                                 scalar_t* __restrict__ out,
-                                 size_t n_pix,
-                                 const float w0,
-                                 const float w1,
-                                 const float w2) {
-     size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
-     size_t stride = blockDim.x * gridDim.x;
+ #ifndef UNROLL
+ #define UNROLL 8
+ #endif
  
-     for (size_t i = idx; i < n_pix; i += stride) {
-         float r = static_cast<float>(in[3*i + 0]);
-         float g = static_cast<float>(in[3*i + 1]);
-         float b = static_cast<float>(in[3*i + 2]);
-         float y = w0 * r + w1 * g + w2 * b;
-         out[i] = static_cast<scalar_t>(y);
+ __launch_bounds__(512, 2)
+ __global__ void rgb2gray_f32_kernel(const float* __restrict__ in,
+                                     float* __restrict__ out,
+                                     size_t n_pix,
+                                     const float w0,
+                                     const float w1,
+                                     const float w2) {
+     const size_t tid = blockIdx.x * blockDim.x + threadIdx.x;
+     const size_t stride = blockDim.x * gridDim.x;
+ 
+     for (size_t base = tid; base < n_pix; base += stride * UNROLL) {
+ #pragma unroll
+         for (int u = 0; u < UNROLL; ++u) {
+             const size_t i = base + (size_t)u * stride;
+             if (i >= n_pix) break;
+ 
+             const size_t o = 3 * i;  // RGB interleaved
+ #if __CUDA_ARCH__ >= 800
+             float r = __ldg(in + o + 0);
+             float g = __ldg(in + o + 1);
+             float b = __ldg(in + o + 2);
+ #else
+             float r = in[o + 0];
+             float g = in[o + 1];
+             float b = in[o + 2];
+ #endif
+             float y = fmaf(w2, b, fmaf(w1, g, w0 * r));
+             out[i] = y;
+         }
      }
  }
  
  ⋯ 2 unchanged lines
                                  double w0,
                                  double w1,
                                  double w2) {
-     TORCH_CHECK(input.is_cuda() && output.is_cuda(), "tensors must be CUDA");
-     TORCH_CHECK(input.is_contiguous() && output.is_contiguous(), "tensors must be contiguous");
-     TORCH_CHECK(input.scalar_type() == output.scalar_type(), "dtype mismatch");
+     TORCH_CHECK(input.is_cuda() && output.is_cuda(), "CUDA only");
+     TORCH_CHECK(input.scalar_type() == at::kFloat && output.scalar_type() == at::kFloat, "fp32 only");
+     TORCH_CHECK(input.is_contiguous() && output.is_contiguous(), "contiguous required");
      TORCH_CHECK(input.size(-1) == 3, "last dim must be 3");
+ 
      const size_t n_pix = static_cast<size_t>(input.numel() / 3);
-     TORCH_CHECK(static_cast<size_t>(output.numel()) == n_pix, "output.numel mismatch");
+     TORCH_CHECK(static_cast<size_t>(output.numel()) == n_pix, "output numel mismatch");
  
-     const int block = 256;
-     const int max_blocks = 32768;
-     const int grid = std::min<int>((n_pix + block - 1) / block, max_blocks);
+     auto* props = at::cuda::getCurrentDeviceProperties();
+     const int sm = props->multiProcessorCount;
  
-     AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.scalar_type(), "rgb2gray_kernel", [&] {
-         const scalar_t* in_ptr = input.data_ptr<scalar_t>();
-         scalar_t* out_ptr = output.data_ptr<scalar_t>();
-         rgb2gray_kernel<scalar_t><<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
-             in_ptr, out_ptr, n_pix, static_cast<float>(w0), static_cast<float>(w1), static_cast<float>(w2));
-     });
+     const int block = 512;
+     int grid = (int)((n_pix + block - 1) / block);
+     grid = max(grid, sm * 8);
+     grid = min(grid, 65535);
+ 
+     auto stream = at::cuda::getCurrentCUDAStream();
+     rgb2gray_f32_kernel<<<grid, block, 0, stream>>>(
+         input.data_ptr<float>(),
+         output.data_ptr<float>(),
+         n_pix,
+         static_cast<float>(w0),
+         static_cast<float>(w1),
+         static_cast<float>(w2));
      C10_CUDA_KERNEL_LAUNCH_CHECK();
      return output;
  }
  """
  
- # Minimal declaration so the autogenerated main.cpp can call the CUDA function.
  _cpp_decl = r"""
  #include <torch/extension.h>
  torch::Tensor grayscale_forward(torch::Tensor, torch::Tensor, double, double, double);
  """
  
  mod = load_inline(
-     name="rgb2gray_inline",
-     cpp_sources=_cpp_decl,          # declaration TU
-     cuda_sources=_cuda,             # definition TU
+     name="rgb2gray_inline_fp32_p2",
+     cpp_sources=_cpp_decl,
+     cuda_sources=_cuda,
      functions=["grayscale_forward"],
-     extra_cuda_cflags=["-O3", "-lineinfo", "--use_fast_math"],
+     extra_cuda_cflags=[
+         "-O3",
+         "-lineinfo",
+         "--use_fast_math",
+         "-Xptxas","-O3",
+         "-Xptxas","-dlcm=ca"
+     ],
      verbose=False,
  )
  
  _WEIGHTS = (0.2989, 0.5870, 0.1140)
  
- def _ensure_contig(x: torch.Tensor) -> torch.Tensor:
+ def _contig(x: torch.Tensor) -> torch.Tensor:
      return x if x.is_contiguous() else x.contiguous()
  
  @torch.inference_mode()
  def custom_kernel(data: input_t) -> output_t:
      x, out = data
-     assert x.is_cuda and out.is_cuda, "move tensors to CUDA"
-     assert x.size(-1) == 3, "last dim must be 3"
-     x = _ensure_contig(x)
-     out = _ensure_contig(out)
-     # Flatten views
-     xv = x.view(-1, 3)
-     ov = out.view(-1)
+     assert x.is_cuda and out.is_cuda
+     assert x.dtype == torch.float32 and out.dtype == torch.float32
+     assert x.size(-1) == 3
+     x = _contig(x); out = _contig(out)
+     xv = x.view(-1, 3); ov = out.view(-1)
      mod.grayscale_forward(xv, ov, *_WEIGHTS)
      return out