Decouple creation, allocation and binding of images from Wayland backend

A lot of the steps involved in creating a VkImage, allocating backing
memory, and binding the VkImage to the memory do not rely specifically
on Wayland features and could be common to other WSI backends.

This commit attempts to decouple the create_and_bind_swapchain_image()
function in the wsi/wayland/swapchain.cpp from Wayland and wsialloc
specific parts, and instead move the independent parts related to
external memory management to a new external_memory class. This will
allow the independent code to be re-used in future WSI backends.

Signed-off-by: Dennis Tsiang <dennis.tsiang@arm.com>
Change-Id: I8c13ccbbfaef0d345fcd06192e6de484dd53645f

7 files changed, 561 insertions, 322 deletions

diff --git a/CMakeLists.txt b/CMakeLists.txt
index f07ace8..62a9ccd 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -194,6 +194,7 @@ add_library(${PROJECT_NAME} SHARED
    util/extension_list.cpp
    util/log.cpp
    util/format_modifiers.cpp
+   wsi/external_memory.cpp
    wsi/surface_properties.cpp
    wsi/swapchain_base.cpp
    wsi/synchronization.cpp
diff --git a/util/helpers.hpp b/util/helpers.hpp
index c477619..a72591d 100644
--- a/util/helpers.hpp
+++ b/util/helpers.hpp
@@ -56,16 +56,15 @@
  *    return VK_SUCCESS;
  * }
  */
-#define TRY(expression) \
-   do \
-   { \
+#define TRY(expression)                 \
+   do                                   \
+   {                                    \
       VkResult try_result = expression; \
-      if (try_result != VK_SUCCESS) \
-      { \
-         return try_result; \
-      } \
-   } \
-   while (0)
+      if (try_result != VK_SUCCESS)     \
+      {                                 \
+         return try_result;             \
+      }                                 \
+   } while (0)
 
 namespace util
 {
@@ -99,6 +98,19 @@ protected:
 
 private:
    noncopyable(const noncopyable &) = delete;
-   noncopyable& operator=(const noncopyable &) = delete;
+   noncopyable &operator=(const noncopyable &) = delete;
 };
+
+static constexpr uint32_t MAX_PLANES = 4;
+
+/**
+ * @brief Helper variable for plane image aspect flag bits.
+ */
+const VkImageAspectFlagBits PLANE_FLAG_BITS[MAX_PLANES] = {
+   VK_IMAGE_ASPECT_MEMORY_PLANE_0_BIT_EXT,
+   VK_IMAGE_ASPECT_MEMORY_PLANE_1_BIT_EXT,
+   VK_IMAGE_ASPECT_MEMORY_PLANE_2_BIT_EXT,
+   VK_IMAGE_ASPECT_MEMORY_PLANE_3_BIT_EXT,
+};
+
 } // namespace util
