CATLASS EVG扩展说明

EVG 扩展说明

【免费下载链接】catlass本项目是CANN的算子模板库，提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass

本文整理当前 EVG 的扩展边界，主要回答两件事：什么时候新增ComputeFn，什么时候新增节点；扩展时需要遵守哪些约束。接入方式见 evg_api，整体设计见 01_evg_design。

先判断扩展类型

优先做这一步，避免把本来只需新增一个算子的事情做成新节点。

可以直接按下面的口径判断：

• 只扩“怎么算”，加ComputeFn
• 涉及“怎么取、怎么放、怎么占资源”，加节点

当前最典型的参考实现：

• ComputeFn：Add、Muls、LeakyRelu、AddRelu
• 节点：VisitorAuxLoad、VisitorAuxStore、VisitorRowBroadcast

新增 ComputeFn

放置位置

当前ComputeFn统一放在：

• include/catlass/epilogue/fusion/operations.hpp

VisitorCompute通过模板参数ComputeFn实例化这里的算子。新增逐元素算子时，直接沿用这里的写法即可，通常不再单独新增节点。

需要保持的形态

最小形态如下：

template <typename T> struct SomeOp { CATLASS_DEVICE void operator()( AscendC::LocalTensor<T>& dst, uint32_t compute_length, AscendC::LocalTensor<T> const& src0, AscendC::LocalTensor<T> const& src1 ) const { // do compute } };

固定规则：

• 第一个参数是输出dst
• 第二个参数是compute_length
• 后续参数才是输入
• 输入路数由算子语义决定

如果算子带标量参数，保持聚合类型即可：

template <typename T> struct ClampMin { T min_value; CATLASS_DEVICE void operator()( AscendC::LocalTensor<T>& dst, uint32_t compute_length, AscendC::LocalTensor<T> const& src ) const { AscendC::Maxs(dst, src, min_value, compute_length); } };

对应接入方式：

using ClampMinOp = Epilogue::Fusion::VisitorCompute<Epilogue::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::PipeBarrier<PIPE_V>()。

简单说：

• 单条原子指令通常不用额外补
• 多步串联计算通常要补齐

4. 多输入算子沿用现有展开方式

像Add、Mul这类多输入算子，当前实现采用链式展开。新增多输入算子时，尽量沿用现有模式，避免单独引入新的调用约定。

5. 保持聚合初始化友好

VisitorCompute当前按Op<ElementCompute>{...}构造算子。新增ComputeFn时，保持简单字段和聚合初始化会更贴合现有实现。

新增节点

放置位置

当前节点实现统一放在：

• include/catlass/epilogue/fusion/visitor_*.hpp

新增节点后，还应接入：

• include/catlass/epilogue/fusion/fusion.hpp

优先参考的同级实现

新增节点时，与现有同级实现保持相同的结构和职责：

• VisitorAuxLoad：典型叶子节点，负责从 GM 读取 tile
• VisitorAuxStore：典型根节点，负责最终写回
• 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::Resource<ArchTag>& resource, uint32_t& ub_offset, uint32_t compute_length ) { return Callbacks(/* ... */); } Params params; };

必备组成与职责

当前约束

1. 阶段语义不能打乱

节点逻辑按三阶段组织：

• LOAD：准备输入、从 GM 读入 UB
• COMPUTE：UB 内计算、变换、广播
• STORE：把结果真正写出

明确规则：

• 真正写回外部地址的动作放在STORE
• 跨阶段节点可以存在，但每一步分开放置
• 本该在STORE的动作通常不前置到COMPUTE

2. 节点职责保持单一

当前实现更适合让一个节点只承担一类职责，例如读、算、写、广播。
“读 GM + 算 + 写 GM”这类复合行为拆成多个节点后，更贴合现有图组织方式。

3. layout 一律按完整张量理解

如果节点带layout，它描述的是完整 GM 张量，不是当前 tile。当前实现都按这个口径处理，新增节点时保持一致即可。

4. UB 分配口径保持一致

get_callbacks中如果申请 UB，通常按下面的口径处理：

• 谁申请，谁推进ub_offset
• 只申请自己需要的那一段
• 大小按compute_length和元素大小推导
• 分配结果不超过当前架构允许的 UB 上限

5. 不在节点里接管外层事件

EVG 的双缓冲和事件同步由BlockEpilogue统一管理。节点内部只负责按阶段执行；如果节点内有多步 V 计算，可以补AscendC::PipeBarrier<PIPE_V>()，外层同步节奏仍由BlockEpilogue管理。

6.Arguments继续支持直写

新增节点后，用户侧通常仍可直接写出整张图的聚合初始化：

typename EVG::Arguments args{ {}, {deviceX, layoutX}, {{2.0f}}, {deviceD, layoutD} };

如果一个节点让整张图的参数初始化显著变复杂，通常说明接口设计和现有规范还有些距离。

7. 先复用现有节点，再决定是否新增

如果诉求只是多一个逐元素算子，通常继续走ComputeFn + VisitorCompute。
只有当现有节点无法表达所需的数据访问、layout 或资源行为时，再新增节点。

visit签名

visit可以写成可变参数，也可以写成固定签名，取决于节点职责：

• 叶子节点、未来可能扩展输入路数的节点：可用可变参数
• 严格单输入或固定多输入节点：可以直接写死输入签名

关键不在于形式统一，而在于节点自身要明确：

• 需要几路输入
• 输入顺序是什么
• 输入类型是否符合预期

实施顺序

新增ComputeFn时：

1. 在operations.hpp中新增算子
2. 保持operator()(dst, compute_length, inputs...)这一形态
3. 多步 V 计算时补PipeBarrier
4. 用VisitorCompute<..., ElementType, Scalars...>接入
5. 用最小图验证参数和类型口径

新增节点时：

1. 新建visitor_xxx.hpp
2. 继承VisitorImpl<>
3. 定义ElementOutput / Arguments / Params
4. 实现to_underlying_arguments / get_workspace_size / can_implement
5. 在Callbacks::visit中按阶段填充逻辑
6. 在get_callbacks中完成 UB 申请
7. 接入fusion.hpp
8. 用最小样例验证能在TreeVisitor或TopologicalVisitor中组合

实现前检查清单

新增ComputeFn前，至少确认：

• 它是不是纯逐元素计算
• 是否真的不需要 GM、layout、workspace
• 是否有多步 V 计算需要 barrier
• 是否没必要新增节点

新增节点前，至少确认：

• 不能只靠ComputeFn解决
• Arguments能否保持简洁并支持花括号初始化
• LOAD / COMPUTE / STORE是否分工明确
• 是否需要 UB 或 workspace
• layout 是否按完整张量解释
• 是否补了基本的can_implement

参考实现

扩展时，可直接对照这些现有实现：

• include/catlass/epilogue/fusion/operations.hpp
• include/catlass/epilogue/fusion/visitor_compute.hpp
• include/catlass/epilogue/fusion/visitor_aux_load.hpp
• include/catlass/epilogue/fusion/visitor_aux_store.hpp
• include/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