CUDAの高速化の復習2023年版 Histogram（主にatomicAdd）編

Reduction、vectrized memory accessに続き、今回はhistogramを題材にして主にatomicAddのパフォーマンスが最近どうなっているのかを見ていきたいと思います。

HistogramはCUDA Samplesの中にもありますが、全然違う実装が、NVIDIAのA100の最適化に関する発表の資料の中で紹介されています。この資料ではatomicAddとL2キャッシュの「persistent data accesses」を利用してhistogramの実装をしています。このpersistent data accessesは発表当時気になっていて、あとで調べようと思って忘れてたのですが、最近調べたのでせっかくなので記事にしました。

今回調査するうえで特に知りたかった点としてpersistent data accessesを使う場合とshared memoryを使う場合だとどっちが速いのかというものがあります。A100の最適化の資料の中にはこれについて特に書いてなかったので、この二つの性能にどれくらい差があるのかの比較を行いました。

検証で使ったコードはこちらにあげてあります。

https://github.com/shu65/cuda-histogram

Histogramと今回対象とする部分に関して

Histogramに関して知っている方も多いと思いますが、どういうものか簡単に紹介します。Histogramはデータの範囲をいくつかのbinに区切り、データの中の各要素がどのbinに含まれるかを計算し、binごとに含まれる要素の個数をカウントするというものになります。

上記の説明の通り、histogramを計算するうえで、大まかに3つくらいのステップに分けることができます。

1. 各binの範囲を決める
2. データの各要素がどのbinに入るのかを計算する
3. bin毎に何個のデータの要素が含まれるかをカウントする

このHistogramは入力データによってどういう風にbinの範囲を決めればよいかが変わるため、データに応じて1,2あたりの処理はデータに応じて変化させる必要があります。また、GPU的にも難しいのは3のところなため、今回は3に注目して説明します。

3の部分は入力としては各要素がどのbinに入るかを表したbinのidの配列を受け取り、bin毎に何個要素があるかをカウントするという処理になります。どういう処理かイメージしやすいようにCPU版のコードを以下に示します。

#include <stdint.h>

void HistogramCPU(const int *data, const uint32_t n, const uint32_t n_bins, uint32_t *bin_counts)
{
  for (uint32_t i = 0; i < n_bins; ++i)
  {
    bin_counts[i] = 0;
  }

  for (uint32_t i = 0; i < n; ++i)
  {
    const int bin_i = data[i];
    ++bin_counts[bin_i];
  }
}

このあと、このコードと同じ結果になるようなGPUコードを紹介していきます。

GPUでHistogramが難しい理由

先にGPUでhistogramを計算する際、難しいポイントに関して触れておきます。GPUに限らず並列処理でhistogramを計算する際、bin毎のカウントをするところで複数のスレッドが同じメモリ領域にアクセスすることになるので、bin毎のカウント部分で排他制御が必要になります。

GPUで簡単に実装するなら、後ほど示す通りatomicAddを使えばいいのですが、atomicAddは遅いという問題があります。特にglobal memoryに対してのatomicAddはshared memoryに対するものよりも遅いです。このため、個人的にはatomicAdd、特にglobal memoryに対するものは注意が必要な計算という認識でした。

それがA100の最適化の資料でL2キャッシュのpersistent data accessesを使うとましになるよ、ということが書かれています。次にこのL2キャッシュのpersistent data accessesについて詳しく説明します。

L2キャッシュのpersistent data accessesについて

L2キャッシュのpersistent data accessesは、L2キャッシュのメモリ領域を分割してpersistent data用の領域を確保して、よくアクセスするものはpersistent data用の領域にキャッシュしてメモリアクセスを高速化するための機能です。

A100最適化の資料の18ページ目あたりからこの機能の紹介があります。

https://developer.download.nvidia.com/video/gputechconf/gtc/2020/presentations/s21819-optimizing-applications-for-nvidia-ampere-gpu-architecture.pdf

CUDAのprogramming guideでは以下の部分に説明があります。

https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#device-memory-l2-access-management

これを使うと、よくデータアクセスする一部の領域とそれ以外の領域のキャッシュを分けることができます。結果として一部だけ何度もアクセスするという場合はこの機能を使うことで高速化が狙えます。

制限としてはL2キャッシュのすべてをpersistent dataにすることはできず、最大値が決まっています。最大値は以下のようにすると確認できます。

  cudaDeviceProp prop;
  CheckCudaErrors(cudaGetDeviceProperties(&prop, device_id));
  cout << "persistingL2CacheMaxSize:" <<   prop.persistingL2CacheMaxSize << endl;

