submission 67773

- # submission.py
- import os
- from task import input_t, output_t
  import torch
  from torch.utils.cpp_extension import load_inline
+ from task import input_t, output_t
  
- _cuda = r"""
- #include <torch/extension.h>
- #include <ATen/cuda/CUDAContext.h>
- #include <cuda.h>
+ grayscale_cuda_source = """
+ #include <cuda_fp16.h>
  
- #ifndef UNROLL
- #define UNROLL 8
- #endif
+ __global__ void rgb_to_grayscale_kernel(
+     const float* __restrict__ rgb,
+     float* __restrict__ gray,
+     const int total_pixels)
+ {
+     const float w_r = 0.2989f;
+     const float w_g = 0.5870f;
+     const float w_b = 0.1140f;
  
- __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;
+     // Each thread processes 1 iteration of 4 pixels = 4 pixels total
+     constexpr int iterations = 1;
+     constexpr int pixels_per_iter = 4;
  
-     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;
+     int base_idx = (blockIdx.x * blockDim.x + threadIdx.x) * pixels_per_iter * iterations;
  
-             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;
+     #pragma unroll
+     for (int iter = 0; iter < iterations; iter++) {
+         int idx = base_idx + iter * pixels_per_iter;
+ 
+         if (idx + 3 < total_pixels) {
+             int rgb_base = idx * 3;
+ 
+             float4 vec0 = *reinterpret_cast<const float4*>(&rgb[rgb_base]);
+             float4 vec1 = *reinterpret_cast<const float4*>(&rgb[rgb_base + 4]);
+             float4 vec2 = *reinterpret_cast<const float4*>(&rgb[rgb_base + 8]);
+ 
+             float r0 = vec0.x, g0 = vec0.y, b0 = vec0.z;
+             float r1 = vec0.w, g1 = vec1.x, b1 = vec1.y;
+             float r2 = vec1.z, g2 = vec1.w, b2 = vec2.x;
+             float r3 = vec2.y, g3 = vec2.z, b3 = vec2.w;
+ 
+             float y0 = fmaf(w_r, r0, fmaf(w_g, g0, w_b * b0));
+             float y1 = fmaf(w_r, r1, fmaf(w_g, g1, w_b * b1));
+             float y2 = fmaf(w_r, r2, fmaf(w_g, g2, w_b * b2));
+             float y3 = fmaf(w_r, r3, fmaf(w_g, g3, w_b * b3));
+ 
+             float4 result = {y0, y1, y2, y3};
+             *reinterpret_cast<float4*>(&gray[idx]) = result;
          }
      }
  }
  
- torch::Tensor grayscale_forward(torch::Tensor input,
-                                 torch::Tensor output,
-                                 double w0,
-                                 double w1,
-                                 double w2) {
-     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");
+ torch::Tensor rgb_to_grayscale_cuda(torch::Tensor rgb, torch::Tensor gray) {
+     TORCH_CHECK(rgb.device().is_cuda(), "RGB tensor must be a CUDA tensor");
+     TORCH_CHECK(gray.device().is_cuda(), "Gray tensor must be a CUDA tensor");
+     TORCH_CHECK(rgb.dim() == 3, "RGB tensor must be 3D (H, W, 3)");
+     TORCH_CHECK(rgb.size(2) == 3, "RGB tensor must have 3 channels");
  
-     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");
+     const int H = rgb.size(0);
+     const int W = rgb.size(1);
+     const int total_pixels = H * W;
  
-     auto* props = at::cuda::getCurrentDeviceProperties();
-     const int sm = props->multiProcessorCount;
+     const int threads = 1024;
+     const int pixels_per_thread = 4;  // 1 iteration × 4 pixels per iteration
+     const int blocks = (total_pixels + threads * pixels_per_thread - 1) / (threads * pixels_per_thread);
  
-     const int block = 512;
-     int grid = (int)((n_pix + block - 1) / block);
-     grid = max(grid, sm * 8);
-     grid = min(grid, 65535);
+     rgb = rgb.contiguous();
+     gray = gray.contiguous();
  
-     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;
+     rgb_to_grayscale_kernel<<<blocks, threads>>>(
+         rgb.data_ptr<float>(),
+         gray.data_ptr<float>(),
+         total_pixels
+     );
+ 
+     cudaError_t err = cudaGetLastError();
+     if (err != cudaSuccess) {
+         throw std::runtime_error(cudaGetErrorString(err));
+     }
+ 
+     return gray;
  }
  """
  
- _cpp_decl = r"""
+ grayscale_cpp_source = """
  #include <torch/extension.h>
- torch::Tensor grayscale_forward(torch::Tensor, torch::Tensor, double, double, double);
+ 
+ torch::Tensor rgb_to_grayscale_cuda(torch::Tensor rgb, torch::Tensor gray);
  """
  
- mod = load_inline(
-     name="rgb2gray_inline_fp32_p2",
-     cpp_sources=_cpp_decl,
-     cuda_sources=_cuda,
-     functions=["grayscale_forward"],
-     extra_cuda_cflags=[
-         "-O3",
-         "-lineinfo",
-         "--use_fast_math",
-         "-Xptxas","-O3",
-         "-Xptxas","-dlcm=ca"
-     ],
+ import sys
+ import os
+ if sys.stdout is None:
+     sys.stdout = open(os.devnull, 'w')
+ if sys.stderr is None:
+     sys.stderr = open(os.devnull, 'w')
+ 
+ grayscale_module = load_inline(
+     name='rgb_to_grayscale_cuda',
+     cpp_sources=grayscale_cpp_source,
+     cuda_sources=grayscale_cuda_source,
+     functions=['rgb_to_grayscale_cuda'],
+     extra_cuda_cflags=['-O3', '--use_fast_math', '-lineinfo', '--generate-code=arch=compute_100,code=sm_100a'],
      verbose=False,
  )
  
- _WEIGHTS = (0.2989, 0.5870, 0.1140)
- 
- 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
-     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
+     rgb, output = data
+     assert rgb.is_cuda and output.is_cuda, "Tensors must be on GPU"
+     assert rgb.dim() == 3 and rgb.size(2) == 3, "RGB must be (H, W, 3)"
+     assert rgb.dtype == torch.float32, "RGB must be float32"
+     return grayscale_module.rgb_to_grayscale_cuda(rgb, output)