From 547af2d744e585e79990bfb3823ccca4594370db Mon Sep 17 00:00:00 2001
From: David Lin <lind@meta.com>
Date: Thu, 21 Nov 2024 09:41:32 -0800
Subject: [PATCH] Add benchmark for training training (#6933)

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
---
 extension/module/module.h | 19 +++++++++----------
 1 file changed, 9 insertions(+), 10 deletions(-)

diff --git a/extension/module/module.h b/extension/module/module.h
index 45ed38a7ff2..fe9a62e0cc2 100644
--- a/extension/module/module.h
+++ b/extension/module/module.h
@@ -38,6 +38,15 @@ class Module {
     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
    * memory locking behavior.
@@ -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};