GCC无法优化掉冗余的坐标转换函数调用,即使在自定义ND数组类中标记为“纯”

k3bvogb1  于 2024-01-08  发布在  其他
关注(0)|答案(1)|浏览(192)

背景和动机

我正在优化一个数值解算器,用有限差分法模拟麦克斯韦方程。
请注意,实际问题与物理、数值模拟甚至相关的优化技术都没有太大关系。实际上,这只是在实现自定义数组 Package 器时遇到的一个小问题。但为了完整起见,这里只描述了最小的细节。
在仿真中,物理空间被离散为三维网格(i, j, k),可以表示为三维阵列。在网格中,每个单元的值是3矢量(x, y, z),表示X,Y或Z轴上的电磁场或材料属性的数值。因此,仿真框可以表示为4D阵列(i, j, k, n)。其中n = 0, 1, 2(或者,它只是一个带有3个向量的3D数组)。
在每个时间步中,四个这样的4D阵列参与新电场的计算。它们是:

  • voltcurr-电场和磁场。
  • vvvi-材料属性

Kokkos并行编程库用于在CPU上编写GPU着色器风格的数据并行代码,还使用其原生N-D数组实现Kokkos::View

自定义数组类

为了实现进一步的优化,我需要做的第一件事是3D区域分解。不是根据标准C顺序[i][j][k]迭代3D数组中的每个单元,而是在3D tile的基本单元上迭代数组,例如20 x20 x20 tile。这被称为循环平铺或循环块。因此,阵列的底层布局也必须以匹配的顺序存储(否则,性能下降可能是显著的,特别是因为跨距-1 × 1 m10n1x维度不再被连续访问)。
为了实现这一点,我实现了一个自定义数组 Package 器,在Kokkos::View之上重载了()运算符。这样, Package 器就可以透明地将单元格(gi, gj, gk)的全局坐标转换为单元格的平铺坐标(ti, tj, tk)。它还允许人们在不更改代码的情况下尝试许多不同的数组布局。
首先定义一个Grid类,这个类只用来表示抽象的仿真网格:

  1. Grid(
  2. std::array<uint32_t, 3> grid_size,
  3. std::array<uint32_t, 3> tile_size
  4. ) :

字符串
它知道网格的尺寸,也知道如何执行坐标转换,使用成员函数公开:

  1. __attribute__((always_inline, pure))
  2. std::array<uint32_t, 4> global_coords_to_tile(
  3. const uint32_t gi,
  4. const uint32_t gj,
  5. const uint32_t gk
  6. ) const


接下来,实现ArrayIJKN类。它使用Grid提供的维度在内存中创建一个模拟框,并使用提供的坐标转换函数。

  1. ArrayIJKN(
  2. const std::string& name,
  3. const Grid& grid
  4. ) :
  5. m_grid(&grid)
  6. {
  7. uint32_t tile_num = grid.tile_num[DIM_I] * grid.tile_num[DIM_J] * grid.tile_num[DIM_K];
  8. // Yes, it's 5-level of indirection and can be considered a bad style
  9. // elsewhere, but this Kokkos::View template is officially designed to be
  10. // used this way. It templated for compile-time support of N-D arrays up
  11. // to 8 dimensions.
  12. //
  13. // The first dimension select a tile, the next 3 dimensions select a
  14. // cell within the 3D tile, the final dimension select the polarization
  15. // (X, Y, Z).
  16. m_ijk = new Kokkos::View<T****[n_num]>(
  17. Kokkos::ViewAllocateWithoutInitializing(name),
  18. tile_num,
  19. grid.tile_size[DIM_I], grid.tile_size[DIM_J], grid.tile_size[DIM_K],
  20. n_num
  21. );
  22. }


这个类是实际的 Package 类。它接受之前创建的grid对象来创建一个底层的5D数组(tile_id, ti, tj, tk, n)。可以在这个“tile坐标系”下访问一个值。但是为了允许透明地试验不同的布局,它还允许通过重载operator ()来访问使用标准4D全局坐标系(gi, gj, gk, n)的数组。在此成员函数中,它使用Grid类提供的grid.global_coords_to_tile()将坐标从全局系统转换到瓦片系统。
这是按照以下方式实现的:

  1. __attribute__((always_inline, pure))
  2. T& operator() (
  3. const uint32_t gi,
  4. const uint32_t gj,
  5. const uint32_t gk,
  6. const uint32_t n
  7. ) const
  8. {
  9. auto coords = (*m_grid).global_coords_to_tile(gi, gj, gk);
  10. return (*m_ijk)(coords[TILE_ID], coords[TILE_I], coords[TILE_J], coords[TILE_K], n);
  11. }

问题

