CANN/catlass GM到UB复制API CopyGm2Ub【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass代码位置[TOC]功能说明CopyGm2Ub实现 epilogue 阶段从 GM 到 UB 的数据搬运操作。用于将最终输出矩阵的 C/X/Y 矩阵从 Global Memory 搬运到 Unified Buffer供 epilogue 后续计算使用。适用范围AtlasA2、Ascend950风格非 TLA直接操作AscendC::LocalTensor/AscendC::GlobalTensor通过AscendC::DataCopyPad实现带 stride 的数据搬运模板原型template class ArchTag, class GmType // Gemm::GemmTypeElement, Layout struct CopyGm2Ub;模板参数说明ArchTag架构标签Arch::AtlasA2或Arch::Ascend950GmTypeGM 数据类型Gemm::GemmTypeElement, LayoutLayout 触发不同偏特化偏特化实现架构GM LayoutUB Layout说明AtlasA2RowMajorRowMajor二维矩阵搬运DataCopyPadAtlasA2VectorLayoutVectorLayout一维向量搬运Ascend950RowMajorRowMajor二维矩阵搬运DataCopyPadAscend950VectorLayoutVectorLayout一维向量搬运调用接口void operator()( AscendC::LocalTensorElement const dstTensor, // 目的 UB LocalTensor AscendC::GlobalTensorElement const srcTensor, // 源 GM GlobalTensor LayoutDst const layoutDst, // 目的 UB layout 描述 LayoutSrc const layoutSrc // 源 GM layout 描述 )参数说明dstTensor目的 UB LocalTensorsrcTensor源 GM GlobalTensorlayoutDst目的 UB 的 layout包含 shape 和 stridelayoutSrc源 GM 的 layout包含 shape 和 stride调用示例RowMajor二维矩阵#include catlass/epilogue/tile/copy_gm_to_ub.hpp using namespace Catlass::Epilogue::Tile; using Element half; using LayoutTagSrc layout::RowMajor; using GmType Gemm::GemmTypeElement, LayoutTagSrc; uint32_t rows 128; uint32_t cols 256; auto layoutSrc LayoutTagSrc::MakeLayoutElement(rows, cols); auto layoutDst LayoutTagSrc::MakeLayoutElement(rows, cols); AscendC::GlobalTensorElement srcTensor; AscendC::LocalTensorElement dstTensor; using CopyOp CopyGm2UbArch::AtlasA2, GmType; CopyOp copyOp; copyOp(dstTensor, srcTensor, layoutDst, layoutSrc);VectorLayout一维向量using Element half; using LayoutTagSrc layout::VectorLayout; uint32_t length 256; auto layoutSrc LayoutTagSrc::MakeLayoutElement(length, 1); auto layoutDst LayoutTagSrc::MakeLayoutElement(length, 1); AscendC::GlobalTensorElement srcTensor; AscendC::LocalTensorElement dstTensor; using GmType Gemm::GemmTypeElement, LayoutTagSrc; using CopyOp CopyGm2UbArch::AtlasA2, GmType; CopyOp copyOp; copyOp(dstTensor, srcTensor, layoutDst, layoutSrc);【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考