Changes to allow the export of functions with no user input. (pytorch#8031)

### Summary
This is a set of changes that I have been using to get executorch
working with Multiple Entrypoints with shared state.
This is related to: pytorch#8030 

### Issue & Changes
Currently trying to export a function:

```python
def init():
  self.shared_state.fill_(0)
```

Will fail, because in to_edge in exir in the program.py the lifted
constants are assumed to be placed before the first user input. This
causes weird behavior when there are no user inputs. I changed it to
place the lifted constant after the last buffer, if there are no user
inputs to place before. However, this relies on my understanding of the
implicit layout of the graph inputs.

At the same time later in the after the memory planning phase of
to_executorch, it validates that memory planning was correct based on
whether the `graph_input_allocated` flag is set, this only applies to
user inputs, of which we have none, so it errors out. I added a check to
bypass this error if there are no user inputs, but I honestly do not
understand enough of the validation check to know if that is
appropriate.

### Comments

In the current executorch with no support for shared state, this case
does not make sense, but pytorch#8030 is my attempt at adding that capability.
and having initialization methods that init the buffers from constants
is useful, especially since their initail state is undefined.

Currently this is not ready, I have no tests/ and have my random notes
in the commit as comments, and other than validating that it worked as I
was working on the shared state export have not done anything in
depth....

But I kind-of want feedback if my solution seems correct, or if I am
missing something. Particularly regarding my understanding of
placeholder ordering and signature logic, and whether bypassing the
graph_input_allocated validation is appropriate.

Out of all the changes I had to make for shared state, this is the one I
am least sure about.

 


cc @JacobSzwejbka @angelayi

Author: cptspacemanspiff

exir/memory_planning.py

@@ -24,7 +24,7 @@
 from executorch.exir.tensor import TensorSpec
 
 from torch import fx
-from torch.export.exported_program import ExportGraphSignature
+from torch.export.exported_program import ExportGraphSignature, InputKind
 from torch.fx import Node
 from torch.utils._pytree import tree_flatten
 
@@ -247,7 +247,19 @@ def verify_graph_input_output(self) -> None:
                     graph_output_allocated = allocated
                     has_dynamic_unbound_output |= has_dynamic_unbound_tensor
 
-        if "placeholder" in check_list:
+        # only check if inputs are allocated if there are user inputs:
+        user_inputs_exist = (
+            len(
+                list(
+                    filter(
+                        lambda input: input.kind == InputKind.USER_INPUT,
+                        self.graph_signature.input_specs,
+                    )
+                )
+            )
+        ) > 0
+
+        if "placeholder" in check_list and user_inputs_exist:
             assert graph_input_allocated is not None, "graph_input_allocated not set"
             if not has_dynamic_unbound_input:
                 assert (

exir/program/_program.py

@@ -322,6 +322,8 @@ def lift_constant_tensor_pass(ep):
             new_input_specs.extend(lifted_constants)
             lifted_constants.clear()
         new_input_specs.append(s)
+    if len(lifted_constants) > 0:
+        new_input_specs = lifted_constants + new_input_specs
     ep.graph_signature.input_specs = new_input_specs
     ep.graph_module.recompile()
     return ep

exir/tests/test_passes.py

@@ -1057,6 +1057,36 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             new_ep.graph_module.code
         )
 
+    def test_pass_no_user_inputs(self) -> None:
+        class NoUserInputs(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("a", torch.ones(1))
+
+            def forward(self) -> torch.Tensor:
+                return 3 + self.a
+
+        mod = NoUserInputs()
+        exported_program = export(mod, (), strict=True)
+        edge = to_edge(
+            exported_program,
+            compile_config=EdgeCompileConfig(_skip_dim_order=False),
+        )
+        ep = edge.exported_program()
+        # because there is no user input, the lifted constant should be the first input.
+        FileCheck().check("_lifted_tensor_constant1").check(
+            "b_a"  # followed by the buffer input.
+        ).run(ep.graph_module.code)
+
+        # the graph signature should also be the same:
+        self.assertEqual(
+            ep.graph_signature.input_specs[0].arg.name, "_lifted_tensor_constant1"
+        )
+        self.assertEqual(ep.graph_signature.input_specs[1].arg.name, "b_a")
+
+        # Validate that the program successfully passes validation to executorch:
+        edge.to_executorch()
+
     def test_constant_prop_pass_for_parameter(self) -> None:
         def count_additions(gm: torch.fx.GraphModule) -> int:
             return sum(