diff --git a/wsi/external_memory.cpp b/wsi/external_memory.cpp
new file mode 100644
index 0000000..80c7438
--- /dev/null
+++ b/wsi/external_memory.cpp
@@ -0,0 +1,254 @@
+/*
+ * Copyright (c) 2022 Arm Limited.
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to
+ * deal in the Software without restriction, including without limitation the
+ * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+ * sell copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+#include "external_memory.hpp"
+
+#include <cassert>
+#include <cstdint>
+#include <unistd.h>
+#include <algorithm>
+
+#include "util/log.hpp"
+#include "util/helpers.hpp"
+
+namespace wsi
+{
+
+external_memory::external_memory(const VkDevice &device, const util::allocator &allocator)
+   : m_device(device)
+   , m_allocator(allocator)
+{
+}
+
+external_memory::~external_memory()
+{
+   auto &device_data = layer::device_private_data::get(m_device);
+   for (uint32_t plane = 0; plane < get_num_planes(); plane++)
+   {
+      auto &memory = m_memories[plane];
+      if (memory != VK_NULL_HANDLE)
+      {
+         device_data.disp.FreeMemory(m_device, memory, m_allocator.get_original_callbacks());
+      }
+      else if (m_buffer_fds[plane] >= 0)
+      {
+         auto it = std::find(std::begin(m_buffer_fds), std::end(m_buffer_fds), m_buffer_fds[plane]);
+         if (std::distance(std::begin(m_buffer_fds), it) == plane)
+         {
+            close(m_buffer_fds[plane]);
+         }
+      }
+   }
+}
+
+uint32_t external_memory::get_num_planes()
+{
+   if (m_num_planes == 0)
+   {
+      for (uint32_t plane = 0; plane < MAX_PLANES; plane++)
+      {
+         if (m_buffer_fds[plane] == -1)
+         {
+            break;
+         }
+         m_num_planes++;
+      }
+   }
+   return m_num_planes;
+}
+
+bool external_memory::is_disjoint()
+{
+   uint32_t planes = get_num_planes();
+   if (planes > 1)
+   {
+      for (uint32_t plane = 1; plane < planes; ++plane)
+      {
+         if (m_buffer_fds[plane] != m_buffer_fds[0])
+         {
+            return true;
+         }
+      }
+   }
+   return false;
+}
+
+VkResult external_memory::get_fd_mem_type_index(uint32_t index, uint32_t *mem_idx)
+{
+   auto &device_data = layer::device_private_data::get(m_device);
+   VkMemoryFdPropertiesKHR mem_props = {};
+   mem_props.sType = VK_STRUCTURE_TYPE_MEMORY_FD_PROPERTIES_KHR;
+
+   VkResult result =
+      device_data.disp.GetMemoryFdPropertiesKHR(m_device, m_handle_type, m_buffer_fds[index], &mem_props);
+   if (result != VK_SUCCESS)
+   {
+      WSI_LOG_ERROR("Error querying Fd properties.");
+      return result;
+   }
+
+   for (*mem_idx = 0; *mem_idx < VK_MAX_MEMORY_TYPES; (*mem_idx)++)
+   {
+      if (mem_props.memoryTypeBits & (1 << *mem_idx))
+      {
+         break;
+      }
+   }
+
+   assert(*mem_idx < VK_MAX_MEMORY_TYPES);
+
+   return VK_SUCCESS;
+}
+
+VkResult external_memory::import_plane_memories()
+{
+   if (is_disjoint())
+   {
+      for (uint32_t plane = 0; plane < get_num_planes(); plane++)
+      {
+         auto it = std::find(std::begin(m_buffer_fds), std::end(m_buffer_fds), m_buffer_fds[plane]);
+         if (std::distance(std::begin(m_buffer_fds), it) == plane)
+         {
+            TRY(import_plane_memory(plane));
+         }
+      }
+      return VK_SUCCESS;
+   }
+   return import_plane_memory(0);
+}
+
+VkResult external_memory::import_plane_memory(uint32_t index)
+{
+   uint32_t mem_index = -1;
+   TRY(get_fd_mem_type_index(index, &mem_index));
+
+   const off_t fd_size = lseek(m_buffer_fds[index], 0, SEEK_END);
+   if (fd_size < 0)
+   {
+      WSI_LOG_ERROR("Failed to get fd size.");
+      return VK_ERROR_OUT_OF_HOST_MEMORY;
+   }
+
+   VkImportMemoryFdInfoKHR import_mem_info = {};
+   import_mem_info.sType = VK_STRUCTURE_TYPE_IMPORT_MEMORY_FD_INFO_KHR;
+   import_mem_info.handleType = m_handle_type;
+   import_mem_info.fd = m_buffer_fds[index];
+
+   VkMemoryAllocateInfo alloc_info = {};
+   alloc_info.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
+   alloc_info.pNext = &import_mem_info;
+   alloc_info.allocationSize = static_cast<uint64_t>(fd_size);
+   alloc_info.memoryTypeIndex = mem_index;
+
+   auto &device_data = layer::device_private_data::get(m_device);
+   VkResult result =
+      device_data.disp.AllocateMemory(m_device, &alloc_info, m_allocator.get_original_callbacks(), &m_memories[index]);
+
+   if (result != VK_SUCCESS)
+   {
+      WSI_LOG_ERROR("Failed to import memory.");
+      return result;
+   }
+
+   return VK_SUCCESS;
+}
+
+VkResult external_memory::bind_swapchain_image_memory(const VkImage &image)
+{
+   auto &device_data = layer::device_private_data::get(m_device);
+   if (is_disjoint())
+   {
+      util::vector<VkBindImageMemoryInfo> bind_img_mem_infos(m_allocator);
+      if (!bind_img_mem_infos.try_resize(get_num_planes()))
+      {
+         return VK_ERROR_OUT_OF_HOST_MEMORY;
+      }
+
+      util::vector<VkBindImagePlaneMemoryInfo> bind_plane_mem_infos(m_allocator);
+      if (!bind_plane_mem_infos.try_resize(get_num_planes()))
+      {
+         return VK_ERROR_OUT_OF_HOST_MEMORY;
+      }
+
+      for (uint32_t plane = 0; plane < get_num_planes(); plane++)
+      {
+
+         bind_plane_mem_infos[plane].planeAspect = util::PLANE_FLAG_BITS[plane];
+         bind_plane_mem_infos[plane].sType = VK_STRUCTURE_TYPE_BIND_IMAGE_PLANE_MEMORY_INFO;
+         bind_plane_mem_infos[plane].pNext = nullptr;
+
+         bind_img_mem_infos[plane].sType = VK_STRUCTURE_TYPE_BIND_IMAGE_MEMORY_INFO;
+         bind_img_mem_infos[plane].pNext = &bind_plane_mem_infos[plane];
+         bind_img_mem_infos[plane].image = image;
+         bind_img_mem_infos[plane].memory = m_memories[plane];
+         bind_img_mem_infos[plane].memoryOffset = m_offsets[plane];
+      }
+
+      return device_data.disp.BindImageMemory2KHR(m_device, bind_img_mem_infos.size(), bind_img_mem_infos.data());
+   }
+
+   return device_data.disp.BindImageMemory(m_device, image, m_memories[0], m_offsets[0]);
+}
+
+VkResult external_memory::import_memory_and_bind_swapchain_image(const VkImage &image)
+{
+   TRY(import_plane_memories());
+   TRY(bind_swapchain_image_memory(image));
+   return VK_SUCCESS;
+}
+
+VkResult external_memory::fill_image_plane_layouts(util::vector<VkSubresourceLayout> &image_plane_layouts)
+{
+   if (!image_plane_layouts.try_resize(get_num_planes()))
+   {
+      return VK_ERROR_OUT_OF_HOST_MEMORY;
+   }
+
+   for (uint32_t plane = 0; plane < get_num_planes(); plane++)
+   {
+      assert(m_strides[plane] >= 0);
+      image_plane_layouts[plane].offset = m_offsets[plane];
+      image_plane_layouts[plane].rowPitch = static_cast<uint32_t>(m_strides[plane]);
+   }
+   return VK_SUCCESS;
+}
+
+void external_memory::fill_drm_mod_info(const void *pNext, VkImageDrmFormatModifierExplicitCreateInfoEXT &drm_mod_info,
+                                        util::vector<VkSubresourceLayout> &plane_layouts, uint64_t modifier)
+{
+   drm_mod_info.sType = VK_STRUCTURE_TYPE_IMAGE_DRM_FORMAT_MODIFIER_EXPLICIT_CREATE_INFO_EXT;
+   drm_mod_info.pNext = pNext;
+   drm_mod_info.drmFormatModifier = modifier;
+   drm_mod_info.drmFormatModifierPlaneCount = get_num_planes();
+   drm_mod_info.pPlaneLayouts = plane_layouts.data();
+}
+
+void external_memory::fill_external_info(VkExternalMemoryImageCreateInfoKHR &external_info, void *pNext)
+{
+   external_info.sType = VK_STRUCTURE_TYPE_EXTERNAL_MEMORY_IMAGE_CREATE_INFO_KHR;
+   external_info.pNext = pNext;
+   external_info.handleTypes = m_handle_type;
+}
+
+} // namespace wsi
diff --git a/wsi/external_memory.hpp b/wsi/external_memory.hpp
new file mode 100644
index 0000000..1bfa5b5
--- /dev/null
+++ b/wsi/external_memory.hpp
@@ -0,0 +1,188 @@
+/*
+ * Copyright (c) 2022 Arm Limited.
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to
+ * deal in the Software without restriction, including without limitation the
+ * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+ * sell copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+/**
+ * @file external_memory.hpp
+ *
+ * @brief Handles importing and binding external memory for swapchain implementations.
+ */
+
+#pragma once
+
+#include <array>
+#include <cstdint>
+#include <vulkan/vulkan.h>
+
+#include "wsi/synchronization.hpp"
+#include "layer/private_data.hpp"
+#include "util/custom_allocator.hpp"
+#include "util/helpers.hpp"
+
+namespace wsi
+{
+
+using util::MAX_PLANES;
+
+class external_memory
+{
+public:
+   external_memory(const VkDevice &device, const util::allocator &allocator);
+   ~external_memory();
+
+   /**
+    * @brief Get the fds representing the externally created memory for each plane.
+    */
+   const std::array<int, MAX_PLANES> &get_buffer_fds(void)
+   {
+      return m_buffer_fds;
+   }
+
+   /**
+    * @brief Get the per plane stride values.
+    */
+   const std::array<int, MAX_PLANES> &get_strides(void)
+   {
+      return m_strides;
+   }
+
+   /**
+    * @brief Get the per plane offset values.
+    */
+   const std::array<uint32_t, MAX_PLANES> &get_offsets(void)
+   {
+      return m_offsets;
+   }
+
+   /**
+    * @brief Set the per plane fd values.
+    */
+   void set_buffer_fds(std::array<int, MAX_PLANES> buffer_fds)
+   {
+      m_buffer_fds = buffer_fds;
+   }
+
+   /**
+    * @brief Set the per plane stride values.
+    */
+   void set_strides(std::array<int, MAX_PLANES> strides)
+   {
+      m_strides = strides;
+   }
+
+   /**
+    * @brief Set the per plane offset values.
+    */
+   void set_offsets(std::array<uint32_t, MAX_PLANES> offsets)
+   {
+      m_offsets = offsets;
+   }
+
+   /**
+    * @brief Get the number of planes the external format uses.
+    */
+   uint32_t get_num_planes(void);
+
+   /**
+    * @brief Returns whether the external memory uses a multi-planar format where each plane is separately bound to
+    *        memory.
+    */
+   bool is_disjoint(void);
+
+   /**
+    * @brief Set the external memory type.
+    */
+   void set_memory_handle_type(VkExternalMemoryHandleTypeFlagBits handle_type)
+   {
+      m_handle_type = handle_type;
+   }
+
+   /**
+    * @brief Binds the external memory to a swapchain image.
+    *
+    * Must only be used after the external memory has been imported.
+    *
+    * @param image     The swapchain image.
+    *
+    * @return VK_ERROR_OUT_OF_HOST_MEMORY when out of memory else returns the result of
+    *         VkBindImageMemory/VkBindImageMemory2KHR.
+    */
+   VkResult bind_swapchain_image_memory(const VkImage &image);
+
+   /**
+    * @brief Imports the externally allocated memory into the Vulkan framework and binds it to a swapchain image.
+    *
+    * @param image     The swapchain image.
+    *
+    * @return VK_ERROR_OUT_OF_HOST_MEMORY when out of memory. Other possible values include the result of calling
+    *         VkAllocateMemory/VkBindImageMemory/VkBindImageMemory2KHR.
+    */
+   VkResult import_memory_and_bind_swapchain_image(const VkImage &image);
+
+   /**
+    * @brief Fills out a list of VkSubresourceLayout for each plane using the stored planes layout data.
+    *
+    * @param image_plane_layouts    A list of plane layouts to fill.
+    *
+    * @return VK_ERROR_OUT_OF_HOST_MEMORY when out of memory else VK_SUCCESS on success.
+    */
+   VkResult fill_image_plane_layouts(util::vector<VkSubresourceLayout> &image_plane_layouts);
+
+   /**
+    * @brief Fills out a VkImageDrmFormatModifierExplicitCreateInfoEXT struct.
+    *
+    * @param pNext         Pointer to the pNext chain to set the VkImageDrmFormatModifierExplicitCreateInfoEXT struct to.
+    * @param drm_mod_info  The struct to fill out.
+    * @param image_layouts Per plane image layouts
+    * @param modifier      Modifier that the DRM format will use.
+    */
+   void fill_drm_mod_info(const void *pNext, VkImageDrmFormatModifierExplicitCreateInfoEXT &drm_mod_info,
+                          util::vector<VkSubresourceLayout> &image_layout, uint64_t modifier);
+
+   /**
+    * @brief Fills out a VkExternalMemoryImageCreateInfoKHR struct.
+    *
+    * @param external_info The struct to fill out.
+    * @param pNext         Pointer to the pNext chain to set the VkExternalMemoryImageCreateInfoKHR struct to.
+    */
+   void fill_external_info(VkExternalMemoryImageCreateInfoKHR &external_info, void *pNext);
+
+private:
+   VkResult get_fd_mem_type_index(uint32_t index, uint32_t *mem_idx);
+
+   VkResult import_plane_memories(void);
+
+   VkResult import_plane_memory(uint32_t index);
+
+   std::array<int, MAX_PLANES> m_buffer_fds{ -1, -1, -1, -1 };
+   std::array<int, MAX_PLANES> m_strides{ 0, 0, 0, 0 };
+   std::array<uint32_t, MAX_PLANES> m_offsets{ 0, 0, 0, 0 };
+   std::array<VkDeviceMemory, MAX_PLANES> m_memories = { VK_NULL_HANDLE, VK_NULL_HANDLE, VK_NULL_HANDLE,
+                                                         VK_NULL_HANDLE };
+   uint32_t m_num_planes{ 0 };
+   VkExternalMemoryHandleTypeFlagBits m_handle_type{ VK_EXTERNAL_MEMORY_HANDLE_TYPE_DMA_BUF_BIT_EXT };
+   const VkDevice &m_device;
+   const util::allocator &m_allocator;
+};
+
+} // namespace wsi
diff --git a/wsi/swapchain_base.hpp b/wsi/swapchain_base.hpp
index c79c7b7..5a53298 100644
--- a/wsi/swapchain_base.hpp
+++ b/wsi/swapchain_base.hpp
@@ -42,9 +42,13 @@
 #include <util/ring_buffer.hpp>
 #include "surface_properties.hpp"
 #include "wsi/synchronization.hpp"
