CUDA编程 - 用向量化访存优化 elementwise 核函数

Cuda elementwise

一、简介
- 1.1、ElementWise
- 1.2、 float4 - 向量化访存
二、实践
- 2.1、如何使用向量化访存
- 2.2、Cuda elementwise - Add
- 2.3、Cuda elementwise - Sigmoid
- - 2.3.1、简单的 Sigmoid 函数
  - 2.3.2、ElementWise Sigmoid+ float4（向量化访存）
- 2.4、Cuda elementwise - relu
- - 2.3.1、简单的 relu 函数
  - 2.3.2、ElementWise relu + float4（向量化访存）
三、不能使用向量化访存的情况
- Cuda elementwise - Histogram（直方图）
四、完整代码
- 4.1、sigmoid.cu
- 4.2、relu.cu
4.3、histogram.cu

一、简介

1.1、ElementWise

Element-wise操作是最基础，最简单的一种核函数的类型，它的计算特点很符合GPU的工作方式:对于每个元素单独做一个算术操作，然后直接输出。虽然简单，但深度学习领域，有很多算子都是这个类型。常见的有Add，Mul，Concat，各种激活Sigmoid，Relu及它们的变体，归一化里的BatchNorm都属于这个类型。

1.2、 float4 - 向量化访存

所谓向量化访存，就是一次性读 4 个 float，而不是单单 1 个

要点：

小数据规模情况下，可以不考虑向量化访存的优化方式
大规模数据情况下，考虑使用向量化访存，且 最好是缩小grid的维度为原来的1/4，避免影响Occupancy
float4 向量化访存只对数据规模大的时候有加速效果，数据规模小的时候没有加速效果

float4 的性能提升主要在于访存指令减少了(同样的数据规模，以前需要4条指令，现在只需1/4的指令)，指令cache里就能存下更多指令，提高指令cache的命中率。

判断是否用上了向量化访存，是在 nsight compute 看生成的SASS代码里会有没有 LDG.E.128 Rx, [Rx.64]或 STG.E.128 [R6.64], Rx这些指令的存在。有则向量化成功，没有则向量化失败。

在这里插入图片描述

官方参考链接1
官方参考链接2

二、实践

2.1、如何使用向量化访存

c ：

#define FLOAT4(value)  *(float4*)(&(value))

宏解释：

对于一个值，先对他取地址，然后再把这个地址解释成 float4
对于这个 float4的指针，对它再取一个值
这样编译器就可以一次读四个 float

c++ ：

#define FLOAT4(value) (reinterpret_cast<float4*>(&(value))[0])

将核函数写成 float4 的形式的时候，首先要先使用宏定义，其次要注意线程数的变化。

线程数变化原因：因为一个线程可以处理 4个float了，所以要减少四倍的线程。（具体看下面的例子）

2.2、Cuda elementwise - Add

参考链接

2.3、Cuda elementwise - Sigmoid

2.3.1、简单的 Sigmoid 函数

a: Nx1, b: Nx1, c: Nx1, c = sigmoid(a, b)

//sigmoid<<<CeilDiv(N, block_size), block_size>>>(d_A, d_B, N)
//
__global__ void sigmoid(float* x, float* y, int N) {
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx < N) y[idx] = 1.0f / (1.0f + expf(-x[idx]));
}

2.3.2、ElementWise Sigmoid+ float4（向量化访存）

Sigmoid x: N, y: N y=1/(1+exp(-x)) float4

__global__ void sigmoid_vec4(float4* x, float4* y, int N) {
  int idx = (blockIdx.x * blockDim.x + threadIdx.x) * 4;
  if (idx < N) {
    float4 tmp_x = FLOAT4(x[idx]);
    float4 tmp_y;
    tmp_y.x = 1.0f / (1.0f + expf(-tmp_x.x));
    tmp_y.y = 1.0f / (1.0f + expf(-tmp_x.y));
    tmp_y.z = 1.0f / (1.0f + expf(-tmp_x.z));
    tmp_y.w = 1.0f / (1.0f + expf(-tmp_x.w));
    FLOAT4(y[idx]) = tmp_y;
  }
}

