submission 513369

+ import triton
+ import triton.language as tl
  import torch
  from task import input_t, output_t
- from torch.utils.cpp_extension import load_inline
   
- cuda_source = """
- #include <cuda_runtime.h>
- #include <torch/extension.h>
   
- __global__ void grayscale_v4_kernel(
-     const float* __restrict__ input,
-     float* __restrict__ output,
-     int n_pixels
- ) {
-     int tid = blockIdx.x * blockDim.x + threadIdx.x;
-     int base_pixel = tid * 4;
+ @triton.autotune(
+     configs=[
+         triton.Config({'BLOCK_SIZE': 256}, num_warps=4, num_stages=3),
+         triton.Config({'BLOCK_SIZE': 512}, num_warps=4, num_stages=3),
+         triton.Config({'BLOCK_SIZE': 512}, num_warps=8, num_stages=3),
+         triton.Config({'BLOCK_SIZE': 1024}, num_warps=4, num_stages=3),
+         triton.Config({'BLOCK_SIZE': 1024}, num_warps=8, num_stages=3),
+         triton.Config({'BLOCK_SIZE': 2048}, num_warps=8, num_stages=3),
+         triton.Config({'BLOCK_SIZE': 2048}, num_warps=16, num_stages=3),
+         triton.Config({'BLOCK_SIZE': 4096}, num_warps=16, num_stages=3),
+         triton.Config({'BLOCK_SIZE': 4096}, num_warps=32, num_stages=3),
+         triton.Config({'BLOCK_SIZE': 256}, num_warps=4, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 512}, num_warps=4, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 512}, num_warps=8, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 1024}, num_warps=4, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 1024}, num_warps=8, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 2048}, num_warps=8, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 2048}, num_warps=16, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 4096}, num_warps=16, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 4096}, num_warps=32, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 1024}, num_warps=16, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 2048}, num_warps=32, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 512}, num_warps=16, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 256}, num_warps=8, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 128}, num_warps=4, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 128}, num_warps=4, num_stages=3),
+         triton.Config({'BLOCK_SIZE': 1024}, num_warps=4, num_stages=5),
+         triton.Config({'BLOCK_SIZE': 2048}, num_warps=8, num_stages=5),
+         triton.Config({'BLOCK_SIZE': 4096}, num_warps=16, num_stages=5),
+         triton.Config({'BLOCK_SIZE': 128}, num_warps=8, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 256}, num_warps=16, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 512}, num_warps=32, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 1024}, num_warps=32, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 1024}, num_warps=8, num_stages=5),
+         triton.Config({'BLOCK_SIZE': 2048}, num_warps=16, num_stages=5),
+         triton.Config({'BLOCK_SIZE': 4096}, num_warps=32, num_stages=5),
+         triton.Config({'BLOCK_SIZE': 512}, num_warps=8, num_stages=5),
+         triton.Config({'BLOCK_SIZE': 256}, num_warps=8, num_stages=5),
+     ],
+     key=['n_pixels'],
+ )
+ @triton.jit
+ def grayscale_kernel_coalesced(
+     rgb_ptr,
+     out_ptr,
+     n_pixels,
+     BLOCK_SIZE: tl.constexpr,
+ ):
+     pid = tl.program_id(0)
+     # Load 3*BLOCK_SIZE contiguous floats
+     base = pid * BLOCK_SIZE * 3
+     offs = tl.arange(0, BLOCK_SIZE * 3)
+     rgb_vals = tl.load(rgb_ptr + base + offs, mask=(base + offs) < n_pixels * 3, other=0.0)
   
-     const float wr = 0.2989f;
-     const float wg = 0.5870f;
-     const float wb = 0.1140f;
+     # Deinterleave: pixel i has R at 3i, G at 3i+1, B at 3i+2
+     i = tl.arange(0, BLOCK_SIZE)
+     r = tl.load(rgb_ptr + (pid * BLOCK_SIZE + i) * 3 + 0, mask=(pid * BLOCK_SIZE + i) < n_pixels, other=0.0)
+     g = tl.load(rgb_ptr + (pid * BLOCK_SIZE + i) * 3 + 1, mask=(pid * BLOCK_SIZE + i) < n_pixels, other=0.0)
+     b = tl.load(rgb_ptr + (pid * BLOCK_SIZE + i) * 3 + 2, mask=(pid * BLOCK_SIZE + i) < n_pixels, other=0.0)
   
-     if (base_pixel + 3 < n_pixels) {
-         const float4* in4 = reinterpret_cast<const float4*>(input + base_pixel * 3);
-         float4 d0 = __ldg(&in4[0]);
-         float4 d1 = __ldg(&in4[1]);
-         float4 d2 = __ldg(&in4[2]);
+     gray = 0.2989 * r + 0.5870 * g + 0.1140 * b
+     pixel_offs = pid * BLOCK_SIZE + i
+     tl.store(out_ptr + pixel_offs, gray, mask=pixel_offs < n_pixels)
   
-         float4 out;
-         out.x = __fmaf_rn(wr, d0.x, __fmaf_rn(wg, d0.y, wb * d0.z));
-         out.y = __fmaf_rn(wr, d0.w, __fmaf_rn(wg, d1.x, wb * d1.y));
-         out.z = __fmaf_rn(wr, d1.z, __fmaf_rn(wg, d1.w, wb * d2.x));
-         out.w = __fmaf_rn(wr, d2.y, __fmaf_rn(wg, d2.z, wb * d2.w));
   
-         reinterpret_cast<float4*>(output + base_pixel)[0] = out;
-     } else {
-         for (int i = 0; i < 4 && base_pixel + i < n_pixels; i++) {
-             int idx = (base_pixel + i) * 3;
-             output[base_pixel + i] = wr * input[idx] + wg * input[idx+1] + wb * input[idx+2];
-         }
-     }
- }
+ @triton.autotune(
+     configs=[
+         triton.Config({'BLOCK_SIZE': 256}, num_warps=4, num_stages=3),
+         triton.Config({'BLOCK_SIZE': 512}, num_warps=4, num_stages=3),
+         triton.Config({'BLOCK_SIZE': 512}, num_warps=8, num_stages=3),
+         triton.Config({'BLOCK_SIZE': 1024}, num_warps=4, num_stages=3),
+         triton.Config({'BLOCK_SIZE': 1024}, num_warps=8, num_stages=3),
+         triton.Config({'BLOCK_SIZE': 2048}, num_warps=8, num_stages=3),
+         triton.Config({'BLOCK_SIZE': 2048}, num_warps=16, num_stages=3),
+         triton.Config({'BLOCK_SIZE': 4096}, num_warps=16, num_stages=3),
+         triton.Config({'BLOCK_SIZE': 4096}, num_warps=32, num_stages=3),
+         triton.Config({'BLOCK_SIZE': 256}, num_warps=4, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 512}, num_warps=4, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 512}, num_warps=8, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 1024}, num_warps=4, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 1024}, num_warps=8, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 2048}, num_warps=8, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 2048}, num_warps=16, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 4096}, num_warps=16, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 4096}, num_warps=32, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 1024}, num_warps=16, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 2048}, num_warps=32, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 512}, num_warps=16, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 256}, num_warps=8, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 128}, num_warps=4, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 128}, num_warps=4, num_stages=3),
+         triton.Config({'BLOCK_SIZE': 1024}, num_warps=4, num_stages=5),
+         triton.Config({'BLOCK_SIZE': 2048}, num_warps=8, num_stages=5),
+         triton.Config({'BLOCK_SIZE': 4096}, num_warps=16, num_stages=5),
+         triton.Config({'BLOCK_SIZE': 128}, num_warps=8, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 256}, num_warps=16, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 512}, num_warps=32, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 1024}, num_warps=32, num_stages=4),
+         triton.Config({'BLOCK_SIZE': 1024}, num_warps=8, num_stages=5),
+         triton.Config({'BLOCK_SIZE': 2048}, num_warps=16, num_stages=5),
+         triton.Config({'BLOCK_SIZE': 4096}, num_warps=32, num_stages=5),
+         triton.Config({'BLOCK_SIZE': 512}, num_warps=8, num_stages=5),
+         triton.Config({'BLOCK_SIZE': 256}, num_warps=8, num_stages=5),
+     ],
+     key=['n_pixels'],
+ )
+ @triton.jit
+ def grayscale_kernel(
+     rgb_ptr,
+     out_ptr,
+     n_pixels,
+     BLOCK_SIZE: tl.constexpr,
+ ):
+     pid = tl.program_id(0)
+     pixel_offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+     mask = pixel_offs < n_pixels
   
- torch::Tensor grayscale_forward(torch::Tensor input, torch::Tensor output) {
-     int n_pixels = input.numel() / 3;
-     int threads = 256;
-     int blocks = (n_pixels / 4 + threads - 1) / threads + 1;
-     grayscale_v4_kernel<<<blocks, threads>>>(
-         input.data_ptr<float>(),
-         output.data_ptr<float>(),
-         n_pixels
-     );
-     return output;
- }
- """
+     r = tl.load(rgb_ptr + pixel_offs * 3 + 0, mask=mask, other=0.0)
+     g = tl.load(rgb_ptr + pixel_offs * 3 + 1, mask=mask, other=0.0)
+     b = tl.load(rgb_ptr + pixel_offs * 3 + 2, mask=mask, other=0.0)
   
