[CINN] Eliminate loops that have a loop-invariant body (#69376)

lshpku · web-flow · commit 882a84a1c252 · 2024-11-18T20:30:27.000+08:00
diff --git a/paddle/cinn/optim/CMakeLists.txt b/paddle/cinn/optim/CMakeLists.txt
@@ -11,6 +11,7 @@ gather_srcs(
   unroll_loops.cc
   transform_polyfor_to_for.cc
   eliminate_broadcast_in_forloop.cc
+  eliminate_invariant_loop.cc
   fold_cinn_call_arguments.cc
   call_arg_list_to_pod_value.cc
   insert_debug_log_callee.cc
diff --git a/paddle/cinn/optim/eliminate_invariant_loop.cc b/paddle/cinn/optim/eliminate_invariant_loop.cc
@@ -0,0 +1,154 @@
+// Copyright (c) 2024 CINN 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.
+
+#include "paddle/cinn/optim/eliminate_invariant_loop.h"
+
+#include "paddle/cinn/common/ir_util.h"
+#include "paddle/cinn/ir/ir_analyzer/ir_analyzer.h"
+#include "paddle/cinn/ir/ir_mutator.h"
+#include "paddle/cinn/ir/utils/ir_replace.h"
+#include "paddle/cinn/optim/replace_var_with_expr.h"
+
+namespace cinn {
+namespace optim {
+
+namespace {
+
+// Check whether var is used in the block's store indices or value.
+bool HasVarInIndicesOrValue(const ir::Expr& block, const ir::Var& var) {
+  auto* block_realize = block.As<ir::ScheduleBlockRealize>();
+  auto* schedule_block = block_realize->schedule_block.As<ir::ScheduleBlock>();
+
+  // collect iter vars that refer to `var`
+  std::set<std::string> ref_iter_vars;
+  for (int i = 0; i < block_realize->iter_values.size(); ++i) {
+    auto var_use = ir::ir_utils::CollectIRNodes(
+        block_realize->iter_values[i],
+        [&](const ir::Expr* x) {
+          return x->is_var() && x->as_var_ref() == var;
+        },
+        /* uniq_target = */ true);
+    if (var_use.size() > 0) {
+      ref_iter_vars.insert(schedule_block->iter_vars[i]->name);
+    }
+  }
+
+  ir::Expr store = ir::analyzer::GetStoreOfSBlock(block);
+  auto var_use = ir::ir_utils::CollectIRNodes(
+      store,
+      [&](const ir::Expr* x) {
+        if (x->is_var()) {
+          if (x->as_var_ref() == var) return true;
+          if (ref_iter_vars.count(x->as_var()->name) > 0) return true;
+        }
+        return false;
+      },
+      /* uniq_target = */ true);
+  return var_use.size() > 0;
+}
+
+// Check whether the block has an inplace update (e.g. a[i] = a[i] + b[i])
+// by comparing between the block's read_buffers and write_buffers.
+bool HasInplaceUpdate(const ir::Expr& block) {
+  auto* schedule_block = block.As<ir::ScheduleBlockRealize>()
+                             ->schedule_block.As<ir::ScheduleBlock>();
+  std::set<std::string> read_buffer_names;
+  for (auto& buffer_range : schedule_block->read_buffers) {
+    read_buffer_names.insert(
+        buffer_range.As<ir::_BufferRange_>()->buffer.as_buffer()->name);
+  }
+  for (auto& buffer_range : schedule_block->write_buffers) {
+    auto& write_buffer_name =
+        buffer_range.As<ir::_BufferRange_>()->buffer.as_buffer()->name;
+    if (read_buffer_names.count(write_buffer_name) > 0) {
+      return true;
+    }
+  }
+  return false;
+}
+
+// Check whether the block is writing to a buffer whose scope is smaller than
+// the For node's scope.
+bool HasWriteToSmallerScope(const ir::Expr& block, const ir::For* for_node) {
+  ir::Expr store_tensor = ir::analyzer::GetStoreTensorOfSBlock(block);
+  ir::MemoryType memory_type = store_tensor.as_tensor()->buffer->memory_type;
+  if (for_node->is_gpu_thread_binded()) {
+    if (memory_type == ir::MemoryType::GPULocal) {
+      return true;
+    }
+  } else if (for_node->is_gpu_block_binded()) {
+    if (memory_type == ir::MemoryType::GPULocal ||
+        memory_type == ir::MemoryType::GPUShared) {
+      return true;
+    }
+  }
+  return false;
+}
+
+struct InvariantLoopEliminator : public ir::IRMutator<> {
+  using ir::IRMutator<>::Visit;
+  void operator()(ir::Expr* expr) { ir::IRMutator<>::Visit(expr, expr); }
+
+  void Visit(const ir::For* op, ir::Expr* expr) override {
+    auto* node = expr->As<ir::For>();
+    ir::IRMutator<>::Visit(&node->body, &node->body);
+
+    // do nothing if we are outside the root schedule block
+    if (!root_) return;
+
+    // check that all child schedule blocks satisfy the four rules
+    std::vector<ir::Expr> child_blocks = ir::analyzer::GetChildBlocks(*expr);
+    ir::Var loop_var = node->loop_var;
+    for (auto& block : child_blocks) {
+      if (HasVarInIndicesOrValue(block, loop_var)) return;
+      if (HasInplaceUpdate(block)) return;
+      if (node->is_binded()) {
+        if (HasWriteToSmallerScope(block, node)) return;
+        if (!ir::analyzer::GetConsumerSBlocks(block, *root_).empty()) return;
+      }
+    }
+
+    // now we can eliminate this For node
+    ir::Expr body = node->body;
+    ir::ir_utils::IrReplaceVarBroadcast(&body, ir::Expr(loop_var), ir::Expr(0));
+
+    if (node->is_binded()) {
+      ir::Expr cond = ir::EQ::Make(loop_var, ir::Expr(0));
+      ir::Expr wrapped_body = ir::IfThenElse::Make(cond, body);
+      node->body = ir::Block::Make({wrapped_body});
+    } else {
+      *expr = body;
+    }
+  }
+
+  void Visit(const ir::ScheduleBlockRealize* op, ir::Expr* expr) override {
+    auto* schedule_block = op->schedule_block.As<ir::ScheduleBlock>();
+    // find and set the root schedule block
+    if (schedule_block->name.substr(0, 4) == "root") {
+      root_ = expr;
+      ir::IRMutator<>::Visit(op, expr);
+      root_ = nullptr;
+    }
+  }
+
+ private:
+  const ir::Expr* root_{nullptr};
+};
+
+}  // namespace
+
+void EliminateInvariantLoop(ir::Expr* expr) { InvariantLoopEliminator()(expr); }
+
+}  // namespace optim
+}  // namespace cinn
diff --git a/paddle/cinn/optim/eliminate_invariant_loop.h b/paddle/cinn/optim/eliminate_invariant_loop.h
@@ -0,0 +1,110 @@
+// Copyright (c) 2024 CINN 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.
+
+#pragma once
+#include "paddle/cinn/ir/ir.h"
+#include "paddle/cinn/ir/lowered_func.h"
+
+namespace cinn {
+namespace optim {
+
+/**
+ * Eliminate serial or bound loops that have a loop-invariant body.
+ *
+ * We can eliminate a serial loop if it satisfies rule (1) and (2):
+ * (1) The loop variable is not used in any load/store indices or computation
+ *     within child schedule blocks. This ensures that the loop writes the same
+ *     value to the same index in each iteration.
+ * (2) It doesn't contain any inplace operations, e.g. a[0] = a[0] + b[0]. In
+ *     the presence of inplace operations, the loop count matters.
+ *
+ * We can eliminate a bound loop if it also satisfies rule (3) and (4):
+ * (3) It doesn't write to the local buffer (for thread-bound loop) or shared
+ *     memory (for block-bound loop), because the index for these storages
+ *     implicitly includes the thread/block index, thereby violating (1).
+ * (4) Its child schedule blocks don't have consumers, otherwise consumers in
+ *     other blocks/threads may read an undefined value.
+ *
+ * When a loop can be eliminated, we:
+ * (a) For serial loop, we just replace the For node with its body, and replace
+ *     all use of the loop variable to 0.
+ * (b) For bound loop, we cannot remove the For node because it contains the
+ *     binding info. Instead, we warp the loop body in `if (loop_var == 0)`.
+ *
+ *
+ * Example 1:
+ *   serial for (k, 0, 32):
+ *     ScheduleBlock(A):
+ *       A[0] = B[0]
+ * =>
+ *   ScheduleBlock(A):
+ *     A[0] = B[0]
+ *
+ *
+ * Example 2:
+ *   bind for (threadIdx.x, 0, 32):
+ *     ScheduleBlock(A):
+ *       A[0] = B[0]  # assume that A is global
+ * =>
+ *   bind for (threadIdx.x, 0, 32):
+ *     if threadIdx.x == 0:
+ *       ScheduleBlock(A):
+ *         A[0] = B[0]
+ * Note:
+ *   We don't remove a bound loop, but wrap its body in `threadIdx.x == 0`.
+ *
+ *
+ * Example 3:
+ *   serial for (k, 0, 32):
+ *     ScheduleBlock(A):
+ *       A[0] = A[0] + B[0]
+ * =>
+ *   DO NOTHING!
+ * Note:
+ *   Even though k is not used in the body, as ScheduleBlock(A) contains an
+ *   inplace operation, each iteration updates A, so we cannot eliminate it.
+ *
+ *
+ * Example 4:
+ *   bind for (blockIdx.x, 0, 32):
+ *     ScheduleBlock(A):
+ *       A[0] = 3.14
+ *   bind for (blockIdx.x, 0, 32):
+ *     ScheduleBlock(B):
+ *       B[blockIdx.x] = A[0]
+ * =>
+ *   DO NOTHING!
+ * Note:
+ *   If we eliminate the For node of A, only the block at `blockIdx.x == 0`
+ *   can get the correct value of A[0] when computing B.
+ *
+ *
+ * Example 5:
+ *   serial for (k, 0, 32):
+ *     if (k * 256) + threadIdx.x < 8000:
+ *       ScheduleBlock(A):
+ *         A[threadIdx.x] = B[threadIdx.x]
+ * =>
+ *   if threadIdx.x < 8000:
+ *     ScheduleBlock(A):
+ *       A[threadIdx.x] = B[threadIdx.x]
+ * Note:
+ *   Although k is used in the if condition, it doesn't affect the index to
+ *   update, but just limits how many times to update. So rewriting k to 0
+ *   doesn't change the semantics.
+ */
+void EliminateInvariantLoop(ir::Expr *expr);
+
+}  // namespace optim
+}  // namespace cinn
diff --git a/paddle/cinn/optim/ir_simplify.cc b/paddle/cinn/optim/ir_simplify.cc
@@ -235,28 +235,25 @@ struct SimplifyIfThenElseMutator : public ir::IRMutator<> {
 
     auto* condition_int = node->condition.As<ir::IntImm>();
     auto* condition_uint = node->condition.As<ir::UIntImm>();
-    int64_t value;
-    if (condition_int || condition_uint) {
-      if (condition_int) {
-        value = condition_int->value;
-      } else {
-        value = condition_uint->value;
-      }
-      if (value) {
-        *expr = op->true_case;
-      } else {
-        if (op->false_case.defined()) {
-          *expr = op->false_case;
-        } else {
-          // null condition
-          *expr = ir::Block::Make({});
-        }
+
+    // not deterministic
+    if (!condition_int && !condition_uint) {
+      Visit(&node->true_case, &node->true_case);
+      if (node->false_case.defined()) {
+        Visit(&node->false_case, &node->false_case);
       }
+      return;
     }
-    if (expr->As<ir::IfThenElse>()) {
-      if (node->true_case.defined()) Visit(&node->true_case, &node->true_case);
-      if (node->false_case.defined())
-        Visit(&node->false_case, &node->false_case);
+
+    bool value = condition_int ? condition_int->value : condition_uint->value;
+    if (value) {
+      *expr = op->true_case;
+      Visit(expr, expr);
+    } else if (op->false_case.defined()) {
+      *expr = op->false_case;
+      Visit(expr, expr);
+    } else {
+      *expr = ir::Block::Make({});
     }
   }
 };
diff --git a/paddle/cinn/optim/optimize.cc b/paddle/cinn/optim/optimize.cc
@@ -20,6 +20,7 @@
 #include "paddle/cinn/optim/call_arg_list_to_pod_value.h"
 #include "paddle/cinn/optim/cast_bool_to_int8.h"
 #include "paddle/cinn/optim/eliminate_broadcast_in_forloop.h"
+#include "paddle/cinn/optim/eliminate_invariant_loop.h"
 #include "paddle/cinn/optim/extern_call_process.h"
 #include "paddle/cinn/optim/fold_cinn_call_arguments.h"
 #include "paddle/cinn/optim/if_fusion.h"
@@ -57,6 +58,8 @@ ir::LoweredFunc Optimize(ir::LoweredFunc fn,
   ReplaceConstParamToInteger(&copied->body);
   // Simplify already contains CastSimplify
   Simplify(&copied->body);
+  EliminateInvariantLoop(&copied->body);
+  VLOG(4) << "After Optimize EliminateInvariantLoop:" << copied;
   ReplaceCrossThreadReduction(copied);
   VLOG(4) << "After Optimize ReplaceCrossThreadReduction:" << copied;
   ReplaceCrossBlockReduction(copied);
diff --git a/paddle/cinn/optim/replace_cross_block_reduction.cc b/paddle/cinn/optim/replace_cross_block_reduction.cc
@@ -17,7 +17,6 @@
 
 #include "paddle/cinn/adt/adt.h"
 #include "paddle/cinn/common/common.h"
-#include "paddle/cinn/hlir/framework/pir/trivial_op_impl.h"
 #include "paddle/cinn/hlir/pe/reduction.h"
 #include "paddle/cinn/ir/ir.h"
 #include "paddle/cinn/ir/ir_mutator.h"
@@ -26,11 +25,6 @@
 
 namespace cinn {
 
-using hlir::framework::pir::trivial_fusion_detail::ExprSetFinderUtils::
-    ChildIfThenElses;
-using hlir::framework::pir::trivial_fusion_detail::ExprSetFinderUtils::
-    ChildTensorLoads;
-
 namespace optim {
 namespace {
 
@@ -81,7 +75,8 @@ ir::Expr GetRightOperand(const ir::Expr& expr) {
 }
 
 void RemoveGridReduceAxisFromIfCondition(ir::Expr* expr) {
-  std::vector<ir::Expr> if_exprs = ChildIfThenElses(*expr);
+  std::vector<ir::Expr> if_exprs = ir::ir_utils::CollectIRNodesInOrder(
+      *expr, [](const ir::Expr* x) { return x->As<ir::IfThenElse>(); });
   for (auto& if_expr : if_exprs) {
     auto& condition = if_expr.As<ir::IfThenElse>()->condition;
     optim::ReplaceVarWithExpr(&condition, ir::Var("blockIdx.y"), ir::Expr(0));
@@ -140,7 +135,9 @@ struct BaseMutator : public ir::IRMutator<> {
 struct CrossBlockReductionReorderer : public BaseMutator {
  private:
   bool IsGridReduceDownstream(const ir::Expr& expr_block) {
-    for (auto& expr_load : ChildTensorLoads(expr_block)) {
+    std::vector<ir::Expr> child_loads = ir::ir_utils::CollectIRNodesInOrder(
+        expr_block, [](const ir::Expr* x) { return x->As<ir::Load>(); });
+    for (auto& expr_load : child_loads) {
       std::string load_tensor_name = expr_load.As<ir::Load>()->name();
       if (downstream_names_.count(load_tensor_name) > 0) {
         return true;
@@ -412,7 +409,10 @@ struct CrossBlockReductionReplacer : public BaseMutator {
         // For trivial ops that are after grid reduce, we need to remove the
         // `blockIdx.y` axis from their if conditions, otherwise they may not
         // be executed in the last done block.
-        RemoveGridReduceAxisFromIfCondition(op);
+        // TODO(liangshuhao): we won't need to remove `blockIdx.y` when we
+        // switch to use cooperative groups.
+        ir::Expr root_compute_body = cur_loops_[0]->body;
+        RemoveGridReduceAxisFromIfCondition(&root_compute_body);
         *op = WrapInLastBlockDone(op);
       }
       return;
