CANN/ops-tensor Block Epilogue 基础框架

发布时间:2026/5/21 4:12:32

CANN/ops-tensor Block Epilogue 基础框架 Block Epilogue 基础框架【免费下载链接】ops-tensorops-tensor 是 CANN Compute Architecture for Neural Networks算子库中提供张量类计算的基础算子库采用模块化设计支持灵活的算子开发和管理。项目地址: https://gitcode.com/cann/ops-tensor公共接口说明概述Block 层后处理组件用于矩阵乘计算后的额外处理。不同实现Empty、StreamK提供不同功能Empty 为空实现StreamK 支持 workspace 汇聚、类型转换、ReLU 激活。详见README.md 查看 API 清单和实现对比。类概述类型别名类型说明BlockShapeBlock 形状类型AscendC::Te::Shapeint64_t, int64_t, int64_t, int64_tBlockCoordBlock 坐标类型AscendC::Te::Coordint64_t, int64_t, int64_t, int64_t核心数据结构Argumentsstruct Arguments { // 具体成员根据实现不同 GM_ADDR cGmAddr; // C 矩阵 GM 地址可选 GM_ADDR workspaceGmAddr; // Workspace 地址可选 };说明Host 端参数结构体传递给 Kernel。Paramsstruct Params { // 具体成员根据实现不同 // 通常与 Arguments 相同或包含更多运行时参数 };说明Kernel 运行时参数结构体。核心成员方法构造函数__aicore__ inline BlockEpilogue()功能构造 BlockEpilogue 对象。析构函数__aicore__ inline ~BlockEpilogue()功能析构 BlockEpilogue 对象。Init函数__aicore__ inline void Init(Params const params, ...)功能初始化 BlockEpilogue 组件设置 GM 地址、Block 形状等参数。 说明参数数量和类型根据实现不同。Run函数__aicore__ inline void Run()功能执行后处理操作。 说明Empty 实现为空直接返回StreamK 实现包含实际计算。operator函数参数版本__aicore__ inline void operator()(Params const params)功能执行后处理操作参数版本。 说明调用Init和Run提供统一调用接口。operator函数Block版本__aicore__ inline void operator()( BlockShape const blockShape, BlockCoord const blockCoord, int64_t dstStartOffset 0, int64_t srcStartOffset 0)功能执行后处理操作Block 坐标版本。 参数说明 | 参数 | 类型 | 说明 | |------|------|------| | blockShape | BlockShape const | Block 形状 | | blockCoord | BlockCoord const | Block 坐标 | | dstStartOffset | int64_t | 目标起始偏移默认 0 | | srcStartOffset | int64_t | 源起始偏移默认 0 |公共约束模板参数要求Empty 无模板参数要求StreamK 需要 WorkspaceType、OutType、DispatchPolicy 参数计算位置Empty无计算空实现StreamK在 AIV 核执行接口一致性所有实现必须提供Init、Run、operator方法保持统一的类型别名公共调用示例组件组装模板// Empty 实现 using BlockEpilogue Blaze::Gemm::Block::BlockEpilogueEmpty; // StreamK 实现 using BlockEpilogue Blaze::Gemm::Block::BlockEpilogueStreamKfloat, half, DispatchPolicy;在 Kernel 中使用using MatmulKernel Blaze::Gemm::Kernel::KernelMatmulBasic ProblemShape, BlockMmad, BlockEpilogue, BlockScheduler; // 或 using MatmulKernel Blaze::Gemm::Kernel::KernelMatmulStreamK ProblemShape, BlockMmad, BlockEpilogue, BlockScheduler;执行模板// Empty无实际效果 BlockEpilogue epilogue; epilogue.Run(); epilogue({}); epilogue(blockShape, blockCoord, 0, 0); // StreamK实际后处理 BlockEpilogue epilogue; epilogue.Init(params, blockShape, tileL1, coord, usedCoreNum, isSkScene); epilogue.Run();设计说明Epilogue 层的作用Epilogue 层位于矩阵乘计算之后用于结果汇聚将 workspace 中 K 轴切分的中间结果累加类型转换将 float 结果转换为 half 或 bfloat16激活函数执行 ReLU 等激活操作其他后处理如量化、缩放等未来扩展为什么采用空实现设计模板参数要求Kernel 模板需要 BlockEpilogue 参数Empty 满足此要求零开销空实现不引入任何计算或内存开销类型安全保持类型系统完整性编译期检查扩展性可无缝替换为实际实现无需修改 Kernel 代码性能影响Epilogue 类型编译时间运行时间内存占用流水线BlockEpilogueEmpty极小无影响无不参与BlockEpilogueStreamK正常有影响有参与计算【免费下载链接】ops-tensorops-tensor 是 CANN Compute Architecture for Neural Networks算子库中提供张量类计算的基础算子库采用模块化设计支持灵活的算子开发和管理。项目地址: https://gitcode.com/cann/ops-tensor创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

相关新闻