CUDA 算子笔记 1

§ 参考资料 共 3 条

Matrix Transposition

参考 《CUDA Programming Guide》2. CUDA C++

Matrix Multiplication

输入 $\boldsymbol A_{M\times N}$、$\boldsymbol B_{N\times K}$,求 $\boldsymbol C_{M\times K}=\boldsymbol A_{M\times N}\times\boldsymbol B_{N\times K}$

最简单的代码

#include <cuda_runtime.h>

__global__ void matrix_multiplication_kernel(const float* A, const float* B, float* C, int M, int N, int K) {
    const auto col = threadIdx.x + blockDim.x * blockIdx.x;
    const auto row = threadIdx.y + blockDim.y * blockIdx.y;

    if (row >= M || col >= K) {
        return;
    }

    float sum = 0;
    for (int i = 0; i < N; ++i) {
        sum += A[row * N + i] * B[i * K + col];
    }

    C[row * K + col] = sum;    
}

// A, B, C are device pointers (i.e. pointers to memory on the GPU)
extern "C" void solve(const float* A, const float* B, float* C, int M, int N, int K) {
    dim3 threadsPerBlock(16, 16);
    dim3 blocksPerGrid((K + threadsPerBlock.x - 1) / threadsPerBlock.x,
                       (M + threadsPerBlock.y - 1) / threadsPerBlock.y);

    matrix_multiplication_kernel<<<blocksPerGrid, threadsPerBlock>>>(A, B, C, M, N, K);
    cudaDeviceSynchronize();
}

使用 Tile 思想优化访存合并

上面的代码中,BC 的访存合并还不错(row 在一个 Warp 中只有两个值,col 连续变化,注意对于同一个 Warp 来说 i 应该在同一个时间不变的),而 A 基本上是 Broadcast

上面代码的问题主要在过度频繁地访问 Global Memory,对于每一个元素,都要访问 A 的一行和 B 的一列

需要使用 Tile 的思想来进行优化

代码:

#include <cuda_runtime.h>

constexpr auto TILE = 32;

__global__ void matrix_multiplication_kernel(const float* A, const float* B, float* C, int M, int N, int K) {
    __shared__ float tileA[TILE][TILE];
    __shared__ float tileB[TILE][TILE];

    const auto x = threadIdx.x;
    const auto y = threadIdx.y;

    const auto bx = blockIdx.x * blockDim.x;
    const auto by = blockIdx.y * blockDim.y;

    float sum = 0.0f;

    for (int t = 0; t < N; t += TILE) {
        // A[y + by][t + x]
        if (y + by < M && t + x < N) {
            tileA[y][x] = A[(y + by) * N + (t + x)];
        } else {
            tileA[y][x] = 0;
        }

        // B[t + y][x + bx]
        if (t + y < N && x + bx < K) {
            tileB[y][x] = B[(t + y) * K + (x + bx)];
        } else {
            tileB[y][x] = 0;
        }

        __syncthreads();

        #pragma unroll
        for (int k = 0; k < TILE; ++k) {
            // C[y + by][x + bx] = sum A[y + by][t + k] * B[t + k][x + bx]
            sum = fmaf(tileA[y][k], tileB[k][x], sum);
        }

        __syncthreads();
    }

    if (y + by < M && x + bx < K) {
        C[(y + by) * K + (x + bx)] = sum;
    }
}

// A, B, C are device pointers (i.e. pointers to memory on the GPU)
extern "C" void solve(const float* A, const float* B, float* C, int M, int N, int K) {
    dim3 threadsPerBlock(TILE, TILE);
    dim3 blocksPerGrid((K + threadsPerBlock.x - 1) / TILE,
                       (M + threadsPerBlock.y - 1) / TILE);

    matrix_multiplication_kernel<<<blocksPerGrid, threadsPerBlock>>>(A, B, C, M, N, K);
    cudaDeviceSynchronize();
}

SoftMax

输入一个 $n$ 个元素的数组 $[x_i]$,计算 Softmax $\sigma(x_i)$,然后放到输出数组中:

$$ \sigma(x_i)=\frac{e^{x_i}}{\sum_{j=1}^ne^{x_j}} $$

例子:

