CANN 亲和调度：AI Core 与计算单元的高效利用

一、昇腾 AI Core 架构

1.1 AI Core 概述

昇腾 NPU 的核心是 AI Core，它是执行神经网络计算的基本单元。不同于传统 GPU 的通用计算架构，昇腾 AI Core 针对深度学习计算进行了专门优化，包含矩阵运算单元（Cube）、向量运算单元（Vector）、标量运算单元（Scalar）和存储队列等组件。

┌──────────────────────────────────────────────────┐ │ 昇腾 AI Core 架构 │ ├──────────────────────────────────────────────────┤ │ │ │ ┌────────────────────────────────────────────┐ │ │ │ AI Core │ │ │ │ ┌─────────┐ ┌─────────┐ ┌─────────┐ │ │ │ │ │ Cube │ │ Vector │ │ Scalar │ │ │ │ │ │ 矩阵运算 │ │向量运算 │ │标量运算 │ │ │ │ │ └────┬────┘ └────┬────┘ └────┬────┘ │ │ │ │ └──────────┬┴──────────┬┘ │ │ │ │ ┌─────┴─────┐ │ │ │ │ │ │ Unified │ │ │ │ │ │ │ Buffer │ │ │ │ │ │ └─────┬─────┘ │ │ │ │ └──────────────────┼───────────┼─────────────┘ │ │ ↓ ↓ │ │ ┌──────────────────┴───────────┴─────────────┐ │ │ │ L1 / L2 Cache │ │ │ └─────────────────────────────────────────────┘ │ │ │ └──────────────────────────────────────────────────┘

1.2 Cube、Vector、Scalar 三单元对比

工作原理：

Cube（矩阵运算单元）: - 输入: [M, K] × [K, N] - 输出: [M, N] - 适合: 全连接层、卷积层 Vector（向量运算单元）: - 输入: 逐元素操作 [N] - 输出: 逐元素操作 [N] - 适合: 激活函数、归一化 Scalar（标量运算单元）: - 输入: 单值操作 - 输出: 单值操作 - 适合: 形状操作、参数更新

1.3 存储层次

昇腾 AI Core 的存储层次直接影响数据访问效率：

数据流动路径：

HBM → L2 Cache → L1 Cache → Unified Buffer → AI Core

优化目标是最小化 HBM 访问，尽量让数据停留在高带宽存储层。

二、Core 亲和配置

2.1 Core Type 概念

昇腾 910 包含不同类型的 AI Core，不同算子适合在不同 Core 上运行：

8.1 及之前（无控制）：

# 使用默认配置，系统自动选择output=torch.matmul(input,weight)

8.2 新增（显式 Core 选择）：

importascend_npu.opsasops# 强制使用 Cube 计算output=ops.matmul(input,weight,core_type="Cube"# 指定矩阵运算用 Cube)# 强制使用 Vector 计算output=ops.relu(input,core_type="Vector")# 自动选择（等价于不指定）output=ops.matmul(input,weight,core_type="Auto")

2.2 Core 亲和场景

场景 1: 大矩阵运算 → Cube

# Transformer 中的 QKV 投影classAttention(nn.Module):defforward(self,x):# Q, K, V 投影都是大矩阵运算q=ops.matmul(x,self.w_q,core_type="Cube")k=ops.matmul(x,self.w_k,core_type="Cube")v=ops.matmul(x,self.w_v,core_type="Cube")# 注意力计算中的 softmax 是向量运算scores=ops.softmax(scores,core_type="Vector")returnoutput

场景 2: 激活函数 → Vector

# 激活层适合 Vector 计算classActivationLayer(nn.Module):defforward(self,x):# ReLU 是逐元素操作，Vector 更高效x=ops.relu(x,core_type="Vector")# Sigmoid 同理x=ops.sigmoid(x,core_type="Vector")# LayerNorm 是向量运算x=ops.layer_norm(x,core_type="Vector")returnx

2.3 Core 亲和配置 API

fromascend_npu.coreimportCoreAffinity# 创建亲和性配置器affinity=CoreAffinity()# 设置全局默认affinity.set_default_core("Auto")# 为特定算子设置亲和affinity.set_op_core("MatMul","Cube")affinity.set_op_core("Conv2d","Cube")affinity.set_op_core("ReLU","Vector")affinity.set_op_core("LayerNorm","Vector")# 应用配置affinity.apply()# 验证配置print(affinity.get_config())

2.4 动态 Core 调度

8.2 支持根据输入 shape 动态选择 Core：

defdynamic_core_selection(input_shape,op_type):"""根据 shape 动态选择最优 Core"""# 小矩阵用 Vector 更高效ifinput_shape[-1]<256:return"Vector"# 大矩阵用 Cubeelse:return"Cube"# 使用动态选择classDynamicCoreModel(nn.Module):defforward(self,x):core=dynamic_core_selection(x.shape,"matmul")returnops.matmul(x,self.weight,core_type=core)

三、内存亲和优化

3.1 内存层级选择

昇腾的内存层级会影响算子性能，合理使用可以显著加速：

3.2 L1/L2 缓存优化

# 启用 L1 缓存优化（8.2 新增）ops.set_cache_mode("L1",enabled=True)# 为特定算子启用 L1 缓存classConvModel(nn.Module):defforward(self,x):# 第一个卷积层启用 L1 缓存（小 kernel 重复访问）x=ops.conv2d(x,self.conv1.weight,cache_policy="L1",# 启用 L1 缓存core_type="Cube")# 中间层用默认策略x=ops.conv2d(x,self.conv2.weight,core_type="Cube")returnx

3.3 地址模式选择

算子计算时数据存放位置影响访问效率：

