[DevLog]24Q4P2 - Lazy Initialize during old dispatcher way
导言
有理有据，步步为营。
背景¶
AI模型在pytorch里执行时，在单算子模式(eager)下执行算子开发时，由于历史开发的原因，一开始选择了复用CANN里的图模式代码，通过对每个算子生成单算子图来执行，被内部称作路径三(对应文件夹op_plugin/ops/aclops # aclop算子适配目录)：
- 优势：快速开发实现功能
- 劣势：性能问题，
- 单算子图不仅无法利用图模式的计算图优化，图并行，算子融合等功能，
- 而且还需要初始化图模式相关（在程序开头的
AclSetCompileOpt）, - 并且每个算子在调用
cmd.Name("Abs")前需要对齐CANN的数据格式（包括数据连续性，数据内存排列 - 非DCHW）。2
张量（tensor）内存常见排列DCHW
- DCHW：这是一种常见的内存排布格式，代表四个维度的排列顺序：
- D: Batch size（批大小）。
- C: Channels（通道数，通常表示图像的颜色通道，比如RGB有3个通道）。
- H: Height（图像的高度）。
- W: Width（图像的宽度）。
为此，后续开发了被内部称作路径五(对应文件夹op_plugin/ops/opapi # aclnn算子适配目录)（路径四并没有代码实现落地）的算子调用方法：
- 抛弃了CANN的图模式代码，重新编写了各个NPU算子，并使用
EXEC_NPU_CMD来调用。 - 大部分算子都实现了从路径三到路径五的迁移，但是还有些算子暂时只能使用路径三。1
目标¶
- Situation：对于只包含支持路径五算子的模型，由于路径三算子引入的公共初始化代码会初始化缓慢，
- Target：计划通过Lazy Initialize的办法，将路径三算子的初始化过程推迟到遇到路径三算子才触发，
- 来避免只包含路径五算子的模型执行路径三算子的初始化，从而实现快速初始化；
- 对于包含路径三算子的模型，等价于延后了初始化处理。
难点与思路¶
Lazy的对象¶
研究NpuSysCtrl::Initialize，探究除开AclSetCompileOpt还有哪些可以挖掘的初始化内容。
// set ACL_PRECISION_MODE by SocVersion("allow_fp32_to_fp16" or "must_keep_origin_dtype").
auto precision_mode = c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend910B1 ?
"must_keep_origin_dtype" : "allow_fp32_to_fp16";
NPU_CHECK_ERROR(at_npu::native::AclSetCompileopt(aclCompileOpt::ACL_PRECISION_MODE, precision_mode));
// set default compile cache mode and dir for users to improve op compile time
MakeCompileCacheDirAndSetOption();
// set default jit_Compile value from Get acl defalut value
GetAndSetDefaultJitCompileByAcl();
SetHF32DefaultValue();
Lazy的触发处¶
有两种基本思路，在上层
- 对于大部分codegen生成的opapi接口，在开头实现单次初始化。
- 在算子执行的公共函数
OpCommandImpl::InnerRun，路径五的判断if (params.customHandler)之后插入算子三的首次判断(只执行一次)，对于第一次的算子三来实现相关初始化。
241008¶
下午
- T4: 和导师对齐了需求
- T5&6: 整理思路。
- T7: 触发的位置很合理，
NpuSysCtrl::Initialize好奇其中的耗时和函数栈（需要GDB）。
晚上
- T1: GDB初始化部分
GDB初始化部分
看上去是没有其余部分的，THNPModule_initExtension的其余对应代码也没有空间了。
#0 c10_npu::NpuSysCtrl::Initialize (this=0xffff08fdb330 <c10_npu::NpuSysCtrl::GetInstance()::instance>, device_id=-1) at /root/document/shaojie/github/pytorch/torch_npu/csrc/core/npu/sys_ctrl/npu_sys_ctrl.cpp:181
#1 0x0000ffff079078a8 in THNPModule_initExtension (self=0xffff09145f40, noargs=0x0) at /root/document/shaojie/github/pytorch/torch_npu/csrc/npu/Module.cpp:267
TO DO
- 可以打印时间，看一下收益。6s的初始化时间。
- 但是
NpuSysCtrl::Initialize也在其他地方被调用，需要仔细分析其组成 - THNPModule_initExtension 代码位置有没有其余的初始化代码被执行？
241010¶
晚上
- T1: 初始化代码，
- AscendCL函数(CANN?)
1. aclInit
2.
aclmdlInitDumpaclrtGetDevice和AclSetCompileopt - 尤其是
aclrtGetDevice在后续 torchnpu里也经常调用。 - IsSupportInfNan，也会利用 CANN的
aclGetCannAttribute - 关键问题：CANN路径五能完全摘出去吗？
- AscendCL函数(CANN?)
1. aclInit
2.
路径五逻辑¶
直接调用底层算子
某些路径五算子Abs可以直接调用底层算子，以此我们可以窥见路径五算子的最简化写法：
at::Tensor abs(const at::Tensor & self)
{
// codegen自动生成 在 StructKernelNpuOpApi.cpp
DO_COMPATIBILITY(aclnnAbs, acl_op::abs(self));
auto output_size_0 = self.sizes();
auto output_dtype_0 = self.scalar_type();
at::Tensor out = npu_preparation::apply_tensor_without_format(output_size_0,
self.options().dtype(output_dtype_0));
EXEC_NPU_CMD(aclnnAbs, self, out);
at::namedinference::propagate_names(out, self);
return out;
}
调用端逻辑：通过函数名在libopapi.so（CANN的lib）里查找
任务队列的实现：
#define EXEC_NPU_CMD(aclnn_api, ...) \
do { \
static const auto task_queue_enable = c10_npu::option::OptionsManager::GetTaskQueueEnable(); \
if (task_queue_enable == 2) { \
EXEC_NPU_CMD_V2(aclnn_api, __VA_ARGS__); \
} else { \
EXEC_NPU_CMD_V1(aclnn_api, __VA_ARGS__); \
} \
} while (false)
核心实现：
at_npu::native::OpCommand cmd; \
cmd.Name(#aclnn_api); \
cmd.SetCustomHandler(acl_call); \
cmd.Run();
// 更深一层到
ACL_REQUIRE_OK_OP(InnerRun(opName, execParam, sync, sync_index, outputTensor), opName.c_str());
// 进一步
aclError OpCommandImpl::InnerRun(
// 其中的hander使用了CANN ？
if (params.customHandler) {
ret = params.customHandler();
// ...
continue;
}
handler 的SetCustomHandler设置逻辑如下：
// 正是之前的 cmd.SetCustomHandler(acl_call); 设置了
api_ret = opApiFunc(workspace_addr, workspace_size, executor, acl_stream);
OpApiFunc opApiFunc = reinterpret_cast<OpApiFunc>(opApiFuncAddr); \
static const auto opApiFuncAddr = GetOpApiFuncAddr(#aclnn_api); \
inline void *GetOpApiFuncAddr(const char *apiName)
static auto opApiHandler = GetOpApiLibHandler(GetOpApiLibName());
// 通过 `dlopen` 是用于加载动态库libopapi.so。
auto handler = dlopen(libName, RTLD_LAZY);
auto funcAddr = GetOpApiFuncAddrInLib(opApiHandler, GetOpApiLibName(), apiName);
// 使用 `dlsym` 是用于从已加载的动态库中获取符号地址（如函数指针或变量地址）的函数。
auto funcAddr = dlsym(handler, apiName);
所以最终使用是CANN_ABI(workspace_addr, workspace_size, executor, acl_stream),后面四个是包装的参数，包括tensor存储空间的传递。
定义端：
路径五算子在定义就是op_plugin, 定义了pytorch如何来对接CANN的接口
路径三逻辑¶
调用端逻辑：通过函数名在libacl_op_compiler.so (CANN的lib)里查找
at::Tensor& abs_out_nocheck(at::Tensor& result, const at::Tensor& self) {
at_npu::native::OpCommand cmd;
cmd.Name("Abs")
.Input(self)
.Output(result)
.Run();
// step 1
void OpCommand::Run() {
aclCmd->Run(sync, sync_index, outputTensor);
}
// step 2
// 和路径五共用前面
aclError OpCommandImpl::InnerRun(
// 分叉
aclError AclopCompileAndExecuteV2(const char *opType,
// 最终
void* FunctionLoader::Get(const std::string& name) {
241011¶
上午
- T1：继续分析CANN路径五能完全摘出去吗？算子三的引入的额外开销是什么？
- T2-5： 发现群空间，有很多资料。整理学习。
- 有很多DEBUG，prof技巧。
- pytorch的基本模块和特性，算子使用的解析，图模式的关系。
下午
- T1&2：开会探讨，Step间重复指令批量下发，来合并一个step里的多个Kernel Launch时间为一个。但是对于动态shape情况不支持。
- T3~6: 继续阅读代码，PTACANN分层解耦方案理解了。
To Do
先把细粒度收尾了
算子是都在op_plugin上定义，
- 为什么libopapi.so找不到，找到了（确实在环境变量里）
- 哪些算子走路径三，有什么特殊的设置。aclinit肯定要有，AclSetCompileopt 等是什么，可以去掉吗？
241015¶
下午
- 现在的问题是路径三有什么特殊的初始化，是路径五用不上的。简单的算子可能都看不出处理。
241025¶
- 发现
torch_npu.npu._lazy_init()
241108¶
- TODO：
- 1212号 迭代四任务：验收指标变成TBE进程触发延迟
- L2 Partition， 测试全部模型
- 预取寻优测试
- 问题：
- 只包含路径五，和路径三的测试样例。
- 监控两者的收益/TBE编译。
- GE初始化，TBE编译触发的位置： 1. 我的理解op_plugin的算子，无论是路径五还是路径三，应该都是PTA编译的时候，一起编译了。 2. 路径三会实时编译TBE算子吗，
TBE相关代码¶
代码并不明显，
// Configure run flag by Session constructor options param,
// its value should be a path
// this option is to obtain the TBE op plugin path
const std::string TBE_PLUGIN_PATH_FLAG = "ge.TBE_plugin_path";
// Configure debug level, its value should be 0(default), 1 or 2.
// 0: close debug; 1: open TBE compiler; 2: open ccec compiler
const std::string OP_DEBUG_LEVEL = "ge.opDebugLevel";
friend class OpRegistrationTbe;
friend class ge::TBEPluginManager;
# `formula; beta1 -1`
#vsub_beta1_1 = tbe.vadds(beta1_broad, sneg_one)
vsub_beta1_1 = torch.add(beta1_broad, sneg_one)
TBE算子 编译与运行
TBE（Tensor Boost Engine）算子是华为昇腾 AI 处理器（如 Ascend 芯片）上的一种算子开发框架，用于加速深度学习模型中的计算过程。TBE 算子是基于自定义的算子开发工具，可以通过定义算子的计算逻辑、性能优化和硬件资源配置来实现高效计算。
在 TBE 中，开发者可以根据特定的 AI 模型需求自定义算子，例如卷积、池化等操作。TBE 算子通常用于自定义高性能的深度学习算子，以优化运行在 Ascend AI 处理器上的计算任务。
在 TBE 框架下，算子的编译运行主要涉及以下几个步骤：
- 编写算子代码：在开发自定义算子时，需要使用 TBE 提供的 DSL（领域特定语言）编写算子的计算逻辑。
- 算子编译：编写完成后，需要将算子代码编译成适用于 Ascend 硬件的可执行文件。编译过程会将高层描述转化为设备上可运行的低层代码，并进行性能优化。
- 运行和调试：编译完成后，可以将编译好的算子加载到 Ascend AI 处理器上运行。通过监控算子的运行情况和性能表现，可以进一步优化算子。
编译和运行的最终目的是使算子能在硬件上高效执行，从而提高整个 AI 模型的计算速度。
TBE相关报错¶
虽然PTA代码里TBE不多，应该是在CANN代码里。
241109¶
- TODO：
- 1212号 迭代四任务：验收指标变成TBE进程触发延迟 1. 保存cpprinter的修改。
- L2 Partition， 脚本测试全部模型
- 科目一练习一道简单和中等。
- 预取寻优测试
- 问题：
- 只包含路径五，和路径三的测试样例。
- 监控两者的收益/TBE编译。
- GE初始化，TBE编译触发的位置： 1. 我的理解op_plugin的算子，无论是路径五还是路径三，应该都是PTA编译的时候，一起编译了。 2. 路径三会实时编译TBE算子吗，