Input: [-10.0, -5.0, 0.0, 5.0, 10.0], N = 5
Output: [2.047e-09, 3.038e-07, 4.509e-05, 6.693e-03, 9.933e-01] (approximately)

Max Trick

如果有一个 $x_i$ 很大,那么 $e^{x_i}$ 会爆炸(例如 $e^{1000}$),为了避免这种情况,取 $m=\max_{i=1}^n x_i$,由于:

$$ \frac{e^{x_i-m}}{\sum_{j=1}^ne^{x_j-m}}=\frac{e^{-m}\times e^{x_i}}{e^{-m}\times\sum_{j=1}^ne^{x_j}}=\frac{e^{x_i}}{\sum_{j=1}^ne^{x_j}}=\sigma(x_i) $$

结果不变,这就是 Max Trick

最简单的代码

把 Softmax 函数分成两部分,一个是最大值 $m=\max_{i=1}^n{x_i}$,一个是求和 $s=\sum_{i=1}^n e^{x_i-m}$,那么 $\sigma(x_i)={e^{x_i-m}}/{s}$

用三个 Kernel:

  • 计算所有元素的最大值 $m$
  • 计算所有元素的 $e^{x_i-m}$,并且拿到 $s=\sum_{i=1}^ne^{x_i-m}$
  • 计算每个元素的 $\sigma(x_i)$

代码如下:

#include <cuda_runtime.h>
#include <float.h>

__device__ float sum, mx;

__device__ float atomicMaxFloat(float* address, float val) {
    int* address_as_int = reinterpret_cast<int*>(address);

    int old = *address_as_int;
    int assumed;

    do {
        assumed = old;
        if (__int_as_float(assumed) >= val){
            break;
        }
        old = atomicCAS(address_as_int, assumed, __float_as_int(val));
    } while (old != assumed);

    return __int_as_float(old);
}

__global__ void get_max_kernel(const float* input, int N) {
    const auto idx = threadIdx.x + blockIdx.x * blockDim.x;
    __shared__ float local_mx;
    if (threadIdx.x == 0) {
        local_mx = 0.0f;
    }
    __syncthreads();
    if (idx < N) {
        atomicMaxFloat(&local_mx, input[idx]);
    }
    __syncthreads();
    if (threadIdx.x == 0) {
        atomicMaxFloat(&mx, local_mx);
    }
}

__global__ void sum_and_exp_kernel(const float* input, float* output, int N) {
    const auto idx = threadIdx.x + blockIdx.x * blockDim.x;
    __shared__ float local_sum;
    if (threadIdx.x == 0) {
        local_sum = 0;
    }
    __syncthreads();
    if (idx < N) {
        const auto expx = __expf(input[idx] - mx);
        output[idx] = expx;
        atomicAdd(&local_sum, expx);
    }
    __syncthreads();
    if (threadIdx.x == 0) {
        atomicAdd(&sum, local_sum);
    }
}

__global__ void softmax_kernel(float* output, int N) {
    const auto idx = threadIdx.x + blockIdx.x * blockDim.x;
    if (idx < N) {
        output[idx] = output[idx] / sum;
    }
}

// input, output are device pointers (i.e. pointers to memory on the GPU)
extern "C" void solve(const float* input, float* output, int N) {
    int threadsPerBlock = 256;
    int blocksPerGrid = (N + threadsPerBlock - 1) / threadsPerBlock;

    float init_sum = 0.0f;
    float init_mx = -FLT_MAX;

    cudaMemcpyToSymbol(sum, &init_sum, sizeof(float));
    cudaMemcpyToSymbol(mx, &init_mx, sizeof(float));

    get_max_kernel<<<blocksPerGrid, threadsPerBlock>>>(input, N);
    sum_and_exp_kernel<<<blocksPerGrid, threadsPerBlock>>>(input, output, N);
    softmax_kernel<<<blocksPerGrid, threadsPerBlock>>>(output, N);
    cudaDeviceSynchronize();
}

CUDA 官方库不支持 float 类型的 atomicMax 函数,必须自己写,这里基于 CAS 实现了一个 atomicMaxFloat,由于 atomicCAS 也不支持 float,需要转换成 int

