DO NOT SUBMIT vulkan logging for vma issue #14776

build_tools/scripts/check_vulkan.sh

@@ -24,3 +24,6 @@ fi
 
 echo "${VULKAN_INSTANCE?}"
 echo "${VK_PHYSICAL_DEVICE_PROPERTIES?}"
+
+cat /tmp/vulkaninfo.stdout
+cat /tmp/vulkaninfo.stderr

runtime/src/iree/hal/drivers/vulkan/base_buffer.c

@@ -45,12 +45,15 @@ iree_status_t iree_hal_vulkan_find_memory_type(
 
   iree_hal_memory_type_t requested_type = params->type;
   if (device_props->deviceType == VK_PHYSICAL_DEVICE_TYPE_INTEGRATED_GPU) {
+    fprintf(stderr, "iree_hal_vulkan_find_memory_type integrated gpu\n");
     // Integrated GPUs have tiny device local heaps commonly used for
     // framebuffers and other bounded resources. We don't currently try to use
     // them but could for very small transients.
     if (iree_all_bits_set(requested_type, IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL)) {
       requested_type &= ~IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL;
       requested_type |= IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE;
+      fprintf(stderr,
+              "iree_hal_vulkan_find_memory_type flip to device visible\n");
     }
   }
 
@@ -61,30 +64,40 @@ iree_status_t iree_hal_vulkan_find_memory_type(
       // Device-local, host-visible.
       require_flags |= VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT;
       prefer_flags |= VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT;
+      fprintf(stderr,
+              "iree_hal_vulkan_find_memory_type req host visible, pref device "
+              "local\n");
     } else {
       // Device-local only.
       require_flags |= VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT;
+      fprintf(stderr, "iree_hal_vulkan_find_memory_type device local only\n");
     }
   } else {
     if (iree_all_bits_set(requested_type,
                           IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE)) {
       // Host-local, device-visible.
       require_flags |= VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT;
+      fprintf(stderr,
+              "iree_hal_vulkan_find_memory_type host local device visible\n");
     } else {
       // Host-local only.
       require_flags |= VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT;
+      fprintf(stderr, "iree_hal_vulkan_find_memory_type host visible only\n");
     }
   }
   if (iree_all_bits_set(requested_type, IREE_HAL_MEMORY_TYPE_HOST_CACHED)) {
     require_flags |= VK_MEMORY_PROPERTY_HOST_CACHED_BIT;
+    fprintf(stderr, "iree_hal_vulkan_find_memory_type host cached\n");
   }
   if (iree_all_bits_set(requested_type, IREE_HAL_MEMORY_TYPE_HOST_COHERENT)) {
     require_flags |= VK_MEMORY_PROPERTY_HOST_COHERENT_BIT;
+    fprintf(stderr, "iree_hal_vulkan_find_memory_type host coherent\n");
   }
   if (iree_any_bit_set(requested_type,
                        IREE_HAL_BUFFER_USAGE_MAPPING_SCOPED |
                            IREE_HAL_BUFFER_USAGE_MAPPING_PERSISTENT)) {
     require_flags |= VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT;
+    fprintf(stderr, "iree_hal_vulkan_find_memory_type mapping req\n");
   }
 
   int most_bits_count = 0;
@@ -95,6 +108,7 @@ iree_status_t iree_hal_vulkan_find_memory_type(
         !iree_hal_vulkan_is_memory_type_usable(flags) ||
         !iree_all_bits_set(allowed_type_indices, 1u << i)) {
       // Excluded (required bits missing or memory type is not usable).
+      fprintf(stderr, "iree_hal_vulkan_find_memory_type excluding bit %u\n", i);
       continue;
     }
     // When all required bits are satisfied try to find the memory type that