最后,执行以下模拟计算内核(物理更新方程仅用于测试数组访问)。在发布版本中,它被自动并行化并被Kokkos内联到平铺循环中。
作为参考,这里有一个内核,它在访问所有数组之前显式地将坐标从全局系统转换到瓦片系统:

  1. void UpdateVoltagesKernel(
  2. Grid grid,
  3. ArrayIJKN<float> volt,
  4. ArrayIJKN<float> curr,
  5. ArrayIJKN<float> vv,
  6. ArrayIJKN<float> vi,
  7. uint32_t gi, uint32_t gj, uint32_t gk
  8. )
  9. {
  10. auto tile_coords = grid.global_coords_to_tile(gi, gj, gk);
  11. uint32_t tile_id = tile_coords[TILE_ID];
  12. uint32_t si = tile_coords[TILE_I];
  13. uint32_t sj = tile_coords[TILE_J];
  14. uint32_t sk = tile_coords[TILE_K];
  15. // 3 FP32 loads
  16. float volt0 = volt(tile_id, si, sj, sk, 0);
  17. float volt1 = volt(tile_id, si, sj, sk, 1);
  18. float volt2 = volt(tile_id, si, sj, sk, 2);
  19. // 3 FP32 loads
  20. float vv0 = vv(tile_id, si, sj, sk, 0);
  21. float vv1 = vv(tile_id, si, sj, sk, 1);
  22. float vv2 = vv(tile_id, si, sj, sk, 2);
  23. // 3 FP32 loads
  24. float vi0 = vi(tile_id, si, sj, sk, 0);
  25. float vi1 = vi(tile_id, si, sj, sk, 1);
  26. float vi2 = vi(tile_id, si, sj, sk, 2);
  27. // 3 FP32 loads
  28. float curr0 = curr(tile_id, si, sj, sk, 0);
  29. float curr1 = curr(tile_id, si, sj, sk, 1);
  30. float curr2 = curr(tile_id, si, sj, sk, 2);
  31. //for x polarization
  32. volt0 *= vv0;
  33. volt0 += vi0 * curr0;
  34. //for y polarization
  35. volt1 *= vv1;
  36. volt1 += vi1 * curr1;
  37. //for z polarization
  38. volt2 *= vv2;
  39. volt2 += vi2 * curr2;
  40. // 3 FP32 stores
  41. volt(tile_id, si, sj, sk, 0) = volt0;
  42. volt(tile_id, si, sj, sk, 1) = volt1;
  43. volt(tile_id, si, sj, sk, 2) = volt2;
  44. }


该内核具有良好的性能。基准测试表明,在双插槽Intel Xeon E5-2680 v4上,(我已经正确处理了NUMA第一次接触),当内核使用Kokkos::parallel_for()在所有内核之间自动并行化时,使用GCC 12.3.1 with -O3 -march=native达到80 GB/s的速度。使用LLVM/clang 18.0,速度更快,超过100 GB/s。显然,这是由于LLVM内联更积极的事实,因为我可以使用__attribute__((always_inline))复制与GCC相同的性能。100+ GB/s非常接近我的CPU的峰值内存带宽,通过类似STREAM的基准测试,如triaddaxpy,因此内核表现出很好的性能。
接下来,还测试了没有显式坐标变换的等效内核:

  1. #define ASSUME(cond) \
  2. do { \
  3. if (!(cond)) \
  4. __builtin_unreachable(); \
  5. } while (0)
  6. void UpdateVoltagesKernel(
  7. ArrayIJKN<float> volt,
  8. ArrayIJKN<float> curr,
  9. ArrayIJKN<float> vv,
  10. ArrayIJKN<float> vi,
  11. uint32_t i, uint32_t j, uint32_t k
  12. )
  13. {
  14. //ASSUME(volt.m_grid == curr.m_grid);
  15. //ASSUME(volt.m_grid == vv.m_grid);
  16. //ASSUME(volt.m_grid == vi.m_grid);
  17. // 3 FP32 loads
  18. float volt0 = volt(i, j, k, 0);
  19. float volt1 = volt(i, j, k, 1);
  20. float volt2 = volt(i, j, k, 2);
  21. // 3 FP32 loads
  22. float vv0 = vv(i, j, k, 0);
  23. float vv1 = vv(i, j, k, 1);
  24. float vv2 = vv(i, j, k, 2);
  25. // 3 FP32 loads
  26. float vi0 = vi(i, j, k, 0);
  27. float vi1 = vi(i, j, k, 1);
  28. float vi2 = vi(i, j, k, 2);
  29. // 3 FP32 loads
  30. float curr0 = curr(i, j, k, 0);
  31. float curr1 = curr(i, j, k, 1);
  32. float curr2 = curr(i, j, k, 2);
  33. //for x polarization
  34. volt0 *= vv0;
  35. volt0 += vi0 * curr0;
  36. //for y polarization
  37. volt1 *= vv1;
  38. volt1 += vi1 * curr1;
  39. //for z polarization
  40. volt2 *= vv2;
  41. volt2 += vi2 * curr2;
  42. // 3 FP32 stores
  43. volt(i, j, k, 0) = volt0;
  44. volt(i, j, k, 1) = volt1;
  45. volt(i, j, k, 2) = volt2;
  46. }