+#include "util/helpers.hpp"
 
 namespace wsi
 {
+
+using util::MAX_PLANES;
+
 struct swapchain_image
 {
    enum status
@@ -57,15 +61,13 @@ struct swapchain_image
    };
 
    /* Implementation specific data */
-   void *data{nullptr};
+   void *data{ nullptr };
 
-   VkImage image{VK_NULL_HANDLE};
-   status status{swapchain_image::INVALID};
+   VkImage image{ VK_NULL_HANDLE };
+   status status{ swapchain_image::INVALID };
    VkSemaphore present_semaphore{ VK_NULL_HANDLE };
 };
 
-constexpr uint32_t MAX_PLANES = 4;
-
 #if WSI_IMAGE_COMPRESSION_CONTROL_SWAPCHAIN
 struct image_compression_control_params
 {
@@ -206,7 +208,6 @@ public:
    }
 
 protected:
-
    layer::device_private_data &m_device_data;
 
    /**
diff --git a/wsi/wayland/swapchain.cpp b/wsi/wayland/swapchain.cpp
index 4067ff9..9a4bf4c 100644
--- a/wsi/wayland/swapchain.cpp
+++ b/wsi/wayland/swapchain.cpp
@@ -49,28 +49,6 @@ namespace wsi
 namespace wayland
 {
 
-const VkImageAspectFlagBits plane_flag_bits[MAX_PLANES] = {
-   VK_IMAGE_ASPECT_MEMORY_PLANE_0_BIT_EXT,
-   VK_IMAGE_ASPECT_MEMORY_PLANE_1_BIT_EXT,
-   VK_IMAGE_ASPECT_MEMORY_PLANE_2_BIT_EXT,
-   VK_IMAGE_ASPECT_MEMORY_PLANE_3_BIT_EXT,
-};
-
-struct swapchain::wayland_image_data
-{
-   int buffer_fd[MAX_PLANES];
-   int stride[MAX_PLANES];
-   uint32_t offset[MAX_PLANES];
-
-   wl_buffer *buffer;
-   VkDeviceMemory memory[MAX_PLANES];
-
-   uint32_t num_planes;
-
-   sync_fd_fence_sync present_fence;
-   bool is_disjoint;
-};
-
 swapchain::swapchain(layer::device_private_data &dev_data, const VkAllocationCallbacks *pAllocator,
                      surface &wsi_surface)
    : swapchain_base(dev_data, pAllocator)
@@ -156,82 +134,6 @@ void swapchain::release_buffer(struct wl_buffer *wayl_buffer)
 
 static struct wl_buffer_listener buffer_listener = { buffer_release };
 
-VkResult swapchain::allocate_plane_memory(int fd, VkDeviceMemory *memory)
-{
-   uint32_t mem_index = -1;
-   VkResult result = get_fd_mem_type_index(fd, mem_index);
-   if (result != VK_SUCCESS)
-   {
-      return result;
-   }
-
-   const off_t dma_buf_size = lseek(fd, 0, SEEK_END);
-   if (dma_buf_size < 0)
-   {
-      WSI_LOG_ERROR("Failed to get DMA Buf size.");
-      return VK_ERROR_OUT_OF_HOST_MEMORY;
-   }
-
-   VkImportMemoryFdInfoKHR import_mem_info = {};
-   import_mem_info.sType = VK_STRUCTURE_TYPE_IMPORT_MEMORY_FD_INFO_KHR;
-   import_mem_info.handleType = VK_EXTERNAL_MEMORY_HANDLE_TYPE_DMA_BUF_BIT_EXT;
-   import_mem_info.fd = fd;
-
-   VkMemoryAllocateInfo alloc_info = {};
-   alloc_info.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
-   alloc_info.pNext = &import_mem_info;
-   alloc_info.allocationSize = static_cast<uint64_t>(dma_buf_size);
-   alloc_info.memoryTypeIndex = mem_index;
-
-   result = m_device_data.disp.AllocateMemory(
-      m_device, &alloc_info, get_allocation_callbacks(), memory);
-
-   if (result != VK_SUCCESS)
-   {
-      WSI_LOG_ERROR("Failed to import memory.");
-      return result;
-   }
-
-   return VK_SUCCESS;
-}
-
-VkResult swapchain::get_fd_mem_type_index(int fd, uint32_t &mem_idx)
-{
-   VkMemoryFdPropertiesKHR mem_props = {};
-   mem_props.sType = VK_STRUCTURE_TYPE_MEMORY_FD_PROPERTIES_KHR;
-
-   VkResult result = m_device_data.disp.GetMemoryFdPropertiesKHR(
-      m_device, VK_EXTERNAL_MEMORY_HANDLE_TYPE_DMA_BUF_BIT_EXT, fd, &mem_props);
-   if (result != VK_SUCCESS)
-   {
-      WSI_LOG_ERROR("Error querying Fd properties.");
-      return result;
-   }
-
-   for (mem_idx = 0; mem_idx < VK_MAX_MEMORY_TYPES; mem_idx++)
-   {
-      if (mem_props.memoryTypeBits & (1 << mem_idx))
-      {
-         break;
-      }
-   }
-
-   assert(mem_idx < VK_MAX_MEMORY_TYPES);
-
-   return VK_SUCCESS;
-}
-
-static uint32_t get_same_fd_index(int fd, int const *fds)
-{
-   uint32_t index = 0;
-   while (fd != fds[index])
-   {
-      index++;
-   }
-
-   return index;
-}
-
 VkResult swapchain::get_surface_compatible_formats(const VkImageCreateInfo &info,
                                                    util::vector<wsialloc_format> &importable_formats,
                                                    util::vector<uint64_t> &exportable_modifers)
@@ -353,7 +255,7 @@ VkResult swapchain::get_surface_compatible_formats(const VkImageCreateInfo &info
    return VK_SUCCESS;
 }
 
-VkResult swapchain::allocate_wsialloc(VkImageCreateInfo &image_create_info, wayland_image_data &image_data,
+VkResult swapchain::allocate_wsialloc(VkImageCreateInfo &image_create_info, wayland_image_data *image_data,
                                       util::vector<wsialloc_format> &importable_formats,
                                       wsialloc_format *allocated_format)
 {
@@ -371,8 +273,11 @@ VkResult swapchain::allocate_wsialloc(VkImageCreateInfo &image_create_info, wayl
                                          image_create_info.extent.width, image_create_info.extent.height,
                                          allocation_flags };
 
-   const auto res = wsialloc_alloc(m_wsi_allocator, &alloc_info, allocated_format, image_data.stride,
-                                   image_data.buffer_fd, image_data.offset);
+   std::array<int, MAX_PLANES> strides{};
+   std::array<int, MAX_PLANES> buffer_fds{ -1, -1, -1, -1 };
+   std::array<uint32_t, MAX_PLANES> offsets{};
+   const auto res =
+      wsialloc_alloc(m_wsi_allocator, &alloc_info, allocated_format, strides.data(), buffer_fds.data(), offsets.data());
    if (res != WSIALLOC_ERROR_NONE)
    {
       WSI_LOG_ERROR("Failed allocation of DMA Buffer. WSI error: %d", static_cast<int>(res));
@@ -382,42 +287,48 @@ VkResult swapchain::allocate_wsialloc(VkImageCreateInfo &image_create_info, wayl
       }
       return VK_ERROR_OUT_OF_HOST_MEMORY;
    }
+   auto &external_memory = image_data->external_mem;
+   external_memory.set_strides(strides);
+   external_memory.set_buffer_fds(buffer_fds);
+   external_memory.set_offsets(offsets);
    return VK_SUCCESS;
 }
 
-VkResult swapchain::allocate_image(VkImageCreateInfo &image_create_info, wayland_image_data *image_data, VkImage *image)
+static VkResult fill_image_create_info(VkImageCreateInfo &image_create_info,
+                                       util::vector<VkSubresourceLayout> &image_plane_layouts,
+                                       VkImageDrmFormatModifierExplicitCreateInfoEXT &drm_mod_info,
+                                       VkExternalMemoryImageCreateInfoKHR &external_info,
+                                       wayland_image_data &image_data, uint64_t modifier)
 {
-   image_data->buffer = nullptr;
-   image_data->num_planes = 0;
-   for (uint32_t plane = 0; plane < MAX_PLANES; plane++)
+   TRY(image_data.external_mem.fill_image_plane_layouts(image_plane_layouts));
+
+   if (image_data.external_mem.is_disjoint())
    {
-      image_data->buffer_fd[plane] = -1;
-      image_data->memory[plane] = VK_NULL_HANDLE;
+      image_create_info.flags |= VK_IMAGE_CREATE_DISJOINT_BIT;
    }
 
-   bool is_disjoint = false;
+   image_data.external_mem.fill_drm_mod_info(image_create_info.pNext, drm_mod_info, image_plane_layouts, modifier);
+   image_data.external_mem.fill_external_info(external_info, &drm_mod_info);
+   image_create_info.pNext = &external_info;
+   image_create_info.tiling = VK_IMAGE_TILING_DRM_FORMAT_MODIFIER_EXT;
+   return VK_SUCCESS;
+}
+
+VkResult swapchain::allocate_image(VkImageCreateInfo &image_create_info, wayland_image_data *image_data, VkImage *image)
+{
    util::vector<wsialloc_format> importable_formats(util::allocator(m_allocator, VK_SYSTEM_ALLOCATION_SCOPE_COMMAND));
    auto &m_allocated_format = m_image_creation_parameters.m_allocated_format;
    if (m_image_create_info.format != VK_FORMAT_UNDEFINED)
    {
-      is_disjoint = (m_image_create_info.flags & VK_IMAGE_CREATE_DISJOINT_BIT) == VK_IMAGE_CREATE_DISJOINT_BIT;
       if (!importable_formats.try_push_back(m_allocated_format))
       {
          return VK_ERROR_OUT_OF_HOST_MEMORY;
       }
-      VkResult result = allocate_wsialloc(m_image_create_info, *image_data, importable_formats, &m_allocated_format);
+      VkResult result = allocate_wsialloc(m_image_create_info, image_data, importable_formats, &m_allocated_format);
       if (result != VK_SUCCESS)
       {
          return result;
       }
-      for (uint32_t plane = 0; plane < MAX_PLANES; plane++)
-      {
-         if (image_data->buffer_fd[plane] == -1)
-         {
-            break;
-         }
-         image_data->num_planes++;
-      }
    }
    else
    {
@@ -436,63 +347,20 @@ VkResult swapchain::allocate_image(VkImageCreateInfo &image_create_info, wayland
       }
 
       wsialloc_format allocated_format = { 0 };
-      result = allocate_wsialloc(image_create_info, *image_data, importable_formats, &allocated_format);
+      result = allocate_wsialloc(image_create_info, image_data, importable_formats, &allocated_format);
       if (result != VK_SUCCESS)
       {
          return result;
       }
 
-      for (uint32_t plane = 0; plane < MAX_PLANES; plane++)
-      {
-         if (image_data->buffer_fd[plane] == -1)
-         {
-            break;
-         }
-         else if (image_data->buffer_fd[plane] != image_data->buffer_fd[0])
-         {
-            is_disjoint = true;
-         }
-         image_data->num_planes++;
-      }
-
-      auto &image_layout = m_image_creation_parameters.m_image_layout;
-      if (!image_layout.try_resize(image_data->num_planes))
-      {
-         return VK_ERROR_OUT_OF_HOST_MEMORY;
-      }
-
-      for (uint32_t plane = 0; plane < image_data->num_planes; plane++)
-      {
-         assert(image_data->stride[plane] >= 0);
-         image_layout[plane].offset = image_data->offset[plane];
-         image_layout[plane].rowPitch = static_cast<uint32_t>(image_data->stride[plane]);
-      }
-
-      if (is_disjoint)
-      {
-         image_create_info.flags |= VK_IMAGE_CREATE_DISJOINT_BIT;
-      }
-
-      auto &drm_mod_info = m_image_creation_parameters.m_drm_mod_info;
-
-      drm_mod_info.sType = VK_STRUCTURE_TYPE_IMAGE_DRM_FORMAT_MODIFIER_EXPLICIT_CREATE_INFO_EXT;
-      drm_mod_info.pNext = image_create_info.pNext;
-      drm_mod_info.drmFormatModifier = allocated_format.modifier;
-      drm_mod_info.drmFormatModifierPlaneCount = image_data->num_planes;
-      drm_mod_info.pPlaneLayouts = image_layout.data();
-
-      auto &external_info = m_image_creation_parameters.m_external_info;
-      external_info.sType = VK_STRUCTURE_TYPE_EXTERNAL_MEMORY_IMAGE_CREATE_INFO_KHR;
-      external_info.pNext = &drm_mod_info;
-      external_info.handleTypes = VK_EXTERNAL_MEMORY_HANDLE_TYPE_DMA_BUF_BIT_EXT;
+      TRY(fill_image_create_info(image_create_info, m_image_creation_parameters.m_image_layout,
+                                 m_image_creation_parameters.m_drm_mod_info,
+                                 m_image_creation_parameters.m_external_info, *image_data, allocated_format.modifier));
 
       m_image_create_info = image_create_info;
-      m_image_create_info.pNext = &external_info;
-      m_image_create_info.tiling = VK_IMAGE_TILING_DRM_FORMAT_MODIFIER_EXT;
-
       m_allocated_format = allocated_format;
    }
-   image_data->is_disjoint = is_disjoint;
+
    VkResult result = m_device_data.disp.CreateImage(m_device, &m_image_create_info, get_allocation_callbacks(), image);
    if (result != VK_SUCCESS)
    {
@@ -502,38 +370,18 @@ VkResult swapchain::allocate_image(VkImageCreateInfo &image_create_info, wayland
    return result;
 }
 
-VkResult swapchain::create_and_bind_swapchain_image(VkImageCreateInfo image_create_info, swapchain_image &image)
+VkResult swapchain::create_wl_buffer(const VkImageCreateInfo &image_create_info, swapchain_image &image,
+                                     wayland_image_data *image_data)
 {
-   /* Create image_data */
-   auto image_data = m_allocator.create<wayland_image_data>(1);
-   if (image_data == nullptr)
-   {
-      return VK_ERROR_OUT_OF_HOST_MEMORY;
-   }
-
-   std::unique_lock<std::recursive_mutex> image_status_lock(m_image_status_mutex);
-
-   image.data = image_data;
-   image.status = swapchain_image::FREE;
-   VkResult result = allocate_image(image_create_info, image_data, &image.image);
-
-   image_status_lock.unlock();
-
-   if (result != VK_SUCCESS)
-   {
-      WSI_LOG_ERROR("Failed to allocate image.");
-      destroy_image(image);
-      return result;
-   }
-
    /* create a wl_buffer using the dma_buf protocol */
    zwp_linux_buffer_params_v1 *params = zwp_linux_dmabuf_v1_create_params(m_wsi_surface->get_dmabuf_interface());
    uint32_t modifier_hi = m_image_creation_parameters.m_allocated_format.modifier >> 32;
    uint32_t modifier_low = m_image_creation_parameters.m_allocated_format.modifier & 0xFFFFFFFF;
