Rainbow Table

原始题目：LeetGPU - Rainbow Table

题目描述

编写一个 GPU 程序，使用提供的哈希函数对 32 位整数数组执行 $R$ 轮并行哈希。哈希函数应迭代应用 $R$ 次——每一轮的输出作为下一轮的输入。

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。
• 最终结果必须存储在数组 output 中。

示例

示例 1

Input:  numbers = [123, 456, 789], R = 2
Output: hashes = [1636807824, 1273011621, 2193987222]

示例 2

Input:  numbers = [0, 1, 2147483647], R = 3
Output: hashes = [96754810, 3571711400, 2006156166]

约束条件

• $1 \le N \le 10{,}000{,}000$。
• $1 \le R \le 100$。
• $0 \le input[i] \le 2{,}147{,}483{,}647$。
• 性能测试在 $N = 5{,}000{,}000$ 的规模下进行。

解题思路

彩虹表是密码学中的经典数据结构，本质是对每个元素独立迭代哈希 $R$ 次。每个线程可以处理一个元素的所有 $R$ 轮迭代，也可以让 $R$ 轮之间通过全局同步交换数据。前者计算密度更高且没有同步开销，是最自然的选择。本次运算需要 $R$ 次哈希调用，每次调用的延迟会累积。

代码实现

CUDA

#include <cuda_runtime.h>

__device__ unsigned int hash(unsigned int x) {
    x = ((x >> 16) ^ x) * 0x45d9f3b;
    x = ((x >> 16) ^ x) * 0x45d9f3b;
    x = (x >> 16) ^ x;
    return x;
}

__global__ void rainbow_kernel(const unsigned int* input, unsigned int* output, int N, int R) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < N) {
        unsigned int val = input[i];
        for (int r = 0; r < R; r++) {
            val = hash(val);
        }
        output[i] = val;
    }
}

extern "C" void solve(const unsigned int* input, unsigned int* output, int N, int R) {
    int threadsPerBlock = 256;
    int blocksPerGrid = (N + threadsPerBlock - 1) / threadsPerBlock;
    rainbow_kernel<<<blocksPerGrid, threadsPerBlock>>>(input, output, N, R);
    cudaDeviceSynchronize();
}

Triton

import triton
import triton.language as tl

@triton.jit
def rainbow_kernel(input_ptr, output_ptr, N: tl.constexpr, R: tl.constexpr, BLOCK_SIZE: tl.constexpr):
    idx = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = idx < N
    val = tl.load(input_ptr + idx, mask=mask)
    for _ in range(R):
        val = ((val >> 16) ^ val) * 0x45d9f3b
        val = ((val >> 16) ^ val) * 0x45d9f3b
        val = (val >> 16) ^ val
    tl.store(output_ptr + idx, val, mask=mask)