Support passing xnnpack weight cache by file descriptor in LiteRT's CpuOptions.

LiteRT-PiperOrigin-RevId: 818880466

Author: ai-edge-bot

litert/c/options/litert_cpu_options.cc

@@ -91,6 +91,9 @@ LiteRtStatus LiteRtSetCpuOptionsXnnPackWeightCachePath(LiteRtCpuOptions options,
                                                        const char* path) {
   LITERT_RETURN_IF_ERROR(options, litert::ErrorStatusBuilder::InvalidArgument())
       << "options is null.";
+  LITERT_RETURN_IF_ERROR(options->xnn.weight_cache_file_descriptor <= 0,
+                         litert::ErrorStatusBuilder::InvalidArgument())
+      << "weight cache file descriptor and path cannot both be set.";
   options->xnn.weight_cache_file_path = path;
   return kLiteRtStatusOk;
 }
@@ -104,3 +107,24 @@ LiteRtStatus LiteRtGetCpuOptionsXnnPackWeightCachePath(
   *path = options->xnn.weight_cache_file_path;
   return kLiteRtStatusOk;
 }
+
+LiteRtStatus LiteRtSetCpuOptionsXnnPackWeightCacheFileDescriptor(
+    LiteRtCpuOptions options, int fd) {
+  LITERT_RETURN_IF_ERROR(options, litert::ErrorStatusBuilder::InvalidArgument())
+      << "options is null.";
+  LITERT_RETURN_IF_ERROR(options->xnn.weight_cache_file_path == nullptr,
+                         litert::ErrorStatusBuilder::InvalidArgument())
+      << "weight cache file descriptor and path cannot both be set.";
+  options->xnn.weight_cache_file_descriptor = fd;
+  return kLiteRtStatusOk;
+}
+
+LiteRtStatus LiteRtGetCpuOptionsXnnPackWeightCacheFileDescriptor(
+    LiteRtCpuOptionsConst options, int* const fd) {
+  LITERT_RETURN_IF_ERROR(options, litert::ErrorStatusBuilder::InvalidArgument())
+      << "options is null.";
+  LITERT_RETURN_IF_ERROR(fd, litert::ErrorStatusBuilder::InvalidArgument())
+      << "fd is null.";
+  *fd = options->xnn.weight_cache_file_descriptor;
+  return kLiteRtStatusOk;
+}

litert/c/options/litert_cpu_options.h

@@ -53,6 +53,7 @@ LiteRtStatus LiteRtGetCpuOptionsXNNPackFlags(LiteRtCpuOptionsConst options,
                                              uint32_t* flags);
 
 // Sets the XNNPack weight cache file path used by the CPU accelerator.
+// Weight cache file path and descriptor must not both be set.
 // The `path` string is owned by the caller and must outlive the `options`
 // object.
 LiteRtStatus LiteRtSetCpuOptionsXnnPackWeightCachePath(LiteRtCpuOptions options,
@@ -64,6 +65,16 @@ LiteRtStatus LiteRtSetCpuOptionsXnnPackWeightCachePath(LiteRtCpuOptions options,
 LiteRtStatus LiteRtGetCpuOptionsXnnPackWeightCachePath(
     LiteRtCpuOptionsConst options, const char** path);
 
+// Sets the XNNPack weight cache file descriptor used by the CPU accelerator.
+// Weight cache file path and descriptor must not both be set.
+LiteRtStatus LiteRtSetCpuOptionsXnnPackWeightCacheFileDescriptor(
+    LiteRtCpuOptions options, int fd);
+
+// Gets the XNNPack weight cache file descriptor used by the CPU accelerator.
+LiteRtStatus LiteRtGetCpuOptionsXnnPackWeightCacheFileDescriptor(
+    LiteRtCpuOptionsConst options, int* fd);
+
+
 #ifdef __cplusplus
 }  // extern "C"
 #endif

litert/c/options/litert_cpu_options_test.cc

@@ -207,4 +207,33 @@ TEST_F(LiteRtCpuOptionsFieldsTest,
               IsError(kLiteRtStatusErrorInvalidArgument));
 }
 