CAS:Compare And Swap,T atomicCAS(T* address, T expected, T desired),原子操作

  • 如果 *address == expected,那么 address 会被赋值为 desired
  • 如果 *address != expected,那么什么都不干

永远返回 address 在该原子操作之前的值

如果要实现 $x\gets f(x,v)$,其中 $v$ 是输入的新值,那么模板就是这样的:

T func(T x, T v);

void atomic_func(T* address, T v) {
    T old = *address;
    while (true) {
        T desired = func(old, v);
        T expected = old;
        old = CAS(address, expected, desired);
        if (old == expected) {
            // 说明当前 CAS 成功,已经替换成 desired 了
            break;
        }
        // CAS 失败,需要重试
    }
}

Online Softmax with Warp Reduction

上面的最大值 $m$ 和求和 $s$ 满足这样的性质:

如果有两个数组 $A=[x_i]$ 和 $B=[y_i]$,它们的最大值和求和分别为:

  • $m_A=\max_{i=1}^nx_i$,$s_A=\sum_{i=1}^ne^{x_i-m_A}$
  • $m_B=\max_{i=1}^ty_i$,$s_A=\sum_{i=1}^te^{y_i-m_B}$

那么对于两个数组的拼接 $C=A+B=[z_i]$,有:$m=\max\{m_A,m_B\}$,$s=e^{m_A-m}s_A+e^{m_B-m}s_B$

这样就可以先计算一部分的结果,然后再“组合”起来了。可以发现,$(m,s)$ 的这个组合运算满足结合律和交换律

一般通过一种叫 Reduction 的方法或者概念来搞或者解释,这里可以分成 Warp Reduction 和 Block Reduction,先用 Warp Reduction 对 Warp 中的所有 Thread 进行组合,然后用 Block Reduction 组合 Warp 的结果

假设我们只划分一个 Block,代码如下:

#include <cuda_runtime.h>
#include <cmath>

struct SoftmaxState {
    float max;
    float sum;
};

__device__ __forceinline__ SoftmaxState combine(SoftmaxState a, SoftmaxState b) {
    SoftmaxState out;
    out.max = fmaxf(a.max, b.max);

    if (out.max == -INFINITY) {
        out.sum = 0.0f;
    } else {
        // x_1, ..., x_m, y_1, ..., y_n -> z_1, ..., z_m+n
        //
        // max_x = max(x_i), sum_x = sum(exp(x_i - max_x))
        // max_y = max(y_i), sum_y = sum(exp(y_i - max_y))
        //
        // max = max(max_x, max_y)
        // sum = sum(exp(z_i - max)) = sum_x * exp(max_x - max) + sum_y * exp(max_y - max)
        out.sum = a.sum * __expf(a.max - out.max) + b.sum * __expf(b.max - out.max);
    }

    return out;
}

__device__ __forceinline__ SoftmaxState warp_reduce(SoftmaxState state) {
    //                         0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31
    // round 1 (offset = 16):  0+16 1+17 2+18 3+19 4+20 5+21 6+22 7+23 8+24 9+25 10+26 11+27 12+28 13+29 14+30 15+31 ...
    // round 2 (offset = 8):   0+16+8+24 1+17+9+25 2+18+10+26 3+19+11+27 4+20+12+28 5+21+13+29 6+22+14+30 7+23+15+31 ...
    // round 3 (offset = 4):   0+16+8+24+4+20+12+28 1+17+9+25+5+21+13+29 2+18+10+26+6+22+14+30 3+19+11+27+7+23+15+31 ...
    // round 4 (offset = 2):   0+16+8+24+4+20+12+28+2+18+10+26+6+22+14+30 1+17+9+25+5+21+13+29+3+19+11+27+7+23+15+31 ...
    // round 5 (offset = 1):   0+16+8+24+4+20+12+28+2+18+10+26+6+22+14+30+1+17+9+25+5+21+13+29+3+19+11+27+7+23+15+31 ...
    for (int offset = 16; offset > 0; offset >>= 1) {
        SoftmaxState other;

        // T __shfl_down_sync(unsigned int mask, T value, unsigned int delta, int width = 32)
        // - mask:指定哪些线程参与此次操作
        // - value:当前线程想要发送/共享的变量值
        // - delta:向下偏移的步长,目标线程 ID = 当前 ID + delta
        // - width:子 Warp 的大小
        other.max = __shfl_down_sync(0xffffffff, state.max, offset);
        other.sum = __shfl_down_sync(0xffffffff, state.sum, offset);

        state = combine(state, other);
    }
    return state;
}

