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

k3bvogb1  于 12个月前  发布在  其他
关注(0)|答案(1)|浏览(143)

背景和动机

我正在优化一个数值解算器,用有限差分法模拟麦克斯韦方程。
请注意,实际问题与物理、数值模拟甚至相关的优化技术都没有太大关系。实际上,这只是在实现自定义数组 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类,这个类只用来表示抽象的仿真网格:

Grid(
            std::array<uint32_t, 3> grid_size,
            std::array<uint32_t, 3> tile_size
    ) :

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

__attribute__((always_inline, pure))
    std::array<uint32_t, 4> global_coords_to_tile(
            const uint32_t gi,
            const uint32_t gj,
            const uint32_t gk
    ) const


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

ArrayIJKN(
        const std::string& name,
        const Grid& grid
    ) :
        m_grid(&grid)
    {
        uint32_t tile_num = grid.tile_num[DIM_I] * grid.tile_num[DIM_J] * grid.tile_num[DIM_K];

        // Yes, it's 5-level of indirection and can be considered a bad style
        // elsewhere, but this Kokkos::View template is officially designed to be
        // used this way. It templated for compile-time support of N-D arrays up
        // to 8 dimensions. 
        // 
        // The first dimension select a tile, the next 3 dimensions select a
        // cell within the 3D tile, the final dimension select the polarization
        // (X, Y, Z).
        m_ijk = new Kokkos::View<T****[n_num]>(
            Kokkos::ViewAllocateWithoutInitializing(name),
            tile_num,
            grid.tile_size[DIM_I], grid.tile_size[DIM_J], grid.tile_size[DIM_K],
            n_num
        );
    }


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

__attribute__((always_inline, pure))
    T& operator() (
            const uint32_t gi,
            const uint32_t gj,
            const uint32_t gk,
            const uint32_t n
    ) const
    {
            auto coords = (*m_grid).global_coords_to_tile(gi, gj, gk);
            return (*m_ijk)(coords[TILE_ID], coords[TILE_I], coords[TILE_J], coords[TILE_K], n);
    }

问题

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

void UpdateVoltagesKernel(
    Grid grid,
    ArrayIJKN<float> volt,
    ArrayIJKN<float> curr,
    ArrayIJKN<float> vv,
    ArrayIJKN<float> vi,
    uint32_t gi, uint32_t gj, uint32_t gk
)
{
    auto tile_coords = grid.global_coords_to_tile(gi, gj, gk);
    uint32_t tile_id = tile_coords[TILE_ID];
    uint32_t si = tile_coords[TILE_I];
    uint32_t sj = tile_coords[TILE_J];
    uint32_t sk = tile_coords[TILE_K];

    // 3 FP32 loads
    float volt0 = volt(tile_id, si, sj, sk, 0);
    float volt1 = volt(tile_id, si, sj, sk, 1);
    float volt2 = volt(tile_id, si, sj, sk, 2);

    // 3 FP32 loads
    float vv0 = vv(tile_id, si, sj, sk, 0);
    float vv1 = vv(tile_id, si, sj, sk, 1);
    float vv2 = vv(tile_id, si, sj, sk, 2);

    // 3 FP32 loads
    float vi0 = vi(tile_id, si, sj, sk, 0);
    float vi1 = vi(tile_id, si, sj, sk, 1);
    float vi2 = vi(tile_id, si, sj, sk, 2);

    // 3 FP32 loads
    float curr0 = curr(tile_id, si, sj, sk, 0);
    float curr1 = curr(tile_id, si, sj, sk, 1);
    float curr2 = curr(tile_id, si, sj, sk, 2);

    //for x polarization
    volt0 *= vv0;
    volt0 += vi0 * curr0;

    //for y polarization
    volt1 *= vv1;
    volt1 += vi1 * curr1;

    //for z polarization
    volt2 *= vv2;
    volt2 += vi2 * curr2;

    // 3 FP32 stores
    volt(tile_id, si, sj, sk, 0) = volt0;
    volt(tile_id, si, sj, sk, 1) = volt1;
    volt(tile_id, si, sj, sk, 2) = volt2;
}


