Removing VMA. (iree-org#14960)

It's been soaking as off by default for a few weeks with no complaints.
This not only removes a lot of code from runtime binaries but also puts
Vulkan on the same field as our other backends so future improvements
can then be shared across all of them.

.gitmodules

@@ -7,9 +7,6 @@
 [submodule "third_party/vulkan_headers"]
 	path = third_party/vulkan_headers
 	url = https://github.com/KhronosGroup/Vulkan-Headers.git
-[submodule "third_party/vulkan_memory_allocator"]
-	path = third_party/vulkan_memory_allocator
-	url = https://github.com/GPUOpen-LibrariesAndSDKs/VulkanMemoryAllocator.git
 [submodule "third_party/spirv_headers"]
 	path = third_party/spirv_headers
 	url = https://github.com/KhronosGroup/SPIRV-Headers.git

CMakeLists.txt

@@ -812,7 +812,7 @@ else()
   # Add default external projects.
   iree_llvm_add_external_project(mlir-iree-dialects ${CMAKE_CURRENT_SOURCE_DIR}/llvm-external-projects/iree-dialects)
   iree_llvm_add_external_project(stablehlo ${CMAKE_CURRENT_SOURCE_DIR}/third_party/stablehlo)
-  
+
   # Ensure that LLVM-based dependencies needed for testing are included.
   add_dependencies(iree-test-deps FileCheck)
   if(IREE_LLD_TARGET)
@@ -835,7 +835,6 @@ include(flatbuffer_c_library)
 add_subdirectory(build_tools/third_party/libyaml EXCLUDE_FROM_ALL)
 add_subdirectory(build_tools/third_party/llvm-project EXCLUDE_FROM_ALL)
 add_subdirectory(build_tools/third_party/tracy_client EXCLUDE_FROM_ALL)
-add_subdirectory(build_tools/third_party/vulkan_memory_allocator EXCLUDE_FROM_ALL)
 
 iree_set_googletest_cmake_options()
 add_subdirectory(third_party/googletest EXCLUDE_FROM_ALL)

build_tools/bazel/workspace.bzl

@@ -114,13 +114,6 @@ def configure_iree_submodule_deps(iree_repo_alias = "@", iree_path = "./"):
         path = paths.join(iree_path, "third_party/vulkan_headers"),
     )
 
-    maybe(
-        native.new_local_repository,
-        name = "vulkan_memory_allocator",
-        build_file = iree_repo_alias + "//:build_tools/third_party/vulkan_memory_allocator/BUILD.overlay",
-        path = paths.join(iree_path, "third_party/vulkan_memory_allocator"),
-    )
-
     maybe(
         native.local_repository,
         name = "spirv_headers",

build_tools/bazel_to_cmake/bazel_to_cmake_targets.py

@@ -99,9 +99,6 @@ def __init__(self, repo_map: Dict[str, str]):
                 "@com_google_googletest//:gtest": ["gmock", "gtest"],
                 "@spirv_cross//:spirv_cross_lib": ["spirv-cross-msl"],
                 "@cpuinfo": ["${IREE_CPUINFO_TARGET}"],
-                "@vulkan_memory_allocator//:impl_header_only": [
-                    "vulkan_memory_allocator"
-                ],
                 "@webgpu_headers": [],
             }
         )

build_tools/scripts/git/runtime_submodules.txt

@@ -8,5 +8,4 @@ third_party/spirv_cross
 third_party/spirv_headers
 third_party/tracy
 third_party/vulkan_headers
-third_party/vulkan_memory_allocator
 third_party/webgpu-headers

runtime/src/iree/hal/drivers/vulkan/BUILD.bazel

@@ -56,10 +56,6 @@ iree_runtime_cc_library(
         "status_util.h",
         "tracing.cc",
         "tracing.h",
-        "vma_allocator.cc",
-        "vma_allocator.h",
-        "vma_impl.cc",
-        "vma_impl.h",
         "vulkan_device.cc",
         "vulkan_driver.cc",
         "vulkan_headers.h",
@@ -89,7 +85,6 @@ iree_runtime_cc_library(
         "//runtime/src/iree/hal/utils:semaphore_base",
         "//runtime/src/iree/schemas:spirv_executable_def_c_fbs",
         "@vulkan_headers",
-        "@vulkan_memory_allocator//:impl_header_only",
     ],
 )
 

runtime/src/iree/hal/drivers/vulkan/CMakeLists.txt

