CATLASS算子UB到GM搬运模块

copy_ub_to_gm

【免费下载链接】catlass 本项目是CANN的算子模板库，提供NPU上高性能矩阵乘及其相关融合类算子模板样例。 项目地址: https://gitcode.com/cann/catlass

代码位置

[TOC]

概述

copy_ub_to_gm 模块实现 epilogue 阶段从 UB（Unified Buffer）到 GM（Global Memory）的数据搬运，将 epilogue 计算结果写回 Global Memory。

TLA 风格的 UB→GM 搬运为独立模块 CopyUb2GmTla。

API 清单

适用硬件

调用示例

CopyUb2Gm

#include "catlass/epilogue/tile/copy_ub_to_gm.hpp"

using namespace Catlass::Epilogue::Tile;

using Element = half;
using LayoutTagDst = layout::RowMajor;

uint32_t rows = 128;
uint32_t cols = 256;

auto layoutSrc = LayoutTagDst::MakeLayout<Element>(rows, cols);
auto layoutDst = LayoutTagDst::MakeLayout<Element>(rows, cols);

AscendC::LocalTensor<Element> srcTensor;
AscendC::GlobalTensor<Element> dstTensor;

using GmType = Gemm::GemmType<Element, LayoutTagDst>;
using CopyOp = CopyUb2Gm<Arch::AtlasA2, GmType>;
CopyOp copyOp;
copyOp(dstTensor, srcTensor, layoutDst, layoutSrc);

【免费下载链接】catlass 本项目是CANN的算子模板库，提供NPU上高性能矩阵乘及其相关融合类算子模板样例。 项目地址: https://gitcode.com/cann/catlass