CATLASS EVG设计概览

EVG 设计概览【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlassEVGEpilogue Visitor Graph是 CATLASS 在 GEMM 尾处理阶段使用的声明式拼装框架。它把“读数据、做逐元素计算、写结果”拆成一组可组合的Visitor节点再由图结构把这些节点连接起来用统一的三阶段执行模型完成尾处理流水。当前实现主要分布在以下位置include/catlass/epilogue/fusion/EVG 图组织与节点实现include/catlass/epilogue/block/block_epilogue_visitor.hppBlock 级执行器include/catlass/gemm/kernel/basic_matmul_tla_visitor.hppGM workspace 路径include/catlass/gemm/kernel/basic_matmul_tla_ub_visitor.hppUB workspace 路径设计目标EVG 要解决的是“尾处理逻辑经常变化但搬运、切块、同步、双缓冲这些代码不想每次重写”的问题。它的目标不是替代 GEMM 主循环而是把尾处理部分从手工组织事件与 UB 空间变成按表达式声明。以D C X为例开发者只需要描述从 GEMM 结果里取C从外部输入里取X做逐元素Add把结果写回D至于 tile 切分、何时搬入、何时计算、何时写回由 EVG 所在的 Block 与 Kernel 组件统一调度。分层关系EVG 在当前代码中的职责分层可以概括为三层。Kernel 层Kernel 层负责把 GEMM 主循环和 EVG 尾处理接起来当前真实入口有两条BasicMatmulTlaVisitor代码路径include/catlass/gemm/kernel/basic_matmul_tla_visitor.hppAIC 先把 MMAD 结果写到 GM workspaceAIV 再从 GM workspace 读出并执行 EVGBasicMatmulTlaUbVisitor代码路径include/catlass/gemm/kernel/basic_matmul_tla_ub_visitor.hppAIC 把结果保留在 UBAIV 直接消费 UB 中的数据并执行 EVG两条路径都会把EVG::Arguments透传到BlockEpilogue但 workspace 组织方式不同GM workspace 路径workspace M * N * sizeof(C)EVG自身 workspaceUB workspace 路径workspace 只为EVG自身保留不再额外申请整块CBlock 层Block 级的执行器是BlockEpilogueEpilogueVisitor..., ArchTag, ComputeLength, EVG, ElementC代码路径include/catlass/epilogue/block/block_epilogue_visitor.hpp。它负责把一个 block 的输出切成更小的 tile为 EVG 分配两套 callback形成双缓冲按LOAD - COMPUTE - STORE次序驱动每个 tile用事件同步串起 MTE2、V、MTE3 三条流水其中有两个重要模板参数EpilogueVisitorfalseGM workspace 路径EpilogueVisitortrueUB workspace 路径ComputeLength决定单次在 UB 内处理多少元素它既影响 tile 大小也直接影响 EVG 节点能够分到的 UB 空间。Fusion 层Fusion 层负责“图怎么描述”和“节点怎么执行”。当前代码提供两种组织方式TreeVisitor代码路径include/catlass/epilogue/fusion/tree_visitor.hpp适合树状表达式先访问子节点再调用父节点TopologicalVisitor代码路径include/catlass/epilogue/fusion/topological_visitor.hpp适合 DAG允许中间结果被多个节点复用图中的基础节点来自visitor_*.hpp例如VisitorAccLoadVisitorAuxLoadVisitorComputeVisitorCastVisitorAuxStoreVisitorRowBroadcast三阶段执行模型EVG 的节点执行统一遵循VisitStage代码路径include/catlass/epilogue/fusion/visitor_impl_base.hpp定义的三阶段LOAD把需要的输入搬入 UBCOMPUTE在 UB 中完成逐元素计算、广播或类型转换STORE把结果写回 GM或者执行真正修改外部输出地址内容的动作这种分阶段设计的意义有两点节点职责清晰。每个节点只需声明自己在哪个阶段做什么。Block 层可以统一组织双缓冲流水而不需要关心节点内部具体做了什么计算。当前BlockEpilogue的流水模型是等待上一轮释放可读 buffer执行当前 tile 的LOAD等待输入与输出依赖满足执行当前 tile 的COMPUTE等待计算结果可写回执行当前 tile 的STORE交替使用两套 callback形成双缓冲双缓冲流水在 tile 级别的核心时序可以按下面理解三阶段和常见节点的关系可以按下面理解图组织方式TreeVisitorTreeVisitorNodeOp, ChildOps...适合表达“结果由若干输入直接合成”的场景例如D C XD silu(C)D cast(add(C, X))它的特点是结构直观和表达式语义接近子节点输出会按顺序传给父节点适合链式或树状组合TopologicalVisitorTopologicalVisitorEdgeTuple, Ops...适合中间结果需要复用的场景例如exp(2x)同时被分子和分母使用一个中间节点被多个后续节点消费它的特点是节点按拓扑序平铺定义最后一个节点视为根节点每次访问都从根节点递归回溯依赖并用缓存避免同一个 tile、同一个阶段里重复计算已访问节点当前仓内matmul_tanh_evg就使用了这种方式。和TreeVisitor相比它更适合显式复用中间结果空间与切块EVG 的资源管理围绕两件事展开ComputeLength每个节点的get_callbacks(...)执行时Block 层会把当前 tile 对齐后交给各节点申请 UB 空间。像VisitorAccLoad、VisitorAuxLoad、VisitorCompute、VisitorCast这类节点会占用 UBVisitorAuxStore主要负责写回不额外分配计算 buffer。因此ComputeLength不能只看结果 tile 大小还要考虑本条 EVG 链路里会同时驻留多少块 UB 数据是否启用了双缓冲当前路径是 GM workspace 还是 UB workspace当前实现边界这篇文档只保留设计层面的边界说明。具体到当前有哪些节点、算子、kernel 入口和样例时以 evg_api 为准fusion目录内文件分布见 fusion/README。【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考