submission 68140

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_runtime.h>
#include <cstdint>

static __device__ __forceinline__ float dot_rgb(float r, float g, float b) {
    // y = 0.2989*r + 0.5870*g + 0.1140*b (FMA で演算数を削減)
    return fmaf(r, 0.2989f, fmaf(g, 0.5870f, b * 0.1140f));
}

// 4画素/スレッド、AoS(3) → float4×3 ロードで取り出し
__global__ void rgb2gray_kernel_vec4(const float* __restrict__ x,
                                     float* __restrict__ y,
                                     int64_t N) {
    // i は「画素インデックス」。ここでは4の倍数から処理する
    int64_t idx4    = ((int64_t)blockIdx.x * blockDim.x + threadIdx.x) * 4;
    int64_t stride4 = (int64_t)blockDim.x * gridDim.x * 4;

    for (int64_t i = idx4; i < N; i += stride4) {
        if (i + 3 < N) {
            // 4画素分の先頭 (float 単位で 3*i)
            const int64_t j = i * 3;

            // 16Bベクトルロード（先頭アドレスが16B整列でなくても機能はしますが、
            // iを4の倍数に保っているので多くのケースで整列します）
            const float4* p4 = reinterpret_cast<const float4*>(x + j);
            float4 v0 = p4[0]; // [r0 g0 b0 r1]
            float4 v1 = p4[1]; // [g1 b1 r2 g2]
            float4 v2 = p4[2]; // [b2 r3 g3 b3]

            float g0 = dot_rgb(v0.x, v0.y, v0.z);
            float g1 = dot_rgb(v0.w, v1.x, v1.y);
            float g2 = dot_rgb(v1.z, v1.w, v2.x);
            float g3 = dot_rgb(v2.y, v2.z, v2.w);

            // 16Bベクトルストア
            reinterpret_cast<float4*>(y)[i >> 2] = make_float4(g0, g1, g2, g3);
        } else {
            // 端数処理（残り 1〜3 画素）
            const int64_t rem = N - i;
            int64_t j = i * 3;
            if (rem >= 1) y[i + 0] = dot_rgb(x[j + 0], x[j + 1], x[j + 2]);
            if (rem >= 2) y[i + 1] = dot_rgb(x[j + 3], x[j + 4], x[j + 5]);
            if (rem >= 3) y[i + 2] = dot_rgb(x[j + 6], x[j + 7], x[j + 8]);
        }
    }
}

torch::Tensor rgb2gray_cuda(torch::Tensor x, torch::Tensor out) {
    // // 必要ならチェックを戻してください
    // TORCH_CHECK(x.device().is_cuda() && out.device().is_cuda(), "x/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 (HWC)");
    // TORCH_CHECK(out.dim() == 2 && out.size(0) == x.size(0) && out.size(1) == x.size(1),
    //             "out must be HxW");
    // TORCH_CHECK(x.scalar_type() == at::kFloat && out.scalar_type() == at::kFloat,
    //             "float32 only");

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

    // 4画素/スレッド用にブロック数を見積もる
    const int threads = 396; // メモリ帯域支配なので 256 or 512 がバランス良
    const int sm = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
    const int64_t N4 = (N + 3) >> 2; // 4画素のグループ数
    const int blocks = (int)std::min<int64_t>((N4 + threads - 1) / threads, (int64_t)sm * 32) * 2;

    auto stream = at::cuda::getCurrentCUDAStream();
    rgb2gray_kernel_vec4<<<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)