Remove the requirement to push the whole Framework Module as input to (…

…#591)

the ForgeModule during compile time. Previously required for parameter
initialization.
- This change is required to remove added constraints during the op test
  generation
- This change cleans the generated module a bit and does serialization
  of
  all required parameters that are just load within the model during
compile time
- This chain represents one of the steps to generate single op tests
  with
  focus on a specific model

Fix #590

forge/forge/python_codegen.py

@@ -259,17 +259,19 @@ def write_forward(self, ops, inputs, outputs):
         self.indent = 0
         self.wl("")
 
-    def write_param_parser(self, param_names, param_file_name):
+    def write_param_parser(
+        self, param_names, param_file_name, names_params_file_name=None, named_buffers_file_name=None
+    ):
         self.indent = 1
 
         if self.framework == "pytorch":
-            self.wl(f"def process_framework_parameters(self, model):")
+            self.wl(f"def process_framework_parameters(self):")
             self.indent += 1
-            self.wl(f"named_parameters = dict(model.state_dict().items())")
+            self.wl(f"named_parameters = torch.load('{names_params_file_name}')")
             if param_file_name is not None:
                 self.wl(f'serialized_params = torch.load("{param_file_name}")')
                 self.wl(f"named_parameters.update(serialized_params)")
-            self.wl("named_buffers = dict(model.named_buffers())")
+            self.wl(f"named_buffers = torch.load('{named_buffers_file_name}')")
             self.wl("named_parameters.update(named_buffers)")
 
             if len(param_names):
@@ -1249,7 +1251,7 @@ def write_param_parser(self, param_names, param_file_name):
         self.indent = 1
 
         if self.framework == "pytorch":
-            self.wl(f"def process_framework_parameters(self, model):")
+            self.wl(f"def process_framework_parameters(self):")
             self.indent += 1
 
             self.wl("named_parameters = dict(model.named_parameters())")

forge/forge/tvm_to_python.py

@@ -2039,7 +2039,12 @@ def generate_forge_module(
             forge_mod.module.process_framework_parameters(framework_mod.module)
         else:
             forge_mod = TestClass(writer.module_name)
-            forge_mod.process_framework_parameters(framework_mod.module)
+
+            if isinstance(framework_mod, forge.module.PyTorchModule):
+                forge_mod.process_framework_parameters()
+            else:
+                forge_mod.process_framework_parameters(framework_mod.module)
+
             assert not any(
                 [param.value() is None for param in forge_mod.get_parameters()]
             ), f"Could not retrieve parameters from framework and tvm"
@@ -2655,8 +2660,24 @@ def delete_unneeded_outputs(ops, returns):
             param_file_name = os.path.join(writer.module_directory, writer.module_name + "_params.pt")
             torch.save(params_from_tvm, param_file_name)
 
-        param_names.update(const_names)
-        writer.write_param_parser(param_names, param_file_name)
+        if framework == "pytorch":
+            # Store named parameters
+            names_params_file_name = os.path.join(writer.module_directory, writer.module_name + "_names_params.pt")
+            named_parameters = dict(framework_mod.module.state_dict().items())
+            torch.save(named_parameters, names_params_file_name)
+
+            # Store named buffers
+            named_buffers_file_name = os.path.join(writer.module_directory, writer.module_name + "_named_buffers.pt")
+            named_buffers = dict(framework_mod.module.named_buffers())
+            torch.save(named_buffers, named_buffers_file_name)
+
+            # Generate Forge module parameter parser
+            param_names.update(const_names)
+            writer.write_param_parser(param_names, param_file_name, names_params_file_name, named_buffers_file_name)
+        else:
+            param_names.update(const_names)
+            writer.write_param_parser(param_names, param_file_name)
+
         writer.close_file()
 
         modules.append(writer)