-   for (uint32_t plane = 0; plane < image_data->num_planes; plane++)
+   for (uint32_t plane = 0; plane < image_data->external_mem.get_num_planes(); plane++)
    {
-      zwp_linux_buffer_params_v1_add(params, image_data->buffer_fd[plane], plane,
-                                     image_data->offset[plane], image_data->stride[plane], modifier_hi, modifier_low);
+      zwp_linux_buffer_params_v1_add(params, image_data->external_mem.get_buffer_fds()[plane], plane,
+                                     image_data->external_mem.get_offsets()[plane],
+                                     image_data->external_mem.get_strides()[plane], modifier_hi, modifier_low);
    }
 
    const auto fourcc = util::drm::vk_to_drm_format(image_create_info.format);
@@ -548,34 +396,41 @@ VkResult swapchain::create_and_bind_swapchain_image(VkImageCreateInfo image_crea
       destroy_image(image);
       return VK_ERROR_INITIALIZATION_FAILED;
    }
+   return VK_SUCCESS;
+}
 
-   if (image_data->is_disjoint)
+VkResult swapchain::create_and_bind_swapchain_image(VkImageCreateInfo image_create_info, swapchain_image &image)
+{
+   /* Create image_data */
+   auto image_data = m_allocator.create<wayland_image_data>(1, m_device, m_allocator);
+   if (image_data == nullptr)
    {
-      for (uint32_t plane = 0; plane < image_data->num_planes; plane++)
-      {
-         const auto fd_index = get_same_fd_index(image_data->buffer_fd[plane], image_data->buffer_fd);
-         if (fd_index == plane)
-         {
-            VkResult result = allocate_plane_memory(image_data->buffer_fd[plane], &image_data->memory[fd_index]);
-            if (result != VK_SUCCESS)
-            {
-               return result;
-            }
-         }
-      }
+      return VK_ERROR_OUT_OF_HOST_MEMORY;
    }
-   else
+
+   std::unique_lock<std::recursive_mutex> image_status_lock(m_image_status_mutex);
+   image.data = image_data;
+   image.status = swapchain_image::FREE;
+   VkResult result = allocate_image(image_create_info, image_data, &image.image);
+   image_status_lock.unlock();
+   if (result != VK_SUCCESS)
    {
-      VkResult result = allocate_plane_memory(image_data->buffer_fd[0], &image_data->memory[0]);
-      if (result != VK_SUCCESS)
-      {
-         return result;
-      }
+      WSI_LOG_ERROR("Failed to allocate image.");
+      return result;
+   }
+
+   result = create_wl_buffer(image_create_info, image, image_data);
+   if (result != VK_SUCCESS)
+   {
+      WSI_LOG_ERROR("Failed to create wl_buffer.");
+      return result;
    }
 
-   result = internal_bind_swapchain_image(m_device, image_data, image.image);
+   image_data->external_mem.set_memory_handle_type(VK_EXTERNAL_MEMORY_HANDLE_TYPE_DMA_BUF_BIT_EXT);
+   result = image_data->external_mem.import_memory_and_bind_swapchain_image(image.image);
    if (result != VK_SUCCESS)
    {
+      WSI_LOG_ERROR("Failed to import memory and bind swapchain image.");
       return result;
    }
 
@@ -583,7 +438,6 @@ VkResult swapchain::create_and_bind_swapchain_image(VkImageCreateInfo image_crea
    auto present_fence = sync_fd_fence_sync::create(m_device_data);
    if (!present_fence.has_value())
    {
-      destroy_image(image);
       return VK_ERROR_OUT_OF_HOST_MEMORY;
    }
    image_data->present_fence = std::move(present_fence.value());
@@ -661,23 +515,6 @@ void swapchain::destroy_image(swapchain_image &image)
       {
          wl_buffer_destroy(image_data->buffer);
       }
-
-      for (uint32_t plane = 0; plane < image_data->num_planes; plane++)
-      {
-         if (image_data->memory[plane] != VK_NULL_HANDLE)
-         {
-            m_device_data.disp.FreeMemory(m_device, image_data->memory[plane], get_allocation_callbacks());
-         }
-         else if (image_data->buffer_fd[plane] >= 0)
-         {
-            const auto same_fd_index = get_same_fd_index(image_data->buffer_fd[plane], image_data->buffer_fd);
-            if (same_fd_index == plane)
-            {
-               close(image_data->buffer_fd[plane]);
-            }
-         }
-      }
-
       m_allocator.destroy(1, image_data);
       image.data = nullptr;
    }
@@ -753,50 +590,12 @@ VkResult swapchain::image_wait_present(swapchain_image &, uint64_t)
    return VK_SUCCESS;
 }
 
