diff --git a/py/torch_tensorrt/dynamo/lowering/passes/repair_input_as_output.py b/py/torch_tensorrt/dynamo/lowering/passes/repair_input_as_output.py index 3773550d99..9da585ed5a 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/repair_input_as_output.py +++ b/py/torch_tensorrt/dynamo/lowering/passes/repair_input_as_output.py @@ -19,29 +19,44 @@ def repair_input_as_output( """ modified_graph = False - # Extract graph placeholder Tensors + # Extract graph placeholder Tensors. Constant folding also turns + # input-independent tensor factories (for example torch.zeros) into + # registered `_frozen_param*` attributes. When such a value is returned + # across a graph break, eager code may legally mutate it in-place. Return + # a clone so the registered constant itself remains immutable. placeholders = get_tensor_placeholders(gm) + folded_constants = [ + node + for node in gm.graph.nodes + if node.op == "get_attr" and str(node.target).startswith("_frozen_param") + ] - for placeholder in placeholders: - # If any placeholder has any users which are direct graph outputs - if len(placeholder.users) >= 1 and any( - user.op == "output" for user in placeholder.users + for source in [*placeholders, *folded_constants]: + # If any source has any users which are direct graph outputs + if len(source.users) >= 1 and any( + user.op == "output" for user in source.users ): modified_graph = True - # Get direct graph outputs which are direct uses of placeholders - direct_outputs = [user for user in placeholder.users if user.op == "output"] + # Get graph outputs which directly use the source + direct_outputs = [ + user for user in source.users if user.op == "output" + ] - # Insert clone node for placeholder to ensure placeholder is not a direct output - with gm.graph.inserting_after(placeholders[-1]): - cloned_placeholder = gm.graph.call_function( + # Insert a clone so the source is not returned directly + insertion_point = ( + placeholders[-1] if source.op == "placeholder" else source + ) + with gm.graph.inserting_after(insertion_point): + cloned_source = gm.graph.call_function( torch.ops.aten.clone.default, - args=(placeholder,), + args=(source,), ) + cloned_source.meta = source.meta.copy() - # Replace placeholder as output with cloned version + # Replace the direct output with the cloned version for output in direct_outputs: - output.replace_input_with(placeholder, cloned_placeholder) + output.replace_input_with(source, cloned_source) if modified_graph: gm = clean_up_graph_after_modifications(gm) diff --git a/tests/py/dynamo/lowering/test_aten_lowering_passes.py b/tests/py/dynamo/lowering/test_aten_lowering_passes.py index 51b31c8da1..79fc4a46a8 100644 --- a/tests/py/dynamo/lowering/test_aten_lowering_passes.py +++ b/tests/py/dynamo/lowering/test_aten_lowering_passes.py @@ -57,6 +57,33 @@ def forward(self, x, y): ) torch._dynamo.reset() + def test_folded_constant_as_output_is_cloned(self): + from torch_tensorrt.dynamo._settings import CompilationSettings + from torch_tensorrt.dynamo.lowering.passes.repair_input_as_output import ( + repair_input_as_output, + ) + + root = torch.nn.Module() + root.register_parameter( + "_frozen_param0", + torch.nn.Parameter(torch.zeros(8), requires_grad=False), + ) + graph = torch.fx.Graph() + folded_constant = graph.get_attr("_frozen_param0") + graph.output((folded_constant,)) + graph_module = torch.fx.GraphModule(root, graph) + + repaired_module = repair_input_as_output( + graph_module, CompilationSettings() + ) + first_output = repaired_module()[0] + first_output.add_(1) + second_output = repaired_module()[0] + + self.assertEqual(first_output, torch.ones(8)) + self.assertEqual(second_output, torch.zeros(8)) + self.assertEqual(repaired_module._frozen_param0, torch.zeros(8)) + class TestLoweringPassMembership(TestCase): def insert_at_end(self):