2.4、Cuda elementwise - relu

2.3.1、简单的 relu 函数

Relu x: N, y: N y=max(0,x)

__global__ void relu(float* x, float* y, int N) {
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx < N) y[idx] = fmaxf(0.0f, x[idx]);
}

2.3.2、ElementWise relu + float4（向量化访存）

Relu x: N, y: N y=max(0,x) float4

__global__ void relu_float4(float* x, float* y, int N) {
  int idx = (blockIdx.x * blockDim.x + threadIdx.x) * 4;
  if (idx < N) {
    float4 tmp_x = FLOAT4(x[idx]);
    float4 tmp_y;
    tmp_y.x = fmaxf(0.0f, tmp_x.x);
    tmp_y.y = fmaxf(0.0f, tmp_x.y);
    tmp_y.z = fmaxf(0.0f, tmp_x.z);
    tmp_y.w = fmaxf(0.0f, tmp_x.w);
    FLOAT4(y[idx]) = tmp_y;
  }
}

三、不能使用向量化访存的情况

Cuda elementwise - Histogram（直方图）

x[i] = i
        -
      - -
    - - -
  - - - -
- - - - -

goal: 统计每个元素出现的次数

①、简单的 Histogram函数

x: Nx1, y: count histogram

__global__ void histogram(int* x, int* y, int N) {
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx < N) atomicAdd(&(y[x[idx]]), 1);
}

其中 int atomicAdd(int* address, int val)；

1、atomicAdd 是 CUDA 中的一个原子加函数，用于实现在多个线程同时修改同一个全局变量的情况下，保证数据一致性和正确性
2、该函数会将原始值和 val 相加，并将结果存储在 address 所指向的内存位置，同时返回原始值
3、当多个线程同时调用 atomicAdd 函数时，CUDA 会自动为它们创建一个硬件级的同步访问机制，从而避免了数据竞争和数据不一致性的问题。

参考链接

②、ElementWise Histogram + float4（向量化访存）

__global__ void histogram_float4(int* x, int* y, int N) {
  int idx = 4 * (blockIdx.x * blockDim.x + threadIdx.x);
  if (idx < N) {
    int4 tmp_y = INT4(x[idx]);
    atomicAdd(&(y[tmp_y.x]), 1);
    atomicAdd(&(y[tmp_y.y]), 1);
    atomicAdd(&(y[tmp_y.z]), 1);
    atomicAdd(&(y[tmp_y.w]), 1);
  }
}

运行 histogram_float4 后发现比 histogram 还慢。为什么呢？

1、nvidia 官方回答: Can float4 be used for atomicAdd efficiently?
2、histogram的写入(load)没有向量化指令，只有读取(load)向量化。
3、使用 nsight compute 看生成的SASS，只有 LDG ，没有STG

4、在一个warp里相邻线程对全局内存y的访问没有合并（因为在一个warp里面不同线程对全局内存是跳着访问的）

四、完整代码

4.1、sigmoid.cu

#include <stdio.h>
#include <stdlib.h>
#include <float.h>
#include <vector>
#include<assert.h>
#include <algorithm>
#include <cublas_v2.h>
#include <cuda_runtime.h>

#define FLOAT4(value)  *(float4*)(&(value))

#define checkCudaErrors(func)               \
{                                   \
    cudaError_t e = (func);         \
    if(e != cudaSuccess)                                        \
        printf ("%s %d CUDA: %s\n", __FILE__,  __LINE__, cudaGetErrorString(e));        \
}

__global__ void relu(float* x, float* y, int N){
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if(idx < N) y[idx] = fmaxf(0.0f, x[idx]);
}