但测试表明,该内核的性能很差,运行速度只有40 GB/s,只有GCC第一个内核(80 GB/s)的一半。通过检查汇编输出,可以立即注意到编译器未能优化Grid::global_coords_to_tile的冗余隐式全局坐标函数调用,因此生成了许多额外的乘法和除法运算。
显式内核的汇编输出只有100行代码和3个部分:

  1. _Z20UpdateVoltagesKernel4Grid9ArrayIJKNIfLj3EES1_S1_S1_jjj:
  2. .LFB8832:
  3. .cfi_startproc
  4. pushq %r14
  5. .cfi_def_cfa_offset 16
  6. .cfi_offset 14, -16
  7. xorl %edx, %edx
  8. movq %rsi, %r8
  9. movq %rcx, %rdi
  10. pushq %r12
  11. .cfi_def_cfa_offset 24
  12. .cfi_offset 12, -24
  13. pushq %rbp
  14. .cfi_def_cfa_offset 32
  15. .cfi_offset 6, -32
  16. pushq %rbx
  17. .cfi_def_cfa_offset 40
  18. .cfi_offset 3, -40
  19. movl 96(%rsp), %eax
  20. movq 8(%r8), %rbx
  21. movq 88(%rsp), %r11
  22. movq 8(%r9), %r12
  23. divl 52(%rsp)
  24. movq 32(%r11), %r14
  25. movl %eax, %esi
  26. movl 104(%rsp), %eax
  27. imull 68(%rsp), %esi
  28. movl %edx, %ecx
  29. xorl %edx, %edx
  30. divl 56(%rsp)
  31. movl %eax, %ebp
  32. movl 112(%rsp), %eax
  33. movl %edx, %r10d
  34. xorl %edx, %edx
  35. addl %ebp, %esi
  36. imull 72(%rsp), %esi
  37. movq 32(%r8), %rbp
  38. divl 60(%rsp)
  39. addl %esi, %eax
  40. movl 24(%r8), %esi
  41. movl %eax, %eax
  42. imull %ecx, %esi
  43. imulq %rax, %rbp
  44. imulq %rax, %r14
  45. addl %r10d, %esi
  46. imull 28(%r8), %esi
  47. movq 32(%r9), %r8
  48. imulq %rax, %r8
  49. addl %edx, %esi
  50. movl %esi, %esi
  51. leaq (%rsi,%rsi,2), %rsi
  52. addq %rbp, %rsi
  53. salq $2, %rsi
  54. leaq (%rbx,%rsi), %rbp
  55. leaq 8(%rbx,%rsi), %rbx
  56. movl 24(%r9), %esi
  57. imull %ecx, %esi
  58. addl %r10d, %esi
  59. imull 28(%r9), %esi
  60. movq 8(%r11), %r9
  61. addl %edx, %esi
  62. movl %esi, %esi
  63. leaq (%rsi,%rsi,2), %rsi
  64. addq %r8, %rsi
  65. movl 24(%r11), %r8d
  66. imull %ecx, %r8d
  67. imull 24(%rdi), %ecx
  68. addl %r10d, %r8d
  69. imull 28(%r11), %r8d
  70. movq 8(%rdi), %r11
  71. addl %r10d, %ecx
  72. imull 28(%rdi), %ecx
  73. addl %edx, %r8d
  74. movl %r8d, %r8d
  75. leaq (%r8,%r8,2), %r8
  76. addq %r14, %r8
  77. addl %ecx, %edx
  78. vmovss 8(%r9,%r8,4), %xmm0
  79. imulq 32(%rdi), %rax
  80. movl %edx, %edx
  81. leaq (%rdx,%rdx,2), %rdx
  82. vmovq (%r9,%r8,4), %xmm4
  83. vmovq 0(%rbp), %xmm1
  84. vmovq (%r12,%rsi,4), %xmm3
  85. vmovss (%rbx), %xmm5
  86. addq %rdx, %rax
  87. vmovq (%r11,%rax,4), %xmm2
  88. vmulss 8(%r11,%rax,4), %xmm0, %xmm0
  89. vfmadd231ss 8(%r12,%rsi,4), %xmm5, %xmm0
  90. vmulps %xmm4, %xmm2, %xmm2
  91. vfmadd132ps %xmm3, %xmm2, %xmm1
  92. vmovlps %xmm1, 0(%rbp)
  93. vmovss %xmm0, (%rbx)
  94. popq %rbx
  95. .cfi_def_cfa_offset 32
  96. popq %rbp
  97. .cfi_def_cfa_offset 24
  98. popq %r12
  99. .cfi_def_cfa_offset 16
  100. popq %r14
  101. .cfi_def_cfa_offset 8
  102. ret
  103. .cfi_endproc


