更新中

在 LeetGPU 的习题列表中，Element-wise 算子指的是对输入张量/数组的每个元素独立执行相同操作、不依赖于其他元素或全局统计量的算子。

以下是明确的 Element-wise 算子：

算子说明#

名称	说明
Vector Addition	两个向量逐元素相加
Matrix Addition	两个矩阵逐元素相加
Matrix Copy	逐元素复制矩阵
Color Inversion	对每个像素独立取反
Reverse Array	反转数组，每个元素独立移动位置
ReLU	逐元素应用 ReLU 函数
Leaky ReLU	逐元素应用 Leaky ReLU
Sigmoid Activation	逐元素应用 Sigmoid 函数
Value Clipping	逐元素裁剪到指定范围
Sigmoid Linear Unit (SiLU)	逐元素 SiLU 激活
Swish-Gated Linear Unit (SWiGLU)	逐元素 SWiGLU（门控部分也为逐元素）
Gaussian Error Gated Linear Unit (GEGLU)	逐元素 GEGLU 激活
RGB to Grayscale	每个像素独立转换，不依赖邻域
Interleave Arrays	交替合并两数组，每个输出元素仅依赖对应位置输入
Rotary Positional Embedding	对每个位置独立应用旋转矩阵
Weight Dequantization	每个权重独立反量化
INT8 Quantized MatMul（仅反量化部分）	反量化部分为逐元素，整体不是
Simple Inference	线性层前向包含矩阵乘，非 element-wise，但其中的激活部分可能是逐元素

Vector Addition#

已讲解

CUDA学习之路：Element Wise详解

1. Matrix Addition题目#

实现一个在 GPU 上对两个包含 32 位浮点数的矩阵进行逐元素相加的程序。程序接收两个相同维度的输入矩阵，输出一个矩阵，其中的每个元素为对应位置元素之和。

示例 1：

1
输入：
2
A = [[1.0, 2.0],
3
     [3.0, 4.0]]
4
B = [[5.0, 6.0],
5
     [7.0, 8.0]]
6

7
输出：
8
C = [[6.0, 8.0],
9
     [10.0, 12.0]]

示例 2：

1
输入：
2
A = [[1.5, 2.5, 3.5],
3
     [4.5, 5.5, 6.5],
4
     [7.5, 8.5, 9.5]]
5
B = [[0.5, 0.5, 0.5],
6
     [0.5, 0.5, 0.5],
7
     [0.5, 0.5, 0.5]]
8

9
输出：
10
C = [[2.0, 3.0, 4.0],
11
     [5.0, 6.0, 7.0],
12
     [8.0, 9.0, 10.0]]

约束条件

输入矩阵 A 和 B 维度相同
1 ≤ N ≤ 4096
所有元素均为 32 位浮点数
性能评测基于 N = 4,096

2. Pytorch题解#

1
import torch
2
# A, B, C are tensors on the GPU
3
def solve(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, N: int):
4
    C.copy_(A + B)

3. Triton题解#

这里应该如何优化呢？

(1) 向量化 (2)二维直接计算

3.1. 向量化如何计算offset？#

在编写Triton Kernel的时候,我们希望一个现成一次处理多个连续元素,例如,每个线程一次读取4个float32, 就等价于一条128位宽的向量加载指令,能够显著提升内存带宽利用率.

那么如何来生成每个线程要访问的全局内存偏移量呢?

一个直观的例子#

假设每个线程块有4个线程,每个线程一次处理2个元素,则BLOCK_SIZE=4，VEC_WIDTH=2，pid=0（起始 0）.

我们希望这4个线程共同覆盖一块连续的内存区域,总元素为BLOCK_SIZE * VEC_WIDTH = 8 . 一种自然的分配方式是:

1
线程0: [0, 1]
2
线程1: [2, 3]
3
线程2: [4, 5]
4
线程3: [6, 7]

因此我们希望Triton能够生成一个形状为(4, 2)的偏移量矩阵,其中第i行第j列的值刚好是全局内存的线性索引.

1
block_start = pid * BLOCK_SIZE * VEC_WIDTH
2

3
offsets = (
4
    block_start # 当前线程块的起始偏移
5
    + tl.arange(0, BLOCK_SIZE)[:, None] * VEC_WIDTH   # 每个线程的基地址偏移
6
    + tl.arange(0, VEC_WIDTH)[None, :]                # 线程内的向量偏移
7
)

按步骤拆解#