A100で確認すると30MBが最大値になっています。

使い方としてはprogramming guideにある通り、以下の手順で使うことができます。

Persistent data accesses用の領域を確保

以下のようにpersistent data accesses用の領域として最大どれくらい使うかを設定します。コード中のsize にpersistent data accesses用の領域のサイズを入れてcudaDeviceSetLimitを呼ぶことで、使用するpersistent data accesses用の領域の最大値を設定します。

cudaDeviceSetLimit(cudaLimitPersistingL2CacheSize, size); 

Persistent data accessesの設定適用

次にstream、もしくはcuda graphのnodeに対してpersistent data accessesの設定を行います。ここではstreamに対しての設定方法を示します。programming guideにある通り、以下のように設定していきます。

cudaStreamAttrValue stream_attribute;                                         // Stream level attributes data structure
stream_attribute.accessPolicyWindow.base_ptr  = reinterpret_cast<void*>(ptr); // Global Memory data pointer
stream_attribute.accessPolicyWindow.num_bytes = num_bytes;                    // Number of bytes for persistence access.
                                                                              // (Must be less than cudaDeviceProp::accessPolicyMaxWindowSize)
stream_attribute.accessPolicyWindow.hitRatio  = 1.0;                          // Hint for cache hit ratio
stream_attribute.accessPolicyWindow.hitProp   = cudaAccessPropertyPersisting; // Type of access property on cache hit
stream_attribute.accessPolicyWindow.missProp  = cudaAccessPropertyStreaming;  // Type of access property on cache miss.

//Set the attributes to a CUDA stream of type cudaStream_t
cudaStreamSetAttribute(stream, cudaStreamAttributeAccessPolicyWindow, &stream_attribute);

注意する点としてはhitRatio の値です。hitRatio はアクセスするglobal memoryのサイズがnum_bytesよりも大きい場合は適切に指定しないとパフォーマンスが落ちることが以下の部分で示されています。

https://docs.nvidia.com/cuda/cuda-c-best-practices-guide/#tuning-the-access-window-hit-ratio

このため、hitRatio は自分のケースでどのくらいのサイズにすべきか？を考えて設定したほうがよさそうです。

ただ、今回のhistogramの例では後ほど示す通り、global memoryで最大20MB分の領域をpersistent data accessesに利用するので、persistent data accessesで指定できるサイズに収まります。このため、hitRatio は1.0でOKです。

HistogramのGPU実装

ここからは今回検証に使うhistogramのGPU実装に関してです。3つありますので、順番にどういうものかを説明していきます。

GPU実装のベースライン

まずはGPU実装のベースラインです。コードとしてはCPUをそのままCUDAで実装したような形になっています。

__global__ void HistogramGPUv1Kernel(const int *data, const uint32_t n, uint32_t *bin_counts)
{
  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
  if (tid >= n)
  {
    return;
  }
  const int bin_i = data[tid];
  atomicAdd(bin_counts + bin_i, 1);
}

こちらのコードがA100の最適化の資料で示されているhistogramのコードとほぼ同じものになっています。bin_countsへのアクセスは全スレッドが同時に行うため、atomicAddを使って排他制御しながらカウントするようにしています。

GPU実装のshared memory版

Histogramの計算でbin_countsがshared memoryに収まる範囲であれば、shared memoryを使うという手があります。

先ほど説明した通り、shared memoryへのatomicAddはglobal memoryに比べて速いので、shared memoryを使ってblock毎に集計し、その後各blockの結果をatomicAddを使ってglobal memoryの領域に加算するという方法で計算します。こうすることでglobal memoryへのatomicAddの回数は減らすことができます。コードとしては以下の通りです。

__global__ void HistogramGPUv2Kernel(const int *data, const uint32_t n, const uint32_t n_bins, uint32_t *bin_counts)
{
  cg::thread_block cta = cg::this_thread_block();
  extern __shared__ uint32_t s_bin_counts[];
  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
  const int stride = gridDim.x * blockDim.x;
  if (threadIdx.x < n_bins)
  {
    s_bin_counts[threadIdx.x] = 0;
  }
  cg::sync(cta);
  for (int i = tid; i < n; i += stride)
  {
    const int bin_i = data[i];
    atomicAdd(s_bin_counts + bin_i, 1);
  }
  cg::sync(cta);
  if (threadIdx.x < n_bins)
  {
    uint32_t sum = s_bin_counts[threadIdx.x];
    atomicAdd(bin_counts + threadIdx.x, sum);
  }
}

