[SYCL][E2E][Bindless] Check and use DedicatedAllocation in vulkan interop test for L0 (#14967)

Vulkan interop in L0 backend currently has an issue that image sizes are
not correctly passed to graphics memory management (GMM). Using
dedicated allocation is an approach to solve the issue. In this case,
image and device memory has 1:1 mapping and GMM is aware of the image
info.
This PR enables VK_KHR_GET_MEMORY_REQUIREMENTS_2_EXTENSION_NAME
extension and checks if
VkMemoryDedicatedRequirements.requiresDedicatedAllocation is true. The
value is false by default in CUDA and L0 backends.

Author: wenju-he

sycl/test-e2e/bindless_images/vulkan_interop/mipmaps.cpp

@@ -272,8 +272,8 @@ bool run_test(sycl::range<NDims> dims, sycl::range<NDims> localSize,
   VkMemoryRequirements memRequirements;
   auto inputImageMemoryTypeIndex = vkutil::getImageMemoryTypeIndex(
       inputImage, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, memRequirements);
-  auto inputMemory = vkutil::allocateDeviceMemory(memRequirements.size,
-                                                  inputImageMemoryTypeIndex);
+  auto inputMemory = vkutil::allocateDeviceMemory(
+      memRequirements.size, inputImageMemoryTypeIndex, inputImage);
   VK_CHECK_CALL(vkBindImageMemory(vk_device, inputImage, inputMemory,
                                   0 /*memoryOffset*/));
 
@@ -286,7 +286,8 @@ bool run_test(sycl::range<NDims> dims, sycl::range<NDims> localSize,
       inputStagingBuffer, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT |
                               VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);
   auto inputStagingMemory = vkutil::allocateDeviceMemory(
-      memRequirements.size, inputStagingMemoryTypeIndex, false /*exportable*/);
+      memRequirements.size, inputStagingMemoryTypeIndex, nullptr /*image*/,
+      false /*exportable*/);
   VK_CHECK_CALL(vkBindBufferMemory(vk_device, inputStagingBuffer,
                                    inputStagingMemory, 0 /*memoryOffset*/));
 

sycl/test-e2e/bindless_images/vulkan_interop/sampled_images.cpp

@@ -213,8 +213,8 @@ bool run_test(sycl::range<NDims> dims, sycl::range<NDims> localSize,
   VkMemoryRequirements memRequirements;
   auto inputImageMemoryTypeIndex = vkutil::getImageMemoryTypeIndex(
       inputImage, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, memRequirements);
-  auto inputMemory =
-      vkutil::allocateDeviceMemory(imageSizeBytes, inputImageMemoryTypeIndex);
+  auto inputMemory = vkutil::allocateDeviceMemory(
+      imageSizeBytes, inputImageMemoryTypeIndex, inputImage);
   VK_CHECK_CALL(vkBindImageMemory(vk_device, inputImage, inputMemory,
                                   0 /*memoryOffset*/));
 
@@ -226,8 +226,9 @@ bool run_test(sycl::range<NDims> dims, sycl::range<NDims> localSize,
   auto inputStagingMemoryTypeIndex = vkutil::getBufferMemoryTypeIndex(
       inputStagingBuffer, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT |
                               VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);
-  auto inputStagingMemory = vkutil::allocateDeviceMemory(
-      imageSizeBytes, inputStagingMemoryTypeIndex, false /*exportable*/);
+  auto inputStagingMemory =
+      vkutil::allocateDeviceMemory(imageSizeBytes, inputStagingMemoryTypeIndex,
+                                   nullptr /*image*/, false /*exportable*/);
   VK_CHECK_CALL(vkBindBufferMemory(vk_device, inputStagingBuffer,
                                    inputStagingMemory, 0 /*memoryOffset*/));
 

sycl/test-e2e/bindless_images/vulkan_interop/sampled_images_USM.cpp

@@ -192,8 +192,8 @@ bool run_test(sycl::range<NDims> dims, sycl::range<NDims> localSize,
   memRequirements.size = imageSizeBytes;
   auto inputImageMemoryTypeIndex = vkutil::getImageMemoryTypeIndex(
       inputImage, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, memRequirements);
-  auto inputMemory =
-      vkutil::allocateDeviceMemory(imageSizeBytes, inputImageMemoryTypeIndex);
+  auto inputMemory = vkutil::allocateDeviceMemory(
+      imageSizeBytes, inputImageMemoryTypeIndex, inputImage);
   VK_CHECK_CALL(vkBindImageMemory(vk_device, inputImage, inputMemory,
                                   0 /*memoryOffset*/));
 
@@ -206,7 +206,7 @@ bool run_test(sycl::range<NDims> dims, sycl::range<NDims> localSize,
       inputStagingBuffer, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT |
                               VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);
   auto inputStagingMemory = vkutil::allocateDeviceMemory(
-      stagingBufferSizeBytes, inputStagingMemoryTypeIndex,
+      stagingBufferSizeBytes, inputStagingMemoryTypeIndex, nullptr /*image*/,
       false /*exportable*/);
   VK_CHECK_CALL(vkBindBufferMemory(vk_device, inputStagingBuffer,
                                    inputStagingMemory, 0 /*memoryOffset*/));

sycl/test-e2e/bindless_images/vulkan_interop/vulkan_common.hpp

@@ -50,6 +50,8 @@ static PFN_vkGetMemoryFdKHR vk_getMemoryFdKHR;
 static PFN_vkGetSemaphoreFdKHR vk_getSemaphoreFdKHR;
 #endif
 
+static PFN_vkGetImageMemoryRequirements2 vk_getImageMemoryRequirements2;
+
 static uint32_t vk_computeQueueFamilyIndex;
 static uint32_t vk_transferQueueFamilyIndex;
 
@@ -59,6 +61,8 @@ static VkCommandPool vk_transferCmdPool;
 static VkCommandBuffer vk_computeCmdBuffer;
 static VkCommandBuffer vk_transferCmdBuffers[2];
 
+static bool requiresDedicatedAllocation = false;
+
 // A static debug callback function that relays messages from the Vulkan
 // validation layer to the terminal.
 static VKAPI_ATTR VkBool32 VKAPI_CALL
@@ -220,7 +224,8 @@ VkResult setupDevice(std::string device) {
   bool foundDevice = false;
 
   // Define the required device extensions to run the tests.
-  static constexpr std::string_view requiredExtensions[] = {
+  static constexpr const char *requiredExtensions[] = {
+      VK_KHR_GET_MEMORY_REQUIREMENTS_2_EXTENSION_NAME,
       VK_KHR_EXTERNAL_MEMORY_EXTENSION_NAME,
       VK_KHR_EXTERNAL_SEMAPHORE_EXTENSION_NAME,
 #ifdef _WIN32
@@ -325,28 +330,15 @@ VkResult setupDevice(std::string device) {
 
   VkPhysicalDeviceFeatures deviceFeatures = {};
 
-  // Store our required device extensions. To be passed to the Vulkan device
-  // creation function.
-  std::vector<const char *> extensions = {
-      VK_KHR_EXTERNAL_MEMORY_EXTENSION_NAME,
-      VK_KHR_EXTERNAL_SEMAPHORE_EXTENSION_NAME,
-#ifdef _WIN32
-      VK_KHR_EXTERNAL_MEMORY_WIN32_EXTENSION_NAME,
-      VK_KHR_EXTERNAL_SEMAPHORE_WIN32_EXTENSION_NAME,
-#else
-      VK_KHR_EXTERNAL_SEMAPHORE_FD_EXTENSION_NAME,
-      VK_KHR_EXTERNAL_MEMORY_FD_EXTENSION_NAME,
-#endif
-  };
-
   // Create the Vulkan device with the above queues, extensions, and layers.
   VkDeviceCreateInfo dci = {};
   dci.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;
   dci.pQueueCreateInfos = qcis.data();
   dci.queueCreateInfoCount = qcis.size();
   dci.pEnabledFeatures = &deviceFeatures;
-  dci.enabledExtensionCount = extensions.size();
-  dci.ppEnabledExtensionNames = extensions.data();
+  dci.enabledExtensionCount =
+      sizeof(requiredExtensions) / sizeof(requiredExtensions[0]);
+  dci.ppEnabledExtensionNames = &requiredExtensions[0];
 
   VK_CHECK_CALL_RET(
       vkCreateDevice(vk_physical_device, &dci, nullptr, &vk_device));
@@ -390,6 +382,15 @@ VkResult setupDevice(std::string device) {
   }
 #endif
 
+  vk_getImageMemoryRequirements2 =
+      reinterpret_cast<PFN_vkGetImageMemoryRequirements2>(
+          vkGetDeviceProcAddr(vk_device, "vkGetImageMemoryRequirements2KHR"));
+  if (!vk_getImageMemoryRequirements2) {
+    std::cerr << "Could not get func pointer to "
+                 "\"vkGetImageMemoryRequirements2KHR\"!\n";
+    return VK_ERROR_UNKNOWN;
+  }
+
   return VK_SUCCESS;
 }
 
@@ -533,21 +534,32 @@ exportable, in which case the appropriate extension struct is populated based on
 the OS the program is compiled for.
 */
 VkDeviceMemory allocateDeviceMemory(size_t size, uint32_t memoryTypeIndex,
-                                    bool exportable = true) {
-  VkMemoryAllocateInfo mai = {};
+                                    VkImage image, bool exportable = true) {
+  VkMemoryAllocateInfo mai{};
   mai.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
   mai.allocationSize = size;
   mai.memoryTypeIndex = memoryTypeIndex;
 
-  VkExportMemoryAllocateInfo emai = {};
+  VkMemoryDedicatedAllocateInfoKHR dedicatedInfo{};
+  if (requiresDedicatedAllocation) {
+    dedicatedInfo.sType = VK_STRUCTURE_TYPE_MEMORY_DEDICATED_ALLOCATE_INFO_KHR;
+    dedicatedInfo.image = image;
+    dedicatedInfo.buffer = VK_NULL_HANDLE;
+    mai.pNext = &dedicatedInfo;
+  }
+
+  VkExportMemoryAllocateInfo emai{};
   if (exportable) {
     emai.sType = VK_STRUCTURE_TYPE_EXPORT_MEMORY_ALLOCATE_INFO;
 #ifdef _WIN32
     emai.handleTypes = VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_WIN32_BIT;
 #else
     emai.handleTypes = VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_FD_BIT;
 #endif
-    mai.pNext = &emai;
+    if (requiresDedicatedAllocation)
+      dedicatedInfo.pNext = &emai;
+    else
+      mai.pNext = &emai;
   }
 
   VkDeviceMemory memory;
@@ -565,11 +577,28 @@ property flags passed.
 */
 uint32_t getImageMemoryTypeIndex(VkImage image, VkMemoryPropertyFlags flags,
                                  VkMemoryRequirements &memRequirements) {
-  vkGetImageMemoryRequirements(vk_device, image, &memRequirements);
+  VkMemoryDedicatedRequirements dedicatedRequirements{};
+  dedicatedRequirements.sType = VK_STRUCTURE_TYPE_MEMORY_DEDICATED_REQUIREMENTS;
+
+  VkMemoryRequirements2 memoryRequirements2{};
+  memoryRequirements2.sType = VK_STRUCTURE_TYPE_MEMORY_REQUIREMENTS_2;
+  memoryRequirements2.pNext = &dedicatedRequirements;
+
+  VkImageMemoryRequirementsInfo2 imageRequirementsInfo{};
+  imageRequirementsInfo.sType =
+      VK_STRUCTURE_TYPE_IMAGE_MEMORY_REQUIREMENTS_INFO_2;
+  imageRequirementsInfo.image = image;
+
+  vk_getImageMemoryRequirements2(vk_device, &imageRequirementsInfo,
+                                 &memoryRequirements2);
+
+  if (dedicatedRequirements.requiresDedicatedAllocation)
+    requiresDedicatedAllocation = true;
 
   VkPhysicalDeviceMemoryProperties memProperties;
   vkGetPhysicalDeviceMemoryProperties(vk_physical_device, &memProperties);
 
+  memRequirements = memoryRequirements2.memoryRequirements;
   for (uint32_t i = 0; i < memProperties.memoryTypeCount; i++) {
     if ((memRequirements.memoryTypeBits & (1 << i)) &&
         (memProperties.memoryTypes[i].propertyFlags & flags) == flags) {
@@ -760,8 +789,8 @@ struct vulkan_image_test_resources_t {
     VkMemoryRequirements memRequirements;
     auto inputImageMemoryTypeIndex = vkutil::getImageMemoryTypeIndex(
         vkImage, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, memRequirements);
-    imageMemory =
-        vkutil::allocateDeviceMemory(imageSizeBytes, inputImageMemoryTypeIndex);
+    imageMemory = vkutil::allocateDeviceMemory(
+        imageSizeBytes, inputImageMemoryTypeIndex, vkImage);
     VK_CHECK_CALL(
         vkBindImageMemory(vk_device, vkImage, imageMemory, 0 /*memoryOffset*/));
 
@@ -772,7 +801,8 @@ struct vulkan_image_test_resources_t {
         stagingBuffer, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT |
                            VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);
     stagingMemory = vkutil::allocateDeviceMemory(
-        imageSizeBytes, inputStagingMemoryTypeIndex, false /*exportable*/);
+        imageSizeBytes, inputStagingMemoryTypeIndex, nullptr /*image*/,
+        false /*exportable*/);
     VK_CHECK_CALL(vkBindBufferMemory(vk_device, stagingBuffer, stagingMemory,
                                      0 /*memoryOffset*/));
   }