EVG API【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlassEVGEpilogue Visitor Graph是 CATLASS 当前用于组织 GEMM 尾处理的图式接口。本文只保留接入、参数顺序和节点口径执行模型见 01_evg_design扩展约束见 02_evg_extension。核心入口层级入口作用KernelBasicMatmulTlaVisitor代码路径include/catlass/gemm/kernel/basic_matmul_tla_visitor.hppAIC 把 MMAD 结果写到 GM workspaceAIV 再执行 EVGKernelBasicMatmulTlaUbVisitor代码路径include/catlass/gemm/kernel/basic_matmul_tla_ub_visitor.hppAIC 把结果保留在 UBAIV 直接在 UB 上执行 EVGBlockBlockEpilogueEpilogueVisitor..., ArchTag, ComputeLength, EVG, ElementC代码路径include/catlass/epilogue/block/block_epilogue_visitor.hpp负责 tile 切分、双缓冲和三阶段调度FusionTreeVisitor代码路径include/catlass/epilogue/fusion/tree_visitor.hpp用树结构描述尾处理FusionTopologicalVisitor代码路径include/catlass/epilogue/fusion/topological_visitor.hpp用 DAG 描述尾处理并复用中间结果接入顺序一个 EVG kernel 的组装顺序通常是选BlockMmad定义EVG用EVG组装BlockEpilogue选择 visitor kernel构造EVG::Arguments把EVG::Arguments填进 kernelArguments最常见的写法如下using EVG Epilogue::Fusion::TreeVisitor Epilogue::Fusion::VisitorAuxStoreElementC, LayoutC, Epilogue::Fusion::TreeVisitor Epilogue::Fusion::VisitorComputeEpilogue::Fusion::Add, ElementC, Epilogue::Fusion::VisitorAccLoadElementC, Epilogue::Fusion::VisitorAuxLoadElementC, LayoutC ; using BlockEpilogue Epilogue::Block::BlockEpilogue Epilogue::EpilogueVisitorfalse, ArchTag, IntcomputeLength, EVG, ElementC ; using MatmulKernel Gemm::Kernel::BasicMatmulTlaVisitorBlockMmad, BlockEpilogue, BlockScheduler;如果使用 UB workspace 路径改动点主要有两个Epilogue::EpilogueVisitortruekernel 换成BasicMatmulTlaUbVisitor并且VisitorAccLoad通常会写成using EpilogueDispatchPolicy Epilogue::EpilogueVisitortrue; using AccLoad Epilogue::Fusion::VisitorAccLoad ElementC, EpilogueDispatchPolicy::USE_UB_WORKSPACE ;TreeVisitorTreeVisitorNodeOp, ChildOps...适合树状表达式。形态using EVG Epilogue::Fusion::TreeVisitor ParentOp, ChildOp1, ChildOp2 ;参数顺序TreeVisitor的Arguments顺序和模板顺序不同规则是“先子后父”。typename EVG::Arguments evg_args{ { ChildOp1::Arguments{}, ChildOp2::Arguments{}, ParentOp::Arguments{} } };如果有嵌套TreeVisitor就按每一层“先子后父”的规则递归书写。Arguments 写法Arguments本质上是聚合结构可以直接用花括号嵌套构造不必先把每一层都显式写成XXX::Arguments变量。例如D C X可以直接写成typename EVG::Arguments evg_args{ { {}, {deviceX, layoutX}, {} }, {deviceD, layoutD} };这里的{}、{deviceX, layoutX}、{deviceD, layoutD}会按当前位置自动匹配到对应节点的Arguments类型。只要嵌套层次和顺序正确就不需要显式写类型名。适用场景D C XD silu(C)D cast(add(C, X))执行流程TopologicalVisitorTopologicalVisitorEdgeTuple, Ops...适合中间结果要复用的场景。形态using Edges tla::tuple tla::seq, tla::seq0, tla::seq1, tla::seq2, tla::seq2, tla::seq3, 4, tla::seq5 ; using EVG Epilogue::Fusion::TopologicalVisitor Edges, Op0, Op1, Op2, Op3, Op4, Op5, Op6 ;参数顺序TopologicalVisitor的Arguments严格按Ops...的平铺顺序书写typename EVG::Arguments evg_args{ Op0::Arguments{}, Op1::Arguments{}, Op2::Arguments{}, Op3::Arguments{}, Op4::Arguments{}, Op5::Arguments{}, Op6::Arguments{} };Arguments 写法TopologicalVisitor同样可以直接用花括号按节点顺序填写不必逐个显式写类型typename EVG::Arguments evg_args{ {}, {{2.0f}}, {}, {{-1.0f}}, {{1.0f}}, {}, {deviceD, layoutD} };判断规则很简单第几个节点就写在第几个位置节点Arguments里有几个字段就按字段顺序写几层花括号没有字段时直接写{}即可适用场景同一个中间结果被多个后继节点消费希望避免在一个 tile 上重复计算执行流程这里的缓存只覆盖当前这次visitStage(...)调用也就是当前 tile 的当前阶段进入下一个阶段时会重新从根节点开始一轮访问。两种组织方式的取舍见设计文档中的“图组织方式”。节点总览当前实现中EVG 常用节点如下。节点头文件作用VisitorAccLoadElement, USE_UB_WORKSPACEvisitor_acc_load.hpp读取 GEMM 结果VisitorAuxLoadElement, Layoutvisitor_aux_load.hpp从外部 GM 读取输入VisitorComputeComputeFn, ElementCompute, Scalars...visitor_compute.hpp做逐元素计算VisitorCastElementTo, ElementFrom, RoundStylevisitor_cast.hpp做类型转换VisitorAuxStoreElement, Layoutvisitor_aux_store.hpp把结果写回 GMVisitorRowBroadcastElement, Layoutvisitor_row_broadcast.hpp读取1 x N行向量并广播到 tile阶段与放置原则所有节点都运行在统一的三阶段模型里LOADCOMPUTESTORE但不是每个节点三个阶段都会做事VisitorAccLoad、VisitorAuxLoad主要工作在LOADVisitorCompute、VisitorCast主要工作在COMPUTEVisitorAuxStore主要工作在STOREVisitorRowBroadcast同时跨LOAD和COMPUTE按图来放时可以先用这个简单规则判断叶子节点放数据源VisitorAccLoad、VisitorAuxLoad、VisitorRowBroadcast中间节点放变换和计算VisitorCast、VisitorCompute根节点放输出VisitorAuxStore如果是TopologicalVisitor也是同样的职责只是节点不再嵌套而是按依赖顺序平铺。阶段和节点职责关系图、BlockEpilogue的双缓冲流水时序图都已经放到设计文档见 01_evg_design 的“三阶段执行模型”。节点说明下面按节点分别说明模板参数、放置位置、Arguments写法和特殊要求。先看一眼速查表会更方便节点常见放置位置输入路数直接写进Arguments的形态特别限制VisitorAccLoad叶子节点0{}UB 通路下直接消费 MMAD 的 UB 结果VisitorAuxLoad叶子节点0{ptr, layout}layout描述完整张量VisitorAuxStore根节点1{ptr, layout}当前实现里真正负责落盘VisitorCast中间节点1{}输入类型与ElementFrom一致VisitorRowBroadcast叶子节点0{ptr, layout}layout使用(1, n)二维 layoutVisitorCompute中间节点1 或多路{}或{{...}}所有输入类型与ElementCompute一致VisitorAccLoadVisitorAccLoadElement, USE_UB_WORKSPACEElement读取出来的元素类型通常与 MMAD 输出类型一致USE_UB_WORKSPACE是否直接从 UB 读取 MMAD 结果放置位置与使用要求通常作为叶子节点使用不接收子节点输入输出当前 tile 的C在 GM workspace 路径下主要在LOAD阶段把数据搬到 UB在 UB workspace 路径下直接从 UB 中取当前 MMAD 结果通常放在VisitorCompute、VisitorCast等计算节点的下游常见写法using AccLoad0 Epilogue::Fusion::VisitorAccLoadElementC; using AccLoad1 Epilogue::Fusion::VisitorAccLoadElementC, true;对应的Argumentstypename AccLoad0::Arguments acc_args{}; typename AccLoad1::Arguments acc_ub_args{};直接写进整张图时就是{}VisitorAuxLoadVisitorAuxLoadElement, LayoutElement外部输入张量的元素类型Layout这个外部输入对应的 layout 类型放置位置与使用要求通常作为叶子节点使用不接收子节点输入在LOAD阶段从 GM 按当前 tile 的全局坐标读取数据适合放在VisitorCompute、VisitorCast等计算节点的下游layout描述完整输入张量而不是当前 tilelayout传的是具体 layout 对象不是 layout taglayout类型和模板里的Layout一致layout和ptr指向的数据真实排布一致常见写法using XLoad Epilogue::Fusion::VisitorAuxLoadElementC, LayoutX;对应的Argumentstypename XLoad::Arguments x_args{deviceX, layoutX};直接写进整张图时就是{deviceX, layoutX}可写成auto layoutX tla::MakeLayoutElementC, layout::RowMajor(m, n); using LayoutX decltype(layoutX); using XLoad Epilogue::Fusion::VisitorAuxLoadElementC, LayoutX;VisitorAuxStoreVisitorAuxStoreElement, LayoutElement最终写回数据的元素类型Layout输出张量的 layout 类型放置位置与使用要求接收一个输入一般作为输出节点使用典型放法是作为整张图的根节点负责最终写回真正写回外部内存的动作发生在STORE阶段当前实现会把输入透传返回因此技术上仍可继续参与组合文档和样例里通常把它放在最后作为结果落盘节点输入元素类型和模板里的Element一致不一致时先插入VisitorCastlayout描述完整输出张量而不是当前 tilelayout传的是具体 layout 对象不是 layout taglayout类型和模板里的Layout一致layout和输出 GM 的真实排布一致常见写法using Store Epilogue::Fusion::VisitorAuxStoreElementC, LayoutD;对应的Argumentstypename Store::Arguments store_args{deviceD, layoutD};直接写进整张图时就是{deviceD, layoutD}可写成auto layoutD tla::MakeLayoutElementC, layout::RowMajor(m, n); using LayoutD decltype(layoutD); using Store Epilogue::Fusion::VisitorAuxStoreElementC, LayoutD;VisitorCastVisitorCastElementTo, ElementFrom, RoundStyleElementTo转换后的类型ElementFrom输入类型RoundStyle舍入方式默认可用AscendC::RoundMode::CAST_NONE放置位置与使用要求接收一个输入通常作为中间父节点使用适合放在某个叶子节点或计算节点之上再把结果交给后续计算节点实际计算发生在COMPUTE阶段输入类型与ElementFrom一致输出类型固定为ElementTo如果上下游已经是同一类型就没必要插这个节点一般放在叶子节点或某个计算节点之上不作为数据源或最终输出节点常见写法using CastFp16ToFp32 Epilogue::Fusion::VisitorCastfloat, half, AscendC::RoundMode::CAST_NONE;对应的Argumentstypename CastFp16ToFp32::Arguments cast_args{};直接写进整张图时就是{}VisitorRowBroadcastVisitorRowBroadcastElement, LayoutElement行向量元素类型Layout这条1 x N输入的 layout 类型这里需要注意当前实现按二维 tensor 处理这一路输入所以Layout使用描述(1, n)的 layout 类型而不是只描述(n)的 vector layout 类型。放置位置与使用要求通常作为叶子节点使用不接收子节点输入LOAD阶段读取当前列范围对应的1 x tile_nCOMPUTE阶段把这一行复制扩展成当前tile_m x tile_n适合用于 bias 这类“按列共享、按行广播”的输入当前实现按(1, n)的二维 layout 处理而不是(n)的一维 vector layoutlayout传的是具体 layout 对象不是 layout taglayout类型和模板里的Layout一致常见写法auto layoutBias tla::MakeLayoutElementC, layout::RowMajor(1, n); using LayoutBias decltype(layoutBias); using BiasLoad Epilogue::Fusion::VisitorRowBroadcastElementC, LayoutBias;对应的Argumentstypename BiasLoad::Arguments bias_args{deviceBias, layoutBias};直接写进整张图时就是{deviceBias, layoutBias}VisitorComputeVisitorComputeComputeFn, ElementCompute, Scalars...三个位置分别表示ComputeFn具体算子例如Add、Exp、MulsElementCompute算子工作的元素类型Scalars...额外标量参数类型没有就不写使用口径通常作为中间计算节点输入来自AccLoad、AuxLoad、RowBroadcast、Cast或其他Compute实际计算发生在COMPUTE阶段输入个数与ComputeFn语义一致所有输入类型都与ElementCompute一致类型不一致时先插VisitorCast常见例子using AddOp Epilogue::Fusion::VisitorComputeEpilogue::Fusion::Add, ElementC; using ExpOp Epilogue::Fusion::VisitorComputeEpilogue::Fusion::Exp, ElementC; using MulsOp Epilogue::Fusion::VisitorComputeEpilogue::Fusion::Muls, ElementC, ElementC; using AddsOp Epilogue::Fusion::VisitorComputeEpilogue::Fusion::Adds, ElementC, ElementC; using LeakyReluOp Epilogue::Fusion::VisitorComputeEpilogue::Fusion::LeakyRelu, ElementC, ElementC;对应的Arguments写法如下typename AddOp::Arguments add_args{}; typename ExpOp::Arguments exp_args{}; typename MulsOp::Arguments muls_args{{2.0f}}; typename AddsOp::Arguments adds_args{{1.0f}}; typename LeakyReluOp::Arguments leaky_args{{0.1f}};如果直接写在整张图里也可以只写花括号typename EVG::Arguments evg_args{ {}, {{2.0f}}, {}, {deviceD, layoutD} };这里{{2.0f}}之所以是两层花括号是因为VisitorCompute::Arguments里包了一层scalars元组。Scalars...有多个标量时就继续按顺序写using SomeOp Epilogue::Fusion::VisitorComputeSomeComputeFn, ElementC, float, int32_t; typename SomeOp::Arguments some_args{{1.0f, 2}};直接写进整张图时对应位置就写{{1.0f, 2}}VisitorCompute会检查所有输入类型是否都等于ElementCompute不一致时先插VisitorCast。常用 ComputeFnVisitorCompute依赖operations.hpp中的算子定义。当前实现里常用的有类型算子一元Exp、Relu、Silu、Sqrt、RsqrtFast带标量LeakyRelu、Muls、Adds二元或多元Add、Sub、Mul、Div、Max、Min组合AddReluBlockEpilogue 关键参数EVG 专用的BlockEpilogue模板实参顺序如下using BlockEpilogue Epilogue::Block::BlockEpilogue Epilogue::EpilogueVisitorfalse, ArchTag, IntcomputeLength, EVG, ElementC ;几个关键点EpilogueVisitorfalse从 GM workspace 取 MMAD 结果EpilogueVisitortrue直接从 UB 取 MMAD 结果computeLength单次 tile 处理的元素数按BYTE_PER_C0对齐EVG整张尾处理图ElementCMMAD 输出元素类型Kernel 侧 Arguments以BasicMatmulTlaVisitor为例Arguments里除了 A、B、C 与 layout 外还需要把evg_args放进去。struct Arguments { GemmCoord problemShape; GM_ADDR ptrA; LayoutA layoutA; GM_ADDR ptrB; LayoutB layoutB; GM_ADDR ptrC; LayoutC layoutC; GM_ADDR ptrBias{nullptr}; typename BlockEpilogue::EVG::Arguments evg_args; };这个Arguments也可以直接用花括号整体构造不需要先单独声明内部类型typename MatmulKernel::Arguments arguments{ problemShape, deviceA, layoutA, deviceB, layoutB, deviceD, layoutD, nullptr, { { {}, {deviceX, layoutX}, {} }, {deviceD, layoutD} } };写时主要对齐两点外层字段顺序和MatmulKernel::Arguments一致evg_args内部嵌套顺序和EVG::Arguments一致layout 相关要求放在这里看最合适layoutA/layoutB传完整输入矩阵的具体 layout 对象传的是 layout 对象不是 layout tag这些 layout 会直接用于构造完整 GM tensorlayoutC在 visitor kernel 的公开接口里仍然保留但当前 visitor 路径真正写回时实际以VisitorAuxStore里传入的 layout 为准需要注意的是当前 visitor kernel 的ToUnderlyingArguments()实现并不消费ptrC/layoutC真正的写回位置由evg_args中的VisitorAuxStore决定。如果使用 GM workspace 路径GetWorkspaceSize()会返回C workspace EVG workspace如果使用 UB workspace 路径GetWorkspaceSize()只返回EVG workspacecomputeLength 选择computeLength表示当前链路在每次迭代里处理的元素个数。这个值理论上越大越有利于减少迭代次数、提升效率但它不能超过当前 UB 空间所能容纳的上限所以实际使用时通常是先计算一个可接受的最大值再据此选定computeLength。计算时看三件事可分给 EVG 的 UB 总量同时驻留的 UB buffer 数量是否启用双缓冲最终结果还要按BYTE_PER_C0向下对齐。计算示例 1GM workspace 通路的D C X这类链路通常会同时驻留三块 UB 数据CXOut并且使用双缓冲所以最大computeLength可以直接写成constexpr uint32_t computeLength (ArchTag::UB_SIZE / 3 / 2 / sizeof(ElementC)) / BYTE_PER_C0 * BYTE_PER_C0;A5架构中EVG 样例里通常不会直接按满额去算按216 * 1024作为可用预算来计算computeLength避免实际运行时触发 UB 空间相关报错。写法可以记成constexpr uint32_t computeLength (216 * 1024 / 3 / 2 / sizeof(ElementC)) / BYTE_PER_C0 * BYTE_PER_C0;这里的216 * 1024是A5架构中的保守可用预算。计算示例 2UB workspace 通路的D C X如果走 UB 通路VisitorAccLoad..., true不再额外申请一块 UB buffer所以这类链路通常只需要再为下面两块数据留空间XOut但这时 EVG 不能使用整块 UB因为前半部分已经预留给 MMAD 结果。当前实现里EVG 的起始分配位置是ArchTag::L0C_SIZE / 2也沿用同样的保守预算口径把可用于计算的总 UB 先按216 * 1024代入所以最大computeLength可写成constexpr uint32_t computeLength ((216 * 1024 - ArchTag::L0C_SIZE / 2) / 2 / 2 / sizeof(ElementC)) / BYTE_PER_C0 * BYTE_PER_C0;使用规则每新增一个会单独申请 UB 的节点就把分母里的 buffer 数量加一VisitorAuxStore一般不单独占计算 buffer通常不计入GM 通路下VisitorCast、VisitorCompute、VisitorAuxLoad、VisitorAccLoad、VisitorRowBroadcast通常都要计入UB 通路下VisitorAccLoad..., true通常不单独计入因为它直接复用 MMAD 已经放在 UB 里的结果UB 通路下要先扣掉留给 MMAD 的那部分 UB再计算最大值【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考