c - 将查找 tables 优化为简单 ALU

问题

假设您有一个简单的函数,它根据 table 的外观返回一个 value,例如:

请参阅有关假设的编辑。

uint32_t
lookup0(uint32_t r) {
    static const uint32_t tbl[] = { 0, 1, 2, 3 };
    if(r >= (sizeof(tbl) / sizeof(tbl[0]))) {
        __builtin_unreachable();
    }

    /* Can replace with: `return r`.  */
    return tbl[r];
}


uint32_t
lookup1(uint32_t r) {
    static const uint32_t tbl[] = { 0, 0, 1, 1 };
    if(r >= (sizeof(tbl) / sizeof(tbl[0]))) {
        __builtin_unreachable();
    }

    /* Can replace with: `return r / 2`.  */
    return tbl[r];
}

是否有任何超级optimization 基础设施或算法可以从查找table 到优化的ALU 实现。

动机

动机是我正在为 NUMA 机器构建一些锁,并希望能够通用地配置我的代码。在 NUMA 锁中您需要执行 cpu_id -> numa_node 是很常见的。我显然可以在配置过程中设置查找table,但是由于我正在争取尽可能多的内存带宽,我希望通常能找到一个能够覆盖大多数布局的解决方案。

看看现代编译器是如何做的:

https://godbolt.org/z/6e1E583vP

如果将 Clang 重写为 switch/case 语句,则可以得到 lookup0

lookup0(unsigned int):                            # @lookup0(unsigned int)
        movl    %edi, %eax
        movl    lookup0(unsigned int)::tbl(,%rax,4), %eax
        retq

...

case0(unsigned int):                              # @case0(unsigned int)
        movl    %edi, %eax
        retq

但无法获得 lookup1

lookup1(unsigned int):                            # @lookup1(unsigned int)
        movl    %edi, %eax
        movl    .Lswitch.table.case1(unsigned int)(,%rax,4), %eax
        retq

...

case1(unsigned int):                              # @case1(unsigned int)
        movl    %edi, %eax
        movl    .Lswitch.table.case1(unsigned int)(,%rax,4), %eax
        retq

Gcc 也得不到。

lookup0(unsigned int):
        movl    %edi, %edi
        movl    lookup0(unsigned int)::tbl(,%rdi,4), %eax
        ret
lookup1(unsigned int):
        movl    %edi, %edi
        movl    lookup1(unsigned int)::tbl(,%rdi,4), %eax
        ret
case0(unsigned int):
        leal    -1(%rdi), %eax
        cmpl    $2, %eax
        movl    $0, %eax
        cmovbe  %edi, %eax
        ret
case1(unsigned int):
        subl    $2, %edi
        xorl    %eax, %eax
        cmpl    $1, %edi
        setbe   %al
        ret

我想我可以用一些自定义的蛮力方法涵盖相当多的必要案例,但希望这是一个已解决的问题。

编辑:

唯一正确的假设是:

  1. 所有输入在 LUT 中都有一个索引。
  2. 所有 values 都是积极的(认为这使事情变得更容易)并且对于几乎所有在线的系统配置都是如此。
  3. (Edit4)会增加一个假设。 LUT 是密集的。也就是说,它涵盖了 [<low_bound>, <bound_bound>] 范围,但没有超出该范围。

对于 CPU 拓扑,我通常期望 sizeof(LUT) >= <max_value_in_lut> 但这仅针对我给出的一个示例,并且会有一些反例。

编辑2:

我编写了一个非常简单的优化器,它为我在https://godbolt.org/z/KKffvbz5f测试过的 CPU 拓扑做了合理的工作。但显然它可能会好很多。

编辑3:

问题/初始示例似乎有些混乱(我应该更清楚)。

lookup0/lookup1 示例是任意的。我希望找到一个可以扩展超过 4 个索引并具有不同 values 的解决方案。

我想到的用例是 CPU 拓扑,所以 ~256 - 1024 是我期望的大小上限,但对于通用 LUT,它显然会变得更大。

回答1

我知道的最好的“通用”解决方案如下:

int compute(int r)
{
    static const int T[] = {0,0,1,1};
    const int lut_size = sizeof(T) / sizeof(T[0]);

    int result = 0;

    for(int i=0 ; i<lut_size ; ++i)
        result += (r == i) * T[i];

    return result;
}

-O3 中,GCC 和 Clang 展开循环,传播常量,并生成类似于以下的中间代码:

int compute(int r)
{
    return (r == 0) * 0 + (r == 1) * 0 + (r == 2) * 1 + (r == 3) * 1;
}

GCC/Clang 优化器知道乘法可以用条件移动代替(因为开发人员经常使用这个技巧来指导编译器生成没有条件分支的汇编代码)。

对于 Clang,生成的程序集如下:

compute:
        xor     ecx, ecx
        cmp     edi, 2
        sete    cl
        xor     eax, eax
        cmp     edi, 3
        sete    al
        add     eax, ecx
        ret

这同样适用于 GCC。没有分支也没有内存访问(至少只要 values 很小)。乘以小 values 也被替换为快速 lea 指令。

https://godbolt.org/z/zzW7Yo4ef 提供了更完整的测试。

请注意,此方法适用于较大的 tables,但如果 table 太大,则不会自动展开循环。由于编译标志,您可以告诉编译器使用更积极的展开。话虽如此,如果 LUT 很大,它可能会更快,因为在这种病态的情况下,要加载和执行大量代码会很慢。

回答2

您可以将数组打包成一个长整数并使用位移和与来提取结果。

例如对于 table {2,0,3,1} 可以处理:

uint32_t lookup0(uint32_t r) {
    static const uint32_t tbl = (2u <<  0) | (0u <<  8) |
                                (3u << 16) | (1u << 24);
    return (tbl >> (8 * r)) & 0xff;
}

它产生了相对不错的组装:

lookup0:                                # @lookup0
        lea     ecx, [8*rdi]
        mov     eax, 16973826
        shr     eax, cl
        movzx   eax, al
        ret

不完美但无分支且没有间接性。

这种方法非常通用,它可以通过同时“查找”多个输入来支持矢量化。

有一些技巧可以允许处理更大的数组,例如使用更长的整数(即 uint64_t__uint128_t 扩展名)。另一种方法是在数组中拆分 value 的位,如高字节和低字节,查找它们并使用按位运算组合。

相似文章

最新文章