11月21日にリリースされた CUDA 10.2 の Release Note を読んでみると、さらっと『Added support for CUDA Virtual Memory Management APIs.』という一文が。
以前から、ManagedなGPUデバイスメモリをマルチプロセスで共有できるようにしてほしいと、機能リクエストを上げていた事もあり、これがそれに該当するのかね？と、少し手を動かして試してみた。
忙しい人のために、先に結論だけ書くと、現状のCUDA Toolkit 10.2の対応範囲では、既存のcuIpcGetMemHandle()とcuIpcOpenMemHandle()で実現できる機能に差はないものの、APIの仕様・拡張可能性としては、将来的にManaged Device Memoryをマルチプロセスで共有するための布石であるようにも見えるので、引き続き注視していきたい。
なお、試行錯誤の跡（サンプルプログラム）はこちら。
https://github.com/heterodb/toybox/tree/master/gpu_mmap
背景 － なぜManaged Device Memoryをマルチプロセスで共有したいか
JOIN処理をGpuJoinで実行する場合、基本的には (N-1) 個の比較的小さな（GB単位の）テーブルと、1個の巨大なテーブル（数十GB以上～）との結合というのが基本線になる。これはマスタ表とファクト表の結合というのを念頭に置いた設計であるが、説明の簡単のため、左の小さいテーブル群をInner側、右の大きなテーブルをOuter側と呼ぶ。
GpuJoinが内部的に使っている Hash Join のアルゴリズム上、Outer側を読み出しながら少しずつ処理を進めていく事はできるのだが、少なくとも Inner 側はオンメモリに、しかもGPUを使用するのでデバイスメモリに載せておく必要がある。*1
そうすると、Inner側でオンメモリにできるデータサイズは自ずとGPUの搭載メモリ容量に規定されてしまうという事がまず前提としてある。難しいのは、ここでInner側のバッファに載せる必要があるのは、Inner側テーブルを条件句付きでスキャンした結果であるので、実行計画では1,000行だと思っていたのに、実際にデータを読んでみたら百万行ありましたという可能性もある事。つまり、必ずしもオプティマイザで弾ける訳ではない。
で、もう一つ考慮すべき事項として、PostgreSQL v9.6で並列クエリがサポートされた事。
これにより、通常のバックエンドプロセスに加えて、ワーカープロセスが、並列に全く同一のGpuJoinを実行する可能性を考慮する必要が出てきたが、これを CUDA から見ると、互いに独立な複数のプロセスがGPUメモリを奪い合いながらGPU kernelを動かしているように見える。
はっきり言えば、デバイスメモリ制限の厳しい GpuJoin ワークロードにとっては最悪である。並列で動作する各プロセス毎にデバイスメモリを確保に走れば、あっという間にデバイスメモリが枯渇してしまわないとも限らない。
そこで現在のPG-Stromでは、Inner側バッファの内容を一個だけ作り、それをcuIpcGetMemHandle()とcuIpcOpenMemHandle()を使って各ワーカープロセスと共有する。Inner側バッファの内容は一度作ってしまえばRead-Onlyなので、デバイスメモリの消費量はほぼ変わらない。
ただし、これらIPC系の機能を使えるのはcuMemAllocで獲得したメモリ、即ち、cuMemAlloc呼び出しの時点で実際にGPUデバイスメモリの物理ページが割り当てられ、決してPage-outしない代わりにOver-subscriptionも不可という制限がある。
どういう事かというと、デバイスメモリの Over-subscription が不能であるために『あと数十MBあればInner側バッファをデバイスメモリに載せられたのに』という状況で涙を飲むしかなくなる。
一方、これが Managed Device Memory と呼ばれる領域であれば、アロケーション時点では仮想アドレス空間だけを確保しておき、プログラムが実際にその領域をアクセスした時点でページを割り当て、必要があればホストメモリ上のバッファとスワップする事で処理を継続する事もできる。
が、その場合、前述のIPC系機能は対応していないため、マルチプロセス下でGpuJoinのInner側バッファとして利用する場合には、同じデータを重複して何ヶ所にも持つ事が必要となり無駄が多い。あっちを立てれば、こっちが立たず。
全体的な流れ
CUDA Driver APIのVirtual Memory Managementの章を読んでみると、全体的な流れは以下のような形である事がわかる。
まず、デバイスメモリをエクスポートする側の流れ：
cuMemCreate()を用いて、memory allocation handleを作成する。この人が物理的に割り当てられたメモリ領域の実体になる。- OSで言えば、
shm_open(3)とfallocate(2)みたいな感じ
- OSで言えば、
cuMemExportToShareableHandle()を用いて、memory allocation handleを指し示すユニークな識別子であるshareable-handleに変換する。
次に、デバイスメモリをインポートする側の流れ
cuMemImportFromShareableHandle()を使って、上記のshareable-handleを、ローカルなmemory allocation handleに変換する。- OSで言えば、他のプロセスから渡された共有メモリのパス名を使って
open(2)するようなもの。
- OSで言えば、他のプロセスから渡された共有メモリのパス名を使って
cuMemAddressReserve()を使って、このデバイスメモリをマッピングする仮想アドレス領域を予めリザーブしておく。cuMemMap()を使って、手順(1)で取得した memory allocation handle を、手順(2)で確保したアドレス空間にマップする。cuMemSetAccess()を使って、上記アドレス空間のアクセス許可を変更する。Read-onlyとRead-Writableの２種類が設定可- OSで言えば、上記の3ステップは
mmap(2)が一度に実行する事になる。
- OSで言えば、上記の3ステップは
なお、cuMemCreate()を使ってデバイスメモリを作成したプロセス自身がそれをマップして使いたい場合は、cuMemExportToShareableHandle()とcuMemImportFromShareableHandle()を使うまでもなく、ローカルのmemory allocation handleを持っているので、単純にcuMemAddressReserve()、cuMemMap()、cuMemSetAccess()を呼び出せばよい。
注目すべきAPI：cuMemCreate()
gpu_mmap.c:L304で呼び出しているcuMemCreateの定義は以下のようになっている。
CUresult cuMemCreate ( CUmemGenericAllocationHandle* handle, size_t size, const CUmemAllocationProp* prop, unsigned long long flags);
第一引数は処理結果の memory allocation handle を返すポインタで、第二引数はサイズ。第三引数に、どういったメモリが欲しいのかという属性をセットする。
CUmemAllocationProp構造体の定義は以下のようになっており、Linuxで関係するのは最初の3つだけ。
typedef struct CUmemAllocationProp_st {
/** Allocation type */
CUmemAllocationType type;
/** requested ::CUmemAllocationHandleType */
CUmemAllocationHandleType requestedHandleTypes;
/** Location of allocation */
CUmemLocation location;
/**
* Windows-specific LPSECURITYATTRIBUTES required when
* ::CU_MEM_HANDLE_TYPE_WIN32 is specified. This security attribute defines
* the scope of which exported allocations may be tranferred to other
* processes. In all other cases, this field is required to be zero.
*/
void *win32HandleMetaData;
/** Reserved for future use, must be zero */
unsigned long long reserved;
} CUmemAllocationProp;で、CUmemAllocationTypeを定義を確認してみると、CUDA 10.2時点で有効なのはCU_MEM_ALLOCATION_TYPE_PINNEDのみ。コメントを読むと、要はcuMemAllocと同じようにふるまうという事である。*2
typedef enum CUmemAllocationType_enum {
CU_MEM_ALLOCATION_TYPE_INVALID = 0x0,
/** This allocation type is 'pinned', i.e. cannot migrate from its current
* location while the application is actively using it
*/
CU_MEM_ALLOCATION_TYPE_PINNED = 0x1,
CU_MEM_ALLOCATION_TYPE_MAX = 0xFFFFFFFF
} CUmemAllocationType;次にCUmemAllocationHandleTypeを確認すると、これは、OS上で交換可能な shareable-handle をどのように表現するかという話で、LinuxであればCU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTORを選択し、識別子はintの幅を持つとされる。
typedef enum CUmemAllocationHandleType_enum {
/**< Allows a file descriptor to be used for exporting. Permitted only on POSIX systems. (int) */
CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR = 0x1,
/**< Allows a Win32 NT handle to be used for exporting. (HANDLE) */
CU_MEM_HANDLE_TYPE_WIN32 = 0x2,
/**< Allows a Win32 KMT handle to be used for exporting. (D3DKMT_HANDLE) */
CU_MEM_HANDLE_TYPE_WIN32_KMT = 0x4,
CU_MEM_HANDLE_TYPE_MAX = 0xFFFFFFFF
} CUmemAllocationHandleType;三番目のCUmemLocationの定義は以下の通り。要はCU_MEM_LOCATION_TYPE_DEVICEを指定して、何番目のGPUでメモリを確保しますかという事になり、現状ではデバイスメモリ以外の選択肢はない。
typedef enum CUmemLocationType_enum {
CU_MEM_LOCATION_TYPE_INVALID = 0x0,
/**< Location is a device location, thus id is a device ordinal */
CU_MEM_LOCATION_TYPE_DEVICE = 0x1,
CU_MEM_LOCATION_TYPE_MAX = 0xFFFFFFFF
} CUmemLocationType;
typedef struct CUmemLocation_st {
/**< Specifies the location type, which modifies the meaning of id. */
CUmemLocationType type;
/**< identifier for a given this location's ::CUmemLocationType. */
int id;
} CUmemLocation;今回、新たに追加されたAPIではあるが、こうして見てみるとCUmemAllocationPropのフラグを変えるだけで、将来的にはManaged Device Memoryのエクスポートとマルチプロセスでの共有も可能にできるようなAPIの切り方をしている。
それを強く感じるのが、次のcuMemAddressReserveである。
注目すべきAPI：cuMemAddressReserve()
CUDAのManaged Memoryのもう一つ大きな特徴として、Over-subscriptionが可能である以外にも、Host側とDevice側で同一の仮想アドレスを利用して、Page Faultをうまくハンドリングする事で、物理ページがHost側なのかDevice側なのか意識せずに（裏でデータ転送が動いている）プログラムを書けるという事があった。
しかし、これをマルチプロセスに拡張するとなると、あるプロセスAでcuMemAllocManagedを呼び出して確保したメモリの仮想アドレスが、他のプロセスBでは既に別の用途に利用されており、そうすると「Host側とDevice側で同一の仮想アドレス」という前提が損なわれるという問題があった。
おそらく、cuMemMapに先立って、わざわざ仮想アドレス空間を予約するAPIを呼び出させるという事は、『Managed Device Memoryをマルチプロセスで共有できるようになっても、当該メモリを共有する全てのプロセスで同じ仮想アドレスに揃うよう保証する事はできないので、アプリケーション側で予め仮想アドレス空間を予約して他が使えないように抑えておいてくれ』というNVIDIAからのメッセージではないだろうか？
OS上の共有メモリでも、別プロセスに貼り付けた時に同一の仮想アドレスを取れるようPROT_NONEでmmapしておくというのはちょくちょく使うテクニックなので、まぁ、そういう事なんだろうと思う。*3
kaigai.hatenablog.com
注目すべきAPI：cuMemImportFromShareableHandle()
コレは注目すべきというか、注意すべきAPI
Linuxの場合、shareable-handleとして使われるのは file descriptor と確かに書いてあったが、当初、cuMemExportToShareableHandleで返ってきた値を、そのまま別プロセスでcuMemImportFromShareableHandleに渡せばよいと思って２～３時間ハマった（汗
これは当然と言えば当然だが、エクスポート元のプロセスにとっての file-descriptor = XX は、エクスポート先のプロセスにとっての file-descriptor = XX ではないので、UNIXドメインソケットとsendmsg(2)/recvmsg(2)にSCM_RIGHTSを使って、file-descriptorの複製を作ってやる必要がある。
gpu_mmap.c:L98以降の一連の処理は、その辺のファイルディスクリプタの受け渡しに関するもの。
折角なので、バッファのサイズとGPUのデバイスIDをメッセージに載せて渡すようにしている。
注目すべきAPI：cuMemSetAccess()
cuMemMapでデバイスメモリをマップするだけでは、この領域に対するアクセス権が無いのでCUDA_ERROR_ILLEGAL_ADDRESSエラーになってしまう。
なので、CU_MEM_ACCESS_FLAGS_PROT_READかCU_MEM_ACCESS_FLAGS_PROT_READWRITEを指定する必要がある。
当初、Over-subscriptionが使えないなら、せめてInner側バッファをread-writableで作成し、その後read-onlyに変えるようにしたら、不慮の事故でメモリ破壊を防げる的な利点はあるかと思ったが…。
↓インポート側プロセスでCU_MEM_ACCESS_FLAGS_PROT_READを指定した場合。
[kaigai@magro gpu_mmap]$ ./gpu_mmap total sum [master]: 6420001848.309442 gpu_mmap.c:259 failed on cuMemSetAccess: CUDA_ERROR_INVALID_VALUE gpu_mmap.c:259 failed on cuMemSetAccess: CUDA_ERROR_INVALID_VALUE
CU_MEM_ACCESS_FLAGS_PROT_READWRITEを指定しないと怒られました。Orz...
結論
サンプルプログラムは、エクスポート側プロセスでデバイスメモリを確保し、それを（なんちゃって）ランダムな値で埋めたあと、インポート側プロセスでそれをマップし、それぞれ合計値を計算するだけのGPUカーネルを実行するというもの。
このように、明示的なデータのコピーは無いものの、同じメモリ領域に対して同じ集計演算を実行した結果、同じ結果を出力している。（浮動小数点同士の加算なので、集計のタイミングによって多少の誤差が出るのは仕方ない）
[kaigai@magro gpu_mmap]$ ./gpu_mmap total sum [master]: 6420001848.309442 total sum [17767]: 6420001848.309443 total sum [17768]: 6420001848.309442
これにより、CUDA 10.2で追加されたVirtual Memory Management機能により、GPUのデバイスメモリを複数のプロセスで共有できる事が確認できた。
駄菓子菓子。現状では Managed Device Memory に対応しておらず、できる事は既存のcuMemAllocとcuIpcGetMemHandle と変わらない。ファイルディスクリプタの受け渡しなど、面倒な要素はあるので、今の時点で積極的に利用する必要性はない。
ただし、おそらく今後の CUDA Toolkit のバージョンアップの中で、どこかのタイミングで Managed Device Memory への対応が（ひっそりと）行われるであろう。その時になって、慌ててアプリケーションの大規模な改修を行わずに済むよう、頭の片隅に当該API群の事を留め置いてもよいだろう。
