在 CUDA 中使用 SIMD 实现位循环运算符

Question

我知道 StackOverflow 不是用来向其他人询问代码的，但请允许我说一下。

我正在尝试在 CUDA C++ 设备代码中实现一些 AES 函数。在尝试实现左字节旋转运算符时，我很不安地看到没有本机 SIMD intrisic。所以我开始了一个幼稚的实现，但是....它很大，虽然我还没有尝试过，但它不会很快，因为解包/打包成本很高......所以，有没有办法做至少有点高效的每字节位循环操作？

如果你不想看，这里是代码。

__inline__ __device__ uint32_t per_byte_bit_left_rotate(uint32_t input, uint8_t amount) {
return ((((input & 0xFF) >> 0) << amount) | (((input & 0xFF) >> 0) >> 7) & ~0x100) << 0 |
     ((((input & 0xFF00) >> 8) << amount) | ((input & 0xFF00 >> 8) >> 7) & ~0x100) << 8 |
     ((((input & 0xFF0000) >> 16) << amount) | ((input & 0xFF0000 >> 16) >> 7) & ~0x100) << 16 |
     ((((input & 0xFF000000) >> 24) << amount) | ((input & 0xFF000000 >> 24) >> 7) & ~0x100) << 24; } // The XORs are for clearing the old 7th bit who is getting pushed to the next byte of the intermediate int

Answer 1

CUDA 有一个 __byte_perm() 内在函数，它直接映射到机器代码 (SASS) 级别的 PRMT 指令，这是一个按字节排列的指令。它可用于有效地提取和合并字节。为了实现逐字节左旋转，我们可以将每个字节加倍，将字节对移动所需的数量，然后提取并合并字节对的四个高字节。

对于逐字节旋转，我们只需要移位量的最低三位，因为s 的旋转与s mod 8 的旋转相同。为了提高效率，最好避免包含少于 32 位的整数类型，因为 C++ 语义要求在表达式中使用之前，将窄于 int 的整数类型扩大到 int。这可能而且确实会在包括 GPU 在内的许多架构上产生转换开销。

PRMT 指令的吞吐量取决于架构，因此使用__byte_perm() 可能会导致代码比使用another answer 中演示的经典SIMD-in-a-register 方法更快或更慢，因此，请务必在部署之前根据您的用例环境进行基准测试。

#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>

__device__ uint32_t per_byte_bit_left_rotate (uint32_t input, uint32_t amount)
{
     uint32_t l = __byte_perm (input, 0, 0x1100) << (amount & 7);
     uint32_t h = __byte_perm (input, 0, 0x3322) << (amount & 7);
     return __byte_perm (l, h, 0x7531);
}

__global__ void rotl_kernel (uint32_t input, uint32_t amount, uint32_t *res)
{
    *res = per_byte_bit_left_rotate (input, amount);
}

uint32_t ref_per_byte_bit_left_rotate (uint32_t input, uint32_t amount)
{
   int s = amount & 7;
   uint8_t b0 = (input >>  0) & 0xff;
   uint8_t b1 = (input >>  8) & 0xff;
   uint8_t b2 = (input >> 16) & 0xff;
   uint8_t b3 = (input >> 24) & 0xff;
   b0 = s ? ((b0 << s) | (b0 >> (8 - s))) : b0;
   b1 = s ? ((b1 << s) | (b1 >> (8 - s))) : b1;
   b2 = s ? ((b2 << s) | (b2 >> (8 - s))) : b2;
   b3 = s ? ((b3 << s) | (b3 >> (8 - s))) : b3;
   return (b3 << 24) | (b2 << 16) | (b1 << 8) | (b0 << 0);
}

// Fixes via: Greg Rose, KISS: A Bit Too Simple. http://eprint.iacr.org/2011/007
static unsigned int z=362436069,w=521288629,jsr=362436069,jcong=123456789;
#define znew (z=36969*(z&0xffff)+(z>>16))
#define wnew (w=18000*(w&0xffff)+(w>>16))
#define MWC  ((znew<<16)+wnew)
#define SHR3 (jsr^=(jsr<<13),jsr^=(jsr>>17),jsr^=(jsr<<5)) /* 2^32-1 */
#define CONG (jcong=69069*jcong+13579)                     /* 2^32 */
#define KISS ((MWC^CONG)+SHR3)

// Macro to catch CUDA errors in CUDA runtime calls
#define CUDA_SAFE_CALL(call)                                          \
do {                                                                  \
    cudaError_t err = call;                                           \
    if (cudaSuccess != err) {                                         \
        fprintf (stderr, "Cuda error in file '%s' in line %i : %s.\n",\
                 __FILE__, __LINE__, cudaGetErrorString(err) );       \
        exit(EXIT_FAILURE);                                           \
    }                                                                 \
} while (0)

