Vulkan Growing Linear Arena

vk_mem_arena.cpp

struct vk_mem_view
{
	VkDeviceMemory	device;
	void*	        host = 0;
};

struct vk_allocation
{
	VkDeviceMemory  deviceMem;
	u8*		hostVisible = 0;
	u64		dataOffset;
};

struct vk_mem_arena
{
	VkDevice				device;
	vector<vk_mem_view>			mem;
	vector<vk_mem_view>			dedicatedAllocs;
	u64					minArenaSize = VK_MIN_DEVICE_BLOCK_SIZE;
	u64					maxParentHeapSize;
	u64					size;
	u64					allocated;
	VkMemoryPropertyFlags			memType;
	VkPhysicalDeviceMemoryProperties	gpuMemProps;
};

inline i32
VkFindMemTypeIdx(
	VkPhysicalDeviceMemoryProperties* pVkMemProps,
	VkMemoryPropertyFlags		requiredProps,
	u32				memTypeBitsRequirement )
{
	for( u64 memIdx = 0; memIdx < pVkMemProps->memoryTypeCount; ++memIdx ){
		u32 memTypeBits = ( 1 << memIdx );
		b32 isRequiredMemType = memTypeBitsRequirement & memTypeBits;

		VkMemoryPropertyFlags props = pVkMemProps->memoryTypes[ memIdx ].propertyFlags;
		b32 hasRequiredProps = ( props & requiredProps ) == requiredProps;
		if( isRequiredMemType && hasRequiredProps )
			return (i32) memIdx;
	}

	VK_CHECK( "Memory type unmatch !" );

	return -1;
}


// TODO: move alloc flags ?
inline static vk_mem_view
VkTryAllocDeviceMem(
	VkDevice				vkDevice,
	u64					size,
	u32					memTypeIdx,
	VkMemoryPropertyFlags			memFlags,
	VkMemoryAllocateFlags			allocFlags,
	const VkMemoryDedicatedAllocateInfo*	dedicated )
{
	VkMemoryAllocateFlagsInfo allocFlagsInfo = {};
	allocFlagsInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_FLAGS_INFO;
	allocFlagsInfo.pNext = dedicated;
	allocFlagsInfo.flags = VK_MEMORY_ALLOCATE_DEVICE_ADDRESS_BIT;// allocFlags;

	VkMemoryAllocateInfo memoryAllocateInfo = {};
	memoryAllocateInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
	memoryAllocateInfo.pNext = &allocFlagsInfo;
	memoryAllocateInfo.allocationSize = size;
	memoryAllocateInfo.memoryTypeIndex = memTypeIdx;

	vk_mem_view vkMemView;
	VK_CHECK( vkAllocateMemory( vkDevice, &memoryAllocateInfo, 0, &vkMemView.device ) );

	if( memFlags & VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT )
		VK_CHECK( vkMapMemory( vkDevice, vkMemView.device, 0, VK_WHOLE_SIZE, 0, &vkMemView.host ) );

	return vkMemView;
}

inline static void
VkArenaInit(
	vk_mem_arena*				vkArena,
	VkPhysicalDeviceMemoryProperties	memProps,
	VkMemoryPropertyFlags			memType,
	VkDevice				vkDevice )
{
	i32 i = 0;
	for( ; i < memProps.memoryTypeCount; ++i )
		if( memProps.memoryTypes[ i ].propertyFlags == memType ) break;

	VK_CHECK( i == memProps.memoryTypeCount );

	vkArena->allocated = 0;
	vkArena->size = 0;
	vkArena->maxParentHeapSize = memProps.memoryHeaps[ memProps.memoryTypes[ i ].heapIndex ].size;
	vkArena->memType = memType;
	vkArena->device = vkDevice;
	vkArena->gpuMemProps = memProps;
	vkArena->minArenaSize =
		( vkArena->maxParentHeapSize < VK_MIN_DEVICE_BLOCK_SIZE ) ?
		( 1 * MB ) :
		VK_MIN_DEVICE_BLOCK_SIZE;
}

