0%

Compute Shader 实现 Pyramid Algorithm

使用 compute shader 生成 mipmap

本文是参考 https://github.com/nvpro-samples/vk_compute_mipmaps 的算法复现。

将输入层级记为 \(L\),将 \(L\) 的 2x2 样本区域生成一个样本则得到 \(L+1\) 层级。下面介绍从 \(L\)\(L+M\) 层级的生成算法,\(M\in[1,6]\),算法输入为 \(L\) 层级的大小为 \(2^M\times 2^M\) 的 tiles。

宽高为 2 的幂次

首先将输入层级划分为多个大小为 \(2^M\times 2^M\) 的 tile,每个 tile 作为 compute shader 中一个 local workgroup 的输入 tile,每个 local workgroup 负责生成其输入 tile 的层级。假设输入层级大小为 \(X,Y\),那么一次 dispatch 产生的 local workgroup 数量为 \((X/2^M, Y/2^M)\),因此使用 workgroup id 作为 tile id,进而得到输入 tile 在输入层级中的起始位置,用以获取输入层级的样本。

  • 假设 workgroup id 为 \((m,n)\),该 workgroup 的输入 tile id 为 \((m,n)\)。该输入 tile 在输入层级中的起始位置设为 (xOffset, yOffset),其计算方式为

    1
    2
    xOffset = m << M;
    yOffset = n << M;

  • 为了能够利用到 subgroup 内部数据同步的高效性,将 workgroup 的线程/invocation 划分为大小为 16 的子组。这里使用 1D 的 invocation ID,gl_LocalInvocationIndex 来划分组,划分方式为

    1
    teamID = (ID & 0x00F0) >> 4;

  • 想要生成当前层级的样本,则需要确定当前层级样本在上一层级中对应的 \(2\times 2\) 区域,即该区域在上一层级中的起始偏移量 (xBlockOffset, yBlockOffset)。该偏移量有三部分组成,当前 workgroup 对应的输入 tile 在输入层级中的起始偏移量 (xOffset, yOffset)、当前 invocation 所属组在输入 tile 中的起始偏移量 (tx, ty)、当前 invocation 在其所属组的起始偏移量 (x, y)。最终有

    1
    2
    xBlockOffset = xOffset + tx + x;
    yBlockOffset = yOffset + ty + y;

1-Level

输入层级划分为 \(2\times 2\) tiles,每个线程加载一个 \(2\times2\) tile 生成一个样本。1 level 算法比较简单,输入层级有多少个 \(2\times 2\) tiles 就分发多少个 invocation,即 global group size 与 tile 数量相等。直接使用 gl_GlobalInvocationID 来索引 \(L\) 层级的样本。以下为一个 tile 中的 4 个样本的索引

1
2
gl_GlobalInvocationID*4+0, gl_GlobalInvocationID*4+1, 
gl_GlobalInvocationID*4+2, gl_GlobalInvocationID*4+3

或者按照下述算法形式实现。

2-Levels

输入层级划分为 \(4\times 4\) tiles,local workgroup size 使用 4。一个输入 tile 可以划分为 4 个 \(2\times 2\) blocks,每个 block 对应 \(L+1\) 层的一个样本。因此每个 local workgroup 负责一个输入 tile,生成 4 个 \(L+1\) 层级样本,4 个 \(L+1\) 层级样本又可生成 1 个 \(L+2\) 层级样本。算法过程如下示意,

1
2
3
4
5
6
   L(4x4)            L+1(2x2)             L+2
+---------+ +-----+-----+ +---------+
| | | 0 | 1 | | |
| 4x4 | ===> +-----+-----+ ====> | 0 |
| | | 2 | 3 | | |
+---------+ +-----+-----+ +---------+