该内核具有良好的性能。基准测试表明,在双插槽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,因此内核表现出很好的性能。
接下来,还测试了没有显式坐标变换的等效内核:

#define ASSUME(cond)                    \
    do {                                \
        if (!(cond))                    \
            __builtin_unreachable();    \
    } while (0)

void UpdateVoltagesKernel(
    ArrayIJKN<float> volt,
    ArrayIJKN<float> curr,
    ArrayIJKN<float> vv,
    ArrayIJKN<float> vi,
    uint32_t i, uint32_t j, uint32_t k
)
{
    //ASSUME(volt.m_grid == curr.m_grid);
    //ASSUME(volt.m_grid == vv.m_grid);
    //ASSUME(volt.m_grid == vi.m_grid);

    // 3 FP32 loads
    float volt0 = volt(i, j, k, 0);
    float volt1 = volt(i, j, k, 1);
    float volt2 = volt(i, j, k, 2);

    // 3 FP32 loads
    float vv0 = vv(i, j, k, 0);
    float vv1 = vv(i, j, k, 1);
    float vv2 = vv(i, j, k, 2);

    // 3 FP32 loads
    float vi0 = vi(i, j, k, 0);
    float vi1 = vi(i, j, k, 1);
    float vi2 = vi(i, j, k, 2);

    // 3 FP32 loads
    float curr0 = curr(i, j, k, 0);
    float curr1 = curr(i, j, k, 1);
    float curr2 = curr(i, j, k, 2);

    //for x polarization
    volt0 *= vv0;
    volt0 += vi0 * curr0;

    //for y polarization
    volt1 *= vv1;
    volt1 += vi1 * curr1;

    //for z polarization
    volt2 *= vv2;
    volt2 += vi2 * curr2;

    // 3 FP32 stores
    volt(i, j, k, 0) = volt0;
    volt(i, j, k, 1) = volt1;
    volt(i, j, k, 2) = volt2;
}


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

