From 0b965bbd7fb9c8ccbf086ae4c59f124254dc5887 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Lindstr=C3=B6m?= Date: Sun, 23 Feb 2025 23:43:55 +0100 Subject: [PATCH] Arm backend: Don't handle const placeholders in match_arg_ranks_pass This pass did previously handle constant placeholders directly by modifying the data in the state_dict, whereas call_functions are handled by inserting view ops. Since the introduction of the FuseConsantOpsPass however, ops acting on constant placeholders are computed AOT and fused into one single node. We can therefore now handle all nodes the same way here by inserting view nodes, removing the need for a special case in this pass. Change-Id: If282eea352205e1b0b10b9b747642fa0d6caa4b3 --- backends/arm/_passes/match_arg_ranks_pass.py | 49 ++++++-------------- 1 file changed, 13 insertions(+), 36 deletions(-) diff --git a/backends/arm/_passes/match_arg_ranks_pass.py b/backends/arm/_passes/match_arg_ranks_pass.py index 941d20c95a1..3fcfe6edd35 100644 --- a/backends/arm/_passes/match_arg_ranks_pass.py +++ b/backends/arm/_passes/match_arg_ranks_pass.py @@ -1,6 +1,6 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. -# Copyright 2024 Arm Limited and/or its affiliates. # All rights reserved. +# Copyright 2024-2025 Arm Limited and/or its affiliates. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -23,7 +23,17 @@ class MatchArgRanksPass(ExportPass): """ For ops in 'targeted_ops', make sure that the inputs share the same rank. - New dimensions are inserted at from the beginning of the + New dimensions are inserted from the beginning of the inputs that have a + lower rank to match the input with the highest rank. + + Example: + input0 = shape(4, 3, 2) + input1 = shape(2) + input2 = shape(3, 1) + Becomes: + input0 = shape(4, 3, 2) + input1 = shape(1, 1, 2) + input2 = shape(1, 3, 1) """ def __init__(self, exported_program): @@ -54,34 +64,6 @@ def _match_op_rank(self, graph_module, node, arg, max_rank): ) node.replace_input_with(arg, view) - def _match_buffer_rank(self, arg, max_rank): - """ - Change arg's fake tensor meta to match max_rank if: - - arg is found in inputs_to_buffers or inputs_to_parameters. - """ - fake_tensor = get_first_fake_tensor(arg) - shape = fake_tensor.shape - rank = len(shape) - new_shape = list([1] * (max_rank - rank) + list(shape)) - - buffer_name = None - if arg.name in self.exported_program.graph_signature.inputs_to_buffers: - buffer_name = self.exported_program.graph_signature.inputs_to_buffers[ - arg.name - ] - elif arg.name in self.exported_program.graph_signature.inputs_to_parameters: - buffer_name = self.exported_program.graph_signature.inputs_to_parameters[ - arg.name - ] - if buffer_name: - new_tensor = self.exported_program.state_dict[buffer_name].reshape( - new_shape - ) - self.exported_program.state_dict[buffer_name] = new_tensor - arg.meta["val"] = fake_tensor.fake_mode.from_tensor( - new_tensor, static_shapes=True - ) - def call(self, graph_module: GraphModule) -> PassResult: for node in graph_module.graph.nodes: node = cast(Node, node) @@ -105,12 +87,7 @@ def call(self, graph_module: GraphModule) -> PassResult: if rank == max_rank: continue - # If the argument is call_function, match shape by inserting view node. - if arg.op == "call_function": - self._match_op_rank(graph_module, node, arg, max_rank) - else: - # If the argument is a buffer or parameter, adjust shape by changing the fake tensor meta. - self._match_buffer_rank(arg, max_rank) + self._match_op_rank(graph_module, node, arg, max_rank) graph_module.recompile() graph_module = super().call(graph_module).graph_module