CANN/catlass EVG UB工作空间kernel入口
Basic Matmul TLA UB Visitor【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass代码路径include/catlass/gemm/kernel/basic_matmul_tla_ub_visitor.hpp功能说明这是当前 EVG 在 UB workspace 路径上的 kernel 入口。执行方式为AIC 在 UB 中完成 MMAD通过跨核同步把时序交给 AIVAIV 直接消费 UB 中的结果并执行 EVG相比 GM workspace 路径它省掉了把 MMAD 结果整体写回 GM 再读出的过程。模板参数template class BlockMmad_, class BlockEpilogue_, class BlockScheduler_ class BasicMatmulTlaUbVisitor;常见搭配BlockEpilogue_使用EpilogueVisitortrueVisitorAccLoad也打开 UB workspace 模式Arguments 关键字段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; };同样需要注意当前 visitor kernel 的ToUnderlyingArguments()实现并不消费ptrC/layoutC最终输出位置由evg_args中的VisitorAuxStore决定。Workspace 规则GetWorkspaceSize()只返回EVG::get_workspace_size(...)因为 MMAD 结果不再额外落到 GM workspace。适用条件当前实现限制为Arch::Ascend950BlockEpilogue::USE_UB_WORKSPACE对应true更适合希望减少中间结果回写开销的场景【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考