_Z20UpdateVoltagesKernel4Grid9ArrayIJKNIfLj3EES1_S1_S1_jjj:
.LFB8832:
    .cfi_startproc
    pushq   %r14
    .cfi_def_cfa_offset 16
    .cfi_offset 14, -16
    xorl    %edx, %edx
    movq    %rsi, %r8
    movq    %rcx, %rdi
    pushq   %r12
    .cfi_def_cfa_offset 24
    .cfi_offset 12, -24
    pushq   %rbp
    .cfi_def_cfa_offset 32
    .cfi_offset 6, -32
    pushq   %rbx
    .cfi_def_cfa_offset 40
    .cfi_offset 3, -40
    movl    96(%rsp), %eax
    movq    8(%r8), %rbx
    movq    88(%rsp), %r11
    movq    8(%r9), %r12
    divl    52(%rsp)
    movq    32(%r11), %r14
    movl    %eax, %esi
    movl    104(%rsp), %eax
    imull   68(%rsp), %esi
    movl    %edx, %ecx
    xorl    %edx, %edx
    divl    56(%rsp)
    movl    %eax, %ebp
    movl    112(%rsp), %eax
    movl    %edx, %r10d
    xorl    %edx, %edx
    addl    %ebp, %esi
    imull   72(%rsp), %esi
    movq    32(%r8), %rbp
    divl    60(%rsp)
    addl    %esi, %eax
    movl    24(%r8), %esi
    movl    %eax, %eax
    imull   %ecx, %esi
    imulq   %rax, %rbp
    imulq   %rax, %r14
    addl    %r10d, %esi
    imull   28(%r8), %esi
    movq    32(%r9), %r8
    imulq   %rax, %r8
    addl    %edx, %esi
    movl    %esi, %esi
    leaq    (%rsi,%rsi,2), %rsi
    addq    %rbp, %rsi
    salq    $2, %rsi
    leaq    (%rbx,%rsi), %rbp
    leaq    8(%rbx,%rsi), %rbx
    movl    24(%r9), %esi
    imull   %ecx, %esi
    addl    %r10d, %esi
    imull   28(%r9), %esi
    movq    8(%r11), %r9
    addl    %edx, %esi
    movl    %esi, %esi
    leaq    (%rsi,%rsi,2), %rsi
    addq    %r8, %rsi
    movl    24(%r11), %r8d
    imull   %ecx, %r8d
    imull   24(%rdi), %ecx
    addl    %r10d, %r8d
    imull   28(%r11), %r8d
    movq    8(%rdi), %r11
    addl    %r10d, %ecx
    imull   28(%rdi), %ecx
    addl    %edx, %r8d
    movl    %r8d, %r8d
    leaq    (%r8,%r8,2), %r8
    addq    %r14, %r8
    addl    %ecx, %edx
    vmovss  8(%r9,%r8,4), %xmm0
    imulq   32(%rdi), %rax
    movl    %edx, %edx
    leaq    (%rdx,%rdx,2), %rdx
    vmovq   (%r9,%r8,4), %xmm4
    vmovq   0(%rbp), %xmm1
    vmovq   (%r12,%rsi,4), %xmm3
    vmovss  (%rbx), %xmm5
    addq    %rdx, %rax
    vmovq   (%r11,%rax,4), %xmm2
    vmulss  8(%r11,%rax,4), %xmm0, %xmm0
    vfmadd231ss 8(%r12,%rsi,4), %xmm5, %xmm0
    vmulps  %xmm4, %xmm2, %xmm2
    vfmadd132ps %xmm3, %xmm2, %xmm1
    vmovlps %xmm1, 0(%rbp)
    vmovss  %xmm0, (%rbx)
    popq    %rbx
    .cfi_def_cfa_offset 32
    popq    %rbp
    .cfi_def_cfa_offset 24
    popq    %r12
    .cfi_def_cfa_offset 16
    popq    %r14
    .cfi_def_cfa_offset 8
    ret
    .cfi_endproc


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

