Merge pull request #3553 from rsuderman:main-to-google
PiperOrigin-RevId: 338166297
diff --git a/SUBMODULE_VERSIONS b/SUBMODULE_VERSIONS
index e9e0de9..d3266bd 100644
--- a/SUBMODULE_VERSIONS
+++ b/SUBMODULE_VERSIONS
@@ -14,7 +14,7 @@
685f86471e9d26b3eb7676695a2e2cefb4551ae9 third_party/spirv_cross
f8bf11a0253a32375c32cad92c841237b96696c0 third_party/spirv_headers
57eb48aed36160c4876bc8310d9ca84d42ee9e2a third_party/swiftshader
-2d6bdab3adb0b8949763d5c63426338f938c9efe third_party/tensorflow
+1914c410b274055078cfb97c1f8db5d98c5f9146 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/build_tools/docker/bazel/Dockerfile b/build_tools/docker/bazel/Dockerfile
index a012e71..a7cda8d 100644
--- a/build_tools/docker/bazel/Dockerfile
+++ b/build_tools/docker/bazel/Dockerfile
@@ -52,9 +52,11 @@
/usr/bin/bazel-real \
/usr/bin/
-# TF requires python2 numpy at configure time...
+# TF requires python and numpy at configure time...
# TODO(#1737): Remove this
RUN apt-get update \
- && apt-get install -y python-pip \
- && python -m pip install --upgrade pip \
- && python -m pip install numpy
+ && apt-get install -y \
+ python3 \
+ python3-pip \
+ && python3 -m pip install --upgrade pip \
+ && python3 -m pip install numpy
diff --git a/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86-swiftshader/bindings/build_kokoro.sh b/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86-swiftshader/bindings/build_kokoro.sh
index 97d5667..199a9a2 100755
--- a/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86-swiftshader/bindings/build_kokoro.sh
+++ b/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86-swiftshader/bindings/build_kokoro.sh
@@ -32,7 +32,7 @@
docker_setup
docker run "${DOCKER_RUN_ARGS[@]?}" \
- gcr.io/iree-oss/bazel-python@sha256:1e46da0e373a9e6636c4252d32b12ce068d2fe81a5bce992b4220e91d1f76259 \
+ gcr.io/iree-oss/bazel-python@sha256:3dc3d72717edec56a7ce8a910a9333aa10bd7c9814858880325968047b3a28d4 \
build_tools/kokoro/gcp_ubuntu/bazel/linux/x86-swiftshader/bindings/build.sh
# Kokoro will rsync this entire directory back to the executor orchestrating the
diff --git a/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86-swiftshader/core/build_kokoro.sh b/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86-swiftshader/core/build_kokoro.sh
index b569c42..e78cd07 100755
--- a/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86-swiftshader/core/build_kokoro.sh
+++ b/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86-swiftshader/core/build_kokoro.sh
@@ -32,7 +32,7 @@
docker_setup
docker run "${DOCKER_RUN_ARGS[@]?}" \
- gcr.io/iree-oss/bazel@sha256:81dc9675efb57204a806b83cd988f6a07a1ca8a778119014ae7057ee53eb68d7 \
+ gcr.io/iree-oss/bazel@sha256:6b83817206384c7e8ad4f522162ee31b927ac9912095dc3ef1d3c8f580feba92 \
build_tools/kokoro/gcp_ubuntu/bazel/linux/x86-swiftshader/core/build.sh
# Kokoro will rsync this entire directory back to the executor orchestrating the
diff --git a/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86-swiftshader/integrations/build_kokoro.sh b/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86-swiftshader/integrations/build_kokoro.sh
index c55206d..9e9d18c 100755
--- a/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86-swiftshader/integrations/build_kokoro.sh
+++ b/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86-swiftshader/integrations/build_kokoro.sh
@@ -32,7 +32,7 @@
docker_setup
docker run "${DOCKER_RUN_ARGS[@]?}" \
- gcr.io/iree-oss/bazel-tensorflow-swiftshader@sha256:269a3a42ca71fa6b040135056f1f5639c2f2c7d651099a88ddd513a011365144 \
+ gcr.io/iree-oss/bazel-tensorflow-swiftshader@sha256:684d38c5fbb4a362476305138d4a78fab8710daa0794afe0f2e41f6b61627fe1 \
build_tools/kokoro/gcp_ubuntu/bazel/linux/x86-swiftshader/integrations/build.sh
# Kokoro will rsync this entire directory back to the executor orchestrating the
diff --git a/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86-turing/integrations/build_kokoro.sh b/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86-turing/integrations/build_kokoro.sh
index 70675bb..5ba3605 100755
--- a/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86-turing/integrations/build_kokoro.sh
+++ b/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86-turing/integrations/build_kokoro.sh
@@ -30,6 +30,10 @@
docker_setup
+# The image hash is pinned to an older version than that reflected by the Docker
+# hierarchy because updating nvidia/vulkan causes everything to fail for some
+# reason.
+# TODO(#3550): Allow this to follow the checked-in Docker hierarchy.
docker run "${DOCKER_RUN_ARGS[@]?}" \
--gpus all \
gcr.io/iree-oss/bazel-tensorflow-nvidia@sha256:754dc09c558157f82e9d53451486951fc096e8d2a2b9a1306a29ebfe9e0772df \
diff --git a/build_tools/kokoro/gcp_ubuntu/check_vulkan.sh b/build_tools/kokoro/gcp_ubuntu/check_vulkan.sh
index 7cfa75c..b91091e 100755
--- a/build_tools/kokoro/gcp_ubuntu/check_vulkan.sh
+++ b/build_tools/kokoro/gcp_ubuntu/check_vulkan.sh
@@ -18,7 +18,6 @@
# Log some information about the environment, initialize the submodules and then
# run the bazel integrations tests.
-set -e
set -x
# Print Vulkan related information: SDK version and GPU ICD version
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 8d9cd9a..1a2daa1 100644
--- a/iree/hal/vulkan/vulkan_driver.cc
+++ b/iree/hal/vulkan/vulkan_driver.cc
@@ -316,10 +316,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;
}
diff --git a/third_party/tensorflow b/third_party/tensorflow
index 2d6bdab..1914c41 160000
--- a/third_party/tensorflow
+++ b/third_party/tensorflow
@@ -1 +1 @@
-Subproject commit 2d6bdab3adb0b8949763d5c63426338f938c9efe
+Subproject commit 1914c410b274055078cfb97c1f8db5d98c5f9146
diff --git a/third_party/vulkan_memory_allocator b/third_party/vulkan_memory_allocator
index 909f36b..3528e2a 160000
--- a/third_party/vulkan_memory_allocator
+++ b/third_party/vulkan_memory_allocator
@@ -1 +1 @@
-Subproject commit 909f36b714c9239ee0b112a321220213a474ba53
+Subproject commit 3528e2aed3e8808f33e1e7d63eeb1560456a605a
