GetWorkspaceSize：CANN opbase 中 L2 接口 workspace 大小的计算与查询指南

GetWorkspaceSizeCANN opbase 中 L2 接口 workspace 大小的计算与查询指南【免费下载链接】opbase本项目是CANN算子库的基础框架库为算子提供公共依赖文件和基础调度能力。项目地址: https://gitcode.com/cann/opbase导读GetWorkspaceSize是 CANN opbase 算子库中aclOpExecutor提供的关键查询接口用于根据 L2二级接口组装阶段调用的各个 L0 算子计算整个 L2 接口实际运行时所需的 workspace工作内存大小。本文以 docs/zh/api/nnopbase/opdev/op_executor/Getworkspacesize.md 为主体结合 op_executor.h、op_executor.cpp、op_cache.cpp 及测试桩代码完整讲解接口原型、调用时机、底层计算原理、缓存机制并给出可直接套用的aclnnXxxGetWorkspaceSize实现模板。功能说明workspace 在 aclnn L2 接口中的作用在 CANN 的 aclnn 接口体系中一个融合的 L2 接口如aclnnAdd通常由多个 L0 算子组合而成。其中部分 L0 算子如需要中间结果暂存的算子在设备侧执行时需要额外的内存缓冲区这个缓冲区就称为 workspace工作内存。它由用户在 Host 侧通过aclrtMalloc分配并在执行阶段传给aclnnXxx接口。GetWorkspaceSize的核心职责是根据 L2 一阶段组装阶段中已调用的 L0 接口计算整个 L2 接口在运行阶段需要的 workspace 大小。这个大小既包括每个 L0 kernel 的 workspace 需求也包括中间张量中间结果在设备内存上的排布空间。从源码看GetWorkspaceSize是aclOpExecutor的公开成员函数声明位于 include/nnopbase/opdev/op_executor.h#L135uint64_t GetWorkspaceSize() const;函数原型与参数说明函数原型uint64_t GetWorkspaceSize()参数说明参数说明无该接口不接收任何参数结果直接从当前 executor 内部已组装好的信息kernel 启动对象列表、图结构、workspace 偏移等计算得到返回值说明返回uint64_t类型的 workspace 大小单位是字节返回值为 L2 接口在运行阶段需要的 workspace 大小在实际 L2 接口实现中该返回值通过出参uint64_t *workspaceSize暴露给调用者供其分配内存aclrtMalloc后在执行阶段把该内存指针传给运行接口使用。约束说明该接口本身无额外约束。需要注意的是它必须在 L2 一阶段组装阶段完成所有 L0 接口调用之后再调用否则计算结果不会包含后续新增 kernel 的 workspace 需求。这一点在下文调用时机小节中会结合源码详细说明。调用时机L2 一阶段 → GetWorkspaceSize → L2 二阶段GetWorkspaceSize的调用位置位于 L2 接口拆分的两个阶段之间L2 一阶段组装阶段调用CREATE_EXECUTOR()创建 executor通过ADD_TO_LAUNCHER_LIST_AICORE、ADD_TO_LAUNCHER_LIST_DSA等宏逐个添加 L0 kernel launcher并完成 shape 推导INFER_SHAPE查询阶段调用uniqueExecutor-GetWorkspaceSize()获得总 workspace 大小L2 二阶段运行阶段用户根据查询到的大小分配 workspace 内存调用aclnnXxx(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)执行其中aclnnXxx内部通过CommonOpExecutorRun真正下发任务。CREATE_EXECUTOR()宏定义于 include/nnopbase/opdev/make_op_executor.h#L19#define CREATE_EXECUTOR() UniqueExecutor(__func__)它创建一个UniqueExecutor对象内部持有std::unique_ptraclOpExecutor。UniqueExecutor的完整定义在 include/nnopbase/opdev/op_executor.h#L176-L208它在构造时记录当前函数名funcName_并提供get()、operator-、ReleaseTo()等接口。当 L2 一阶段结束时需要调用ReleaseTo(executor)把内部的aclOpExecutor所有权移交出去同时触发AddCache()注册缓存。调用示例aclnn 固定写法以下代码来自 Getworkspacesize.md 中的官方示例展示了GetWorkspaceSize在 aclnn L2 接口中的标准使用方式// aclnn固定写法 void aclnnAddGetWorkspaceSize(..., uint64_t *workspaceSize, aclOpExecutor **executor) { auto uniqueExecutor CREATE_EXECUTOR(); ...... *workspaceSize uniqueExecutor-GetWorkspaceSize(); }将这段模板展开成完整的 L2 一阶段实现并结合ReleaseTo移交 executor可以得到可运行的结构aclnnStatus aclnnAddGetWorkspaceSize(const aclTensor* x1, const aclTensor* x2, aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor) { auto uniqueExecutor CREATE_EXECUTOR(); // 1. 添加 L0 kernel 到 executor含 shape 推导、workspace 偏移记录 ADD_TO_LAUNCHER_LIST_AICORE(Add, ...); // 2. 计算总 workspace 大小 *workspaceSize uniqueExecutor-GetWorkspaceSize(); // 3. 移交 executor 所有权注册缓存 uniqueExecutor.ReleaseTo(executor); return ACLNN_SUCCESS; }在仓库测试桩 tests/nnopbase/common/depends/op/aclnn_mul_stub.cpp#L40-L58 中可以找到完全一致的实现范式aclnnStatus aclnnMulStubGetWorkspaceSize(const aclTensor* intput1, const aclTensor* intput2, aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor) { ... *workspaceSize uniqueExecutor-GetWorkspaceSize(); ... } aclnnStatus aclnnMulStub(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, const aclrtStream stream) { ... return CommonOpExecutorRun(workspace, workspaceSize, executor, stream); }这验证了一阶段查大小、二阶段跑任务的完整闭环aclnnMulStubGetWorkspaceSize负责查询大小aclnnMulStub负责携带workspace与workspaceSize调用CommonOpExecutorRun。底层原理workspace 大小是如何计算的GetWorkspaceSize的实现位于 src/nnopbase/composite_op/aclnn_engine/op_executor.cpp#L610-L627uint64_t aclOpExecutor::GetWorkspaceSize() const { if (impl_-GetWorkspaceSize() ! 0) { return impl_-GetWorkspaceSize(); } auto graph (op::mem::KernelGraph*)(impl_-GetGraph()); auto tensors graph-GetSortedKernelTensors(); OP_LOGD(Workspace tensor count: %zu., tensors.size()); auto allocator op::mem::MaxAllocator(); workspaceDeviceAicpuTaskOffset_ allocator.Allocate(tensors); uint64_t workspaceSize workspaceDeviceAicpuTaskOffset_ workspaceDeviceAicpuMem_; if (impl_-GetOpExecCache() ! nullptr) { impl_-OpExecCacheSetWorkspaceSize(workspaceSize); } OP_LOGD(workspaceSize_:%lu, workspaceDeviceAicpuTaskOffset_:%lu, workspaceSize, workspaceDeviceAicpuTaskOffset_); return workspaceSize; }其计算逻辑分为三步优先复用已缓存结果如果OpExecutorImpl中已经记录过 workspace 大小非 0直接返回避免重复计算内存分配器计算从 executor 内部的KernelGraph取出已排序的 kernel 张量列表GetSortedKernelTensors使用op::mem::MaxAllocator对这些张量在 workspace 内做内存分配得到workspaceDeviceAicpuTaskOffset_汇总返回总大小 张量排布偏移workspaceDeviceAicpuTaskOffset_ AICPU 任务区大小workspaceDeviceAicpuMem_。这两个成员定义于 include/nnopbase/opdev/op_executor.h#L165-L166。与之对应的GetLinearWorkspaceSizeop_executor.cpp#L629-L646使用LinearAllocator做线性排布适用于对内存复用要求更严格的场景。与 UpdateTensorAddr 的配合aclOpExecutor还提供UpdateTensorAddr(void* workspaceAddr, const size_t size)op_executor.cpp#L368-L384。在运行阶段拿到用户分配的 workspace 地址后它会通过SetWorkspaceAddr记录 workspace 基地址遍历allocatedTensorList_将所有IsFromWorkspace()为真的张量按各自的GetWorkspaceOffset()偏移把存储地址设置为workspaceAddr offset同步设置张量数据大小并更新 executor 记录的 workspace 大小。也就是说GetWorkspaceSize负责算多大UpdateTensorAddr负责把 workspace 中的中间张量指到对应位置两者配合完成了 workspace 从查询到使用的完整链路。缓存场景下的 workspace 查询当 executor 命中缓存PTAGetExecCache/PTAFindExecCache时workspace 大小直接从缓存对象获取见 op_executor.cpp#L912-L935 与 op_cache.cpp#L1073-L1075uint64_t OpExecCache::GetWorkspaceSize() const { return workspaceSize_; }workspaceSize_是在首次计算时通过OpExecCacheSetWorkspaceSizeop_executor.cpp#L312-L315写入缓存的。这保证了重复执行同一算子时repeatable 场景无需重新计算 workspace直接复用首次查询得到的大小即可。二阶段运行workspace 的校验与使用运行接口CommonOpExecutorRunop_executor.cpp#L962对 workspace 参数有严格的空指针校验if (unlikely(executor nullptr)) { OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, The executor is nullptr.); return ACLNN_ERR_PARAM_NULLPTR; } if (unlikely(workspaceSize 0 workspace nullptr)) { OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, The workspace is nullptr.); DeleteExecutorForError(executor); return ACLNN_ERR_PARAM_NULLPTR; }即当GetWorkspaceSize返回的大小大于 0 时二阶段必须传入有效的 workspace 内存指针否则返回ACLNN_ERR_PARAM_NULLPTR错误。这是使用者最容易忽略的约束务必保证一阶段查询的大小与二阶段实际分配的内存大小一致。在自定义算子与复杂场景中的应用对于需要多段 workspace 的复杂场景如aclnn_bninference_d_kernel测试桩 aclnn_bninference_d_kernel_stub.cpp一阶段查询出的workspaceSize还可以在二阶段被拆分使用。例如在 aclnn_bninference_d_kernel_stub.cpp#L174-L198 中先比较不同子接口的 workspace 大小取较大值再在二阶段从同一块 workspace 中划分子区段if (viewcopyWsSize *workspaceSize) { *workspaceSize viewcopyWsSize; } *workspaceSize inContWorkspaceSize; ... workspaceSize - inContWorkspaceSize; void* inWorkspace (char*)workspace workspaceSize;这表明GetWorkspaceSize返回的大小是整块可用空间的上界二阶段实现可以在该空间内自行规划多个子缓冲区的布局。总结与最佳实践GetWorkspaceSize必须在 L2 一阶段完成所有 L0 kernel 添加之后调用其结果才完整返回值以字节为单位的uint64_t用于指导调用方分配 workspace 内存若返回大小大于 0二阶段必须传入对应大小的有效 workspace 指针否则运行接口会报ACLNN_ERR_PARAM_NULLPTR底层通过MaxAllocator对图中 kernel 张量做内存排布并叠加 AICPU 任务区大小得到总大小已缓存场景直接复用缓存值避免重复计算实现 L2 接口时请严格遵循CREATE_EXECUTOR()→ 添加 L0 →GetWorkspaceSize()→ReleaseTo(executor)→ 二阶段CommonOpExecutorRun的固定范式。相关参考接口声明include/nnopbase/opdev/op_executor.h宏与启动工具include/nnopbase/opdev/make_op_executor.h核心实现src/nnopbase/composite_op/aclnn_engine/op_executor.cpp缓存实现src/nnopbase/common/utils/op_cache.cpp测试桩示例tests/nnopbase/common/depends/op/aclnn_mul_stub.cpp、tests/nnopbase/common/depends/op/aclnn_bninference_d_kernel_stub.cpp【免费下载链接】opbase本项目是CANN算子库的基础框架库为算子提供公共依赖文件和基础调度能力。项目地址: https://gitcode.com/cann/opbase创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考