_Z20UpdateVoltagesKernel9ArrayIJKNIfLj3EES0_S0_S0_jjj:
.LFB8832:
    .cfi_startproc
    pushq   %r15
    .cfi_def_cfa_offset 16
    .cfi_offset 15, -16
    pushq   %r14
    .cfi_def_cfa_offset 24
    .cfi_offset 14, -24
    pushq   %r13
    .cfi_def_cfa_offset 32
    .cfi_offset 13, -32
    pushq   %r12
    .cfi_def_cfa_offset 40
    .cfi_offset 12, -40
    movq    %rsi, %r12
    movq    %rdx, %rsi
    xorl    %edx, %edx
    pushq   %rbp
    .cfi_def_cfa_offset 48
    .cfi_offset 6, -48
    movq    %r9, %rbp
    pushq   %rbx
    .cfi_def_cfa_offset 56
    .cfi_offset 3, -56
    movl    72(%rsp), %eax
    movq    %rcx, %rbx
    movq    8(%r12), %r15
    movq    56(%rsp), %rcx
    movq    64(%rsp), %r13
    divl    12(%rdi)
    movl    %eax, %r11d
    movl    80(%rsp), %eax
    imull   28(%rdi), %r11d
    movl    %edx, %r10d
    xorl    %edx, %edx
    imull   24(%r12), %r10d
    divl    16(%rdi)
    movl    %eax, %r14d
    movl    88(%rsp), %eax
    movl    %edx, %r9d
    xorl    %edx, %edx
    addl    %r14d, %r11d
    imull   32(%rdi), %r11d
    addl    %r9d, %r10d
    divl    20(%rdi)
    imull   28(%r12), %r10d
    addl    %r11d, %eax
    addl    %r10d, %edx
    movl    %eax, %eax
    imulq   32(%r12), %rax
    movl    %edx, %edx
    leaq    (%rdx,%rdx,2), %rdx
    addq    %rdx, %rax
    xorl    %edx, %edx
    salq    $2, %rax
    leaq    (%r15,%rax), %r11
    leaq    8(%r15,%rax), %r10
    movl    72(%rsp), %eax
    vmovq   (%r11), %xmm1
    divl    12(%r8)
    movl    %eax, %r9d
    movl    80(%rsp), %eax
    movl    %edx, %edi
    xorl    %edx, %edx
    divl    16(%r8)
    movl    %eax, %r15d
    movl    88(%rsp), %eax
    movl    %edx, %r14d
    xorl    %edx, %edx
    divl    20(%r8)
    movq    8(%rbp), %r12
    imull   28(%r8), %r9d
    imull   24(%rbp), %edi
    addl    %r15d, %r9d
    imull   32(%r8), %r9d
    movq    8(%r13), %r15
    addl    %r14d, %edi
    imull   28(%rbp), %edi
    addl    %eax, %r9d
    leal    (%rdi,%rdx), %eax
    movl    %r9d, %r9d
    imulq   32(%rbp), %r9
    leaq    (%rax,%rax,2), %rax
    addq    %rax, %r9
    movl    72(%rsp), %eax
    xorl    %edx, %edx
    vmovq   (%r12,%r9,4), %xmm3
    divl    12(%rcx)
    movl    %eax, %r8d
    movl    80(%rsp), %eax
    imull   28(%rcx), %r8d
    movl    %edx, %edi
    xorl    %edx, %edx
    imull   24(%r13), %edi
    divl    16(%rcx)
    movl    %eax, %r14d
    movl    88(%rsp), %eax
    movl    %edx, %ebp
    xorl    %edx, %edx
    addl    %r14d, %r8d
    imull   32(%rcx), %r8d
    addl    %ebp, %edi
    divl    20(%rcx)
    imull   28(%r13), %edi
    addl    %r8d, %eax
    addl    %edi, %edx
    movl    %eax, %eax
    imulq   32(%r13), %rax
    movl    %edx, %edx
    leaq    (%rdx,%rdx,2), %rdx
    addq    %rdx, %rax
    xorl    %edx, %edx
    vmovq   (%r15,%rax,4), %xmm4
    vmovss  8(%r15,%rax,4), %xmm0
    movl    72(%rsp), %eax
    divl    12(%rsi)
    movl    %eax, %edi
    movl    80(%rsp), %eax
    movl    %edx, %ecx
    xorl    %edx, %edx
    divl    16(%rsi)
    movl    %eax, %r13d
    movl    88(%rsp), %eax
    movl    %edx, %r8d
    xorl    %edx, %edx
    divl    20(%rsi)
    movq    8(%rbx), %rbp
    vmovss  (%r10), %xmm5
    imull   28(%rsi), %edi
    imull   24(%rbx), %ecx
    addl    %r13d, %edi
    imull   32(%rsi), %edi
    addl    %r8d, %ecx
    imull   28(%rbx), %ecx
    addl    %edi, %eax
    addl    %ecx, %edx
    movl    %eax, %eax
    imulq   32(%rbx), %rax
    movl    %edx, %edx
    leaq    (%rdx,%rdx,2), %rdx
    addq    %rdx, %rax
    vmovq   0(%rbp,%rax,4), %xmm2
    vmulss  8(%rbp,%rax,4), %xmm0, %xmm0
    vfmadd231ss 8(%r12,%r9,4), %xmm5, %xmm0
    vmulps  %xmm4, %xmm2, %xmm2
    vfmadd132ps %xmm3, %xmm2, %xmm1
    vmovlps %xmm1, (%r11)
    vmovss  %xmm0, (%r10)
    popq    %rbx
    .cfi_def_cfa_offset 48
    popq    %rbp
    .cfi_def_cfa_offset 40
    popq    %r12
    .cfi_def_cfa_offset 32
    popq    %r13
    .cfi_def_cfa_offset 24
    popq    %r14
    .cfi_def_cfa_offset 16
    popq    %r15
    .cfi_def_cfa_offset 8
    ret
    .cfi_endproc

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

