🏠

Bitonic Sorting Network 算法深度解析——基于 Vitis-Tutorials AI Engine 实现

1. 问题陈述

我们将问题形式化定义如下：

定义 1.1（排序问题）：给定长度为 \(N = 2^k\)（\(k \in \mathbb{N}^*\)）的输入序列 \(\mathbf{x} = [x_0, x_1, \dots, x_{N-1}]\)，其中 \(x_i \in \mathbb{R}\)，寻找输出序列 \(\mathbf{y} = \text{sort}(\mathbf{x})\)，满足：

  1. \(\mathbf{y}\) 是 \(\mathbf{x}\) 的一个排列：\(\{y_0, \dots, y_{N-1}\} = \{x_0, \dots, x_{N-1}\}\)
  2. 单调性约束：对于 \(0 \leq i < j < N\)，有 \(y_i \leq y_j\)（非降序）或 \(y_i \geq y_j\)（非升序）

定义 1.2（排序网络）：排序网络是一种仅由**比较交换器（Compare-Exchange, CE）**组成的计算图，其中每个比较交换器 \(\text{CE}(i,j,d)\) 执行以下操作： $$ (x_i, x_j) \leftarrow \begin{cases} (\min(x_i,x_j), \max(x_i,x_j)) & \text{若 } d = \text{up} \ (\max(x_i,x_j), \min(x_i,x_j)) & \text{若 } d = \text{down} \end{cases} $$ 排序网络的执行不依赖于输入数据的值，仅依赖于输入长度 \(N\)，因此适合硬件并行实现。

2. 直觉

2.1 朴素方法的局限性

经典的基于比较的排序算法（如快速排序、归并排序）的时间复杂度下界为 \(\Omega(N \log N)\)，但这些算法的控制流依赖于输入数据（如快速排序的枢轴选择），无法直接映射为固定拓扑的硬件电路，难以利用 AI Engine 等并行架构的 SIMD 或多核资源。

简单的并行排序网络（如奇偶排序网络）虽然具有固定拓扑，但其深度为 \(O(N)\)，对于大 \(N\)（如 \(1024\)）而言延迟过高。

2.2 关键洞察

Bitonic Sorting Network 的核心是基于**双调序列（Bitonic Sequence）**的性质：

定义 2.1（双调序列）：一个序列 \(\mathbf{a}\) 是双调的，当且仅当存在索引 \(0 \leq m < N\)，使得：

  1. \(a_0 \leq a_1 \leq \dots \leq a_m \geq a_{m+1} \geq \dots \geq a_{N-1}\)（先增后减），或
  2. 上述序列的循环移位，或
  3. 上述序列的反转（先减后增）

定理 2.1（双调排序定理，Batcher 1968）：给定长度为 \(N=2^k\) 的双调序列 \(\mathbf{a}\)，可以通过 \(k\) 级比较交换操作将其排序为单调序列，每级包含 \(N/2\) 个比较交换器。

利用该定理，Bitonic Sort 通过递归地将两个单调序列合并为双调序列，再对双调序列进行排序，最终得到完全有序的序列。

3. 形式化定义

我们将 Bitonic Sorting Network 形式化如下：

3.1 基本操作：Bitonic Merge

设 \(\mathbf{x}\) 是长度为 \(N=2^k\) 的双调序列，\(\text{BitonicMerge}(\mathbf{x}, N, d)\) 输出长度为 \(N\) 的单调序列，其中 \(d \in \{\text{up}, \text{down}\}\) 表示排序方向。其递归定义为： $$ \text{BitonicMerge}(\mathbf{x}, N, d) = \begin{cases} \mathbf{x} & \text{若 } N = 1 \ \text{concat}\left( \text{BitonicMerge}(\mathbf{x}{\text{low}}, N/2, d), \text{BitonicMerge}(\mathbf{x}{\text{high}}, N/2, d) \right) & \text{若 } N > 1 \end{cases} $$ 其中 \(\mathbf{x}_{\text{low}}\) 和 \(\mathbf{x}_{\text{high}}\) 由 \(\text{HalfClean}(\mathbf{x}, N, d)\) 生成：

