CANN opbase 算子开发：OP_WORKSPACE 宏详解——显式封装算子 Workspace 参数

CANN opbase 算子开发OP_WORKSPACE 宏详解——显式封装算子 Workspace 参数【免费下载链接】opbase本项目是CANN算子库的基础框架库为算子提供公共依赖文件和基础调度能力。项目地址: https://gitcode.com/cann/opbase导读OP_WORKSPACE是 CANN opbase算子库基础框架为算子开发者提供的参数封装宏用于在算子执行任务中显式声明计算过程所需的 Device 内存workspace参数。本文以 OP_WORKSPACE 官方文档 为主线结合 op_arg_def.h 等仓库源码与测试用例系统讲解该宏的语法、底层实现、与运行时框架的联动机制以及典型调用场景帮助读者在自定义算子开发中正确使用 workspace 参数声明。宏功能显式封装算子 Workspace 参数OP_WORKSPACE的核心作用是将算子计算过程中显式指定的 Device 内存参数封装为框架可识别的参数对象。在 AI Core 算子执行流程中算子除了输入Input、输出Output、属性Attr之外往往还需要一块额外的临时内存用于存放中间计算结果这块内存即 workspace工作空间。通过OP_WORKSPACE宏开发者可以把这块显式指定的内存以标准参数的形式挂接到算子执行上下文中供后续的 kernel 启动、建图BuildGraph、非有限值检查Non-Finite Check等环节统一访问。从算子开发框架的职责分工看该宏属于opdev 公共宏与类common macros and classes中的参数封装一族与 OP_INPUT、OP_OUTPUT、OP_ATTR、OP_OPTION 等并列共同构成ADD_TO_LAUNCHER_LIST_AICORE宏可识别的算子参数集合详见 ADD_TO_LAUNCHER_LIST_AICORE 文档 的关联接口列表。宏原型OP_WORKSPACE(x...)该宏为变参宏x...表示可接收一个或多个待封装的 workspace 参数单个参数时直接传入变量多个参数时以逗号分隔。宏展开后构造出一个op::OpWorkspace类型的参数对象其定义位于 op_arg_def.h#define OP_WORKSPACE(x...) op::OpWorkspace(std::make_tuple(x))参数说明参数输入/输出说明x...输入算子的 workspace 参数表示算子计算过程中需要的 Device 内存。支持 aclTensor 或 aclTensorList 类型。参数需要澄清的两个要点类型范围文档明确规定支持aclTensor或aclTensorList两种类型分别对应单个 workspace 张量与 workspace 张量列表多段连续内存场景。输入属性workspace 由算子声明、由框架或调用侧提供在参数分类中属于输入侧声明但不参与算子本身的语义计算仅提供临时存储。从源码实现看OpWorkspace参数类型的 tag 值为OP_WORKSPACE_ARG 3定义于 op_arg_def.h并通过DEFINE_OP_ARG(OpWorkspace, OP_WORKSPACE_ARG)op_arg_def.h自动生成参数类型模板。该模板继承自OpArgBaseT...内部以std::tupleT...保存实际参数并支持通过VisitBy、VisitByNoReturn、VisitAt等接口遍历元组元素见 op_arg_def.h。约束说明官方文档明确该宏无额外约束。结合源码与测试可以进一步说明两点使用前提当算子不需要显式指定 workspace 时可以不传该宏框架会在内部自动向执行上下文追加 workspace见下文“运行时联动”。若需要显式指定则应在ADD_TO_LAUNCHER_LIST_AICORE宏的参数列表中按框架要求的位置传入该宏被调用前若算子需要INFER_SHAPE应遵循 INFER_SHAPE 的使用约束先做 shape 推导再创建执行任务。调用示例// 封装算子的workspace参数a OP_WORKSPACE(a);多参数形式与测试中的实际用法// 多参数形式同时封装多个workspace张量来自测试用例 tests/nnopbase/ut/composite_op/test_kernel_launch.cpp auto workspace OP_WORKSPACE(tensorPtr4, workspaceTensorList); // 结合输入输出一起构造完整的OpArgContext来自测试用例 tests/nnopbase/st/composite_op/test_op_cache.cpp auto ctx op::MakeOpArgContext(OP_WORKSPACE(workspace), OP_INPUT(self.get()), OP_OUTPUT(out.get(), idx.get()));底层实现与参数体系参数类型体系中的位置在op::OpArgDef枚举中算子参数被划分为 9 类op_arg_def.h枚举值含义OP_INPUT_ARG 0输入参数OP_OUTPUT_ARG 1输出参数OP_ATTR_ARG 2属性参数OP_WORKSPACE_ARG 3workspace 参数OP_OUTSHAPE_ARG 4输出 shape 参数OP_OPTION_ARG 5option 参数OP_EXEC_MODE_ARG 6执行模式参数OP_MEMSET_ARG 7memset 参数OP_EMPTY_ARG 8空参数占位OP_WORKSPACE_ARG与输入、输出、属性等并列说明 workspace 是一类一等公民参数会在OpArgContext中被单独建槽存储。OpArgContext内部维护一个std::arrayOpArgList, OP_ARG_TYPE_NUM argLists数组op_arg_def.h按参数类型索引每个槽位是一个OpArgList参数数组 数量 类型标识。宏展开后的类型擦除与遍历OP_WORKSPACE(x...)展开后得到op::OpWorkspace(std::make_tuple(x))。OpWorkspace作为OpArgBase的派生类用std::tuple保存任意数量的参数通过OpArgContextSize递归统计所有 OpArgBase 派生对象的元素总数op_arg_def.h在OpArgContextInit中按参数类型值写入对应槽位并逐个元素调用AppendOpArg完成类型擦除——每个元素被转换为OpArg{type, value}结构其中type来自OpArgType枚举OPARG_ACLTENSOR、OPARG_ACLTENSOR_LIST等value是OpArgValue联合体op_arg_def.h。AppendOpArg针对不同 C 类型有多个重载op_arg_def.h其中与 workspace 相关的关键重载为inline void AppendOpArg([[maybe_unused]] size_t idx, aclTensor* tensor, OpArg* currArg) { currArg-type OpArgType::OPARG_ACLTENSOR; currArg-value OpArgValue(tensor); currArg; } inline void AppendOpArg([[maybe_unused]] size_t idx, aclTensorList* tensorList, OpArg* currArg) { currArg-type OpArgType::OPARG_ACLTENSOR_LIST; currArg-value OpArgValue(tensorList); currArg; }这两条路径与文档“支持 aclTensor 或 aclTensorList 类型”的说明一一对应单个张量被擦除为OPARG_ACLTENSOR张量列表被擦除为OPARG_ACLTENSOR_LIST且均以指针形式存入OpArgValue联合体。OpArgContext 的预留机制MakeOpArgContext在分配OpArgContext内存时除了按参数总数分配OpArg数组还会额外预留一个 OpArg 槽位op_arg_def.h// add 1 for append workspace size_t size sizeof(OpArgContext) (OpArgContextSize(ts...) 1) * sizeof(OpArg); ... // reserve first OpArg for append workspace OpArg* currArg reinterpret_castOpArg*(ctx 1) 1;该预留槽位的用途在OpArgContext::AppendOpWorkspaceArgop_arg_def.cpp中揭示当算子没有显式声明 workspace 时框架可在此槽位动态追加一个由内部生成的 workspace 张量列表void OpArgContext::AppendOpWorkspaceArg(aclTensorList* tensorList) { if (ContainsOpArgType(OP_WORKSPACE_ARG)) { return; } OpArg* arg reinterpret_castOpArg*(this 1); arg-type OpArgType::OPARG_ACLTENSOR_LIST; arg-value OpArgValue(tensorList); argLists[OP_WORKSPACE_ARG].args arg; argLists[OP_WORKSPACE_ARG].count 1; }可以看到若算子已通过OP_WORKSPACE显式声明了 workspaceAppendOpWorkspaceArg会直接返回以显式声明为准只有未声明时才会用内部生成的 workspace 填充。这正是该宏在运行时框架中“显式指定”语义的落点。运行时框架中的联动机制执行器创建阶段的两种路径在 op_executor.cpp 的CreatAiCoreKernelLauncher中能否检测到OP_WORKSPACE_ARG直接决定了执行路径if (args-ContainsOpArgType(op::OP_WORKSPACE_ARG)) { // 路径一算子显式声明了 workspace直接用于建图 auto* launcher new op::AiCoreKernelLauncher{opType, op::AI_CORE, profilingInfoId, executor, args}; executor-AddToKernelLauncherList(launcher); op::internal::BuildGraph(executor-GetGraph(), opType, *args-GetOpArg(op::OP_INPUT_ARG), *args-GetOpArg(op::OP_OUTPUT_ARG), *args-GetOpArg(op::OP_WORKSPACE_ARG)); } else { // 路径二未显式声明由框架通过 GetWorkspace 获取并追加 aclTensorList* workspace nullptr; op::internal::GetLauncherCtx().ClearTilingCache(); addToLaunchRet op::internal::GetWorkspace(opType, workspace, executor, *args-GetOpArg(op::OP_INPUT_ARG), *args-GetOpArg(op::OP_OUTPUT_ARG), *args-GetOpArg(op::OP_ATTR_ARG)); args-AppendOpWorkspaceArg(workspace); ... }路径二中GetWorkspace定义于 kernel_workspace.cpp通过 KernelManager 查询算子 tiling 计算出的 workspace 大小并分配 Device 内存再以aclTensorList形式通过AppendOpWorkspaceArg补入上下文。同理DSADevice Side Accelerator任务创建时也通过GetDSATaskWorkspace获取 workspace 并追加op_executor.cpp。这一机制保证了无论算子是否显式声明OP_WORKSPACE_ARG槽位最终都能拿到有效的 workspacekernel 启动逻辑可以统一按“存在 workspace”处理。Kernel 启动参数的组织在 launch_arg_info.h 中AiCoreKernelLauncher在组织启动参数时会检测OP_WORKSPACE_ARG并把其中的每个元素通过AppendWorkspaceLaunchArg追加到设备侧参数列表kernel_launcher.h 同样会对 workspace 参数做访问处理。也就是说OP_WORKSPACE声明的张量会与输入、输出张量一起被编排进最终下发到 Device 侧 kernel 的参数布局。其他框架环节的消费点建图BuildGraph接收OP_WORKSPACE_ARG槽位将 workspace 纳入计算图见 op_executor.cpp 与 dsa_task.h。DSA 任务DSAKernelLauncher通过VisitWorkspaceArg(*opArgCtx-GetOpArg(op::OP_WORKSPACE_ARG))消费 workspacedsa_task.h。算子信息序列化/DFXop_dfx.cpp在 dump 算子信息时若存在OP_WORKSPACE_ARG会把 workspace 张量与索引记录下来op_dfx.cpp。非有限值检查non_finite_check_op.h中非有限值检查算子同样读取 workspace 槽位non_finite_check_op.h。kernel 二进制查询op_kernel_bin.cpp通过ContainsOpArgType(op::OP_WORKSPACE_ARG)判断并读取 workspace 参数列表op_kernel_bin.cpp。AICPU 任务aicpu_task.cpp在 AICPU 算子任务构建时同样取用OP_WORKSPACE_ARG槽位aicpu_task.cpp。测试用例验证仓库的单元测试与系统测试覆盖了OP_WORKSPACE的多种使用形态可作为实践参考单张量 workspaceauto ws OP_WORKSPACE(out.get());test_kernel_launch.cpp。张量列表 workspaceauto wsArg OP_WORKSPACE(wsList);test_kernel_launch.cpp。多参数 workspaceauto workspace OP_WORKSPACE(tensorPtr4, workspaceTensorList);test_kernel_launch.cpp。与输入输出组合构造上下文op::MakeOpArgContext(OP_WORKSPACE(workspace), OP_INPUT(self.get()), OP_OUTPUT(out.get(), idx.get()))test_op_cache.cpp。参数类型断言test_op_arg_def.cpp通过ContainsOpArgType(op::OpArgDef::OP_WORKSPACE_ARG)验证未声明时槽位为空、声明后为真并进一步取出OpArgList检查test_op_arg_def.cpp。上述用例印证OP_WORKSPACE既可以单独声明单个/多个 workspace 张量也可以与OP_INPUT、OP_OUTPUT等宏并列传入MakeOpArgContext其最终效果是让OpArgContext的OP_WORKSPACE_ARG槽位携带有效的aclTensor或aclTensorList类型参数。小结OP_WORKSPACE是 CANN opbase 算子开发中声明 workspace 参数的标准入口其背后是一套完整的参数封装与类型擦除机制声明层OP_WORKSPACE(x...)将 aclTensor / aclTensorList 封装为OpWorkspace参数对象存储层OpArgContext按OP_WORKSPACE_ARG类型槽位统一存储支持显式声明与运行时追加两种填充方式消费层kernel 启动参数组织、建图、DFX dump、非有限值检查、AICPU/AICore 任务构建等多个环节均可读取该槽位。对于算子开发者而言若算子计算需要一块显式提供的 Device 临时内存应使用OP_WORKSPACE在ADD_TO_LAUNCHER_LIST_AICORE的参数列表中声明若无需显式指定可省略该宏由框架内部自动申请并追加 workspace两种方式均能保证后续执行流程获得一致的参数视图。相关资源OP_WORKSPACE 英文文档 | OP_WORKSPACE 中文文档 | 参数定义头文件 | 参数上下文实现 | 执行器实现【免费下载链接】opbase本项目是CANN算子库的基础框架库为算子提供公共依赖文件和基础调度能力。项目地址: https://gitcode.com/cann/opbase创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考