但是,隐式内核的汇编输出只有150行代码和12个分区,这些分区大大降低了CPU的吞吐量。

  1. _Z20UpdateVoltagesKernel9ArrayIJKNIfLj3EES0_S0_S0_jjj:
  2. .LFB8832:
  3. .cfi_startproc
  4. pushq %r15
  5. .cfi_def_cfa_offset 16
  6. .cfi_offset 15, -16
  7. pushq %r14
  8. .cfi_def_cfa_offset 24
  9. .cfi_offset 14, -24
  10. pushq %r13
  11. .cfi_def_cfa_offset 32
  12. .cfi_offset 13, -32
  13. pushq %r12
  14. .cfi_def_cfa_offset 40
  15. .cfi_offset 12, -40
  16. movq %rsi, %r12
  17. movq %rdx, %rsi
  18. xorl %edx, %edx
  19. pushq %rbp
  20. .cfi_def_cfa_offset 48
  21. .cfi_offset 6, -48
  22. movq %r9, %rbp
  23. pushq %rbx
  24. .cfi_def_cfa_offset 56
  25. .cfi_offset 3, -56
  26. movl 72(%rsp), %eax
  27. movq %rcx, %rbx
  28. movq 8(%r12), %r15
  29. movq 56(%rsp), %rcx
  30. movq 64(%rsp), %r13
  31. divl 12(%rdi)
  32. movl %eax, %r11d
  33. movl 80(%rsp), %eax
  34. imull 28(%rdi), %r11d
  35. movl %edx, %r10d
  36. xorl %edx, %edx
  37. imull 24(%r12), %r10d
  38. divl 16(%rdi)
  39. movl %eax, %r14d
  40. movl 88(%rsp), %eax
  41. movl %edx, %r9d
  42. xorl %edx, %edx
  43. addl %r14d, %r11d
  44. imull 32(%rdi), %r11d
  45. addl %r9d, %r10d
  46. divl 20(%rdi)
  47. imull 28(%r12), %r10d
  48. addl %r11d, %eax
  49. addl %r10d, %edx
  50. movl %eax, %eax
  51. imulq 32(%r12), %rax
  52. movl %edx, %edx
  53. leaq (%rdx,%rdx,2), %rdx
  54. addq %rdx, %rax
  55. xorl %edx, %edx
  56. salq $2, %rax
  57. leaq (%r15,%rax), %r11
  58. leaq 8(%r15,%rax), %r10
  59. movl 72(%rsp), %eax
  60. vmovq (%r11), %xmm1
  61. divl 12(%r8)
  62. movl %eax, %r9d
  63. movl 80(%rsp), %eax
  64. movl %edx, %edi
  65. xorl %edx, %edx
  66. divl 16(%r8)
  67. movl %eax, %r15d
  68. movl 88(%rsp), %eax
  69. movl %edx, %r14d
  70. xorl %edx, %edx
  71. divl 20(%r8)
  72. movq 8(%rbp), %r12
  73. imull 28(%r8), %r9d
  74. imull 24(%rbp), %edi
  75. addl %r15d, %r9d
  76. imull 32(%r8), %r9d
  77. movq 8(%r13), %r15
  78. addl %r14d, %edi
  79. imull 28(%rbp), %edi
  80. addl %eax, %r9d
  81. leal (%rdi,%rdx), %eax
  82. movl %r9d, %r9d
  83. imulq 32(%rbp), %r9
  84. leaq (%rax,%rax,2), %rax
  85. addq %rax, %r9
  86. movl 72(%rsp), %eax
  87. xorl %edx, %edx
  88. vmovq (%r12,%r9,4), %xmm3
  89. divl 12(%rcx)
  90. movl %eax, %r8d
  91. movl 80(%rsp), %eax
  92. imull 28(%rcx), %r8d
  93. movl %edx, %edi
  94. xorl %edx, %edx
  95. imull 24(%r13), %edi
  96. divl 16(%rcx)
  97. movl %eax, %r14d
  98. movl 88(%rsp), %eax
  99. movl %edx, %ebp
  100. xorl %edx, %edx
  101. addl %r14d, %r8d
  102. imull 32(%rcx), %r8d
  103. addl %ebp, %edi
  104. divl 20(%rcx)
  105. imull 28(%r13), %edi
  106. addl %r8d, %eax
  107. addl %edi, %edx
  108. movl %eax, %eax
  109. imulq 32(%r13), %rax
  110. movl %edx, %edx
  111. leaq (%rdx,%rdx,2), %rdx
  112. addq %rdx, %rax
  113. xorl %edx, %edx
  114. vmovq (%r15,%rax,4), %xmm4
  115. vmovss 8(%r15,%rax,4), %xmm0
  116. movl 72(%rsp), %eax
  117. divl 12(%rsi)
  118. movl %eax, %edi
  119. movl 80(%rsp), %eax
  120. movl %edx, %ecx
  121. xorl %edx, %edx
  122. divl 16(%rsi)
  123. movl %eax, %r13d
  124. movl 88(%rsp), %eax
  125. movl %edx, %r8d
  126. xorl %edx, %edx
  127. divl 20(%rsi)
  128. movq 8(%rbx), %rbp
  129. vmovss (%r10), %xmm5
  130. imull 28(%rsi), %edi
  131. imull 24(%rbx), %ecx
  132. addl %r13d, %edi
  133. imull 32(%rsi), %edi
  134. addl %r8d, %ecx
  135. imull 28(%rbx), %ecx
  136. addl %edi, %eax
  137. addl %ecx, %edx
  138. movl %eax, %eax
  139. imulq 32(%rbx), %rax
  140. movl %edx, %edx
  141. leaq (%rdx,%rdx,2), %rdx
  142. addq %rdx, %rax
  143. vmovq 0(%rbp,%rax,4), %xmm2
  144. vmulss 8(%rbp,%rax,4), %xmm0, %xmm0
  145. vfmadd231ss 8(%r12,%r9,4), %xmm5, %xmm0
  146. vmulps %xmm4, %xmm2, %xmm2
  147. vfmadd132ps %xmm3, %xmm2, %xmm1
  148. vmovlps %xmm1, (%r11)
  149. vmovss %xmm0, (%r10)
  150. popq %rbx
  151. .cfi_def_cfa_offset 48
  152. popq %rbp
  153. .cfi_def_cfa_offset 40
  154. popq %r12
  155. .cfi_def_cfa_offset 32
  156. popq %r13
  157. .cfi_def_cfa_offset 24
  158. popq %r14
  159. .cfi_def_cfa_offset 16
  160. popq %r15
  161. .cfi_def_cfa_offset 8
  162. ret
  163. .cfi_endproc

