三亩地 三亩地SAN MU DI · CODE DIARY
ARTICLE DETAIL

日记详情

真实记录编程学习的某一天,欢迎挑你感兴趣的翻一翻。

CANN/catlass GEMV AIV搬运模板

CANN/catlass GEMV AIV搬运模板

TileCopyGemvAiv

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

代码位置

[TOC]

功能说明

TileCopyGemvAiv为 AIV(AI Vector)芯片提供 GEMV 搬运子组件的聚合模板。数据通路为 GM↔UB↔GM,引用 GEMV 专有的搬运组件。

  • 适用范围:AtlasA2
  • 不直接执行算子,以类型成员方式暴露子组件引用

模板原型

template <class ArchTag, class AType, class XType, class YType, class BiasType = void> struct TileCopyGemvAiv;
模板参数说明
ArchTag架构标签
ATypeA 矩阵类型GemmType<ElementA, RowMajor/ColumnMajor>
XTypeX 向量类型GemmType<ElementX, VectorLayout>
YTypeY 向量类型GemmType<ElementY, VectorLayout>
BiasType偏置类型,默认void

成员类型定义

成员类型对应子组件说明
VecCopyGmToUbGemv::Tile::VecCopyGmToUB向量 X: GM→UB
VecCopyUbToGmGemv::Tile::VecCopyUBToGm向量 Y: UB→GM(可选 atomic add)
MatrixCopyGmToUbGemv::Tile::MatrixCopyGmToUB矩阵 A: GM→UB

调用示例

#include "catlass/gemv/tile/tile_copy.hpp" using namespace Catlass::Gemv::Tile; using ElementA = half; using ElementX = half; using ElementY = half; using AType = Gemm::GemmType<ElementA, layout::RowMajor>; using XType = Gemm::GemmType<ElementX, layout::VectorLayout>; using YType = Gemm::GemmType<ElementY, layout::VectorLayout>; using Copy = TileCopyGemvAiv<Arch::AtlasA2, AType, XType, YType>; // 子组件: // typename Copy::VecCopyGmToUb // typename Copy::VecCopyUbToGm // typename Copy::MatrixCopyGmToUb

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

创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

← 返回列表