__device__ __forceinline__ SoftmaxState block_reduce(SoftmaxState state) {
    __shared__ SoftmaxState warp_states[32];

    const auto lane = threadIdx.x & 31;
    const auto warp = threadIdx.x >> 5;
    const auto num_warps = (blockDim.x + 31) >> 5;

    // 对每个 warp 执行:将 warp 中的所有线程的输入进行 combine,结果放到 lane=0 的线程对应的位置
    // 每个线程都需要执行 warp_reduce(...),不只是 lane=0 的线程,因为所有线程都必须共享自己持有的变量
    state = warp_reduce(state);

    if (lane == 0) {
        // 每个 warp 只有 lane=0 的线程结果有意义
        warp_states[warp] = state;
    }

    __syncthreads();

    SoftmaxState block_state{ .max = -INFINITY, .sum = 0.0f };
    if (warp == 0) {
        // 由第一个 warp 来汇总所有 warp 的结果
        if (lane < num_warps) {
            block_state = warp_states[lane];
        }

        // 仍然用 warp reduce 的方法做,只有第一个线程的结果有意义
        block_state = warp_reduce(block_state);
    }

    return block_state;
}

__device__ __forceinline__ void norm(const float* input, SoftmaxState state, float* output, int N) {
    for (int i = threadIdx.x; i < N; i += blockDim.x) {
        output[i] = __expf(input[i] - state.max) / state.sum;
    }
}

__global__ void softmax_kernel(const float* input, float* output, int N) {
    // 每个线程计算对应列所有元素 combine 的结果
    SoftmaxState state = { .max = -INFINITY, .sum = 0.0f };
    for (int i = threadIdx.x; i < N; i += blockDim.x) {
        const SoftmaxState cur_state = { .max = input[i], .sum = 1.0f };
        state = combine(state, cur_state);
    }

    // 计算得到所有元素的 combine 结果,只有第一个线程的结果有意义
    state = block_reduce(state);

    __shared__ SoftmaxState final_state;
    if (threadIdx.x == 0) {
        final_state = state;
    }

    __syncthreads();

    // 计算每个元素的结果并且放到 output 中
    norm(input, final_state, output, N);
}

// input, output are device pointers (i.e. pointers to memory on the GPU)
extern "C" void solve(const float* input, float* output, int N) {
    int threadsPerBlock = 256;
    int blocksPerGrid = (N + threadsPerBlock - 1) / threadsPerBlock;

    softmax_kernel<<<1, threadsPerBlock>>>(input, output, N);
    cudaDeviceSynchronize();
}

其中,__shfl_down_sync 是 Warp Reduction 的关键

如果一个 Warp 中,每个线程自己的寄存器里都有一个变量 x,一般来说一个线程 A 不能拿到另一个线程 B 的 x 的值,因为这属于线程自己的上下文。但 NVIDIA 提供了一组 Warp Shuffle 指令,让 Warp 内线程进行直接的数据交换,它可以在 Warp 的非退出线程之间交换变量,无需使用 Shared Memory,例如:

__shfl_sync(...)
__shfl_up_sync(...)
__shfl_down_sync(...)
__shfl_xor_sync(...)

__shfl_down_sync 的签名为 T __shfl_down_sync(unsigned int mask, T value, unsigned int delta, int width = 32),它的意思是:当前 Lane 从 Lane ID + delta 那个线程取得它的 value。这样的话,值整体向 Lane 编号更小的方向“移动”,所以得名 “down”

这样比 Shared Memory 更快,因为这仅仅是寄存器之间的数据移动,而 Shared Memory 方案需要:寄存器 -> Shared Memory -> 同步 -> 寄存器

__shfl_down_sync 的一个经典应用就是求和 Warp Reduction,实际上任何满足交换律和结合律的运算都可以这样 Reduction(这里的组合运算就是这样的),求和 Warp Reduction 一般这么实现:

__device__ void warp_reduction(int value) {
    for (int offset = 16; offset > 0; offset >>= 1) {
        other = __shfl_down_sync(0xffffffff, value, offset);
        value += other;
    }
}

原理图如下:

img

对于整个 Block 来说,需要进行两步 Reduction:

  • 每个 Warp 内部的 Reduction
  • 所有 Warp 的内部 Reduction 结果的 Reduction

像上面代码中写的那样,先把每个 Warp 内部的 Reduction 写到一个 Shared Memory 数组中(这里需要同步),然后对这个 Shared Memory 数组看成一个 Warp 再进行一次 Reduction

Kernel 本身就很简单了

Online Softmax 改进

上一个版本的 Softmax 是只有一个 Block,肯定不能这样,考虑将整个数组分成多个 Block,每个 Block 负责自己的那部分数据(Tile 的思想)

由于 Block 之间没有简单的同步机制和共享内存,我们需要借助 Global Memory,分成多个 Kernel:

  • 计算 Block 负责的数据中的最大值 $m$ 和求和 $s$,放到一个 Global Memory 数组中
  • 将多个 Block 的结果组合一下得到最终的 $m$ 和 $s$,同样放到 Global Memory 中
  • 计算每个元素的 $\sigma(x_i)$

代码:

#include <cuda_runtime.h>
#include <cmath>

constexpr auto BLOCK_SIZE = 256;
constexpr auto ITEMS_PER_THREAD = 32;
constexpr auto MAX_BLOCKS = 1024;
constexpr auto MAX_WARPS = 32;

struct SoftmaxState {
    float max;
    float sum;
};

__device__ SoftmaxState g_partial_states[MAX_BLOCKS];
__device__ SoftmaxState g_final_state;

__device__ __forceinline__ SoftmaxState identity_state() {
    return { .max = -INFINITY, .sum = 0.0f };
}

__device__ __forceinline__ SoftmaxState mono_state(const float value) {
    return { .max = value, .sum = 1.0f };
}

__device__ __forceinline__ SoftmaxState combine(SoftmaxState a, SoftmaxState b) {
    SoftmaxState out;
    out.max = fmaxf(a.max, b.max);

    if (out.max == -INFINITY) {
        out.sum = 0.0f;
    } else {
        // x_1, ..., x_m, y_1, ..., y_n -> z_1, ..., z_m+n
        //
        // max_x = max(x_i), sum_x = sum(exp(x_i - max_x)), sum_x_final = exp(max_x) * sum_x
        // max_y = max(y_i), sum_y = sum(exp(y_i - max_y)), sum_y_final = exp(max_y) * sum_y
        //
        // max = max(max_x, max_y)
        // sum = sum(exp(z_i - max)) = sum_x * exp(max_x - max) + sum_y * exp(max_y - max)
        out.sum = a.sum * __expf(a.max - out.max) + b.sum * __expf(b.max - out.max);
    }

    return out;
}

__device__ __forceinline__ SoftmaxState warp_reduce(SoftmaxState state) {
    //                         0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31
    // round 1 (offset = 16):  0+16 1+17 2+18 3+19 4+20 5+21 6+22 7+23 8+24 9+25 10+26 11+27 12+28 13+29 14+30 15+31 ...
    // round 2 (offset = 8):   0+16+8+24 1+17+9+25 2+18+10+26 3+19+11+27 4+20+12+28 5+21+13+29 6+22+14+30 7+23+15+31 ...
    // round 3 (offset = 4):   0+16+8+24+4+20+12+28 1+17+9+25+5+21+13+29 2+18+10+26+6+22+14+30 3+19+11+27+7+23+15+31 ...
    // round 4 (offset = 2):   0+16+8+24+4+20+12+28+2+18+10+26+6+22+14+30 1+17+9+25+5+21+13+29+3+19+11+27+7+23+15+31 ...
    // round 5 (offset = 1):   0+16+8+24+4+20+12+28+2+18+10+26+6+22+14+30+1+17+9+25+5+21+13+29+3+19+11+27+7+23+15+31 ...
    for (int offset = 16; offset > 0; offset >>= 1) {
        SoftmaxState other;

        // T __shfl_down_sync(unsigned int mask, T var, unsigned int delta, int width = 32)
        // - mask:指定哪些线程参与此次操作
        // - var:当前线程想要发送/共享的变量值
        // - delta:向下偏移的步长,目标线程 ID = 当前 ID + delta
        // - width:子 Warp 的大小
        other.max = __shfl_down_sync(0xffffffff, state.max, offset);
        other.sum = __shfl_down_sync(0xffffffff, state.sum, offset);

        state = combine(state, other);
    }
    return state;
}

