昇腾AI原生创新算子挑战赛(S1赛季)ScatterSub算子实现简介
收藏回复举报
昇腾AI原生创新算子挑战赛(S1赛季)ScatterSub算子实现简介
新人帖
发表于2024-06-07 10:27:33
0 查看

本文主要是对ScatterSub算子进行简单的介绍。

ScatterSub 用于在给定的索引位置上减去一个指定的值。它的用法类似于 ScatterUpdate 操作,不同之处在于 ScatterUpdate 是在指定索引位置上直接覆盖原来的值,而 ScatterSub 是对原来的值进行减法操作。

%E5%9B%BE%E7%89%871.jpg

参考tensorflow的用法:

tf.raw_ops.ScatterSub(
    ref, indices, updates, use_locking=False, name=None
)

其中,x 是要进行更新操作的张量,indices 是要更新的位置的索引,updates 是要减去的值。use_locking 参数表示在更新时是否使用锁来确保同一时间只有一个操作在更新变量的值。name 参数为操作的名称。

ScatterSub实现思路:

对齐(矢量计算)

非对齐(标量计算)

大循环

Indices元素数量

小循环

X[1:]按照最大Tiling切分

X[1:]的元素

根据实现思路,在Ascend C实现过程中,有两个问题需要进行解决:

1、最大Tiling如何计算?

2Half如何实现标量减法?

实现方案关键代码:

1、最大Tiling计算

// 最后一维是对其的情况处理
uint32_t firstTiling = lastDim;
if (lastDim * sizeOfDataType % BLOCK_SIZE == 0) {
    // bufferSize是初始化的Buffer数量
    uint32_t bufferSize = 6;
    
    // Int8类型比其他类型多了两块
    if (tilingKey == 4) {
        bufferSize = 10;
    }

    // Ub大小如果大于Buffer区域大小则直接按照lastDim分块
    // 否则,则按照32B依次减小,知道满足Ub大小
    if (ub_size > bufferSize * lastDim * sizeOfDataType) {
        firstTiling = lastDim;
    } else {
        while(ub_size < bufferSize * firstTiling * sizeOfDataType) {
            firstTiling -= 32 / sizeOfDataType;
        }
    }
    tilingKey += 10;
}

2Half类型标量减法

%E5%9B%BE%E7%89%872.jpg

%E5%9B%BE%E7%89%873.jpg

// 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;

算子详细实现请参考

https://gitee.com/ascend/samples/tree/master/operator_contrib/ScatterSubSample

我要发帖子