@@ -57,10 +57,6 @@ iree_cc_library(
     "status_util.h"
     "tracing.cc"
     "tracing.h"
-    "vma_allocator.cc"
-    "vma_allocator.h"
-    "vma_impl.cc"
-    "vma_impl.h"
     "vulkan_device.cc"
     "vulkan_driver.cc"
     "vulkan_headers.h"
@@ -83,7 +79,6 @@ iree_cc_library(
     iree::hal::utils::resource_set
     iree::hal::utils::semaphore_base
     iree::schemas::spirv_executable_def_c_fbs
-    vulkan_memory_allocator
   PUBLIC
 )
 

runtime/src/iree/hal/drivers/vulkan/api.h

@@ -207,24 +207,12 @@ enum iree_hal_vulkan_device_flag_bits_t {
   // IREE execution to run asynchronously with the graphics workloads.
   // See: https://gpuopen.com/learn/concurrent-execution-asynchronous-queues/
   IREE_HAL_VULKAN_DEVICE_FLAG_DEDICATED_COMPUTE_QUEUE = 1u << 0,
-
-  // Whether to use the VMA allocator instead of native Vulkan API memory
-  // allocations.
-  // NOTE: this is temporary and VMA is slated for removal in the future.
-  IREE_HAL_VULKAN_DEVICE_FLAG_VMA_ALLOCATOR = 1u << 1,
 };
 typedef uint32_t iree_hal_vulkan_device_flags_t;
 
 typedef struct iree_hal_vulkan_device_options_t {
   // Flags controlling device behavior.
   iree_hal_vulkan_device_flags_t flags;
-
-  // Sets the VMA preferredLargeHeapBlockSize field to control the preferred
-  // size of a large heap block allocation. This effectively specifies the
-  // minimum amount of memory required and will always allocate at least this
-  // much.
-  // NOTE: this is temporary and VMA is slated for removal in the future.
-  iree_device_size_t large_heap_block_size;
 } iree_hal_vulkan_device_options_t;
 
 IREE_API_EXPORT void iree_hal_vulkan_device_options_initialize(
@@ -324,8 +312,6 @@ IREE_API_EXPORT iree_status_t iree_hal_vulkan_driver_create_using_instance(
 // iree_hal_vulkan_*_buffer_t
 //===----------------------------------------------------------------------===//
 
-// EXPERIMENTAL: until VMA is removed this is doing a shady reinterpret cast.
-//
 // TODO(benvanik): make this safer (dyn_cast-like, lookup allocated buffer).
 // Returns the backing device memory and logical buffer handle of a HAL buffer
 // managed by the Vulkan HAL. Invalid to call on any buffer but a base allocated

runtime/src/iree/hal/drivers/vulkan/descriptor_set_arena.cc

@@ -44,11 +44,11 @@ static void PopulateDescriptorSetWriteInfos(
     if (binding.length == IREE_WHOLE_BUFFER) {
       buffer_info.range = VK_WHOLE_SIZE;
     } else {
-      // Round up to a multiple of 32-bit. 32-bit is the most native bitwidth on
-      // GPUs; it has the best support compared to other bitwidths. We use VMA
-      // to manage GPU memory for us and VMA should already handled proper
-      // alignment when performing allocations; here we just need to provide the
-      // proper "view" to Vulkan drivers over the allocated memory.
+      // Round up to a multiple of 32-bit. 32-bit is the defacto native bitwidth
+      // on GPUs; it has the best support compared to other bitwidths. The
+      // allocator should already handled proper alignment when performing
+      // allocations; here we just need to provide the proper "view" to Vulkan
+      // drivers over the allocated memory.
       //
       // Note this is needed because we can see unusal buffers like
       // tensor<3xi8>. Depending on GPU capabilities, this might not always be

runtime/src/iree/hal/drivers/vulkan/registration/driver_module.cc

@@ -45,13 +45,6 @@ IREE_FLAG(bool, vulkan_buffer_device_addresses, true,
 IREE_FLAG(
     bool, vulkan_dedicated_compute_queue, false,
     "Use a dedicated queue with VK_QUEUE_COMPUTE_BIT for dispatch workloads.");
-IREE_FLAG(bool, vulkan_vma_allocator, false,
-          "Whether to use the VMA allocator instead of native Vulkan API "
-          "memory allocations.");
-IREE_FLAG(
-    int64_t, vulkan_large_heap_block_size, 0,
-    "Preferred allocator block size for large allocations in bytes. Sets the\n"
-    "minimum bound on memory consumption.");
 
 static iree_status_t iree_hal_vulkan_create_driver_with_flags(
     iree_string_view_t identifier, iree_allocator_t host_allocator,
@@ -106,17 +99,6 @@ static iree_status_t iree_hal_vulkan_create_driver_with_flags(
     driver_options.device_options.flags |=
         IREE_HAL_VULKAN_DEVICE_FLAG_DEDICATED_COMPUTE_QUEUE;
   }
-  if (FLAG_vulkan_vma_allocator) {
-    driver_options.device_options.flags |=
-        IREE_HAL_VULKAN_DEVICE_FLAG_VMA_ALLOCATOR;
-  } else {
-    driver_options.device_options.flags &=
-        ~IREE_HAL_VULKAN_DEVICE_FLAG_VMA_ALLOCATOR;
-  }
-  if (FLAG_vulkan_large_heap_block_size) {
-    driver_options.device_options.large_heap_block_size =
-        FLAG_vulkan_large_heap_block_size;
-  }
 
   // Load the Vulkan library. This will fail if the library cannot be found or
   // does not have the expected functions.