submission 68049

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);
    }
}

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;
    const int sm = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
    const int max_blocks = sm * 128;
    const int blocks = (int)std::min<int64_t>((N + threads - 1) / threads, (int64_t)max_blocks);
    auto stream = at::cuda::getCurrentCUDAStream();

    rgb2gray_kernel_fp<float><<<blocks, threads, 0, stream>>>(
        x.data_ptr<float>(), out.data_ptr<float>(), N);

    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        throw std::runtime_error(cudaGetErrorString(err));
    }
    return out;
}
""",
    extra_cuda_cflags=_cuda_cflags,
    extra_cflags=["-O3", "--fast-math"],
    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:
    x, out = data
    return rgb2gray(x, out)


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