Scatter
产品支持情况
|
产品 |
是否支持 |
|---|---|
|
Atlas 350 加速卡 |
√ |
|
|
x |
|
|
x |
|
|
x |
|
|
x |
|
|
x |
|
|
x |
功能说明
给定源操作数的寄存器张量和索引张量,以及结果操作数在UB中的基地址,Scatter指令将源操作数按元素根据索引位置分散到UB中。分散过程如下图所示:

定义原型
template <typename T = DefaultType, typename U = DefaultType, typename S, typename V> __simd_callee__ inline void Scatter(__ubuf__ T* baseAddr, S& srcReg, V& index, MaskReg& mask)
参数说明
|
参数名 |
描述 |
|---|---|
|
T |
目的操作数和源操作数的数据类型。 |
|
U |
索引的数据类型。 |
|
S |
源操作数的RegTensor类型,例如RegTensor<half>,由编译器自动推导,用户不需要填写。 |
|
V |
索引值的RegTensor类型,例如RegTensor<uint16_t>,由编译器自动推导,用户不需要填写。 |
|
参数名 |
输入/输出 |
描述 |
|---|---|---|
|
baseAddr |
输出 |
目的操作数在UB中的基地址。 类型为UB指针。 Atlas 350 加速卡,支持的数据类型详见表3。 |
|
srcReg |
输入 |
源操作数。 类型为RegTensor。 Atlas 350 加速卡,支持的数据类型详见表3。 |
|
index |
输入 |
srcReg中的每个元素在UB中相对于baseAddr的索引位置。索引值要大于等于0。 类型为RegTensor。 IndexT数据类型需要与目的操作数和源操作数的数据类型T配套使用。类型配套对应表详见约束说明。 Atlas 350 加速卡,支持的数据类型详见表3。 |
|
mask |
输入 |
src element操作有效指示,详细说明请参考MaskReg。 |
约束说明
- 目的操作数和源操作数的数据类型T和U数据类型需要配套使用。类型配套对应表如下:
- 当T为b64数据类型时,T,U,V数据类型只支持以下组合:
T数据类型
IndexT数据类型
RegT数据类型
RegIndexT数据类型
备注
b64
uint32_t
RegTensor<uint64_t, RegTraitNumOne>
RegTensor<uint32_t>
index前32个数有效
RegTensor<int64_t, RegTraitNumOne>
RegTensor<uint64_t, RegTraitNumTwo>
RegTensor<uint32_t>
-
RegTensor<int64_t, RegTraitNumTwo>
b64
uint64_t
RegTensor<uint64_t, RegTraitNumOne>
RegTensor<uint64_t, RegTraitNumOne>
-
RegTensor<int64_t, RegTraitNumOne>
RegTensor<uint64_t, RegTraitNumOne>
RegTensor<uint64_t, RegTraitNumTwo>
index前32个数有效
RegTensor<int64_t, RegTraitNumOne>
RegTensor<uint64_t,RegTraitNumTwo>
RegTensor<uint64_t, RegTraitNumTwo>
-
RegTensor<int64_t,RegTraitNumTwo>
- 当T为int8或者uint8数据类型时,源操作数Tensor中仅偶数位Byte有效。最终存入UB结果地址的数据仅为源操作数Tensor偶数位Byte,即srcReg中的第0,2,4,..., 252,254位置数据会被分散存储到目的操作数中。
- index中不能有相同的值。若有2个或者2个以上的index中的数据相同,则只有其中一个值对应的数据是有效的,具体哪一个数据可能是未知的。
- 索引位置需要按照dtype_size(单个元素的Byte数)对齐,否则可能会造成数据分散结果错乱。
调用示例
template<typename T, typename U>
__simd_vf__ inline void ScatterVF(__ubuf__ T* dstAddr, __ubuf__ T* src0Addr, __ubuf__ U* src1Addr, uint32_t count, uint32_t oneRepeatSize, uint16_t repeatTimes)
{
AscendC::Reg::RegTensor<T> srcReg0;
AscendC::Reg::RegTensor<U> srcReg1;
AscendC::Reg::MaskReg mask;
for (uint16_t i = 0; i < repeatTimes; i++) {
mask = AscendC::Reg::UpdateMask<T>(count);
AscendC::Reg::LoadAlign(srcReg0, src0Addr + i * oneRepeatSize);
AscendC::Reg::LoadAlign(srcReg1, src1Addr + i * oneRepeatSize);
AscendC::Reg::Scatter(dstAddr, srcReg0, srcReg1, mask);
}
}