submission 68006

from utils import make_match_reference, DeterministicContext
import torch
from task import input_t, output_t


def ref_kernel(data: input_t) -> output_t:
    """
    Reference implementation of RGB to grayscale conversion using PyTorch.
    Uses the standard coefficients: Y = 0.2989 R + 0.5870 G + 0.1140 B

    Args:
        data: RGB tensor of shape (H, W, 3) with values in [0, 1]
    Returns:
        Grayscale tensor of shape (H, W) with values in [0, 1]
    """
    with DeterministicContext():
        data, output = data
        # Standard RGB to Grayscale coefficients
        weights = torch.tensor(
            [0.2989, 0.5870, 0.1140], device=data.device, dtype=data.dtype
        )
        output[...] = torch.sum(data * weights, dim=-1)
        return output


def generate_input(size: int, seed: int) -> input_t:
    """
    Generates random RGB image tensor of specified size.
    Returns:
        Tensor of shape (size, size, 3) with values in [0, 1]
    """
    gen = torch.Generator(device="cuda")
    gen.manual_seed(seed)

    x = torch.rand(
        size, size, 3, device="cuda", dtype=torch.float32, generator=gen
    ).contiguous()

    y = torch.empty(size, size, device="cuda", dtype=torch.float32).contiguous()

    return x, y


check_implementation = make_match_reference(ref_kernel, rtol=1e-4, atol=1e-4)

import torch
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t

# Choose arch flags to generate native SASS for your GPU (e.g., B200 -> sm_100).
_cc = "".join(map(str, torch.cuda.get_device_capability()))
_cuda_cflags = [
    "-O3",
    "--use_fast_math",
    f"-gencode=arch=compute_{_cc},code=sm_{_cc}",
    f"-gencode=arch=compute_{_cc},code=compute_{_cc}",
]

rgb2gray_cpp_source = r"""
#include <torch/extension.h>
torch::Tensor rgb2gray_cuda(torch::Tensor x, torch::Tensor out);
"""

# Build once; if arch flag isn't recognized, retry without it.
rgb2gray_module = load_inline(
    name="rgb2gray_cuda",
    cpp_sources=rgb2gray_cpp_source,
    cuda_sources="""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_fp16.h>
#include <cstdint>

template <typename scalar_t>
__global__ void rgb2gray_kernel_fp(const scalar_t* __restrict__ x,
                                   scalar_t* __restrict__ y,
                                   int64_t N) {
    const int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int64_t stride = (int64_t)blockDim.x * gridDim.x;
    for (int64_t i = idx; i < N; i += stride) {
        const int64_t j = i * 3; // HWC contiguous: 3 scalars per pixel
        scalar_t r = x[j + 0];
        scalar_t g = x[j + 1];
        scalar_t b = x[j + 2];
        y[i] = r * scalar_t(0.2989) + g * scalar_t(0.5870) + b * scalar_t(0.1140);
    }
}

__global__ void rgb2gray_kernel_half(const __half* __restrict__ x,
                                     __half* __restrict__ y,
                                     int64_t N) {
    const int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int64_t stride = (int64_t)blockDim.x * gridDim.x;
    for (int64_t i = idx; i < N; i += stride) {
        const int64_t j = i * 3;
        float r = __half2float(x[j + 0]);
        float g = __half2float(x[j + 1]);
        float b = __half2float(x[j + 2]);
        float val = __fmaf_rn(r, 0.2989f, __fmaf_rn(g, 0.5870f, b * 0.1140f));
        y[i] = __float2half(val);
    }
}

torch::Tensor rgb2gray_cuda(torch::Tensor x, torch::Tensor out) {
    TORCH_CHECK(x.device().is_cuda(), "x must be CUDA");
    TORCH_CHECK(out.device().is_cuda(), "out must be CUDA");
    TORCH_CHECK(x.is_contiguous() && out.is_contiguous(), "x/out must be contiguous");
    TORCH_CHECK(x.dim() == 3 && x.size(2) == 3, "x must be HxWx3");
    TORCH_CHECK(out.dim() == 2 && out.size(0) == x.size(0) && out.size(1) == x.size(1),
                "out must be HxW matching x's H,W");
    TORCH_CHECK(x.scalar_type() == out.scalar_type(), "dtype mismatch x/out");

    const int64_t N = x.size(0) * x.size(1);
    if (N == 0) return out;

    const int threads = 1024;  // good bandwidth setting on B200
    const int sm = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
    const int max_blocks = sm * 16;
    const int blocks = (int)std::min<int64_t>((N + threads - 1) / threads, (int64_t)max_blocks);
    auto stream = at::cuda::getCurrentCUDAStream();

    if (x.scalar_type() == at::kHalf) {
        const __half* xp = reinterpret_cast<const __half*>(x.data_ptr<at::Half>());
        __half* yp = reinterpret_cast<__half*>(out.data_ptr<at::Half>());
        rgb2gray_kernel_half<<<blocks, threads, 0, stream>>>(xp, yp, N);
    } else {
        AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), "rgb2gray_kernel_fp", ([&] {
            const scalar_t* xp = x.data_ptr<scalar_t>();
            scalar_t* yp = out.data_ptr<scalar_t>();
            rgb2gray_kernel_fp<scalar_t><<<blocks, threads, 0, stream>>>(xp, yp, N);
        }));
    }

    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        throw std::runtime_error(cudaGetErrorString(err));
    }
    return out;
}
""",
    extra_cuda_cflags=_cuda_cflags,
    extra_cflags=["-O3"],
    functions=["rgb2gray_cuda"],
    verbose=False,
)


def rgb2gray(A: torch.Tensor, Out: torch.Tensor) -> torch.Tensor:
    return rgb2gray_module.rgb2gray_cuda(A, Out)


def custom_kernel(data: input_t) -> output_t:
    """
    Fast CUDA RGB->grayscale: single pass, no temporaries.
    Works with float16/float32/float64. Expects HxWx3 (contiguous) -> HxW (contiguous).
    """
    x, out = data
    if not x.is_cuda or not out.is_cuda:
        raise RuntimeError("Both tensors must be on GPU")
    if x.dim() != 3 or x.size(-1) != 3:
        raise RuntimeError("Input must be HxWx3")
    if not x.is_contiguous() or not out.is_contiguous():
        x = x.contiguous()
        out = out.contiguous()
    return rgb2gray(x, out)


x, y = generate_input(1024, seed=42)
results = check_implementation((x, y), custom_kernel((x, y)))
print("Check implementation result:", results)