// Macro to catch CUDA errors in kernel launches
#define CHECK_LAUNCH_ERROR()                                          \
do {                                                                  \
    /* Check synchronous errors, i.e. pre-launch */                   \
    cudaError_t err = cudaGetLastError();                             \
    if (cudaSuccess != err) {                                         \
        fprintf (stderr, "Cuda error in file '%s' in line %i : %s.\n",\
                 __FILE__, __LINE__, cudaGetErrorString(err) );       \
        exit(EXIT_FAILURE);                                           \
    }                                                                 \
    /* Check asynchronous errors, i.e. kernel failed (ULF) */         \
    err = cudaThreadSynchronize();                                    \
    if (cudaSuccess != err) {                                         \
        fprintf (stderr, "Cuda error in file '%s' in line %i : %s.\n",\
                 __FILE__, __LINE__, cudaGetErrorString( err) );      \
        exit(EXIT_FAILURE);                                           \
    }                                                                 \
} while (0)

int main (void)
{
    uint32_t arg, ref, res = 0, *res_d = 0;
    uint32_t shft;

    CUDA_SAFE_CALL (cudaMalloc ((void**)&res_d, sizeof(*res_d)));
    for (int i = 0; i < 100000; i++) {
        arg  = KISS;
        shft = KISS;
        ref = ref_per_byte_bit_left_rotate (arg, shft);
        rotl_kernel <<<1,1>>>(arg, shft, res_d);
        CHECK_LAUNCH_ERROR();
        CUDA_SAFE_CALL (cudaMemcpy (&res, res_d, sizeof (res), 
                                    cudaMemcpyDeviceToHost));
        if (res != ref) {
            printf ("!!!! arg=%08x shft=%d  res=%08x  ref=%08x\n", 
                    arg, shft, res, ref);
        }
    }
    CUDA_SAFE_CALL (cudaFree (res_d));
    CUDA_SAFE_CALL (cudaDeviceSynchronize());
    return EXIT_SUCCESS;
}

Answer 2

所有元素的旋转计数都相同，对吧？

将整个输入向左和向右移动，然后将那些带有掩码的那些跨越字节边界的所有位都归零，将所有 4 个字节放在一个 AND 中。我认为amount 始终是 AES 中的编译时常量，因此您不必担心动态生成掩码的运行时成本。让编译器来做。 （IDK CUDA，但这似乎与为普通 C++ 编写带有 32 位整数的 SWAR bit-hack 相同的问题）

这基于通常的(x << count) | (x >> (32-count)) rotate idiom，带有掩码和不同的右移计数，以使其成为单独的 8 位循环。

inline
uint32_t per_byte_bit_left_rotate(uint32_t input, unsigned amount)
{
    // With constant amount, the left/right masks are constants
    uint32_t rmask = 0xFF >> ((8 - amount) & 7);
    rmask = (rmask<<24 | rmask<<16 | rmask<<8 | rmask);
    uint32_t lmask = ~rmask;

    uint32_t lshift = input << amount;
    lshift &= lmask;
    if (amount == 1) {  // special case left-shift by 1 using an in-lane add instead of shift&mask
        lshift = __vadd4(input, input);
    }
    uint32_t rshift = input >> ((8 - amount) & 7);
    rshift &= rmask;

    uint32_t rotated = lshift | rshift;
    return rotated;
}

在移位前以一种方式屏蔽输入，并在移位后屏蔽输出（(in&lmask)<<amount | ((in>>(8-amount))&rmask)，使用不同的 lmask）可能会更有效。 NVidia 硬件是有序超标量，shifts have limited throughput。这样做更有可能作为两个独立的 shift+m​​ask 对来执行。

（这并没有试图避免数量>=32 的 C++ UB。请参阅 Best practices for circular shift (rotate) operations in C++。在这种情况下，我认为更改为 lshift = input << (amount & 7) 可以解决问题。

为了测试这个编译是否有效，我查看了 x86-64 的 clang -O3 asm output 和常量 amount。 Godbolt 编译器资源管理器具有适用于各种架构的编译器（虽然不是 CUDA），因此如果您比 x86 更容易阅读这些 asm 语言，请单击该链接并切换到 ARM、MIPS 或 PowerPC。

uint32_t rol7(uint32_t a) {
    return per_byte_bit_left_rotate(a, 7);
}
    mov     eax, edi
    shl     eax, 7
    shr     edi
    and     eax, -2139062144   # 0x80808080
    and     edi, 2139062143    # 0x7F7F7F7F
    lea     eax, [rdi + rax]   # ADD = OR when no bits intersect
    ret

完美，正是我所希望的。

几个测试用例：

uint32_t test_rol() {
    return per_byte_bit_left_rotate(0x02ffff04, 0);
}
    // yup, returns the input with count=0
    // return 0x2FFFF04


uint32_t test2_rol() {
    return per_byte_bit_left_rotate(0x02f73804, 4);
}
    // yup, swaps nibbles
    // return 0x207F8340

这与使用 x86 SSE2 / AVX2 进行 8 位移位需要做的事情相同，因为硬件支持的最小位移粒度是 16 位。