\[ \begin{align*} \mathbf{x}_{\text{low}}[i] &= \text{CMPElement}(x_i, x_{i+N/2}, d)_{\text{min}} \\ \mathbf{x}_{\text{high}}[i] &= \text{CMPElement}(x_i, x_{i+N/2}, d)_{\text{max}} \end{align*} \]
这里 \(\text{CMPElement}(a,b,d)\) 返回 \((min(a,b), max(a,b))\) 当 \(d=\text{up}\)，否则返回 \((max(a,b), min(a,b))\)。

3.2 完整排序：Bitonic Sort

设 \(\mathbf{x}\) 是任意长度为 \(N=2^k\) 的序列，\(\text{BitonicSort}(\mathbf{x}, N, d)\) 输出单调序列，递归定义为： $$ \text{BitonicSort}(\mathbf{x}, N, d) = \begin{cases} \mathbf{x} & \text{若 } N = 1 \ \text{BitonicMerge}\left( \text{concat}\left( \text{BitonicSort}(\mathbf{x}{0:N/2-1}, N/2, \text{up}), \text{BitonicSort}(\mathbf{x}{N/2:N-1}, N/2, \text{down}) \right), N, d \right) & \text{若 } N > 1 \end{cases} $$

4. 算法

4.1 伪代码

// Bitonic Sorting Network 迭代实现（非递归，更适合硬件）
// 输入：data[0..N-1]，N=2^k；dir：排序方向，true=up，false=down
// 输出：data 被原地排序
procedure BitonicSort(data: array[float], N: int, dir: bool)
    // 外层循环：k = log2(N) 阶段
    for k_stage from 1 to log2(N) do
        m = 2^k_stage  // 当前阶段合并的双调序列长度
        // 中层循环：对于每个双调序列执行 BitonicMerge
        for j from m/2 downto 1 do
            // 内层循环：对所有并行的比较交换器
            for i from 0 to N-1 do
                if (i & j) == 0 then
                    // 比较 i 和 i+j，方向由 (i & m) == 0 决定
                    up = (dir == ((i & m) == 0))
                    CompareExchange(data[i], data[i+j], up)
                end if
            end for
        end for
    end for
end procedure

// 比较交换操作
procedure CompareExchange(a: float, b: float, up: bool)
    if (up and a > b) or (not up and a < b) then
        swap(a, b)
    end if
end procedure

4.2 执行流程图

flowchart TD Start[Start: Input Unsorted Array N=2^k] --> Init[Initialize k_stage from 1 to log2 N] Init --> m_calc["m = 2^k_stage (bitonic seq length)"] m_calc --> j_init["j from m/2 downto 1 (merge step size)"] j_init --> i_init["i from 0 to N-1"] i_init --> cond1{"(i AND j) == 0?"} cond1 -- No --> i_next[i = i+1] cond1 -- Yes --> up_calc["up = dir XOR ((i AND m) != 0)"] up_calc --> CE["CompareExchange(data[i], data[i+j], up)"] CE --> i_next i_next --> cond_i{"i < N?"} cond_i -- Yes --> i_init cond_i -- No --> j_next[j = j/2] j_next --> cond_j{"j >= 1?"} cond_j -- Yes --> j_init cond_j -- No --> k_next[k_stage = k_stage+1] k_next --> cond_k{"k_stage <= log2 N?"} cond_k -- Yes --> m_calc cond_k -- No --> End[End: Output Sorted Array]

5. 复杂度分析

5.1 比较交换器数量（工作复杂度）

定理 5.1：Bitonic Sorting Network 对于 \(N=2^k\) 的输入，包含的比较交换器总数为：

\[ C(N) = \frac{1}{2} N \log_2 N (\log_2 N + 1) \]
证明：由递归关系，\(C(N) = 2C(N/2) + (N/2)\log_2 N\)，边界条件 \(C(1)=0\)，展开可得上述闭式解。

5.2 网络深度（延迟复杂度）

定理 5.2：Bitonic Sorting Network 的深度（即串行执行的比较交换器级数）为：

\[ D(N) = \frac{1}{2} \log_2 N (\log_2 N + 1) \]
证明：由递归关系，\(D(N) = D(N/2) + \log_2 N\)，边界条件 \(D(1)=0\)，展开可得上述闭式解。

