ATVOSS kernel层调度执行引擎深度解析BaseKernelSchedule::Run 的调度策略落地机制【免费下载链接】atvossATVOSSAscend C Templates for Vector Operator Subroutines是一套基于Ascend C开发的Vector算子库致力于为昇腾硬件上的Vector类融合算子提供极简、高效、高性能、高拓展的编程方式。项目地址: https://gitcode.com/cann/atvossBaseKernelSchedule 是 CANN atvoss 项目中 kernel 层调度的基类而Run则是其核心执行入口它负责把 tiling 配置与用户输入参数转化为每个 AI Core 实际要计算的数据分片并驱动 block 层完成计算。本文基于仓库源码include/elewise/kernel/schedule.h逐行拆解BaseKernelSchedule::Run的执行流水线并结合MakeScheduleConfig、DefaultKernelConfig与DeviceAdapter完整还原host 侧算 tiling → device 侧执行调度的端到端链路帮助读者掌握如何在 ATVOSS 中继承该基类、理解默认调度策略并编写自定义调度。一、定位kernel 层调度基类在 ATVOSS 分层架构中的角色ATVOSSAscend C Templates for Vector Operator Subroutines采用Device 层 → Kernel 层 → Block 层 → Compute 层的分层结构。其中 kernel 层负责回答两个问题切多少块根据输入 shape 与硬件核数计算需要启用多少个 AI CoreblockNum、每个核处理多少基本块unitNumPerCore每块算哪里在 device 侧根据当前核的 block ID从全局内存 GM 中定位当前核需要处理的数据区间再交给 block 层执行。BaseKernelSchedule就是 kernel 层调度策略的公共基类。按照官方接口文档的定义见 docs/api/README.md 的 Kernel 层接口列表默认调度策略和用户自定义调度策略都必须继承自该类Run基类接口负责执行调度策略MakeScheduleConfig基类接口负责根据传入的参数信息生成 scheduleCfg 配置信息。两者分工明确MakeScheduleConfig在host 侧tiling 阶段运行产出调度配置Run在device 侧AI Core 上运行消费调度配置并驱动计算。Run对应的类模板原型如下摘自 include/elewise/kernel/schedule.htemplate typename BlockOp, const auto Policy, typename ScheduleCfg class BaseKernelSchedule { template typename OpParam, typename... Args aicore inline void Run(OpParam cfg, Args... args) };二、函数原型与参数说明Run是BaseKernelSchedule的公开成员函数函数原型为template typename BlockOp, const auto Policy, typename ScheduleCfg class BaseKernelSchedule { template typename OpParam, typename... Args __aicore__ inline void Run(OpParam cfg, Args... args) };模板参数说明参数名称参数类型输入/输出数据类型参数说明默认值BlockOp模板参数输入NAblock 层对象类型与 kernel 层是被包含关系NAPolicy模板参数输入NAkernel 层的用户静态策略类型NAScheduleCfg模板参数输入NAkernel 层调度配置类型NAOpParam模板参数输入NAscheduleCfg根据用户设置的类型实例化NAArgs模板参数输入NA用户的输入参数列表类型根据用户传入的参数实例化NA函数形参说明参数名称参数类型输入/输出数据类型参数说明默认值cfg函数形参输入OpParam用户定义的 schedule 配置NAargs函数形参输入Args用户的输入参数列表NA返回值说明返回值数据类型返回值说明voidNARun无返回值其作用是把调度动作就地完成修改cfg引用的配置、把参数打包后传递给 block 层执行不向调用方回传结果。需要特别说明的是OpParam并不是裸的ScheduleCfg。从 kernel 层对象构建类 include/elewise/kernel/builder.h 可以看到其完整定义struct OpParam { ScheduleCfg kernelParam; // kernel 层调度配置tiling 信息 typename BlockOp::ScheduleCfgClz blockParam; // block 层调度配置 };也就是说Run收到的cfg里同时装着 kernel 层与 block 层两份调度配置这为Run内部调用BlockOp::Run时传递 block 层参数提供了数据来源。三、Run 的执行流程源码级拆解Run在源码 include/elewise/kernel/schedule.h 中的实现非常精炼核心只有四步template typename OpParam, typename... Args __aicore__ inline void Run(OpParam cfg, Args... args) { // 1. 配置 block 调度参数计算当前核需要处理的元素总数 typename BlockOp::ScheduleClz::ParamStruct configBlock cfg.blockParam; configBlock.totalElemCnt CalCurCoreEleCnt(cfg.kernelParam); // 2. 将用户传入的 GM 地址/标量参数打包为元组 auto argTuple AscendC::Std::forward_as_tuple(AscendC::Std::forwardArgs(args)...); // 3. 根据表达式的参数占位信息构造 device 侧真正使用的参数含 GM 偏移 auto params PrepareParamsParams(cfg.kernelParam, argTuple); auto convertArgs ConvertArgsParams(params, argTuple); // 4. 实例化 block 层对象并执行 BlockOp blockOp{}; blockOp.Run(configBlock, convertArgs); }3.1 计算当前核的处理量CalCurCoreEleCnt调度最关键的一步是让每个 AI Core 知道自己要处理多少元素。CalCurCoreEleCntschedule.h根据ScheduleCfg中的 tiling 字段计算__aicore__ inline auto CalCurCoreEleCnt(ScheduleCfg cfg) { uint64_t blockNum cfg.blockNum; // 总核数 uint64_t actualNum cfg.unitNum * cfg.unitNumPerCore; // 基础每核基本块数 × 基本块元素数 if (AscendC::GetBlockIdx() cfg.moreUnitCoreNum) { actualNum cfg.unitNum; // 靠前的核多分一个完整基本块 } if (AscendC::GetBlockIdx() blockNum - 1) { actualNum cfg.tailNum; // 最后一个核追加尾块元素 } return actualNum; }这里用到DefaultKernelConfig的全部字段其定义与默认值见 include/elewise/kernel/builder.h成员名称成员类型成员说明默认值blockNumuint32_t启用的核的数量1unitNumPerCoreuint64_t平均每个核处理的基本块个数0moreUnitCoreNumuint64_t核均分后需要处理额外多出来的基本块的核的数量0tailNumuint64_t最后一个核要处理的尾块元素数量0unitNumuint64_t基本块的元素数量1这种均匀切分 余块均摊 尾块兜底的模式正是DefaultKernelPolicy中UniformSegment均匀切分策略在 device 侧的落地形态策略定义见 DefaultKernelPolicy。3.2 组装参数PrepareParams 与 ConstructParamPrepareParamsschedule.h利用编译期展开的索引序列index_sequence对表达式推导出的参数列表Params逐个调用ConstructParam构造 device 侧实参。ConstructParamschedule.h区分两类参数Tensor 类参数取出对应的 GM 地址叠加CalGMOffset计算出的当前核数据偏移后返回偏移后的__gm__ uint8_t*指针从而让每个核只访问自己负责的数据区间标量参数若传入的是指针则解引用取值否则原样传递。template typename ParamType, typename ArgTup __aicore__ inline auto ConstructParam(ScheduleCfg config, ArgTup args) { auto arg AscendC::Std::getParamType::number - 1(args); if constexpr (!std::is_scalar_vtypename ParamType::Type) { using DTypeTmp typename ParamType::Type::PrimType; uint64_t offset CalGMOffset(config); auto ptr (uint64_t)(arg) sizeof(DTypeTmp) * offset; return reinterpret_cast__gm__ uint8_t*(ptr); } else { using argType decltype(arg); if constexpr (std::is_pointer_vargType) { return *arg; } else { return arg; } } }3.3 计算 GM 偏移CalGMOffsetCalGMOffsetschedule.h与CalCurCoreEleCnt严格对称负责回答当前核的数据从 GM 哪里开始__aicore__ inline auto CalGMOffset(ScheduleCfg config) { if (AscendC::GetBlockIdx() config.moreUnitCoreNum) { return AscendC::GetBlockIdx() * (config.unitNumPerCore * config.unitNum config.unitNum); } return config.unitNumPerCore * config.unitNum * AscendC::GetBlockIdx() config.moreUnitCoreNum * config.unitNum; }其语义是前moreUnitCoreNum个核每个多处理一个基本块因此它们的起始偏移要多乘一个unitNum其余核则按前面的核多占的量累加补偿。这一对函数计算元素数与计算偏移共同保证了多核之间数据区间不重叠、不遗漏。3.4 参数重排与下发ConvertArgs 与 BlockOp::RunConvertArgsschedule.h将表达式声明的参数序与用户传入的参数序对齐通过Atvoss::Util::Find_v在Params中按占位号查找匹配的参数找到则用构造好的 device 侧参数否则回退到用户原始参数。这样即使PlaceHolder的编号顺序与用户传参顺序不完全一致也能正确映射。最后Run实例化 block 层对象BlockOp blockOp{}并调用blockOp.Run(configBlock, convertArgs)把计算好的元素总数configBlock.totalElemCnt与参数元组交给 block 层调度对应接口 BaseBlockSchedule::Run至此 kernel 层调度完成使命。四、Run 的数据从哪来与 MakeScheduleConfig 的分工Run消费的cfg.kernelParam全部字段由基类的另一个静态接口MakeScheduleConfig在 host 侧生成详见 docs/api/BaseKernelSchedule_MakeScheduleConfig.md 与源码 schedule.h。其核心逻辑从arguments元组第一个元素取出输入 Tensor 的shape_vector()累乘得到总元素数totalEleNumkernelParam.unitNum ACTUAL_N_ASSIGN即单个基本块的实际元素数。ACTUAL_N_ASSIGN与 tile 形状、BASIC_BLOCK相关在 schedule.h 中计算TILE_SHAPE_SIZE 1时取 32否则取 TileShape 最后一维的值BASIC_CORE_ELE_NUM按BASIC_BLOCK向上对齐到基本块整数倍表示单核能处理的最大元素数若totalEleNum BASIC_CORE_ELE_NUM则只启用 1 个核blockNum 1全部元素作为尾块交给最后一个核unitNumPerCore moreUnitCoreNum 0否则按每核基本块数均摊blockNum取(totalUnitCnt basicCoreUnitNum - 1) / basicCoreUnitNum并受ArchTag::CORE_NUM上限约束unitNumPerCore为商moreUnitCoreNum为余数tailNum为totalEleNum % ACTUAL_N_ASSIGNshape 为空或总元素数为 0 时打印错误并返回false其余情况返回true。可以清晰看到MakeScheduleConfighost/tiling与Rundevice/执行通过ScheduleCfg这份契约完全解耦这也是自定义调度策略只需分别实现生成配置与执行配置两个环节即可的原因。五、端到端调用链从 DeviceAdapter::Run 到 BaseKernelSchedule::RunBaseKernelSchedule::Run不是被用户直接调用的它处在一条完整的调用链末端。以默认用法为例完整的链路如下关键实现见 include/elewise/device/device_adapter.h用户调用DeviceAdapter::Run(arguments, stream)device_adapter.h内部通过CalculateTilingKernelOp触发CalcParamdevice_adapter.h依次调用KernelOp::ScheduleClz::MakeScheduleConfig(arguments, opParam.kernelParam)生成 kernel 层 tiling再调用 block 层的MakeScheduleConfig生成 block 层配置LaunchKernelWithDataTuple以opParam.kernelParam.blockNum为网格规模启动KernelCustomkerneldevice_adapter.h每个核一个 blockkernel 入口经KernelWrapper转发到KernelBuilder::Runinclude/elewise/kernel/builder.h其中通过static_assert限定参数只能是标量或 Tensor 类型KernelBuilder::Run实例化ScheduleClz schedule即ScheduleBlockOp, Policy, ScheduleCfg默认是DefaultKernelSchedule调用schedule.Run(cfg, args...)最终进入本文主角BaseKernelSchedule::Run完成元素数计算、参数偏移与 block 层下发。也就是说BaseKernelSchedule::Run是整个 ATVOSS 分层调度流水线上 kernel 层的最后一公里执行引擎。六、如何使用继承默认调度与自定义调度6.1 直接使用默认调度DefaultKernelSchedule是对BaseKernelSchedule的空实现继承schedule.h能力完全来自基类template typename BlockOp, const auto Policy, typename ScheduleCfg class DefaultKernelSchedule : public BaseKernelScheduleBlockOp, Policy, ScheduleCfg {};同时KernelBuilder的模板参数已把Schedule默认绑定为DefaultKernelSchedulebuilder.h因此绝大多数算子无需显式书写调度类。完整可运行示例见 docs/api/DefaultKernelSchedule.mdtemplate typename InputDtype, typename OutputDtype struct AddSubConfig { struct AddSubCompute { template template typename class Tensor __host_aicore__ constexpr auto Compute() const { auto in1 Atvoss::PlaceHolder1, TensorInputDtype, Atvoss::ParamUsage::IN(); auto in2 Atvoss::PlaceHolder2, TensorInputDtype, Atvoss::ParamUsage::IN(); auto in3 Atvoss::PlaceHolder3, InputDtype, Atvoss::ParamUsage::IN(); auto out Atvoss::PlaceHolder4, TensorOutputDtype, Atvoss::ParamUsage::OUT(); return (out in1 in2 - in3); }; }; static constexpr Atvoss::Ele::DefaultKernelPolicy kernelPolicy{Atvoss::Ele::DefaultSegmentPolicy::UniformSegment}; using ArchTag Atvoss::Arch::DAV_3510; using BlockOp Atvoss::Ele::BlockBuilderAddSubCompute, ArchTag; using KernelOp Atvoss::Ele::KernelBuilder BlockOp, kernelPolicy, Atvoss::Ele::DefaultKernelConfig, Atvoss::Ele::DefaultKernelSchedule; // ← 显式指定默认 kernel 层调度 using DeviceOp Atvoss::DeviceAdapterKernelOp; }; template typename InputDtype, typename OutputDtype static void Run() { /* ACL init and stream create */ ... Atvoss::TensorInputDtype in1(deviceIn1, {{3, 4, 0, 0, 0, 0, 0, 0}}, 2); Atvoss::TensorInputDtype in2(deviceIn2, {{3, 4, 0, 0, 0, 0, 0, 0}}, 2); InputDtype in3 5.0; Atvoss::TensorOutputDtype out(deviceOut, {{3, 4, 0, 0, 0, 0, 0, 0}}, 2); auto arguments Atvoss::ArgumentsBuilder{}.inputOutput(in1, in2, in3, out).attr(dim, 5).build(); using DeviceOp typename AddSubConfigInputDtype, OutputDtype::DeviceOp; DeviceOp deviceOp; deviceOp.Run(arguments, stream); } int main(int argc, char const* argv[]) { Runfloat, float(); return 0; }6.2 自定义调度策略的两种写法从 docs/api/KernelBuilder.md 的模板签名可知Schedule是一个模板模板参数template typename BlockOp, const auto Policy defaultKernelPolicy, typename ScheduleCfg DefaultKernelConfig, template typename, const auto, typename class Schedule DefaultKernelSchedule class KernelBuilder因此自定义调度有两种途径继承BaseKernelScheduleBlockOp, Policy, ScheduleCfg覆写Rundevice 侧执行逻辑与MakeScheduleConfighost 侧 tiling 逻辑可复用基类实现或完全重写再将自定义类作为KernelBuilder的第 4 个模板实参传入直接继承DefaultKernelSchedule在保留默认 tiling 与执行能力的基础上做增量覆写适合只调整某一环节如修改CalGMOffset的切分策略的场景。无论哪种方式自定义调度类都必须保证host 侧MakeScheduleConfig填满ScheduleCfg中Run会用到的字段blockNum、unitNumPerCore、moreUnitCoreNum、tailNum、unitNum否则CalCurCoreEleCnt与CalGMOffset会得到错误的分片结果。七、约束与注意事项运行环境约束Run及CalCurCoreEleCnt、PrepareParams、ConvertArgs、CalGMOffset等实现均带有__aicore__修饰且被#if !defined(__ATVOSS_HOST_ONLY__)宏保护schedule.h只能在 AI Core 侧编译执行host 侧编译时__ATVOSS_HOST_ONLY__定义这些实现会被裁剪。参数类型约束KernelBuilder::Run通过static_assert限定传入参数只能是标量或 Tensor 类型builder.h其他类型在编译期即报错。shape 有效性MakeScheduleConfig对空 shape 或总元素数为 0 的输入会打印[ERROR]: [Atvoss][Kernel] Shape info error并返回falseDeviceAdapter::Run检测到失败会返回 -1 并终止 kernel 启动。tiling 上限blockNum会被限制在ArchTag::CORE_NUM以内数据量超出单核处理能力时按核数封顶这是均匀切分策略的固有行为若需更精细的多核调度应自定义Schedule。测试覆盖仓库 tests/st 目录下的test_op_*系列与test_tile_rms_norm_*系列用例均通过默认链路DeviceAdapter::Run→BaseKernelSchedule::Run执行验证可作为默认调度行为正确性的参考样例。八、总结BaseKernelSchedule::Run是 ATVOSS kernel 层调度策略的执行引擎它用极简的模板化接口封装了多核数据切分的全部细节CalCurCoreEleCnt决定算多少CalGMOffset决定从哪算PrepareParams/ConvertArgs负责参数映射与 GM 指针偏移最后把控制权交给 block 层。理解它就理解了 ATVOSS 从 host tiling 到 AI Core 实际执行之间最关键的一环而通过继承BaseKernelSchedule或DefaultKernelSchedule开发者可以在不触碰上层框架的前提下自由定制自己的 kernel 层调度策略。【免费下载链接】atvossATVOSSAscend C Templates for Vector Operator Subroutines是一套基于Ascend C开发的Vector算子库致力于为昇腾硬件上的Vector类融合算子提供极简、高效、高性能、高拓展的编程方式。项目地址: https://gitcode.com/cann/atvoss创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考