RGB to Grayscale

原始题目：LeetGPU - RGB to Grayscale

题目描述

编写一个 GPU 程序，在 GPU 上将 RGB 图像转换为灰度图像。输入 RGB 图像表示为 32 位浮点值的一维数组，使用标准的 RGB 到灰度转换公式计算对应的灰度图像：

$$gray = 0.299 \times R + 0.587 \times G + 0.114 \times B$$

输入数组 input 包含 $height \times width \times 3$ 个元素，每个像素的 RGB 值连续存储 $[R, G, B, R, G, B, \dots]$。输出数组 output 应包含 $height \times width$ 个灰度值。

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。
• 最终结果必须存储在数组 output 中。
• 使用精确的系数：红色 0.299、绿色 0.587、蓝色 0.114。

示例

示例 1

Input:  input = [255.0, 0.0, 0.0, 0.0, 255.0, 0.0, 0.0, 0.0, 255.0, 128.0, 128.0, 128.0]
        width = 2, height = 2
Output: [76.245, 149.685, 29.07, 128.0]

示例 2

Input:  input = [100.0, 150.0, 200.0], width = 1, height = 1
Output: [140.75]

约束条件

• $1 \le width \le 4{,}096$，$1 \le height \le 4{,}096$。
• $width \times height \le 4{,}194{,}304$。
• 所有 RGB 值在 $[0.0, 255.0]$ 范围内。
• 性能测试在 $height = 2{,}048,\ width = 2{,}048$ 的规模下进行。

解题思路

RGB 转灰度是典型的逐像素映射操作。每个线程处理一个像素（读取 3 个 float，计算加权和，写入 1 个 float）。算术强度低，属于内存带宽受限内核。float3 向量化读取可将 3 次 LDG 指令合并为 1 次 LDG.128，减少指令发射数。

代码实现

CUDA

#include <cuda_runtime.h>

// 解法一：基础逐像素转换
__global__ void rgb_to_grayscale_kernel(const float* input, float* output, int total_pixels) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < total_pixels) {
        int idx = i * 3;
        float r = input[idx];
        float g = input[idx + 1];
        float b = input[idx + 2];
        output[i] = 0.299f * r + 0.587f * g + 0.114f * b;
    }
}

// 解法二：Grid-Stride Loop
__global__ void rgb_to_grayscale_grid_stride(const float* input, float* output, int total_pixels) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = gridDim.x * blockDim.x;
    for (int i = idx; i < total_pixels; i += stride) {
        int off = i * 3;
        output[i] = 0.299f * input[off] + 0.587f * input[off + 1] + 0.114f * input[off + 2];
    }
}

// 解法三：float3 向量化读取（一次 LDG.128 搬运 12 字节 = R+G+B）
__global__ void rgb_to_grayscale_float3(const float* input, float* output, int total_pixels) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = gridDim.x * blockDim.x;
    const float3* in3 = (const float3*)input;
    for (int i = idx; i < total_pixels; i += stride) {
        float3 pixel = in3[i];
        output[i] = 0.299f * pixel.x + 0.587f * pixel.y + 0.114f * pixel.z;
    }
}

extern "C" void solve(const float* input, float* output, int width, int height) {
    int total_pixels = width * height;
    int threadsPerBlock = 256;
    int blocksPerGrid = (total_pixels + threadsPerBlock - 1) / threadsPerBlock;
    rgb_to_grayscale_kernel<<<blocksPerGrid, threadsPerBlock>>>(input, output, total_pixels);
    cudaDeviceSynchronize();
}

Triton

import triton
import triton.language as tl
import torch


@triton.jit
def rgb_to_grayscale_kernel(
    input_ptr, output_ptr,
    total_pixels: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(0)
    idx = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = idx < total_pixels

    off = idx * 3
    r = tl.load(input_ptr + off, mask=mask)
    g = tl.load(input_ptr + off + 1, mask=mask)
    b = tl.load(input_ptr + off + 2, mask=mask)

    gray = 0.299 * r + 0.587 * g + 0.114 * b
    tl.store(output_ptr + idx, gray, mask=mask)