Interpreter for decomposing aten -> prims (#79989)

If an aten -> prim decomposition is needed *after* the initial trace
with make_fx, this interpreter can be used to perform the decomposition.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/79989
Approved by: https://github.com/SherlockNoMad

test/test_proxy_tensor.py

diff --git a/test/test_proxy_tensor.py b/test/test_proxy_tensor.py
index 09aabb1..d8d3150 100644
--- a/test/test_proxy_tensor.py
+++ b/test/test_proxy_tensor.py
@@ -9,8 +9,9 @@
 from torch.testing._internal.common_methods_invocations import op_db, wrapper_set_seed
 from torch._subclasses.fake_tensor import DynamicOutputShapeException
 
+from torch._decomp import decomposition_table
 from torch.testing._internal.common_device_type import ops
-from torch.fx.experimental.proxy_tensor import make_fx
+from torch.fx.experimental.proxy_tensor import make_fx, DecompositionInterpreter
 
 # Copied from functorch
 def xfail(op_name, variant_name='', *, device_type=None, dtypes=None):
@@ -156,6 +157,35 @@
         self.assertTrue(all([isinstance(node.target, torch._ops.OpOverload)
                              for node in traced.graph.nodes if node.op == 'call_function']))
 
+    def test_decomposition_interpreter(self):
+        def fn(x):
+            return torch.nn.functional.silu(x)
+
+        x = torch.rand((4, 4))
+        fx_module = make_fx(fn, decomposition_table=None)(x)
+
+        found_silu = False
+        for n in fx_module.graph.nodes:
+            if n.target == torch.ops.aten.silu or n.target == torch.ops.aten.silu.default:
+                found_silu = True
+
+        self.assertTrue(found_silu)
+
+        new_graph = torch.fx.Graph()
+        silu_decomp_table = {torch.ops.aten.silu.default: decomposition_table[torch.ops.aten.silu.default]}
+        DecompositionInterpreter(
+            fx_module,
+            new_graph=new_graph,
+            decomposition_table=silu_decomp_table,
+        ).run(x)
+
+        decomposed_module = torch.fx.GraphModule(fx_module, new_graph)
+
+        for n in decomposed_module.graph.nodes:
+            self.assertTrue(n.target != torch.ops.aten.silu)
+            self.assertTrue(n.target != torch.ops.aten.silu.default)
+
+        self.assertEqual(fx_module(x), decomposed_module(x))
 
 make_fx_failures = {
     # unknown