图内设置AI Core和Vector Core核数

功能简介

多流场景下，会出现所有核（Core）都被一个流占用的情况，导致算子执行并行度降低，因此需要把核分给不同的流用，从而保证算子并行执行收益。

目前提供了两种指定核数的方法，算子级核数和全局核数（session）配置，其中算子级优先级高于全局核数的配置。用户需根据实际情况进行核数分配。

使用约束

目前仅适用于如下产品：
- Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件
- Atlas A3 训练系列产品/Atlas A3 推理系列产品
配置算子级核数或全局核数时，不能超过AI处理器本身允许的最大AI Core数量（设为max_aicore）与最大Vector Core数量（设为max_vectorcore）。您可以通过“CANN软件安装目录/<arch>-linux/data/platform_config/<soc_version>.ini”文件查看，如下所示，说明AI处理器上存在32个AI Core，不存在Vector Core。
1 2 3
[SoCInfo] ai_core_cnt=32 vector_core_cnt=0
如果配置的算子级核数或全局核数超过AI处理器本身允许的最大核数，系统默认采用最大核数作为实际运行核数。

使用方法

用户自行分析模型脚本中需要指定核数的算子。
（可选）指定算子级核数。
使用如下with语句块（limit_core_num），语句块内的算子均按照入参指定核数。
1
with torchair.scope.limit_core_num (op_aicore_num: int, op_vectorcore_num: int):
- op_aicore_num：表示该算子运行时的aicore核数，取值范围为[1, max_aicore]。
- op_vectorcore_num：表示该算子运行时的vectorcore核数，取值范围为[1, max_vectorcore]。当AI处理器上仅存在AI Core不存在Vector Core时，此时仅支持取值为0。
配置结果可通过如下操作查看：

参考图结构dump功能获取完整的图结构信息，如config.debug.graph_dump.type="txt"，再查看生效的算子级核数，一般在“attr”属性中，key名分别为“_op_aicore_num”和“_op_vectorcore_num”。

（可选）指定全局核数。

该功能通过torchair.get_npu_backend中compiler_config配置来指定全局核数，示例如下，参数介绍参见表1。

import torch_npu, torchair
config = torchair.CompilerConfig()
# 全局核数配置项
config.ge_config.aicore_num = "24|100"
npu_backend = torchair.get_npu_backend(compiler_config=config)
opt_model = torch.compile(model, backend=npu_backend)

表1 参数说明
参数名	参数说明
aicore_num	指定全局aicore和vectorcore核数，字符串类型，形如“${aicore_num}\|${vectorcore_num}”，必须用“\|”来分隔。 ${aicore_num}：表示全局aicore核数，整数类型，取值范围为[1, max_aicore]。 ${vectorcore_num}：表示全局vectorcore核数，整数类型，取值范围为[1, max_vectorcore]。当AI处理器上仅存在AI Core不存在Vector Core时，本参数可不配置，即config.ge_config.aicore_num = "24\|"，若配置了其它数值也不生效。

配置结果可通过开启Python侧日志获取，假设config.ge_config.aicore_num="24|100"，日志信息如下：