__global__ void relu_float4(float* x, float* y, int N){

    int idx = (blockIdx.x * blockDim.x + threadIdx.x) * 4;
    if(idx < N){
        float4 tmp_x = FLOAT4(x[idx]);
        float4 tmp_y;
        tmp_y.x = fmaxf(0.0f, tmp_x.x);
        tmp_y.y = fmaxf(0.0f, tmp_x.y);
        tmp_y.z = fmaxf(0.0f, tmp_x.z);
        tmp_y.w = fmaxf(0.0f, tmp_x.w);
        FLOAT4(y[idx]) = tmp_y;
    }
}

//nvcc -o sigmoid sigmoid.cu && ./sigmoid
//sigmoid<<<CeilDiv(N, block_size), block_size>>>(d_A, d_B, N)
//a: Nx1, b: Nx1, c: Nx1, c = sigmoid(a, b)
__global__ void sigmoid(float* a, float* b, int N) {
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx < N) b[idx] = 1.0f / (1.0f + expf(-a[idx]));
}

__global__ void sigmoid_float4(float* a, float* b, int N)
{
    int idx = (blockDim.x * blockIdx.x + threadIdx.x) * 4;
    if(idx < N){
        float4 tmp_a = FLOAT4(a[idx]);
        float4 tmp_b;
        tmp_b.x = 1.0f /(1.0f + expf(-tmp_a.x));
        tmp_b.y = 1.0f /(1.0f + expf(-tmp_a.y));
        tmp_b.z = 1.0f /(1.0f + expf(-tmp_a.z));
        tmp_b.w = 1.0f /(1.0f + expf(-tmp_a.w));
        FLOAT4(b[idx]) = tmp_b;
    }
}

template <typename T> 
inline T CeilDiv(const T& a, const T& b) {
    return (a + b - 1) / b;
}

int main(){

    size_t block_size = 128;
    size_t N = 1 * 1024;
    size_t bytes_A = sizeof(float) * N;
    size_t bytes_B = sizeof(float) * N;

    float* h_A = (float*)malloc(bytes_A);
    float* h_B = (float*)malloc(bytes_B);

    for( int i = 0; i < N; i++ ){
        h_A[i] = (i / 666) * ((i % 2 == 0) ? 1: -1);
    }

    float* d_A;
    float* d_B;

    checkCudaErrors(cudaMalloc(&d_A, bytes_A));
    checkCudaErrors(cudaMalloc(&d_B, bytes_B));

    checkCudaErrors(cudaMemcpy( d_A, h_A, bytes_A, cudaMemcpyHostToDevice));

    cudaEvent_t start, stop;
    checkCudaErrors(cudaEventCreate(&start));
    checkCudaErrors(cudaEventCreate(&stop));
    float msec = 0;
    int iteration = 1000;
    checkCudaErrors(cudaEventRecord(start));
    for(int i = 0; i < iteration; i++)
    {
        //sigmoid<<<CeilDiv(N, block_size), block_size>>>(d_A, d_B, N);
        //sigmoid_float4<<<CeilDiv(N, block_size), block_size/4>>>(d_A, d_B, N);
        sigmoid_float4<<<CeilDiv(N/4,block_size), block_size>>>(d_A, d_B, N);
    }

    checkCudaErrors(cudaEventRecord(stop));
    checkCudaErrors(cudaEventSynchronize(stop));
    checkCudaErrors(cudaEventElapsedTime(&msec, start, stop));
    printf("sigmoid takes %.3f msec\n", msec/iteration);
    checkCudaErrors(cudaMemcpy(h_B, d_B, bytes_B, cudaMemcpyDeviceToHost));

    for(int i = 0; i < N; i++){
        double err = fabs(h_B[i] - 1.0f/(1.0f + expf(-h_A[i])));
        if(err > 1.e-6) {
            printf("wrong answer!\n");
            break;
        }
    }
    
    cudaFree(d_A);
    cudaFree(d_B);

    free(h_A);
    free(h_B);

    return 0;
}