#define ASSUME(cond)                    \
    do {                                \
        if (!(cond))                    \
            __builtin_unreachable();    \
    } while (0)

    ASSUME(volt.m_grid == curr.m_grid);
    ASSUME(volt.m_grid == vv.m_grid);
    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

enum {
    DIM_I = 0,
    DIM_J = 1,
    DIM_K = 2
};

enum {
    TILE_ID = 0,
    TILE_I = 1,
    TILE_J = 2,
    TILE_K = 3
};

struct Grid
{
    Grid(
        std::array<uint32_t, 3> grid_size,
        std::array<uint32_t, 3> tile_size
    ) :
        grid_size(grid_size),
        tile_size(tile_size),
        tile_num{
            grid_size[DIM_I] / tile_size[DIM_I],
                        grid_size[DIM_J] / tile_size[DIM_J],
                        grid_size[DIM_K] / tile_size[DIM_K]
        }
    {
    }

    __attribute__((always_inline, pure))
    std::array<uint32_t, 4> global_coords_to_tile(
        const uint32_t gi,
        const uint32_t gj,
        const uint32_t gk
    ) const
    {
        uint32_t tile_id_i = gi / tile_size[DIM_I];
        uint32_t tile_id_j = gj / tile_size[DIM_J];
        uint32_t tile_id_k = gk / tile_size[DIM_K];
        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;

        uint32_t ti = gi % tile_size[DIM_I];
        uint32_t tj = gj % tile_size[DIM_J];
        uint32_t tk = gk % tile_size[DIM_K];

        return {tile_id, ti, tj, tk};
    }

    const std::array<uint32_t, 3> grid_size;
    const std::array<uint32_t, 3> tile_size;
    const std::array<uint32_t, 3> tile_num;
};

template <typename T, uint32_t n_num=3>
struct ArrayIJKN
{
    ArrayIJKN(
        const std::string& name,
        const Grid& grid
    ) :
        m_grid(&grid)
    {
        uint32_t tile_num = grid.tile_num[DIM_I] * grid.tile_num[DIM_J] * grid.tile_num[DIM_K];

        m_ijk = new Kokkos::View<T****[n_num]>(
            Kokkos::ViewAllocateWithoutInitializing(name),
            tile_num,
            grid.tile_size[DIM_I], grid.tile_size[DIM_J], grid.tile_size[DIM_K], n_num
        );
    }

    T& operator() (
        const uint32_t tile_id,
        const uint32_t ti,
        const uint32_t tj,
        const uint32_t tk,
        const uint32_t n
    ) const
    {
        return (*m_ijk)(tile_id, ti, tj, tk, n);
    }

    __attribute__((always_inline, pure))
    T& operator() (
        const uint32_t gi,
        const uint32_t gj,
        const uint32_t gk,
        const uint32_t n
    ) const
    {
        auto coords = (*m_grid).global_coords_to_tile(gi, gj, gk);
        return (*m_ijk)(coords[TILE_ID], coords[TILE_I], coords[TILE_J], coords[TILE_K], n);
    }

    const Grid* __restrict__ m_grid;
    Kokkos::View<T****[n_num]>* __restrict__ m_ijk;
};

explicit.cpp

#include <cstdio>
#include <cstdint>
#include <Kokkos_Core.hpp>

#include "kernel.hpp"