__device__ __forceinline__ SoftmaxState block_reduce(SoftmaxState state) {
    __shared__ SoftmaxState warp_states[MAX_WARPS];

    const auto lane = threadIdx.x & 31;
    const auto warp = threadIdx.x >> 5;
    const auto num_warps = (blockDim.x + 31) >> 5;

    // 对每个 warp 执行:将 warp 中的所有线程的输入进行 combine,结果放到 lane=0 的线程对应的位置
    // 每个线程都需要执行 warp_reduce(...),不只是 lane=0 的线程,因为所有线程都必须共享自己持有的变量
    state = warp_reduce(state);

    if (lane == 0) {
        // 每个 warp 只有 lane=0 的线程结果有意义
        warp_states[warp] = state;
    }

    __syncthreads();

    auto block_state = identity_state();
    if (warp == 0) {
        // 由第一个 warp 来汇总所有 warp 的结果
        if (lane < num_warps) {
            block_state = warp_states[lane];
        }

        // 仍然用 warp reduce 的方法做,只有第一个线程的结果有意义
        block_state = warp_reduce(block_state);
    }

    return block_state;
}

__global__ void reduce_input_kernel(const float* input, int N) {
    const auto total_threads = blockDim.x * gridDim.x;
    const auto idx = threadIdx.x + blockDim.x * blockIdx.x;

    auto state = identity_state();
    for (int i = idx; i < N; i += total_threads) {
        // 每个线程负责计算 ITEMS_PER_THREAD 个元素的 combine 结果
        const auto cur_state = mono_state(input[i]);
        state = combine(state, cur_state);
    }

    // 计算得到 block 中负责的所有元素的 combine 结果,只有 block 内的第一个线程的结果有意义
    state = block_reduce(state);

    // 记录到 global 数组中
    if (threadIdx.x == 0) {
        g_partial_states[blockIdx.x] = state;

        // 如果只有一个 block,它就是最终结果,直接把它放到结果变量中,后续就不用启动第二个 kernel 了
        if (gridDim.x == 1) {
            g_final_state = state;
        }
    }
}

__global__ void reduce_partials_kernel(int num_partials) {
    auto state = identity_state();
    for (int i = threadIdx.x; i < num_partials; i += blockDim.x) {
        // 仍然是每个线程负责计算至多 4 个元素的 combine 结果
        // num_partials <= 1024, BLOCK_SIZE = 256, 1024/256 = 4
        state = combine(state, g_partial_states[i]);
    }

    state = block_reduce(state);

    if (threadIdx.x == 0) {
        g_final_state = state;
    }
}

__global__ void normalize_kernel(const float* input, float* output, int N) {
    // 拿到最终的 state,这里是一个 broadcast 机制
    __shared__ SoftmaxState state;
    if (threadIdx.x == 0) {
        state = g_final_state;
    }
    __syncthreads();

    const auto total_threads = blockDim.x * gridDim.x;
    const auto idx = threadIdx.x + blockDim.x * blockIdx.x;

    for (int i = idx; i < N; i += total_threads) {
        output[i] = __expf(input[i] - state.max) / state.sum;
    }
}

// input, output are device pointers (i.e. pointers to memory on the GPU)
extern "C" void solve(const float* input, float* output, int N) {
    const auto num_blocks = (N + BLOCK_SIZE * ITEMS_PER_THREAD - 1) / (BLOCK_SIZE * ITEMS_PER_THREAD);

    reduce_input_kernel<<<num_blocks, BLOCK_SIZE>>>(input, N);
    if (num_blocks > 1) {
        reduce_partials_kernel<<<1, BLOCK_SIZE>>>(num_blocks);
    }
    normalize_kernel<<<num_blocks, BLOCK_SIZE>>>(input, output, N);
    cudaDeviceSynchronize();
}

基于 Cooperative Groups 的改进