[Inductor-FX] Support symbol and dynamic scalar graph inputs and outputs (pytorch#163596)

# Problems
This PR fixes a few edge cases that the FX converter missed related to dynamic shapes.

1. Inductor graphs can sometimes take `sympy.Symbol` inputs. We have logic to convert these to FX placeholder nodes. However, this logic did not update the `self.expr_to_proxy` table mapping symbols to proxy nodes. (There was existing logic to do this for `ir.TensorBox` inputs, but not `sympy.Symbol`.) This caused sympy tracing to fail when these symbol inputs were used in other expressions.

2. We lacked codegen for `ShapeAsConstantBuffer`. This IR node is seen when the graph input or output is a scalar computed from dynamic shapes.

# Fixes
a. Update `self.expr_to_proxy` when generating placeholders for `sympy.Symbol` inputs. Change `SymbolBuffer.get_example` to convert the symbol to a `torch.SymInt`, so we can populate `meta["val"]` correctly and use the value in other computations.
b. Support `ShapeAsConstantBuffer` by tracing the sympy expression.
c. Move output generation inside the metadata hook, allowing us to populate `meta["val"]` for the nodes computing `ShapeAsConstantBuffer`.

# Test plan
Added several new CI tests:
 1. `torch.cond` with dynamic shapes. This exposes both issues, as the predicate is a `ShapeAsConstantBuffer` and one of the subgraphs uses a symbol input, due to the closure. Also tests when the parent and subgraphs have different input shapes.
 2. Output dynamic shape scalar. This tests `ShapeAsConstantBuffer` as an output.

Pull Request resolved: pytorch#163596
Approved by: https://github.com/angelayi, https://github.com/jansel

Author: blaine-rister

test/inductor/test_fxir_backend.py

@@ -990,6 +990,48 @@ def forward(self, arg0_1, arg1_1, arg2_1):
     return [buf1, buf2]""",  # noqa: B950
         )
 
+    @parametrize("length", (4, 8))
+    def test_cond_dynamic_shape_pred_scalar_closure(self, length: int):
+        """
+        Test cond using a predicate computed from dynamic shapes.
+        Also test a dynamic scalar computed outside the branches.
+        """
+
+        class M(torch.nn.Module):
+            def forward(self, x, y):
+                z = x.reshape(-1)
+                a = y.shape[0]
+
+                def true_fn(x):
+                    return x + a
+
+                def false_fn(x):
+                    return true_fn(x) / 2
+
+                return torch.cond(x.shape[0] > 5, true_fn, false_fn, (z,))
+
+        (x, y) = [
+            torch.randn(shape, device=self.device)
+            for shape in [(length // 2,) * 2, (length,)]
+        ]
+        dynamic_shapes = {
+            "x": {0: Dim.DYNAMIC},
+            "y": {0: Dim.DYNAMIC},
+        }
+        self.check(M(), (x, y), dynamic_shapes=dynamic_shapes)
+
+    def test_dynamic_scalar_output(self):
+        """
+        Test an output scalar from dynamic shapes.
+        """
+
+        class M(torch.nn.Module):
+            def forward(self, x):
+                return x.shape[0] * 3
+
+        x = torch.randn(7, device=self.device)
+        self.check(M(), (x,), dynamic_shapes=({0: Dim.DYNAMIC},))
+
 
 if __name__ == "__main__":
     from torch._inductor.test_case import run_tests

torch/_inductor/codegen/wrapper_fxir.py

@@ -23,7 +23,11 @@
 from torch._inductor.codecache import LambdaFuture, PyCodeCache
 from torch._inductor.runtime.triton_heuristics import CachingAutotuner
 from torch._inductor.select_algorithm import extern_kernels  # noqa: F401
-from torch._inductor.utils import convert_shape_to_symint, sympy_product
+from torch._inductor.utils import (
+    convert_shape_to_symint,
+    convert_to_symint,
+    sympy_product,
+)
 from torch._inductor.virtualized import V
 from torch._library.triton import wrap_triton
 from torch.fx import GraphModule
@@ -89,8 +93,10 @@ class SymbolBuffer(CodegenSymbol):
     def get_name(self) -> str:
         return str(self.symbol)
 
-    def get_example(self) -> Union[torch.Tensor, sympy.Symbol]:
-        return self.symbol
+    def get_example(self) -> Union[torch.Tensor, torch.SymInt]:
+        sym_int = convert_to_symint(self.symbol)
+        assert isinstance(sym_int, torch.SymInt)
+        return sym_int
 
 
 CodegenBuffer = Union[BufferLike, SymbolBuffer]
@@ -386,6 +392,13 @@ def _get_buffer(self, node: ir.IRNode) -> CodegenBuffer:
         else:
             raise NotImplementedError(f"Unable to extract buffer from node: {node}")
 
+    def _generate_size_proxy(
+        self, node: torch.fx.Node, expr: sympy.Expr
+    ) -> torch.fx.Proxy:
+        proxy = torch.fx.Proxy(node, tracer=self.tracer)
+        self.expr_to_proxy[expr] = proxy
+        return proxy
+
     def _generate_graph_inputs(self) -> None:
         """
         Converts graph inputs to FX placeholders.
@@ -398,15 +411,22 @@ def _generate_graph_inputs(self) -> None:
                 continue
 
             # Introduce a new symbol for constant inputs.
+            is_constant = isinstance(ir_node, (int, float, sympy.Integer, sympy.Float))
             buffer = (
                 SymbolBuffer(sympy.Symbol(name, is_integer=True))
-                if isinstance(ir_node, (int, float, sympy.Integer, sympy.Float))
+                if is_constant
                 else self._get_buffer(ir_node)
             )
             placeholder_node = self.gm.graph.placeholder(buffer.get_name())
-            placeholder_node.meta["val"] = buffer.get_example()
+            placeholder_node.meta["val"] = (
+                ir_node if is_constant else buffer.get_example()
+            )
             self._record_allocation(buffer, placeholder_node)
 
+            # Record symbol definitions for dynamic shapes.
+            if isinstance(ir_node, sympy.Symbol):
+                self._generate_size_proxy(placeholder_node, ir_node)
+
     def _generate_graph_input_shapes(self) -> None:
         """
         Generate nodes creating symints that are part of graph input
@@ -421,8 +441,7 @@ def _codegen_symbol(
         ) -> None:
             def codegen_proxy() -> torch.fx.Proxy:
                 size_node = self.gm.graph.call_function(target, (base_node, dim))
-                size_proxy = torch.fx.Proxy(size_node, tracer=self.tracer)
-                self.expr_to_proxy[sym_or_exp] = size_proxy
+                size_proxy = self._generate_size_proxy(size_node, sym_or_exp)
                 return size_proxy
 
             if isinstance(sym_or_exp, sympy.Symbol):
@@ -475,9 +494,7 @@ def codegen_proxy() -> torch.fx.Proxy:
                     undefined_symbol_expr
                 ]
 
-        for node in V.graph.module.graph.find_nodes(op="placeholder"):  # type: ignore[operator, union-attr]
-            name = node.name
-            ir_node = self.graph_inputs.get(name)
+        for ir_node in self.graph_inputs.values():
             if isinstance(ir_node, ir.TensorBox):
                 buffer = self._get_buffer(ir_node)
                 placeholder_node = self.buffer_to_node[buffer.get_name()]
@@ -504,6 +521,10 @@ def _generate_buffer(self, node: ir.IRNode) -> Optional[torch.fx.Node]:
         Does nothing if no such transformations are present.
         """
 
+        if isinstance(node, ir.ShapeAsConstantBuffer):
+            # Generate FX nodes to compute the shape expression.
+            return self._sympy_interp(node.expr).node
+
         def generate_to_buffer(node: ir.IRNode) -> Optional[BufferLike]:
             if isinstance(node, (ir.Buffer, WorkspaceArg)):
                 return node
@@ -539,7 +560,9 @@ def generate_to_buffer(node: ir.IRNode) -> Optional[BufferLike]:
         buffer = generate_to_buffer(node)
         return self.buffer_to_node[buffer.get_name()] if buffer is not None else None
 
-    def _generate_output(self) -> None:
+    def _generate_outputs(
+        self,
+    ) -> Union[Optional[torch.fx.Node], list[Optional[torch.fx.Node]]]:
         """
         Generate FX IR for graph outputs.
         """
@@ -554,7 +577,7 @@ def _generate_output(self) -> None:
             else output_nodes
         )
 
-        self.gm.graph.output(output_value)
+        return output_value
 
     def _generate_subgm_getattrs(self) -> None:
         """
@@ -614,7 +637,9 @@ def generate(self) -> torch.fx.GraphModule:
                         )
                     )
 
-        self._generate_output()
+            output = self._generate_outputs()
+
+        self.gm.graph.output(output)
         self.gm.recompile()
         return self.gm
 

torch/_inductor/ir.py

@@ -4220,7 +4220,7 @@ def get_name(self) -> str:
         assert self.name, self
         return self.name
 
-    def get_example(self) -> Union[torch.Tensor, sympy.Symbol]:
+    def get_example(self) -> Union[torch.Tensor, torch.SymInt]:
         if isinstance(self.layout, Layout):
             return self.layout.get_example()
         raise NotImplementedError(type(self.layout).__name__)