void UpdateVoltagesKernel(
    Grid grid,
    ArrayIJKN<float> volt,
    ArrayIJKN<float> curr,
    ArrayIJKN<float> vv,
    ArrayIJKN<float> vi,
    uint32_t gi, uint32_t gj, uint32_t gk
)
{
    auto tile_coords = grid.global_coords_to_tile(gi, gj, gk);
    uint32_t tile_id = tile_coords[TILE_ID];
    uint32_t si = tile_coords[TILE_I];
    uint32_t sj = tile_coords[TILE_J];
    uint32_t sk = tile_coords[TILE_K];

    // 3 FP32 loads
    float volt0 = volt(tile_id, si, sj, sk, 0);
    float volt1 = volt(tile_id, si, sj, sk, 1);
    float volt2 = volt(tile_id, si, sj, sk, 2);

    // 3 FP32 loads
    float vv0 = vv(tile_id, si, sj, sk, 0);
    float vv1 = vv(tile_id, si, sj, sk, 1);
    float vv2 = vv(tile_id, si, sj, sk, 2);

    // 3 FP32 loads
    float vi0 = vi(tile_id, si, sj, sk, 0);
    float vi1 = vi(tile_id, si, sj, sk, 1);
    float vi2 = vi(tile_id, si, sj, sk, 2);

    // 3 FP32 loads
    float curr0 = curr(tile_id, si, sj, sk, 0);
    float curr1 = curr(tile_id, si, sj, sk, 1);
    float curr2 = curr(tile_id, si, sj, sk, 2);

    //for x polarization
    volt0 *= vv0;
    volt0 += vi0 * curr0;

    //for y polarization
    volt1 *= vv1;
    volt1 += vi1 * curr1;

    //for z polarization
    volt2 *= vv2;
    volt2 += vi2 * curr2;

    // 3 FP32 stores
    volt(tile_id, si, sj, sk, 0) = volt0;
    volt(tile_id, si, sj, sk, 1) = volt1;
    volt(tile_id, si, sj, sk, 2) = volt2;
}

int main(int argc, char** argv)
{
    std::array<uint32_t, 3> numLines = {
        200, 200, 200
    };
    std::array<uint32_t, 3> dimTiles = {
        20, 20, 20
    };

    const int timesteps = 100;

    Kokkos::initialize(argc, argv);

    auto grid = Grid(numLines, dimTiles);

    auto volt = ArrayIJKN<float>("volt", grid);
    auto curr = ArrayIJKN<float>("curr", grid);

    auto vv = ArrayIJKN<float>("vv", grid);
    auto vi = ArrayIJKN<float>("vi", grid);
    auto ii = ArrayIJKN<float>("ii", grid);
    auto iv = ArrayIJKN<float>("iv", grid);

    Kokkos::parallel_for("First Touch",
        Kokkos::MDRangePolicy<Kokkos::Rank<3>> (
            {0, 0, 0},
            {numLines[0], numLines[1], numLines[2]},
            {dimTiles[0], dimTiles[1], dimTiles[2]}
        ),
        [=](
            const uint32_t i,
            const uint32_t j,
            const uint32_t k
        ) {
            for (uint32_t n = 0; n < 3; n++) {
                volt(i, j, k, n) = 0;
                curr(i, j, k, n) = 0;

                vv(i, j, k, n) = 0;
                vi(i, j, k, n) = 0;
                iv(i, j, k, n) = 0;
                ii(i, j, k, n) = 0;
            }
        }
    );
    Kokkos::fence();

    auto t1 = std::chrono::high_resolution_clock().now();

    for (size_t i = 0; i < timesteps; i++) {
        Kokkos::parallel_for("UpdateVoltages",
            Kokkos::MDRangePolicy<Kokkos::Rank<3>> (
                {0, 0, 0},
                {numLines[0], numLines[1], numLines[2]},
                {dimTiles[0], dimTiles[1], dimTiles[2]}
            ),
            [=](
                const uint32_t gi,
                const uint32_t gj,
                const uint32_t gk
            ) {
                UpdateVoltagesKernel(grid, volt, curr, vv, vi, gi, gj, gk);
            }
        );

        Kokkos::fence();
    }

    Kokkos::fence();
    auto t2 = std::chrono::high_resolution_clock().now();
    double dt = std::chrono::duration_cast<std::chrono::microseconds>(t2 - t1).count() / 1e6;

    size_t bytes_per_iteration = numLines[0] * numLines[1] * numLines[2] * 3 * sizeof(float);
    bytes_per_iteration *= 5;  /* volt r+w, curr, vv, vi arrays */
    fprintf(stderr, "traffic: %.1f GB/s\n", (double) bytes_per_iteration * timesteps / 1000 / 1000 / 1000);
    fprintf(stderr, "speed: %.1f GB/s\n",
        (double) bytes_per_iteration * timesteps / dt / 1000 / 1000 / 1000
    );

    Kokkos::finalize();
}