log/debug/plog/plog-67380_20250418094019557.log:15072:[INFO] GE(67380,python3.x):2025-04-18-09:40:20.131.488 [session_manager.cc:60]67380 CreateSession:GE option: ge.aicoreNum, value: 24|100.
log/debug/plog/plog-67380_20250418094019557.log:15083:[INFO] GE(67380,python3.x):2025-04-18-09:40:20.131.526 [platform_info_util.cc:141]67380 parseAicoreNumOption:origin ge.aicoreNum in options, value: 24|100.
......
log/debug/plog/plog-67380_20250418094019557.log:42713:[INFO] GE(67380,python3.x):2025-04-18-09:40:27.333.983 [model_helper.cc:1703]67380 UpdateCoreCountWithOption:ge.aicoreNum in ThreadLocalContext, value: 24.
log/debug/plog/plog-67380_20250418094019557.log:42714:[INFO] GE(67380,python3.x):2025-04-18-09:40:27.333.987 [model_helper.cc:1715]67380 UpdateCoreCountWithOption:Change ge.aicoreNum from platform 20 to rts 20.
log/debug/plog/plog-67380_20250418094014712.log:54036:[INFO] GE(67380,python3.x):2025-04-18-09:40:16.442.320 [model_helper.cc:1703]67380 UpdateCoreCountWithOption:ge.aicoreNum in ThreadLocalContext, value: .
log/debug/plog/plog-67380_20250418094014712.log:54037:[INFO] GE(67380,python3.x):2025-04-18-09:40:16.442.323 [model_helper.cc:1715]67380 UpdateCoreCountWithOption:Change ge.aicoreNum from platform 20 to rts 20.
log/debug/plog/plog-67380_20250418094014712.log:54577:[DEBUG] TEFUSION(67380,python3.x):2025-04-18-09:40:16.663.665 [te_config_info.cc:283]67380 InitConfigItemsFromOptions The value of param[ge.aicoreNum] is [20].
log/debug/plog/plog-67380_20250418094019557.log:15103:[INFO] GE(67380,python3.x):2025-04-18-09:40:20.131.766 [model_helper.cc:1700]67380 UpdateCoreCountWithOption:ge.vectorcoreNum in options, value: 100.
log/debug/plog/plog-67380_20250418094019557.log:15104:[INFO] GE(67380,python3.x):2025-04-18-09:40:20.131.769 [model_helper.cc:1715]67380 UpdateCoreCountWithOption:Change ge.vectorcoreNum from platform 40 to rts 40.
log/debug/plog/plog-67380_20250418094019557.log:42715:[INFO] GE(67380,python3.x):2025-04-18-09:40:27.333.990 [model_helper.cc:1703]67380 UpdateCoreCountWithOption:ge.vectorcoreNum in ThreadLocalContext, value: 100.
log/debug/plog/plog-67380_20250418094019557.log:42716:[INFO] GE(67380,python3.x):2025-04-18-09:40:27.333.993 [model_helper.cc:1715]67380 UpdateCoreCountWithOption:Change ge.vectorcoreNum from platform 40 to rts 40.
log/debug/plog/plog-67380_20250418094014712.log:54038:[INFO] GE(67380,python3.x):2025-04-18-09:40:16.442.326 [model_helper.cc:1703]67380 UpdateCoreCountWithOption:ge.vectorcoreNum in ThreadLocalContext, value: .
log/debug/plog/plog-67380_20250418094014712.log:54039:[INFO] GE(67380,python3.x):2025-04-18-09:40:16.442.328 [model_helper.cc:1715]67380 UpdateCoreCountWithOption:Change ge.vectorcoreNum from platform 40 to rts 40.

根据提示“Change ge.aicoreNum from xx”和“Change ge.vectorcoreNum from xx”发现核数均超过AI处理器允许的最大核数，默认采用最大核数作为实际运行核数，因此本案例中真正生效的核数为aicore_num=20、vectorcore_num=40。

使用示例

import torch, os
import torch_npu
import torchair
from torchair.configs.compiler_config import CompilerConfig
from torchair.core.utils import logger
import logging
logger.setLevel(logging.DEBUG)

# 定义模型model
class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, in1, in2, in3, in4):
        # 指定算子级核数
        with torchair.scope.limit_core_num(4, 5): 
            mm_result = torch.mm(in3, in4)
            add_result = torch.add(in1, in2)
        mm1_result = torch.mm(in3, in4)
        return add_result, mm_result,mm1_result

model = Model()
config = CompilerConfig()
config.debug.graph_dump.type = "pbtxt"
# 指定全局核数
config.ge_config.aicore_num = "24|48"     
npu_backend = torchair.get_npu_backend(compiler_config=config)
model = torch.compile(model, backend=npu_backend, dynamic=False, fullgraph=True)
in1 = torch.randn(1000, 1000, dtype = torch.float16).npu()
in2 = torch.randn(1000, 1000, dtype = torch.float16).npu()
in3 = torch.randn(1000, 1000, dtype = torch.float16).npu()
in4 = torch.randn(1000, 1000, dtype = torch.float16).npu()
result = model(in1, in2, in3, in4)
print(f"Result:\n{result}\n")

父主题： max-autotune模式功能