Update VulkanMemoryAllocator to 3528e2a. PiperOrigin-RevId: 337963255
diff --git a/SUBMODULE_VERSIONS b/SUBMODULE_VERSIONS index e9e0de9..2f33cb7 100644 --- a/SUBMODULE_VERSIONS +++ b/SUBMODULE_VERSIONS
@@ -17,4 +17,4 @@ 2d6bdab3adb0b8949763d5c63426338f938c9efe third_party/tensorflow a9a09ab0940408898fccfdcfe2bb8dc19b50f13c third_party/tracy 9bd3f561bcee3f01d22912de10bb07ce4e23d378 third_party/vulkan_headers -909f36b714c9239ee0b112a321220213a474ba53 third_party/vulkan_memory_allocator +3528e2aed3e8808f33e1e7d63eeb1560456a605a third_party/vulkan_memory_allocator
diff --git a/iree/hal/vulkan/internal_vk_mem_alloc.cc b/iree/hal/vulkan/internal_vk_mem_alloc.cc index d61f108..325c522 100644 --- a/iree/hal/vulkan/internal_vk_mem_alloc.cc +++ b/iree/hal/vulkan/internal_vk_mem_alloc.cc
@@ -25,8 +25,8 @@ #include "absl/synchronization/mutex.h" #include "iree/base/logging.h" -// Use std::vector instead of the VMA version. -#define VMA_USE_STL_VECTOR 1 +// Uncomment to try using std::vector instead of the VMA version. +// #define VMA_USE_STL_VECTOR 1 // TODO(benvanik): figure out why std::list cannot be used. // #define VMA_USE_STL_LIST 1 @@ -54,14 +54,22 @@ public: void LockRead() ABSL_SHARED_LOCK_FUNCTION() { mutex_.ReaderLock(); } void UnlockRead() ABSL_UNLOCK_FUNCTION() { mutex_.ReaderUnlock(); } + bool TryLockRead() ABSL_SHARED_TRYLOCK_FUNCTION(true) { + return mutex_.ReaderTryLock(); + } void LockWrite() ABSL_EXCLUSIVE_LOCK_FUNCTION() { mutex_.WriterLock(); } void UnlockWrite() ABSL_UNLOCK_FUNCTION() { mutex_.WriterUnlock(); } + bool TryLockWrite() ABSL_EXCLUSIVE_TRYLOCK_FUNCTION(true) { + return mutex_.WriterTryLock(); + } private: absl::Mutex mutex_; }; #define VMA_RW_MUTEX AbslVmaRWMutex +#define VMA_DYNAMIC_VULKAN_FUNCTIONS 0 + #define VMA_IMPLEMENTATION #include "vk_mem_alloc.h"
diff --git a/iree/hal/vulkan/vma_allocator.cc b/iree/hal/vulkan/vma_allocator.cc index dd18cc2..f1f9225 100644 --- a/iree/hal/vulkan/vma_allocator.cc +++ b/iree/hal/vulkan/vma_allocator.cc
@@ -37,7 +37,7 @@ // static StatusOr<std::unique_ptr<VmaAllocator>> VmaAllocator::Create( VkPhysicalDevice physical_device, - const ref_ptr<VkDeviceHandle>& logical_device) { + const ref_ptr<VkDeviceHandle>& logical_device, VkInstance instance) { IREE_TRACE_SCOPE0("VmaAllocator::Create"); const auto& syms = logical_device->syms(); @@ -79,6 +79,7 @@ create_info.flags = 0; create_info.physicalDevice = physical_device; create_info.device = *logical_device; + create_info.instance = instance; create_info.preferredLargeHeapBlockSize = 64 * 1024 * 1024; create_info.pAllocationCallbacks = logical_device->allocator(); create_info.pDeviceMemoryCallbacks = nullptr;
diff --git a/iree/hal/vulkan/vma_allocator.h b/iree/hal/vulkan/vma_allocator.h index 2c0054c..fde9b98 100644 --- a/iree/hal/vulkan/vma_allocator.h +++ b/iree/hal/vulkan/vma_allocator.h
@@ -49,7 +49,7 @@ public: static StatusOr<std::unique_ptr<VmaAllocator>> Create( VkPhysicalDevice physical_device, - const ref_ptr<VkDeviceHandle>& logical_device); + const ref_ptr<VkDeviceHandle>& logical_device, VkInstance instance); ~VmaAllocator() override;
diff --git a/iree/hal/vulkan/vulkan_device.cc b/iree/hal/vulkan/vulkan_device.cc index 4df227e..05fa4ab 100644 --- a/iree/hal/vulkan/vulkan_device.cc +++ b/iree/hal/vulkan/vulkan_device.cc
@@ -346,8 +346,9 @@ // Create the device memory allocator. // TODO(benvanik): allow other types to be plugged in. - IREE_ASSIGN_OR_RETURN(auto allocator, - VmaAllocator::Create(physical_device, logical_device)); + IREE_ASSIGN_OR_RETURN( + auto allocator, + VmaAllocator::Create(physical_device, logical_device, instance)); // Create command pools for each queue family. If we don't have a transfer // queue then we'll ignore that one and just use the dispatch pool. @@ -407,7 +408,7 @@ // static StatusOr<ref_ptr<VulkanDevice>> VulkanDevice::Wrap( - ref_ptr<Driver> driver, const DeviceInfo& device_info, + ref_ptr<Driver> driver, VkInstance instance, const DeviceInfo& device_info, VkPhysicalDevice physical_device, VkDevice logical_device, const ExtensibilitySpec& extensibility_spec, const QueueSet& compute_queue_set, const QueueSet& transfer_queue_set, @@ -442,8 +443,9 @@ // Create the device memory allocator. // TODO(benvanik): allow other types to be plugged in. - IREE_ASSIGN_OR_RETURN(auto allocator, - VmaAllocator::Create(physical_device, device_handle)); + IREE_ASSIGN_OR_RETURN( + auto allocator, + VmaAllocator::Create(physical_device, device_handle, instance)); bool has_dedicated_transfer_queues = transfer_queue_count > 0;
diff --git a/iree/hal/vulkan/vulkan_device.h b/iree/hal/vulkan/vulkan_device.h index 273d407..d7b31bb 100644 --- a/iree/hal/vulkan/vulkan_device.h +++ b/iree/hal/vulkan/vulkan_device.h
@@ -60,9 +60,9 @@ // Creates a device that wraps an externally managed VkDevice. static StatusOr<ref_ptr<VulkanDevice>> Wrap( - ref_ptr<Driver> driver, const DeviceInfo& device_info, - VkPhysicalDevice physical_device, VkDevice logical_device, - const ExtensibilitySpec& extensibility_spec, + ref_ptr<Driver> driver, VkInstance instance, + const DeviceInfo& device_info, VkPhysicalDevice physical_device, + VkDevice logical_device, const ExtensibilitySpec& extensibility_spec, const QueueSet& compute_queue_set, const QueueSet& transfer_queue_set, const ref_ptr<DynamicSymbols>& syms);
diff --git a/iree/hal/vulkan/vulkan_driver.cc b/iree/hal/vulkan/vulkan_driver.cc index 6941bb6..576145e 100644 --- a/iree/hal/vulkan/vulkan_driver.cc +++ b/iree/hal/vulkan/vulkan_driver.cc
@@ -311,10 +311,10 @@ // Attempt to create the device. // This may fail if the VkDevice does not support all necessary features. IREE_ASSIGN_OR_RETURN( - auto device, - VulkanDevice::Wrap(add_ref(this), device_info, physical_device, - logical_device, device_extensibility_spec_, - compute_queue_set, transfer_queue_set, syms())); + auto device, VulkanDevice::Wrap( + add_ref(this), instance(), device_info, physical_device, + logical_device, device_extensibility_spec_, + compute_queue_set, transfer_queue_set, syms())); return device; }