编译和运行代码：

nvcc -o sigmoid sigmoid.cu
./sigmoid

4.2、relu.cu

#include <stdio.h>
#include <stdlib.h>
#include <float.h>
#include <vector>
#include<assert.h>
#include <algorithm>
#include <cublas_v2.h>
#include <cuda_runtime.h>

#define FLOAT4(value)  *(float4*)(&(value))

#define checkCudaErrors(func)               \
{                                   \
    cudaError_t e = (func);         \
    if(e != cudaSuccess)                                        \
        printf ("%s %d CUDA: %s\n", __FILE__,  __LINE__, cudaGetErrorString(e));        \
}

//nvcc -o relu relu.cu && ./relu
//sigmoid<<<CeilDiv(N, block_size), block_size>>>(d_A, d_B, N)
//a: Nx1, b: Nx1, c: Nx1, y = relu(x)
__global__ void relu(float* x, float* y, int N){
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if(idx < N) y[idx] = fmaxf(0.0f, x[idx]);
}

__global__ void relu_float4(float* x, float* y, int N){

    int idx = (blockIdx.x * blockDim.x + threadIdx.x) * 4;
    if(idx < N){
        float4 tmp_x = FLOAT4(x[idx]);
        float4 tmp_y;
        tmp_y.x = fmaxf(0.0f, tmp_x.x);
        tmp_y.y = fmaxf(0.0f, tmp_x.y);
        tmp_y.z = fmaxf(0.0f, tmp_x.z);
        tmp_y.w = fmaxf(0.0f, tmp_x.w);
        FLOAT4(y[idx]) = tmp_y;
    }
}


template <typename T> 
inline T CeilDiv(const T& a, const T& b) {
    return (a + b - 1) / b;
}


int main(){

    size_t block_size = 128;
    size_t N = 1 * 1024;
    size_t bytes_A = sizeof(float) * N;
    size_t bytes_B = sizeof(float) * N;

    float* h_A = (float*)malloc(bytes_A);
    float* h_B = (float*)malloc(bytes_B);

    for( int i = 0; i < N; i++ ){
        h_A[i] = (i / 666) * ((i % 2 == 0) ? 1: -1);
    }

    float* d_A;
    float* d_B;

    checkCudaErrors(cudaMalloc(&d_A, bytes_A));
    checkCudaErrors(cudaMalloc(&d_B, bytes_B));

    checkCudaErrors(cudaMemcpy( d_A, h_A, bytes_A, cudaMemcpyHostToDevice));

    cudaEvent_t start, stop;
    checkCudaErrors(cudaEventCreate(&start));
    checkCudaErrors(cudaEventCreate(&stop));
    float msec = 0;
    int iteration = 1000;
    checkCudaErrors(cudaEventRecord(start));
    for(int i = 0; i < iteration; i++)
    {
        //relu<<<CeilDiv(N, block_size), block_size>>>(d_A, d_B, N);
        //relu_float4<<<CeilDiv(N, block_size), block_size/4>>>(d_A, d_B, N);
        relu_float4<<<CeilDiv(N/4, block_size), block_size>>>(d_A, d_B, N);
    }

    checkCudaErrors(cudaEventRecord(stop));
    checkCudaErrors(cudaEventSynchronize(stop));
    checkCudaErrors(cudaEventElapsedTime(&msec, start, stop));
    printf("relu takes %.3f msec\n", msec/iteration);
    checkCudaErrors(cudaMemcpy(h_B, d_B, bytes_B, cudaMemcpyDeviceToHost));

    for(int i = 0; i < N; i++){
        double err = fabs(h_B[i] - fmaxf(0.0f, h_A[i]));

        if(err > 1.e-6) {
            printf("wrong answer!\n");
            break;
        }
    }
    
    cudaFree(d_A);
    cudaFree(d_B);

    free(h_A);
    free(h_B);

    return 0;
}

