Expose a method getter api in Module (pytorch#11929)

Summary:
Pull Request resolved: pytorch#11929

It is sometimes useful for modules to be passed around in user space for user convenience, but for utilities those users interact with to operate on the methods directly.
This is a bit of a layering violation but it allows us to define narrower scoped helpers so I think its worth the risk that the user finds a way to trash the internal method object. In general method should be fairly robust to being trashed anyway since itself is a blessed front end api from the team.

Differential Revision: D77248983

Reviewed By: larryliu0820

Author: JacobSzwejbka

extension/module/module.cpp

@@ -228,6 +228,16 @@ runtime::Error Module::load_method(
   return runtime::Error::Ok;
 }
 
+ET_NODISCARD inline runtime::Result<Method*> Module::method(
+    const std::string& method_name) {
+  ET_CHECK_OR_RETURN_ERROR(
+      methods_.count(method_name) > 0,
+      InvalidArgument,
+      "no such method in program: %s",
+      method_name.c_str());
+  return methods_[method_name].method.get();
+}
+
 runtime::Result<MethodMeta> Module::method_meta(
     const std::string& method_name) {
   ET_CHECK_OK_OR_RETURN_ERROR(load());

extension/module/module.h

@@ -194,6 +194,19 @@ class Module {
     return load_method(method_name, nullptr, event_tracer);
   }
 
+  /**
+   * Get a method by it's name. Not recommended to use this method directly as
+   * an end user. It's exposed to allow for composability of module in apis that
+   * operate on method.
+   *
+   * @param[in] method_name The name of the method to get.
+   *
+   * @returns A Result object containing either a pointer to the requested
+   *          method or an error to indicate failure.
+   */
+  ET_NODISCARD inline runtime::Result<Method*> method(
+      const std::string& method_name);
+
   /**
    * Load the 'forward' method from the program and set up memory management if
    * needed. The loaded method is cached to reuse the next time it's executed.