input/output spec update of the exported program after graph transformation

[mhs4670go]

Hello.

I'm trying to transform the graph before passing the graph as inputs of the my backend's serialization process. Below is very similar with my circumstances.

executorch/backends/vulkan/vulkan_preprocess.py

Line 47 in 51f6455

AddmmToLinearTransform(),

The problem is that there is output spec change of the graph after running AddmmToLinearTransform pass. In the serialization, I refer to the edge_program.graph_signature but it failed because of the mismatch between the graph module and the exported program input/output signature.

class SimpleLinear(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(3, 3)

    def forward(self, arg):
        return (self.linear(arg),)

    def get_example_inputs(self):
        return (torch.randn(3, 3),)

### Before the transformation
class GraphModule(torch.nn.Module):
        def forward(self, p_linear_weight: "f32[3, 3]", p_linear_bias: "f32[3]", arg: "f32[3, 3]"):
            # File: /home/seongwoo/-exir/test/modules/single/op/linear.py:10 in forward, code: return (self.linear(arg),)
            aten_permute_copy_default: "f32[3, 3]" = executorch_exir_dialects_edge__ops_aten_permute_copy_default(p_linear_weight, [1, 0]);  p_linear_weight = None
            aten_addmm_default: "f32[3, 3]" = executorch_exir_dialects_edge__ops_aten_addmm_default(p_linear_bias, arg, aten_permute_copy_default);  p_linear_bias = arg = aten_permute_copy_default = None
            return (aten_addmm_default,)
            
Graph signature: ExportGraphSignature(input_specs=[InputSpec(kind=<InputKind.PARAMETER: 2>, arg=TensorArgument(name='p_linear_weight'), target='linear.weight', persistent=None), InputSpec(kind=<InputKind.PARAMETER: 2>, arg=TensorArgument(name='p_linear_bias'), target='linear.bias', persistent=None), InputSpec(kind=<InputKind.USER_INPUT: 1>, arg=TensorArgument(name='arg'), target=None, persistent=None)], output_specs=[OutputSpec(kind=<OutputKind.USER_OUTPUT: 1>, arg=TensorArgument(name='aten_addmm_default'), target=None)])
Range constraints: {}

### After the transformation
class GraphModule(torch.nn.Module):
        def forward(self, p_linear_weight: "f32[3, 3]", p_linear_bias: "f32[3]", arg: "f32[3, 3]"):
            # File: /home/seongwoo/-exir/test/modules/single/op/linear.py:10 in forward, code: return (self.linear(arg),)
            aten_linear_default: "f32[3, 3]" = executorch_exir_dialects_edge__ops_aten_linear_default(arg, p_linear_weight, p_linear_bias);  arg = p_linear_weight = p_linear_bias = None
            return [aten_linear_default]
            
Graph signature: ExportGraphSignature(input_specs=[InputSpec(kind=<InputKind.PARAMETER: 2>, arg=TensorArgument(name='p_linear_weight'), target='linear.weight', persistent=None), InputSpec(kind=<InputKind.PARAMETER: 2>, arg=TensorArgument(name='p_linear_bias'), target='linear.bias', persistent=None), InputSpec(kind=<InputKind.USER_INPUT: 1>, arg=TensorArgument(name='arg'), target=None, persistent=None)], output_specs=[OutputSpec(kind=<OutputKind.USER_OUTPUT: 1>, arg=TensorArgument(name='aten_addmm_default'), target=None)])
Range constraints: {}
###############
# Not changed despite of the transoformation
# arg=TensorArgument(name='aten_addmm_default'), target=None)]

Should I just update the exported program like below or is there any related API?

for node in edge_program.graph.nodes:
        if node.op != "output":
            continue
        edge_program.graph_signature.output_specs = [
            OutputSpec(
                kind=OutputKind.USER_OUTPUT,
                arg=TensorArgument(name=arg.name),
                target=None,
            )
            for arg in node.args[0]
        ]

[guangy10]

@mhs4670go What backend serialization process are you trying to lower the graph to? Is it Vulkan?

[mhs4670go]

@guangy10 Ah, it's internal backend of my company. The graph IR follows this schema. Technically speaking, when the outputs are registered in the circle graph, it refers to the name of the output spec in the exported program, which is addmm. But the name of the created linear op in the transformation is aten_linear_default. Then, there is mismatch between them.

So, I think there are two options for this case.

1. Update the output spec of the exported program after the transformation.
2. Set the name of the linear op to the name of the existing addmm operator during transformation.

constant_prop_pass updates the input spec according to the propagation. But replace_addmm_mm_with_linear doesn't update the spec. So, I wonder that I should manually update the spec.

[SS-JIA]

@mhs4670go recently, #3483 was landed that allows backends to specify certain operators to avoid decomposing; this will make it so that the AddmmToLinearTransform() pass is not necessary as aten.linear will never be decomposed in the first place.

To see an example of how it works please take a look at the NonDecompTestPartitioner in op_partitioner_demo.py. I would recommend going this route instead of using the linear recomposition pass. Please try it out!

cc: @tarun292

[angelayi]

Instead of manually updating the graph signature, you can also try wrapping the code that runs the passes with the following context manager, which will automatically update the signature for you

with ep.graph_module._set_replace_hook(ep.graph_signature.get_replace_hook()):
    AddmmToLinearTransform()(ep.graph_module)

Here is an example where we use it: code

But separately, a work item for the executorch team should be to update the passes so that it will do this automatic updating. cc @kimishpatel who wrote the pass

[mhs4670go]

@SS-JIA @angelayi Thanks for the helps! Both way are very helpful for me.

[kimishpatel]

_set_replace_hook

I did not know about that magic context manager. Yeah when I wrote that pass, I definitely missed this case, but also it was a long time ago so I dont know if we even had graph signature. Will follow up on updating the pass, but I wont be surprised if I forget about it. lol.