确定 \(L+1\) 层每个样本在 \(L\) 层对应的 \(2\times 2\) block 的起始位置,

  • 输入 tile 在 \(L\) 层级的起始位置:\(M=2\) 代入

    1
    2
    xOffset = m << M;
    yOffset = n << M;

  • local workgroup 内线程的分组 ID:由于 local workgroup size 为 4,因此只有一个分组,teamID 恒为 0,

    1
    teamID = (ID & 0x00F0) >> 4;

    当前分组在输入 tile 中的偏移量也恒为 tx=0, ty=0

  • 当前 invocation 在其所属分组的偏移量:gl_LocalInvocationIndex 为当前 invocation 的索引,记为 id。\(L+1\) 层的 4 个样本如下,以及对应的组内偏移量

    1
    2
    3
    4
    5
    6
     L+1 样本编号       L+1 样本编号二进制    L+1 样本的组内起始偏移量
    +-----+-----+ +-----+-----+ +-----+-----+
    | 0 | 1 | | 00 | 01 | |(0,0)|(0,2)|
    +-----+-----+ +-----+-----+ +-----+-----+
    | 2 | 3 | | 10 | 11 | |(2,0)|(2,2)|
    +-----+-----+ +-----+-----+ +-----+-----+

    当前 invocation 组内的偏移量的计算方式如下

    1
    2
    x = (id & 2);
    y = (id & 1) << 1;
  • 最终得到 \(L+1\) 层样本在 \(L\) 层对应的 \(2\times 2\) block 的起始偏移量:

    1
    2
    xBlockOffset = xOffset + x;
    yBlockOffset = yOffset + y;

得到 (xBlockOffset, yBlockOffset) 后,则可加载该 2x2 区域的 4 个 \(L\) 层样本,即

1
2
(xBlockOffset+0, yBlockOffset+0), (xBlockOffset+0, yBlockOffset+1)
(xBlockOffset+1, yBlockOffset+0), (xBlockOffset+1, yBlockOffset+1)

4 个 \(L\) 层样本生成一个 \(L+1\) 层样本,这样下来,每个 local workgroup 都得到了 \(L+1\) 层级的 4 个样本。\(L+1\) 层样本的写入位置为 (xBlockOffset >> 1, yBlockOffset >> 1)。

由于 local workgroup size 为 4,即 local workgroup 的所有 invocation 执行于同一个 subgroup,因此可利用 subgroup 内部的高效数据同步。每个 local workgroup 的 0 号线程再使用 subgroupShuffleXor 得到 1、2、3 号线程的样本,生成 \(L+2\) 层级的一个样本,写入位置为 (xBlockOffset >> 2, yBlockOffset >> 2)。

3-Levels

3 levels 通过不同区域分别应用 2 levels 算法。输入层级划分为 \(8\times 8\) tiles,local workgroup size 使用 16。将 local workgroup 的 16 个线程分为 4 个 subtiles,每个 subtile 大小为 \(2\times 2\),具有 4 个线程,如下所示

1
2
3
4
5
6
7
8
9
10
workgroup 内的 subtile,"|| =" 为 subtile 边界
+---+---++---+---+
| 0 | 1 || 4 | 5 |
+---+---++---+---+
| 2 | 3 || 6 | 7 |
+===+===++===+===+
| 8 | 9 ||12 |13 |
+---+---++---+---+
|10 |11 ||14 |15 |
+---+---++---+---+

local workgroup 的每个 subtile 应用 2 levels 算法来生成 \(L+1\)\(L+2\) 的样本,在执行完 2 levels 算法后,每个 subtile 得到一个 \(L+2\) 层样本。每个 workgroup 共得到 4 个 \(L+2\) 层样本,最后生成 1 个 \(L+3\) 层样本,过程如下所示,

1
2
3
4
5
6
7
8
9
10
     L(8x8)                  L+1(4x4)                  L+2(2x2)
+---------------+ +---+---++---+---+ +-------++-------+
| | | 0 | 1 || 4 | 5 | | || |
| | +---+---++---+---+ | 0 || 4 |
| | | 2 | 3 || 6 | 7 | | || |
| 8x8 | ====> +===+===++===+===+ ====> +=======++=======+
| | | 8 | 9 ||12 |13 | | || |
| | +---+---++---+---+ | 8 || 12 |
| | |10 |11 ||14 |15 | | || |
+---------------+ +---+---++---+---+ +-------++-------+

