[HOP][print][dynamo]Add dynamo for hop print (pytorch#167571)

Following the previous implementation of HOP print, this continues to enable HOP print for dynamo so as to enable eager full graph and aot_eager backend for torch compile. At the end of the the implementation, the HOP print is able to enable stateful print without causing graph break. In the prior built in print, dynamo is able to reduce the graph break but unable to eliminate it. This enable the format-based printing for such purpose in dynamo.

Pull Request resolved: pytorch#167571
Approved by: https://github.com/angelayi
ghstack dependencies: pytorch#167016

Author: fxdawnn

test/higher_order_ops/test_print.py

@@ -3,12 +3,17 @@
 from unittest.mock import patch
 
 import torch
-from torch._dynamo.utils import counters
 from torch._functorch.aot_autograd import aot_export_module
 from torch.fx.experimental.proxy_tensor import make_fx
-from torch.testing._internal.common_utils import run_tests, TestCase
+from torch.testing._internal.common_utils import (
+    instantiate_parametrized_tests,
+    parametrize,
+    run_tests,
+    TestCase,
+)
 
 
+@instantiate_parametrized_tests
 class TestHopPrint(TestCase):
     def test_base_print(self):
         def f(x):
@@ -18,7 +23,6 @@ def f(x):
             torch._higher_order_ops.print("moo")
             return x
 
-        counters.clear()
         x = torch.randn(3, 3)
         with patch("sys.stdout", new_callable=io.StringIO) as mock_stdout:
             f(x)
@@ -33,7 +37,6 @@ def f(x):
             x = x * x
             return x
 
-        counters.clear()
         x = torch.randn(3, 3)
         with patch("sys.stdout", new_callable=io.StringIO) as mock_stdout:
             f(x)
@@ -184,6 +187,62 @@ def test_print_gen_schema(self):
             """print(str format_str) -> ()""",
         )
 
+    @parametrize("backend", ["eager", "aot_eager"])
+    def test_reorder_print_no_graph_break(self, backend):
+        def f(x):
+            x1 = x + x
+            torch._higher_order_ops.print("moo {x}", x=x1)
+            x2 = x1 * x1
+            torch._higher_order_ops.print("moo {x}", x=x2)
+            x3 = x2 + x2
+            return (x1, x3)
+
+        # Eager and aot_eager backend for dynamo tracing testing
+        x = torch.randn(3, 3)
+        opt_f = torch.compile(backend=backend, fullgraph=True)(f)
+        with patch("sys.stdout", new_callable=io.StringIO) as mock_stdout:
+            opt_out = opt_f(x)
+            printed_output = mock_stdout.getvalue().strip()
+            orig_out = f(x)
+
+        self.assertEqual(
+            printed_output,
+            f"moo {x * 2}\nmoo {x * 2 * x * 2}",
+        )
+        self.assertEqual(orig_out, opt_out)
+
+        x_new = torch.randn(2, 2)
+        with patch("sys.stdout", new_callable=io.StringIO) as mock_stdout:
+            opt_out = opt_f(x_new)
+            printed_output = mock_stdout.getvalue().strip()
+
+        self.assertEqual(
+            printed_output,
+            f"moo {x_new * 2}\nmoo {x_new * 2 * x_new * 2}",
+        )
+
+    @parametrize("backend", ["eager", "aot_eager"])
+    def test_constant_mutation(self, backend):
+        def f(x):
+            alist = [x]
+            alist.append(x + 1)
+            torch._higher_order_ops.print("moo {x}", x=alist[-1])
+            alist[0].sum().item()  # graph break
+            res = alist.pop()
+            torch._higher_order_ops.print("moo {x}", x=alist[-1])
+            res.sum().item()  # graph break
+            return res
+
+        inputs = (torch.tensor([1]),)
+        opt_f = torch.compile(backend=backend, fullgraph=True)(f)
+        with patch("sys.stdout", new_callable=io.StringIO) as mock_stdout:
+            opt_out = opt_f(*inputs)
+            printed_output = mock_stdout.getvalue().strip()
+            orig_out = f(*inputs)
+
+        self.assertEqual(printed_output, "moo tensor([2])\nmoo tensor([1])")
+        self.assertEqual(orig_out, opt_out)
+
 
 if __name__ == "__main__":
     run_tests()

torch/_dynamo/variables/higher_order_ops.py

@@ -2673,6 +2673,30 @@ def _call_function(
         )
 
 
+class PrintHigherOrderVariable(TorchHigherOrderOperatorVariable):
+    def _call_function(
+        self,
+        tx: "InstructionTranslator",
+        args: "list[VariableTracker]",
+        kwargs: "dict[str, VariableTracker]",
+    ) -> "VariableTracker":
+        from .builder import wrap_fx_proxy
+
+        args, kwargs = LazyVariableTracker.realize_all((args, kwargs))
+
+        args_proxy = [arg.as_proxy() for arg in args]
+        kwargs_proxy = {k: v.as_proxy() for k, v in kwargs.items()}
+        return wrap_fx_proxy(
+            tx=tx,
+            proxy=tx.output.create_proxy(
+                "call_function",
+                self.value,
+                args=tuple(args_proxy),
+                kwargs=kwargs_proxy,
+            ),
+        )
+
+
 class ExecutorchCallDelegateHigherOrderVariable(TorchHigherOrderOperatorVariable):
     def _call_function(
         self,
@@ -4537,6 +4561,7 @@ def make_error_msg(*args):
     "associative_scan": AssociativeScanHigherOrderVariable,
     "scan": ScanHigherOrderVariable,
     "call_torchbind": CallTorchbindHigherOrderVariable,
+    "print": PrintHigherOrderVariable,
     "wrap_with_set_grad_enabled": WrapWithSetGradEnabledHigherOrderVariable,
     "wrap_with_autocast": WrapWithAutocastHigherOrderVariable,
     "dynamo_bypassing_wrapper": DynamoBypassingWrapperHigherOrderVariable,