Reduction cache (pytorch#487)

* enable cse for reductions * self.cse
chunyuan-w · Jun 30, 2022 · b22b6a5 · b22b6a5
1 parent b3c3b9e
commit b22b6a5
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 20 deletions.
diff --git a/torchinductor/codegen/common.py b/torchinductor/codegen/common.py
@@ -429,12 +429,14 @@ def __init__(
         name_prefix="tmp",
         iter_buffers=None,
         store_cache=None,
+        reduction_cache=None,
     ):
         self.prefix = prefix
         self.suffix = suffix
         self.cache = {}
         self.name_prefix = name_prefix
         self.store_cache = store_cache or {}
+        self.reduction_cache = reduction_cache or {}
         self.iter_buffer_ids = iter_buffers or itertools.count()
 
     def invalidate(self, keep_vars: typing.Set[str]):

diff --git a/torchinductor/codegen/triton.py b/torchinductor/codegen/triton.py
@@ -592,29 +592,34 @@ def reduction(self, name, dtype, reduction_type, index, value):
 
         dim = len(self.range_trees) - 1
         result_var = self.cse.newvar()
-        accumulator = f"_{result_var}"
-        self.body.writeline(
-            f"{accumulator} = tl.zeros({self.dense_size_str()}, {triton_compute_type(dtype)}) + {default}"
-        )
+        if (dtype, reduction_type, value) not in self.cse.reduction_cache:
+            self.cse.reduction_cache[(dtype, reduction_type, value)] = result_var
+            accumulator = f"_{result_var}"
+            self.body.writeline(
+                f"{accumulator} = tl.zeros({self.dense_size_str()}, {triton_compute_type(dtype)}) + {default}"
+            )
 
-        updated = value
-        if reduction_type == "min":
-            masks.append(f"({accumulator} > {value})")
-        elif reduction_type == "max":
-            masks.append(f"({accumulator} < {value})")
-        elif reduction_type == "sum":
-            updated = f"{accumulator} + {value}"
-        else:
-            raise NotImplementedError(f"reduction_type {reduction_type}")
+            updated = value
+            if reduction_type == "min":
+                masks.append(f"({accumulator} > {value})")
+            elif reduction_type == "max":
+                masks.append(f"({accumulator} < {value})")
+            elif reduction_type == "sum":
+                updated = f"{accumulator} + {value}"
+            else:
+                raise NotImplementedError(f"reduction_type {reduction_type}")
 
-        cond = " & ".join(masks)
-        self.compute.writeline(
-            f"{accumulator} = tl.where({cond}, {updated}, {accumulator})"
-        )
+            cond = " & ".join(masks)
+            self.compute.writeline(
+                f"{accumulator} = tl.where({cond}, {updated}, {accumulator})"
+            )
 
-        self.suffix.writeline(
-            f"{result_var} = tl.reshape(tl.{reduction_type}({accumulator}, {dim}), [{', '.join(sizes)}])"
-        )
+            self.suffix.writeline(
+                f"{result_var} = tl.reshape(tl.{reduction_type}({accumulator}, {dim}), [{', '.join(sizes)}])"
+            )
+        else:
+            var_name = self.cse.reduction_cache[(dtype, reduction_type, value)]
+            self.suffix.writeline(f"{result_var} = {var_name}")
         self.inside_reduction = False
         index, mask = self.indexing(index, result_var)
         assert "rmask" not in index