在 OpenCL 中计算部分和

Computing partial sums in OpenCL

一个一维数据集被分成多个段,每个工作项处理一个段。它从段中读取了一些元素?元素的数量事先未知,每个段都不同。

例如:

+----+----+----+----+----+----+----+----+----+     <-- segments
  A    BCD  E    FG  HIJK   L    M        N        <-- elements in this segment

在所有段都被处理后,他们应该将 elements 写入 连续 输出内存,如

A B C D E F G H I J K L M N

所以一个段中元素的绝对输出位置取决于前面段中元素的个数。 E 位于位置 4,因为段包含 1 个元素 (A),段 2 包含 3 个元素。


OpenCL 内核将每个段的元素数量写入 local/shared 内存缓冲区,并像这样工作(伪代码)

kernel void k(
    constant uchar* input,
    global int* output,
    local int* segment_element_counts
) {
    int segment = get_local_id(0);
    int count = count_elements(&input[segment * segment_size]);

    segment_element_counts[segment] = count;

    barrier(CLK_LOCAL_MEM_FENCE);

    ptrdiff_t position = 0;
    for(int previous_segment = 0; previous_segment < segment; ++previous_segment)
        position += segment_element_counts[previous_segment];

    global int* output_ptr = &output[position];
    read_elements(&input[segment * segment_size], output_ptr);
}

因此每个工作项都必须使用循环计算部分和,其中具有较大 id 的工作项进行更多迭代。

在 OpenCL 1.2 中是否有更有效的方法来实现这一点(每个工作项计算序列的部分和,直到其索引)? OpenCL 2 似乎为此提供了 work_group_scan_inclusive_add

您可以使用类似这样的方法在 log2(N) 次迭代中进行 N 部分(前缀)求和:

offsets[get_local_id(0)] = count;
barrier(CLK_LOCAL_MEM_FENCE);

for (ushort combine = 1; combine < total_num_segments; combine *= 2)
{
    if (get_local_id(0) & combine)
    {
        offsets[get_local_id(0)] +=
            offsets[(get_local_id(0) & ~(combine * 2u - 1u)) | (combine - 1u)];
    }
    barrier(CLK_LOCAL_MEM_FENCE);
}

给定的段元素计数为

a     b     c        d

连续迭代将产生:

a     b+a   c        d+c

a     b+a   c+(b+a)  (d+c)+(b+a)

这就是我们想要的结果。

所以在第一次迭代中,我们将段元素计数分成 2 组,并在其中求和。然后我们一次将 2 个组合并为 4 个元素,并将结果从第一组传播到第二组。我们再次将组增加到 8 个,依此类推。

关键的观察是这个模式也匹配每个段索引的二进制表示:

0: 0b00  1: 0b01  2: 0b10  3: 0b11

索引 0 不执行求和。索引 1 和 3 在第一次迭代中执行求和(位 0/LSB = 1),而索引 2 和 3 在第二次迭代中执行求和(位 1 = 1)。这解释了这一行:

    if (get_local_id(0) & combine)

另一个真正需要解释的说法当然是

        offsets[get_local_id(0)] +=
            offsets[(get_local_id(0) & ~(combine * 2u - 1u)) | (combine - 1u)];

计算我们想要累加到工作项总和上的前一个前缀总和的索引有点棘手。子表达式 (combine * 2u - 1u) 在每次迭代中取值 (2n-1)(n 从 1 开始):

1 = 0b001
3 = 0b011
7 = 0b111
…

通过按位屏蔽这些位后缀(即 i & ~x)工作项索引,这将为您提供当前组中 第一个 项的索引.

(combine - 1u) 子表达式然后为您提供 上半部分的最后一项 当前组中的索引。把两者放在一起,就得到了你要累加到当前段的item的整体索引。

结果有点丑陋:它向左移动了一位:所以段 1 需要使用 offsets[0],依此类推,而段 0 的偏移量当然是 0。您可以将偏移量数组过度分配 1,并对从索引 1 开始的子数组执行前缀和并将索引 0 初始化为 0,或使用条件。

您可以对上述代码进行分析驱动的微优化。