Add Current Mask Var To CSE Cache Key (#140838)

eellison · pytorchmergebot · commit eff22171d22d · 2024-11-20T00:55:56.000Z
This torch.cat kernel has multiple subblocks which load from the same input. We were incorrectly reusing the mask vars from the first load for the second load. Pull Request resolved: #140838 Approved by: https://github.com/jansel ghstack dependencies: #140841
diff --git a/test/inductor/test_cuda_repro.py b/test/inductor/test_cuda_repro.py
@@ -1494,6 +1494,32 @@ def foo(x, y):
             device_stats2["active.all.peak"] <= device_stats["active.all.peak"]
         )
 
+    def test_repeated_masked_load(self):
+        target_size = (8, 2)
+        mem_eff_temporal_upsampling_interp_chunks = 2
+        from functorch.einops import rearrange
+
+        x = torch.randn(1, 8, 12, 12, 4, dtype=torch.float16, device="cuda")
+        x = x.permute(0, 1, 4, 2, 3)  # make non-contiguous
+        x = rearrange(x, "b c t h w -> b c t (h w)")
+
+        def interpolate_chunked(x):
+            # chunk along c
+            chunks = x.chunk(chunks=mem_eff_temporal_upsampling_interp_chunks, dim=1)
+            r = []
+            for t in chunks:
+                r.append(
+                    torch.nn.functional.interpolate(
+                        t.float(), size=target_size, mode="nearest"
+                    ).to(t.dtype)
+                )
+            out_chunked = torch.cat(r, dim=1)
+            return out_chunked
+
+        out_eager = interpolate_chunked(x)
+        out_compiled = torch.compile(interpolate_chunked)(x)
+        self.assertEqual(out_eager, out_compiled)
+
 
 if __name__ == "__main__":
     from torch._inductor.test_case import run_tests
diff --git a/torch/_inductor/codegen/common.py b/torch/_inductor/codegen/common.py
@@ -1693,7 +1693,7 @@ def __init__(
     ):
         self.prefix = prefix
         self.suffix = suffix
-        self.cache = {}
+        self._cache = {}
         self.name_prefix = name_prefix
         self.store_cache = store_cache or {}
         self.reduction_cache = reduction_cache or {}
@@ -1706,11 +1706,11 @@ def invalidate(self, keep_vars: OrderedSet[str]):
             if tmp not in keep_vars:
                 del self.store_cache[name]
                 self.invalidated_stores.add(name)
-        self.cache = {k: v for k, v in self.cache.items() if v in keep_vars}
+        self._cache = {k: v for k, v in self._cache.items() if v in keep_vars}
 
     def clone(self):
         # Note(fdrocha): reduction_cache is not being cloned, not sure if this is intentional
