Added unit test.

Simplified code, In the case of no user input lifted constants are placed before 'none' node, aka the begining. We update the graph node to place these at the beginning as well.

fixed formatting.

Author: cptspacemanspiff

exir/memory_planning.py

@@ -248,7 +248,16 @@ def verify_graph_input_output(self) -> None:
                     has_dynamic_unbound_output |= has_dynamic_unbound_tensor
 
         # 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
+        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"

exir/program/_program.py

@@ -268,42 +268,13 @@ def lift_constant_tensor_pass(ep):
     buffers = list(graph_signature.buffers)
 
     fake_mode = list(ep.graph.nodes)[0].meta["val"].fake_mode
-    insert_before_node = None
+    first_user_input = None
     lifted_constants = []
     for node in ep.graph.nodes:
         if node.op == "placeholder" and node.name in graph_signature.user_inputs:
-            insert_before_node = node  # first user input
+            first_user_input = node
             break
 
-    if insert_before_node is None:
-        # we have no user inputs, find the node after the last buffer
-        # (that we will insert the lifted constants before).
-        # this is a bit hacky, but I am not certain of what the contract is for
-        # node ordering. is the first non-placeholder node guranteed to be the
-        # first node after input paramters? what if there is no op, and it is
-        # just placeholders? Easier to just find the last buffer, and insert after.
-
-        # also error if we have no buffers and no user inputs... if that is an issue, fix it later?
-        last_buffer = None
-        for node in ep.graph.nodes:
-            node_buffer_fqn = graph_signature.inputs_to_buffers.get(node.name, None)
-            # not sure if both cases are needed, if is it possible to encounter a
-            # buffer that is not a user input?
-            if (
-                node_buffer_fqn is not None
-                and node_buffer_fqn in graph_signature.buffers
-            ):
-                last_buffer = node
-                continue
-            if node.op == "placeholder" and node.name in graph_signature.buffers:
-                last_buffer = node
-                continue
-        # we have our last buffer, grab the node after it, to insert the lifted constants before.
-        insert_before_node = last_buffer.next
-
-    if insert_before_node is None:
-        raise ValueError("No user inputs and no buffers found. Cannot lift constants.")
-
     for node in ep.graph.nodes:
         if node.op == "get_attr":
             constant_tensor = getattr(ep.graph_module, node.target)
@@ -312,7 +283,7 @@ def lift_constant_tensor_pass(ep):
 
             constant_tensor_fqn = f"_lifted_tensor_constant{len(buffers)}"
 
-            with ep.graph.inserting_before(insert_before_node):
+            with ep.graph.inserting_before(first_user_input):
                 # Insert the constant node before the first user input
                 const_placeholder_node = ep.graph.placeholder(constant_tensor_fqn)
                 for k, v in node.meta.items():
@@ -345,9 +316,8 @@ def lift_constant_tensor_pass(ep):
             new_input_specs.extend(lifted_constants)
             lifted_constants.clear()
         new_input_specs.append(s)
-    # Add remaining lifted constants if no user inputs exist.
     if len(lifted_constants) > 0:
-        new_input_specs.extend(lifted_constants)
+        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,35 @@ 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:
+        assert ep.graph_signature.input_specs[0].arg.name == "_lifted_tensor_constant1"
+        assert ep.graph_signature.input_specs[1].arg.name == "b_a"
+
+        executorch_program = edge.to_executorch()
+        # # the graph signature should also be the same:
+        # executorch_program.graph_signature.input_specs[0].arg.name == "_lifted_tensor_constant1"
+        # executorch_program.graph_signature.input_specs[1].arg.name == "b_a"
+
     def test_constant_prop_pass_for_parameter(self) -> None:
         def count_additions(gm: torch.fx.GraphModule) -> int:
             return sum(