很明显,编译器没有删除对Grid::global_coords_to_tile的重复调用。因此,坐标转换被重复了很多次,而不是严格必要的。
实际上,当将以下代码添加到隐式内核中以告诉编译器volt.m_gridcurr.m_gridvv.m_gridvi.m_grid是相同的东西时,LLVM/clang能够接受此提示并像显式内核一样快速地生成代码:

  1. #define ASSUME(cond) \
  2. do { \
  3. if (!(cond)) \
  4. __builtin_unreachable(); \
  5. } while (0)
  6. ASSUME(volt.m_grid == curr.m_grid);
  7. ASSUME(volt.m_grid == vv.m_grid);
  8. ASSUME(volt.m_grid == vi.m_grid);


将这些假设添加到隐式内核中,在LLVM/clang上,速度超过100 GB/s,与显式内核一样快。不幸的是,使用GCC 12编译时,它无法理解这些暗示值得注意的是,即使将ArrayIJKN::operator()Grid::global_coords_to_tile标记为__attribute__((always_inline, pure)),情况仍然如此。但GCC在我的情况下仍然无法做到这一点。
从GCC 13开始,GCC终于正式添加了对assume的支持,但这样的GCC新版本在许多系统中不可用。
问题
是否有类似但更有效的方法来实现隐式数组访问语法array(gi, gj, gk, n)以隐藏底层坐标转换细节?
是否有一种变通方法允许GCC优化重复调用,使其不再导致性能问题?如果没有,是否有一种替代的编程风格,实现相同的语法suger,但没有这种性能损失?

附录:完整源代码

kernel.hpp

  1. enum {
  2. DIM_I = 0,
  3. DIM_J = 1,
  4. DIM_K = 2
  5. };
  6. enum {
  7. TILE_ID = 0,
  8. TILE_I = 1,
  9. TILE_J = 2,
  10. TILE_K = 3
  11. };
  12. struct Grid
  13. {
  14. Grid(
  15. std::array<uint32_t, 3> grid_size,
  16. std::array<uint32_t, 3> tile_size
  17. ) :
  18. grid_size(grid_size),
  19. tile_size(tile_size),
  20. tile_num{
  21. grid_size[DIM_I] / tile_size[DIM_I],
  22. grid_size[DIM_J] / tile_size[DIM_J],
  23. grid_size[DIM_K] / tile_size[DIM_K]
  24. }
  25. {
  26. }
  27. __attribute__((always_inline, pure))
  28. std::array<uint32_t, 4> global_coords_to_tile(
  29. const uint32_t gi,
  30. const uint32_t gj,
  31. const uint32_t gk
  32. ) const
  33. {
  34. uint32_t tile_id_i = gi / tile_size[DIM_I];
  35. uint32_t tile_id_j = gj / tile_size[DIM_J];
  36. uint32_t tile_id_k = gk / tile_size[DIM_K];
  37. uint32_t tile_id = tile_id_i * tile_num[DIM_J] * tile_num[DIM_K] + tile_id_j * tile_num[DIM_K] + tile_id_k;
  38. uint32_t ti = gi % tile_size[DIM_I];
  39. uint32_t tj = gj % tile_size[DIM_J];
  40. uint32_t tk = gk % tile_size[DIM_K];
  41. return {tile_id, ti, tj, tk};
  42. }
  43. const std::array<uint32_t, 3> grid_size;
  44. const std::array<uint32_t, 3> tile_size;
  45. const std::array<uint32_t, 3> tile_num;
  46. };
  47. template <typename T, uint32_t n_num=3>
  48. struct ArrayIJKN
  49. {
  50. ArrayIJKN(
  51. const std::string& name,
  52. const Grid& grid
  53. ) :
  54. m_grid(&grid)
  55. {
  56. uint32_t tile_num = grid.tile_num[DIM_I] * grid.tile_num[DIM_J] * grid.tile_num[DIM_K];
  57. m_ijk = new Kokkos::View<T****[n_num]>(
  58. Kokkos::ViewAllocateWithoutInitializing(name),
  59. tile_num,
  60. grid.tile_size[DIM_I], grid.tile_size[DIM_J], grid.tile_size[DIM_K], n_num
  61. );
  62. }
  63. T& operator() (
  64. const uint32_t tile_id,
  65. const uint32_t ti,
  66. const uint32_t tj,
  67. const uint32_t tk,
  68. const uint32_t n
  69. ) const
  70. {
  71. return (*m_ijk)(tile_id, ti, tj, tk, n);
  72. }
  73. __attribute__((always_inline, pure))
  74. T& operator() (
  75. const uint32_t gi,
  76. const uint32_t gj,
  77. const uint32_t gk,
  78. const uint32_t n
  79. ) const
  80. {
  81. auto coords = (*m_grid).global_coords_to_tile(gi, gj, gk);
  82. return (*m_ijk)(coords[TILE_ID], coords[TILE_I], coords[TILE_J], coords[TILE_K], n);
  83. }
  84. const Grid* __restrict__ m_grid;
  85. Kokkos::View<T****[n_num]>* __restrict__ m_ijk;
  86. };