- cpp_source = """
- torch::Tensor grayscale_forward(torch::Tensor input, torch::Tensor output);
- """
+     gray = 0.2989 * r + 0.5870 * g + 0.1140 * b
+     tl.store(out_ptr + pixel_offs, gray, mask=mask)
   
- _module = None
   
- def _get_module():
-     global _module
-     if _module is None:
-         _module = load_inline(
-             name="grayscale_v4_fma_v3",
-             cpp_sources=cpp_source,
-             cuda_sources=cuda_source,
-             functions=["grayscale_forward"],
-             extra_cuda_cflags=["-O3", "--use_fast_math"],
-             verbose=False,
-         )
-     return _module
+ def custom_kernel(data: input_t) -> output_t:
+     rgb, output = data
+     H, W, _ = rgb.shape
+     n_pixels = H * W
   
+     rgb_flat = rgb.view(-1)
+     out_flat = output.view(-1)
   
- def custom_kernel(data: input_t) -> output_t:
-     x, output = data
-     mod = _get_module()
-     mod.grayscale_forward(x.contiguous().view(-1), output.view(-1))
-     return output
+     def grid(meta):
+         return (triton.cdiv(n_pixels, meta['BLOCK_SIZE']),)
+  
+     grayscale_kernel[grid](rgb_flat, out_flat, n_pixels)
+     return output
   No newline at end of file