implicit.cpp

#include <cstdio>
#include <cstdint>
#include <Kokkos_Core.hpp>

#include "kernel.hpp"

#define ASSUME(cond)                    \
    do {                                \
        if (!(cond))                    \
            __builtin_unreachable();    \
    } while (0)

void UpdateVoltagesKernel(
    ArrayIJKN<float> volt,
    ArrayIJKN<float> curr,
    ArrayIJKN<float> vv,
    ArrayIJKN<float> vi,
    uint32_t i, uint32_t j, uint32_t k
)
{
    ASSUME(volt.m_grid == curr.m_grid);
    ASSUME(volt.m_grid == vv.m_grid);
    ASSUME(volt.m_grid == vi.m_grid);

    // 3 FP32 loads
    float volt0 = volt(i, j, k, 0);
    float volt1 = volt(i, j, k, 1);
    float volt2 = volt(i, j, k, 2);

    // 3 FP32 loads
    float vv0 = vv(i, j, k, 0);
    float vv1 = vv(i, j, k, 1);
    float vv2 = vv(i, j, k, 2);

    // 3 FP32 loads
    float vi0 = vi(i, j, k, 0);
    float vi1 = vi(i, j, k, 1);
    float vi2 = vi(i, j, k, 2);

    // 3 FP32 loads
    float curr0 = curr(i, j, k, 0);
    float curr1 = curr(i, j, k, 1);
    float curr2 = curr(i, j, k, 2);

    //for x polarization
    volt0 *= vv0;
    volt0 += vi0 * curr0;

    //for y polarization
    volt1 *= vv1;
    volt1 += vi1 * curr1;

    //for z polarization
    volt2 *= vv2;
    volt2 += vi2 * curr2;

    // 3 FP32 stores
    volt(i, j, k, 0) = volt0;
    volt(i, j, k, 1) = volt1;
    volt(i, j, k, 2) = volt2;
}

int main(int argc, char** argv)
{
    std::array<uint32_t, 3> numLines = {
        200, 200, 200
    };
    std::array<uint32_t, 3> dimTiles = {
        20, 20, 20
    };

    const int timesteps = 100;

    Kokkos::initialize(argc, argv);

    auto grid = Grid(numLines, dimTiles);

    auto volt = ArrayIJKN<float>("volt", grid);
    auto curr = ArrayIJKN<float>("curr", grid);

    auto vv = ArrayIJKN<float>("vv", grid);
    auto vi = ArrayIJKN<float>("vi", grid);
    auto ii = ArrayIJKN<float>("ii", grid);
    auto iv = ArrayIJKN<float>("iv", grid);

    Kokkos::parallel_for("First Touch",
        Kokkos::MDRangePolicy<Kokkos::Rank<3>> (
            {0, 0, 0},
            {numLines[0], numLines[1], numLines[2]},
            {dimTiles[0], dimTiles[1], dimTiles[2]}
        ),
        [=](
            const uint32_t i,
            const uint32_t j,
            const uint32_t k
        ) {
            for (uint32_t n = 0; n < 3; n++) {
                volt(i, j, k, n) = 0;
                curr(i, j, k, n) = 0;

                vv(i, j, k, n) = 0;
                vi(i, j, k, n) = 0;
                iv(i, j, k, n) = 0;
                ii(i, j, k, n) = 0;
            }
        }
    );
    Kokkos::fence();

    auto t1 = std::chrono::high_resolution_clock().now();

    for (size_t i = 0; i < timesteps; i++) {
        Kokkos::parallel_for("UpdateVoltages",
            Kokkos::MDRangePolicy<Kokkos::Rank<3>> (
                {0, 0, 0},
                {numLines[0], numLines[1], numLines[2]},
                {dimTiles[0], dimTiles[1], dimTiles[2]}
            ),
            [=](
                const uint32_t i,
                const uint32_t j,
                const uint32_t k
            ) {
                UpdateVoltagesKernel(volt, curr, vv, vi, i, j, k);
            }
        );

        Kokkos::fence();
    }

    Kokkos::fence();
    auto t2 = std::chrono::high_resolution_clock().now();
    double dt = std::chrono::duration_cast<std::chrono::microseconds>(t2 - t1).count() / 1e6;

    size_t bytes_per_iteration = numLines[0] * numLines[1] * numLines[2] * 3 * sizeof(float);
    bytes_per_iteration *= 5;  /* volt r+w, curr, vv, vi arrays */
    fprintf(stderr, "traffic: %.1f GB/s\n", (double) bytes_per_iteration * timesteps / 1000 / 1000 / 1000);
    fprintf(stderr, "speed: %.1f GB/s\n",
        (double) bytes_per_iteration * timesteps / dt / 1000 / 1000 / 1000
    );

    Kokkos::finalize();
}

