submission 779864

- import triton
- import triton.language as tl
  from task import input_t, output_t
+ from torch.utils.cpp_extension import load_inline
   
+ cuda_source = """
+ #include <cuda_runtime.h>
   
- @triton.jit
- def _grayscale_kernel(
-     data_ptr,
-     output_ptr,
-     n_pixels,
-     w0,
-     w1,
-     w2,
-     BLOCK_SIZE: tl.constexpr,
- ):
-     pid = tl.program_id(0)
-     offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
-     mask = offsets < n_pixels
+ __global__ void grayscale_kernel(
+     const float* __restrict__ data,
+     float* __restrict__ output
+ ) {
+     // each thread handles 4 pixels = 12 floats = 3 x float4 loads
+     int base = (blockIdx.x * blockDim.x + threadIdx.x) * 4;
   
-     base = offsets * 3
-     r = tl.load(data_ptr + base, mask=mask)
-     g = tl.load(data_ptr + base + 1, mask=mask)
-     b = tl.load(data_ptr + base + 2, mask=mask)
+     float4 a = __ldg((const float4*)&data[base * 3 + 0]); // R0 G0 B0 R1
+     float4 b = __ldg((const float4*)&data[base * 3 + 4]); // G1 B1 R2 G2
+     float4 c = __ldg((const float4*)&data[base * 3 + 8]); // B2 R3 G3 B3
   
-     tl.store(output_ptr + offsets, r * w0 + g * w1 + b * w2, mask=mask)
+     output[base + 0] = a.x * 0.2989f + a.y * 0.5870f + a.z * 0.1140f;
+     output[base + 1] = a.w * 0.2989f + b.x * 0.5870f + b.y * 0.1140f;
+     output[base + 2] = b.z * 0.2989f + b.w * 0.5870f + c.x * 0.1140f;
+     output[base + 3] = c.y * 0.2989f + c.z * 0.5870f + c.w * 0.1140f;
+ }
   
+ torch::Tensor grayscale_cuda(torch::Tensor data, torch::Tensor output) {
+     int n_pixels = output.numel();
+     const int BLOCK_SIZE = 256;
+     int grid = n_pixels / (BLOCK_SIZE * 4);
+     grayscale_kernel<<<grid, BLOCK_SIZE>>>(
+         data.data_ptr<float>(),
+         output.data_ptr<float>()
+     );
+     return output;
+ }
+ """
   
+ cpp_source = "torch::Tensor grayscale_cuda(torch::Tensor data, torch::Tensor output);"
+  
+ _ext = load_inline(
+     name="grayscale_float4",
+     cpp_sources=cpp_source,
+     cuda_sources=cuda_source,
+     functions=["grayscale_cuda"],
+     verbose=False,
+ )
+  
+  
  def custom_kernel(data: input_t) -> output_t:
      data, output = data
-     n_pixels = output.numel()
-     grid = (triton.cdiv(n_pixels, 1024),)
-     _grayscale_kernel[grid](
-         data.contiguous().view(-1),
-         output.view(-1),
-         n_pixels,
-         0.2989,
-         0.5870,
-         0.1140,
-         BLOCK_SIZE=1024,
-     )
+     _ext.grayscale_cuda(data, output)
      return output