首先是计算线程块的起始偏移: block_start = pid * BLOCK_SIZE * VEC_WIDTH. 每个线程块整体负责的数量为BLOCK_SIZE * VEC_WIDTH, 乘以pid后,就得到当前线程块在整个数据中的起始索引.

其次计算每个线程负责的第一个元素偏移: tl.arange(0, BLOCK_SIZE)[:, None] * VEC_WIDTH. tl.arange(0, BLOCK_SIZE) 生成一个一维向量 [0, 1, 2, 3]，代表线程的局部 ID。

[:, None] 将其转换为列向量 (BLOCK_SIZE, 1)，即形状变为 4×1.

1
[[0],
2
 [1],
3
 [2],
4
 [3]]

乘以VEC_WIDTH后得到:

1
[[0],   # 线程0负责的第一个元素距离块起始的偏移
2
 [2],   # 线程1负责的第一个元素偏移
3
 [4],
4
 [6]]

最后再加上线程内的连续向量偏移: tl.arange(0, VEC_WIDTH)[None, :] 同样先生成一个[0, 1]的范围,然后[None,:]将其转换为行向量,即形状为(1, 2).

现在,将三者相加,根据Triton的广播机制,列会沿着列方向复制,行会沿着行方向复制,最终的到一个 $4 \times 2$ 的矩阵:

1
[[0 + 0, 0 + 1],     → [[0, 1],
2
 [2 + 0, 2 + 1],        [2, 3],
3
 [4 + 0, 4 + 1],        [4, 5],
4
 [6 + 0, 6 + 1]]        [6, 7]]

计算完成偏移量之后,我们将其展平为一维数组:

1
offsets = tl.reshape(offsets, (BLOCK_SIZE * VEC_WIDTH,))

这样我们就可以直接用这个偏移量数组调用tl.load或者tl.store,一次性进行向量化访问.

3.2. 二维块指针构造#

实际上,在处理矩阵、图像这类天然具有二维结构的数据时，将并行任务也按二维方式划分，往往能让代码思路更加清晰。此时，每个线程块直接对应数据中的一个子矩阵（或子图像区域），省去了将二维坐标手动转换为一维线性 ID 的麻烦。更重要的是，这种划分方式让我们可以灵活调整行、列两个方向上的并行度，更好地适配不同硬件架构与问题规模。

Triton 提供了 tl.make_block_ptr 来优雅地描述这种二维块访问模式。我们先通过一个直观的例子感受它的便利，再逐一拆解每个参数的含义。

一个直观的例子#

假设我们有一个4*4的矩阵A，在GPU以行优先存储，内存布局如下：

我们有一个 4×4 的矩阵 A，在 GPU 内存中按行优先（row-major）方式存储，其内存布局如下图所示（数字表示元素在内存中的线性偏移）：

1
列 0  列 1  列 2  列 3
2
┌────┬────┬────┬────┐
3
│ 0  │ 1  │ 2  │ 3  │  行 0
4
├────┼────┼────┼────┤
5
│ 4  │ 5  │ 6  │ 7  │  行 1
6
├────┼────┼────┼────┤
7
│ 8  │ 9  │10  │11  │  行 2
8
├────┼────┼────┼────┤
9
│12  │13  │14  │15  │  行 3
10
└────┴────┴────┴────┘

现在我们希望在 kernel 中每次处理一个 2×2 的子块。例如，假设当前线程块负责处理图中右下角的子块（包含元素 10, 11, 14, 15），它的左上角位于矩阵的第 2 行、第 2 列。

我们使用tl.make_block_ptr来描述这个子块:

1
a_block_ptr = tl.make_block_ptr(
2
    base=a_ptr,                     # ① 内存基地址
3
    shape=(4, 4),                   # ② 完整矩阵形状
4
    strides=(4, 1),                 # ③ 行步长、列步长
5
    offsets=(2, 2),                 # ④ 子块左上角坐标
6
    block_shape=(2, 2),             # ⑤ 子块大小
7
    order=(1, 0)                    # ⑥ 线程映射顺序
8
)

创建好这个块指针后，我们只需一行 tl.load(a_block_ptr)，Triton 便会自动计算出所有需要访问的全局内存地址，并安全地处理边界情况。接下来，我们详细看看每个参数的具体作用。

参数详解#

