Make sure we can set binary artifacts to cache mfa related shaders.

Author: liuliu

lib/nnc/ccv_nnc.h

@@ -929,6 +929,15 @@ CCV_WARN_UNUSED(int) ccv_nnc_queue_watermark(void);
  * @param size The size of the array, only first 64 will be used.
  */
 void ccv_nnc_set_device_permutation(const int type, const int* const device_map, const int size);
+/**
+ * Set the path to binary artifacts that would accelerate command compilations. Note that the binary
+ * artifacts path are separated into read one and write one. They could be the same, but would be
+ * better to be a separate one to avoid competing with each other.
+ * @param paths_to_read The file paths to read binary artifacts. Whether it is a file or directory is implementation dependent.
+ * @param paths_to_read_size How many paths in 1.
+ * @param path_to_write The file path to write binary artifacts. Whether it is a file or directory is implementation dependent.
+ */
+void ccv_nnc_set_binary_artifacts(const char** const paths_to_read, const int paths_to_read_size, const char* const path_to_write);
 /**
  * Quantize a given memory region of a given datatype / memory resides, into nbits palette.
  * @param input The input memory region, it can be CCV_64F, CCV_32F or CCV_16F.

lib/nnc/ccv_nnc_cmd.c

@@ -964,3 +964,10 @@ void ccv_nnc_set_device_permutation(const int type, const int* const device_map,
 	cusetdevicemap(device_map, size);
 #endif
 }
+
+void ccv_nnc_set_binary_artifacts(const char** const paths_to_read, const int paths_to_read_size, const char* const path_to_write)
+{
+#ifdef HAVE_MPS
+	ccv_nnc_mps_set_binary_artifacts(paths_to_read, paths_to_read_size, path_to_write);
+#endif
+}

lib/nnc/mfa/ccv_nnc_mfa.cpp

@@ -55,6 +55,15 @@ mtl_buffer_t* ccv_nnc_mfa_request_scratch(ccv_nnc_mfa_context_t* context, const
   return context->request_scratch(size);
 }
 
+void ccv_nnc_mfa_set_binary_archives(ccv_nnc_mfa_context_t* const context, const char** const paths_to_read, const int paths_to_read_size, const char* const path_to_write) {
+  std::vector<std::string> paths_to_read_vec;
+  for (int i = 0; i < paths_to_read_size; i++) {
+    paths_to_read_vec.push_back(std::string(paths_to_read[i]));
+  }
+  std::string path_to_write_str = path_to_write != nullptr ? std::string(path_to_write) : std::string();
+  context->v2_cache.setBinaryArchives(context->device.get(), paths_to_read_vec, path_to_write_str);
+}
+
 void ccv_nnc_mfa_log_message(const char* message) {
   std::cerr << METAL_LOG_HEADER << message << std::endl;
 }

lib/nnc/mfa/ccv_nnc_mfa.hpp

@@ -65,6 +65,8 @@ class context {
 } // namespace nnc
 } // namespace ccv
 
+std::pair<std::string, std::string> ccv_nnc_mfa_get_binary_artifacts(void);
+
 extern "C" {
 #endif // __cplusplus
 
@@ -80,6 +82,7 @@ void ccv_nnc_mfa_log_message(const char* message);
 mtl_command_batch_t* ccv_nnc_start_command_batch(mtl_command_queue_t* command_queue);
 void ccv_nnc_finish_command_batch(mtl_command_batch_t* command_batch);
 mtl_buffer_t* ccv_nnc_mfa_request_scratch(ccv_nnc_mfa_context_t* context, const uint64_t size);
+void ccv_nnc_mfa_set_binary_archives(ccv_nnc_mfa_context_t* context, const char** paths_to_read, const int paths_to_read_size, const char* path_to_write);
 
 #ifdef __cplusplus
 } // extern "C"

lib/nnc/mfa/v2/AddDescriptor.cpp

@@ -19,7 +19,7 @@ std::size_t std::hash<AddDescriptor>::operator()(const AddDescriptor& hash) cons
   return seed;
 }
 
-std::pair<AddKernelDescriptor, PipelineValue<AddKernel> *> AddDescriptor::findKernel(MTL::Device *const device, const DeviceProperties &dprops, std::unordered_map<AddKernelDescriptor, std::unique_ptr<AddKernel>> *const libraryCache) const noexcept {
+std::pair<AddKernelDescriptor, PipelineValue<AddKernel> *> AddDescriptor::findKernel(MTL::Device *const device, const DeviceProperties &dprops, NS::Array* const binaryArchivesToRead, MTL::BinaryArchive* const binaryArchiveToWrite, const std::string& pathToWrite, std::unordered_map<AddKernelDescriptor, std::unique_ptr<AddKernel>> *const libraryCache) const noexcept {
   // The caller is not responsible for calling 'delete' on this pointer. The
   // reference is saved in the 'libraryCache'. It will be deallocated whenever
   // the shader cache itself is cleaned up.

lib/nnc/mfa/v2/AddDescriptor.hpp

@@ -33,7 +33,7 @@ struct AddDescriptor {
 
   bool operator==(const AddDescriptor& rhs) const;
 
-  std::pair<AddKernelDescriptor, PipelineValue<AddKernel> *> findKernel(MTL::Device* const device, const DeviceProperties &dprops, std::unordered_map<AddKernelDescriptor, std::unique_ptr<AddKernel>> *const libraryCache) const noexcept;
+  std::pair<AddKernelDescriptor, PipelineValue<AddKernel> *> findKernel(MTL::Device* const device, const DeviceProperties &dprops, NS::Array* const binaryArchivesToRead, MTL::BinaryArchive* const binaryArchiveToWrite, const std::string& pathToWrite, std::unordered_map<AddKernelDescriptor, std::unique_ptr<AddKernel>> *const libraryCache) const noexcept;
 };
 
 template<>

lib/nnc/mfa/v2/AttentionDescriptor.cpp

@@ -127,7 +127,7 @@ AttentionKernelDescriptor AttentionDescriptor::kernelDescriptor(MTL::Device *con
   }
 }
 
-std::pair<AttentionKernelDescriptor, PipelineValue<AttentionKernel> *> AttentionDescriptor::findKernel(MTL::Device *const device, const DeviceProperties &dprops, std::unordered_map<AttentionKernelDescriptor, std::unique_ptr<AttentionKernel>> *const libraryCache) const noexcept {
+std::pair<AttentionKernelDescriptor, PipelineValue<AttentionKernel> *> AttentionDescriptor::findKernel(MTL::Device *const device, const DeviceProperties &dprops, NS::Array* const binaryArchivesToRead, MTL::BinaryArchive* const binaryArchiveToWrite, const std::string& pathToWrite, std::unordered_map<AttentionKernelDescriptor, std::unique_ptr<AttentionKernel>> *const libraryCache) const noexcept {
   auto createPipeline =
   [=](MTL::Library* library) -> MTL::ComputePipelineState* {
     // Set the function constants.

lib/nnc/mfa/v2/AttentionDescriptor.hpp

@@ -61,7 +61,7 @@ struct AttentionDescriptor {
 
   bool operator==(const AttentionDescriptor& rhs) const;
 
-  std::pair<AttentionKernelDescriptor, PipelineValue<AttentionKernel> *> findKernel(MTL::Device* const device, const DeviceProperties &dprops, std::unordered_map<AttentionKernelDescriptor, std::unique_ptr<AttentionKernel>> *const libraryCache) const noexcept;
+  std::pair<AttentionKernelDescriptor, PipelineValue<AttentionKernel> *> findKernel(MTL::Device* const device, const DeviceProperties &dprops, NS::Array* const binaryArchivesToRead, MTL::BinaryArchive* const binaryArchiveToWrite, const std::string& pathToWrite, std::unordered_map<AttentionKernelDescriptor, std::unique_ptr<AttentionKernel>> *const libraryCache) const noexcept;
 
 private:
   AttentionKernelDescriptor kernelDescriptor(MTL::Device *const device, const DeviceProperties &dprops) const noexcept;

lib/nnc/mfa/v2/CMulDescriptor.cpp

@@ -27,7 +27,7 @@ std::size_t std::hash<CMulDescriptor>::operator()(const CMulDescriptor& hash) co
   return seed;
 }
 
-std::pair<CMulKernelDescriptor, PipelineValue<CMulKernel> *> CMulDescriptor::findKernel(MTL::Device *const device, const DeviceProperties &dprops, std::unordered_map<CMulKernelDescriptor, std::unique_ptr<CMulKernel>> *const libraryCache) const noexcept {
+std::pair<CMulKernelDescriptor, PipelineValue<CMulKernel> *> CMulDescriptor::findKernel(MTL::Device *const device, const DeviceProperties &dprops, NS::Array* const binaryArchivesToRead, MTL::BinaryArchive* const binaryArchiveToWrite, const std::string& pathToWrite, std::unordered_map<CMulKernelDescriptor, std::unique_ptr<CMulKernel>> *const libraryCache) const noexcept {
   // The caller is not responsible for calling 'delete' on this pointer. The
   // reference is saved in the 'libraryCache'. It will be deallocated whenever
   // the shader cache itself is cleaned up.

lib/nnc/mfa/v2/CMulDescriptor.hpp

@@ -39,7 +39,7 @@ struct CMulDescriptor {
 
   bool operator==(const CMulDescriptor& rhs) const;
 
-  std::pair<CMulKernelDescriptor, PipelineValue<CMulKernel> *> findKernel(MTL::Device* const device, const DeviceProperties &dprops, std::unordered_map<CMulKernelDescriptor, std::unique_ptr<CMulKernel>> *const libraryCache) const noexcept;
+  std::pair<CMulKernelDescriptor, PipelineValue<CMulKernel> *> findKernel(MTL::Device* const device, const DeviceProperties &dprops, NS::Array* const binaryArchivesToRead, MTL::BinaryArchive* const binaryArchiveToWrite, const std::string& pathToWrite, std::unordered_map<CMulKernelDescriptor, std::unique_ptr<CMulKernel>> *const libraryCache) const noexcept;
 };
 
 template<>