[ET-VK] Save and load VkPipelineCache file if path is specified

Pull Request resolved: #3546

## Context
Pipeline creation involves the compilation of shader SPIR-V code into machine-specific code. This makes the application's first model-load via the `Program::load_method()` ET-API very slow, due to the creation of compute pipelines via the `vkCreateComputePipelines()` VK-API. To amortize that cost, Vulkan offers a [Compute Pipeline Cache API](https://docs.vulkan.org/guide/latest/pipeline_cache.html). Following [this Vulkan example](https://github.com/KhronosGroup/Vulkan-Samples/tree/main/samples/performance/pipeline_cache), we can (A) retrieve the compiled machine-specific code saving it to a file and (B) load it to a file next time. For an internal model executing on a resource-constrained device, this improves model-load time from ~1200ms to ~500ms.

## This change
We implement both (A)+(B) ET-VK logic. Note that these changes are actually no-op unless you initialize the `pipeline_cache_file_path` manually. The expectation is for the client application to specify the file path of their pipeline cache data if they want to leverage this optimization. In a future ET-wide change, we will expose the file_path config parameter to the ET-API.
ghstack-source-id: 225755565

Differential Revision: [D57085276](https://our.internmc.facebook.com/intern/diff/D57085276/)

Author: Jorge Pineda

backends/vulkan/runtime/api/Adapter.cpp

@@ -292,7 +292,8 @@ DeviceHandle::~DeviceHandle() {
 Adapter::Adapter(
     VkInstance instance,
     PhysicalDevice physical_device,
-    const uint32_t num_queues)
+    const uint32_t num_queues,
+    const std::string& pipeline_cache_file_path)
     : queue_usage_mutex_{},
       physical_device_(std::move(physical_device)),
       queues_{},
@@ -307,7 +308,7 @@ Adapter::Adapter(
       shader_layout_cache_(device_.handle_),
       shader_cache_(device_.handle_),
       pipeline_layout_cache_(device_.handle_),
-      compute_pipeline_cache_(device_.handle_),
+      compute_pipeline_cache_(device_.handle_, pipeline_cache_file_path),
       sampler_cache_(device_.handle_),
       vma_(instance_, physical_device_.handle, device_.handle_) {}
 

backends/vulkan/runtime/api/Adapter.h

@@ -101,7 +101,8 @@ class Adapter final {
   explicit Adapter(
       VkInstance instance,
       PhysicalDevice physical_device,
-      const uint32_t num_queues);
+      const uint32_t num_queues,
+      const std::string& pipeline_cache_file_path);
 
   Adapter(const Adapter&) = delete;
   Adapter& operator=(const Adapter&) = delete;

backends/vulkan/runtime/api/Pipeline.cpp

@@ -8,6 +8,8 @@
 
 #include <executorch/backends/vulkan/runtime/api/Pipeline.h>
 
+#include <fstream>
+
 namespace vkcompute {
 namespace api {
 
@@ -358,17 +360,24 @@ void PipelineLayoutCache::purge() {
 // ComputePipelineCache
 //
 
-ComputePipelineCache::ComputePipelineCache(VkDevice device)
+ComputePipelineCache::ComputePipelineCache(
+    VkDevice device,
+    const std::string& file_path)
     : cache_mutex_{},
       device_(device),
       pipeline_cache_{VK_NULL_HANDLE},
-      cache_{} {
-  const VkPipelineCacheCreateInfo pipeline_cache_create_info{
+      cache_{},
+      file_path_(file_path) {
+  VkPipelineCacheCreateInfo pipeline_cache_create_info{};
+
+  auto buffer = load_cache();
+
+  pipeline_cache_create_info = {
       VK_STRUCTURE_TYPE_PIPELINE_CACHE_CREATE_INFO, // sType
       nullptr, // pNext
       0u, // flags
-      0u, // initialDataSize
-      nullptr, // pInitialData
+      buffer.size(), // initialDataSize
+      buffer.data(), // pInitialData
   };
 
   VK_CHECK(vkCreatePipelineCache(
@@ -392,6 +401,9 @@ ComputePipelineCache::~ComputePipelineCache() {
   if (VK_NULL_HANDLE == pipeline_cache_) {
     return;
   }
+
+  save_cache();
+
   vkDestroyPipelineCache(device_, pipeline_cache_, nullptr);
   pipeline_cache_ = VK_NULL_HANDLE;
 }
@@ -416,5 +428,37 @@ void ComputePipelineCache::purge() {
   cache_.clear();
 }
 
+std::vector<char> ComputePipelineCache::load_cache() {
+  // Return if path is not specified; this means the optimization is disabled
+  if (file_path_.empty()) {
+    return {};
+  }
+
+  // Return if file doesn't exist; this is expected on the first model-load
+  std::ifstream file(file_path_, std::ios::binary | std::ios::ate);
+  if (file.fail()) {
+    return {};
+  }
+
+  auto size = file.tellg();
+  file.seekg(0, std::ios::beg);
+
+  std::vector<char> buffer(size);
+  file.read(buffer.data(), size);
+
+  return buffer;
+}
+
+void ComputePipelineCache::save_cache() {
+  size_t size{};
+  vkGetPipelineCacheData(device_, pipeline_cache_, &size, nullptr);
+
+  std::vector<char> buffer(size);
+  vkGetPipelineCacheData(device_, pipeline_cache_, &size, buffer.data());
+
+  std::ofstream file(file_path_, std::ios::binary);
+  file.write(buffer.data(), buffer.size());
+}
+
 } // namespace api
 } // namespace vkcompute

backends/vulkan/runtime/api/Pipeline.h

@@ -216,7 +216,7 @@ class PipelineLayoutCache final {
 
 class ComputePipelineCache final {
  public:
-  explicit ComputePipelineCache(VkDevice device);
+  explicit ComputePipelineCache(VkDevice device, const std::string& file_path);
 
   ComputePipelineCache(const ComputePipelineCache&) = delete;
   ComputePipelineCache& operator=(const ComputePipelineCache&) = delete;
@@ -266,13 +266,17 @@ class ComputePipelineCache final {
   };
 
  private:
+  std::vector<char> load_cache();
+  void save_cache();
+
   // Multiple threads could potentially be adding entries into the cache, so use
   // a mutex to manage access
   std::mutex cache_mutex_;
 
   VkDevice device_;
   VkPipelineCache pipeline_cache_;
   std::unordered_map<Key, Value, Hasher> cache_;
+  const std::string file_path_;
 
  public:
   VkPipeline retrieve(const Key&);

backends/vulkan/runtime/api/Runtime.cpp

@@ -253,12 +253,14 @@ std::unique_ptr<Runtime> init_global_vulkan_runtime() {
 #endif /* VULKAN_DEBUG */
   const bool init_default_device = true;
   const uint32_t num_requested_queues = 1; // TODO: raise this value
+  const std::string pipeline_cache_file_path = ""; // TODO: expose to client
 
   const RuntimeConfiguration default_config{
       enable_validation_messages,
       init_default_device,
       AdapterSelector::First,
       num_requested_queues,
+      pipeline_cache_file_path,
   };
 
   try {
@@ -351,7 +353,10 @@ uint32_t Runtime::create_adapter(const Selector& selector) {
   // Otherwise, create an adapter for the selected physical device
   adapter_i = utils::safe_downcast<int32_t>(adapters_.size());
   adapters_.emplace_back(new Adapter(
-      instance_, device_mapping.first, config_.num_requested_queues));
+      instance_,
+      device_mapping.first,
+      config_.num_requested_queues,
+      config_.pipeline_cache_file_path));
   device_mapping.second = adapter_i;
 
   return adapter_i;

backends/vulkan/runtime/api/Runtime.h

@@ -39,6 +39,7 @@ struct RuntimeConfiguration final {
   bool init_default_device;
   AdapterSelector default_selector;
   uint32_t num_requested_queues;
+  std::string pipeline_cache_file_path;
 };
 
 class Runtime final {