EVG 扩展说明【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass本文整理当前 EVG 的扩展边界主要回答两件事什么时候新增ComputeFn什么时候新增节点扩展时需要遵守哪些约束。接入方式见 evg_api整体设计见 01_evg_design。先判断扩展类型优先做这一步避免把本来只需新增一个算子的事情做成新节点。场景处理方式只是逐元素计算输入输出都在 UB不需要 GM、layout、workspace新增ComputeFn需要读 GM、写 GM、管理 layout、申请 UB或需要自己的Arguments / Params新增节点可以直接按下面的口径判断只扩“怎么算”加ComputeFn涉及“怎么取、怎么放、怎么占资源”加节点当前最典型的参考实现ComputeFnAdd、Muls、LeakyRelu、AddRelu节点VisitorAuxLoad、VisitorAuxStore、VisitorRowBroadcast新增 ComputeFn放置位置当前ComputeFn统一放在include/catlass/epilogue/fusion/operations.hppVisitorCompute通过模板参数ComputeFn实例化这里的算子。新增逐元素算子时直接沿用这里的写法即可通常不再单独新增节点。需要保持的形态最小形态如下template typename T struct SomeOp { CATLASS_DEVICE void operator()( AscendC::LocalTensorT dst, uint32_t compute_length, AscendC::LocalTensorT const src0, AscendC::LocalTensorT const src1 ) const { // do compute } };固定规则第一个参数是输出dst第二个参数是compute_length后续参数才是输入输入路数由算子语义决定如果算子带标量参数保持聚合类型即可template typename T struct ClampMin { T min_value; CATLASS_DEVICE void operator()( AscendC::LocalTensorT dst, uint32_t compute_length, AscendC::LocalTensorT const src ) const { AscendC::Maxs(dst, src, min_value, compute_length); } };对应接入方式using ClampMinOp Epilogue::Fusion::VisitorComputeEpilogue::Fusion::ClampMin, ElementC, ElementC; typename ClampMinOp::Arguments args{{0.0f}};当前约束1. 只做纯计算当前实现中ComputeFn只承载 UB 内的计算下面这些职责仍放在节点层访问 GM管理 layout申请 workspace处理事件同步依赖全局坐标这些职责不放在ComputeFn内处理。2. 类型规则沿用VisitorCompute当前VisitorCompute要求所有输入类型都等于ElementCompute。因此混合精度场景通常先插VisitorCast输入类型兼容仍放在图里处理不放在ComputeFn内展开3. 多步 V 计算自己补 barrier如果一个operator()内包含多步 V 计算应像AddRelu一样在步骤之间补AscendC::PipeBarrierPIPE_V()。简单说单条原子指令通常不用额外补多步串联计算通常要补齐4. 多输入算子沿用现有展开方式像Add、Mul这类多输入算子当前实现采用链式展开。新增多输入算子时尽量沿用现有模式避免单独引入新的调用约定。5. 保持聚合初始化友好VisitorCompute当前按OpElementCompute{...}构造算子。新增ComputeFn时保持简单字段和聚合初始化会更贴合现有实现。新增节点放置位置当前节点实现统一放在include/catlass/epilogue/fusion/visitor_*.hpp新增节点后还应接入include/catlass/epilogue/fusion/fusion.hpp优先参考的同级实现新增节点时与现有同级实现保持相同的结构和职责VisitorAuxLoad典型叶子节点负责从 GM 读取 tileVisitorAuxStore典型根节点负责最终写回VisitorRowBroadcast典型跨阶段节点LOAD取数、COMPUTE扩展VisitorCompute纯计算节点如果需求和它一致通常直接复用即可最小骨架template class Element struct VisitorSomeNode : VisitorImpl { using VisitorImpl::VisitorImpl; using ElementOutput Element; struct Arguments {}; struct Params {}; template class ProblemShape static constexpr Params to_underlying_arguments(ProblemShape const, Arguments const args, void* workspace) { return Params(/* ... */); } template class ProblemShape static size_t get_workspace_size(ProblemShape const, Arguments const) { return 0; } template class ProblemShape static bool can_implement(ProblemShape const, Arguments const) { return true; } struct Callbacks : EmptyCallbacks { template VisitStage Stage, class ArchTag, class TensorC, typename... Inputs CATLASS_DEVICE auto visit( TensorC const tensorTile, MatrixCoord const alignedTileShape, MatrixCoord const globalOffset, Inputs const... inputs ) { // stage-specific work } }; template class ArchTag CATLASS_DEVICE auto get_callbacks( Arch::ResourceArchTag resource, uint32_t ub_offset, uint32_t compute_length ) { return Callbacks(/* ... */); } Params params; };必备组成与职责组成作用约束ElementOutput告诉图组织器节点输出类型与实际输出一致Arguments用户侧直接填写的参数保持简单支持花括号初始化Paramsdevice 侧真正使用的参数尽量只保留执行需要的信息to_underlying_arguments把Arguments转成Params需要 workspace 时在这里完成映射get_workspace_size声明节点自有 workspace只计算自己那一段can_implement轻量合法性检查至少检查容易误用的参数Callbacks::visit执行节点动作严格按阶段写逻辑get_callbacks申请 UB 并构造回调谁申请 UB谁推进ub_offset当前约束1. 阶段语义不能打乱节点逻辑按三阶段组织LOAD准备输入、从 GM 读入 UBCOMPUTEUB 内计算、变换、广播STORE把结果真正写出明确规则真正写回外部地址的动作放在STORE跨阶段节点可以存在但每一步分开放置本该在STORE的动作通常不前置到COMPUTE2. 节点职责保持单一当前实现更适合让一个节点只承担一类职责例如读、算、写、广播。“读 GM 算 写 GM”这类复合行为拆成多个节点后更贴合现有图组织方式。3. layout 一律按完整张量理解如果节点带layout它描述的是完整 GM 张量不是当前 tile。当前实现都按这个口径处理新增节点时保持一致即可。4. UB 分配口径保持一致get_callbacks中如果申请 UB通常按下面的口径处理谁申请谁推进ub_offset只申请自己需要的那一段大小按compute_length和元素大小推导分配结果不超过当前架构允许的 UB 上限5. 不在节点里接管外层事件EVG 的双缓冲和事件同步由BlockEpilogue统一管理。节点内部只负责按阶段执行如果节点内有多步 V 计算可以补AscendC::PipeBarrierPIPE_V()外层同步节奏仍由BlockEpilogue管理。6.Arguments继续支持直写新增节点后用户侧通常仍可直接写出整张图的聚合初始化typename EVG::Arguments args{ {}, {deviceX, layoutX}, {{2.0f}}, {deviceD, layoutD} };如果一个节点让整张图的参数初始化显著变复杂通常说明接口设计和现有规范还有些距离。7. 先复用现有节点再决定是否新增如果诉求只是多一个逐元素算子通常继续走ComputeFn VisitorCompute。只有当现有节点无法表达所需的数据访问、layout 或资源行为时再新增节点。visit签名visit可以写成可变参数也可以写成固定签名取决于节点职责叶子节点、未来可能扩展输入路数的节点可用可变参数严格单输入或固定多输入节点可以直接写死输入签名关键不在于形式统一而在于节点自身要明确需要几路输入输入顺序是什么输入类型是否符合预期实施顺序新增ComputeFn时在operations.hpp中新增算子保持operator()(dst, compute_length, inputs...)这一形态多步 V 计算时补PipeBarrier用VisitorCompute..., ElementType, Scalars...接入用最小图验证参数和类型口径新增节点时新建visitor_xxx.hpp继承VisitorImpl定义ElementOutput / Arguments / Params实现to_underlying_arguments / get_workspace_size / can_implement在Callbacks::visit中按阶段填充逻辑在get_callbacks中完成 UB 申请接入fusion.hpp用最小样例验证能在TreeVisitor或TopologicalVisitor中组合实现前检查清单新增ComputeFn前至少确认它是不是纯逐元素计算是否真的不需要 GM、layout、workspace是否有多步 V 计算需要 barrier是否没必要新增节点新增节点前至少确认不能只靠ComputeFn解决Arguments能否保持简洁并支持花括号初始化LOAD / COMPUTE / STORE是否分工明确是否需要 UB 或 workspacelayout 是否按完整张量解释是否补了基本的can_implement参考实现扩展时可直接对照这些现有实现include/catlass/epilogue/fusion/operations.hppinclude/catlass/epilogue/fusion/visitor_compute.hppinclude/catlass/epilogue/fusion/visitor_aux_load.hppinclude/catlass/epilogue/fusion/visitor_aux_store.hppinclude/catlass/epilogue/fusion/visitor_row_broadcast.hpp可按下面顺序对照只加逐元素算子先看operations.hpp加读写类节点先看visitor_aux_load.hpp、visitor_aux_store.hpp加跨阶段节点先看visitor_row_broadcast.hpp【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考