XJTU 宇宙不怎么闪烁

在 S7 赛季中，西安交通大学“宇宙不怎么闪烁”编写了挑战赛题 Fmax, PdistGrad 和基础赛题 L1Loss, Softplus, Trace 共五题代码。均进入排行榜前 10 名。其中 Fmax 获校内赛第一名，L1Loss 获整题第一名（而且是断档领先，不知道怎么回事），Softplus 获整题第三名。

此仓库包含了最好成绩的源代码，并对于一些优化点做出解释。

通用的代码

在本赛季中，我们力图写出通用的代码，即一个模板函数在任何类型与任何广播情况下都有良好性能。所以我们的代码都比较短，尽可能不对输入情况做分类。这种设计思路只能说是有利有弊，能够逼迫编写者深入思考性能的关键，但是维护难度和掌控难度其实比分开写更高。而且，越复杂的场景越难达成。比如 cublas 对应的场景，要求极致的性能，那么对于每一种形状比例都要特化核函数。

手动 Double Buffer

Ascend C 允许我们手动管理内存。因此可以不使用队列实现 double buffer，这里以单输入的 Softplus 为例：

// in Init()
InitSocState(); 

// in Process()
LocalTensor<T> x_local[MANUAL_BUFFER_NUM] = {
    LocalTensor<T>(TPosition::VECIN, 0, tile_len),
    LocalTensor<T>(TPosition::VECIN, tile_len * 2 * sizeof(calc_t<T>), tile_len),
};
LocalTensor<T> y_local[MANUAL_BUFFER_NUM] = {
    LocalTensor<T>(TPosition::VECOUT, tile_len * sizeof(calc_t<T>), tile_len),
    LocalTensor<T>(TPosition::VECOUT, tile_len * 3 * sizeof(calc_t<T>), tile_len),
};

for (int i = 0; i < tile_num; i += MANUAL_BUFFER_NUM) {
    for (int j = 0; j < MANUAL_BUFFER_NUM; j++) {
        // ================ BEGIN Tiling ================
        int tile = i + j;
        if (tile >= tile_num) break;
        int cur_len = (tile == tile_num - 1) ? last_len : tile_len;
        bool set_backward  = (tile + MANUAL_BUFFER_NUM < tile_num);
        bool wait_backward = (tile >= MANUAL_BUFFER_NUM);
        // ================ END Tiling ================
        if (wait_backward) WaitFlag<HardEvent::V_MTE2>(j);
        // ================ BEGIN CopyIn ================
        DataCopy(x_local[j], x_global[tile * tile_len], cur_len);
        // ================ END CopyIn ================
        SetFlag<HardEvent::MTE2_V>(j);
        if (wait_backward) WaitFlag<HardEvent::MTE3_V>(j);
        WaitFlag<HardEvent::MTE2_V>(j);
        // ================ BEGIN Compute ================
        // Compute()
        // ================ END Compute ================
        SetFlag<HardEvent::V_MTE3>(j);
        WaitFlag<HardEvent::V_MTE3>(j);
        // ================ BEGIN CopyOut ================
        DataCopy(y_global[tile * tile_len], y_local[j], cur_len);
        // ================ END CopyOut ================
        if (set_backward) SetFlag<HardEvent::MTE3_V>(j);
    }
};

思路上就是第 j 个 buffer 持有一把叫做 j 的锁，用 SetFlag 和 WaitFlag 同步。需要注意对于多核的 910b，MANUAL_BUFFER_NUM 需要乘上 GetBlockNum()，得到正确的步长。

Fmax

广播枚举顺序优化

使用上一赛季 Fmin 的开源代码即可获得 16000 左右的成绩。这一份代码的关键优化在于对广播的处理。假设输入维度为 $(N_1, 1, N3, 1),(N_1, N_2, N_3, N_4)$，如果枚举输出元素进行处理，只能减少 $N_4$ 一维的输入量；但如果枚举第一个输入的元素进行处理，则可以减少 $N_1, N_3$ 两维的读入。尽管这会让输出不是连续写，但是数据量的减少对优化更有效。我们的代码选择枚举实际元素数最少的输入：

shape[0][len] = cnt[0];
shape[1][len] = cnt[1];
for (int j = 0; j < 2; j++) {
    throughput[j][j] *= shape[j][len];
    if (!len) throughput[j][j ^ 1] *= shape[j][len];
    else throughput[j][j ^ 1] *= std::max(shape[j][len], shape[j ^ 1][len]);
}
cnt[0] = cnt[1] = 1;
len++;
std::reverse(shape[0], shape[0] + MAX_SHAPE);
std::reverse(shape[1], shape[1] + MAX_SHAPE);

int iter_idx = 0;
if (throughput[0][0] + throughput[1][0] > throughput[0][1] + throughput[1][1]) iter_idx = 1;

实际枚举过程就是带余除法，其实跟 $k$ 进制的枚举过程很像，但是每一位的进制都不一样：

__aicore__ inline void calc_vec(int64_t iter, int64_t &input_vec, int64_t &other_vec,
                                    int64_t &out_vec, int64_t &tile) {
    input_vec = other_vec = out_vec = 0;
    for (int i = 0; i < DIM_LAST; i++) {
        if (input_stride[i]) {
            int64_t id = iter / iter_prod_input_order[i];
            iter %= iter_prod_input_order[i];
            input_vec += id * input_prod[i];
            other_vec += id * other_stride[i] * other_prod[i];
            out_vec += id * out_prod[i];
        }
    }
    tile = iter / iter_prod_input_order[DIM_LAST];
    iter %= iter_prod_input_order[DIM_LAST];
    for (int i = 0; i < DIM_LAST; i++) {
        if (!input_stride[i]) {
            int64_t id = iter / iter_prod_input_order[i];
            iter %= iter_prod_input_order[i];
            other_vec += id * other_stride[i] * other_prod[i];
            out_vec += id * out_prod[i];
        }
    }
}

bflaot16 处理优化

传统 bfloat16 的处理思路是 Cast 到 float32 进行处理，但是这里给出如下引理：

对于 bfloat16 的两个非 nan/inf 变量 a, b，如果 a < b，那么 &reinterept_cast<half*>(a) < &reinterept_cast<half*>(b)

这是浮点数的特性决定的。按照符号位，指数位，尾数位排列的 $(S, E, D)$ 任何浮点数，把它视为原码表示的整数，序是不会变的。进而可以推出，无论指数位和尾数位的长度如何分布，序是不会变的。

因此，可以将 bfloat16 视作 half 直接比较。

template <class T> struct reduce {
    using type = T;
};
template <> struct reduce<bool> {
    using type = uint8_t;
};
template <> struct reduce<bfloat16_t> {
    using type = float16_t;
};
template <class T> using reduce_t = typename reduce<T>::type;

L1Loss

我是求和的大师 —— ShwStone

精度问题

对于 bfloat16，由于尾数只有 7 位，实际上无法达到千分之一的精度要求，因此需要还原 CPU 上的舍入机制。经过测试，torch 的舍入机制是 RoundMode::CAST_RINT，也即“四舍六入五成双”。

如何求和

在 API 文档里，ReduceSum 宣称自己可以对任意长度元素求和。但实际上，其内部至多调用两次 WholeReduceSum。一次调用可以对 256B 数据求和，也就是 64 个 float。因此，ReduceSum 其实至多只能对 4096 个元素求和。这个 bug 卡了我很久，后来在文档的一个不知名地方发现了端倪。现在已经不知道在哪找到了。

根据需求，我最终实现了分段求和的设计：WholeReduceSum 每个 repeat 求和 64 个数，最多可以设置 255 个 repeat，将 $64 \times 255$ 个数求和成 255 个数。对于这 255 个数，我们不进行进一步求和，而是准备一个大小为 256 元素的缓冲区，将这 255 个数用 Add 加到缓冲区中，耗时 4 个 repeat。如果有超过 $64 \times 255$ 个数，就进行循环。最后，对于缓冲区里 255 个数，再调用一次 WholeReduceSum，求出最终的和。

template <class T> constexpr static int64_t block_n  = 32 / sizeof(T);
template <class T> constexpr static int64_t repeat_n = 256 / sizeof(T);

inline __aicore__ void ReduceSum255(const LocalTensor<float> &res, const LocalTensor<float> &x,
                                    int64_t len) {
    int64_t repeats = len / repeat_n<float>;
    int64_t count   = repeats / 255;

    for (int64_t i = 0; i < count; i++) {
        int64_t pos = i * 255 * repeat_n<float>;
        WholeReduceSum(x[pos], x[pos], 64, 255, 1, 1, 8);
        Add(res, res, x[pos], 255);
    }

    int64_t residue = repeats - count * 255;
    if (residue) {
        int64_t pos = count * 255 * repeat_n<float>;
        WholeReduceSum(x[pos], x[pos], 64, residue, 1, 1, 8);
        Add(res, res, x[pos], residue);
    }

    residue = len - repeats * repeat_n<float>;
    if (residue) {
        int64_t pos = repeats * repeat_n<float>;
        Add(res, res, x[pos], residue);
    }
}

几乎可以肯定有更好的求和方法，因为我没进行很多尝试。但是仅这一种方法就超越了所有人，在榜首呆了两个月。也就没动力优化了。

PdistGrad

这个我花的时间不多。有趣的发现是 Axpy 并不能加快计算速度，其速度和 Mul 后 Add 是一样的。

首先要进行推导。我们对 p=1,2,3 的情况进行了特化，对于其它的情况，使用 Exp,Ln 实现幂函数。

主要优化点在于 tiling 的设计。我们的设计如下：把向量拆成长度为 4096 的段，这是我的设计中 UB 能存下的最大值（因为要很多临时空间）。尽可能让每个核分到均匀数量的连续的段。grad 和 pdist 直接进行标量读，不是瓶颈。

从排行榜上来看也有更好的设计。

Softplus

中规中矩的吞吐量优化题。我不太擅长吞吐量优化，感觉没有头绪，只会调块长、buffer 数量这种，还要继续学习。

不过首先要对计算进行改良。首先，如果精度实现得当，threshold 其实可以直接不管。对于 float 类型，直接实现的精度是够的。对于 half 类型，技巧在于把输出和原始输入取最大值。因为 half 精度的主要问题是 Exp 函数会上溢，最多只能到 65506，那么对于更大的输入就会偏小。

if (beta == 1) {
    Exp(y_calc, x_calc, cur_len);
    if constexpr (is_same_v<U, float>)
        if (set_backward) SetFlag<HardEvent::V_MTE2>(j);
    Adds(y_calc, y_calc, SCAST<U>(1), cur_len);
    Ln(y_calc, y_calc, cur_len);
} else {
    Muls(y_calc, x_calc, SCAST<U>(beta), cur_len);
    if constexpr (is_same_v<U, float>)
        if (set_backward) SetFlag<HardEvent::V_MTE2>(j);
    Exp(y_calc, y_calc, cur_len);
    Adds(y_calc, y_calc, SCAST<U>(1), cur_len);
    Ln(y_calc, y_calc, cur_len);
    Muls(y_calc, y_calc, SCAST<U>(1 / beta), cur_len);
}
if constexpr (!is_same_v<U, float>) {
    Max(y_calc, x_calc, y_calc, cur_len);
    if (set_backward) SetFlag<HardEvent::V_MTE2>(j);
}

Trace

简单写了一点的题。关键在于两点：使用 DataCopyPad 读入，尽管为了一个数要读入 32B，但是仍然比标量快；使用 Gather 把数从 32B 一个抓成连续的，最后调用 ReduceSum 即可。

似乎有很多更厉害的人。我还要继续学习。我估计是在读入上进一步优化？

对于 Gather 的下标从哪来，虽然 310b 用不了 CreateVecIndex，但是从文档的描述可以知道，CreateVecIndex 的内部是反复调用 Add 函数进行倍增。

祝昇腾越来越好，祝大家代码水平越来越高