submission 66367

import os, tempfile, hashlib, time
import torch
from torch.utils.cpp_extension import load
from task import input_t, output_t

CUDA_SRC = r"""
#include <ATen/ATen.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <torch/library.h>
#include <cstdio>
#include <cstdlib>

#ifndef CUDA_CHECK
#define CUDA_CHECK(x) do { auto _e=(x); if(_e!=cudaSuccess){ \
  fprintf(stderr,"CUDA error %s:%d: %s\n", __FILE__, __LINE__, cudaGetErrorString(_e)); \
  fflush(stderr); abort(); }} while(0)
#endif

// -------- Tunables (via -D flags) --------
#ifndef WAVES
#define WAVES 32            // grid oversubscription (try 24/32)
#endif
#ifndef FORCE_BLOCK
#define FORCE_BLOCK 256     // try 256 vs 512
#endif
#ifndef RGBGRAY_VALIDATE
#define RGBGRAY_VALIDATE 1  // set 0 to compile out guards
#endif
#ifndef PREFETCH_DIST
#define PREFETCH_DIST 192   // try 128/160/192/224
#endif
#ifndef USE_PTX
#define USE_PTX 1           // 1: .cg/.cs PTX + prefetch; 0: float4 path only
#endif
#ifndef ILP
#define ILP 2               // 2 tiles per loop to raise ILP; try 1 vs 2
#endif

#if __CUDA_ARCH__ >= 800 && USE_PTX
__device__ __forceinline__ void ld_v4_cg(const float* p, float& a,float& b,float& c,float& d){
  asm volatile("ld.global.cs.v4.f32 {%0,%1,%2,%3}, [%4];"
               : "=f"(a),"=f"(b),"=f"(c),"=f"(d) : "l"(p));
}
__device__ __forceinline__ void st_v4_cs(float* p, float a,float b,float c,float d){
  asm volatile("st.global.cs.v4.f32 [%0],{%1,%2,%3,%4};"::"l"(p),"f"(a),"f"(b),"f"(c),"f"(d));
}
__device__ __forceinline__ void prefetch_L2(const float* p){
  asm volatile("prefetch.global.L2 [%0];"::"l"(p));
}
#else
__device__ __forceinline__ void ld_v4_cg(const float* p, float& a,float& b,float& c,float& d){
  float4 v = reinterpret_cast<const float4*>(p)[0]; a=v.x; b=v.y; c=v.z; d=v.w;
}
__device__ __forceinline__ void st_v4_cs(float* p, float a,float b,float c,float d){
  reinterpret_cast<float4*>(p)[0] = make_float4(a,b,c,d);
}
__device__ __forceinline__ void prefetch_L2(const float*){}
#endif

// Process one 4-pixel tile at index i (assumes i+3 < HW)
__device__ __forceinline__ void tile4(const float* __restrict__ rgb,
                                      float* __restrict__ gray,
                                      int i)
{
  const float* p = rgb + 3*i;   // aligned to 16B when i%4==0
  float r0,g0,b0,r1;  float g1,b1,r2,g2;  float b2,r3,g3,b3;

  // 3 × 16B loads cover 12 floats for 4 pixels (HWC3 interleaved)
  ld_v4_cg(p + 0, r0,g0,b0,r1);
  ld_v4_cg(p + 4, g1,b1,r2,g2);
  ld_v4_cg(p + 8, b2,r3,g3,b3);

  // Y = 0.2989 R + 0.5870 G + 0.1140 B
  const float rC=0.2989f, gC=0.5870f, bC=0.1140f;
  float y0 = fmaf(bC, b0, fmaf(gC, g0, rC*r0));
  float y1 = fmaf(bC, b1, fmaf(gC, g1, rC*r1));
  float y2 = fmaf(bC, b2, fmaf(gC, g2, rC*r2));
  float y3 = fmaf(bC, b3, fmaf(gC, g3, rC*r3));

  st_v4_cs(gray + i, y0,y1,y2,y3);
}

__launch_bounds__(FORCE_BLOCK, 2)
__global__ void rgb2gray_hwc4px_kernel(const float* __restrict__ rgb,
                                       float* __restrict__ gray,
                                       int H, int W)
{
  const int HW  = H * W;
  const int HW4 = (HW & ~3);  // branch-free main path limit

  const int t = blockIdx.x * blockDim.x + threadIdx.x;
  const int step = (blockDim.x * gridDim.x) << 2;  // 4 px per iteration

  // ===== Main vectorized region [0 .. HW4) =====
#if ILP==2
  // ILP=2: each loop does two stride-separated tiles (i0, i1=i0+step)
  int i = (t << 2);
  while (i < HW4) {

    // Prefetch for i0 and i1 (if exist)
  #if PREFETCH_DIST > 0
    if (i + PREFETCH_DIST < HW)              prefetch_L2(rgb + 3 * (i + PREFETCH_DIST));
    int i1 = i + step;
    if (i1 < HW4 && (i1 + PREFETCH_DIST < HW)) prefetch_L2(rgb + 3 * (i1 + PREFETCH_DIST));
  #endif

    // Tile 0
    tile4(rgb, gray, i);

    // Tile 1 (if available)
    i += step;
    if (i < HW4) {
      tile4(rgb, gray, i);
      i += step;
    } else {
      break;
    }
  }
#else
  // ILP=1 (baseline): one tile per loop
  for (int i = (t << 2); i < HW4; i += step) {
  #if PREFETCH_DIST > 0
    if (i + PREFETCH_DIST < HW) prefetch_L2(rgb + 3 * (i + PREFETCH_DIST));
  #endif
    tile4(rgb, gray, i);
  }
#endif

  // ===== Tail region [HW4 .. HW) (<=3 pixels total per thread over all iters) =====
  // Start tail at the first index for this thread that is >= HW4
  int first_tail = (t << 2);
  if (first_tail < HW4) {
    int delta = HW4 - first_tail;
    int m = (delta + step - 1) / step;   // ceil_div(delta, step)
    first_tail += m * step;
  }
  for (int itail = first_tail; itail < HW; itail += step) {
    #pragma unroll
    for (int k=0; k<4; ++k) {
      const int idx = itail + k;
      if (idx < HW) {
        const float* q = rgb + 3*idx;
        float y = fmaf(0.1140f,q[2], fmaf(0.5870f,q[1], 0.2989f*q[0]));
        gray[idx] = y;
      }
    }
  }
}

// ---------------- C++ wrappers (OUT + return) ----------------
static at::Tensor hwc2gray_out_cuda(const at::Tensor& x, at::Tensor y){
#if RGBGRAY_VALIDATE
  TORCH_CHECK(x.is_cuda() && y.is_cuda(), "tensors must be CUDA");
  TORCH_CHECK(x.dtype()==at::kFloat && y.dtype()==at::kFloat, "dtype must be float32");
  TORCH_CHECK(x.dim()==3 && x.size(2)==3, "x must be (H,W,3)");
  TORCH_CHECK(y.dim()==2 && y.size(0)==x.size(0) && y.size(1)==x.size(1), "y must be (H,W)");
  TORCH_CHECK(x.is_contiguous() && y.is_contiguous(), "inputs must be contiguous");
#endif
  const int H = (int)x.size(0), W = (int)x.size(1);

  int sm=0; CUDA_CHECK(cudaDeviceGetAttribute(&sm, cudaDevAttrMultiProcessorCount, 0));

  int block = FORCE_BLOCK ? FORCE_BLOCK : 512;
  if(block<256) block=256; if(block>512) block=512;

  long long HWll = (long long)H * (long long)W;
  long long per_iter = (long long)block * 4 * (long long)(ILP);
  int need  = (int)((HWll + per_iter - 1) / per_iter);
  int grid  = need; if (grid < sm) grid=sm; if (grid > sm*WAVES) grid = sm*WAVES;

  rgb2gray_hwc4px_kernel<<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
      x.data_ptr<float>(), y.data_ptr<float>(), H, W);
  CUDA_CHECK(cudaGetLastError());
  return y;
}

static at::Tensor hwc2gray_cuda(const at::Tensor& x){
  auto y = at::empty({x.size(0), x.size(1)}, x.options().dtype(at::kFloat));
  return hwc2gray_out_cuda(x, y);
}

// Register ops (no pybind11)
TORCH_LIBRARY(rgb, m) {
  m.def("hwc2gray(Tensor x) -> Tensor");
  m.def("hwc2gray_out(Tensor x, Tensor(a!) y) -> Tensor(a!)");
}
TORCH_LIBRARY_IMPL(rgb, CUDA, m) {
  m.impl("hwc2gray", hwc2gray_cuda);
  m.impl("hwc2gray_out", hwc2gray_out_cuda);
}
"""

