[MPSInductor] Fix nested loop var elimination (pytorch#156566)

As reduction resuts must be kept around
Add regression test that is specific for this issue

Fixes pytorch#156426

Pull Request resolved: pytorch#156566
Approved by: https://github.com/dcci

Author: malfet

test/inductor/test_torchinductor.py

@@ -5530,6 +5530,13 @@ def fn(x):
                 (torch.randn([1, 2, 4, 8]),),
             )
 
+    def test_var_mean_div_by(self):
+        def fn(x):
+            var, mean = torch.var_mean(x, dim=2, keepdim=True)
+            return x / var, var, mean
+
+        self.common(fn, (torch.rand([1, 17, 2048]),))
+
     def test_var_correction(self):
         def fn(x):
             dim = -1

torch/_inductor/codegen/mps.py

@@ -754,7 +754,16 @@ def codegen_body(self) -> None:
                 self.body.splice(self.compute)
             self.body.writeline("}" * len(self.multistage_reduction_entry))
             # Invalidate variables instantiated inside loop
-            self.cse.invalidate(OrderedSet(self.cse.reduction_cache.values()))
+            # But results of reduction alive. Reduction cache values can be
+            # either CSEVariable or tuple of CSEVariables, in which case all
+            # variables in the tuple must be preserved
+            self.cse.invalidate(
+                OrderedSet(
+                    v
+                    for item in self.cse.reduction_cache.values()
+                    for v in (item if isinstance(item, tuple) else (item,))
+                )
+            )
             # And loop codegen
             while self.multistage_reduction_entry:
                 self.multistage_reduction_entry.pop().cache_clear()