CMakeLists.txt

cmake_minimum_required(VERSION 3.25)

project(fdtd-kokkos CXX)

set(KOKKOS_IN_TREE "/home/fdtd/kokkos")
add_subdirectory(${KOKKOS_IN_TREE} ${PROJECT_BINARY_DIR}/kokkos)

file (GLOB project_SOURCE ${project_SOURCE_DIR}/*.cpp)
foreach (sourcefile ${project_SOURCE})
    get_filename_component(sourcename ${sourcefile} NAME_WE)
    add_executable(${sourcename} ${sourcefile})
    target_link_libraries(
        ${sourcename} Kokkos::kokkoscore Kokkos::kokkossimd

    )
    target_include_directories(
        ${sourcename} PUBLIC ${project_SOURCE_DIR}
    )
endforeach (sourcefile ${project_SOURCE_DIR})


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

export CFLAGS="-O3 -march=native"
export CXXFLAGS="-O3 -march=native"
mkdir build && cd build
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();

#define ASSUME(cond)                    \
    do {                                \
        if (!(cond)) {                  \
            abort();                    \
            __builtin_unreachable();    \
        }                               \
    } while (0)

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

模板grid

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

template <typename T, uint32_t n_num=3>
struct ArrayIJKN
{
        ArrayIJKN(
                const std::string& name,
                const Grid& grid
        ) :


我用途:

template <typename T, Grid& grid, uint32_t n_num=3>
struct ArrayIJKN
{
        ArrayIJKN(
                const std::string& name
        )


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

std::array<uint32_t, 3> numLines = {
        200, 200, 200
};
std::array<uint32_t, 3> dimTiles = {
        20, 20, 20
};

auto grid = Grid(numLines, dimTiles);

// ...

int main(int argc, char** argv)
{
        Kokkos::initialize(argc, argv);

        auto volt = ArrayIJKN<float, grid>("volt");
        auto curr = ArrayIJKN<float, grid>("curr");

        auto vv = ArrayIJKN<float, grid>("vv");
        auto vi = ArrayIJKN<float, grid>("vi");
        auto ii = ArrayIJKN<float, grid>("ii");
        auto iv = ArrayIJKN<float, grid>("iv");
}


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

void UpdateVoltagesKernel(
        const ArrayIJKN<float, grid>& __restrict__ volt,
        const ArrayIJKN<float, grid>& __restrict__ curr,
        const ArrayIJKN<float, grid>& __restrict__ vv,
        const ArrayIJKN<float, grid>& __restrict__ vi,
        uint32_t i, uint32_t j, uint32_t k
)


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

相关问题