submission 617563

  ⋯ 5 unchanged lines
  #include <torch/extension.h>
  #include <cuda_runtime.h>
   
+ // Branchless version: for power-of-2 sizes that are multiples of 4,
+ // n_quads is exactly divisible by 1024 so we skip the branch
  __global__ void __launch_bounds__(1024, 1)
- grayscale_kernel(const float* __restrict__ input,
-                  float* __restrict__ output,
-                  const int n_quads) {
+ grayscale_kernel_exact(const float* __restrict__ input,
+                        float* __restrict__ output) {
      const int tid = blockIdx.x * 1024 + threadIdx.x;
+     const int base = tid * 12;
   
+     const float2 a = __ldg((const float2*)(input + base));
+     const float2 b = __ldg((const float2*)(input + base + 2));
+     const float2 c = __ldg((const float2*)(input + base + 4));
+     const float2 d = __ldg((const float2*)(input + base + 6));
+     const float2 e = __ldg((const float2*)(input + base + 8));
+     const float2 f = __ldg((const float2*)(input + base + 10));
+  
+     float4 out;
+     out.x = __fmaf_rn(0.2989f, a.x, __fmaf_rn(0.5870f, a.y, 0.1140f * b.x));
+     out.y = __fmaf_rn(0.2989f, b.y, __fmaf_rn(0.5870f, c.x, 0.1140f * c.y));
+     out.z = __fmaf_rn(0.2989f, d.x, __fmaf_rn(0.5870f, d.y, 0.1140f * e.x));
+     out.w = __fmaf_rn(0.2989f, e.y, __fmaf_rn(0.5870f, f.x, 0.1140f * f.y));
+  
+     float* dst = (float*)output + tid * 4;
+     asm volatile(
+         "st.global.cs.v4.f32 [%0], {%1,%2,%3,%4};\n\t"
+         :: "l"(dst), "f"(out.x), "f"(out.y), "f"(out.z), "f"(out.w)
+         : "memory"
+     );
+ }
+  
+ // Guarded version for non-exact sizes
+ __global__ void __launch_bounds__(1024, 1)
+ grayscale_kernel_guarded(const float* __restrict__ input,
+                          float* __restrict__ output,
+                          const int n_quads) {
+     const int tid = blockIdx.x * 1024 + threadIdx.x;
      if (tid < n_quads) {
          const int base = tid * 12;
-         const float4 v0 = __ldg((const float4*)(input + base));
-         const float4 v1 = __ldg((const float4*)(input + base + 4));
-         const float4 v2 = __ldg((const float4*)(input + base + 8));
   
+         const float2 a = __ldg((const float2*)(input + base));
+         const float2 b = __ldg((const float2*)(input + base + 2));
+         const float2 c = __ldg((const float2*)(input + base + 4));
+         const float2 d = __ldg((const float2*)(input + base + 6));
+         const float2 e = __ldg((const float2*)(input + base + 8));
+         const float2 f = __ldg((const float2*)(input + base + 10));
+  
          float4 out;
-         out.x = __fmaf_rn(0.2989f, v0.x, __fmaf_rn(0.5870f, v0.y, 0.1140f * v0.z));
-         out.y = __fmaf_rn(0.2989f, v0.w, __fmaf_rn(0.5870f, v1.x, 0.1140f * v1.y));
-         out.z = __fmaf_rn(0.2989f, v1.z, __fmaf_rn(0.5870f, v1.w, 0.1140f * v2.x));
-         out.w = __fmaf_rn(0.2989f, v2.y, __fmaf_rn(0.5870f, v2.z, 0.1140f * v2.w));
+         out.x = __fmaf_rn(0.2989f, a.x, __fmaf_rn(0.5870f, a.y, 0.1140f * b.x));
+         out.y = __fmaf_rn(0.2989f, b.y, __fmaf_rn(0.5870f, c.x, 0.1140f * c.y));
+         out.z = __fmaf_rn(0.2989f, d.x, __fmaf_rn(0.5870f, d.y, 0.1140f * e.x));
+         out.w = __fmaf_rn(0.2989f, e.y, __fmaf_rn(0.5870f, f.x, 0.1140f * f.y));
   
          float* dst = (float*)output + tid * 4;
          asm volatile(
  ⋯ 6 unchanged lines
   
  torch::Tensor launch_grayscale(torch::Tensor input, torch::Tensor output) {
      const int n_quads = (input.size(0) * input.size(1)) >> 2;
-     grayscale_kernel<<<(n_quads + 1023) / 1024, 1024>>>(
-         input.data_ptr<float>(),
-         output.data_ptr<float>(),
-         n_quads
-     );
+     const int n_blocks = (n_quads + 1023) / 1024;
+  
+     // If exact multiple of 1024, use branchless kernel
+     if ((n_quads & 1023) == 0) {
+         grayscale_kernel_exact<<<n_blocks, 1024>>>(
+             input.data_ptr<float>(),
+             output.data_ptr<float>()
+         );
+     } else {
+         grayscale_kernel_guarded<<<n_blocks, 1024>>>(
+             input.data_ptr<float>(),
+             output.data_ptr<float>(),
+             n_quads
+         );
+     }
      return output;
  }
  '''
  ⋯ 1 unchanged lines
  cpp_src = 'torch::Tensor launch_grayscale(torch::Tensor input, torch::Tensor output);'
   
  module = load_inline(
-     name='grayscale_v7_cs',
+     name='grayscale_v10',
      cpp_sources=cpp_src,
      cuda_sources=cuda_src,
      functions=['launch_grayscale'],
      verbose=False,
-     extra_cuda_cflags=['-O3', '--use_fast_math', '-maxrregcount=28'],
+     extra_cuda_cflags=['-O3', '--use_fast_math'],
  )
   
  _launch = module.launch_grayscale