昇腾 CANN ops-blas GEMM 性能调优：分块与向量化的深度实践
·
前言
矩阵乘法（GEMM）是深度学习里计算量最大的算子。ResNet 的 Conv2D、Transformer 的 Attention、LLaMA 的 FFN，底层全是 GEMM。GEMM 优化做好了，整个模型的训练和推理速度能翻倍。这篇文章深入 ops-blas 仓库，从分块策略、Cache 优化、到硬件单元调度，把 GEMM 性能调优讲透。
GEMM 的数学本质
标准定义
GEMM（General Matrix Multiply）的一般形式：
C = α · (A × B) + β · C
其中：
- A: (M, K) 矩阵
- B: (K, N) 矩阵
- C: (M, N) 输出矩阵
- α, β: 标量系数
深度学习中的 GEMM
神经网络里的 GEMM 有几个特点：
# 深度学习中常见的 GEMM 场景
# 场景1：全连接层（Linear）
# Y = X × W^T + B
# X: (batch, in_features) = (B, K)
# W: (out_features, in_features) = (N, K)
# Y: (batch, out_features) = (B, N)
# 即 GEMM: C = A × B, A=(B,K), B=(K,N), C=(B,N)
# 场景2：批量矩阵乘法（Batch GEMM）
# BMM: 多个矩阵对同时做乘法
# Y[i] = X[i] × W[i] for i in batch
# 通常 batch_size = B, seq_len = M, hidden = K, out = N
# 场景3：分组卷积（Grouped Conv）
# 通道被分成 G 组，每组独立做 GEMM
# 适合 MobileNet、ShuffleNet 等轻量化网络
# 验证：手动实现一个 GEMM
import numpy as np
def naive_gemm(A, B, C=None, alpha=1.0, beta=0.0):
"""朴素 GEMM 实现，O(M·N·K)"""
M, K = A.shape
K2, N = B.shape
assert K == K2
C_out = np.zeros((M, N), dtype=np.float32) if C is None else C.copy() * beta
for m in range(M):
for n in range(N):
for k in range(K):
C_out[m, n] += alpha * A[m, k] * B[k, n]
return C_out
# 测试
A = np.random.randn(1024, 512).astype(np.float32)
B = np.random.randn(512, 2048).astype(np.float32)
C = np.random.randn(1024, 2048).astype(np.float32)
result = naive_gemm(A, B, C, alpha=1.0, beta=1.0)
# 这个朴素实现极慢，深度学习框架用分块 + 向量化加速
昇腾 NPU 的计算单元
Cube 和 Vector 双单元
昇腾芯片有两种计算单元，分工不同：
| 单元 | 擅长 | 计算方式 | 适合操作 |
|---|---|---|---|
| Cube | 大矩阵乘法 | 3D Tensor Core（类似 NVIDIA Tensor Core） | GEMM、Conv |
| Vector | 逐元素操作 | 1D/2D 向量指令 | ReLU、Softmax、Add |
# ops-blas 中指定计算单元
import cann
from cann import ops_blas
# 方式1：自动选择（框架自动选择 Cube 或 Vector）
config = ops_blas.GemmConfig()
config.auto_backend = True
# 方式2：强制使用 Cube（适合大矩阵）
config.force_cube = True
config.cube_policy = "max_throughput"
# 方式3：强制使用 Vector（适合小矩阵）
config.force_vector = True
config.vector_policy = "low_latency"
# 创建 GEMM 算子
gemm = ops_blas.create_gemm(config)
result = gemm(A, B)
# 查询实际使用了哪个单元
print(f"实际使用计算单元: {gemm.backend_used}")
# 输出：Cube 或 Vector
Cube 计算单元的微架构
┌─────────────────────────────────────────────┐
│ Cube 计算单元 │
│ │
│ 输入 A: (M, K) 输入 B: (K, N) │
│ ↓ ↓ │
│ ┌─────────────────────────────────────┐ │
│ │ Tensor Core (16x16x16) │ │
│ │ 一次完成 16×16×16 次乘加运算 │ │
│ │ 输入 FP16，累加 FP32，输出 FP32 │ │
│ └─────────────────────────────────────┘ │
│ ↓ │
│ 输出 C: (M, N) │
└─────────────────────────────────────────────┘
# Tensor Core 的数据格式（Cube 特定）：
# A 矩阵：按 (MTile, KTile) 分块存储
# B 矩阵：按 (KTile, NTile) 分块存储
# C 矩阵：按 (MTile, NTile) 分块存储
# 典型的 tile 大小：
# MTile = 64 或 128
# NTile = 64 或 128
# KTile = 16（固定，与 Tensor Core 宽度匹配）
分块策略：三维拆分 M/N/K
为什么要分块
不分块的 GEMM 有两个问题：
- Cache 不友好：矩阵太大，无法全部放入 Cache，数据被反复从显存读取
- 并行度不够：单个 GPU Core 无法高效调度大规模矩阵乘法
分块策略把 M、N、K 三个维度切分成小块（Tile），让每个 Tile 能放进 L2 Cache：
# ops-blas 分块配置
config = ops_blas.GemmConfig()
# M 方向分块（行方向）
config.m_block_size = 64 # 每个 block 处理 64 行
# N 方向分块（列方向）
config.n_block_size = 256 # 每个 block 处理 256 列
# K 方向分块（累加方向）
config.k_block_size = 64 # 每个 K 块处理 64 列
# 实际效果：矩阵被切成 64×256 的小块
# 每个小块可以完全放入 L2 Cache
分块后的计算流程
原始矩阵 A (M=1024, K=512) 和 B (K=512, N=4096)
设定：m_block=64, n_block=256, k_block=64
计算 C = A × B
Block 循环：
for m_block in range(0, M, 64):
for n_block in range(0, N, 256):
C[m_block:m_block+64, n_block:n_block+256] = 0
for k_block in range(0, K, 64):
# 每个 inner loop 处理一个 K 分块
A_tile = A[m_block:m_block+64, k_block:k_block+64]
B_tile = B[k_block:k_block+64, n_block:n_block+256]
C[m_block:m_block+64, n_block:n_block+256] += A_tile × B_tile
# 关键点：K 方向的分块决定了 C 的累加次数
# K=512, k_block=64 -> 需要 8 次累加
实战：找到最优分块参数
# gemm_tune.py
import cann
from cann import ops_blas
import numpy as np
import time
def benchmark_gemm(M, N, K, m_block, n_block, k_block, n_iters=100):
"""测试特定分块参数的性能"""
config = ops_blas.GemmConfig()
config.m_block_size = m_block
config.n_block_size = n_block
config.k_block_size = k_block
gemm = ops_blas.create_gemm(config)
# 分配内存
A = np.random.randn(M, K).astype(np.float16)
B = np.random.randn(K, N).astype(np.float16)
C = np.zeros((M, N), dtype=np.float32)
# Warmup
for _ in range(10):
_ = gemm(A, B)
# 正式测试
times = []
for _ in range(n_iters):
start = time.time()
_ = gemm(A, B)
times.append((time.time() - start) * 1000)
times.sort()
return np.median(times)
# 测试不同的 M/N 分块组合
configs = [
# (m_block, n_block, k_block)
(16, 64, 16),
(32, 128, 16),
(64, 256, 16),
(64, 512, 16),
(128, 256, 16),
(128, 512, 16),
(256, 256, 16),
]
M, N, K = 1024, 4096, 512
print(f"GEMM 基准测试: M={M}, N={N}, K={K}")
print("-" * 60)
print(f"{'m_block':>10} {'n_block':>10} {'k_block':>10} {'延迟(ms)':>10} {'吞吐(GF/s)':>12}")
print("-" * 60)
for m, n, k in configs:
try:
t = benchmark_gemm(M, N, K, m, n, k)
# 计算理论 GFLOPS
flops = 2 * M * N * K / 1e9 # GFLOPS
gf = flops / (t / 1000)
print(f"{m:10d} {n:10d} {k:10d} {t:10.2f} {gf:12.1f}")
except Exception as e:
print(f"{m:10d} {n:10d} {k:10d} {'ERROR':>10} {str(e):>12}")
# 输出示例：
# m_block n_block k_block 延迟(ms) 吞吐(GF/s)
# ------------------------------------------------------------
# 16 64 16 12.34 425.6
# 32 128 16 8.92 589.3
# 64 256 16 5.67 926.8
# 64 512 16 4.23 1241.5 <- 最优
# 128 256 16 4.89 1074.2
# 128 512 16 3.98 1319.8 <- 最优
# 256 256 16 5.12 1026.3
Cache 优化：L1/L2 策略
内存层级与 Tile 大小
┌─────────────────────────────────────────┐
│ HBM（显存）~32GB │
│ 访问延迟：~500-700 cycles │
│ 带宽：~256 GB/s │
└─────────────────────────────────────────┘
↑
L2 Cache（每 Core）~2MB
访问延迟：~50-80 cycles
↑（Cache miss 时才访问 HBM）
┌─────────────────────────────────────────┐
│ L1 Cache（每 Core）~256KB │
│ 访问延迟：~5-10 cycles │
└─────────────────────────────────────────┘
↑
Cube 计算单元（Tensor Core）
16×16×16 矩阵乘法单元
Tile 大小与 Cache 容量匹配
Tile 大小需要精心设计，既要充分利用 Cache 空间，又不能溢出：
# cache_tuning.py
import cann
from cann import ops_blas
# Cache 容量（昇腾 910 典型值）
L1_SIZE = 256 * 1024 # 256KB
L2_SIZE = 2 * 1024 * 1024 # 2MB
# L1 Cache 策略：存放 A 和 C 的 tile
# L1 可以容纳的 tile 数量（考虑 FP16 数据）
bytes_per_element = 2 # FP16
l1_capacity = L1_SIZE / bytes_per_element # 131072 元素
# 每个 tile 需要的元素：(m_tile * k_tile) + (m_tile * n_tile)
# 对于 m_tile=64, n_tile=64, k_tile=64：
# L1 使用 = 64*64 + 64*64 = 8192 元素（仅占 L1 的 6.2%）
# 实际上还需要存 B 的 k_tile*n_tile = 4096 元素
# 总共约 12288 元素 < 131072，安全
# L2 Cache 策略：存放所有 tile 的累加结果
# L2 可以容纳：2MB / 2B = 1M 元素
def compute_optimal_tile_sizes():
"""根据 Cache 容量自动计算最优 tile 大小"""
# 经验公式：L1 tile 约 16-32KB（数据 + 部分结果）
# L2 tile 约 256KB - 512KB
configs = {
# (m_tile, n_tile, k_tile) -> L1 使用估算(KB), L2 使用估算(KB)
(32, 128, 16): (32, 256),
(64, 128, 16): (64, 512),
(64, 256, 16): (128, 1024),
(128, 128, 16): (128, 512),
(128, 256, 16): (256, 1024),
}
print("Tile 大小 vs Cache 使用:")
for tile, (l1_kb, l2_kb) in configs.items():
status_l1 = "OK" if l1_kb < 256 else "溢出"
status_l2 = "OK" if l2_kb < 2048 else "溢出"
print(f" {tile}: L1={l1_kb}KB [{status_l1}], L2={l2_kb}KB [{status_l2}]")
compute_optimal_tile_sizes()
预取策略
预取（Prefetch）让下一个 tile 的数据在当前 tile 计算时提前加载到 Cache：
# prefetch_demo.py
import cann
from cann import ops_blas
config = ops_blas.GemmConfig()
# 开启预取（默认开启）
config.enable_prefetch = True
# 预取距离（多少个 tile 提前加载）
config.prefetch_distance = 2 # 提前加载 2 个 tile
# 预取策略
config.prefetch_policy = "double_buffer" # 双缓冲：计算和加载并行
# 测试预取效果
def test_prefetch():
M, N, K = 2048, 4096, 1024
# 无预取
config_no_prefetch = ops_blas.GemmConfig()
config_no_prefetch.enable_prefetch = False
gemm_no_pf = ops_blas.create_gemm(config_no_prefetch)
# 有预取
config_pf = ops_blas.GemmConfig()
config_pf.enable_prefetch = True
config_pf.prefetch_distance = 2
gemm_pf = ops_blas.create_gemm(config_pf)
# 测试
A = np.random.randn(M, K).astype(np.float16)
B = np.random.randn(K, N).astype(np.float16)
# Warmup
for _ in range(10):
_ = gemm_no_pf(A, B)
_ = gemm_pf(A, B)
import time
t_no_pf = time.time()
for _ in range(50):
_ = gemm_no_pf(A, B)
t_no_pf = (time.time() - t_no_pf) / 50 * 1000
t_pf = time.time()
for _ in range(50):
_ = gemm_pf(A, B)
t_pf = (time.time() - t_pf) / 50 * 1000
print(f"无预取延迟: {t_no_pf:.2f}ms")
print(f"有预取延迟: {t_pf:.2f}ms")
print(f"提升: {(t_no_pf - t_pf) / t_no_pf * 100:.1f}%")
# 输出示例：
# 无预取延迟: 15.67ms
# 有预取延迟: 12.34ms
# 提升: 21.3%
向量化与数据排布
数据排布格式（Layout）
昇腾 NPU 主要使用两种数据排布：
# layout_demo.py
# 格式1：NCHW（RowMajor / 行优先）
# 最直接的内存布局，但 Cube 计算效率低
# 适用于：中间结果传递、Debug
A_nchw = np.ones((1024, 512), dtype=np.float16) # 行优先存储
# 格式2：FRN（Fixed Rotation N）/ NC1HWC0（昇腾专用格式）
# Cube 计算单元的原生格式，高效利用带宽
# NC1HWC0 解释：
# N: batch
# C1: ceil(C/16)，C 向上取整到 16 的倍数
# H: height
# W: width
# C0: 16，C 的最小计算单位
def to_nc1hwc0(tensor, c0=16):
"""把 NCHW 转成昇腾的 NC1HWC0 格式"""
N, C, H, W = tensor.shape
C1 = (C + c0 - 1) // c0
output = np.zeros((N, C1, H, W, c0), dtype=tensor.dtype)
for c in range(C):
c1 = c // c0
c0_idx = c % c0
output[:, c1, :, :, c0_idx] = tensor[:, c, :, :]
return output
# 示例：Conv 输入转换
conv_input = np.random.randn(1, 64, 224, 224).astype(np.float16)
conv_input_nc1hwc0 = to_nc1hwc0(conv_input)
print(f"原始形状: {conv_input.shape}")
print(f"NC1HWC0 形状: {conv_input_nc1hwc0.shape}")
# 输出：
# 原始形状: (1, 64, 224, 224)
# NC1HWC0 形状: (1, 4, 224, 224, 16)
昇腾 GEMM 的原生数据格式
ops-blas 对 GEMM 做了专门的格式适配：
# ops_blas_gemm.py
import cann
from cann import ops_blas
# A 矩阵：(M, K)
# B 矩阵：(K, N)
# C 矩阵：(M, N)
# 昇腾 GEMM 原生接口
class GemmOp:
def __init__(self, trans_a=False, trans_b=False, trans_c=False):
"""
trans_a: 是否转置 A
trans_b: 是否转置 B
trans_c: 是否转置 C
"""
self.trans_a = trans_a
self.trans_b = trans_b
self.trans_c = trans_c
def __call__(self, A, B, C=None, alpha=1.0, beta=0.0):
"""
计算: C = alpha * (A @ B) + beta * C
A, B, C 必须是 npu tensor
"""
pass
# 示例：计算 Y = X × W^T + B（Linear 层）
def linear_layer_npu(x, weight, bias=None):
"""
x: (batch, in_features) = (B, K)
weight: (out_features, in_features) = (N, K)
返回: (batch, out_features) = (B, N)
"""
gemm = GemmOp(trans_b=True) # B 转置 -> W^T
if bias is not None:
# C = X @ W^T + bias
return gemm(x, weight, bias, alpha=1.0, beta=1.0)
else:
# C = X @ W^T
return gemm(x, weight, alpha=1.0, beta=0.0)
# 示例：计算 Attention 的 QK^T
def attention_score_npu(Q, K):
"""
Q: (B, N_heads, T, D) query
K: (B, N_heads, T, D) key
返回: (B, N_heads, T, T) attention score
"""
B, H, T, D = Q.shape
Q_2d = Q.reshape(B * H, T, D) # (B*H, T, D)
K_2d = K.reshape(B * H, T, D) # (B*H, T, D)
gemm = GemmOp(trans_b=True)
scores = gemm(Q_2d, K_2d) # (B*H, T, T)
return scores.reshape(B, H, T, T)
混合精度与量化
FP16/BF16/FP32 混用策略
# mixed_precision.py
import cann
from cann import ops_blas
config = ops_blas.GemmConfig()
# 输入：A/B 用 FP16（计算快）
config.dtype_a = "float16"
config.dtype_b = "float16"
# 累加：中间结果用 FP32（防止溢出）
config.accum_dtype = "float32"
# 输出：C 可以是 FP16 或 FP32
config.dtype_c = "float32" # 推荐用 FP32 输出，再转 FP16
# 计算：Cube Tensor Core 在 FP16 输入时效率最高
# Tensor Core 自动做 FP16*FP16 -> FP32 累加 -> FP16/FP32 输出
# BF16 支持（昇腾 910B+）
config.dtype_a = "bfloat16"
config.dtype_b = "bfloat16"
config.accum_dtype = "float32"
# BF16 vs FP16 对比：
# FP16: 1 sign + 5 exp + 10 mantissa = 16 bits, range: 6e-5 ~ 65504
# BF16: 1 sign + 8 exp + 7 mantissa = 16 bits, range: 9e-5 ~ 3e38
# BF16 动态范围更大，适合训练；FP16 精度更高，适合推理
INT8 量化 GEMM
# int8_gemm.py
import cann
from cann import ops_blas
# 量化配置
quant_config = ops_blas.QuantizedGemmConfig()
# 动态量化（每行/每列独立 scale）
quant_config.quant_mode = "dynamic"
quant_config.activation_dtype = "int8"
quant_config.weight_dtype = "int8"
# 量化因子
quant_config.scale_a = np.ones((1024, 1), dtype=np.float32) / 127.0
quant_config.scale_b = np.ones((1, 4096), dtype=np.float32) / 127.0
# 创建量化 GEMM
qgemm = ops_blas.create_quantized_gemm(quant_config)
# 输入
A_int8 = (np.random.randn(1024, 512) * 80).astype(np.int8)
B_int8 = (np.random.randn(512, 4096) * 80).astype(np.int8)
# 计算
C_int32 = qgemm(A_int8, B_int8)
# 反量化到 FP32
scale = quant_config.scale_a @ quant_config.scale_b
C_fp32 = C_int32.astype(np.float32) * scale * 127.0
print(f"量化 GEMM: A={A_int8.shape} (int8), B={B_int8.shape} (int8)")
print(f"输出: C={C_fp32.shape} (fp32)")
print(f"INT8 带宽节省: {(1 - 2/8) * 100:.0f}%（vs FP32）")
# 输出：
# 量化 GEMM: A=(1024, 512) (int8), B=(512, 4096) (int8)
# 输出: C=(1024, 4096) (fp32)
# INT8 带宽节省: 75%（vs FP32）
性能实测：端到端优化
BERT GEMM 性能分析
# bert_gemm_profile.py
import cann
from cann import ops_blas
import numpy as np
import time
def profile_bert_gemms():
"""Profile BERT 中各类 GEMM 的性能"""
configs = {
# BERT 中的 GEMM 形状
"QKV projection": {"M": 384, "N": 3072, "K": 768},
"Attention scores": {"M": 384, "N": 768, "K": 64}, # per head
"FFN first layer": {"M": 384, "N": 3072, "K": 768},
"FFN second layer": {"M": 384, "N": 768, "K": 3072},
}
print("BERT GEMM 性能 Profile:")
print("-" * 70)
for name, shape in configs.items():
M, N, K = shape["M"], shape["N"], shape["K"]
# 准备数据
A = np.random.randn(M, K).astype(np.float16)
B = np.random.randn(K, N).astype(np.float16)
# 基准配置
config = ops_blas.GemmConfig()
gemm = ops_blas.create_gemm(config)
# Warmup
for _ in range(20):
_ = gemm(A, B)
# 测试 100 次
times = []
for _ in range(100):
start = time.time()
_ = gemm(A, B)
times.append((time.time() - start) * 1000)
median_t = np.median(times)
# 计算吞吐
flops = 2 * M * N * K / 1e9
throughput = flops / (median_t / 1000)
# 计算利用率（假设昇腾 910 峰值 256 GFLOPS FP16）
peak_fp16 = 256
utilization = throughput / peak_fp16 * 100
print(f"{name:25s} | Shape: ({M}, {N}, {K}) | "
f"Latency: {median_t:6.2f}ms | "
f"GFLOPS: {throughput:6.1f} | "
f"Util: {utilization:5.1f}%")
print("-" * 70)
# profile_bert_gemms()
# 输出示例：
# BERT GEMM 性能 Profile:
# ----------------------------------------------------------------------
# QKV projection | Shape: (384, 3072, 768) | Latency: 2.34ms | GFLOPS: 1215.2 | Util: 87.3%
# Attention scores | Shape: (384, 768, 64) | Latency: 0.45ms | GFLOPS: 423.1 | Util: 30.5%
# FFN first layer | Shape: (384, 3072, 768) | Latency: 2.31ms | GFLOPS: 1218.5 | Util: 87.5%
# FFN second layer | Shape: (384, 768, 3072) | Latency: 4.12ms | GFLOPS: 1220.3 | Util: 87.6%
# ----------------------------------------------------------------------
# 分析：Attention scores 的利用率只有 30%，因为矩阵太小，N=768 不适合 Cube 单元
Attention scores 的特殊优化
# attention_small_gemm.py
# 小矩阵 GEMM（Attention scores: M=384, N=768, K=64）
# 直接用 FP16 Cube 不够高效，有几种优化方案：
# 方案1：多 batch 拼接
# 把多个 head 的小矩阵拼成一个大矩阵
def batch_attention_scores(Q_all, K_all, num_heads=12):
"""
Q_all: (B, T, H*D) -> (B*H, T, D)
K_all: (B, T, H*D) -> (B*H, T, D)
"""
B, T, HD = Q_all.shape
D = HD // num_heads
# reshape: (B, H, T, D)
Q = Q_all.view(B, num_heads, T, D)
K = K_all.view(B, num_heads, T, D)
# transpose: (B, H, D, T)
K_t = K.transpose(2, 3)
# 批量 GEMM：Q @ K^T
# Q: (B, H, T, D), K_t: (B, H, D, T)
# 结果: (B, H, T, T)
config = ops_blas.GemmConfig()
config.force_cube = True
config.m_block_size = 384 # T
config.n_block_size = 384 # T
config.k_block_size = 64 # D
gemm = ops_blas.create_gemm(config)
scores = gemm(Q.view(-1, T, D), K_t.view(-1, D, T))
return scores.view(B, num_heads, T, T)
# 方案2：切换到 Vector 单元（极小矩阵）
def vector_attention_scores(Q, K):
"""
当矩阵极小时，用 Vector 单元反而更快
Vector 单元无 tile 开销
"""
config = ops_blas.GemmConfig()
config.force_vector = True # 强制 Vector
gemm = ops_blas.create_gemm(config)
# K 转置 + GEMM
K_t = K.transpose(0, 1) # (D, T)
scores = gemm(Q, K_t) # (B*H*T, T)
return scores.view(-1, T, T)
# 方案3：自动选择（推荐）
def auto_attention_scores(Q, K):
"""
ops-blas 自动选择最优 backend
"""
config = ops_blas.GemmConfig()
config.auto_backend = True # 自动选择 Cube 或 Vector
gemm = ops_blas.create_gemm(config)
scores = gemm(Q, K.transpose(0, 1))
return scores.view(-1, T, T)
分布式 GEMM：多卡并行
数据并行 GEMM
# distributed_gemm.py
import cann
from cann import ops_blas
import torch.distributed as dist
def distributed_gemm_row(A, B, rank, world_size):
"""
行并行：按 M 维度切分 A
每个 rank 持有 A 的 (M/world_size, K) 部分
"""
M, K = A.shape
M_local = M // world_size
# 每个 rank 只拿自己那部分 A
A_local = A[rank * M_local:(rank + 1) * M_local, :]
# 全量 B（每个 rank 都有 B 的副本）
B_full = B
# 本地 GEMM
C_local = ops_blas.gemm(A_local, B_full)
# AllGather：收集所有 rank 的结果
C_full = [None for _ in range(world_size)]
dist.all_gather_object(C_full, C_local)
return np.concatenate(C_full, axis=0)
def distributed_gemm_col(A, B, rank, world_size):
"""
列并行：按 N 维度切分 B
每个 rank 持有 B 的 (K, N/world_size) 部分
需要 AllReduce 汇总
"""
M, K = A.shape
K2, N = B.shape
N_local = N // world_size
# 全量 A（每个 rank 都有 A 的副本）
A_full = A
# 每个 rank 只拿自己那部分 B
B_local = B[:, rank * N_local:(rank + 1) * N_local]
# 本地 GEMM
C_local = ops_blas.gemm(A_full, B_local)
# AllReduce：所有 rank 的 C_local 加起来
# C = A @ B = A @ (B0 + B1 + ... + Bn)
# = A @ B0 + A @ B1 + ... + A @ Bn
C = np.zeros_like(C_local)
dist.all_reduce(C_local)
C = C_local
return C
总结：ops-blas GEMM 调优检查清单
| 优化项 | 默认值 | 建议调优 | 效果 |
|---|---|---|---|
| M/N 分块大小 | 64/256 | 128/512（适配 L2 Cache） | +20-30% |
| K 分块大小 | 16 | 固定（Tensor Core 要求） | 基准 |
| 计算单元 | 自动 | 小矩阵用 Vector，大矩阵用 Cube | +10-50% |
| 数据格式 | NCHW | NC1HWC0（Cube 原生） | +15-25% |
| 预取 | 开启 | 保持开启 | +15-20% |
| 混合精度 | FP32 | FP16 输入 + FP32 累加 | +2x |
| 量化 | 无 | INT8 动态量化 | +2-4x |
调优步骤：
- 先 profiling：找到瓶颈 GEMM 是哪个
- 优化 tile 大小：适配 L2 Cache 容量
- 选择计算单元：小矩阵用 Vector，大矩阵用 Cube
- 开启混合精度：FP16 是昇腾的甜点精度
- 验证正确性：numerical error < 0.1%
GEMM 没有银弹，但 profiling + 调参是最有效的方法。
仓库地址：https://atomgit.com/cann/ops-blas
鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者，聚合华为+生态”的社区，内容涵盖鲲鹏、昇腾资源，帮助开发者快速获取所需的知识、经验、软件、工具、算力，支撑开发者易学、好用、成功，成为核心开发者。
更多推荐
7
0- 0
已为社区贡献24条内容
所有评论(0)