LeetGPU习题03：Color inversion

1
import torch
2
import triton
3
import triton.language as tl
4

5

6
@triton.jit
7
def invert_kernel(
8
    image_ptr,
9
    width,
10
    height,
11
    BLOCK_SIZE_X: tl.constexpr,
12
    BLOCK_SIZE_Y: tl.constexpr
13
):
14
    pid_x = tl.program_id(0)
15
    pid_y = tl.program_id(1)
16

17
    rx = pid_x * BLOCK_SIZE_X + tl.arange(0, BLOCK_SIZE_X)
18
    ry = pid_y * BLOCK_SIZE_Y + tl.arange(0, BLOCK_SIZE_Y)
19

20
    # rx[:, None] 是 (BLOCK_SIZE_X, 1)，ry[None, :] 是 (1, BLOCK_SIZE_Y)
21
    mask = (rx[:, None] < height) & (ry[None, :] < width)
22

23
    # 计算 1D 数组中的基础偏移量：(row * width + col) * 4（因为是RGBA）
24
    # 这里利用了广播机制
25
    base_offsets = (rx[:, None] * width + ry[None, :]) * 4
26

27
    for i in range(3):
28
        ptr = image_ptr + base_offsets + i
29
        val = tl.load(ptr, mask=mask)
30
        tl.store(ptr, 255 - val, mask=mask)
31

32

33
def solve(image: torch.Tensor, width: int, height: int):
34
    BLOCK_SIZE_X = 32
35
    BLOCK_SIZE_Y = 32
36

37
    # 计算网格大小
38
    grid = (
39
        triton.cdiv(height, BLOCK_SIZE_X),
40
        triton.cdiv(width, BLOCK_SIZE_Y)
41
    )
42

43
    # 启动核函数
44
    invert_kernel[grid](
45
        image,
46
        width,
47
        height,
48
        BLOCK_SIZE_X=BLOCK_SIZE_X,
49
        BLOCK_SIZE_Y=BLOCK_SIZE_Y,
50
        num_warps=8
51
    )
52

53
    return image

1
__global__ void invert_naive(unsigned char* image, int N) {
2
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
3
    if (tid < N) {
4
        image[tid] = 255 - image[tid];
5
    }
6
}
7

8
__global__ void invert_stride(unsigned char* image, int N) {
9
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
10
    int stride = gridDim.x * blockDim.x;
11

12
    for (int i = tid; i < N; i += stride) {
13
        image[i] = ~image[i];  // ~x == 255 - x for unsigned char
14
    }
15
}
16

17
__global__ void invert_rgba_uchar4(unsigned char* image, int num_pixels) {
18
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
19
    int stride = gridDim.x * blockDim.x;
20

21
    uchar4* img4 = (uchar4*)image;
22

23
    for (int i = tid; i < num_pixels; i += stride) {
24
        uchar4 p = img4[i];
25

26
        p.x = 255 - p.x;
27
        p.y = 255 - p.y;
28
        p.z = 255 - p.z;
29
        // p.w (alpha) 保持不变
30

31
        img4[i] = p;
32
    }
33
}
34

35
__global__ void invert_2D(unsigned char* image, int width, int height) {
36
    int row = blockIdx.y * blockDim.y + threadIdx.y;
37
    int col = blockIdx.x * blockDim.x + threadIdx.x;
38

39
    if (col < width && row < height) {
40
        int idx = (row * width + col) * 4;
41
        image[idx + 0] = 255 - image[idx + 0];
42
        image[idx + 1] = 255 - image[idx + 1];
43
        image[idx + 2] = 255 - image[idx + 2];
44
    }
45
}