# 设置地址模式ops.matmul(a,b,addr_mode="L1",# 数据放 L1，更快但有限core_type="Cube")# 自动选择（默认）ops.matmul(a,b,addr_mode="Auto")

四、Stream 调度优化

4.1 Stream 概念

Stream 是昇腾异步执行的基础，类似于 CUDA Stream。不同 Stream 可以并行执行独立的计算任务：

┌──────────────────────────────────────────────────┐ │ Stream 并行 │ ├──────────────────────────────────────────────────┤ │ │ │ Stream 0: ──────────────────────────→ │ │ Stream 1: ──────────────────────────→ │ │ Stream 2: ──────────────────────────→ │ │ │ │ 同一时刻可并行执行多个 kernel │ └──────────────────────────────────────────────────┘

4.2 Stream 创建与使用

importtorch# 创建多个 Streamstream0=torch.npu.Stream(priority=0)# 高优先级stream1=torch.npu.Stream(priority=1)# 低优先级# 在指定 Stream 上执行withtorch.npu.stream(stream0):result0=model0(data)# 模型0 在 stream0 执行withtorch.npu.stream(stream1):result1=model1(data)# 模型1 在 stream1 执行

4.3 计算与通信Overlap

分布式训练中，计算和通信可以通过 Stream 重叠来隐藏延迟：

8.1 及之前（串行）：

# 计算和通信串行，无法重叠forbatchindataloader:output=model(data)loss.backward()# 计算dist.all_reduce(grad)# 等待通信完成optimizer.step()

8.2 新增（异步 Overlap）：

# 创建独立 Stream 用于通信compute_stream=torch.npu.Stream()comm_stream=torch.npu.Stream()forbatchindataloader:# 在计算 Stream 中执行前向和反向withtorch.npu.stream(compute_stream):output=model(data)loss.backward()# 计算完成# 异步触发通信withtorch.npu.stream(comm_stream):handle=dist.all_reduce(grad,async_op=True)# 等待通信完成handle.wait()# 计算 Stream 和通信 Stream 并行执行optimizer.step()

4.4 同步点优化

过多的同步点会打断流水线，降低硬件利用率：

# 8.1 及之前（频繁同步）forbatchindataloader:loss1=model1(data)# 无依赖但同步torch.npu.synchronize()# 等待loss2=model2(data)# 等前一个完成torch.npu.synchronize()# 8.2 新增（延迟同步）forbatchindataloader:loss1=model1(data)loss2=model2(data)# 最后一个同步点torch.npu.synchronize()# 只同步一次# 两个计算并行提交，最后统一等待

五、亲和性自动调优

5.1 Auto-Tuning 机制

8.2 引入了基于 Profiling 的自动调优功能，可以自动找到最优的 Core 和内存配置：

# 启用 Auto-TuningexportASCEND_AUTOTUNE=1exportASCEND_AUTOTUNE_CONFIG=/workspace/autotune_config.json# 运行训练python train.py# 自动生成调优结果cat/workspace/autotune_config.json

Auto-Tuning 配置：

{"autotune":{"enabled":true,"profile_iterations":100,"warmup_iterations":10,"tune_ops":["MatMul","Conv2d","LayerNorm"],"core_types":["Auto","Cube","Vector"],"cache_modes":["Auto","L1","L2"]}}

5.2 Auto-Tuning 结果解读

importjson# 读取调优结果withopen("/workspace/autotune_config.json")asf:results=json.load(f)# 分析结果forop,configinresults["best_configs"].items():print(f""" 算子:{op}Core Type:{config['core_type']}Cache Mode:{config['cache_mode']}Speedup:{config['speedup']:.2f}x vs baseline """)

示例输出：

算子: MatMul Core Type: Cube Cache Mode: L1 Speedup: 1.35x vs baseline 算子: Conv2d Core Type: Cube Cache Mode: L2 Speedup: 1.28x vs baseline 算子: LayerNorm Core Type: Vector Cache Mode: Auto Speedup: 1.12x vs baseline

5.3 手动调优流程

在没有 Auto-Tuning 时，可以手动调优：

# 手动调优脚本defmanual_tune(model,input_shape,config_grid):best_config=Nonebest_latency=float('inf')# 网格搜索最优配置forcore_typeinconfig_grid['core_types']:forcache_modeinconfig_grid['cache_modes']:foraddr_modeinconfig_grid['addr_modes']:# 配置ops.set_core_type(core_type)ops.set_cache_mode(cache_mode)ops.set_addr_mode(addr_mode)# 测量延迟latencies=[]for_inrange(100):start=time.time()_=model(input_data)latencies.append(time.time()-start)avg_latency=np.mean(latencies)ifavg_latency<best_latency:best_latency=avg_latency best_config={'core_type':core_type,'cache_mode':cache_mode,'addr_mode':addr_mode}returnbest_config# 调优配置网格config_grid={'core_types':['Auto','Cube','Vector'],'cache_modes':['Auto','L1','L2'],'addr_modes':['Auto','HBM','L1']}best=manual_tune(model,(1,3,224,224),config_grid)print(f"最优配置:{best}")

六、性能数据参考

6.1 不同 Core 的性能对比

6.2 缓存命中率对性能的影响

6.3 Stream 并行效率

“Core 亲和调优的核心收益：合理选择 Cube/Vector 可获得 2-8 倍算子加速，Stream 并行可提升 2-3 倍吞吐。”

七、常见问题

相关仓库

• ascend-npu- Core 亲和配置接口 https://gitee.com/ascend/ascend-npu
• torch_npu- Stream 管理接口 https://gitee.com/ascend/torch_npu
• ascend-toolkit- Auto-Tuning 工具 https://gitee.com/ascend/ascend-toolkit