首先确定 \(L+1\) 层每个样本在 \(L\) 层对应的 \(2\times 2\) block 的起始位置,

  • 输入 tile 在 \(L\) 层级的起始位置:\(M=3\) 代入 (xOffset, yOffset) 的计算

  • local workgroup 内线程的分组编号:local workgroup size 为 16,只有一个分组,teamID 恒为 0,当前分组在输入 tile 中的偏移量也恒为 tx=0, ty=0

  • 当前 invocation 在其所属分组的偏移量:gl_LocalInvocationIndex 为当前 invocation 的索引,记为 id。\(L+1\) 层的 16 个样本编号以及对应的组内偏移量,如下

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
             L+1 层编号二进制              L+1 层编号对应的输入 tile 的偏移量
    +------+------++------+------+ +------+------++------+------+
    | 0000 | 0001 || 0100 | 0101 | |(0,0) |(0,2) ||(0,4) |(0,6) |
    +------+------++------+------+ +------+------++------+------+
    | 0010 | 0011 || 0110 | 0111 | |(2,0) |(2,2) ||(2,4) |(2,6) |
    +======+======++======+======+ +======+======++======+======+
    | 1000 | 1001 || 1100 | 1101 | |(4,0) |(4,2) ||(4,4) |(4,6) |
    +------+------++------+------+ +------+------++------+------+
    | 1010 | 1011 || 1110 | 1111 | |(6,0) |(6,2) ||(6,4) |(6,6) |
    +------+------++------+------+ +------+------++------+------+

    当前 invocation 组内的偏移量的计算方式如下

    1
    2
    x = (id & 2) | (id & 8) >> 1;
    y = (id & 1) << 1 | (id & 4);
  • 最终得到 \(L+1\) 层样本在 \(L\) 层对应的 \(2\times 2\) block 的起始偏移量:

    1
    2
    xBlockOffset = xOffset + x;
    yBlockOffset = yOffset + y;

与 2 levels 算法相似,得到 (xBlockOffset, yBlockOffset) 后,即可得到 \(L\) 层级的 4 个样本,从而生成 \(L+1\) 层级样本。每个 local workgroup 共生成 16 个 \(L+1\) 层级样本,写入位置为 (xBlockOffset >> 1, yBlockOffset >> 1)。

每个 local workgroup 在同一个 subgroup 内执行,因此利用 subgroup 内部的高效数据同步,每个 subtile 的第一个线程收集其他线程生成的 \(L+1\) 层级样本,然后生成一个 \(L+2\) 层级样本。每个 subtile 的第一个线程为 id&3 == 0,通过 使用 mask 参数分别为 1、2、3 的subgroupShuffleXor ,得到 subtile 的其他线程生成的 \(L+1\) 层级样本,过程如下

1
2
3
4
0 (0000) ^ (0001, 0010, 0011) = (0001, 0010, 0011) = (1, 2, 3)
4 (0100) ^ (0001, 0010, 0011) = (0101, 0110, 0111) = (5, 6, 7)
8 (1000) ^ (0001, 0010, 0011) = (1001, 1010, 1011) = (9, 10, 11)
12(1100) ^ (0001, 0010, 0011) = (1101, 1110, 1111) = (13, 14, 15)

每个 local workgroup 生成 4 个 \(L+2\) 层级样本,写入位置为 (xBlockOffset >> 2, yBlockOffset >> 2)。每个 local workgroup 中的第一个线程使用 subgroupShuffleXor 得到 4、8、12 号线程的 \(L+2\) 层级样本,生成一个 \(L+3\) 层级样本,写入位置为 (xBlockOffset >> 3, yBlockOffset >> 3)。每个 local workgroup 的第一个线程为 id&15 == 0,过程如下:

1
0(0000) ^ (0100, 1000, 1100) = (0100, 1000, 1100)

4/5-Levels

理论上,上述算法可以继续递归应用得到 4/5 levels 算法,例如对于 4 levels 算法,输入层级划分为 \(16\times 16\) tiles,local workgroup 大小使用 64,划分为 4 组,每组 16 个线程,分别应用 3 levels 算法,然后 0 号线程使用 shuffle 得到 16、32、48 号线程的样本,最终生成 \(L+4\) 的一个样本。但实际上,subgroup 大小有限,例如 NVIDIA 显卡的 subgroup 大小为 32,无法使用 32 及之后的 gl_SubgroupInvocationID,因此不能通过 subgroup 得到 32、48 号线程的样本。对于此,改用 shared memory。