编译和运行代码：

nvcc -o relu relu.cu
./relu

4.3、histogram.cu

#include <stdio.h>
#include <stdlib.h>
#include <float.h>
#include <vector>
#include<assert.h>
#include <algorithm>
#include <cublas_v2.h>
#include <cuda_runtime.h>

#define INT4(value) *(int4*)(&(value))
#define FLOAT4(value)  *(float4*)(&(value))

#define checkCudaErrors(func)               \
{                                   \
    cudaError_t e = (func);         \
    if(e != cudaSuccess)                                        \
        printf ("%s %d CUDA: %s\n", __FILE__,  __LINE__, cudaGetErrorString(e));        \
}


/*
x[i] = i
        -
      - -
    - - -
  - - - -
- - - - -
考虑一个warp里相邻线程对全局内存y的访问是否合并(coalesced global access)
warp thread[0]: 0, 1, 2, 3, 
warp thread[1]: 4, 5, 6, 7
*/
__global__ void histogram(int* x, int* y, int N){

    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if(idx < N) atomicAdd(&y[x[idx]], 1);
}

__global__ void histogram_int4(int* x, int* y, int N) {
  int idx = 4 * (blockIdx.x * blockDim.x + threadIdx.x);
  if (idx < N) {
    int4 tmp_y = INT4(x[idx]);
    atomicAdd(&(y[tmp_y.x]), 1);
    atomicAdd(&(y[tmp_y.y]), 1);
    atomicAdd(&(y[tmp_y.z]), 1);
    atomicAdd(&(y[tmp_y.w]), 1);
  }
}

template <typename T> 
inline T CeilDiv(const T& a, const T& b) {
    return (a + b - 1) / b;
}

int main(){

    size_t block_size = 128;
    size_t N = 32 * 1024 * 1024;
    size_t bytes_A = sizeof(int) * N;
    size_t bytes_B = sizeof(int) * N;

    int* h_A = (int*)malloc(bytes_A);
    int* h_B = (int*)malloc(bytes_B);

    for( int i = 0; i < N; i++ ){
        h_A[i] = i;
    }

    int* d_A;
    int* d_B;

    checkCudaErrors(cudaMalloc(&d_A, bytes_A));
    checkCudaErrors(cudaMalloc(&d_B, bytes_B));

    checkCudaErrors(cudaMemcpy( d_A, h_A, bytes_A, cudaMemcpyHostToDevice));

    cudaEvent_t start, stop;
    checkCudaErrors(cudaEventCreate(&start));
    checkCudaErrors(cudaEventCreate(&stop));
    float msec = 0;

    int iteration = 1;
    checkCudaErrors(cudaEventRecord(start));
    for(int i = 0; i < iteration; i++){
        //histogram<<<CeilDiv(N, block_size), block_size>>>(d_A, d_B, N);
        //histogram_int4<<<CeilDiv(N, block_size), block_size/4>>>(d_A, d_B, N);
        histogram_int4<<<CeilDiv(N/4, block_size), block_size>>>(d_A, d_B, N);
    }

    checkCudaErrors(cudaEventRecord(stop));
    checkCudaErrors(cudaEventSynchronize(stop));
    checkCudaErrors(cudaEventElapsedTime(&msec, start, stop));
    printf("histogram takes %.3f msec\n", msec/iteration);

    checkCudaErrors(cudaMemcpy(h_B, d_B, bytes_B, cudaMemcpyDeviceToHost));

    for(int i = 0; i < N; i++){
        //all ones;
        double err = fabs(h_B[i] - iteration * 1.0f);
        if(err > 1.e-6) {
            printf("wrong answer!\n");
            break;
        }
    }

    cudaFree(d_A);
    cudaFree(d_B);

    free(h_A);
    free(h_B);

    return 0;
}

编译和运行代码：

nvcc -o histogram histogram.cu
./histogram