Ascend C 算子实现简介:ScatterAdd 算子
本文主要对 ScatterAdd 算子进行简要介绍。
ScatterAdd 用于在给定的索引位置上增加一个指定的值。它的用法类似于 ScatterUpdate 操作,但不同的是,ScatterUpdate 是在指定索引位置上直接覆盖原来的值,而 ScatterAdd 则是在原来的值上执行加法操作。
参考 TensorFlow 的用法:
tf.raw_ops.ScatterAdd(ref, indices, updates, use_locking=False, name=None
)
其中,ref 是要进行更新操作的张量,indices 是要更新的位置的索引,updates 是要加上的值。use_locking 参数表示在更新时是否使用锁来确保同一时间只有一个操作在更新变量的值。name 参数为操作的名称。
实现思路:
在 Ascend C 上实现 ScatterAdd 算子的过程中,有几个核心问题需要解决:
-
如何优化最大 Tiling?
对于大规模数据的处理,需确保尽可能高效地利用硬件资源。 -
数据类型的处理
特别是对于 Half 类型的处理,如何保证精度与性能的平衡。
关键代码实现:
- 最大 Tiling 计算
uint32_t firstTiling = lastDim;if (lastDim * sizeOfDataType % BLOCK_SIZE == 0) {uint32_t bufferSize = 6;if (tilingKey == 4) {bufferSize = 10;}if (ub_size > bufferSize * lastDim * sizeOfDataType) {firstTiling = lastDim;} else {while (ub_size < bufferSize * firstTiling * sizeOfDataType) {firstTiling -= 32 / sizeOfDataType;}}tilingKey += 10;
}
通过这种方式,我们确保了在硬件上能够高效处理大规模数据,减少不必要的计算浪费。
- Half 类型的标量加法
// half转float
int16_t varInt16 = *(int16_t*)&var;
int32_t fltInt32 = ((varInt16 & 0x8000) << 16);
fltInt32 |= ((varInt16 & 0x7fff) << 13) + 0x38000000;
float varFp32 = *(float*)&fltInt32;int16_t updateInt16 = *(int16_t*)&update;
fltInt32 = ((updateInt16 & 0x8000) << 16);
fltInt32 |= ((updateInt16 & 0x7fff) << 13) + 0x38000000;
float updateFp32 = *(float*)&fltInt32;float res = varFp32 + updateFp32;// float 转 half
int16_t fltInt16;
fltInt32 = *(int32_t*)&res;
fltInt16 = ((fltInt32 & 0x7fffffff) >> 13) - (0x38000000 >> 13);
fltInt16 |= ((fltInt32 & 0x80000000) >> 16);half resHalf = *(half*)&fltInt16;
通过上述代码,我们确保了对于 Half 类型的加法操作能够准确高效地完成,保证了在硬件加速场景下的性能表现。
以上就是 ScatterAdd 算子的简单实现介绍,核心在于对齐操作、最大 Tiling 的优化以及特殊数据类型(如 Half)的处理。这些技术点在 Ascend C 硬件架构上被充分利用,实现了高效的算子执行。