Star410
2703
代码
代码
Issues334
Pull Requests895
Actions
项目讨论
Wiki
分析
Star410
2703
已开启
【代码侦探·Challenge 09】 CATLASS DSL MatmulAddRelu 算子练习题 #4270
JeffDing创建于 28 天前
【代码侦探·Challenge 09】 CATLASS DSL MatmulAddRelu 算子练习题 #4270
已开启
JeffDing创建于 28 天前
从JeffDing/cann-outreachsubmit/Developer_JeffDing/challenge09_matmuladdrelu合入到cann/cann-outreachmaster
共 2 个文件变更+233-0
A2026/CANN-Code-Detective/Challenge09-MatmulAddRelu/PR_code/JeffDing/matmul_add_relu.py+170-0
@@ -0,0 +1,170 @@
1+"""Challenge 09：固定 256×256×256 的 Catlass Python 融合核骨架。"""
2+ 
3+import catlass.tla as tla
4+ 
5+ 
6+M = N = K = 256
7+L0_K = 64
8+AIV_M = M // 2
9+X_CHUNK_M = 64
10+SIMD_LANES = 64 # 256 字节 SIMD 寄存器可容纳 64 个 FP32 元素。
11+ 
12+ 
13+@tla.kernel
14+def matmul_add_relu_kernel(
15+ gm_a: tla.Tensor,
16+ gm_b: tla.Tensor,
17+ gm_x: tla.Tensor,
18+ gm_y: tla.Tensor,
19+) -> None:
20+ """仅计算 A(256,256) @ B(256,256) + X(256,256) 后的 ReLU。"""
21+ # mutex 管理核内缓冲的访问顺序；A/B 成组搬运和使用，共用一把锁。
22+ mutex_l1 = tla.mutex(resource="l1_ab", id=0)
23+ mutex_l0 = tla.mutex(resource="l0_ab", id=1)
24+ mutex_l0c = tla.mutex(resource="l0_c", id=2)
25+ mutex_acc = tla.mutex(resource="ub_acc", id=3)
26+ mutex_x = tla.mutex(resource="ub_x", id=4)
27+ 
28+ # v2.1.0 的 mutex 不支持跨核同步，FIX → AIV 仍需跨核通知。
29+ fix_done = tla.cross_flag("fix_done")
30+ 
31+ # L1: A/B 各 128 KiB；L0A/L0B: 各 32 KiB；L0C: 256 KiB。
32+ l1_a_ptr = tla.allocate(M * K, tla.Float16, tla.AddressSpace.l1, 512)
33+ l1_b_ptr = tla.allocate(K * N, tla.Float16, tla.AddressSpace.l1, 512)
34+ l0_a_ptr = tla.allocate(M * L0_K, tla.Float16, tla.AddressSpace.l0a, 512)
35+ l0_b_ptr = tla.allocate(L0_K * N, tla.Float16, tla.AddressSpace.l0b, 512)
36+ l0_c_ptr = tla.allocate(M * N, tla.Float32, tla.AddressSpace.l0c, 512)
37+ 
38+ # 每个 AIV 的 UB: 累加结果 128 KiB + X 分片 64 KiB = 192 KiB。
39+ ub_acc_ptr = tla.allocate(AIV_M * N, tla.Float32, tla.AddressSpace.ub, 256)
40+ ub_x_ptr = tla.allocate(X_CHUNK_M * N, tla.Float32, tla.AddressSpace.ub, 256)
41+ 
42+ with tla.cube():
43+ l1_a = tla.make_tensor_like(l1_a_ptr, gm_a, tla.arch.zN)
44+ l1_b = tla.make_tensor_like(l1_b_ptr, gm_b, tla.arch.zN)
45+ l0_c = tla.make_tensor_like(l0_c_ptr, gm_y, tla.arch.L0Clayout)
46+ 
47+ mutex_l1.lock(pipe=tla.arch.MTE2)
48+ # --------------------------------
49+ # TODO 1（GM → L1）：输入 gm_a/gm_b 是连续行优先的 FP16 矩阵。
50+ # 按 A、B 的顺序调用 tla.copy，分别搬到 L1 中 zN 布局的
51+ # l1_a/l1_b；mutex_l1 保证 MTE1 在两次搬运完成后读取。
52+ tla.copy(l1_a, gm_a)
53+ tla.copy(l1_b, gm_b)
54+ # --------------------------------
55+ mutex_l1.unlock(pipe=tla.arch.MTE2)
56+ 
57+ for k_tile in tla.range(0, K // L0_K, 1):
58+ l1_a_k = tla.tile_view(
59+ l1_a, tla.make_shape(M, L0_K), tla.make_coord(0, k_tile)
60+ )
61+ l1_b_k = tla.tile_view(
62+ l1_b, tla.make_shape(L0_K, N), tla.make_coord(k_tile, 0)
63+ )
64+ l0_a = tla.make_tensor_like(l0_a_ptr, l1_a_k, tla.arch.zN)
65+ l0_b = tla.make_tensor_like(l0_b_ptr, l1_b_k, tla.arch.nZ)
66+ 
67+ mutex_l1.lock(pipe=tla.arch.MTE1)
68+ mutex_l0.lock(pipe=tla.arch.MTE1)
69+ # --------------------------------
70+ # TODO 2（L1 → L0）：输入为 L1 的 l1_a_k/l1_b_k。
71+ # 按 A、B 的顺序调用 tla.copy，分别搬到 L0A 的 l0_a（zN）
72+ # 和 L0B 的 l0_b（nZ）；mutex_l0 保证 Cube 读取前搬运完成，
73+ # 并保证下一次搬运在当前矩阵乘读完 L0A/L0B 后开始。
74+ tla.copy(l0_a, l1_a_k)
75+ tla.copy(l0_b, l1_b_k)
76+ # --------------------------------
77+ mutex_l0.unlock(pipe=tla.arch.MTE1)
78+ mutex_l1.unlock(pipe=tla.arch.MTE1)
79+ 
80+ mutex_l0.lock(pipe=tla.arch.CUBE)
81+ mutex_l0c.lock(pipe=tla.arch.CUBE)
82+ # --------------------------------
83+ # TODO 3（矩阵乘累加）：用 tla.mmad 将 FP16 的 l0_a×l0_b
84+ # 累加到 L0C 中的 FP32 l0_c。k_tile=0 时 init_c=True，
85+ # 后续三次必须为 False，以保留前面 K 分片的累加结果。
86+ # 此处只计算矩阵乘；X 的相加和 ReLU 留给 Vector。
87+ init_c = k_tile == 0
88+ tla.mmad(l0_c, l0_a, l0_b, init_c=init_c)
89+ # --------------------------------
90+ mutex_l0c.unlock(pipe=tla.arch.CUBE)
91+ mutex_l0.unlock(pipe=tla.arch.CUBE)
92+ 
93+ ub_acc_full = tla.make_tensor_like(ub_acc_ptr, gm_y, tla.arch.RowMajor)
94+ mutex_l0c.lock(pipe=tla.arch.FIX)
95+ # --------------------------------
96+ # TODO 4（L0C → UB）：输入是 l0_c 中完整的 FP32 乘积累加结果。
97+ # 用 tla.copy 搬到 ub_acc_full，并传入 CopyL0C2DstParams，
98+ # 将 l0c2ub_mode 设为 SPLIT_M，使两个 AIV 各收到 128×256。
99+ # 此时尚未加 X，不能在 FixPipe 中提前启用 ReLU。
100+ tla.copy(
101+ ub_acc_full,
102+ l0_c,
103+ tla.params.CopyL0C2DstParams(
104+ l0c2ub_mode=tla.params.L0C2UBMode.SPLIT_M
105+ ),
106+ )
107+ # --------------------------------
108+ mutex_l0c.unlock(pipe=tla.arch.FIX)
109+ tla.cross_core_set_flag(fix_done, tla.arch.FIX)
110+ 
111+ with tla.vector():
112+ aiv_id = tla.arch.sub_block_idx()
113+ gm_x_aiv = tla.tile_view(
114+ gm_x, tla.make_shape(AIV_M, N), tla.make_coord(aiv_id, 0)
115+ )
116+ gm_y_aiv = tla.tile_view(
117+ gm_y, tla.make_shape(AIV_M, N), tla.make_coord(aiv_id, 0)
118+ )
119+ ub_acc_aiv = tla.make_tensor_like(ub_acc_ptr, gm_y_aiv, tla.arch.RowMajor)
120+ 
121+ tla.cross_core_wait_flag(fix_done, tla.arch.VECTOR)
122+ for chunk in tla.range(0, AIV_M // X_CHUNK_M, 1):
123+ gm_x_chunk = tla.tile_view(
124+ gm_x_aiv, tla.make_shape(X_CHUNK_M, N), tla.make_coord(chunk, 0)
125+ )
126+ gm_y_chunk = tla.tile_view(
127+ gm_y_aiv, tla.make_shape(X_CHUNK_M, N), tla.make_coord(chunk, 0)
128+ )
129+ ub_acc_chunk = tla.tile_view(
130+ ub_acc_aiv, tla.make_shape(X_CHUNK_M, N), tla.make_coord(chunk, 0)
131+ )
132+ ub_x = tla.make_tensor_like(ub_x_ptr, gm_x_chunk, tla.arch.RowMajor)
133+ 
134+ mutex_x.lock(pipe=tla.arch.MTE2)
135+ tla.copy(ub_x, gm_x_chunk)
136+ mutex_x.unlock(pipe=tla.arch.MTE2)
137+ 
138+ mutex_acc.lock(pipe=tla.arch.VECTOR)
139+ mutex_x.lock(pipe=tla.arch.VECTOR)
140+ with tla.vec.func(mode="simd"):
141+ for row in tla.range(0, X_CHUNK_M, 1):
142+ for col in tla.range(0, N // SIMD_LANES, 1):
143+ acc_reg = tla.tile_view(
144+ ub_acc_chunk,
145+ tla.make_shape(1, SIMD_LANES),
146+ tla.make_coord(row, col),
147+ )
148+ x_reg = tla.tile_view(
149+ ub_x,
150+ tla.make_shape(1, SIMD_LANES),
151+ tla.make_coord(row, col),
152+ )
153+ # --------------------------------
154+ # TODO 5（Vector Add + ReLU）：从 UB 的 acc_reg/x_reg
155+ # 载入 FP32 向量，先用 tla.add 相加，再用 tla.max
156+ # 与 0.0 取最大值，最后原位写回 acc_reg。
157+ # 每行固定 256 列，共 4 个完整 SIMD 向量，无需尾掩码。
158+ # 下方 tla.copy 将处理后的 ub_acc_chunk 写到 GM 的 Y。
159+ acc_data = acc_reg.load()
160+ x_data = x_reg.load()
161+ relu_data = tla.max(tla.add(acc_data, x_data), 0.0)
162+ acc_reg.store(relu_data)
163+ # --------------------------------
164+ 
165+ mutex_x.unlock(pipe=tla.arch.VECTOR)
166+ mutex_acc.unlock(pipe=tla.arch.VECTOR)
167+ 
168+ mutex_acc.lock(pipe=tla.arch.MTE3)
169+ tla.copy(gm_y_chunk, ub_acc_chunk)
170+ mutex_acc.unlock(pipe=tla.arch.MTE3)
A2026/CANN-Code-Detective/Challenge09-MatmulAddRelu/PR_code/JeffDing/test_matmul_add_relu.py+63-0
@@ -0,0 +1,63 @@
1+"""补全模板后，在 Ascend 950 上验证固定 256×256×256 的算子。"""
2+ 
3+import argparse
4+ 
5+ 
6+def main() -> int:
7+ parser = argparse.ArgumentParser(description="验证 Catlass Python MatmulAddRelu")
8+ parser.add_argument("--device", type=int, default=0, help="NPU 设备号，默认 0")
9+ args = parser.parse_args()
10+ if args.device < 0:
11+ parser.error("--device 不能为负数")
12+ 
13+ # 延迟导入，使 --help 在未安装 NPU 软件时也能使用。
14+ import catlass.tla as tla
15+ from catlass.tla.runtime import from_dlpack
16+ import torch
17+ import torch_npu # noqa: F401，注册 torch.npu 后端。
18+ 
19+ from matmul_add_relu import matmul_add_relu_kernel
20+ 
21+ torch.npu.set_device(args.device)
22+ torch.manual_seed(2026)
23+ size = 256
24+ a = (torch.rand(size, size) * 2 - 1).half()
25+ b = (torch.rand(size, size) * 2 - 1).half()
26+ x = torch.linspace(-2.0, 2.0, size * size).reshape(size, size)
27+ expected = torch.relu(a.float() @ b.float() + x)
28+ 
29+ device_a, device_b, device_x = a.npu(), b.npu(), x.npu()
30+ device_y = torch.full((size, size), -9999.0, dtype=torch.float32, device="npu")
31+ 
32+ def wrap(tensor):
33+ return from_dlpack(
34+ tensor, layout_tag=tla.arch.RowMajor, origin_shape=(size, size)
35+ )
36+ 
37+ ta, tb, tx, ty = map(wrap, (device_a, device_b, device_x, device_y))
38+ artifact = tla.compile(
39+ matmul_add_relu_kernel, ta, tb, tx, ty, options="--npu-arch 3510"
40+ )
41+ artifact(ta, tb, tx, ty, block_num=1)
42+ # 主机等待整个核函数完成，再读取结果。
43+ torch.npu.synchronize()
44+ 
45+ actual = device_y.cpu()
46+ print(f"Catlass Python 计算结果：\n{actual}")
47+ print(f"PyTorch 参考结果：\n{expected}")
48+ tolerance = torch.maximum(torch.ones_like(expected), expected.abs()) / 256.0
49+ valid = torch.isfinite(actual) & (actual >= 0)
50+ valid &= (actual - expected).abs() <= tolerance
51+ if bool(valid.all()):
52+ print("PASS: M=N=K=256, blocks=1")
53+ return 0
54+ print(f"FAIL: M=N=K=256, mismatches={int((~valid).sum())}")
55+ return 1
56+ 
57+ 
58+if __name__ == "__main__":
59+ try:
60+ raise SystemExit(main())
61+ except Exception as exc:
62+ print(f"FAIL: {exc}")
63+ raise