base —— 内存基地址,即矩阵在 GPU 内存中的起始地址，通常通过 PyTorch Tensor 的 .data_ptr() 方法获得。
shape —— 完整数据的逻辑形状. 此处的 (4, 4) 告诉 Triton 矩阵共有4行4列。
strides —— 各维度的内存跨步. 这是连接“逻辑坐标”与“物理地址”的关键。对于行优先存储的矩阵：

strides详解

沿着第 0 维（行）移动 1，意味着从当前行跳到下一行，需要跨越一整行的元素，因此步长为 4（即一行有 4 个元素）。

沿着第 1 维（列）移动 1，意味着在同一行内移动到下一列，相邻元素在内存中紧挨着，因此步长为 1。所以 strides = (4, 1)。反之，如果矩阵是列优先存储，则 strides 应为 (1, 4)。

offsets —— 子块左上角的起始坐标. 它是一个元组 (row_start, col_start)。

offsets计算详解

在实际 kernel 中，offsets 会根据线程块的 ID 动态计算，例如：

1
pid_m = tl.program_id(0)   # 行方向的块 ID
2
pid_n = tl.program_id(1)   # 列方向的块 ID
3
offsets = (pid_m * BLOCK_M, pid_n * BLOCK_N)

block_shape —— 要加载的块大小.

这是我们在 kernel 中设定的超参数，例如 BLOCK_M = 2、BLOCK_N = 2。它决定了每次 tl.load 会从内存中读取多少个元素。合适的块大小需要根据共享内存容量、计算资源以及数据复用程度进行权衡，常见取值如 64×64、128×128 等。

order —— 线程数据分布映射顺序

这个参数稍微有些进阶。简单来说，它告诉编译器数据在块内是如何分配给各个线程的，会影响生成的 PTX 代码结构。在大多数常规场景下（例如非 Tensor Core 的简单加载/存储），order=(0,1) 与 order=(1,0) 的性能差异微乎其微，因此你不必过分纠结。

不过，当你使用 NVIDIA Hopper 架构（H100 等）并希望启用 TMA（Tensor Memory Accelerator）硬件加速时，正确设置 order 就显得至关重要了，它可以带来数倍的访存性能提升。一般情况下，对于矩阵乘法中的 A 矩阵加载，常用 order=(1,0)；对于 B 矩阵，常用 order=(0,1)。若想深入探究，可以借我一张H100测试一下!

3.3. Naive, Vec, 2D 完整实现#

1
import torch
2
import triton
3
import triton.language as tl
4
import time
5
import numpy as np
6

7
# 方案 1：直接加法
8
@triton.jit
9
def matrix_add_kernel(a, b, c, n_elements, BLOCK_SIZE: tl.constexpr):
10
    pid = tl.program_id(axis=0)
11

12
    block_start = pid * BLOCK_SIZE
13
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
14
    mask = offsets < n_elements
15

16
    ga = tl.load(a + offsets, mask=mask)
17
    gb = tl.load(b + offsets, mask=mask)
18
    gc = ga + gb
19
    tl.store(c + offsets, gc, mask=mask)
20

21
def solve_triton_naive(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, N: int):
22
    BLOCK_SIZE = 1024
23
    n_elements = N * N
24
    grid = (triton.cdiv(n_elements, BLOCK_SIZE),)
25
    matrix_add_kernel[grid](a, b, c, n_elements, BLOCK_SIZE)
26

27

28
# 方案 2：Triton 一维向量化 + Autotune
29
@triton.autotune(
30
    configs=[
31
        triton.Config({'BLOCK_SIZE': 1024, 'VEC_WIDTH': 1}, num_warps=4),
32
        triton.Config({'BLOCK_SIZE': 1024, 'VEC_WIDTH': 2}, num_warps=4),
33
        triton.Config({'BLOCK_SIZE': 2048, 'VEC_WIDTH': 2}, num_warps=8),
34
        triton.Config({'BLOCK_SIZE': 4096, 'VEC_WIDTH': 4}, num_warps=8),
35
        triton.Config({'BLOCK_SIZE': 4096, 'VEC_WIDTH': 8}, num_warps=16),
36
    ],
37
    key=['n_elements'],
38
)
39
@triton.jit
40
def matrix_add_kernel_1d(
41
    a_ptr, b_ptr, c_ptr,
42
    n_elements: tl.constexpr,
43
    BLOCK_SIZE: tl.constexpr,
44
    VEC_WIDTH: tl.constexpr,
45
):
46
    pid = tl.program_id(axis=0)
47
    block_start = pid * BLOCK_SIZE * VEC_WIDTH
48

