Fix get_free_symbol_uses for several nodes. (#160134)

laithsakka · pytorchmergebot · commit db78943a1ca1 · 2025-08-09T18:15:46.000Z
get_free_symbol_uses is used to know what unbacked symbols are used by a given node. not having correct get_free_symbol_uses defined properly leads to : 1. eliminating of some nodes due to not detection of any users. (See the added unit test) 2. Incorrect topological sort. Fix get_free_symbol_uses , NopKernel , ConcarKernel, InputsKerenl, external kernel. for ComputedBuffer with NonOwningLayout its interesting case. when layout is NonOwningLayout we need to access the actual view op base layout and use detect symbols in it. Because when we codegen the ComputedBuffer we uses those symbols. Pull Request resolved: #160134 Approved by: https://github.com/bobrenjc93
diff --git a/test/test_dynamic_shapes.py b/test/test_dynamic_shapes.py
@@ -3616,6 +3616,17 @@ def func3(x, y):
     def test_unbacked_select_index_cpp_wrapper(self):
         self.test_unbacked_select_index()
 
+    @torch._dynamo.config.patch("capture_scalar_outputs", True)
+    def test_unbacked_select2(self):
+        def f(idx, x):
+            x = x.select(0, idx.item())
+            return x @ x
+
+        x = torch.randn(3, 3, 3)
+        idx = torch.tensor(1, dtype=torch.int64)
+        out = torch.compile(f)(idx, x)
+        self.assertEqual(out, f(idx, x))
+
 
 instantiate_parametrized_tests(TestUnbacked)
 
diff --git a/torch/_inductor/ir.py b/torch/_inductor/ir.py
@@ -4443,14 +4443,30 @@ def get_free_symbol_uses(
         # unusual reason: we only need accurate dependencies for item() call,
         # but it's impossible to end up with a reduction over i0 from an
         # item() call without a regular non-reduction buffer first.
-        return (
+
+        result = (
             get_free_symbols(self.get_size(), unbacked_only)
             | get_free_symbols(self.get_stride(), unbacked_only)
             | get_free_symbols(self.get_offset(), unbacked_only)
             | self.data.get_free_symbol_uses(unbacked_only)
             | self.get_read_writes().get_free_symbol_uses(unbacked_only)
         )
 
+        if isinstance(self.layout, NonOwningLayout):
+            assert isinstance(self.layout.view, ReinterpretView)
+            box = self.layout.view.data
+            assert isinstance(box, StorageBox), type(box)
+            input_buffer = box.data
+            assert isinstance(input_buffer, Buffer), type(box)
+            result = (
+                result
+                | get_free_symbols(input_buffer.get_size(), unbacked_only)
+                | get_free_symbols(input_buffer.get_stride(), unbacked_only)
+                | get_free_symbols(input_buffer.get_offset(), unbacked_only)
+            )
+
+        return result
+
     def make_loader(self) -> Callable[[Sequence[Expr]], OpsValue]:
         if (
             not self.get_reduction_type()
@@ -5126,6 +5142,18 @@ def get_read_writes(self) -> dependencies.ReadWrites:
     def get_reads(self) -> OrderedSet[Dep]:
         return self.get_read_writes().reads
 
+    def get_free_symbol_uses(
+        self, unbacked_only: bool = False
+    ) -> OrderedSet[sympy.Symbol]:
+        r = OrderedSet[sympy.Symbol]()
+        for inp in self.inputs:
+            if isinstance(inp, IRNode):
+                r |= inp.get_free_symbol_uses(unbacked_only)
+            else:
+                for inner_inp in inp:
+                    r |= inner_inp.get_free_symbol_uses(unbacked_only)
+        return r
+
     @classmethod
     def unwrap_storage_for_input(cls, x: IRNode) -> IRNode:
         if isinstance(x, TensorBox):
@@ -5172,6 +5200,11 @@ def is_no_op(self) -> bool:
     def get_reads(self) -> OrderedSet[Dep]:
         return OrderedSet()
 
+    def get_free_symbol_uses(
+        self, unbacked_only: bool = False
+    ) -> OrderedSet[sympy.Symbol]:
+        return InputsKernel.get_free_symbol_uses(self, unbacked_only)
+
 
 class ConcatKernel(NopKernel):
     """
@@ -5326,6 +5359,11 @@ def can_realize_into_without_copy(
             and not isinstance(src.data, ExternKernelAlloc)
         )
 
+    def get_free_symbol_uses(
+        self, unbacked_only: bool = False
+    ) -> OrderedSet[sympy.Symbol]:
+        return NopKernel.get_free_symbol_uses(self, unbacked_only)
+
     @classmethod
     def realize_into(cls, src: IRNode, dst: IRNode) -> IRNode:
         # Attempt to turn this into a ReinterpretView rather than assert.
@@ -6221,12 +6259,10 @@ def canonicalize(self) -> tuple[Expr, Sequence[Expr]]:
     def get_free_symbol_uses(
         self, unbacked_only: bool = False
     ) -> OrderedSet[sympy.Symbol]:
-        # NB: It's not necessary to check regular inputs as we automatically
-        # have dependencies on them
         maybe_get_symbols = (
             maybe_free_unbacked_symbols if unbacked_only else maybe_free_symbols
         )
-        r = OrderedSet[sympy.Symbol]()
+        r = InputsKernel.get_free_symbol_uses(self, unbacked_only)
         for arg in self.constant_args:
             r |= maybe_get_symbols(arg)
         for arg in self.kwargs.values():
