Color Inversion

原始题目：LeetGPU - Color Inversion

题目描述

编写一个 GPU 程序来反转图像的颜色。图像表示为一个包含 RGBA（红、绿、蓝、Alpha）值的一维数组，每个分量均为 8 位无符号整数（unsigned char）。

颜色反转通过将每个颜色分量（R、G、B）从 255 中减去来实现，Alpha 分量保持不变：

$$\begin{aligned} R' &= 255 - R \\ G' &= 255 - G \\ B' &= 255 - B \\ A' &= A \end{aligned}$$

输入数组 image 包含 $width \times height \times 4$ 个元素。前 4 个元素表示左上角像素的 RGBA 值，接下来 4 个表示其右侧的像素，以此类推。

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。
• 最终结果必须存储在 image 数组中。

示例

示例 1

Input:  image = [255, 0, 128, 255, 0, 255, 0, 255], width = 1, height = 2
Output: [0, 255, 127, 255, 255, 0, 255, 255]

示例 2

Input:  image = [10, 20, 30, 255, 100, 150, 200, 255], width = 2, height = 1
Output: [245, 235, 225, 255, 155, 105, 55, 255]

约束条件

• $1 \le width \le 4{,}096$，$1 \le height \le 4{,}096$。
• $width \times height \le 8{,}388{,}608$。
• 每个像素 4 字节（RGBA），Alpha 通道需要保留。
• 性能测试在 $height = 5{,}120,\ width = 4{,}096$ 的规模下进行。

解题思路

颜色反转属于典型的逐元素映射（Element-wise Map）操作，每个像素独立计算，完全数据并行。可以直接将每个线程映射到一个像素（4 个颜色分量），也可以将每个线程映射到一个分量以提供更大的并行度。由于每个分量只需一次减法操作，计算强度极低，这是一个典型的内存带宽受限（Memory Bandwidth Bound）内核。

代码实现

CUDA

#include <cuda_runtime.h>

__global__ void color_invert_kernel(unsigned char* image, int total_pixels) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < total_pixels) {
        int idx = i * 4;
        image[idx]     = 255 - image[idx];      // R
        image[idx + 1] = 255 - image[idx + 1];  // G
        image[idx + 2] = 255 - image[idx + 2];  // B
        // Alpha (idx+3) unchanged
    }
}

extern "C" void solve(unsigned char* image, int width, int height) {
    int total_pixels = width * height;
    int threadsPerBlock = 256;
    int blocksPerGrid = (total_pixels + threadsPerBlock - 1) / threadsPerBlock;
    color_invert_kernel<<<blocksPerGrid, threadsPerBlock>>>(image, total_pixels);
    cudaDeviceSynchronize();
}

Triton

import triton
import triton.language as tl

@triton.jit
def color_invert_kernel(image_ptr, total_pixels: tl.constexpr, BLOCK_SIZE: tl.constexpr):
    i = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = i < total_pixels
    off = i * 4
    tl.store(image_ptr + off,     255 - tl.load(image_ptr + off,     mask=mask), mask=mask)
    tl.store(image_ptr + off + 1, 255 - tl.load(image_ptr + off + 1, mask=mask), mask=mask)
    tl.store(image_ptr + off + 2, 255 - tl.load(image_ptr + off + 2, mask=mask), mask=mask)
    # Alpha unchanged