49
    offsets = block_start + tl.arange(0, BLOCK_SIZE)[:, None] * VEC_WIDTH + tl.arange(0, VEC_WIDTH)[None, :]
50
    offsets = tl.reshape(offsets, (BLOCK_SIZE * VEC_WIDTH,))
51

52
    mask = offsets < n_elements
53

54
    a_vals = tl.load(a_ptr + offsets, mask=mask, other=0.0)
55
    b_vals = tl.load(b_ptr + offsets, mask=mask, other=0.0)
56
    c_vals = a_vals + b_vals
57

58
    tl.store(c_ptr + offsets, c_vals, mask=mask)
59

60
def solve_triton_1d(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, N: int):
61
    n_elements = N * N
62
    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE'] * meta['VEC_WIDTH']),)
63
    matrix_add_kernel_1d[grid](A, B, C, n_elements)
64

65
# 方案 3：Triton 二维块指针 + Autotune
66
@triton.autotune(
67
    configs=[
68
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128}, num_warps=4),
69
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 256}, num_warps=4),
70
        triton.Config({'BLOCK_M': 256, 'BLOCK_N': 128}, num_warps=8),
71
        triton.Config({'BLOCK_M': 256, 'BLOCK_N': 256}, num_warps=8),
72
        triton.Config({'BLOCK_M': 512, 'BLOCK_N': 128}, num_warps=8),
73
        triton.Config({'BLOCK_M': 512, 'BLOCK_N': 256}, num_warps=8),
74
        triton.Config({'BLOCK_M': 512, 'BLOCK_N': 512}, num_warps=8),
75
    ],
76
    key=['N'],
77
)
78
@triton.jit
79
def matrix_add_kernel_2d(
80
    a_ptr, b_ptr, c_ptr,
81
    N,
82
    BLOCK_M: tl.constexpr,
83
    BLOCK_N: tl.constexpr,
84
):
85
    pid_m = tl.program_id(axis=0)
86
    pid_n = tl.program_id(axis=1)
87

88
    a_block_ptr = tl.make_block_ptr(
89
        base=a_ptr,
90
        shape=(N, N),
91
        strides=(N, 1),
92
        offsets=(pid_m * BLOCK_M, pid_n * BLOCK_N),
93
        block_shape=(BLOCK_M, BLOCK_N),
94
        order=(1, 0),
95
    )
96
    b_block_ptr = tl.make_block_ptr(
97
        base=b_ptr,
98
        shape=(N, N),
99
        strides=(N, 1),
100
        offsets=(pid_m * BLOCK_M, pid_n * BLOCK_N),
101
        block_shape=(BLOCK_M, BLOCK_N),
102
        order=(1, 0),
103
    )
104
    c_block_ptr = tl.make_block_ptr(
105
        base=c_ptr,
106
        shape=(N, N),
107
        strides=(N, 1),
108
        offsets=(pid_m * BLOCK_M, pid_n * BLOCK_N),
109
        block_shape=(BLOCK_M, BLOCK_N),
110
        order=(1, 0),
111
    )
112

113
    a = tl.load(a_block_ptr, boundary_check=(0, 1))
114
    b = tl.load(b_block_ptr, boundary_check=(0, 1))
115
    c = a + b
116
    tl.store(c_block_ptr, c, boundary_check=(0, 1))
117

118
def solve_triton_2d(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, N: int):
119
    grid = lambda meta: (
120
        triton.cdiv(N, meta['BLOCK_M']),
121
        triton.cdiv(N, meta['BLOCK_N']),
122
    )
123
    matrix_add_kernel_2d[grid](A, B, C, N)
124

125
# ------------------------------------------------------------
126
# 性能测试工具
127
# ------------------------------------------------------------
128
def benchmark(func, A, B, C, N, warmup=10, repeat=100):
129
    """
130
    运行指定函数多次，返回平均耗时（毫秒）。
131
    """
132
    # 预热
133
    for _ in range(warmup):
134
        func(A, B, C, N)
135
    torch.cuda.synchronize()
136

137
    # 计时
138
    start = time.perf_counter()
139
    for _ in range(repeat):
140
        func(A, B, C, N)
141
    torch.cuda.synchronize()
142
    end = time.perf_counter()
143

144
    avg_time_ms = (end - start) / repeat * 1000
145
    return avg_time_ms
146

147
def verify_results(C_triton_naive, C_triton_1d, C_triton_2d):
148
    """
149
    验证三种方案结果是否一致。
150
    """
