Vulkan: Sampler_ycbcr_conversion step 2

Add helper function: getAspectKeysFromAspectFlags(), given an aspect
mask flag, returns a list of aspect bit keys, which can be used in
indexing image->aspect->layer->level hierarchy and also the potential
separated memory binding info hierarchy.

Index the memory requirements by aspect(plane) bits

Add helper function: getImagePlaneMemoryRequirments() to manage getting
image plane memory requirements. All uses of image requirements should
get the requirement info from this helper function.

gapii/cc/vulkan_extras.cpp

@@ -584,8 +584,11 @@ VulkanSpy::fetchPhysicalDeviceFormatProperties(
 gapil::Ref<ImageMemoryRequirements> VulkanSpy::fetchImageMemoryRequirements(
     CallObserver* observer, VkDevice device, VkImage image, bool hasSparseBit) {
   auto reqs = gapil::Ref<ImageMemoryRequirements>::create(arena());
+  VkMemoryRequirements rawReq{0};
   mImports.mVkDeviceFunctions[device].vkGetImageMemoryRequirements(
-      device, image, &reqs->mMemoryRequirements);
+      device, image, &rawReq);
+  // TODO: Handle multi-planar images
+  reqs->mPlaneBitsToMemoryRequirements[0] = rawReq;
   if (hasSparseBit) {
     uint32_t sparse_mem_req_count = 0;
     mImports.mVkDeviceFunctions[device].vkGetImageSparseMemoryRequirements(

gapii/cc/vulkan_mid_execution.cpp

@@ -630,7 +630,8 @@ void VulkanSpy::serializeGPUBuffers(StateSerializer *serializer) {
         bool is_valid = true;
         // If this is a sparsely resident image, then at least ALL metadata
         // must be bound.
-        for (const auto &requirements : img->mSparseMemoryRequirements) {
+        for (const auto &requirements :
+             img->mMemoryRequirements->mAspectBitsToSparseMemoryRequirements) {
           const auto &prop = requirements.second.mformatProperties;
           if (prop.maspectMask ==
               VkImageAspectFlagBits::VK_IMAGE_ASPECT_METADATA_BIT) {
@@ -648,8 +649,12 @@ void VulkanSpy::serializeGPUBuffers(StateSerializer *serializer) {
       } else {
         // If we are not sparsely-resident, then all memory must
         // be bound before we are used.
-        if (!IsFullyBound(0, img->mMemoryRequirements.msize,
-                          img->mOpaqueSparseMemoryBindings)) {
+        // TODO: Handle multi-planar images
+        if (!IsFullyBound(
+                0,
+                img->mMemoryRequirements->mPlaneBitsToMemoryRequirements[0]
+                    .msize,
+                img->mOpaqueSparseMemoryBindings)) {
           continue;
         }
       }
@@ -674,7 +679,8 @@ void VulkanSpy::serializeGPUBuffers(StateSerializer *serializer) {
     if (denseBound || !sparseResidency) {
       walkImageSubRng(img, img_whole_rng, append_image_level_to_opaque_pieces);
     } else {
-      for (const auto &req : img->mSparseMemoryRequirements) {
+      for (const auto &req :
+           img->mMemoryRequirements->mAspectBitsToSparseMemoryRequirements) {
         const auto &prop = req.second.mformatProperties;
         if (prop.maspectMask == img->mImageAspect) {
           if (prop.mflags & VkSparseImageFormatFlagBits::

gapis/api/vulkan/api/coherent_memory.api

@@ -222,7 +222,7 @@ sub void accessImageSubresourceSlice(ref!ImageObject image, VkImageSubresourceRa
   blockWidth := as!u64(elementAndTexelBlockSize.TexelBlockSize.Width)
   blockHeight := as!u64(elementAndTexelBlockSize.TexelBlockSize.Height)
 
-  for _ , _ , aspectBit in unpackImageAspectFlags(rng.aspectMask) {
+  for _ , _ , aspectBit in getAspectKeysWithAspectFlags(image, rng.aspectMask) {
     elementSize := switch (aspectBit) {
       case VK_IMAGE_ASPECT_COLOR_BIT:
         as!u64(elementAndTexelBlockSize.ElementSize)

gapis/api/vulkan/api/copy_clear_commands.api

@@ -245,7 +245,9 @@ sub void trackVkCmdCopyImage(ref!vkCmdCopyImageArgs args) {
     srcMipLevel := region.srcSubresource.mipLevel
     dstMipLevel := region.dstSubresource.mipLevel
 
-    for _ , _ , aspectBit in unpackImageAspectFlags(region.srcSubresource.aspectMask) {
+    // TODO: Handle multi-planar image correctly. The dst image aspect must
+    // ether match or contains only COLOR_BIT with a compatible format.
+    for _ , _ , aspectBit in getAspectKeysWithAspectFlags(srcImageObject, region.srcSubresource.aspectMask) {
       srcElementSize := switch (aspectBit) {
         case VK_IMAGE_ASPECT_COLOR_BIT:
           as!u64(srcElementAndTexelBlockSize.ElementSize)
@@ -421,7 +423,9 @@ sub void trackVkCmdBlitImage(ref!vkCmdBlitImageArgs args) {
     srcMipLevel := region.srcSubresource.mipLevel
     dstMipLevel := region.dstSubresource.mipLevel
 
-    for _ , _ , aspectBit in unpackImageAspectFlags(region.srcSubresource.aspectMask) {
+    // TODO: Handle multi-planar image correctly. The dst image aspect must
+    // ether match or contains only COLOR_BIT with a compatible format.
+    for _ , _ , aspectBit in getAspectKeysWithAspectFlags(srcImageObject, region.srcSubresource.aspectMask) {
       srcElementSize := switch (aspectBit) {
         case VK_IMAGE_ASPECT_COLOR_BIT:
           as!u64(srcElementAndTexelBlockSize.ElementSize)
@@ -598,7 +602,8 @@ sub void copyImageBuffer(VkBuffer buffer, VkImage image, VkImageLayout layout, m
     rowLength := as!u64(rowLengthAndImageHeight.RowLength / elementAndTexelBlockSize.TexelBlockSize.Width)
     imageHeight := as!u64(rowLengthAndImageHeight.ImageHeight / elementAndTexelBlockSize.TexelBlockSize.Height)
     // The VkImageSubresourceLayer used for buffer image copy should specify only one aspect bit.
-    for _ , _ , aspectBit in unpackImageAspectFlags(region.imageSubresource.aspectMask) {
+    // TODO: Handle multi-planar image correctly.
+    for _ , _ , aspectBit in getAspectKeysWithAspectFlags(imageObject, region.imageSubresource.aspectMask) {
       elementSize := switch (aspectBit) {
         case VK_IMAGE_ASPECT_COLOR_BIT:
           as!u64(elementAndTexelBlockSize.ElementSize)

gapis/api/vulkan/api/image.api

@@ -71,7 +71,7 @@
   VkImageAspectFlags                           ImageAspect
   map!(VkImageAspectFlagBits, ref!ImageAspect) Aspects
   @unused ref!VulkanDebugMarkerInfo            DebugInfo
-  VkMemoryRequirements                         MemoryRequirements
+  ref!ImageMemoryRequirements                  MemoryRequirements
   map!(u32, VkSparseImageMemoryRequirements)   SparseMemoryRequirements
   // Vulkan 1.1 promoted from extension: VK_KHR_dedicated_allocation
   ref!DedicatedRequirements                    DedicatedRequirements
@@ -123,8 +123,8 @@
 }
 
 @internal class ImageMemoryRequirements {
-  VkMemoryRequirements                       MemoryRequirements
-  map!(u32, VkSparseImageMemoryRequirements) AspectBitsToSparseMemoryRequirements
+  map!(VkImageAspectFlagBits, VkMemoryRequirements) PlaneBitsToMemoryRequirements
+  map!(u32, VkSparseImageMemoryRequirements)        AspectBitsToSparseMemoryRequirements
 }
 
 @internal class LinearImageLayouts {
@@ -170,26 +170,7 @@ cmd VkResult vkCreateImage(
     imageInfo.QueueFamilyIndices[i] = queueFamilyIndices[i]
   }
 
-  imageAspect := as!VkImageAspectFlags(
-    switch info.format {
-      // Depth only
-      case VK_FORMAT_D16_UNORM,
-          VK_FORMAT_X8_D24_UNORM_PACK32,
-          VK_FORMAT_D32_SFLOAT:
-        VK_IMAGE_ASPECT_DEPTH_BIT
-      // Stencil only
-      case VK_FORMAT_S8_UINT:
-        VK_IMAGE_ASPECT_STENCIL_BIT
-      // Depth and Stencil
-      case VK_FORMAT_D16_UNORM_S8_UINT,
-          VK_FORMAT_D24_UNORM_S8_UINT,
-          VK_FORMAT_D32_SFLOAT_S8_UINT:
-        VK_IMAGE_ASPECT_DEPTH_BIT |
-        VK_IMAGE_ASPECT_STENCIL_BIT
-      // Others
-      default:
-        VK_IMAGE_ASPECT_COLOR_BIT
-    })
+  imageAspect := as!VkImageAspectFlags(getAspectBitsFromImageFormat(info.format))
 
   hasSparseBit := (as!u32(info.flags) & as!u32(VK_IMAGE_CREATE_SPARSE_BINDING_BIT)) != 0
 
@@ -241,14 +222,12 @@ cmd VkResult vkCreateImage(
     }
   }
 
-  memRequirements := fetchImageMemoryRequirements(device, handle, hasSparseBit)
+  object.MemoryRequirements = new!ImageMemoryRequirements()
   // If the vkCreateImage is inserted by GAPID (e.g. the staging image for
   // reading framebuffer), NO memory requirements will be returned.
-  if memRequirements != null {
-    object.MemoryRequirements = memRequirements.MemoryRequirements
-    if (hasSparseBit) {
-      object.SparseMemoryRequirements = memRequirements.AspectBitsToSparseMemoryRequirements
-    }
+  fetchedRequirements := fetchImageMemoryRequirements(device, handle, hasSparseBit)
+  if fetchedRequirements != null {
+    object.MemoryRequirements = fetchedRequirements
   }
 
   // If the image tiling is LINEAR, get the VkSubresourceLayout for each linear image level
@@ -704,7 +683,7 @@ sub void setQueueInRange(ref!ImageObject image, VkImageSubresourceRange rng,
     ref!QueueObject queue) {
   layerCount := imageSubresourceLayerCount(image, rng)
   levelCount := imageSubresourceLevelCount(image, rng)
-  for _ , _ , aspectBit in unpackImageAspectFlags(rng.aspectMask) {
+  for _ , _ , aspectBit in getAspectKeysWithAspectFlags(image, rng.aspectMask) {
     if aspectBit in image.Aspects {
       layers := image.Aspects[aspectBit].Layers
       for _, i, layer in  layers {
@@ -753,7 +732,7 @@ sub u32 imageSubresourceLevelCount(ref!ImageObject image, VkImageSubresourceRang
 sub void transitionImageLayout(ref!ImageObject img, VkImageSubresourceRange rng, VkImageLayout oldLayout, VkImageLayout newLayout) {
   layerCount := imageSubresourceLayerCount(img, rng)
   levelCount := imageSubresourceLevelCount(img, rng)
-  for _ , _ , aspectBit in unpackImageAspectFlags(rng.aspectMask) {
+  for _ , _ , aspectBit in getAspectKeysWithAspectFlags(img, rng.aspectMask) {
     if (aspectBit in img.Aspects) {
       aspects := img.Aspects[aspectBit]
       for layer in (rng.baseArrayLayer .. rng.baseArrayLayer + layerCount) {
@@ -914,3 +893,96 @@ cmd void vkDestroySamplerYcbcrConversion(
     const VkAllocationCallbacks*                pAllocator) {
   DestroySamplerYcbcrConversion(device, ycbcrConversion, pAllocator)
 }
+
+///////////
+// Utils //
+///////////
+
+sub bool isDisjointImage(ref!ImageObject img) {
+  return (as!u32(img.Info.Flags) & as!u32(VK_IMAGE_CREATE_DISJOINT_BIT)) != as!u32(0)
+}
+
+sub VkImageAspectFlagBits getAspectBitsFromImageFormat(VkFormat fmt) {
+  return switch fmt {
+    // Depth only
+    case VK_FORMAT_D16_UNORM,
+        VK_FORMAT_X8_D24_UNORM_PACK32,
+        VK_FORMAT_D32_SFLOAT:
+      VK_IMAGE_ASPECT_DEPTH_BIT
+    // Stencil only
+    case VK_FORMAT_S8_UINT:
+      VK_IMAGE_ASPECT_STENCIL_BIT
+    // Depth and stencil
+    case VK_FORMAT_D16_UNORM_S8_UINT,
+        VK_FORMAT_D24_UNORM_S8_UINT,
+        VK_FORMAT_D32_SFLOAT_S8_UINT:
+      VK_IMAGE_ASPECT_DEPTH_BIT | VK_IMAGE_ASPECT_STENCIL_BIT
+    // Multi-planar
+    case VK_FORMAT_G8_B8_R8_3PLANE_420_UNORM_KHR,
+        VK_FORMAT_G8_B8_R8_3PLANE_422_UNORM_KHR,
+        VK_FORMAT_G8_B8_R8_3PLANE_444_UNORM_KHR,
+        VK_FORMAT_G10X6_B10X6_R10X6_3PLANE_420_UNORM_3PACK16_KHR,
+        VK_FORMAT_G10X6_B10X6_R10X6_3PLANE_422_UNORM_3PACK16_KHR,
+        VK_FORMAT_G10X6_B10X6_R10X6_3PLANE_444_UNORM_3PACK16_KHR,
+        VK_FORMAT_G12X4_B12X4_R12X4_3PLANE_420_UNORM_3PACK16_KHR,
+        VK_FORMAT_G12X4_B12X4_R12X4_3PLANE_422_UNORM_3PACK16_KHR,
+        VK_FORMAT_G12X4_B12X4_R12X4_3PLANE_444_UNORM_3PACK16_KHR,
+        VK_FORMAT_G16_B16_R16_3PLANE_420_UNORM_KHR,
+        VK_FORMAT_G16_B16_R16_3PLANE_422_UNORM_KHR,
+        VK_FORMAT_G16_B16_R16_3PLANE_444_UNORM_KHR:
+      VK_IMAGE_ASPECT_PLANE_0_BIT | VK_IMAGE_ASPECT_PLANE_1_BIT | VK_IMAGE_ASPECT_PLANE_2_BIT
+
+    case VK_FORMAT_G8_B8R8_2PLANE_422_UNORM_KHR,
+        VK_FORMAT_G8_B8R8_2PLANE_420_UNORM_KHR,
+        VK_FORMAT_G10X6_B10X6R10X6_2PLANE_420_UNORM_3PACK16_KHR,
+        VK_FORMAT_G10X6_B10X6R10X6_2PLANE_422_UNORM_3PACK16_KHR,
+        VK_FORMAT_G12X4_B12X4R12X4_2PLANE_420_UNORM_3PACK16_KHR,
+        VK_FORMAT_G12X4_B12X4R12X4_2PLANE_422_UNORM_3PACK16_KHR,
+        VK_FORMAT_G16_B16R16_2PLANE_420_UNORM_KHR,
+        VK_FORMAT_G16_B16R16_2PLANE_422_UNORM_KHR:
+      VK_IMAGE_ASPECT_PLANE_0_BIT | VK_IMAGE_ASPECT_PLANE_1_BIT
+
+    case VK_FORMAT_G8B8G8R8_422_UNORM,
+        VK_FORMAT_B8G8R8G8_422_UNORM,
+        VK_FORMAT_R10X6G10X6B10X6A10X6_UNORM_4PACK16,
+        VK_FORMAT_G10X6B10X6G10X6R10X6_422_UNORM_4PACK16,
+        VK_FORMAT_B10X6G10X6R10X6G10X6_422_UNORM_4PACK16,
+        VK_FORMAT_R12X4G12X4B12X4A12X4_UNORM_4PACK16,
+        VK_FORMAT_G12X4B12X4G12X4R12X4_422_UNORM_4PACK16,
+        VK_FORMAT_B12X4G12X4R12X4G12X4_422_UNORM_4PACK16,
+        VK_FORMAT_G16B16G16R16_422_UNORM,
+        VK_FORMAT_B16G16R16G16_422_UNORM:
+      VK_IMAGE_ASPECT_PLANE_0_BIT
+    // Others
+    default:
+      VK_IMAGE_ASPECT_COLOR_BIT
+  }
+}
+
+sub dense_map!(u32, VkImageAspectFlagBits) getAspectKeysWithAspectFlags(ref!ImageObject img, VkImageAspectFlags flags) {
+  keys := new!ImageAspectBits()
+  numPlanes := numberOfPlanes(img.Info.Format)
+  if numPlanes == as!u32(0) {
+    // single planar image
+    keys.Bits = unpackImageAspectFlags(flags)
+  } else {
+    // multiple planar image
+    if (as!u32(flags) & as!u32(VK_IMAGE_ASPECT_COLOR_BIT)) != as!u32(0) {
+      // Has COLOR bit, expand to all planes
+      keys.Bits = unpackImageAspectFlags(img.ImageAspect)
+    } else {
+      keys.Bits = unpackImageAspectFlags(flags)
+    }
+  }
+  return keys.Bits
+}
+
+sub VkMemoryRequirements getImagePlaneMemoryRequirements(ref!ImageObject img, VkImageAspectFlagBits plane) {
+  disjoint := isDisjointImage(img)
+  return switch disjoint {
+    case true:
+      img.MemoryRequirements.PlaneBitsToMemoryRequirements[plane]
+    default:
+      img.MemoryRequirements.PlaneBitsToMemoryRequirements[as!VkImageAspectFlagBits(0)]
+  }
+}

gapis/api/vulkan/api/properties_features_requirements.api

@@ -273,7 +273,7 @@ cmd void vkGetImageMemoryRequirements(
     requirements := ?
     if pMemoryRequirements == null { vkErrorNullPointer("VkMemoryRequirements") }
     pMemoryRequirements[0] = requirements
-    Images[image].MemoryRequirements = requirements
+    Images[image].MemoryRequirements.PlaneBitsToMemoryRequirements[as!VkImageAspectFlagBits(0)] = requirements
   }
 }
 
@@ -298,7 +298,7 @@ cmd void vkGetImageSparseMemoryRequirements(
     for i in (0 .. count) {
       requirements[i] = ?
       aspect := requirements[i].formatProperties.aspectMask
-      Images[image].SparseMemoryRequirements[as!u32(aspect)] = requirements[i]
+      Images[image].MemoryRequirements.AspectBitsToSparseMemoryRequirements[as!u32(aspect)] = requirements[i]
     }
     pSparseMemoryRequirementCount[0] = count
   }
@@ -850,13 +850,20 @@ sub void GetImageMemoryRequirements2(
     const VkImageMemoryRequirementsInfo2* pInfo,
     VkMemoryRequirements2*                pMemoryRequirements) {
   if !(device in Devices) { vkErrorInvalidDevice(device) }
+  plane := MutableU32(0)
   info := pInfo[0]
   memReqs := pMemoryRequirements[0]
   if memReqs.pNext != null {
     nPNext := numberOfPNext(as!const void*(memReqs.pNext))
     next := MutableVoidPtr(as!void*(memReqs.pNext))
     for i in (0 .. nPNext) {
-      _ = as!const VkStructureType*(next.Ptr)[0]
+      sType := as!const VkStructureType*(next.Ptr)[0]
+      switch sType {
+        case VK_STRUCTURE_TYPE_IMAGE_PLANE_MEMORY_REQUIREMENTS_INFO: {
+          ext := as!VkImagePlaneMemoryRequirementsInfo*(next.Ptr)[0]
+          plane.Val = as!u32(ext.planeAspect)
+        }
+      }
       next.Ptr = as!VulkanStructHeader*(next.Ptr)[0].PNext
     }
   }
@@ -871,7 +878,7 @@ sub void GetImageMemoryRequirements2(
   if !(info.image in Images) { vkErrorInvalidImage(info.image) }
   // TODO: Drop the touch of the image object once we extract the memory
   // requirement info out of the image object.
-  Images[info.image].MemoryRequirements = memReq.memoryRequirements
+  Images[info.image].MemoryRequirements.PlaneBitsToMemoryRequirements[as!VkImageAspectFlagBits(plane.Val)] = memReq.memoryRequirements
 
   if memReq.pNext != null {
     numPNext := numberOfPNext(as!const void*(memReq.pNext))
@@ -937,7 +944,7 @@ sub void GetImageSparseMemoryRequirements2(
       // TODO: handle pNext for 'requirements[i]'
       aspect := requirements[i].memoryRequirements.formatProperties.aspectMask
       // TODO: Drop the touch of the image object.
-      Images[image].SparseMemoryRequirements[as!u32(aspect)] = requirements[i].memoryRequirements
+      Images[image].MemoryRequirements.AspectBitsToSparseMemoryRequirements[as!u32(aspect)] = requirements[i].memoryRequirements
     }
     pSparseMemoryRequirementCount[0] = count
   }

gapis/api/vulkan/api/queue.api

@@ -259,9 +259,10 @@ sub void addSparseImageMemoryBinding(VkImage image, VkSparseImageMemoryBind bind
   }
   levelInfo := layerInfo.Levels[level]
 
-  blockSize := img.MemoryRequirements.alignment
-  if (aspect in img.SparseMemoryRequirements) {
-    gran := img.SparseMemoryRequirements[aspect].formatProperties.imageGranularity
+  // TODO: Handle multi-planar images
+  blockSize := getImagePlaneMemoryRequirements(img, as!VkImageAspectFlagBits(0)).alignment
+  if (aspect in img.MemoryRequirements.AspectBitsToSparseMemoryRequirements) {
+    gran := img.MemoryRequirements.AspectBitsToSparseMemoryRequirements[aspect].formatProperties.imageGranularity
     mipWidth := getMipSize(img.Info.Extent.Width, level)
     mipHeight := getMipSize(img.Info.Extent.Height, level)
     widthInBlocks := roundUpTo(mipWidth, gran.Width)

gapis/api/vulkan/api/queued_command_tracking.api

@@ -106,7 +106,8 @@ sub void execPendingCommands(VkQueue queue, bool isRoot) {
           for _ , img , binds in cmd.SparseBinds.OpaqueImageBinds {
             if (!img in Images) { vkErrorInvalidImage(img) }
             image := Images[img]
-            blockSize := image.MemoryRequirements.alignment
+            // TODO: Handle multi-planar images
+            blockSize := getImagePlaneMemoryRequirements(image, as!VkImageAspectFlagBits(0)).alignment
             for _ , _ , bind in binds.SparseMemoryBinds {
               numBlocks := (bind.size + blockSize - 1) / blockSize
               memOffset := bind.memoryOffset