-VkResult swapchain::internal_bind_swapchain_image(VkDevice &device, wayland_image_data *image_data,
-                                                  const VkImage &image)
-{
-   auto &device_data = layer::device_private_data::get(device);
-   if (image_data->is_disjoint)
-   {
-      util::vector<VkBindImageMemoryInfo> bind_img_mem_infos(m_allocator);
-      if (!bind_img_mem_infos.try_resize(image_data->num_planes))
-      {
-         return VK_ERROR_OUT_OF_HOST_MEMORY;
-      }
-
-      util::vector<VkBindImagePlaneMemoryInfo> bind_plane_mem_infos(m_allocator);
-      if (!bind_plane_mem_infos.try_resize(image_data->num_planes))
-      {
-         return VK_ERROR_OUT_OF_HOST_MEMORY;
-      }
-
-      for (uint32_t plane = 0; plane < image_data->num_planes; plane++)
-      {
-
-         bind_plane_mem_infos[plane].planeAspect = plane_flag_bits[plane];
-         bind_plane_mem_infos[plane].sType = VK_STRUCTURE_TYPE_BIND_IMAGE_PLANE_MEMORY_INFO;
-         bind_plane_mem_infos[plane].pNext = NULL;
-
-         bind_img_mem_infos[plane].sType = VK_STRUCTURE_TYPE_BIND_IMAGE_MEMORY_INFO;
-         bind_img_mem_infos[plane].pNext = &bind_plane_mem_infos[plane];
-         bind_img_mem_infos[plane].image = image;
-         bind_img_mem_infos[plane].memory = image_data->memory[plane];
-         bind_img_mem_infos[plane].memoryOffset = image_data->offset[plane];
-      }
-
-      return device_data.disp.BindImageMemory2KHR(device, bind_img_mem_infos.size(), bind_img_mem_infos.data());
-   }
-
-   return device_data.disp.BindImageMemory(device, image, image_data->memory[0], image_data->offset[0]);
-}
-
 VkResult swapchain::bind_swapchain_image(VkDevice &device, const VkBindImageMemoryInfo *bind_image_mem_info,
                                          const VkBindImageMemorySwapchainInfoKHR *bind_sc_info)
 {
    const wsi::swapchain_image &swapchain_image = m_swapchain_images[bind_sc_info->imageIndex];
    auto image_data = reinterpret_cast<wayland_image_data *>(swapchain_image.data);
-   return internal_bind_swapchain_image(device, image_data, bind_image_mem_info->image);
+   return image_data->external_mem.bind_swapchain_image_memory(bind_image_mem_info->image);
 }
 
 } // namespace wayland
