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

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

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

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

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

29

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

53
    offsets = block_start + tl.arange(0, BLOCK_SIZE)[:, None] * VEC_WIDTH + tl.arange(0, VEC_WIDTH)[None, :]
54
    offsets = tl.reshape(offsets, (BLOCK_SIZE * VEC_WIDTH,))
55

56
    mask = offsets < n_elements
57

58
    a_vals = tl.load(a_ptr + offsets, mask=mask, other=0.0)
59
    b_vals = tl.load(b_ptr + offsets, mask=mask, other=0.0)
60
    c_vals = a_vals + b_vals
61

62
    tl.store(c_ptr + offsets, c_vals, mask=mask)
63

64
def solve_triton_1d(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, N: int):
65
    n_elements = N * N
66
    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE'] * meta['VEC_WIDTH']),)
67
    matrix_add_kernel_1d[grid](A, B, C, n_elements)
68

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

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

118
    a = tl.load(a_block_ptr, boundary_check=(0, 1))
119
    b = tl.load(b_block_ptr, boundary_check=(0, 1))
120
    c = a + b
121
    tl.store(c_block_ptr, c, boundary_check=(0, 1))
122

123
def solve_triton_2d(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, N: int):
124
    grid = lambda meta: (
125
        triton.cdiv(N, meta['BLOCK_M']),
126
        triton.cdiv(N, meta['BLOCK_N']),
127
    )
128
    matrix_add_kernel_2d[grid](A, B, C, N)
129

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

142
    # 计时
143
    start = time.perf_counter()
144
    for _ in range(repeat):
145
        func(A, B, C, N)
146
    torch.cuda.synchronize()
147
    end = time.perf_counter()
148

149
    avg_time_ms = (end - start) / repeat * 1000
150
    return avg_time_ms
151

152
def verify_results(C_triton_naive, C_triton_1d, C_triton_2d):
153
    """
154
    验证三种方案结果是否一致。
155
    """
156
    if torch.allclose(C_triton_naive, C_triton_1d, atol=1e-5):
157
        print("✅ PyTorch 与 Triton 1D 结果一致")
158
    else:
159
        print("❌ PyTorch 与 Triton 1D 结果不一致")
160

161
    if torch.allclose(C_triton_naive, C_triton_2d, atol=1e-5):
162
        print("✅ PyTorch 与 Triton 2D 结果一致")
163
    else:
164
        print("❌ PyTorch 与 Triton 2D 结果不一致")
165

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

173
    # 问题规模
174
    N = 4096
175

176
    # 分配 GPU 张量
177
    A = torch.randn(N, N, device=device, dtype=torch.float32)
178
    B = torch.randn(N, N, device=device, dtype=torch.float32)
179
    C_triton_naive = torch.empty_like(A)
180
    C_triton_1d = torch.empty_like(A)
181
    C_triton_2d = torch.empty_like(A)
182

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

185
    # 验证正确性（单独运行一次）
186
    solve_triton_naive(A, B, C_triton_naive, N)
187
    solve_triton_1d(A, B, C_triton_1d, N)
188
    solve_triton_2d(A, B, C_triton_2d, N)
189
    verify_results(C_triton_naive, C_triton_1d, C_triton_2d)
190

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

194
    time_pytorch = benchmark(solve_triton_naive, A, B, C_triton_naive, N)
195
    time_triton_1d = benchmark(solve_triton_1d, A, B, C_triton_1d, N)
196
    time_triton_2d = benchmark(solve_triton_2d, A, B, C_triton_2d, N)
197

198
    # 输出结果
199
    print(f"PyTorch 直接加法:      {time_pytorch:.4f} ms")
200
    print(f"Triton 1D 向量化:      {time_triton_1d:.4f} ms")
201
    print(f"Triton 2D 块指针:      {time_triton_2d:.4f} ms")
202

203
    # 计算加速比
204
    baseline = time_pytorch
205
    print(f"\n相对于 PyTorch 的加速比:")
206
    print(f"  Triton 1D: {baseline / time_triton_1d:.2f}x")
207
    print(f"  Triton 2D: {baseline / time_triton_2d:.2f}x")
208

209
    # 计算内存带宽
210
    bytes_per_element = 4  # float32
211
    total_bytes = 3 * N * N * bytes_per_element  # A读 + B读 + C写
212
    bw_pytorch = total_bytes / (time_pytorch / 1000) / 1e9
213
    bw_triton_1d = total_bytes / (time_triton_1d / 1000) / 1e9
214
    bw_triton_2d = total_bytes / (time_triton_2d / 1000) / 1e9
215

216
    print(f"\n估算内存带宽 (GB/s):")
217
    print(f"  PyTorch:  {bw_pytorch:.2f} GB/s")
218
    print(f"  Triton 1D: {bw_triton_1d:.2f} GB/s")
219
    print(f"  Triton 2D: {bw_triton_2d:.2f} GB/s")
220

221
if __name__ == "__main__":
222
    main()