4-levels

对于 4 levels 算法,输入层级划分为 \(16\times 16\) tiles,local workgroup 大小使用 64,划分为 4 组,每组 16 个线程。local workgroup 中的线程使用 gl_LocalInvocationIndex 分组,0~15 为组 0,16~31 为组 1,32~47 为组 2,48~63 为组 3,如下所示:

1
2
3
4
5
6
7
8
9
10
11
workgroup 内的分组,"|| =" 为分组边界
每组中间数字为分组编号
+---+---++---+---+
| 0 | 4 ||16 |20 |
+---0---++---1---+
| 8 |12 ||24 |28 |
+===+===++===+===+
|32 |36 ||48 |52 |
+---2---++---3---+
|40 |44 ||56 |60 |
+---+---++---+---+

每组分别执行 3 levels 算法。但每组的第一个线程不仅生成 \(L+3\) 的一个样本,还要将该样本写入 shared memory中。然后发出一个 barrier,barrier 之后在 shared memory 中存在 4 组线程分别写入的一个 \(L+3\) 层级样本,即 \(L+3\) 层级的一个 \(2\times 2\) tile。最后 0 号线程使用 shared memory 中的 4 个 \(L+3\) 层级样本生成一个 \(L+4\) 层级样本。过程如下:

1
2
3
4
5
6
7
8
9
10
11
12
     L(16x16)               L+2(8x8)                       L+3(2x2)                   L+4
+-------++-------+ +---+---++---+---+ +-------++-------+ +---------------+
| || | | 0 | 4 ||16 |20 | | || | | |
| 8x8 || 8x8 | +---+---++---+---+ | 0 || 16 | | |
| || |(L+1)| 8 |12 ||24 |28 | | || | | |
+=======++=======+ ==> +===+===++===+===+ =>omitted=> +=======++=======+ ==|==> | 0 |
| || | |32 |36 ||48 |52 | | | || | | | |
| 8x8 || 8x8 | +---+---++---+---+ | | 32 || 48 | | | |
| || | |40 |44 ||56 |60 | | | || | | | |
+-------++-------+ +---+---++---+---+ | +-------++-------+ | +---------------+
| |
16-thread team handles each 8x8 sub-tile (barrier)

