LoadData
产品支持情况
| 产品 | 是否支持 |
|---|---|
| Ascend 950PR/Ascend 950DT | √ |
功能说明
LoadData接口支持Local Memory内部的数据搬运,能够完成L1 Buffer -> L0A Buffer或L1 Buffer -> L0B Buffer的矩阵数据搬运。 该接口通过L1 Buffer与L0A Buffer、L0B Buffer数据排布格式,自动完成是否转置的推导。
函数原型
-
不支持坐标偏移的接口
template<const LoadDataTrait& trait = DEFAULT_LOAD_DATA_TRAIT, typename T, typename U> __aicore__ inline void LoadData(const T& dst, const U& src) -
支持坐标偏移的接口
template<const LoadDataTrait& trait = DEFAULT_LOAD_DATA_TRAIT, typename T, typename U, class Coord> __aicore__ inline void LoadData(const T& dst, const U& src, const Coord& coord)
参数说明
表 1 模板参数说明
| 参数名 | 描述 |
|---|---|
| LoadDataTrait | LoadData相关配置参数结构体,类型为LoadDataTrait: struct LoadDataTrait { __aicore__ constexpr LoadDataTrait() {} __aicore__ constexpr LoadDataTrait(bool transposedIn) : transposed(transposedIn) {} __aicore__ constexpr LoadDataTrait(const LoadDataTrait& trait, bool transposedIn) : transposed(transposedIn) {} bool transposed = false; }; 具体参考表3。DEFAULT_LOAD_DATA_TRAIT为默认取值。 |
| T | 目的操作数的数据类型,通过MakeTensor构造的LocalTensor类型。 |
| U | 源操作数的数据类型,通过MakeTensor构造的LocalTensor类型。 |
| Coord | 坐标偏移的数据类型,通过MakeCoord构造的Coord类型。 |
表 2 参数说明
| 参数名 | 输入/输出 | 含义 |
|---|---|---|
| dst | 输出 | 目的操作数。 |
| src | 输入 | 源操作数。 |
| coord | 输入 | 源操作数上的偏移坐标。 |
| 参数名 | 含义 |
|---|---|
| transposed | 是否启用转置功能,对每个分形矩阵进行转置,默认为false,该参数无需设置,接口内部自动推导,具体参考数据排布支持。 |
数据类型支持
表 4 L1->L0A/L0B数据类型支持
| 源操作数类型 | 目的操作数类型 | 是否支持 |
|---|---|---|
| half | half | √ |
| bfloat16_t | bfloat16_t | √ |
| uint32_t | uint32_t | √ |
| int32_t | int32_t | √ |
| float | float | √ |
| uint8_t | uint8_t | √ |
| int8_t | int8_t | √ |
| fp8_e4m3fn_t | fp8_e4m3fn_t | √ |
| fp8_e5m2_t | fp8_e5m2_t | √ |
| hifloat8_t | hifloat8_t | √ |
| fp4x2_e2m1_t | fp4x2_e2m1_t | √ |
| fp4x2_e1m2_t | fp4x2_e1m2_t | √ |
| uint16_t | uint16_t | √ |
| int16_t | int16_t | √ |
数据排布支持
表 5 L1->L0A数据排布支持
| 源操作数排布 | 目的操作数排布 | 是否转置 |
|---|---|---|
| NZ | NZ | x |
| ZN | NZ | √ |
表 6 L1->L0B数据排布支持
| 源操作数排布 | 目的操作数排布 | 是否转置 |
|---|---|---|
| NZ | ZN | √ |
| ZN | ZN | x |
约束说明
- 地址重叠约束:源操作数与目的操作数位于不同物理位置,不存在地址重叠场景。
- 地址对齐约束:源操作数需要满足32字节对齐,目的操作数需要512字节对齐。
- 参数组合约束:无特殊约束。
- 环境影响约束:无约束。
- 特殊数据类型约束:L1->L0A和L1->L0B通路使能转置时,即ZN2NZ或NZ2ZN,支持的数据类型约束如下:
- b32数据类型要求源矩阵Shape在K轴方向16对齐。
- b8数据类型要求源矩阵Shape在M轴方向32对齐。
- b4数据类型要求源矩阵Shape在M轴方向64对齐。
- 性能约束:无特殊约束。
- 异常和边界值处理:无约束。
- Tensor Layout相关约束:
- Shape、Stride仅支持四维,针对不同的物理存储位置,四个维度的配置均有不同的约束,部分维度为固定值,不可配置。详见层次化表述法。
- Shape、Stride具体维度的数据,仅支持size_t和Std::Int类型。
- LoadData要求L1 Buffer上Layout排布格式为NZ或ZN,L0A Buffer上Layout排布为NZ,L0B Buffer上Layout排布为ZN。
返回值说明
无
流水类型
PIPE_MTE1
调用示例
// 以L1->L0A,half类型为例
using namespace AscendC::Te;
using SrcT = half;
using DstT = half;
size_t m = 128;
size_t k = 64;
auto dstLayout = MakeNZLayout<DstT>(m, k);
auto srcLayout = MakeNZLayout<SrcT>(m, k);
__cbuf__ half l1Addr[256*256*2];
__ca__ half l0AAddr[128*128*2];
auto l1Ptr = MakeL1memPtr(l1Addr);
auto l0APtr = MakeL0AmemPtr(l0AAddr);
auto l1Buff = l1Ptr + 0;
auto l0ABuff = l0APtr + 0;
auto srcTensor = MakeTensor(l1Buff, srcLayout);
auto dstTensor = MakeTensor(l0ABuff, dstLayout);
// 调用方式1
Copy(CopyAtom<CopyTraits<CopyL12L0, LoadDataTraitDefault>>{}, dstTensor, srcTensor);
// 调用方式2
CopyAtom<CopyTraits<CopyL12L0, LoadDataTraitDefault>>{}.Call(dstTensor, srcTensor);
// 调用方式3
auto atomCopyL1A2L0A = MakeCopy(CopyL12L0{}, LoadDataTraitDefault{});
atomCopyL1A2L0A.Call(dstTensor, srcTensor);
完整样例请参考TensorAPI样例代码。