Ability to set external instance + devices (#11393)

Summary:

Adds a new extension to ETVK, set_and_get_external_adapter, which allows clients to share their Vulkan instance/devices with ETVK.  This is useful when using volk, which does not support multiple devices when using volkLoadDevice.

Differential Revision: D71372344

Author: SamGondelman

backends/vulkan/runtime/VulkanBackend.cpp

@@ -13,6 +13,8 @@
 
 #include <executorch/backends/vulkan/runtime/graph/ops/OperatorRegistry.h>
 
+#include <executorch/extension/vulkan/external_adapter_wrapper.h>
+
 #include <executorch/runtime/backend/interface.h>
 #include <executorch/runtime/core/error.h>
 #include <executorch/runtime/core/evalue.h>
@@ -517,7 +519,9 @@ class VulkanBackend final : public ::executorch::runtime::BackendInterface {
       return Error::MemoryAllocationFailed;
     }
 
-    new (compute_graph) ComputeGraph(get_graph_config(compile_specs));
+    GraphConfig graph_config = get_graph_config(compile_specs);
+    graph_config.external_adapter = extension::set_and_get_external_adapter();
+    new (compute_graph) ComputeGraph(graph_config);
 
     Error err = compileModel(processed->data(), compute_graph);
 

backends/vulkan/runtime/api/Context.cpp

@@ -24,10 +24,10 @@
 namespace vkcompute {
 namespace api {
 
-Context::Context(size_t adapter_i, const ContextConfig& config)
+Context::Context(vkapi::Adapter* adapter, const ContextConfig& config)
     : config_(config),
       // Important handles
-      adapter_p_(vkapi::runtime()->get_adapter_p(adapter_i)),
+      adapter_p_(adapter),
       device_(adapter_p_->device_handle()),
       queue_(adapter_p_->request_queue()),
       // Resource pools
@@ -256,7 +256,7 @@ Context* context() {
           query_pool_config,
       };
 
-      return new Context(vkapi::runtime()->default_adapter_i(), config);
+      return new Context(vkapi::runtime()->get_adapter_p(), config);
     } catch (...) {
     }
 

backends/vulkan/runtime/api/Context.h

@@ -42,7 +42,7 @@ struct ContextConfig final {
 
 class Context final {
  public:
-  explicit Context(size_t adapter_i, const ContextConfig&);
+  explicit Context(vkapi::Adapter*, const ContextConfig&);
 
   Context(const Context&) = delete;
   Context& operator=(const Context&) = delete;

backends/vulkan/runtime/graph/ComputeGraph.cpp

@@ -122,7 +122,8 @@ ComputeGraph::ComputeGraph(GraphConfig config)
       prepack_descriptor_counts_{},
       execute_descriptor_counts_{},
       context_{new api::Context(
-          vkapi::runtime()->default_adapter_i(),
+          config.external_adapter ? config.external_adapter
+            : vkapi::runtime()->get_adapter_p(),
           config_.context_config)},
       shared_objects_{},
       values_{},

backends/vulkan/runtime/graph/GraphConfig.cpp

@@ -63,6 +63,8 @@ GraphConfig::GraphConfig() {
 
   enable_local_wg_size_override = false;
   local_wg_size_override = {};
+
+  external_adapter = nullptr;
 }
 
 void GraphConfig::set_storage_type_override(utils::StorageType storage_type) {

backends/vulkan/runtime/graph/GraphConfig.h

@@ -33,6 +33,8 @@ struct GraphConfig final {
   bool enable_local_wg_size_override;
   utils::uvec3 local_wg_size_override;
 
+  vkapi::Adapter* external_adapter;
+
   // Generate a default graph config with pre-configured settings
   explicit GraphConfig();
 

backends/vulkan/runtime/vk_api/Adapter.cpp

@@ -17,17 +17,12 @@ namespace vkapi {
 
 namespace {
 
-VkDevice create_logical_device(
+void find_compute_queues(
     const PhysicalDevice& physical_device,
     const uint32_t num_queues_to_create,
-    std::vector<Adapter::Queue>& queues,
-    std::vector<uint32_t>& queue_usage) {
-  // Find compute queues up to the requested number of queues
-
-  std::vector<VkDeviceQueueCreateInfo> queue_create_infos;
+    std::vector<VkDeviceQueueCreateInfo>& queue_create_infos,
+    std::vector<std::pair<uint32_t, uint32_t>>& queues_to_get) {
   queue_create_infos.reserve(num_queues_to_create);
-
-  std::vector<std::pair<uint32_t, uint32_t>> queues_to_get;
   queues_to_get.reserve(num_queues_to_create);
 
   uint32_t remaining_queues = num_queues_to_create;
@@ -60,12 +55,44 @@ VkDevice create_logical_device(
       break;
     }
   }
+}
 
+void populate_queue_info(
+    const PhysicalDevice& physical_device,
+    VkDevice logical_device,
+    const std::vector<std::pair<uint32_t, uint32_t>>& queues_to_get,
+    std::vector<Adapter::Queue>& queues,
+    std::vector<uint32_t>& queue_usage) {
   queues.reserve(queues_to_get.size());
   queue_usage.reserve(queues_to_get.size());
 
-  // Create the VkDevice
+  // Obtain handles for the created queues and initialize queue usage heuristic
+
+  for (const std::pair<uint32_t, uint32_t>& queue_idx : queues_to_get) {
+    VkQueue queue_handle = VK_NULL_HANDLE;
+    VkQueueFlags flags =
+        physical_device.queue_families.at(queue_idx.first).queueFlags;
+    vkGetDeviceQueue(
+      logical_device, queue_idx.first, queue_idx.second, &queue_handle);
+    queues.push_back({queue_idx.first, queue_idx.second, flags, queue_handle});
+    // Initial usage value
+    queue_usage.push_back(0);
+  }
+}
+
+VkDevice create_logical_device(
+    const PhysicalDevice& physical_device,
+    const uint32_t num_queues_to_create,
+    std::vector<Adapter::Queue>& queues,
+    std::vector<uint32_t>& queue_usage) {
+  // Find compute queues up to the requested number of queues
 
+  std::vector<VkDeviceQueueCreateInfo> queue_create_infos;
+  std::vector<std::pair<uint32_t, uint32_t>> queues_to_get;
+  find_compute_queues(
+    physical_device, num_queues_to_create, queue_create_infos, queues_to_get);
+
+  // Create the VkDevice
   std::vector<const char*> requested_device_extensions{
 #ifdef VK_KHR_portability_subset
       VK_KHR_PORTABILITY_SUBSET_EXTENSION_NAME,
@@ -143,19 +170,42 @@ VkDevice create_logical_device(
   volkLoadDevice(handle);
 #endif /* USE_VULKAN_VOLK */
 
-  // Obtain handles for the created queues and initialize queue usage heuristic
+  populate_queue_info(
+    physical_device, handle, queues_to_get, queues, queue_usage);
 
-  for (const std::pair<uint32_t, uint32_t>& queue_idx : queues_to_get) {
-    VkQueue queue_handle = VK_NULL_HANDLE;
-    VkQueueFlags flags =
-        physical_device.queue_families.at(queue_idx.first).queueFlags;
-    vkGetDeviceQueue(handle, queue_idx.first, queue_idx.second, &queue_handle);
-    queues.push_back({queue_idx.first, queue_idx.second, flags, queue_handle});
-    // Initial usage value
-    queue_usage.push_back(0);
+  return handle;
+}
+
+bool test_linear_tiling_3d_image_support(VkDevice device) {
+  // Test creating a 3D image with linear tiling to see if it is supported.
+  // According to the Vulkan spec, linear tiling may not be supported for 3D
+  // images.
+  VkExtent3D image_extents{1u, 1u, 1u};
+  const VkImageCreateInfo image_create_info{
+      VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO, // sType
+      nullptr, // pNext
+      0u, // flags
+      VK_IMAGE_TYPE_3D, // imageType
+      VK_FORMAT_R32G32B32A32_SFLOAT, // format
+      image_extents, // extents
+      1u, // mipLevels
+      1u, // arrayLayers
+      VK_SAMPLE_COUNT_1_BIT, // samples
+      VK_IMAGE_TILING_LINEAR, // tiling
+      VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_STORAGE_BIT, // usage
+      VK_SHARING_MODE_EXCLUSIVE, // sharingMode
+      0u, // queueFamilyIndexCount
+      nullptr, // pQueueFamilyIndices
+      VK_IMAGE_LAYOUT_UNDEFINED, // initialLayout
+  };
+  VkImage image = VK_NULL_HANDLE;
+  VkResult res = vkCreateImage(device, &image_create_info, nullptr, &image);
+
+  if (res == VK_SUCCESS) {
+    vkDestroyImage(device, image, nullptr);
   }
 
-  return handle;
+  return res == VK_SUCCESS;
 }
 
 } // namespace
@@ -186,37 +236,44 @@ Adapter::Adapter(
       compute_pipeline_cache_(device_.handle, cache_data_path),
       sampler_cache_(device_.handle),
       vma_(instance_, physical_device_.handle, device_.handle),
-      linear_tiling_3d_enabled_{true} {
-  // Test creating a 3D image with linear tiling to see if it is supported.
-  // According to the Vulkan spec, linear tiling may not be supported for 3D
-  // images.
-  VkExtent3D image_extents{1u, 1u, 1u};
-  const VkImageCreateInfo image_create_info{
-      VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO, // sType
-      nullptr, // pNext
-      0u, // flags
-      VK_IMAGE_TYPE_3D, // imageType
-      VK_FORMAT_R32G32B32A32_SFLOAT, // format
-      image_extents, // extents
-      1u, // mipLevels
-      1u, // arrayLayers
-      VK_SAMPLE_COUNT_1_BIT, // samples
-      VK_IMAGE_TILING_LINEAR, // tiling
-      VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_STORAGE_BIT, // usage
-      VK_SHARING_MODE_EXCLUSIVE, // sharingMode
-      0u, // queueFamilyIndexCount
-      nullptr, // pQueueFamilyIndices
-      VK_IMAGE_LAYOUT_UNDEFINED, // initialLayout
-  };
-  VkImage image = VK_NULL_HANDLE;
-  VkResult res =
-      vkCreateImage(device_.handle, &image_create_info, nullptr, &image);
-  if (res != VK_SUCCESS) {
-    linear_tiling_3d_enabled_ = false;
-  } else {
-    vkDestroyImage(device_.handle, image, nullptr);
+      linear_tiling_3d_enabled_{
+        test_linear_tiling_3d_image_support(device_.handle)},
+      owns_device_{true} {}
+
+Adapter::Adapter(
+    VkInstance instance,
+    VkPhysicalDevice physical_device,
+    VkDevice logical_device,
+    const uint32_t num_queues,
+    const std::string& cache_data_path)
+    : queue_usage_mutex_{},
+      physical_device_(physical_device),
+      queues_{},
+      queue_usage_{},
+      queue_mutexes_{},
+      instance_(instance),
+      device_(logical_device),
+      shader_layout_cache_(device_.handle),
+      shader_cache_(device_.handle),
+      pipeline_layout_cache_(device_.handle),
+      compute_pipeline_cache_(device_.handle, cache_data_path),
+      sampler_cache_(device_.handle),
+      vma_(instance_, physical_device_.handle, device_.handle),
+      linear_tiling_3d_enabled_{
+        test_linear_tiling_3d_image_support(device_.handle)},
+      owns_device_{false} {
+  std::vector<VkDeviceQueueCreateInfo> queue_create_infos;
+  std::vector<std::pair<uint32_t, uint32_t>> queues_to_get;
+  find_compute_queues(
+    physical_device_, num_queues, queue_create_infos, queues_to_get);
+  populate_queue_info(
+    physical_device_, device_.handle, queues_to_get, queues_, queue_usage_);
+}
+
+Adapter::~Adapter() {
+  if (!owns_device_) {
+    device_.handle = VK_NULL_HANDLE;
   }
-  return;
 }
 
 Adapter::Queue Adapter::request_queue() {

backends/vulkan/runtime/vk_api/Adapter.h

@@ -56,13 +56,20 @@ class Adapter final {
       const uint32_t num_queues,
       const std::string& cache_data_path);
 
+  explicit Adapter(
+      VkInstance instance,
+      VkPhysicalDevice physical_device,
+      VkDevice logical_device,
+      const uint32_t num_queues,
+      const std::string& cache_data_path);
+
   Adapter(const Adapter&) = delete;
   Adapter& operator=(const Adapter&) = delete;
 
   Adapter(Adapter&&) = delete;
   Adapter& operator=(Adapter&&) = delete;
 
-  ~Adapter() = default;
+  ~Adapter();
 
   struct Queue {
     uint32_t family_index;
@@ -94,6 +101,7 @@ class Adapter final {
   Allocator vma_;
   // Miscellaneous
   bool linear_tiling_3d_enabled_;
+  bool owns_device_;
 
  public:
   // Physical Device metadata

backends/vulkan/runtime/vk_api/Runtime.cpp

@@ -14,6 +14,12 @@
 #include <iostream>
 #include <sstream>
 
+#ifdef USE_VOLK_HEADER_ONLY
+// For volk.h, define this before including volk.h in exactly one CPP file.
+#define VOLK_IMPLEMENTATION
+#include <volk.h>
+#endif /* USE_VOLK_HEADER_ONLY */
+
 namespace vkcompute {
 namespace vkapi {
 

backends/vulkan/runtime/vk_api/Types.h

@@ -17,6 +17,10 @@
 #include <cstddef>
 #include <cstdint>
 
+#if defined(__linux__)
+#undef Bool
+#endif
+
 #ifdef USE_VULKAN_FP16_INFERENCE
 #define VK_FORMAT_FLOAT4 VK_FORMAT_R16G16B16A16_SFLOAT
 #else