inline vk_allocation
VkArenaAlignAlloc( 
	vk_mem_arena*				vkArena, 
	u64					size, 
	u64					align, 
	u32					memTypeIdx,
	VkMemoryAllocateFlags			allocFlags,
	const VkMemoryDedicatedAllocateInfo*	dedicated )
{
	u64 allocatedWithOffset = FwdAlign( vkArena->allocated, align );
	vkArena->allocated = allocatedWithOffset;

	assert( size <= vkArena->maxParentHeapSize );

	if( dedicated ){
		vkArena->dedicatedAllocs.push_back( 
			VkTryAllocDeviceMem( vkArena->device, size, memTypeIdx, vkArena->memType, allocFlags, dedicated ) );
		return { vkArena->dedicatedAllocs.back().device, (u8*) vkArena->dedicatedAllocs.back().host, 0 };
	}

	if( ( vkArena->allocated + size ) > vkArena->size ){
		u64 newArenaSize = max( size, vkArena->minArenaSize );
		vkArena->mem.push_back( 
			VkTryAllocDeviceMem( vkArena->device, newArenaSize, memTypeIdx, vkArena->memType, allocFlags, 0 ) );
		vkArena->size = newArenaSize;
		vkArena->allocated = 0;
	}

	assert( ( vkArena->allocated + size ) <= vkArena->size );
	assert( vkArena->allocated % align == 0 );

	vk_allocation allocId = { vkArena->mem.back().device, (u8*) vkArena->mem.back().host, vkArena->allocated };
	vkArena->allocated += size;

	return allocId;
}

inline void
VkArenaTerimate( vk_mem_arena* vkArena )
{
	for( u64 i = 0; i < vkArena->mem.size(); ++i )
		vkFreeMemory( vkArena->device, vkArena->mem[ i ].device, 0 );
	for( u64 i = 0; i < vkArena->dedicatedAllocs.size(); ++i )
		vkFreeMemory( vkArena->device, vkArena->dedicatedAllocs[ i ].device, 0 );
}
 
inline static void
VkInitMemory( VkPhysicalDevice vkPhysicalDevice, VkDevice vkDevice )
{
	VkPhysicalDeviceMemoryProperties memProps;
	vkGetPhysicalDeviceMemoryProperties( vkPhysicalDevice, &memProps );

	VkArenaInit( &vkRscArena,
				 memProps,
				 VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,
				 vkDevice );
	VkArenaInit( &vkAlbumArena,
				 memProps,
				 VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,
				 vkDevice );
	VkArenaInit( &vkStagingArena,
				 memProps,
				 VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT |
				 VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
				 vkDevice );
	VkArenaInit( &vkHostComArena,
				 memProps,
				 VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT |
				 VK_MEMORY_PROPERTY_HOST_COHERENT_BIT |
				 VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,
				 vkDevice );

	VkArenaInit( &vkDbgArena,
				 memProps,
				 VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,
				 vkDevice );

}

__forceinline u64 VkGetBufferDeviceAddress( VkDevice vkDevice, VkBuffer hndl )
{
#ifdef _VK_DEBUG_
	static_assert( std::is_same<VkDeviceAddress, u64>::value );
#endif

	VkBufferDeviceAddressInfo deviceAddrInfo = {};
	deviceAddrInfo.sType = VK_STRUCTURE_TYPE_BUFFER_DEVICE_ADDRESS_INFO;
	deviceAddrInfo.buffer = hndl;

	return vkGetBufferDeviceAddress( vkDevice, &deviceAddrInfo );
}

