L1 Buffer到L0A Buffer数据搬运

【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言，原生支持C和C++标准规范，主要由类库和语言扩展层构成，提供多层级API，满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit

产品支持情况

• Ascend 950PR/Ascend 950DT：支持

• Atlas A3 训练系列产品/Atlas A3 推理系列产品：不支持

• Atlas A2 训练系列产品/Atlas A2 推理系列产品：不支持

• Atlas 200I/500 A2 推理产品：不支持

• Atlas 推理系列产品AI Core：不支持

• Atlas 推理系列产品Vector Core：不支持

• Atlas 训练系列产品：不支持

功能说明

头文件路径为：tensor_api/tensor.h。

Tensor API通过Copy接口统一执行不同通路数据搬运。该接口用于将L1 Buffer中的左矩阵数据搬运到L0A Buffer。

该通路以512Byte数据分形为基本搬运单位。不同数据类型对应的分形矩阵形态如下：

接口支持非转置搬运和转置搬运。源张量和目的张量的Layout需要匹配当前通路支持的格式组合。

接口支持Batch模式。Batch模式下，源张量和目的张量的Layout需要在原有分形Layout最前面增加Batch维，Shape形态为(B, 单矩阵Shape)，其中B表示Batch数量。用户可使用MakeFrameLayout<NZLayoutPtn, DataType>(B, m, k)或MakeFrameLayout<ZNLayoutPtn, DataType>(B, m, k)构造带Batch维的Layout。

函数原型

• 执行L1 Buffer到L0A Buffer的数据搬运。

  template <typename AtomType, typename DstTensor, typename SrcTensor> __aicore__ inline void Copy(const CopyAtom<AtomType>& atomCopy, const DstTensor& dst, const SrcTensor& src)

• 构造默认搬运原子对象。

  template <typename CopyOperationType> __aicore__ inline constexpr auto MakeCopy(const CopyOperationType& copyOperation)

• 构造指定Trait的搬运原子对象。

  template <typename CopyOperationType, typename CopyTraitType> __aicore__ inline constexpr auto MakeCopy( const CopyOperationType& copyOperation, const CopyTraitType& copyTrait)

参数说明

表1Copy接口参数说明

表2MakeCopy接口参数说明

数据类型

支持的数据类型包括：

fp4x2_e2m1_t、fp4x2_e1m2_t、int8_t、uint8_t、hifloat8_t、fp8_e5m2_t、fp8_e4m3fn_t、half、bfloat16_t、int16_t、uint16_t、int32_t、uint32_t、float。

源张量和目的张量的数据类型需要保持一致。

返回值说明

Copy无返回值。MakeCopy返回CopyAtom对象。

约束说明

• 目的地址位于L0A Buffer时，起始地址需要满足512Byte对齐要求。
• 源地址位于L1 Buffer时，起始地址需要满足32Byte对齐要求。
• 特殊数据类型约束：L1 Buffer到L0A Buffer通路使能转置时，即ZN2NZ，支持的数据类型约束如下：
  • b32数据类型要求源矩阵Shape在K轴方向16对齐。
  • b8数据类型要求源矩阵Shape在M轴方向32对齐。
  • b4数据类型要求源矩阵Shape在M轴方向64对齐。
• Batch模式要求源张量和目的张量的Batch数量一致，且在L1 Buffer和L0A Buffer上张量数据连续排布。

关键特性说明

非转置搬运

当源张量和目的张量数据格式为NZ时，接口按非转置方式从L1 Buffer搬运到L0A Buffer。

图1非转置搬运示意图

转置搬运

当源张量数据格式为ZN、目的张量数据格式为NZ时，接口按转置方式搬运。转置搬运会对源张量中每个分形矩阵进行转置。

不同数据位宽的转置搬运示意如下。

b4转置搬运

图2b4转置搬运示意图

b8转置搬运

图3b8转置搬运示意图

b16转置搬运

图4b16转置搬运示意图

b32转置搬运

图5b32转置搬运示意图

Batch搬运

当源张量和目的张量均使用带Batch维的Layout时，Copy接口会按照Batch维描述的数据范围完成L1 Buffer到L0A Buffer的数据搬运。非转置搬运时源张量、目的张量均为NZ格式，转置搬运时源张量为ZN格式，目的张量为NZ格式。

调用示例

#include "tensor_api/tensor.h" using namespace AscendC::Te; __aicore__ inline void CopyL1ToL0AExample() { using AType = half; constexpr uint32_t m = 16; constexpr uint32_t k = 16; __cbuf__ AType l1Buf[m * k]; __ca__ AType l0aBuf[m * k]; auto l1Tensor = MakeTensor(MakeMemPtr(l1Buf), MakeFrameLayout<NZLayoutPtn, AType>(m, k)); auto l0aTensor = MakeTensor(MakeMemPtr(l0aBuf), MakeFrameLayout<NZLayoutPtn, AType>(m, k)); auto copyAtom = MakeCopy(CopyL12L0A{}, CopyL12L0ATraitDefault{}); Copy(copyAtom, l0aTensor, l1Tensor); }

Batch模式示例如下。

#include "tensor_api/tensor.h" using namespace AscendC::Te; __aicore__ inline void CopyL1ToL0ABatchExample() { using AType = half; constexpr uint32_t batch = 2; constexpr uint32_t m = 16; constexpr uint32_t k = 16; __cbuf__ AType l1Buf[batch * m * k]; __ca__ AType l0aBuf[batch * m * k]; auto layout = MakeFrameLayout<NZLayoutPtn, AType>(batch, m, k); auto l1Tensor = MakeTensor(MakeMemPtr(l1Buf), layout); auto l0aTensor = MakeTensor(MakeMemPtr(l0aBuf), layout); auto copyAtom = MakeCopy(CopyL12L0A{}, CopyL12L0ATraitDefault{}); Copy(copyAtom, l0aTensor, l1Tensor); }

【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言，原生支持C和C++标准规范，主要由类库和语言扩展层构成，提供多层级API，满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit