Skip to content

TVM-AF-DS

Chendi Li, Yufan Xu, Sina Mahdipour Saravani, and Ponnuswamy Sadayappan. 2024. Accelerated Auto-Tuning of GPU Kernels for Tensor Computations. In Proceedings of the 38th ACM International Conference on Supercomputing (ICS '24). Association for Computing Machinery, New York, NY, USA, 549–561.

https://doi.org/10.1145/3650200.3656626

Introduction

TVM/Ansor 需要 15-30 分钟来进行 1000 次测量，调优一层。备选配置搜索的不确定性，导致 1000 次测量后的最佳性能可能存在相当大的变化。要进行超过 1000 次测量的搜索，或进行次进行 1000 次搜索，才能对最优配置到达最佳可能性能有足够的信心。

减少编译时间的编译器 [11,18,33]，没有在广泛的张量计算中表现出与 TVM/Ansor 相当的性能。本文在不牺牲性能的情况下，将测量的次数减少一个数量级。

贡献：

  • New Features for Proxy Performance Model：TVM/Ansor 使用 XGBoost 模型来预测性能，使用从代码的 AST（抽象语法树）派生的 164 个组件的大型特征向量。本文使用 7 个特征作为替代，如共享内存和寄存器级别的数据重用、指令级别和 warp 级别并发性以及 SM（流式多处理器）之间的负载不平衡估计。
  • New Dynamic Search Space Exploration Strategy： TVM/Ansor 使用遗传算法来搜索调度空间。本文开发了一种在代码配置的 tile 大小多维空间中进行动态梯度下降的搜索策略。
  • Demonstration of Order-of-Magnitude Speedup over TVM/Ansor Auto-tuning：本文只需 1-2 分钟进行搜索。

Overview of Paper

TVM/Ansor 对输入的算子配置，创建参数化的多级 tile 搜索空间。搜索空间中的一个配置可描述为一个 tile-size 向量 T 和一个特征向量 F。Ansor 从 AST 生成 164 个特征的特征向量，包括数据移动的状态（ global 和 share memory 之间的，share memory 和 register 之间的），表达算子执行的计算状态。特征向量 F 用于训练在线 ML model（XGBoost），model 用于预测性能。Ansor 使用遗传算法进行搜索，多轮迭代后，将为选出的配置生成代码、编译、执行。执行时间重新训练 ML model。

本文使用 Ansor-AF 来指代使用了 7 个特征的 Ansor 版本。i）数据移动，ii）并发/占用，iii）SM 间的负载不平衡。

基于多维 tile 空间的动态梯度下降的搜索策略。本文的 proxy ML model 用于预测空间中相邻点的性能。预测性能好的邻居将编译执行，如果性能更好，搜索轨迹将移动到性能更好的邻居。这个变体叫做 Ansor-DS。

两种策略结合起来，形成 Ansor-AF-DS。

Background on GPU Kernel Optimization by TVM/Ansor Auto-tuning

略

Ansor-AF: ML Performance Modeling with Analytical Features

Analytical performance modeling features

  1. 数据移动。以 Operational Intensities 来计算。OI = 算术运算总数/数据移动事务数。
    • OI_Global_Mem：从全局内存到共享内存的数据移动的 OI。
    • OI_Shared_Mem：从共享内存到寄存器的数据移动的 OI。(warp 间共享内存访问，可能会有 bank 冲突，通过步长估计 bank 冲突）
    • Reg_Reuse_Factor：寄存器到 global memory 的 OI。
  2. 并发性。
    • ILP：指令级并行。线程内并行度。寄存器 tile 大小可估计 ILP。
    • WLP： warp 级并行。WLP 表达了每个 SM 中同时活跃的 warp 数量。WLP 主要取决于资源限制——每线程的寄存器使用量和每 block 的共享内存使用量。使用量通过三个 array（？？？）的寄存器级和共享内存的 buffer 大小来估计。
    • Estimated_Occupancy：实际活跃 warp 数量和每个 SM 可加载的并发 warp 数量上限之比。活跃 warp 根据寄存器和共享内存使用情况来估计。
  3. 负载不平衡。根据 tail effect 来估计这个特征。
    • Wave_Efficiency：估计 SM 利用率。若 m 个 block，n 个 SM，那么 Wave_Efficiency = \(\frac{m/n}{\lceil{m/n}\rceil}\)

kernel 性能预测模型是基于如上 7 个特征的回归模型。使用这 7 个特征替换掉 Ansor 使用的 164 个特征来使用 XGBoost 决策树模型，该方法称为 Ansor-AF。

Performance model evaluation

对比 Ansor 和 Ansor-AF 获得的最佳代码的 TFLOP（取三次独立搜索的平均值）。发现：

  1. 仅执行 2 min 的搜索，绝大多数 baseline 算子（27/30），Ansor-AF 获得的代码更好
  2. 若进行 1000 次 measure，大多数 baseline 测试（18/30），Ansor-AF 性能优于 Ansor。
  3. 若 2 min 的 Ansor-AF 搜索结果，与 1000 次 measure 的 Ansor 的搜索结果对比，仅 2 个 baseline 获得更优性能。

Ansor-DS: Dynamic Gradient Descent Search Space Exploration

搜索空间构建

在多维搜索空间中，每个维度的 iterator 都应该是迭代空间对应维度范围的完美除数。本文将一个 iterator 的一个 hop （跳）定义为从当前 factor 跳跃到 factor list 中的相邻值。例如，如果 I 是 32，那么它的 factor list 为 (1,2,4,8,16,32)。当 iterator i 的值为 8 时，在一跳内，它可以变为 4 或者 16。对于一个三维的空间，某一点的一跳节点有 6 个，而两跳节点有 12 个（三维空间中变化某两维的值为相邻值），三跳节点有 8 个（三个维度的值都改变为相邻值）。

动态搜索算法

基于轻量级在线测量的动态梯度下降算法 —— DGD_Search。

首先随机 64 个代码进行初始性能模型的构建。每一轮会从随机节点开始，探索 n 跳距离的候选者（从 1 开始），并使用性能模型预测候选者性能。DGD_Move 对预测结果得分较高的候选者节点进行性能测量，测量数据会被用于更新性能模型。如果在 n 跳距离的候选者中测量发现了性能更优点，那么更优点作为新的当前节点，重复上述过程，否则继续增加 n 的值直到阈值（hop_threshold）。不断移动轨迹，直到达到局部最优点（n 达到阈值，且 DGD_Move 无法找到性能更优点），本轮搜索结束。下轮搜索将从新的随机节点开始，直到进行足够多的轮数（search_budget）。

def DGD_Search(search_budget, hop_threshold):
    # construct the init model
    model.init(sample_size=64)
    global best_measured = 0
    while i < search_budget:
        cur_point = random_sample_point()
        while cur_point != null:
            n = 1
            meas_list.clear()
            do:
                points = n_hop_cands(cur_point, n)
                scores = model.predict(points)
                sort(scores)
                cur_point, meas_list = DGD_Move(points, scores)
                n += 1
            while (n < hop_threshold and cur_point == null)
            # update model with measured code
            model.update(meas_list)
        # start a new point after hiting local min
        i += 1
return

DGD_Move 函数通过使用大小为 3 的滑动窗口决定移动位置。即在按预测分数排好序的候选节点列表中，从分数最高的一端开始，用滑动窗口选择需要在设备上进行测量的候选者。首先选择预测分数最高的 3 个候选者进行测量，若测量获得的性能超过了当前节点，搜索轨迹移动（即将性能更优的候选节点返回上层函数），否则继续向后滑动窗口寻找是否存在性能更优点。若窗口中最好性能比全局最优性能的 0.6 倍还要小，则 early stop，认为在当前候选者中无更优节点。

def DGD_Move(points, scores, best_measured):
    base_score = scores[0]
    score_threshold = base_score * 0.6
    measure_threshold = best_measured * 0.6
    slide = 3    # slide window size for measurement
    i = 1
    meas_list = []
    while (all(scores[i:i+wd]) >= score_threshold and i+wd < points.size):
        local_meas_list = measure(points[i:i+wd])
        if max(local_meas_list) > base_performance:
            next_base = local_meas_list[max_index]
            meas_list.append(local_meas_list)
            break
        else :
            meas_list.append(local_meas_list)
        best_measured = max(local_meas_list)
        # early stop
        if max(local_meas_list) < measure_threshold:
            break
        i += wd
return next_base , meas_list

动态搜索算法有效性评估

  • 经过 2 分钟的调优后，Ansor-DS 比 Ansor 表现更好
  • 21/30 的基准测试，2 分钟的 Ansor-DS 优于 1000 measure 的 Ansor。
  • 1000 measure 的 Ansor DS 在 22/30 的基准测试上，比 Ansor 好 5% 以上，而其余的相当（<5%提升）。

Ansor-AF-DS: Evaluation of Integrated Optimizations

References

[11] Perry Gibson and José Cano. 2022. Transfer-Tuning: Reusing Auto-Schedules for Efficient Tensor Program Code Generation. In Proceedings of the International Conference on Parallel Architectures and Compilation Techniques, PACT 2022, Chicago, Illinois, October 8-12, 2022. ACM, 28–39. https://doi.org/10.1145/3559009.3569682

[18] Jaehun Ryu, Eunhyeok Park, and Hyojin Sung. 2022. One-shot tuner for deep learning compilers. In CC ’22: 31st ACM SIGPLAN International Conference on Compiler Construction, Seoul, South Korea, April 2- 3, 2022. ACM, 89–103. https://doi.org/10.1145/3497776.3517774

[33] ] Hongyu Zhu, Ruofan Wu, Yijia Diao, Shanbin Ke, Haoyu Li, Chen Zhang, Jilong Xue, Lingxiao Ma, Yuqing Xia, Wei Cui, Fan Yang, Mao Yang, Lidong Zhou, Asaf Cidon, and Gennady Pekhimenko. 2022. ROLLER: Fast and Efficient Tensor Compilation for Deep Learning. In 16th USENIX Symposium on Operating Systems Design and Implementation, OSDI 2022, Carlsbad, CA, USA, July 11-13, 2022. USENIX Association, 233–248. https://www.usenix.org/conference/osdi22/presentation/zhu