Add benchmark for training training

Summary: This adds benchmarking for training. It loads the 3 methods required for training (param names, param index, gradient index) and runs the exported forward+backwards, then runs the optimizer.

Reviewed By: JacobSzwejbka

Differential Revision: D65905482

Author: David Lin

extension/module/module.h

@@ -37,6 +37,15 @@ class Module {
     /// Use memory locking and ignore errors.
     MmapUseMlockIgnoreErrors,
   };
+  
+  struct MethodHolder {
+    std::vector<std::vector<uint8_t>> planned_buffers;
+    std::vector<runtime::Span<uint8_t>> planned_spans;
+    std::unique_ptr<runtime::HierarchicalAllocator> planned_memory;
+    std::unique_ptr<runtime::MemoryManager> memory_manager;
+    std::unique_ptr<runtime::Method> method;
+    std::vector<runtime::EValue> inputs;
+  };
 
   /**
    * Constructs an instance by loading a program from a file with specified
@@ -423,16 +432,6 @@ class Module {
     return event_tracer_.get();
   }
 
- private:
-  struct MethodHolder {
-    std::vector<std::vector<uint8_t>> planned_buffers;
-    std::vector<runtime::Span<uint8_t>> planned_spans;
-    std::unique_ptr<runtime::HierarchicalAllocator> planned_memory;
-    std::unique_ptr<runtime::MemoryManager> memory_manager;
-    std::unique_ptr<runtime::Method> method;
-    std::vector<runtime::EValue> inputs;
-  };
-
  private:
   std::string file_path_;
   LoadMode load_mode_{LoadMode::MmapUseMlock};