将查找表优化为简单的 ALU

Question

问题

假设您有一个简单的函数，它根据查找表返回一个值，例如：

请参阅有关假设的编辑。

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];
}

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

动机

动机是我正在为 NUMA 机器构建一些锁，并且希望能够通用地配置我的代码。在 NUMA 锁中，您需要执行cpu_id->操作，这很常见numa_node。显然，我可以在配置过程中设置查找表，但由于我正在努力争取每一滴内存带宽，因此我希望能够普遍达成一个能够覆盖大多数布局的解决方案。

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

目前两者clang都gcc无法做到这一点。

lookup0如果将其重写为switch/语句，Clang 就可以获取case。

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. 所有值都是正值（认为这会让事情变得更容易）并且对于任何在线系统配置都是正确的。
3. （编辑4）将添加一个假设。LUT 很密集。也就是说，它覆盖了一个范围[<low_bound>, <bound_bound>]，但没有超出该范围。

就我的 CPU 拓扑而言，我通常会期望sizeof(LUT) >= <max_value_in_lut>，但这特定于我给出的一个示例，并且会有一些反例。

编辑2：

我编写了一个非常简单的优化器，它对我在这里测试的 CPU 拓扑做了合理的工作。但显然它可能会好得多。

编辑3：

这个问题/最初的例子似乎有些混乱（我应该更清楚）。

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

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

Answer 1

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

例如，对于表{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扩展）。其他方法是拆分数组中的值位（例如高字节和低字节），查找它们并使用按位运算进行组合。