151
    if torch.allclose(C_triton_naive, C_triton_1d, atol=1e-5):
152
        print("✅ Naive Triton 与 Triton 1D 结果一致")
153
    else:
154
        print("❌ Naive Triton 与 Triton 1D 结果不一致")
155

156
    if torch.allclose(C_triton_naive, C_triton_2d, atol=1e-5):
157
        print("✅ Naive Triton 与 Triton 2D 结果一致")
158
    else:
159
        print("❌ Naive Triton 与 Triton 2D 结果不一致")
160

161
def main():
162
    # 检查 CUDA 可用性
163
    if not torch.cuda.is_available():
164
        raise RuntimeError("CUDA 不可用，请在有 GPU 的环境下运行")
165
    device = torch.device("cuda")
166
    print(f"运行设备: {torch.cuda.get_device_name(device)}")
167

168
    # 问题规模
169
    N = 4096
170

171
    # 分配 GPU 张量
172
    A = torch.randn(N, N, device=device, dtype=torch.float32)
173
    B = torch.randn(N, N, device=device, dtype=torch.float32)
174
    C_triton_naive = torch.empty_like(A)
175
    C_triton_1d = torch.empty_like(A)
176
    C_triton_2d = torch.empty_like(A)
177

178
    print(f"\n矩阵大小: {N} x {N} ({N*N} 个元素)")
179

180
    # 验证正确性
181
    solve_triton_naive(A, B, C_triton_naive, N)
182
    solve_triton_1d(A, B, C_triton_1d, N)
183
    solve_triton_2d(A, B, C_triton_2d, N)
184
    verify_results(C_triton_naive, C_triton_1d, C_triton_2d)
185

186
    # 性能测试
187
    print("\n开始性能测试 (预热 10 次，计时 100 次取平均)...\n")
188

189
    time_pytorch = benchmark(solve_triton_naive, A, B, C_triton_naive, N)
190
    time_triton_1d = benchmark(solve_triton_1d, A, B, C_triton_1d, N)
191
    time_triton_2d = benchmark(solve_triton_2d, A, B, C_triton_2d, N)
192

193
    # 输出结果
194
    print(f"Triton 直接加法:      {time_pytorch:.4f} ms")
195
    print(f"Triton 1D 向量化:      {time_triton_1d:.4f} ms")
196
    print(f"Triton 2D 块指针:      {time_triton_2d:.4f} ms")
197

198
    # 计算加速比
199
    baseline = time_pytorch
200
    print(f"\n相对于 Naive 的加速比:")
201
    print(f"  Triton 1D: {baseline / time_triton_1d:.2f}x")
202
    print(f"  Triton 2D: {baseline / time_triton_2d:.2f}x")
203

204
    # 计算内存带宽
205
    bytes_per_element = 4  # float32
206
    total_bytes = 3 * N * N * bytes_per_element  # A读 + B读 + C写
207
    bw_pytorch = total_bytes / (time_pytorch / 1000) / 1e9
208
    bw_triton_1d = total_bytes / (time_triton_1d / 1000) / 1e9
209
    bw_triton_2d = total_bytes / (time_triton_2d / 1000) / 1e9
210

211
    print(f"\n估算内存带宽 (GB/s):")
212
    print(f"  Triton Naive:  {bw_pytorch:.2f} GB/s")
213
    print(f"  Triton 1D: {bw_triton_1d:.2f} GB/s")
214
    print(f"  Triton 2D: {bw_triton_2d:.2f} GB/s")
215

216
if __name__ == "__main__":
217
    main()

音乐

音乐

算子说明#

Vector Addition#

1. Matrix Addition题目#

2. Pytorch题解#

3. Triton题解#

3.1. 向量化如何计算offset？#

一个直观的例子#

按步骤拆解#

3.2. 二维块指针构造#

一个直观的例子#

参数详解#

3.3. Naive, Vec, 2D 完整实现#

4. CUDA 题解#

支持与分享

音乐

目录

音乐

音乐

LeetGPU习题01：Matrix Addition

算子说明#

Vector Addition#

1. Matrix Addition题目#

2. Pytorch题解#

3. Triton题解#

3.1. 向量化如何计算offset？#

一个直观的例子#

按步骤拆解#

3.2. 二维块指针构造#

一个直观的例子#

参数详解#

3.3. Naive, Vec, 2D 完整实现#

4. CUDA 题解#

支持与分享

音乐

目录