+TEST_F(LiteRtCpuOptionsFieldsTest, SetAndGetXNNPackWeightCacheDescriptor) {
+  const int expected_fd = 1234;
+  int fd;
+
+  // Avoid a no-op test.
+  LITERT_EXPECT_OK(
+      LiteRtGetCpuOptionsXnnPackWeightCacheFileDescriptor(cpu_options_, &fd));
+  ASSERT_NE(fd, expected_fd);
+
+  // Actual test.
+  LITERT_EXPECT_OK(LiteRtSetCpuOptionsXnnPackWeightCacheFileDescriptor(
+      cpu_options_, expected_fd));
+  LITERT_EXPECT_OK(
+      LiteRtGetCpuOptionsXnnPackWeightCacheFileDescriptor(cpu_options_, &fd));
+  ASSERT_EQ(fd, expected_fd);
+}
+
+TEST_F(LiteRtCpuOptionsFieldsTest,
+       SetXNNPackWeightCacheFailsIfBothPathAndDescriptorAreSet) {
+  const absl::string_view path = "a/path/to/the/cache";
+  const int fd = 1234;
+
+  LITERT_EXPECT_OK(
+      LiteRtSetCpuOptionsXnnPackWeightCachePath(cpu_options_, path.data()));
+  EXPECT_THAT(
+      LiteRtSetCpuOptionsXnnPackWeightCacheFileDescriptor(cpu_options_, fd),
+      IsError(kLiteRtStatusErrorInvalidArgument));
+}
+
 }  // namespace

litert/cc/options/litert_cpu_options.cc

@@ -96,4 +96,21 @@ Expected<absl::string_view> CpuOptions::GetXNNPackWeightCachePath() const {
   return absl::NullSafeStringView(path);
 }
 
+Expected<void> CpuOptions::SetXNNPackWeightCacheFileDescriptor(int fd) {
+  LiteRtCpuOptions cpu_options;
+  LITERT_RETURN_IF_ERROR(LiteRtFindCpuOptions(Get(), &cpu_options));
+  LITERT_RETURN_IF_ERROR(
+      LiteRtSetCpuOptionsXnnPackWeightCacheFileDescriptor(cpu_options, fd));
+  return {};
+}
+
+Expected<int> CpuOptions::GetXNNPackWeightCacheFileDescriptor() const {
+  LiteRtCpuOptions cpu_options;
+  LITERT_RETURN_IF_ERROR(LiteRtFindCpuOptions(Get(), &cpu_options));
+  int fd;
+  LITERT_RETURN_IF_ERROR(
+      LiteRtGetCpuOptionsXnnPackWeightCacheFileDescriptor(cpu_options, &fd));
+  return fd;
+}
+
 }  // namespace litert

litert/cc/options/litert_cpu_options.h

@@ -40,6 +40,9 @@ class CpuOptions : public OpaqueOptions {
 
   Expected<void> SetXNNPackWeightCachePath(const char* path);
   Expected<absl::string_view> GetXNNPackWeightCachePath() const;
+
+  Expected<void> SetXNNPackWeightCacheFileDescriptor(int fd);
+  Expected<int> GetXNNPackWeightCacheFileDescriptor() const;
 };
 
 }  // namespace litert

litert/cc/options/litert_cpu_options_test.cc

@@ -97,6 +97,20 @@ TEST(CpuOptions, SetAndGetXNNPackWeighCachePathWorks) {
               IsOkAndHolds(StrEq("a/path")));
 }
 
+TEST(CpuOptions, CheckXNNPackWeightCacheFileDescriptorDefaultValue) {
+  LITERT_ASSERT_OK_AND_ASSIGN(CpuOptions options, CpuOptions::Create());
+  EXPECT_THAT(options.GetXNNPackWeightCacheFileDescriptor(),
+              IsOkAndHolds(-1));
+}
+
+TEST(CpuOptions, SetAndGetXNNPackWeighCacheFileDescriptorWorks) {
+  LITERT_ASSERT_OK_AND_ASSIGN(CpuOptions options, CpuOptions::Create());
+
+  LITERT_EXPECT_OK(options.SetXNNPackWeightCacheFileDescriptor(1234));
+  EXPECT_THAT(options.GetXNNPackWeightCacheFileDescriptor(),
+              IsOkAndHolds(1234));
+}
+
 TEST(CpuOptions, CheckXNNPackFlagsDefaultValue) {
   LITERT_ASSERT_OK_AND_ASSIGN(CpuOptions options, CpuOptions::Create());
   // Note: we can't check the default value for this as XNNPack compile options