昇腾算子挑战赛S5的WriteUp。本次是抽空参加，就只重点完成了Lcm，因为可以直接在上一次决赛的GCD算子上面修改，这一题也获得了本次的单题第一，其他两题只是在前十，做的一般。

1 BitwiseLeftShift

  • 对每个输入input[i]，左移other[i]位。

  • 输入Tensor Input和Other可能维度不对应，需要广播。

  • 最多三个维度

1.1 任意维度的广播支持

广播实际上是因为other在某个维度上只有1个单位，因此只能取该单位中的值的场景。例如对于输入数据input(14，10，30，9)，other(14，1，30，0)，other的维度小于输入数据，最终输出维度与input的维度相同，则需要操作的元素数是input的每一个元素，并且在other中缺失的维度只能取0，例如input的(3，3，3，3)，对应other的数据是(3，0，3，0)。由此可以根据输入数据的坐标以及总体维度信息，得到other的坐标，取值进行计算。

对于BitwiseLeftShift，input的维度一定大于等于other的维度，所以输出形状与input相对应，可以直接根据input的元素数量进行多核数据划分。但这样会将所有数据展开至一维划分，失去维度上的含义，每个核在初始化阶段只知道当前元素在所有元素中的偏移，所以在计算过程中，需要根据偏移计算每个操作的元素在input中的坐标，其每一维坐标为：

其中Dim存储的是每一个维度的大小，而S[i]是包含$i$维度在内的维度大小后缀和。根据上式可计算出每一维在原input Tensor的坐标，然后可以根据此坐标得到Other中的对应元素坐标，取出该元素计算即可，就实现了广播。

1.2 单维度广播优化

如果只有一个维度需要广播，可以将other的数据也通过向量化搬运的方式拷入进行batch计算。例如对于input（256，30，99），other数据（256，1，99），可以按照batch进行划分，即将256*30个组计算分配给40个核，每个核内则需要拷入分配到的batch数量组input进行计算，每轮input的size是99，而other的每轮数据则需要根据当前计算的batchID来拷入。（更严谨来说，应该考虑other数据是否能够一次拷入UB Buffer，如果不能，还需要再进行划分，多次拷入处理）。通过这样的优化，可以将单维度广播这种特殊情况的访存开销和索引计算开销降低。

为了进一步优化搬运，可以一次搬运多个batch来进行计算。对于单维度广播，维度总是可以抽象为三部分，以输入数据为input（256，30，99），other（256，1，99），中间维度需要广播为例：

  • batch部分，是截至广播维度的维度大小乘积，例如对于上述提到的（256，30，99），共有256*30个batch
  • broadCastDim，需要广播的维度大小，这个大小表示每连续的broadCastDim组数据，使用相同的数据进行计算，例如上述提到的（256，30，99），broadCastDim为30，因为other的中间维度是1，所以对于输入的每连续30组数据，使用的是相同的99个元素
  • batchDataSize，每组计算的数据量，对于输入张量（256，30，99）为99

每个核可以知道自己的起始batchID，最后一个batchID，而对应的otherBatchID为batchIdx / broadCastDim，由于起始batchID可能不是broadCastDim的整数倍，所以对于一组other数据，要重用的次数是(otherBatchId + 1) * broadCastDim - batchIdx，再与剩余batch数量取较小的那一个。

为了减少搬运次数，增加搬运量，可以一次搬多行。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
__aicore__ inline void Process1() {
    for (int batchIdx = startBatch; batchIdx < endBatch;) {
        uint32_t otherBatchId = batchIdx / broadCastDim;
        uint32_t nSameOther = (otherBatchId + 1) * broadCastDim - batchIdx; //当前other要用多少次
        nSameOther = min(endBatch - batchIdx, nSameOther);                  //不能超过endBatch
        CopyInOther(otherBatchId);
        auto otherLocal = otherQue.DeQue<DTYPE_INPUT>();
        int i = 0;
        for (i = 0; i+nPerCopy <= nSameOther; i+=nPerCopy) {
            CopyIn(batchIdx + i, nPerCopy);
            auto inputLocal = inputQue.DeQue<DTYPE_INPUT>();
            //一次读入nPerCopy个batch组数据
            for(int j=0;j<nPerCopy;j++) {
                Compute(inputLocal, otherLocal, j*batchDataSize);
                CopyOut(batchIdx + i + j);
            }
            inputQue.FreeTensor(inputLocal);
        }
        for (; i < nSameOther; i++) {
            CopyIn(batchIdx + i, 1);
            auto inputLocal = inputQue.DeQue<DTYPE_INPUT>();
            Compute(inputLocal, otherLocal, 0);
            inputQue.FreeTensor(inputLocal);
            CopyOut(batchIdx + i);
        }
        batchIdx += nSameOther; //更新batchIdx
        otherQue.FreeTensor(otherLocal);
    }
}

1.3 向量化计算

对于不需要广播的情况，首先仍然可以将所有维度展至一维进行数据的划分。输入、输出全部都可以通过DataCopy直接搬入，只需要按照UB buffer的大小，逐个迭代完成单核内所有分配到的数据的计算。

在计算方面，ai core支持int16，int32类型的左移指令，因此对于这两种数据类型，可以实现向量化计算，而另外两种类型只能标量计算。以int32为例，移位数依次为32、16、8、4、2、1，最多6次移位就可以完成所有输入元素的移位，每次移位前后要进行Compare，判断当前移位位数是否小于剩余待移位位数，通过Compare后的bitmask选择移位结果或之前的结果。

由于Compare指令操作类型只有half/float，所以需要预先将Other输入Cast到half或float，而Select指令也只支持half和float，所以要将临时结果和移位结果的Tensor Reinterpret得到一个half或float的指针进行操作，由于Select操作只进行数据选择，所以int16和int32与half或float进行的操作是完全相同的。由于向量化需要进行Cast，指针重解释等操作，模板推导可能会有问题，所以最好不使用模板，对int16和int32的数据编写两个kernel处理。

2 CopySign

与BitwiseLeftShift基本相同。

3 Lcm

  • 求两个输入的最小公倍数

  • 输入Tensor Input和Other可能维度不对应，需要广播。

  • 最多三个维度

3.1 形状与数据类型

需要广播的场景与题目一BitwiseLeftShift完全一致，类似题目一：

  • 任意维度广播，全标量处理
  • 单维度广播，向量数据搬运
  • 不需要广播

而根据不同的数据类型，计算又可以分三种情况：

  • int8 int16：向量数据搬运(展开至一维进行数据划分)
  • int32：cast至float进行向量化计算
  • int64：stein+gcd结合的标量+向量化计算得到公因数，再标量计算lcm。

3.2 向量化计算优化

对于int32类型的数据，可以cast至float进行向量化计算，通过CompareScalar、Select、Fmod进行向量化的Gcd计算。

对于int64类型的数据，可以先采用stein算法将数据降至int32范围，然后再使用向量化的gcd来完成计算。在stein算法过程中，需要把降低了数值范围的两个值存入一个tmpBuffer中，后续gcd使用。经过一些测试，当a<(1<<24)，gcd迭代次数24的时候，是可以保证计算的正确性的。可以将stein的a降低到更小，让gcd迭代次数更多，来确保在全数据范围都能计算正确。

赛题的时间最长的case刚好是单一维度广播的int64类型，对于单一维度广播，相较于BitwiseLeftShift，还更改了循环的模式，一次读取other数据，处理多个相应的input数据，减少other tensor的标量读取，以及所有的标量计算部分，都可以通过循环展开来提升性能。经实际测试，展开4次的性能是最好的。