explicit.cpp

  1. #include <cstdio>
  2. #include <cstdint>
  3. #include <Kokkos_Core.hpp>
  4. #include "kernel.hpp"
  5. void UpdateVoltagesKernel(
  6. Grid grid,
  7. ArrayIJKN<float> volt,
  8. ArrayIJKN<float> curr,
  9. ArrayIJKN<float> vv,
  10. ArrayIJKN<float> vi,
  11. uint32_t gi, uint32_t gj, uint32_t gk
  12. )
  13. {
  14. auto tile_coords = grid.global_coords_to_tile(gi, gj, gk);
  15. uint32_t tile_id = tile_coords[TILE_ID];
  16. uint32_t si = tile_coords[TILE_I];
  17. uint32_t sj = tile_coords[TILE_J];
  18. uint32_t sk = tile_coords[TILE_K];
  19. // 3 FP32 loads
  20. float volt0 = volt(tile_id, si, sj, sk, 0);
  21. float volt1 = volt(tile_id, si, sj, sk, 1);
  22. float volt2 = volt(tile_id, si, sj, sk, 2);
  23. // 3 FP32 loads
  24. float vv0 = vv(tile_id, si, sj, sk, 0);
  25. float vv1 = vv(tile_id, si, sj, sk, 1);
  26. float vv2 = vv(tile_id, si, sj, sk, 2);
  27. // 3 FP32 loads
  28. float vi0 = vi(tile_id, si, sj, sk, 0);
  29. float vi1 = vi(tile_id, si, sj, sk, 1);
  30. float vi2 = vi(tile_id, si, sj, sk, 2);
  31. // 3 FP32 loads
  32. float curr0 = curr(tile_id, si, sj, sk, 0);
  33. float curr1 = curr(tile_id, si, sj, sk, 1);
  34. float curr2 = curr(tile_id, si, sj, sk, 2);
  35. //for x polarization
  36. volt0 *= vv0;
  37. volt0 += vi0 * curr0;
  38. //for y polarization
  39. volt1 *= vv1;
  40. volt1 += vi1 * curr1;
  41. //for z polarization
  42. volt2 *= vv2;
  43. volt2 += vi2 * curr2;
  44. // 3 FP32 stores
  45. volt(tile_id, si, sj, sk, 0) = volt0;
  46. volt(tile_id, si, sj, sk, 1) = volt1;
  47. volt(tile_id, si, sj, sk, 2) = volt2;
  48. }
  49. int main(int argc, char** argv)
  50. {
  51. std::array<uint32_t, 3> numLines = {
  52. 200, 200, 200
  53. };
  54. std::array<uint32_t, 3> dimTiles = {
  55. 20, 20, 20
  56. };
  57. const int timesteps = 100;
  58. Kokkos::initialize(argc, argv);
  59. auto grid = Grid(numLines, dimTiles);
  60. auto volt = ArrayIJKN<float>("volt", grid);
  61. auto curr = ArrayIJKN<float>("curr", grid);
  62. auto vv = ArrayIJKN<float>("vv", grid);
  63. auto vi = ArrayIJKN<float>("vi", grid);
  64. auto ii = ArrayIJKN<float>("ii", grid);
  65. auto iv = ArrayIJKN<float>("iv", grid);
  66. Kokkos::parallel_for("First Touch",
  67. Kokkos::MDRangePolicy<Kokkos::Rank<3>> (
  68. {0, 0, 0},
  69. {numLines[0], numLines[1], numLines[2]},
  70. {dimTiles[0], dimTiles[1], dimTiles[2]}
  71. ),
  72. [=](
  73. const uint32_t i,
  74. const uint32_t j,
  75. const uint32_t k
  76. ) {
  77. for (uint32_t n = 0; n < 3; n++) {
  78. volt(i, j, k, n) = 0;
  79. curr(i, j, k, n) = 0;
  80. vv(i, j, k, n) = 0;
  81. vi(i, j, k, n) = 0;
  82. iv(i, j, k, n) = 0;
  83. ii(i, j, k, n) = 0;
  84. }
  85. }
  86. );
  87. Kokkos::fence();
  88. auto t1 = std::chrono::high_resolution_clock().now();
  89. for (size_t i = 0; i < timesteps; i++) {
  90. Kokkos::parallel_for("UpdateVoltages",
  91. Kokkos::MDRangePolicy<Kokkos::Rank<3>> (
  92. {0, 0, 0},
  93. {numLines[0], numLines[1], numLines[2]},
  94. {dimTiles[0], dimTiles[1], dimTiles[2]}
  95. ),
  96. [=](
  97. const uint32_t gi,
  98. const uint32_t gj,
  99. const uint32_t gk
  100. ) {
  101. UpdateVoltagesKernel(grid, volt, curr, vv, vi, gi, gj, gk);
  102. }
  103. );
  104. Kokkos::fence();
  105. }
  106. Kokkos::fence();
  107. auto t2 = std::chrono::high_resolution_clock().now();
  108. double dt = std::chrono::duration_cast<std::chrono::microseconds>(t2 - t1).count() / 1e6;
  109. size_t bytes_per_iteration = numLines[0] * numLines[1] * numLines[2] * 3 * sizeof(float);
  110. bytes_per_iteration *= 5; /* volt r+w, curr, vv, vi arrays */
  111. fprintf(stderr, "traffic: %.1f GB/s\n", (double) bytes_per_iteration * timesteps / 1000 / 1000 / 1000);
  112. fprintf(stderr, "speed: %.1f GB/s\n",
  113. (double) bytes_per_iteration * timesteps / dt / 1000 / 1000 / 1000
  114. );
  115. Kokkos::finalize();
  116. }

