[CINN] enable eliminate common global memory read (#65804)

Dmovic · web-flow · commit bde4498bafde · 2024-07-30T15:10:56.000+08:00
* [CINN] enable eliminate common global memory read

* update common factor of local

* update recursive simplify

* update local index simplify

* fix code style

* add detailed log

* fix symbolic simplify

* fix add expr simplify

* add constrant for expr can not simplify

* add transform from index to iter var

* polish code for common global var

* fix codestyle

* fix codestyle

* enable transform

* add test for pass

* add constraint for select

* polish code

* update test name, fix code

* add local buffer size limit
diff --git a/paddle/cinn/optim/eliminate_common_factor_of_local_index.cc b/paddle/cinn/optim/eliminate_common_factor_of_local_index.cc
@@ -172,8 +172,28 @@ int gcd(int a, int b) {
   return gcd(b, a % b);
 }
 
+ir::Expr ExtractSymbolicFromExpr(const ir::Expr& expr) {
+  ir::Expr simplied_expr = cinn::common::AutoSimplify(expr);
+  if (simplied_expr.is_constant()) {
+    return ir::Expr(0);
+  } else if (expr.As<ir::_Var_>()) {
+    auto var = expr.As<ir::_Var_>();
+    if (var->is_symbolic_constant) {
+      VLOG(6) << "Extract symbolic constant, name = " << var->name;
+      return ir::ir_utils::IRCopy(expr);
+    }
+    return ir::Expr(0);
+  } else {
+    VLOG(6) << "Not supported for calculating symbolic, expr = " << expr;
+    return ir::Expr(0);
+  }
+  PADDLE_THROW(phi::errors::Fatal(
+      "Dead code. Fail to extract symbolic from expression."));
+}
+
 class Gcd {};
 class Offset {};
+class Symbolic {};
 
 template <typename Op>
 struct CommonFactorTrait;
@@ -217,6 +237,33 @@ struct CommonFactorTrait<Offset> {
 
 const ir::Expr CommonFactorTrait<Offset>::unit = ir::Expr(0);
 
+template <>
+struct CommonFactorTrait<Symbolic> {
+  static const ir::Expr unit;
+
+  static ir::Expr Calculate(const ir::Expr& expr1, const ir::Expr& expr2) {
+    auto IsSymbolicNotEqual = [&](const ir::Expr& expr1,
+                                  const ir::Expr& expr2) -> bool {
+      return cinn::common::AutoSimplify(
+                 ir::Sub::Make(ExtractSymbolicFromExpr(expr1),
+                               ExtractSymbolicFromExpr(expr2))) != ir::Expr(0);
+    };
+    if (IsSymbolicNotEqual(expr1, expr2)) {
+      return ir::Expr(0);
+    }
+    return ExtractSymbolicFromExpr(expr1);
+  }
+
+  static ir::Expr Simplify(const ir::Expr& expr, const ir::Expr& factor) {
+    if (factor != unit) {
+      return cinn::common::AutoSimplify(ir::Sub::Make(expr, factor));
+    }
+    return expr;
+  }
+};
+
+const ir::Expr CommonFactorTrait<Symbolic>::unit = ir::Expr(0);
+
 template <typename DoEachT>
 void VisitEachRowExpr(const std::vector<std::vector<ir::Expr>>& indexes,
                       std::size_t var_idx,
@@ -359,10 +406,114 @@ void EliminateCommonFactorHelper(ir::Expr* expr) {
   eliminate_common_factor_visitor(expr);
 }
 
+class TransformLocalIndicesVisitor : public ir::IRMutator<> {
+ public:
+  void operator()(ir::Expr* expr) { ir::IRMutator<>::Visit(expr, expr); }
+
+ private:
+  template <typename OpType>
+  void ExtractIterHelper(
+      const ir::Expr& expr,
+      std::unordered_map<std::string, ir::Expr>* name_to_iter) {
+    const auto op = expr.As<OpType>();
+    ExtractIterFromIndice(op->a(), name_to_iter);
+    ExtractIterFromIndice(op->b(), name_to_iter);
+  }
+
+  void ExtractIterFromIndice(
+      const ir::Expr& expr,
+      std::unordered_map<std::string, ir::Expr>* name_to_iter) {
+    if (expr.As<ir::_Var_>()) {
+      const auto var = expr.As<ir::_Var_>();
+      if (name_to_iter->count(var->name) == 0) {
+        (*name_to_iter)[var->name] = expr;
+      }
+    } else if (expr.As<ir::Add>()) {
+      ExtractIterHelper<ir::Add>(expr, name_to_iter);
+    } else if (expr.As<ir::Sub>()) {
+      ExtractIterHelper<ir::Sub>(expr, name_to_iter);
+    } else if (expr.As<ir::Mul>()) {
+      ExtractIterHelper<ir::Mul>(expr, name_to_iter);
+    } else if (expr.As<ir::Div>()) {
+      ExtractIterHelper<ir::Div>(expr, name_to_iter);
+    } else if (expr.As<ir::Mod>()) {
+      ExtractIterHelper<ir::Mod>(expr, name_to_iter);
+    } else {
+      VLOG(4) << "Not support for extract iter: \n" << expr;
+      return;
+    }
+    return;
+  }
+
+  std::vector<ir::Expr> ConvertIndicesToIters(
+      const std::vector<ir::Expr>& indices) {
+    auto CopyIndiceItersToLocalBuffer =
+        [&](const std::unordered_map<std::string, ir::Expr>& name_to_iter,
+            const std::vector<ir::Expr>& indices) -> std::vector<ir::Expr> {
+      std::vector<ir::Expr> local_buffer_iters;
+      for (std::size_t i = 0; i < loop_vars_.size(); ++i) {
+        VLOG(6) << "loop var name: " << loop_vars_[i]->name;
+        if (name_to_iter.count(loop_vars_[i]->name) > 0) {
+          local_buffer_iters.push_back(name_to_iter.at(loop_vars_[i]->name));
+        }
+      }
+
+      while (local_buffer_iters.size() < indices.size()) {
+        local_buffer_iters.insert(local_buffer_iters.begin(), ir::Expr(0));
+      }
+      return local_buffer_iters;
+    };
+
+    std::unordered_map<std::string, ir::Expr> name_to_iter;
+    for (const auto& indice : indices) {
+      ExtractIterFromIndice(indice, &name_to_iter);
+      VLOG(6) << "extract iter: " << indice
+              << " iter_set size: " << name_to_iter.size();
+    }
+    return CopyIndiceItersToLocalBuffer(name_to_iter, indices);
+  }
+
+  void Visit(const ir::For* op, ir::Expr* expr) override {
+    auto* for_ir = expr->As<ir::For>();
+    loop_vars_.push_back(for_ir->loop_var);
+    IRMutator<>::Visit(op, expr);
+    loop_vars_.pop_back();
+  }
+
+  void Visit(const ir::Store* op, ir::Expr* expr) override {
+    auto store = expr->As<ir::Store>();
+    if (store->tensor.as_tensor_ref()->buffer->memory_type ==
+        ir::MemoryType::GPULocal) {
+      store->indices = ConvertIndicesToIters(store->indices);
+    }
+    ir::IRMutator<>::Visit(op, expr);
+  }
+
+  void Visit(const ir::Load* op, ir::Expr* expr) override {
+    auto load = expr->As<ir::Load>();
+    if (load->tensor.as_tensor_ref()->buffer->memory_type ==
+        ir::MemoryType::GPULocal) {
+      load->indices = ConvertIndicesToIters(load->indices);
+    }
+    ir::IRMutator<>::Visit(op, expr);
+  }
+
+  std::vector<ir::Var> loop_vars_;
+};
+
+void TransformLocalIndicesToIters(ir::Expr* expr) {
+  TransformLocalIndicesVisitor transform_local_indices_visitor;
+  transform_local_indices_visitor(expr);
+}
+
 void EliminateCommonFactorOfLocalIndex(ir::Expr* expr) {
   VLOG(4) << "Before EliminateCommonFactorOfLocalIndex, Expr = \n" << *expr;
   EliminateCommonFactorHelper<Gcd>(expr);
   EliminateCommonFactorHelper<Offset>(expr);
+  EliminateCommonFactorHelper<Symbolic>(expr);
+
+  TransformLocalIndicesToIters(expr);
+
   VLOG(4) << "After EliminateCommonFactorOfLocalIndex, Expr = \n" << *expr;
 }
 
diff --git a/paddle/cinn/optim/eliminate_common_global_memory_read.cc b/paddle/cinn/optim/eliminate_common_global_memory_read.cc
@@ -15,6 +15,7 @@
 #include "paddle/cinn/optim/eliminate_common_global_memory_read.h"
 
 #include "paddle/cinn/common/cas.h"
+#include "paddle/cinn/common/integer_set.h"
 #include "paddle/cinn/ir/ir_mutator.h"
 #include "paddle/cinn/ir/ir_printer.h"
 #include "paddle/cinn/ir/utils/ir_compare.h"
@@ -125,9 +126,60 @@ struct GlobalTensorInfoCollector : public ir::IRMutator<Expr*> {
         [&](const std::vector<IndicesAndExtent>& indice_and_extent) -> bool {
       if (indice_and_extent.size() <= 1) return false;
       if (IndiceContainsLoad(indice_and_extent[0])) return false;
+      if (contains_select_) return false;
       return AllIndiceAndExtentEqual(indice_and_extent);
     };
 
+    const auto GetIterVarNames =
+        [](const std::vector<ir::Expr>& indices) -> std::set<std::string> {
+      std::set<std::string> iter_var_names;
+      for (const ir::Expr& e : indices) {
+        ir::ir_utils::CollectIRNodes(e, [&](const ir::Expr* x) {
+          if (x->as_var() && !x->as_var()->is_symbolic_constant) {
+            iter_var_names.insert(x->as_var()->name);
+          }
+          return false;
+        });
+      }
+      return iter_var_names;
+    };
+
+    auto CalculateBufferSize =
+        [&](const std::vector<ir::Expr>& indices) -> ir::Expr {
+      ir::Expr buffer_size(1);
+      std::set<std::string> iter_var_names = GetIterVarNames(indices);
+      for (const auto& iter_var_name : iter_var_names) {
+        if (iter_var_name_to_extent_.find(iter_var_name) ==
+            iter_var_name_to_extent_.end()) {
+          continue;
+        }
+        VLOG(6) << "Iter var name: " << iter_var_name << " with extent: "
+                << iter_var_name_to_extent_.at(iter_var_name);
+        buffer_size = cinn::common::AutoSimplify(ir::Mul::Make(
+            buffer_size, iter_var_name_to_extent_.at(iter_var_name)));
+      }
+      return buffer_size;
+    };
+
+    auto LocalBufferSizeLimit =
+        [&](const std::unordered_set<std::string>& global_buffer_name) -> bool {
+      ir::Expr size(0);
+      for (const auto& name : global_buffer_name) {
+        const std::vector<IndicesAndExtent>& indices_and_extent =
+            buffer_to_indice_and_extent_.at(name);
+        const ir::Expr buffer_size =
+            CalculateBufferSize(indices_and_extent[0].indices);
+        VLOG(6) << "Global buffer name: " << name
+                << " with size: " << buffer_size;
+        size = cinn::common::AutoSimplify(ir::Add::Make(size, buffer_size));
+      }
+      VLOG(6) << "Total buffer size: " << size;
+      common::cas_intervals_t var_intervals;
+      common::SymbolicExprAnalyzer analyzer(var_intervals);
+      std::optional<bool> prove_gt = analyzer.ProveGT(size, ir::Expr(128));
+      return prove_gt.value_or(false);
+    };
+
     std::unordered_set<std::string> global_buffer_name;
     for (const auto& [buffer_name, indice_and_extent] :
          buffer_to_indice_and_extent_) {
@@ -140,6 +192,13 @@ struct GlobalTensorInfoCollector : public ir::IRMutator<Expr*> {
         global_buffer_name.insert(buffer_name);
       }
     }
+    // When local buffer size too large, it will cause
+    // out of memory error, use global buffer instead.
+    // Fuse for loop will relax this constraints.
+    if (LocalBufferSizeLimit(global_buffer_name)) {
+      VLOG(6) << "Local buffer size too large, use global instead.\n";
+      global_buffer_name.clear();
+    }
     return global_buffer_name;
   }
 
@@ -168,6 +227,9 @@ struct GlobalTensorInfoCollector : public ir::IRMutator<Expr*> {
     CHECK(node);
     for_var_extents_.push_back(
         {node->loop_var, ir::ir_utils::IRCopy(node->extent)});
+    if (!node->is_binded()) {
+      iter_var_name_to_extent_[node->loop_var->name] = node->extent;
+    }
     ir::IRMutator<>::Visit(op, expr);
     for_var_extents_.pop_back();
   }
@@ -197,13 +259,21 @@ struct GlobalTensorInfoCollector : public ir::IRMutator<Expr*> {
     if (store_buffer->memory_type == ir::MemoryType::Heap) {
       global_store_buffer_names_.insert(store_buffer->name);
     }
+    ir::IRMutator<>::Visit(op, expr);
+  }
+
+  void Visit(const ir::Select* op, ir::Expr* expr) override {
+    contains_select_ = true;
+    ir::IRMutator<>::Visit(op, expr);
   }
 
   std::vector<ForVarExtent> for_var_extents_;
   std::unordered_map<ir::Var, ir::Expr> var_to_sb_expr_;
+  std::unordered_map<std::string, ir::Expr> iter_var_name_to_extent_;
   std::unordered_map<std::string, std::vector<IndicesAndExtent>>
       buffer_to_indice_and_extent_;
   std::unordered_set<std::string> global_store_buffer_names_;
+  bool contains_select_ = false;
 };
 
 struct CommonGlobalMemoryEliminator : public ir::IRMutator<Expr*> {
diff --git a/test/ir/pir/cinn/test_cinn_multiple_reduce.py b/test/ir/pir/cinn/test_cinn_multiple_reduce.py
@@ -0,0 +1,62 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import unittest
+
+import numpy as np
+import utils
+
+import paddle
+from paddle.base import core
+
+
+class MultipleReduceSubGraph(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        x1 = x.sum(axis=-1)
+        x2 = (x * x).sum(axis=-1)
+        return x1 + x2
+
+
+class TestMultipleReduceSubGraph(unittest.TestCase):
+    def setUp(self):
+        paddle.seed(2024)
+        self.shape = [32, 32]
+        self.dtype = "float32"
+        self.prepare_data()
+
+    def prepare_data(self):
+        self.x = paddle.randn(self.shape, dtype=self.dtype)
+        self.x.stop_gradient = False
+
+    def eval(self, use_cinn, use_prim=False):
+        if use_prim:
+            core._set_prim_all_enabled(True)
+        net = MultipleReduceSubGraph()
+        net = utils.apply_to_static(net, use_cinn=use_cinn)
+        net.eval()
+        out = net(self.x)
+
+        core._set_prim_all_enabled(False)
+        return out
+
+    def test_cinn(self):
+        cinn_out = self.eval(use_cinn=True, use_prim=True)
+        dy_out = self.eval(use_cinn=False, use_prim=True)
+        np.testing.assert_allclose(cinn_out, dy_out, atol=1e-2)
+
+
+if __name__ == '__main__':
+    unittest.main()
