当前位置：首页 > article >正文

ScatterAdd算子实现简介

article 2024/12/29 9:55:49

在这里插入图片描述

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;