diff --git a/wsi/wayland/swapchain.hpp b/wsi/wayland/swapchain.hpp
index 99d9d7f..5e1c544 100644
--- a/wsi/wayland/swapchain.hpp
+++ b/wsi/wayland/swapchain.hpp
@@ -27,7 +27,8 @@
 #include "wsi/swapchain_base.hpp"
 #include "wl_helpers.hpp"
 
-extern "C" {
+extern "C"
+{
 #include <vulkan/vk_icd.h>
 }
 
@@ -37,12 +38,26 @@ extern "C" {
 #include "util/custom_allocator.hpp"
 #include "wl_object_owner.hpp"
 #include "surface.hpp"
+#include "wsi/external_memory.hpp"
 
 namespace wsi
 {
 namespace wayland
 {
 
+struct wayland_image_data
+{
+   wayland_image_data(const VkDevice &device, const util::allocator &allocator)
+      : external_mem(device, allocator)
+      , buffer(nullptr)
+   {
+   }
+
+   external_memory external_mem;
+   wl_buffer *buffer;
+   sync_fd_fence_sync present_fence;
+};
+
 struct image_creation_parameters
 {
    wsialloc_format m_allocated_format;
@@ -50,8 +65,8 @@ struct image_creation_parameters
    VkExternalMemoryImageCreateInfoKHR m_external_info;
    VkImageDrmFormatModifierExplicitCreateInfoEXT m_drm_mod_info;
 
-   image_creation_parameters(wsialloc_format allocated_format,
-                             util::allocator allocator, VkExternalMemoryImageCreateInfoKHR external_info,
+   image_creation_parameters(wsialloc_format allocated_format, util::allocator allocator,
+                             VkExternalMemoryImageCreateInfoKHR external_info,
                              VkImageDrmFormatModifierExplicitCreateInfoEXT drm_mod_info)
       : m_allocated_format(allocated_format)
       , m_image_layout(allocator)
@@ -142,13 +157,11 @@ protected:
                                  const VkBindImageMemorySwapchainInfoKHR *bind_sc_info) override;
 
 private:
-   struct wayland_image_data;
-
+   VkResult create_wl_buffer(const VkImageCreateInfo &image_create_info, swapchain_image &image,
+                             wayland_image_data *image_data);
    VkResult allocate_image(VkImageCreateInfo &image_create_info, wayland_image_data *image_data, VkImage *image);
-   VkResult allocate_wsialloc(VkImageCreateInfo &image_create_info, wayland_image_data &image_data,
+   VkResult allocate_wsialloc(VkImageCreateInfo &image_create_info, wayland_image_data *image_data,
                               util::vector<wsialloc_format> &importable_formats, wsialloc_format *allocated_format);
-   VkResult internal_bind_swapchain_image(VkDevice &device, wayland_image_data *swapchain_image,
-                                          const VkImage &image);
 
    struct wl_display *m_display;
    struct wl_surface *m_surface;
@@ -169,34 +182,6 @@ private:
     */
    struct image_creation_parameters m_image_creation_parameters;
 
-   /*
-    * @brief Allocate memory for an image plane.
-    *
-    * Allocates a VkDeviceMemory object from a given fd for an image plane. First
-    * it makes a call to get_fd_mem_type_index() to acquire the memory type for
-    * the given fd and then it allocates device memory by calling vkAllocateMemory().
-    *
-    * @param      fd     The plane's fd.
-    * @param[out] memory The allocated VkDeviceMemory object.
-    *
-    * @return VK_SUCCESS on success. If one of the functions that are being called
-    * fails its return value is returned. VK_ERROR_OUT_OF_HOST_MEMORY is returned
-    * when the host gets out of memory.
-    */
-   VkResult allocate_plane_memory(int fd, VkDeviceMemory *memory);
-
-   /*
-    * @brief Get the memory type which the specified file descriptor can be
-    * imported as.
-    *
-    * @param      fd      The given fd.
-    * @param[out] mem_idx The index of the supported memory type.
-    *
-    * @return VK_SUCCESS on success. On failure the error value of
-    * vkGetMemoryFdPropertiesKHR is returned.
-    */
-   VkResult get_fd_mem_type_index(int fd, uint32_t &mem_idx);
-
    /**
     * @brief Finds what formats are compatible with the requested swapchain image Vulkan Device and Wayland surface.
     *
@@ -209,7 +194,6 @@ private:
    VkResult get_surface_compatible_formats(const VkImageCreateInfo &info,
                                            util::vector<wsialloc_format> &importable_formats,
                                            util::vector<uint64_t> &exportable_modifers);
-
 };
 } // namespace wayland
 } // namespace wsi