[DevLog]24Q4P2 - Lazy Initialize during old dispatcher way

导言

有理有据，步步为营。

背景

AI模型在pytorch里执行时，在单算子模式(eager)下执行算子开发时，由于历史开发的原因，一开始选择了复用CANN里的图模式代码，通过对每个算子生成单算子图来执行，被内部称作路径三(对应文件夹op_plugin/ops/aclops # aclop算子适配目录)：

1. 优势：快速开发实现功能
2. 劣势：性能问题，
   1. 单算子图不仅无法利用图模式的计算图优化，图并行，算子融合等功能，
   2. 而且还需要初始化图模式相关（在程序开头的AclSetCompileOpt）,
   3. 并且每个算子在调用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的触发处

有两种基本思路，在上层

1. 对于大部分codegen生成的opapi接口，在开头实现单次初始化。
2. 在算子执行的公共函数OpCommandImpl::InnerRun，路径五的判断if (params.customHandler)之后插入算子三的首次判断(只执行一次)，对于第一次的算子三来实现相关初始化。

241008

下午

1. T4: 和导师对齐了需求
2. T5&6: 整理思路。
3. T7: 触发的位置很合理， NpuSysCtrl::Initialize好奇其中的耗时和函数栈（需要GDB）。

晚上

1. 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

晚上

1. T1: 初始化代码，
   1. AscendCL函数(CANN?)
      1. aclInit
      2. aclmdlInitDump aclrtGetDevice 和AclSetCompileopt
   2. 尤其是 aclrtGetDevice 在后续 torchnpu里也经常调用。
   3. IsSupportInfNan，也会利用 CANN的 aclGetCannAttribute
   4. 关键问题：CANN路径五能完全摘出去吗？

路径五逻辑

直接调用底层算子

某些路径五算子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

上午

1. T1：继续分析CANN路径五能完全摘出去吗？算子三的引入的额外开销是什么？
2. T2-5： 发现群空间，有很多资料。整理学习。
   1. 有很多DEBUG，prof技巧。
   2. pytorch的基本模块和特性，算子使用的解析，图模式的关系。

下午

1. T1&2：开会探讨，Step间重复指令批量下发，来合并一个step里的多个Kernel Launch时间为一个。但是对于动态shape情况不支持。
2. T3~6: 继续阅读代码，PTACANN分层解耦方案理解了。

To Do

先把细粒度收尾了

算子是都在op_plugin上定义，

1. 为什么libopapi.so找不到，找到了（确实在环境变量里）
2. 哪些算子走路径三，有什么特殊的设置。aclinit肯定要有，AclSetCompileopt 等是什么，可以去掉吗？

241015

下午

1. 现在的问题是路径三有什么特殊的初始化，是路径五用不上的。简单的算子可能都看不出处理。

241025

1. 发现torch_npu.npu._lazy_init()

241108

1. TODO：
   1. 1212号 迭代四任务：验收指标变成TBE进程触发延迟
   2. L2 Partition， 测试全部模型
   3. 预取寻优测试
2. 问题：
   1. 只包含路径五，和路径三的测试样例。
   2. 监控两者的收益/TBE编译。
   3. 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

1. TODO：
   1. 1212号 迭代四任务：验收指标变成TBE进程触发延迟
      1. 保存cpprinter的修改。
   2. L2 Partition， 脚本测试全部模型
   3. 科目一练习一道简单和中等。
   4. 预取寻优测试
2. 问题：
   1. 只包含路径五，和路径三的测试样例。
   2. 监控两者的收益/TBE编译。
   3. GE初始化，TBE编译触发的位置：
      1. 我的理解op_plugin的算子，无论是路径五还是路径三，应该都是PTA编译的时候，一起编译了。
      2. 路径三会实时编译TBE算子吗，

参考文献

[^1]: Ascend/op-plugin API适配开发流程

[^2]: CANN TBE&AI CPU 自定义算子开发指南