[AOTI] Fix a unbacked symint retrieve bug (#134670)

desertfire · pytorchmergebot · commit 50d1e3707986 · 2024-09-05T11:34:14.000Z
Summary: Fix #134081. When a unbacked symint is computed as the shape of a tensor from a tuple, generated C++ code needs to use std::get<> to extract the tensor. Differential Revision: [D62142113](https://our.internmc.facebook.com/intern/diff/D62142113) Pull Request resolved: #134670 Approved by: https://github.com/angelayi, https://github.com/22quinn, https://github.com/chenyang78
diff --git a/test/inductor/test_aot_inductor.py b/test/inductor/test_aot_inductor.py
@@ -3421,6 +3421,20 @@ def forward(self, x, y):
                     count,
                 ).run(code)
 
+    def test_size_from_multi_output(self):
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.relu = torch.nn.ReLU()
+
+            def forward(self, x):
+                _x, _i = torch.unique(x, sorted=True, return_inverse=True)
+                _x = _x.clone().detach()
+                return self.relu(_x), _i
+
+        example_inputs = (torch.randn(8, device=self.device),)
+        self.check_model(Model(), example_inputs)
+
 
 common_utils.instantiate_parametrized_tests(AOTInductorTestsTemplate)
 
@@ -3609,6 +3623,7 @@ def fail_non_abi_compatible_cuda(is_skip=False):
     "test_custom_op_missing_arg_with_default_value": fail_minimal_arrayref_interface(
         is_skip=True
     ),
+    "test_size_from_multi_output": fail_stack_allocation(is_skip=True),
 }
 
 # test_failures, xfail by default, set is_skip=True to skip
diff --git a/torch/_inductor/codegen/cpp_wrapper_cuda.py b/torch/_inductor/codegen/cpp_wrapper_cuda.py
@@ -63,7 +63,7 @@ def _new_line(self, line):
 
 class DeferredCudaDefaultGrid:
     """
-    A marker to
+    A container for the default grid, which may be used by DeferredCudaGridLine
     """
 
     def __init__(
diff --git a/torch/_inductor/ir.py b/torch/_inductor/ir.py
@@ -5864,7 +5864,11 @@ def go(expr, keypath):
                 elif isinstance(keypath[0], CallMethodKey):
                     return go(f"{expr}.{keypath[0].name}()", keypath[1:])
                 elif isinstance(keypath[0], pytree.SequenceKey):
-                    return go(f"{expr}[{keypath[0].idx}]", keypath[1:])
+                    return (
+                        go(f"std::get<{keypath[0].idx}>({expr})", keypath[1:])
+                        if V.graph.cpp_wrapper
+                        else go(f"{expr}[{keypath[0].idx}]", keypath[1:])
+                    )
                 elif isinstance(keypath[0], DivideByKey):
                     # TODO: need to assert divisibility
                     # TODO: this is invalid C++ codegen