5.3 空间复杂度

  • 理论空间复杂度：\(O(N)\)，仅需存储输入序列本身，所有比较交换操作均为原地执行。
  • AI Engine 实现空间复杂度：需额外分配向量寄存器文件（Vector Register File, VRF）用于 SIMD 操作，但仍为 \(O(N)\) 量级。

5.4 情况分析

由于排序网络的拓扑完全由 \(N\) 决定，不依赖于输入数据的值，因此：

  • 最好情况：与最坏情况、平均情况的比较交换器数量和深度完全相同，即 \(C(N)\) 和 \(D(N)\)。

6. 实现说明

基于 Vitis-Tutorials 14-Bitonic-Sorting 的 AI Engine 实现，我们总结以下工程实现要点：

6.1 从理论到实现的调整

  1. 迭代替代递归：理论分析通常使用递归定义，但 AI Engine 实现采用 4.1 节的迭代版本，避免递归调用开销，更适合硬件流水线。
  2. SIMD 向量化：AI Engine 原生支持 128/256/512 位 SIMD 操作，实现中使用 aie::vector 类型同时处理多个 float 元素，将内层循环的串行比较交换转换为向量操作。
  3. 向量寄存器管理：对于 N=16 的小案例，直接使用 VRF 存储整个序列；对于 N=1024 的大案例，采用分块策略，将序列分段装入 VRF 处理，通过 DMA 或数据广播在 AI Engine 本地内存和 VRF 之间传输数据。

6.2 工程权衡

  1. 输入长度约束：严格要求 \(N=2^k\)，对于非 2 幂的输入需先进行填充（Padding）。
  2. 浮点比较处理：正确处理 IEEE 754 浮点数的特殊值（如 NaN、Inf），教程实现中假设输入为正常浮点数。
  3. 性能与资源权衡：增加并行度（如使用多个 AI Engine 核）可以提升吞吐量，但会增加核间通信开销和资源占用；教程中 N=1024 的实现可通过多核并行进一步扩展。

6.3 教程代码结构

  • bitonic_fp16/：N=16 的小案例演示，重点展示向量化和 VRF 管理。
  • bitonic_fp1024/：N=1024 的大案例实现，展示分块处理和 AI Engine 优化。
  • sortcpp_fp1024/：使用 std::sort 的参考实现，用于性能和正确性对比。

7. 对比分析

7.1 与经典排序算法对比

算法 时间复杂度（比较次数） 空间复杂度 数据依赖 并行友好性
Bitonic Sort \(O(N \log^2 N)\) \(O(N)\) 无 极高（固定拓扑）
快速排序 \(O(N \log N)\) 平均 \(O(\log N)\) 栈 有 低
归并排序 \(O(N \log N)\) \(O(N)\) 无 中
奇偶排序网络 \(O(N^2)\) \(O(N)\) 无 高

7.2 与其他排序网络对比

根据 Batcher 的原始论文（1968）和 Knuth 的《计算机程序设计艺术》第 3 卷：

  • 奇偶转置排序网络：深度为 \(N-1\)，比较次数为 \(N(N-1)/2\)，对于 \(N=1024\)，深度为 1023，远大于 Bitonic Sort 的深度约 512。
  • 奇数偶数排序网络（Batcher 另一网络）：比较次数与 Bitonic Sort 相当，但拓扑结构不同，对于 AI Engine 而言，Bitonic Sort 的规律步长更适合 SIMD 向量化。

7.3 教程实现与 std::sort 对比

根据 Vitis-Tutorials 的 profiling 结果：

  • 正确性：AI Engine 实现的输出与 std::sort 完全一致。
  • 吞吐量：AI Engine SIMD 实现的吞吐量显著高于 Arm A72 核心上的 std::sort，尤其适合大规模数据流排序。

参考文献

  1. Batcher, K. E. (1968). Sorting networks and their applications. AFIPS Spring Joint Computer Conference, 32, 307–314.
  2. Knuth, D. E. (1998). The Art of Computer Programming, Volume 3: Sorting and Searching (2nd ed.). Addison-Wesley.
  3. AMD (2024). Vitis AI Engine Development Tutorials: 14-Bitonic-Sorting. Retrieved from https://github.com/Xilinx/Vitis-Tutorials
On this page