implicit.cpp

  1. #include <cstdio>
  2. #include <cstdint>
  3. #include <Kokkos_Core.hpp>
  4. #include "kernel.hpp"
  5. #define ASSUME(cond) \
  6. do { \
  7. if (!(cond)) \
  8. __builtin_unreachable(); \
  9. } while (0)
  10. void UpdateVoltagesKernel(
  11. ArrayIJKN<float> volt,
  12. ArrayIJKN<float> curr,
  13. ArrayIJKN<float> vv,
  14. ArrayIJKN<float> vi,
  15. uint32_t i, uint32_t j, uint32_t k
  16. )
  17. {
  18. ASSUME(volt.m_grid == curr.m_grid);
  19. ASSUME(volt.m_grid == vv.m_grid);
  20. ASSUME(volt.m_grid == vi.m_grid);
  21. // 3 FP32 loads
  22. float volt0 = volt(i, j, k, 0);
  23. float volt1 = volt(i, j, k, 1);
  24. float volt2 = volt(i, j, k, 2);
  25. // 3 FP32 loads
  26. float vv0 = vv(i, j, k, 0);
  27. float vv1 = vv(i, j, k, 1);
  28. float vv2 = vv(i, j, k, 2);
  29. // 3 FP32 loads
  30. float vi0 = vi(i, j, k, 0);
  31. float vi1 = vi(i, j, k, 1);
  32. float vi2 = vi(i, j, k, 2);
  33. // 3 FP32 loads
  34. float curr0 = curr(i, j, k, 0);
  35. float curr1 = curr(i, j, k, 1);
  36. float curr2 = curr(i, j, k, 2);
  37. //for x polarization
  38. volt0 *= vv0;
  39. volt0 += vi0 * curr0;
  40. //for y polarization
  41. volt1 *= vv1;
  42. volt1 += vi1 * curr1;
  43. //for z polarization
  44. volt2 *= vv2;
  45. volt2 += vi2 * curr2;
  46. // 3 FP32 stores
  47. volt(i, j, k, 0) = volt0;
  48. volt(i, j, k, 1) = volt1;
  49. volt(i, j, k, 2) = volt2;
  50. }
  51. int main(int argc, char** argv)
  52. {
  53. std::array<uint32_t, 3> numLines = {
  54. 200, 200, 200
  55. };
  56. std::array<uint32_t, 3> dimTiles = {
  57. 20, 20, 20
  58. };
  59. const int timesteps = 100;
  60. Kokkos::initialize(argc, argv);
  61. auto grid = Grid(numLines, dimTiles);
  62. auto volt = ArrayIJKN<float>("volt", grid);
  63. auto curr = ArrayIJKN<float>("curr", grid);
  64. auto vv = ArrayIJKN<float>("vv", grid);
  65. auto vi = ArrayIJKN<float>("vi", grid);
  66. auto ii = ArrayIJKN<float>("ii", grid);
  67. auto iv = ArrayIJKN<float>("iv", grid);
  68. Kokkos::parallel_for("First Touch",
  69. Kokkos::MDRangePolicy<Kokkos::Rank<3>> (
  70. {0, 0, 0},
  71. {numLines[0], numLines[1], numLines[2]},
  72. {dimTiles[0], dimTiles[1], dimTiles[2]}
  73. ),
  74. [=](
  75. const uint32_t i,
  76. const uint32_t j,
  77. const uint32_t k
  78. ) {
  79. for (uint32_t n = 0; n < 3; n++) {
  80. volt(i, j, k, n) = 0;
  81. curr(i, j, k, n) = 0;
  82. vv(i, j, k, n) = 0;
  83. vi(i, j, k, n) = 0;
  84. iv(i, j, k, n) = 0;
  85. ii(i, j, k, n) = 0;
  86. }
  87. }
  88. );
  89. Kokkos::fence();
  90. auto t1 = std::chrono::high_resolution_clock().now();
  91. for (size_t i = 0; i < timesteps; i++) {
  92. Kokkos::parallel_for("UpdateVoltages",
  93. Kokkos::MDRangePolicy<Kokkos::Rank<3>> (
  94. {0, 0, 0},
  95. {numLines[0], numLines[1], numLines[2]},
  96. {dimTiles[0], dimTiles[1], dimTiles[2]}
  97. ),
  98. [=](
  99. const uint32_t i,
  100. const uint32_t j,
  101. const uint32_t k
  102. ) {
  103. UpdateVoltagesKernel(volt, curr, vv, vi, i, j, k);
  104. }
  105. );
  106. Kokkos::fence();
  107. }
  108. Kokkos::fence();
  109. auto t2 = std::chrono::high_resolution_clock().now();
  110. double dt = std::chrono::duration_cast<std::chrono::microseconds>(t2 - t1).count() / 1e6;
  111. size_t bytes_per_iteration = numLines[0] * numLines[1] * numLines[2] * 3 * sizeof(float);
  112. bytes_per_iteration *= 5; /* volt r+w, curr, vv, vi arrays */
  113. fprintf(stderr, "traffic: %.1f GB/s\n", (double) bytes_per_iteration * timesteps / 1000 / 1000 / 1000);
  114. fprintf(stderr, "speed: %.1f GB/s\n",
  115. (double) bytes_per_iteration * timesteps / dt / 1000 / 1000 / 1000
  116. );
  117. Kokkos::finalize();
  118. }