-        return CSE(
+        return type(self)(
             prefix=self.prefix,
             suffix=self.suffix,
             name_prefix=self.name_prefix,
@@ -1719,6 +1719,22 @@ def clone(self):
             varname_map=self.varname_map,
         )
 
+    def augment_key(self, cache_key: object) -> object:
+        "Override this method to augment cache key with backend specifics"
+        return cache_key
+
+    def put(self, cache_key: object, val: CSEVariable) -> None:
+        self._cache[self.augment_key(cache_key)] = val
+
+    def contains(self, cache_key) -> bool:
+        return self.augment_key(cache_key) in self._cache
+
+    def try_get(self, cache_key: object) -> Optional[CSEVariable]:
+        return self._cache.get(self.augment_key(cache_key), None)
+
+    def get(self, cache_key: object) -> CSEVariable:
+        return self._cache[self.augment_key(cache_key)]
+
     def generate(
         self,
         buffer: IndentedBuffer,
@@ -1747,10 +1763,10 @@ def generate(
         else:
             assert isinstance(expr, str)
             cache_key = expr
-        var = self.cache.get(cache_key, None)
+        var = self.try_get(cache_key)
         if not var:
             var = self.newvar(bounds, dtype)
-            self.cache[cache_key] = var
+            self.put(cache_key, var)
             if write:
                 if V.kernel.current_node:
                     V.kernel.current_node.codegen_originating_info(
@@ -2078,7 +2094,7 @@ def set_current_node(self, node):
     def swap_buffers(self, lb, cb=None, sb=None):
         def scope_cse(cse):
             new_cse = cse.clone()
-            new_cse.cache = ScopedDict(cse.cache)
+            new_cse._cache = ScopedDict(cse._cache)
             new_cse.reduction_cache = ScopedDict(cse.reduction_cache)
             new_cse.store_cache = ScopedDict(cse.store_cache)
             return new_cse
diff --git a/torch/_inductor/codegen/cpp.py b/torch/_inductor/codegen/cpp.py
@@ -776,8 +776,8 @@ def copysign(x, y):
     @staticmethod
     def frexp(x):
         cache_keys = f"frexp({x})[0]", f"frexp({x})[1]"
-        if all(cache_key in V.kernel.cse.cache for cache_key in cache_keys):
-            return tuple(V.kernel.cse.cache[cache_key] for cache_key in cache_keys)
+        if all(V.kernel.cse.try_get(cache_key) is not None for cache_key in cache_keys):
+            return tuple(V.kernel.cse.try_get(cache_key) for cache_key in cache_keys)
 
         code = BracesBuffer()
         exponent = V.kernel.cse.newvar()
@@ -787,7 +787,7 @@ def frexp(x):
         V.kernel.compute.splice(code)
         cse_vars = (mantissa, exponent)
         for cache_key, cse_var in zip(cache_keys, cse_vars):
-            V.kernel.cse.cache[cache_key] = cse_var
+            V.kernel.cse.put(cache_key, cse_var)
         return mantissa, exponent
 
     @staticmethod
@@ -1574,8 +1574,8 @@ def index_expr(expr, dtype):
     @staticmethod
     def frexp(x):
         cache_keys = f"frexp({x})[0]", f"frexp({x})[1]"
-        if all(cache_key in V.kernel.cse.cache for cache_key in cache_keys):
-            return tuple(V.kernel.cse.cache[cache_key] for cache_key in cache_keys)
+        if all(V.kernel.cse.try_get(cache_key) is not None for cache_key in cache_keys):
+            return tuple(V.kernel.cse.try_get(cache_key) for cache_key in cache_keys)
 
         cdtype = DTYPE_TO_CPP[x.dtype]
         size = V.kernel.tail_size if V.kernel.tail_size else V.kernel.tiling_factor
@@ -1625,7 +1625,7 @@ def frexp(x):
         V.kernel.compute.splice(code)
         cse_vars = (mantissa, exponent)
         for cache_key, cse_var in zip(cache_keys, cse_vars):
-            V.kernel.cse.cache[cache_key] = cse_var
+            V.kernel.cse.put(cache_key, cse_var)
         return mantissa, exponent
 
     @classmethod
@@ -2198,7 +2198,7 @@ def get_to_dtype_expr(self, src, dtype, src_dtype):
 
     def cache_dtype_convert(self, dst, dst_dtype, src, src_dtype):
         expr = self.get_to_dtype_expr(src, dst_dtype, src_dtype)
-        self.cse.cache[expr] = dst
+        self.cse.put(expr, dst)
 
 
 class CppVecKernel(CppKernel):
@@ -3086,11 +3086,11 @@ def gen_transposed_tile_load_store(self, name, var, index, is_store):
             )
         if is_store:
             tile_var = self.cse.newvar()
-        elif load_or_store not in self.cse.cache:
+        elif not self.cse.contains(load_or_store):
             tile_var = self.cse.generate(self.preloads, load_or_store, write=False)
         else:
             need_define = False
-            tile_var = self.cse.cache[load_or_store]
+            tile_var = self.cse.get(load_or_store)
 
         if need_define:
             cpp_dtype = DTYPE_TO_CPP[dtype]
diff --git a/torch/_inductor/codegen/triton.py b/torch/_inductor/codegen/triton.py
@@ -1219,6 +1219,7 @@ def masked(mask, body, other):
                     need_where = True
 
         value = None if need_where else other
+
         with V.kernel.mask_loads(mask, value=value) as new_mask:
             result = body()
 
@@ -1250,15 +1251,15 @@ def load_seed(name, offset):
     @staticmethod
     def frexp(x):
         cache_key = f"frexp({x})"
-        if cache_key in V.kernel.cse.cache:
-            return V.kernel.cse.cache[cache_key]
+        if cse_val := V.kernel.cse.try_get(cache_key):
+            return cse_val
 
         mantissa = V.kernel.cse.newvar(dtype=x.dtype)
         exponent = V.kernel.cse.newvar(dtype=x.dtype)
         V.kernel.compute.writeline(
             f"{mantissa}, {exponent} = triton_helpers.frexp({x})"
         )
-        V.kernel.cse.cache[cache_key] = (mantissa, exponent)
+        V.kernel.cse.put(cache_key, (mantissa, exponent))
         return (mantissa, exponent)
 
 
@@ -1370,6 +1371,19 @@ def __getitem__(self, item):
         return self.config[item]
 
 
+class TritonCSE(CSE):
+    """
+    Subclasses CSE to apply the current load mask to the cache key to avoid CSEing
+    variables across separate masked blocks.
+    """
+
+    def augment_key(self, cache_key: object) -> object:
+        if mask := V.kernel._load_mask:
+            return (cache_key, mask.name)
+        else:
+            return cache_key
+
+
 class TritonKernel(SIMDKernel):
     overrides = TritonKernelOverrides  # type: ignore[assignment]
     helper_functions: HelperFunctions
@@ -1387,6 +1401,7 @@ def __init__(
         self.optimize_mask: bool = optimize_mask
         self.fixed_config = fixed_config
         super().__init__(*groups, **kwargs)
+        self.cse = TritonCSE(self.newvar_prefix, self.suffix)
         self.post_loop_combine: IndentedBuffer = IndentedBuffer()
         self.post_loop_store: IndentedBuffer = IndentedBuffer()
         self.outside_loop_vars: OrderedSet[Any] = OrderedSet()
@@ -2603,16 +2618,16 @@ def csv(values):
         def cse_multiple(line, values, masks, dtypes):
             n = len(values)
             cache_keys = [f"{line}, {i}, {masks}" for i in range(n)]
-            if all(cache_key in self.cse.cache for cache_key in cache_keys):
-                return [self.cse.cache[cache_key] for cache_key in cache_keys]
+            if all(self.cse.contains(cache_key) for cache_key in cache_keys):
+                return [self.cse.get(cache_key) for cache_key in cache_keys]
             result_vars = [self.cse.newvar(dtype=_dtype) for _dtype in dtypes]
             self.compute.writeline(
                 f"{csv(result_vars)} = {line}",
             )
             for result_var, cache_key in zip(result_vars, cache_keys):
                 if masks:
                     result_var.mask_vars = masks  # type: ignore[attr-defined]
-                self.cse.cache[cache_key] = result_var
+                self.cse.put(cache_key, result_var)
             return tuple(result_vars)
 
         partial_scan_vars = cse_multiple(
@@ -2690,16 +2705,16 @@ def csv(values):
 
         def cse_multiple(line, n, masks, dtypes):
             cache_keys = [f"{line}, {i}, {masks}" for i in range(n)]
-            if all(cache_key in self.cse.cache for cache_key in cache_keys):
-                return [self.cse.cache[cache_key] for cache_key in cache_keys]
+            if all(self.cse.contains(cache_key) for cache_key in cache_keys):
+                return [self.cse.get(cache_key) for cache_key in cache_keys]
             result_vars = [self.cse.newvar(dtype=dtypes[i]) for i in range(n)]  # type: ignore[attr-defined]
             self.compute.writeline(
                 f"{csv(result_vars)} = {line}",
             )
             for result_var, cache_key in zip(result_vars, cache_keys):
                 if masks:
                     result_var.mask_vars = masks  # type: ignore[attr-defined]
-                self.cse.cache[cache_key] = result_var
+                self.cse.put(cache_key, result_var)
             return tuple(result_vars)
 
         assert self.range_trees[-1].prefix == "r"
