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;
}
