[inductor] allow clone cse cache during vectorized indirect load

ghstack-source-id: 3d5eb12fdc823729df0a6b67f8bf042c5d6c9b73
Pull Request resolved: #124597

test/inductor/test_cpu_repro.py

@@ -3648,6 +3648,43 @@ def forward(self, x):
         x = torch.randn(1, 4, 2, 2)
         self.common(fn, (x,))
 
+    def test_vec_indirect_load_cse_cache(self):
+        # https://github.com/pytorch/pytorch/issues/123502.
+        from math import inf
+
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, arg0_1):
+                full_default = torch.ops.aten.full.default([209985], 1)
+                select = torch.ops.aten.select.int(arg0_1, 0, 0)
+                select_1 = torch.ops.aten.select.int(arg0_1, 0, 1)
+                view = torch.ops.aten.reshape.default(select_1, [-1])
+                expand = torch.ops.aten.expand.default(view, [209985])
+                full_default_1 = torch.ops.aten.full.default([10000], 0)
+                scatter_add = torch.ops.aten.scatter_add.default(
+                    full_default_1, 0, expand, full_default
+                )
+                pow_1 = torch.ops.aten.pow.Tensor_Scalar(scatter_add, -0.5)
+                eq = torch.ops.aten.eq.Scalar(pow_1, inf)
+                full_default_2 = torch.ops.aten.full.default([], 0.0)
+                where = torch.ops.aten.where.self(eq, full_default_2, pow_1)
+                index = torch.ops.aten.index.Tensor(where, [select])
+                index_1 = torch.ops.aten.index.Tensor(where, [select_1])
+                mul_1 = torch.ops.aten.mul.Tensor(index, index_1)
+                return (mul_1,)
+
+        fn = Model()
+        x = torch.zeros(2, 209985).to(torch.int64)
+        _fn_opt = torch.compile()(fn)
+        _, code = run_and_get_cpp_code(_fn_opt, x)
+        FileCheck().check_count(
+            "return at::vec::VectorizedN<int64_t,2>::loadu(tmpbuf.data(), 16);",
+            2,
+            exactly=True,
+        ).run(code)
+
 
 if __name__ == "__main__":
     from torch._inductor.test_case import run_tests

torch/_inductor/codegen/common.py

@@ -1348,9 +1348,11 @@ def set_current_node(self, node):
             self.current_node = prior
 
     @contextlib.contextmanager
-    def swap_buffers(self, lb, cb=None, sb=None):
+    def swap_buffers(self, lb, cb=None, sb=None, cse_cache=None):
         if cb is None:
             cb = lb
+        if cse_cache is None:
+            cse_cache = {}
         loads = self.loads
         compute = self.compute
         stores = self.stores
@@ -1359,6 +1361,7 @@ def swap_buffers(self, lb, cb=None, sb=None):
         self.compute = cb
         self.stores = sb
         self.cse = cse.clone()
+        self.cse.cache = cse_cache
         try:
             yield
         finally:

torch/_inductor/codegen/cpp.py

@@ -2318,7 +2318,7 @@ def vec_to_array(vec_var: CppCSEVariable) -> CppCSEVariable:
             assert vec_var.is_vec
             code = BracesBuffer()
             code.writeline("[&]")
-            with self.swap_buffers(code), code.indent():
+            with self.swap_buffers(code, cse_cache=self.cse.cache), code.indent():
                 vec_dtype = vec_var.dtype
                 assert vec_dtype is not None
                 if vec_dtype == torch.bool:
@@ -2339,7 +2339,7 @@ def vec_to_array(vec_var: CppCSEVariable) -> CppCSEVariable:
         assert opt_ctx is not None
         code = BracesBuffer()
         code.writeline("[&]")
-        with self.swap_buffers(code), code.indent():
+        with self.swap_buffers(code, cse_cache=self.cse.cache), code.indent():
             result_size = get_result_size(dtype)
             result_declare = (
                 f"__at_align__ std::array<{DTYPE_TO_CPP[dtype]}, {result_size}> tmpbuf;"