注意点としてはshared memoryのサイズは最大でA100の場合でも164KBらしいので、bin_counts に必要なサイズがこれ以上のときはこの戦略はそのまま使うことができません。

GPU実装のshared memory + reduction版

shared memory版では最後global memoryへの加算はatomicAddを利用しましたが、この部分もatomicAddなしで実行するようにします。具体的にはCUDAにおけるreductionのような戦略をとり、各blockが計算した結果をCUDAにおけるparallel reductionに似たアルゴリズムで集計します。コードとしては以下の通り。

__global__ void HistogramGPUv3Kernel(const int *data, const uint32_t n, const uint32_t n_bins, uint32_t *tmp_bin_counts)
{
  cg::thread_block cta = cg::this_thread_block();
  extern __shared__ uint32_t s_bin_counts[];
  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
  const int tmp_bin_counts_offset = blockIdx.x * n_bins;
  const int stride = gridDim.x * blockDim.x;
  if (threadIdx.x < n_bins)
  {
    s_bin_counts[threadIdx.x] = 0;
  }
  cg::sync(cta);
  for (int i = tid; i < n; i += stride)
  {
    const int bin_i = data[i];
    atomicAdd(s_bin_counts + bin_i, 1);
  }
  cg::sync(cta);
  if (threadIdx.x < n_bins)
  {
    uint32_t sum = s_bin_counts[threadIdx.x];
    tmp_bin_counts[tmp_bin_counts_offset + threadIdx.x] = sum;
  }
}

__global__ void HistogramGPUv3MergeKernel(const uint32_t *tmp_bin_counts, const int n, uint32_t *bin_counts)
{
  cg::thread_block cta = cg::this_thread_block();
  extern __shared__ uint32_t s_data[];

  uint32_t sum = 0;
  for (int i = threadIdx.x; i < n; i += blockDim.x)
  {
    sum += tmp_bin_counts[blockIdx.x + i * blockDim.x];
  }
  s_data[threadIdx.x] = sum;
  for (uint stride = blockDim.x / 2; stride > 0; stride >>= 1)
  {
    cg::sync(cta);
    if (threadIdx.x < stride)
    {
      s_data[threadIdx.x] += s_data[threadIdx.x + stride];
    }
  }

  if (threadIdx.x == 0)
  {
    bin_counts[blockIdx.x] = s_data[0];
  }
}

ちなみに、アルゴリズム的にこれはほぼCUDA Samplesに含まれるhistogramと同じになります。

評価

今回、binの数で傾向が変わったので、以下の2種類のデータで比較します。

1. 入力データ数は256M個、binの数が256
2. 入力データ数は256M個、binの数が5M個

1つ目のほうがCUDA Samplesに含まれるhistogramの条件に近いもので、2つ目がA100の最適化の資料に書かれている条件になります。また、1のほうは先ほど紹介したアルゴリズムすべてが実行できますが、2つ目のほうはshared memoryが足りないのでベースラインのみとなっています。

また、persistent data accessesのありなしでどれくらい計算結果が変化するのかも知りたいので、各アルゴリズムでbin_countsの部分にpersistent data accessesを使う場合と使わなかった場合も比較します。

時間の計測方法としては10回の平均時間を算出して比較します。実行環境としてはCUDA 12.0、A100を利用しています。

計測した計算時間はそれぞれ以下の通りです。

表からわかる通り、1のデータではglobal memoryへのatomicAddが少なければ少ないほど高速化できていることがわかります。また、persistent data accessesは1のデータでは効果がありませんでした。

また、2のデータに関してはpersistent data accessesありなしで若干差がありますが、今回は約7%の向上と効果は小さいという結果になりました。A100の最適化のほうの資料では43%向上とあるのでどこか設定を間違えているのかもしれません。（いろいろ試しましたがわからなかったのでご存じの方いたら教えていただけるとありがたいです。）

ただ、結果からshared memoryが使える状況下であればL2キャッシュのチューニングするよりもshared memoryを使ったほうが速くなりそうという印象を持ちました。

終わりに

今回、気になっていたatomicAddとL2キャッシュのpersistent data accessesのパフォーマンスについて調べました。結果としてやっぱりglobal memoryへのatomicAddはできるだけ避けたほうがいいということが確認できてよかったです。

これまでも最近のGPU、CUDAを使ってreduction、vectorized memory accessに関しても調査したまとめを書いたのでもしよろしければそちらもご覧ください。