Add get_method API to module

Summary: It is for reusing the exention/module in the module definition in pybindings https://github.com/pytorch/executorch/blob/1a918c779e16c0ee903a08b30c1c666d1efb2c57/extension/pybindings/pybindings.cpp#L172

Differential Revision: D71135352

Author: cccclai

extension/module/module.h

@@ -459,6 +459,15 @@ class Module {
     return event_tracer_.get();
   }
 
+  ET_NODISCARD
+  runtime::Result<runtime::Method*> get_method(const std::string& method_name) {
+    if (methods_.count(method_name) == 0) {
+      ET_LOG(Info, "Method %s not found", method_name.c_str());
+      return runtime::Error::NotFound;
+    }
+    return methods_[method_name].method.get();
+  }
+
  private:
   struct MethodHolder {
     std::vector<std::vector<uint8_t>> planned_buffers;

extension/module/test/module_test.cpp

@@ -53,6 +53,7 @@ TEST_F(ModuleTest, TestLoadNonExistent) {
   EXPECT_FALSE(module.is_loaded());
 }
 
+
 TEST_F(ModuleTest, TestLoadCorruptedFile) {
   Module module("/dev/null");
   const auto error = module.load();
@@ -95,6 +96,30 @@ TEST_F(ModuleTest, TestLoadNonExistentMethod) {
   EXPECT_TRUE(module.is_loaded());
 }
 
+TEST_F(ModuleTest, TestGetMethod) {
+  Module module(model_path_);
+
+  const auto error = module.load_method("forward");
+  EXPECT_EQ(error, Error::Ok);
+  auto method_res = module.get_method("forward");
+  EXPECT_EQ(method_res.error(), Error::Ok);
+  auto method = method_res.get();
+  EXPECT_EQ(strcmp(method->method_meta().name(), "forward"), 0);
+
+}
+
+TEST_F(ModuleTest, TestGetNonExistMethod) {
+  Module module(model_path_);
+
+  const auto error = module.load_method("forward");
+  EXPECT_EQ(error, Error::Ok);
+
+  // Try to get a method that doesn't exist
+  auto method_res = module.get_method("backward");
+  EXPECT_EQ(method_res.error(), Error::NotFound);
+}
+
+
 TEST_F(ModuleTest, TestMethodMeta) {
   Module module(model_path_);