mirai/src/gpu_resource/allocator.cpp

#include "allocator.h"

#include "resource_types_vulkan.h"
#include "core/logger.h"

#define VMA_IMPLEMENTATION
#include <vk_mem_alloc.h>
#include <vk_mem_alloc.hpp>

namespace mirai {
	result_t<vma::Allocator> create_vma_allocator(const allocator_config& config) {
		// 假设你已经有了 instance, physicalDevice, device 和 dl (调度器)
		auto& d = config.disp;

		// 1. 填充 VMA 需要的函数指针结构体
		vma::VulkanFunctions functions{};
		functions.vkGetInstanceProcAddr = d.vkGetInstanceProcAddr;
		functions.vkGetDeviceProcAddr   = d.vkGetDeviceProcAddr;
		// 还可以填充其他具体函数，但在 Dynamic 模式下，主要给这两个入口，
		// VMA 会尝试自己获取，或者你需要把 d 里面的函数逐个填进去（视 VMA 版本而定）
		// 比较稳妥的做法是把 VMA 需要的核心函数都填上：
		functions.vkGetPhysicalDeviceProperties = d.vkGetPhysicalDeviceProperties;
		functions.vkGetPhysicalDeviceMemoryProperties = d.vkGetPhysicalDeviceMemoryProperties;
		functions.vkAllocateMemory = d.vkAllocateMemory;
		functions.vkFreeMemory = d.vkFreeMemory;
		functions.vkMapMemory = d.vkMapMemory;
		functions.vkUnmapMemory = d.vkUnmapMemory;
		functions.vkFlushMappedMemoryRanges = d.vkFlushMappedMemoryRanges;
		functions.vkInvalidateMappedMemoryRanges = d.vkInvalidateMappedMemoryRanges;
		functions.vkBindBufferMemory = d.vkBindBufferMemory;
		functions.vkBindImageMemory = d.vkBindImageMemory;
		functions.vkGetBufferMemoryRequirements = d.vkGetBufferMemoryRequirements;
		functions.vkGetImageMemoryRequirements = d.vkGetImageMemoryRequirements;
		functions.vkCreateBuffer = d.vkCreateBuffer;
		functions.vkDestroyBuffer = d.vkDestroyBuffer;
		functions.vkCreateImage = d.vkCreateImage;
		functions.vkDestroyImage = d.vkDestroyImage;
		functions.vkCmdCopyBuffer = d.vkCmdCopyBuffer;
		// 如果开启了 bufferDeviceAddress，还需要相关的函数...

		// 2. 创建 Allocator
		vma::AllocatorCreateInfo allocator_info{};
		allocator_info.vulkanApiVersion = VK_API_VERSION_1_3;
		allocator_info.physicalDevice = config.physical_device;
		allocator_info.device = config.device;
		allocator_info.instance = config.instance;
		allocator_info.pVulkanFunctions = &functions; // <--- 关键：告诉 VMA 用动态加载的函数
		allocator_info.preferredLargeHeapBlockSize = config.preferred_large_heap_block_size;
		// allocatorInfo.flags = vma::AllocatorCreateFlagBits::eBufferDeviceAddress; // 如果你要用 BDA

		// 启用功能标志
		vma::AllocatorCreateFlags flags;

		if (config.enable_buffer_device_address) {
			flags |= vma::AllocatorCreateFlagBits::eBufferDeviceAddress;
		}

		if (config.enable_memory_budget) {
			flags |= vma::AllocatorCreateFlagBits::eExtMemoryBudget;
		}
		allocator_info.flags = flags;

		// 创建 RAII 风格的 Allocator
		auto [result, allocator] = vma::createAllocator(allocator_info);
		if (result != vk::Result::eSuccess) {
			return MAKE_ERROR_INFO(error_code::vulkan_init_failed, "VMA分配器创建失败: {}", to_string(result));
		}
		return allocator;
	}

	void vma_allocator::setup(const allocator_config& config) {
		instance_                     = config.instance;
		physical_device_              = config.physical_device;
		device_                       = config.device;
		enable_buffer_device_address_ = config.enable_buffer_device_address;

		auto result = create_vma_allocator(config);
		if (!result) {
			MIRAI_LOG_ERROR("VMA 分配器初始化失败: {}", result.error().full_description());
			throw std::runtime_error("VMA 分配器初始化失败");
		}
		vma_allocator_ = result.value();
	}