首先确定 \(L+1\) 层每个样本在 \(L\) 层对应的 \(2\times 2\) block 的起始位置,

  • 输入 tile 在 \(L\) 层级的起始位置:\(M=4\) 代入 (xOffset, yOffset) 的计算

  • local workgroup 内线程的分组编号:local workgroup size 为 64,分为 4 个组,0~15 为组 0,16~31 为组 1,32~47 为组 2,48~63 为组 3。如下所示,

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
        L+1 层编号十六进制        L+1层编号对应的输入 tile 的偏移量           L+1 层分组编号
    +----+----++----+----+ +------+------++------+------+ +-----------++-----------+
    |0000|0004||0010|0014| |(0,0) |(0,4) ||(0,8) |(0,12)| | || |
    +----+----++----+----+ +------+------++------+------+ | 0 || 1 |
    |0008|000C||0018|001C| |(4,0) |(4,4) ||(4,8) |(4,12)| | || |
    +====+====++====+====+ +======+======++======+======+ +===========++===========+
    |0020|0024||0030|0034| |(8,0) |(8,4) ||(8,8) |(8,12)| | || |
    +----+----++----+----+ +------+------++------+------+ | 2 || 3 |
    |0028|002C||0038|003C| |(12,0)|(12,4)||(12,8)|(12,12| | || |
    +----+----++----+----+ +------+------++------+------+ +-----------++-----------+

    invocation 索引为 id,所属分组编号计算方式如下:

    1
    teamID = (id & 0x00F0) >> 4;

    当前分组在输入 tile 中的偏移量为

    1
    2
    tx = (teamID & 2) << 2;
    ty = (teamID & 1) << 3;

  • 当前 invocation 在其所属分组的偏移量:

    1
    2
    x = (id & 2) | (id & 8) >> 1;
    y = (id & 1) << 1 | (id & 4);

  • 最终得到 \(L+1\) 层样本在 \(L\) 层对应的 \(2\times 2\) block 的起始偏移量:

    1
    2
    xBlockOffset = xOffset + tx + x;
    yBlockOffset = yOffset + ty + y;

得到 (xBlockOffset, yBlockOffset) 后即可在每个分组内分别执行 3-levels 算法。每个分组都会生成一个 \(L+3\) 层级样本,即 workgroup 生成了 4 个 \(L+3\) 层级样本。在生成最终的 \(L+4\) 层样本时,不能再使用 subgroupShuffleXor 来得到其他分组生成的 \(L+3\) 层级样本,因为 subgroup size 为 32,其他分组不一定位于同一个 subgroup 内,因此这时需要使用 shared memory。

每个分组的第一个线程不仅需要生成一个 \(L+3\) 层级样本,还需要将该样本写入 shared memory,组内的第一个线程为 (id & 0x00F0)==id,写入位置为 teamID。加一个 barrier 等待 shared memory 的写操作完成。最后,workgroup 的第一个线程负责生成 \(L+4\) 层样本。

5-levels

对于 5 levels 算法,输入层级划分为 \(32\times 32\) tiles,local workgroup 大小使用 256,划分为 16 组,每组 16 个线程。分组如下所示,

1
2
3
4
5
6
7
8
9
10
 workgroup的每个分组编号      workgroup内每组的invocation起始编号    
+----+----++----+----+ +----+----++----+----+
| 0 | 1 || 4 | 5 | | 0 | 16 || 64 | 80 |
+----+----++----+----+ +----+----++----+----+
| 2 | 3 || 6 | 7 | | 32 | 48 || 96 | 112|
+====+====++====+====+ +====+====++====+====+
| 8 | 9 || 12 | 13 | |128 | 144||192 | 208|
+----+----++----+----+ +----+----++----+----+
| 10 | 11 || 14 | 15 | |160 | 176||224 | 240|
+----+----++----+----+ +----+----++----+----+

与 4 levels 算法相似,每组分别执行 3 levels 算法,并且每组的第一个线程不仅生成 \(L+3\) 层级的一个样本,还要写入 shared memory 中。在 barrier 之后,shared memory 中存在 16 组线程分别写入的一个 \(L+3\) 层级样本,即 \(L+3\) 层级的一个 \(4\times 4\) tile。最后使用 4 个线程执行一次 2 levels 算法,最终得到 4 个 \(L+4\) 层级样本与 1 个 \(L+5\) 层级样本。过程如下所示,

1
2
3
4
5
6
7
8
9
10
11
12
  L(32x32)                 L+3 (+ copy in smem)            L+4(2x2)                  L+5
+---+---++---+---+ +---+---++---+---+ +-------++-------+ +---------------+
|8x8|8x8||8x8|8x8| | 0| 16|| 64| 80| | || | | |
+---+---++---+---+ +---+---++---+---+ | 0 || 1 | | |
|8x8|8x8||8x8|8x8| | 32| 48|| 96|112| | || | | |
+===+===++===+===+ =>omitted=> +===+===++===+===+ ==|==> +=======++=======+ ===> | 0 |
|8x8|8x8||8x8|8x8| | |128|144||192|208| | | || | | |
+---+---++---+---+ | +---+---++---+---+ | | 2 || 3 | | |
|8x8|8x8||8x8|8x8| | |160|176||224|240| | | || | | |
+---+---++---+---+ | +---+---++---+---+ | +-------++-------+ +---------------+
| |
16-thread team handles each 8x8 sub-tile (barrier)

首先确定 \(L+1\) 层每个样本在 \(L\) 层对应的 \(2\times 2\) block 的起始位置,

  • 输入 tile 在 \(L\) 层级的起始位置:\(M=5\) 代入 (xOffset, yOffset) 的计算

  • local workgroup 内线程的分组编号:local workgroup size 为 256,分为 16 个组,分组如下

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
            分组编号                每组的起始编号十六进制            起始编号在输入 tile 中对应的偏移量
    +----+----++----+----+ +----+----++----+----+ +-------+-------++-------+-------+
    | 0 | 1 || 4 | 5 | |0000|0010||0040|0050| | (0,0) | (0,8) ||(0,16) | (0,24)|
    +----+----++----+----+ +----+----++----+----+ +-------+-------++-------+-------+
    | 2 | 3 || 6 | 7 | |0020|0030||0060|0070| | (8,0) | (8,8) ||(8,16) | (8,24)|
    +====+====++====+====+ +====+====++====+====+ +=======+=======++=======+=======+
    | 8 | 9 || 12 | 13 | |0080|0090||00C0|00D0| |(16,0) | (16,8)||(16,16)|(16,24)|
    +----+----++----+----+ +----+----++----+----+ +-------+-------++-------+-------+
    | 10 | 11 || 14 | 15 | |00A0|00B0||00E0|00F0| |(24,0) | (24,8)||(24,16)|(24,24)|
    +----+----++----+----+ +----+----++----+----+ +-------+-------++-------+-------+

    invocation 索引为 id,所属分组编号计算方式如下:

    1
    teamID = (id & 0x00F0) >> 4;

    当前分组在输入 tile 中的偏移量为

    1
    2
    tx = (teamID & 2) << 2 | (teamID & 8) << 1;
    ty = (teamID & 1) << 3 | (teamID & 4) << 2;
  • 当前 invocation 在其所属分组的偏移量:

    1
    2
    x = (id & 2) | (id & 8) >> 1;
    y = (id & 1) << 1 | (id & 4);
  • 最终得到 \(L+1\) 层样本在 \(L\) 层对应的 \(2\times 2\) block 的起始偏移量:

    1
    2
    xBlockOffset = xOffset + tx + x;
    yBlockOffset = yOffset + ty + y;

得到 (xBlockOffset, yBlockOffset) 后,与 4-levels 算法类似,在每个分组内分别执行 3-levels 算法。每组生成一个 \(L+3\) 层级样本,且组内第一个线程 (id & 0x00F0)==id 负责写入 shared memory 中索引 teamID 的位置。workgroup 共生成 16 个 \(L+3\) 层级样本,得到一个 4x4 大小的 shared memory。

为了能够继续使用 subgroup 特性,使用一个 subgroup 来生成最后两层。假设使用 ID 为 0~15 的 invocation,ID 用来索引 shared memory 中的样本。shared memory 中的数组分为四部分,03、47、811、1215,每个部分包含 4 个 \(L+3\) 层级样本。 id == id & 0x000C 的 invocation 负责收集 4 个 \(L+3\) 层级样本,生成 1 个 \(L+4\) 层级样本。最后 0 号 invocation 通过 subgroupShuffleXor 收集其他 3 个 \(L+4\) 层级样本,生成最终的 1 个 L+5 层级样本。

宽高为非 2 的次幂

对于图像宽高不是 2 的次幂的情况,\(L+1\) 层级的样本在 \(L\) 层级对应的 block 可能是 \(1\times 1\)\(1\times 2\)\(2\times 1\)、$2$2、\(2\times 3\)\(3\times 2\)\(3\times 3\) 等。在划分为 tile 计算之后层级样本时,某些样本会用到多次,如下述 \(5 \times 5\) 大小的 mip level 生成 \(2\times 2\) 大小的 mip level,其中数字表示对应样本被使用到的次数。

1
2
3
4
5
6
7
8
9
10
11
+---+---+---+---+---+        +---------+---------+
| | | 2 | | | | | |
+---+---+---+---+---+ | | |
| | | 2 | | | | | |
+---+---+---+---+---+ | | |
| 2 | 2 | 4 | 2 | 2 | =====> +---------+---------+ (each sample generated with 3x3 kernel)
+---+---+---+---+---+ | | |
| | | 2 | | | | | |
+---+---+---+---+---+ | | |
| | | 2 | | | | | |
+---+---+---+---+---+ +---------+---------+

这意味着,无法像前述 2 的次幂的算法那样,每个 workgroup 负责的输入 tile,不会和其他 workgroup 有任何交集。这种由非 2 的次幂引入的多次引用,会导致一些 workgroup 生成下一层级样本时, 也可能使用到其他 workgroup 生成的样本。为了提高并行性,去除不同 workgroup 之间的同步,可以推算出生成 \(M\) 层级所需的输入层级的大小范围。在划分 workgroup 时,不同 workgroup 之间的输入 tile 允许有重叠部分。