Add ynn_subgraph_optimize

Ultimately, we should not call this automatically from `ynn_create_runtime`, but it is still there for now, so we can replace the call in client uses of YNNPACK.

PiperOrigin-RevId: 822833148

Author: dsharletg

ynnpack/include/ynnpack.h

@@ -35,6 +35,10 @@ enum ynn_status {
 
 typedef struct ynn_subgraph* ynn_subgraph_t;
 
+// This type is an alias for `slinky::thread_pool`. A `slinky::thread_pool`
+// instance may be casted to `ynn_threadpool` and passed to YNNPACK APIs.
+typedef struct ynn_threadpool* ynn_threadpool_t;
+
 // Create a new subgraph, with `external_value_ids` reserved ids for external
 // values.
 enum ynn_status ynn_create_subgraph(uint32_t external_value_ids, uint32_t flags,
@@ -43,6 +47,10 @@ enum ynn_status ynn_create_subgraph(uint32_t external_value_ids, uint32_t flags,
 // Delete a subgraph previously created with `ynn_create_subgraph`.
 void ynn_delete_subgraph(ynn_subgraph_t subgraph);
 
+// Apply subgraph rewrites and other optimizations to the subgraph.
+ynn_status ynn_optimize_subgraph(ynn_subgraph_t subgraph,
+                                 ynn_threadpool_t threadpool, uint32_t flags);
+
 // Describes a type for a value.
 enum ynn_type {
   ynn_type_invalid = -1,
@@ -398,10 +406,6 @@ struct ynn_scheduler {
 
 typedef const struct ynn_scheduler* ynn_scheduler_t;
 
-// This type is an alias for `slinky::thread_pool`. A `slinky::thread_pool`
-// instance may be casted to `ynn_threadpool` and passed to YNNPACK APIs.
-typedef struct ynn_threadpool* ynn_threadpool_t;
-
 // Create a threadpool that uses `scheduler` to start work on other threads.
 // `scheduler` is not copied, since it is stateless we expect it to be stored
 // globally.

ynnpack/subgraph/runtime.cc

@@ -640,23 +640,14 @@ ynn_status ynn_runtime::setup() {
 ynn_status ynn_create_runtime(ynn_subgraph_t subgraph,
                               ynn_threadpool_t threadpool, uint32_t flags,
                               ynn_runtime_t* runtime_out) {
-#if YNN_LOG_LEVEL >= YNN_LOG_LEVEL_DEBUG
-  YNN_LOG_DEBUG() << "subgraph before optimization:\n";
-  subgraph->dump(std::cout);
-#endif
-
-  slinky::thread_pool* slinky_threadpool =
-      reinterpret_cast<slinky::thread_pool*>(threadpool);
-  ynn_status status = subgraph->optimize(slinky_threadpool);
+  // TODO(b/454450773): Don't call this here.
+  ynn_status status = ynn_optimize_subgraph(subgraph, threadpool, 0);
   if (status != ynn_status_success) {
     return status;
   }
 
-#if YNN_LOG_LEVEL >= YNN_LOG_LEVEL_DEBUG
-  YNN_LOG_DEBUG() << "subgraph after optimization:\n";
-  subgraph->dump(std::cout);
-#endif
-
+  slinky::thread_pool* slinky_threadpool =
+      reinterpret_cast<slinky::thread_pool*>(threadpool);
   std::unique_ptr<ynn_runtime> runtime(
       new ynn_runtime(*subgraph, slinky_threadpool, flags));
   status = runtime->build();

ynnpack/subgraph/subgraph.cc

@@ -806,6 +806,28 @@ ynn_status ynn_create_subgraph(uint32_t external_value_ids, uint32_t flags,
   return ynn_status_success;
 }
 
+ynn_status ynn_optimize_subgraph(ynn_subgraph_t subgraph,
+                                 ynn_threadpool_t threadpool, uint32_t flags) {
+#if YNN_LOG_LEVEL >= YNN_LOG_LEVEL_DEBUG
+  YNN_LOG_DEBUG() << "subgraph before optimization:\n";
+  subgraph->dump(std::cout);
+#endif
+
+  slinky::thread_pool* slinky_threadpool =
+      reinterpret_cast<slinky::thread_pool*>(threadpool);
+  ynn_status status = subgraph->optimize(slinky_threadpool);
+  if (status != ynn_status_success) {
+    return status;
+  }
+
+#if YNN_LOG_LEVEL >= YNN_LOG_LEVEL_DEBUG
+  YNN_LOG_DEBUG() << "subgraph after optimization:\n";
+  subgraph->dump(std::cout);
+#endif
+
+  return ynn_status_success;
+}
+
 void ynn_delete_subgraph(ynn_subgraph_t subgraph) { delete subgraph; }
 
 }  // extern "C"

ynnpack/subgraph/test/subgraph_builder.cc

@@ -290,6 +290,12 @@ Runtime::Runtime(ynn_subgraph_t subgraph, TestScheduler* scheduler,
         std::unique_ptr<ynn_threadpool, decltype(&ynn_delete_threadpool)>(
             threadpool, ynn_delete_threadpool);
   }
+
+  status_ = ynn_optimize_subgraph(subgraph, threadpool_.get(), 0);
+  if (status_ != ynn_status_success) {
+    return;
+  }
+
   status_ = ynn_create_runtime(subgraph, threadpool_.get(), flags, &runtime);
   runtime_ = std::unique_ptr<ynn_runtime, decltype(&ynn_delete_runtime)>(
       runtime, ynn_delete_runtime);

ynnpack/xnnpack/runtime.cc

@@ -122,9 +122,17 @@ static xnn_status create_runtime_impl(xnn_subgraph_t subgraph,
     ynn_flags |= YNN_RUNTIME_FLAG_NO_SCHEDULE;
   }
 
+  ynn_threadpool_t ynn_threadpool =
+      xnn_threadpool ? xnn_threadpool->ynn : nullptr;
+
+  ynn_status status =
+      ynn_optimize_subgraph(subgraph->ynn, ynn_threadpool, ynn_flags);
+  if (status != ynn_status_success) {
+    return ynn::xnn_status_from_ynn(status);
+  }
+
   return ynn::xnn_status_from_ynn(ynn_create_runtime(
-      subgraph->ynn, xnn_threadpool ? xnn_threadpool->ynn : nullptr, ynn_flags,
-      (ynn_runtime_t*)runtime_out));
+      subgraph->ynn, ynn_threadpool, ynn_flags, (ynn_runtime_t*)runtime_out));
 }
 
 xnn_status xnn_create_runtime_with_threadpool(xnn_subgraph_t subgraph,