From 6e4e09df3e2eb0709d458f5e3ffc03a5a9155f13 Mon Sep 17 00:00:00 2001
From: Regan Stehle <reganstehle@meta.com>
Date: Fri, 6 Dec 2024 11:56:01 -0800
Subject: [PATCH] Add num_instructions program function (#7206)

Summary:

Add num_instructions program function to executorch program methods, because some embedded systems this information useful to estimate the required method allocator buffer size.

Reviewed By: JacobSzwejbka, dbort, cmt0

Differential Revision: D66504180
---
 runtime/executor/method_meta.cpp           | 20 ++++++++++++++++++++
 runtime/executor/method_meta.h             |  7 +++++++
 runtime/executor/test/method_meta_test.cpp |  3 +++
 3 files changed, 30 insertions(+)

diff --git a/runtime/executor/method_meta.cpp b/runtime/executor/method_meta.cpp
index 5acf055a89f..f43398d0ab7 100644
--- a/runtime/executor/method_meta.cpp
+++ b/runtime/executor/method_meta.cpp
@@ -210,5 +210,25 @@ Result<int64_t> MethodMeta::memory_planned_buffer_size(size_t index) const {
   return s_plan_->non_const_buffer_sizes()->Get(index + 1);
 }
 
+size_t MethodMeta::num_instructions() const {
+  const auto chains = s_plan_->chains();
+  if (chains == nullptr) {
+    return 0;
+  }
+  const auto num_chains = chains->size();
+  auto num_instructions = 0;
+  for (size_t i = 0; i < num_chains; ++i) {
+    auto s_chain = chains->Get(i);
+    if (s_chain == nullptr) {
+      continue;
+    }
+    auto s_instructions = s_chain->instructions();
+    if (s_instructions != nullptr) {
+      num_instructions += s_instructions->size();
+    }
+  }
+  return num_instructions;
+}
+
 } // namespace runtime
 } // namespace executorch
diff --git a/runtime/executor/method_meta.h b/runtime/executor/method_meta.h
index 569b93c0f9d..a5e932981c1 100644
--- a/runtime/executor/method_meta.h
+++ b/runtime/executor/method_meta.h
@@ -185,6 +185,13 @@ class MethodMeta final {
    */
   Result<int64_t> memory_planned_buffer_size(size_t index) const;
 
+  /**
+   * Get the number of instructions in this method.
+   *
+   * @returns The number of instructions.
+   */
+  ET_EXPERIMENTAL size_t num_instructions() const;
+
   /**
    * DEPRECATED: Use num_memory_planned_buffers() instead.
    */
diff --git a/runtime/executor/test/method_meta_test.cpp b/runtime/executor/test/method_meta_test.cpp
index bd48f64d98f..1cb20496f5a 100644
--- a/runtime/executor/test/method_meta_test.cpp
+++ b/runtime/executor/test/method_meta_test.cpp
@@ -92,6 +92,9 @@ TEST_F(MethodMetaTest, MethodMetaApi) {
       method_meta->non_const_buffer_size(1).error(),
       Error::InvalidArgument); // Deprecated API
 
+  // Number instructions in method is nonzero
+  EXPECT_NE(method_meta->num_instructions(), 0);
+
   // Missing method fails
   EXPECT_EQ(
       program_->method_meta("not_a_method").error(), Error::InvalidArgument);