@@ -103,9 +117,13 @@ iree_status_t iree_hal_vulkan_find_memory_type(
     if (most_bits_idx == -1) {
       most_bits_count = bit_count;
       most_bits_idx = (int)i;
+      fprintf(stderr,
+              "iree_hal_vulkan_find_memory_type first bit found at %u\n", i);
     } else if (bit_count > most_bits_count) {
       most_bits_count = bit_count;
       most_bits_idx = (int)i;
+      fprintf(stderr, "iree_hal_vulkan_find_memory_type better found at %u\n",
+              i);
     }
   }
   if (most_bits_idx == -1) {

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

@@ -268,6 +268,8 @@ static iree_status_t iree_hal_vulkan_native_allocator_commit_and_wrap(
   VkMemoryRequirements requirements = {0};
   logical_device->syms()->vkGetBufferMemoryRequirements(*logical_device, handle,
                                                         &requirements);
+  fprintf(stderr, "vkGetBufferMemoryRequirements alignment %u bits %08X\n",
+          (uint32_t)requirements.alignment, requirements.memoryTypeBits);
   uint32_t memory_type_index = 0;
   IREE_RETURN_IF_ERROR(iree_hal_vulkan_find_memory_type(
       &allocator->device_props, &allocator->memory_props, params,
@@ -287,13 +289,23 @@ static iree_status_t iree_hal_vulkan_native_allocator_commit_and_wrap(
         allocator->device_props_11.maxMemoryAllocationSize, out_buffer);
   }
 
+  fprintf(stderr, "vkAllocateMemory\nreq size %u\nmemory type index %u\n",
+          (uint32_t)requirements.size, memory_type_index);
+
+
   // Allocate the device memory we'll attach the buffer to.
   VkMemoryAllocateInfo allocate_info = {};
   allocate_info.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
   allocate_info.pNext = NULL;
   allocate_info.allocationSize = requirements.size;
   allocate_info.memoryTypeIndex = memory_type_index;
   VkDeviceMemory device_memory = VK_NULL_HANDLE;
+  fprintf(stderr,
+          "vkAllocateMemory(%p, allocationSize=%" PRIu64
+          ", memoryTypeIndex=%u, pNext=%p, %p, %p)\n",
+          (void*)logical_device->value(),
+          (uint64_t)allocate_info.allocationSize, allocate_info.memoryTypeIndex,
+          allocate_info.pNext, logical_device->allocator(), &device_memory);
   VK_RETURN_IF_ERROR(logical_device->syms()->vkAllocateMemory(
                          *logical_device, &allocate_info,
                          logical_device->allocator(), &device_memory),

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

@@ -261,6 +261,19 @@ static void VKAPI_PTR iree_hal_vulkan_vma_free_callback(
 static void iree_hal_vulkan_vma_allocator_destroy(
     iree_hal_allocator_t* IREE_RESTRICT base_allocator);
 
+static PFN_vkAllocateMemory allocate_memory_ptr = NULL;
+static VkResult HOOKED_vkAllocateMemory(
+    VkDevice device, const VkMemoryAllocateInfo* pAllocateInfo,
+    const VkAllocationCallbacks* pAllocator, VkDeviceMemory* pMemory) {
+  fprintf(stderr,
+          "vkAllocateMemory(%p, allocationSize=%" PRIu64
+          ", memoryTypeIndex=%u, pNext=%p, %p, %p)\n",
+          (void*)device, (uint64_t)pAllocateInfo->allocationSize,
+          pAllocateInfo->memoryTypeIndex, pAllocateInfo->pNext, pAllocator,
+          pMemory);
+  return allocate_memory_ptr(device, pAllocateInfo, pAllocator, pMemory);
+}
+
 iree_status_t iree_hal_vulkan_vma_allocator_create(
     const iree_hal_vulkan_device_options_t* options, VkInstance instance,
     VkPhysicalDevice physical_device, VkDeviceHandle* logical_device,
@@ -287,7 +300,8 @@ iree_status_t iree_hal_vulkan_vma_allocator_create(
       syms->vkGetPhysicalDeviceProperties;
   vulkan_fns.vkGetPhysicalDeviceMemoryProperties =
       syms->vkGetPhysicalDeviceMemoryProperties;
-  vulkan_fns.vkAllocateMemory = syms->vkAllocateMemory;
+  allocate_memory_ptr = syms->vkAllocateMemory;
+  vulkan_fns.vkAllocateMemory = HOOKED_vkAllocateMemory;
   vulkan_fns.vkFreeMemory = syms->vkFreeMemory;
   vulkan_fns.vkMapMemory = syms->vkMapMemory;
   vulkan_fns.vkUnmapMemory = syms->vkUnmapMemory;

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

@@ -18,6 +18,12 @@
 // to be omitted and not have VMA poking around where it shouldn't.
 #define VMA_DYNAMIC_VULKAN_FUNCTIONS 0
 
+#define VMA_DEBUG_LOG_FORMAT(format, ...)   \
+  do {                                      \
+    fprintf(stderr, (format), __VA_ARGS__); \
+    fprintf(stderr, "\n");                  \
+  } while (false)
+
 #include <vk_mem_alloc.h>  // IWYU pragma: export
 
 #endif  // IREE_HAL_DRIVERS_VULKAN_VMA_IMPL_H_

tests/e2e/models/mnist_train_test/CMakeLists.txt

@@ -39,18 +39,28 @@ if(IREE_TARGET_BACKEND_CUDA AND IREE_HAL_DRIVER_CUDA)
   )
 endif()
 
-# TODO(#14776): enable Vulkan when it passes on the bots - currently fails with
-# an unexpected out-of-memory.
-# if(IREE_TARGET_BACKEND_VULKAN_SPIRV AND IREE_HAL_DRIVER_VULKAN)
-#   iree_py_test(
-#     NAME
-#       mnist_train_test_vulkan
-#     SRCS
-#       "mnist_train_test.py"
-#     ARGS
-#       "--target_backend=vulkan-spirv"
-#       "--driver=vulkan"
-#     LABELS
-#       "driver=vulkan"
-#   )
-# endif()
+if(IREE_TARGET_BACKEND_VULKAN_SPIRV AND IREE_HAL_DRIVER_VULKAN)
+  iree_py_test(
+    NAME
+      mnist_train_test_vulkan_vma
+    SRCS
+      "mnist_train_test.py"
+    ARGS
+      "--target_backend=vulkan-spirv"
+      "--driver=vulkan"
+      "--vma"
+    LABELS
+      "driver=vulkan"
+  )
+  iree_py_test(
+    NAME
+      mnist_train_test_vulkan
+    SRCS
+      "mnist_train_test.py"
+    ARGS
+      "--target_backend=vulkan-spirv"
+      "--driver=vulkan"
+    LABELS
+      "driver=vulkan"
+  )
+endif()

tests/e2e/models/mnist_train_test/mnist_train_test.py

@@ -17,6 +17,7 @@
 import numpy as np
 
 from iree.compiler.tools import InputType, compile_file
+from iree import runtime as rt
 from iree.runtime import load_vm_flatbuffer_file
 
 MODEL_ARTIFACTS_URL = "https://storage.googleapis.com/iree-model-artifacts/mnist_train.a49ba1535a45ac0f3e6be22a7ed5dddf4a53cd1f41126af938f0667b998f8e11.tar"
@@ -76,6 +77,7 @@ def parse_args():
     parser = argparse.ArgumentParser()
     parser.add_argument("--target_backend", type=str, default="llvm-cpu")
     parser.add_argument("--driver", type=str, default="local-task")
+    parser.add_argument("--vma", default=False, action="store_true")
     return parser.parse_known_args()
 
 
@@ -115,6 +117,13 @@ def extract_test_data(archive_path: str, out_dir: str):
 
 class MnistTrainTest(unittest.TestCase):
     def test_mnist_training(self):
+        if args.vma:
+          rt.flags.parse_flags("--vulkan_vma_allocator=true")
+        else:
+          rt.flags.parse_flags("--vulkan_vma_allocator=false")
+        rt.flags.parse_flags("--vulkan_validation_layers=true")
+        rt.flags.parse_flags("--vulkan_debug_utils=true")
+        rt.flags.parse_flags("--vulkan_debug_verbosity=4")
         with tempfile.TemporaryDirectory() as tmp_dir:
             archive_path = os.path.join(tmp_dir, "mnist_train.tar")
             download_test_data(archive_path)

tests/e2e/tensor_ops/tensor_cast.mlir

@@ -1,7 +1,7 @@
 // RUN: iree-run-mlir --Xcompiler,iree-hal-target-backends=llvm-cpu %s | FileCheck %s
 // RUN: [[ $IREE_VMVX_DISABLE == 1 ]] || (iree-run-mlir --Xcompiler,iree-hal-target-backends=vmvx %s | FileCheck %s)
-// TODO(#14776): enable Vulkan when it passes on the bots - currently fails with an unexpected out-of-memory.
-// XRUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir --Xcompiler,iree-hal-target-backends=vulkan-spirv %s | FileCheck %s)
+// RUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir --Xcompiler,iree-hal-target-backends=vulkan-spirv --trace_execution --vulkan_debug_verbosity=4 --vulkan_debug_utils=true --vulkan_validation_layers=true --vulkan_vma_allocator=true %s | FileCheck %s)
+// RUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir --Xcompiler,iree-hal-target-backends=vulkan-spirv --trace_execution --vulkan_debug_verbosity=4 --vulkan_debug_utils=true --vulkan_validation_layers=true %s | FileCheck %s)
 
 func.func @tensor_cast() -> tensor<2x?xf32> {
   %input = util.unfoldable_constant dense<[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]> : tensor<2x3xf32>