	result_t<buffer_allocation> vma_allocator::alloc_buffer(const buffer_create_info& info) {
		buffer_allocation alloc;

		#if MIRAI_DEBUG
		alloc.debug_name = info.debug_name;
		#endif

		vk::BufferCreateInfo buffer_info{};
		buffer_info.size        = info.size;
		buffer_info.usage       = to_vulkan_buffer_usage(info.usage);
		buffer_info.sharingMode = to_vulkan_sharing_mode(info.sharing_mode);

		if (enable_buffer_device_address_ && has_flag(info.usage, buffer_usage::shader_device_address)) {
			buffer_info.usage |= vk::BufferUsageFlagBits::eShaderDeviceAddress;
		}

		vma::AllocationCreateInfo alloc_info{};
		alloc_info.usage = vma::MemoryUsage::eAuto;
		alloc_info.flags = vma::AllocationCreateFlagBits::eHostAccessSequentialWrite;
		alloc_info.flags |= vma::AllocationCreateFlagBits::eMapped;
		alloc_info.preferredFlags = vk::MemoryPropertyFlagBits::eDeviceLocal;

		{
			std::lock_guard lock(mutex_);
			auto [result, pair] = vma_allocator_.createBuffer(buffer_info, alloc_info);
			if (result != vk::Result::eSuccess) {
				return MAKE_ERROR_INFO(error_code::vulkan_allocation_failed,
									   "VMA 分配 Buffer 失败: {}",
									   vk::to_string(result));
			}
			const auto& [allocation, buffer] = pair;
			alloc.buffer = buffer;
			alloc.allocation = allocation;
		}
		auto res_info  = vma_allocator_.getAllocationInfo(alloc.allocation);
		auto mem_flags = vma_allocator_.getAllocationMemoryProperties(alloc.allocation);
		auto is_direct_access = static_cast<bool>(mem_flags & vk::MemoryPropertyFlagBits::eDeviceLocal);
		// 获取映射好的指针（因为加了 eMapped 标志）
		const auto mapped_ptr = vma_allocator_.getAllocationInfo(alloc.allocation).pMappedData;

		alloc.info.size = res_info.size;
		alloc.info.offset = res_info.offset;
		alloc.info.device_memory = res_info.deviceMemory;
		alloc.info.mapped_data = mapped_ptr;
		alloc.info.is_direct_access = is_direct_access;
		alloc.info.memory_type_index = res_info.memoryType;

		return alloc;
	}

	result_t<buffer_allocation> vma_allocator::alloc_staging_buffer(const buffer_create_info& info) {
		buffer_allocation alloc;

		vma::AllocationCreateInfo alloc_info{};
		alloc_info.usage          = vma::MemoryUsage::eAuto;
		alloc_info.flags          = vma::AllocationCreateFlagBits::eHostAccessSequentialWrite;
		alloc_info.flags          |= vma::AllocationCreateFlagBits::eMapped;
		alloc_info.preferredFlags = vk::MemoryPropertyFlagBits::eHostVisible;

		vk::BufferCreateInfo buffer_create_info{};
		buffer_create_info.setSize(info.size);
		buffer_create_info.setUsage(vk::BufferUsageFlagBits::eTransferSrc);

		{
			auto [result, pair] = vma_allocator_.createBuffer(buffer_create_info, alloc_info);
			if (result != vk::Result::eSuccess) {
				return MAKE_ERROR_INFO(error_code::vulkan_allocation_failed,
									   "VMA 分配暂存 Buffer 失败: {}",
									   vk::to_string(result));
			}
			const auto& [allocation, buffer] = pair;
			alloc.buffer = buffer;
			alloc.allocation = allocation;
		}

		auto res_info = vma_allocator_.getAllocationInfo(alloc.allocation);
		auto mem_flags = vma_allocator_.getAllocationMemoryProperties(alloc.allocation);
		auto is_direct_access = static_cast<bool>(mem_flags & vk::MemoryPropertyFlagBits::eHostVisible);
		// 获取映射好的指针（因为加了 eMapped 标志）
		const auto mapped_ptr = vma_allocator_.getAllocationInfo(alloc.allocation).pMappedData;
		alloc.info.size = res_info.size;
		alloc.info.offset = res_info.offset;
		alloc.info.device_memory = res_info.deviceMemory;
		alloc.info.mapped_data = mapped_ptr;
		alloc.info.is_direct_access = is_direct_access;
		alloc.info.memory_type_index = res_info.memoryType;

		return alloc;
	}

	void_result_t vma_allocator::free_buffer(const buffer_allocation& buffer) {
		{
			std::lock_guard lock(mutex_);
			vma_allocator_.destroyBuffer(buffer.buffer, buffer.allocation);
		}
		return {};
	}

	void_result_t vma_allocator::free_staging_buffer(const buffer_allocation& buffer) {
		{
			std::lock_guard lock(mutex_);
			vma_allocator_.destroyBuffer(buffer.buffer, buffer.allocation);
		}
		return {};
	}

	void vma_allocator::on_destroying() {
		object::on_destroying();
		MIRAI_LOG_INFO("VMA 分配器正在销毁");
		if (vma_allocator_) {
			vma_allocator_.destroy();
			vma_allocator_ = nullptr;
		}
	}
}