// TODO: return/use alignmed size ?
static buffer_data
VkCreateAllocBindBuffer(
	u64			sizeInBytes,
	VkBufferUsageFlags	usage,
	vk_mem_arena*		vkArena )
{
	VkBufferCreateInfo bufferInfo = {};
	bufferInfo.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
	bufferInfo.size = sizeInBytes;
	bufferInfo.usage = usage;
	bufferInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
	VK_CHECK( vkCreateBuffer( vkArena->device, &bufferInfo, 0, &buffData.hndl ) );


	VkMemoryDedicatedRequirements dedicatedReqs = { VK_STRUCTURE_TYPE_MEMORY_DEDICATED_REQUIREMENTS_KHR };
	VkMemoryRequirements2 memReqs2 = {};
	memReqs2.sType = VK_STRUCTURE_TYPE_MEMORY_REQUIREMENTS_2;
	memReqs2.pNext = &dedicatedReqs;

	VkBufferMemoryRequirementsInfo2 buffMemReqs2 = {};
	buffMemReqs2.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_REQUIREMENTS_INFO_2;
	buffMemReqs2.buffer = buffData.hndl;
	vkGetBufferMemoryRequirements2( vkArena->device, &buffMemReqs2, &memReqs2 );

	i32 memTypeIdx = VkFindMemTypeIdx( &vkArena->gpuMemProps, vkArena->memType, memReqs2.memoryRequirements.memoryTypeBits );

	VkMemoryAllocateFlags allocFlags = 
		( usage & VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT ) ? VK_MEMORY_ALLOCATE_DEVICE_ADDRESS_BIT : 0;

	VkMemoryDedicatedAllocateInfo dedicatedAllocateInfo = {};
	dedicatedAllocateInfo.sType = VK_STRUCTURE_TYPE_MEMORY_DEDICATED_ALLOCATE_INFO;
	dedicatedAllocateInfo.buffer = buffMemReqs2.buffer;

	b32 dedicatedAlloc = dedicatedReqs.prefersDedicatedAllocation || dedicatedReqs.requiresDedicatedAllocation; 

	vk_allocation bufferMem = VkArenaAlignAlloc( vkArena,
												 memReqs2.memoryRequirements.size,
												 memReqs2.memoryRequirements.alignment,
												 memTypeIdx,
												 allocFlags,
												 dedicatedAlloc ? &dedicatedAllocateInfo : 0 );

	buffData.mem = bufferMem.deviceMem;
	buffData.hostVisible = bufferMem.hostVisible;
	buffData.size = sizeInBytes;

	VK_CHECK( vkBindBufferMemory( vkArena->device, buffData.hndl, buffData.mem, bufferMem.dataOffset ) );

	if( allocFlags == VK_MEMORY_ALLOCATE_DEVICE_ADDRESS_BIT )
		buffData.devicePointer = VkGetBufferDeviceAddress( vkArena->device, buffData.hndl );

	return buffData;
}

inline static VkImageView
VkMakeImgView(
	VkDevice	vkDevice,
	VkImage		vkImg,
	VkFormat	imgFormat,
	u32		mipLevel,
	u32		levelCount,
	VkImageViewType imgViewType = VK_IMAGE_VIEW_TYPE_2D,
	u32		arrayLayer = 0,
	u32		layerCount = 1 )
{
	VkImageAspectFlags aspectMask =
		( imgFormat == VK_FORMAT_D32_SFLOAT ) ? VK_IMAGE_ASPECT_DEPTH_BIT : VK_IMAGE_ASPECT_COLOR_BIT;

	VkImageViewCreateInfo viewInfo = {};
	viewInfo.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
	viewInfo.viewType = imgViewType;
	viewInfo.format = imgFormat;
	viewInfo.subresourceRange.aspectMask = aspectMask;
	viewInfo.subresourceRange.baseMipLevel = mipLevel;
	viewInfo.subresourceRange.levelCount = levelCount;
	viewInfo.subresourceRange.baseArrayLayer = arrayLayer;
	viewInfo.subresourceRange.layerCount = layerCount;
	viewInfo.image = vkImg;

	VkImageView view;
	VK_CHECK( vkCreateImageView( vkDevice, &viewInfo, 0, &view ) );

	return view;
}