def _build_and_load():
    key = CUDA_SRC.encode()
    tag = hashlib.sha1(key).hexdigest()[:10]
    build_dir = os.path.join(tempfile.gettempdir(), f"rgb_min_fix_{tag}")
    os.makedirs(build_dir, exist_ok=True)
    cu_path = os.path.join(build_dir, "gray_hwc_op.cu")
    if not os.path.exists(cu_path):
        with open(cu_path, "w") as f: f.write(CUDA_SRC)

    flags = [
        "-O3","-lineinfo",
        "-DWAVES=1536",           # tune: 24/32
        "-DFORCE_BLOCK=512",    # tune: 256 vs 512
        "-DPREFETCH_DIST=192",  # tune: 128/160/192/224
        "-DUSE_PTX=0",
        "-DILP=2",              # ILP=2 for ~1–3% gain
        "-DRGBGRAY_VALIDATE=0",
        # default L2-only for any non-PTX loads:
        "-Xptxas","-dlcm=cg",
        # keep occupancy with ILP=2:
        #"-maxrregcount=64",
        "-use_fast_math",
    ]
    return load(
        name=f"rgb_gray_min_fix_{tag}",
        sources=[cu_path],
        is_python_module=False,
        extra_cuda_cflags=flags,
        extra_cflags=["-O3"],
        verbose=False,
    )

_build_and_load()


def custom_kernel(data):
    x, y = data
    torch.ops.rgb.hwc2gray_out(x, y)
    return y