A Simple Vulkan Memory Manager

As part of my journey towards understanding the vulkan api, i tried to implement a simple memory management scheme from scratch. The core idea of the system is to divide the available memory into pools and then to sub-allocate from the appropriate pool.

A pool has the following structure:

struct xeVulkanMemoryPool
{
  VkDeviceMemory deviceMemory;
  VkDeviceSize poolSize;
  VkDeviceSize blockSize;

  xeQueue blockQueue;
};

The device memory is allocated via vkAllocateMemory and divided into blocks of equal size. The queue maintains the offsets of free blocks.

Allocating a block is then simply:

static xeResult xe_allocate_pool_memory(xeVulkanMemoryPool* pool, xeVulkanMemoryAllocation* allocation)
{
  VkDeviceSize offset;
  xeResult result = XE_SUCCESS;
  if (xe_pop_queue(&pool->blockQueue, &offset) == XE_QUEUE_EMPTY)
    result = XE_OUT_OF_MEMORY;
  else {
    allocation->deviceMemory = pool->deviceMemory;
    allocation->offset = offset;
  }
  return result;
}

The used queue is lockless, so this function is thread-safe.

Freeing a block is done by pushing the offset into the queue.

The more interesting part comes from managing the pools. This is the responsiblity of the “memory manager”. Its structure is shown below:

struct xeVulkanMemoryManager
{
  VkDevice device;
  VkPhysicalDeviceMemoryProperties memoryProperties;

  uint32_t maxMemoryPoolCount; 

  /* we track the amount of memory we allocated */
  uint32_t deviceLocalBudget;
  uint32_t deviceLocalHostVisibleBudget;
  uint32_t usedDeviceLocalBudget;
  uint32_t usedDeviceLocalHostVisibleBudget;

  VkDeviceSize maxAllocationSize;

  /* Hash from blockSize + memory type + required properties to the index
   * of the first pool that satisfies these requirements.
   * If the pool is full check memoryPoolLinks[index] for another pool.
   * If no suitable pool with available memory is found, create a new pool.
   */
  xeHash poolHash;
  xeVulkanMemoryPool* memoryPools;
  uint32_t* memoryPoolLinks;
};

The algorithm for allocating a block of device memory works as follows (pseudo-c)

xeResult xe_allocate_vulkan_memory(xeVulkanMemoryManager* memoryManager,
                                   uint32_t memoryTypeBits,
                                   VkMemoryPropertyFlags requiredProperties,
                                   VkDeviceSize size,
                                   VkDeviceSize alignment,
                                   xeVulkanMemoryAllocation* allocation)
{
  /* Determine block size from size+alignment. 
   * and round to nearest multiple of 64 kb.
   * This keeps the number of different pools smaller. */
  size = XE_MAX(size, alignment); 
  size = round_to_multiple(size, XE_KB(64));

  /* Construct the lookup key.
   * The lower 32 bits are made up of the memory type and the block size,
   * the upper 32 bits contain the required memory properties */
  uint64_t log2BlockSize = log2(size);
  uint32_t memoryTypeIndex = find_vulkan_memory_type_index(memoryManager, memoryTypeBits, requiredProperties);
  uint64_t key = memoryTypeIndex | (log2BlockSize << VK_MAX_MEMORY_TYPES) | (uint64_t)requiredProperties << 32);
  
  uint32_t poolIndex = xe_hash_lookup(memoryManager->poolHash, key, NO_SUITABLE_POOL);
  if (poolIndex == NO_SUITABLE_POOL) {
    /* Create a new pool with the required parameters, insert into the hash table and allocate from that pool */
  }
  else {
    /* Try to allocate from that pool */
    if (xe_allocate_pool_memory(&memoryManager->memoryPools[poolIndex], allocation) == XE_SUCCESS)
      return XE_SUCCESS;
    
    /* Try the next suitable pool */
    uint32_t lastValidPoolIndex = poolIndex;
    while (1) {
      if (poolIndex == NO_SUITABLE_POOL)
        break;
      poolIndex = memoryManager->memoryPoolLinks[poolIndex];
      if (xe_allocate_pool_memory(&memoryManager->memoryPools[poolIndex], allocation) == XE_SUCCESS)
        return XE_SUCCESS;
      lastValidPoolIndex = poolIndex;
    }

    /* Create a new pool, update memoryPoolLinks[lastValidPoolIndex] and allocate from the new pool */
  }

  /* Allocation failed */
  return XE_OUT_OF_MEMORY;
}

This system is probably not as sophisticated as the one used by VMA but it is easy to understand, which in my opinion is the most valuable property of a learning project.