// TODO: add more VkImageCreateInfo stuff
static image
VkCreateAllocBindImage(
	VkFormat		format,
	VkImageUsageFlags	usageFlags,
	VkExtent3D		extent,
	u32			mipCount,
	vk_mem_arena*		vkArena,
	VkPhysicalDevice	gpu = dc.gpu )
{
	VkFormatFeatureFlags imgType = VK_FORMAT_FEATURE_FLAG_BITS_MAX_ENUM;
	switch( usageFlags ){
	case VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT: imgType = VK_FORMAT_FEATURE_COLOR_ATTACHMENT_BIT;
	break;

	case VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT: imgType = VK_FORMAT_FEATURE_DEPTH_STENCIL_ATTACHMENT_BIT;
	break;

	case VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT: 
	imgType = VK_FORMAT_FEATURE_TRANSFER_DST_BIT | VK_FORMAT_FEATURE_SAMPLED_IMAGE_BIT;
	break;
	}

	VkFormatProperties formatProps;
	vkGetPhysicalDeviceFormatProperties( gpu, format, &formatProps );
	VK_CHECK( !( formatProps.optimalTilingFeatures & imgType ) );


	image img = {};

	VkImageCreateInfo imgInfo = {};
	imgInfo.sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO;
	imgInfo.imageType = VK_IMAGE_TYPE_2D;
	imgInfo.format = img.nativeFormat = format;
	imgInfo.extent = img.nativeRes = extent;
	imgInfo.mipLevels = mipCount;
	imgInfo.arrayLayers = 1;
	imgInfo.samples = VK_SAMPLE_COUNT_1_BIT;
	imgInfo.tiling = VK_IMAGE_TILING_OPTIMAL;
	imgInfo.usage = usageFlags;
	imgInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
	imgInfo.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
	VK_CHECK( vkCreateImage( vkArena->device, &imgInfo, 0, &img.img ) );

	VkImageMemoryRequirementsInfo2 imgReqs2 = {};
	imgReqs2.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_REQUIREMENTS_INFO_2;
	imgReqs2.image = img.img;

	VkMemoryDedicatedRequirements dedicatedReqs = {};
	dedicatedReqs.sType = VK_STRUCTURE_TYPE_MEMORY_DEDICATED_REQUIREMENTS_KHR;

	VkMemoryRequirements2 memReqs2 = {};
	memReqs2.sType = VK_STRUCTURE_TYPE_MEMORY_REQUIREMENTS_2;
	memReqs2.pNext = &dedicatedReqs;

	vkGetImageMemoryRequirements2( vkArena->device, &imgReqs2, &memReqs2 );

	VkMemoryDedicatedAllocateInfo dedicatedAllocateInfo = {};
	dedicatedAllocateInfo.sType = VK_STRUCTURE_TYPE_MEMORY_DEDICATED_ALLOCATE_INFO;
	dedicatedAllocateInfo.image = imgReqs2.image;

	b32 dedicatedAlloc = dedicatedReqs.prefersDedicatedAllocation || dedicatedReqs.requiresDedicatedAllocation;

	i32 memTypeIdx = VkFindMemTypeIdx( &vkArena->gpuMemProps, vkArena->memType, memReqs2.memoryRequirements.memoryTypeBits );
	vk_allocation imgMem = VkArenaAlignAlloc( vkArena,
											  memReqs2.memoryRequirements.size,
											  memReqs2.memoryRequirements.alignment,
											  memTypeIdx, 0,
											  dedicatedAlloc ? &dedicatedAllocateInfo : 0 );

	img.mem = imgMem.deviceMem;
	
	VK_CHECK( vkBindImageMemory( vkArena->device, img.img, img.mem, imgMem.dataOffset ) );

	VkImageViewType viewType = VK_IMAGE_VIEW_TYPE_MAX_ENUM;
	switch( imgInfo.imageType ){
		case VK_IMAGE_TYPE_1D: viewType = VK_IMAGE_VIEW_TYPE_1D; break;
		case VK_IMAGE_TYPE_2D: viewType = VK_IMAGE_VIEW_TYPE_2D; break;
		case VK_IMAGE_TYPE_3D: viewType = VK_IMAGE_VIEW_TYPE_3D; break;
		default: VK_CHECK( "Uknown image type !" ); break;
	}
	img.view = VkMakeImgView( vkArena->device, img.img, imgInfo.format, 0, imgInfo.mipLevels, viewType, 0, imgInfo.arrayLayers );

	return img;
}