CMakeLists.txt

  1. cmake_minimum_required(VERSION 3.25)
  2. project(fdtd-kokkos CXX)
  3. set(KOKKOS_IN_TREE "/home/fdtd/kokkos")
  4. add_subdirectory(${KOKKOS_IN_TREE} ${PROJECT_BINARY_DIR}/kokkos)
  5. file (GLOB project_SOURCE ${project_SOURCE_DIR}/*.cpp)
  6. foreach (sourcefile ${project_SOURCE})
  7. get_filename_component(sourcename ${sourcefile} NAME_WE)
  8. add_executable(${sourcename} ${sourcefile})
  9. target_link_libraries(
  10. ${sourcename} Kokkos::kokkoscore Kokkos::kokkossimd
  11. )
  12. target_include_directories(
  13. ${sourcename} PUBLIC ${project_SOURCE_DIR}
  14. )
  15. endforeach (sourcefile ${project_SOURCE_DIR})


要复制构建,请将源代码保存到project/src/,将CMakeLists.txt保存到project/CMakeLists.txt,将Kokkos的副本(从https://github.com/kokkos/kokkos.git)保存到/home/fdtd/kokkos。然后:

  1. export CFLAGS="-O3 -march=native"
  2. export CXXFLAGS="-O3 -march=native"
  3. mkdir build && cd build
  4. cmake ../ -DKokkos_ENABLE_THREADS=ON -DKokkos_ENABLE_HWLOC=ON

ni65a41a

ni65a41a1#

这里的困难是允许GCC推断出每个ArrayIJKN(包括voltcurrvvvi)持有的成员变量m_grid指向Grid的同一个示例。
到目前为止,我无法找到一个满意的解决方案,但以下是部分解决方案。

abort()return(如果违反假设)

一个能够更好地生成GCC代码的部分解决方案是改变ASSUME的定义。当假设被打破时,我们也可以首先从内核函数调用abort()return;,而不是调用__builtin_unreachable();

  1. #define ASSUME(cond) \
  2. do { \
  3. if (!(cond)) { \
  4. abort(); \
  5. __builtin_unreachable(); \
  6. } \
  7. } while (0)

字符串
此更改允许隐式计算内核达到80 GB/s的速度,与显式版本的非内联版本一样快。
不幸的是,虽然90 GB/s是可能的,但这两种解决方案都无法达到完整的110 GB/s,无论是否在GCC 12中强制内联。同时,在GCC 13中,这种解决方案允许GCC的assume像clang一样工作,并且能够达到110 GB/s。

模板grid

到目前为止,我发现的最好但不令人满意的解决方法是将gridArrayIJKN的运行时参数更改为可以在编译时确定的模板参数。
因此,而不是:

  1. template <typename T, uint32_t n_num=3>
  2. struct ArrayIJKN
  3. {
  4. ArrayIJKN(
  5. const std::string& name,
  6. const Grid& grid
  7. ) :


我用途:

  1. template <typename T, Grid& grid, uint32_t n_num=3>
  2. struct ArrayIJKN
  3. {
  4. ArrayIJKN(
  5. const std::string& name
  6. )


然后将对象Grid重新定义为全局变量:

  1. std::array<uint32_t, 3> numLines = {
  2. 200, 200, 200
  3. };
  4. std::array<uint32_t, 3> dimTiles = {
  5. 20, 20, 20
  6. };
  7. auto grid = Grid(numLines, dimTiles);
  8. // ...
  9. int main(int argc, char** argv)
  10. {
  11. Kokkos::initialize(argc, argv);
  12. auto volt = ArrayIJKN<float, grid>("volt");
  13. auto curr = ArrayIJKN<float, grid>("curr");
  14. auto vv = ArrayIJKN<float, grid>("vv");
  15. auto vi = ArrayIJKN<float, grid>("vi");
  16. auto ii = ArrayIJKN<float, grid>("ii");
  17. auto iv = ArrayIJKN<float, grid>("iv");
  18. }


然后将核函数重新定义为:

  1. void UpdateVoltagesKernel(
  2. const ArrayIJKN<float, grid>& __restrict__ volt,
  3. const ArrayIJKN<float, grid>& __restrict__ curr,
  4. const ArrayIJKN<float, grid>& __restrict__ vv,
  5. const ArrayIJKN<float, grid>& __restrict__ vi,
  6. uint32_t i, uint32_t j, uint32_t k
  7. )


在这个版本中,GCC的目标代码已经恢复到80 GB/s的速度,但无法达到100 GB/s的全速。似乎数组 Package 函数中的强制内联在这种情况下适得其反。在从所有成员函数中删除__attribute__((always_inline, pure)),同时仍然保留always_inline用于内核函数之后,现在,GCC能够生成以110+ GB/s运行的高性能目标代码,没有问题,与LLVM或显式版本相当。
但是这种解决方法远远不能令人满意--变量Grid必须全局定义,并且必须在定义任何其他函数之前出现。

展开查看全部

相关问题