RAPIDS Memory Allocator

개요

RAPIDS Memory Manager(RMM)는 High-frequency memory allocation/deallocation에서 발생하는 병목현상을 줄이고, 다양한 메모리 할당 전략을 제공하는 Memory management library이다. Pooling을 사용하는 경우, 큰 메모리 영역을 확보한 뒤 개별 Allocation 요청에 따라 Suballocation하고 재사용함으로써 cudaMalloc이나 cudaFree의 호출 빈도와 관련 Overhead를 줄인다.

RMM은 공통된 Memory resource interface를 통해서 단순한 CUDA Allocation부터 Pooling, 크기별 Allocation, Logging 등 다양한 방식을 선택하고 조합할 수 있도록 한다.

배경

GPU workflow를 최적화하기 위해서는 Host와 Device memory의 할당 방식도 함께 고려해야 한다. 예를 들어서 빈번한 Device memory allocation의 비용을 줄이기 위해 Memory pool을 사용하거나, 비동기 Host-device memory transfer를 위해 Pinned host memory를 사용할 수 있다.

더 나아가서 Memory diagnostics, logging, leak detection, profiling, debugging 등의 기능을 메모리 할당 과정에 적용할 필요가 있다. 이러한 기능을 각각 구현하기보다는 공통된 Memory allocation interface를 통해서 제공하면 메모리 관리를 쉽게 확장할 수 있다.

RMM은 이런 환경에서 Device와 Host memory allocation을 쉽고 효과적으로 Customizing할 수 있는 Interface와 이를 사용하는 자료구조를 제공한다. RMM은 이미 NVIDIA RAPIDS 라이브러리 전반에 걸쳐서 사용되고 있다.

내부 구조

Stream-ordered Memory Allocation

RMM은 Stream을 인자로 받는 Allocation과 Deallocation interface를 제공한다. 전통적인 cudaMalloc과 cudaFree는 Stream을 인자로 받지 않으며, 특히 cudaMalloc으로 할당한 메모리를 cudaFree로 해제하는 과정에서는 Device 전체의 Synchronization이 발생할 수 있다.

Stream-ordered allocation은 메모리의 할당, 사용, 해제, 재사용 순서를 Stream의 실행 순서와 연결한다. 같은 Stream에서는 앞선 작업이 끝난 뒤 다음 작업이 실행되므로, 해제된 메모리를 재사용할 때 별도로 Stream 전체가 끝나기를 기다리는 Synchronization을 피할 수 있다. 즉, 매번 해당 Stream을 Synchronize하는 방식이 아니라 Stream ordering을 이용해서 불필요한 동기화를 줄이는 방식이다.

Suballocation in RMM

RMM은 여러 종류의 Suballocator를 제공한다. Suballocator는 Upstream memory resource에서 큰 메모리 영역을 확보한 뒤, 이를 작은 영역으로 나누어서 개별 Allocation 요청을 처리한다. 따라서 매번 CUDA Allocation을 수행하지 않고도 Pool 내부에서 메모리를 할당하고 재사용할 수 있다.

pool_memory_resource는 Coalescing best-fit 방식의 Suballocator이다. 메모리가 Free되면 주소상 인접한 Free block들과 병합해서 Fragmentation을 줄인다. 또한 Stream마다 Free list를 관리해서 같은 Stream에서 해제된 Block을 동기화 없이 재사용할 수 있도록 한다. 해당 Stream의 Free list에 충분한 크기의 Block이 없으면 다른 Stream의 Free list를 검색할 수 있으며, 이 경우 Event를 통해서 필요한 실행 순서를 보장한다. 따라서 Stream마다 완전히 독립된 Pool이 존재한다기보다는, 재사용 가능한 Block을 Stream별 Free list로 관리한다고 보는 것이 정확하다.

binning_memory_resource는 Allocation 크기에 따라 서로 다른 Memory resource로 요청을 전달한다. 예를 들어서 작은 Allocation은 fixed_size_memory_resource로 처리하고, 큰 Allocation은 Upstream memory resource로 전달하는 Hybrid approach를 구성할 수 있다.

arena_memory_resource는 Global arena와 Default stream을 위한 Per-thread arena, Non-default stream을 위한 Per-stream arena를 사용한다. 각 Local arena는 Global arena에서 Superblock 단위로 메모리를 받아서 관리한다. 작은 Allocation은 Local arena에서 처리하고 큰 Allocation은 Global arena를 활용하는 방식으로, 여러 Thread와 Stream이 함께 사용하는 환경에서 동시성과 Fragmentation을 고려한다.