wehub-resource-sync
290 lines
9.4 KiB
C++
290 lines
9.4 KiB
C++
// Copyright (c) 2025 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
|
|
//
|
|
// http://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/entail_loop_condition_pass.h"
|
|
#include "paddle/cinn/common/ir_util.h"
|
|
#include "paddle/cinn/ir/ir_mutator.h"
|
|
#include "paddle/cinn/ir/stmt_visitors.h"
|
|
#include "paddle/cinn/ir/utils/ir_copy.h"
|
|
|
|
namespace cinn {
|
|
namespace optim {
|
|
|
|
using ir::stmt::BlockRef;
|
|
using ir::stmt::Evaluate;
|
|
using ir::stmt::For;
|
|
using ir::stmt::IfThenElse;
|
|
using ir::stmt::Let;
|
|
using ir::stmt::Schedule;
|
|
using ir::stmt::StmtRef;
|
|
using ir::stmt::Store;
|
|
|
|
namespace {
|
|
|
|
struct LoopCondLowerboundMutator : public ir::IRMutator<> {
|
|
explicit LoopCondLowerboundMutator(const ir::Var& var, const ir::Expr& extent)
|
|
: var_(var), extent_(extent) {}
|
|
|
|
void operator()(ir::Expr* expr) { ir::IRMutator<>::Visit(expr, expr); }
|
|
|
|
private:
|
|
void Visit(const ir::_Var_* op, ir::Expr* expr) override {
|
|
if (!op->is_symbolic_constant) {
|
|
if (expr->as_var_ref() == var_) {
|
|
*expr = extent_;
|
|
} else {
|
|
*expr = common::make_const(op->type(), 0);
|
|
}
|
|
}
|
|
}
|
|
|
|
ir::Var var_;
|
|
ir::Expr extent_;
|
|
};
|
|
|
|
bool CanProveEntailment(const ir::Var& loop_var,
|
|
const ir::Expr& extent,
|
|
const ir::LT* if_cond) {
|
|
// 1. For static shape, it's easy to prove by setting `loop_var = extent` and
|
|
// check whether the if-condition is violated
|
|
if (extent.is_constant()) {
|
|
ir::Expr lhs = ir::ir_utils::IRCopy(if_cond->a());
|
|
ir::Expr rhs = if_cond->b();
|
|
LoopCondLowerboundMutator(loop_var, extent)(&lhs);
|
|
if (!lhs.is_index()) return false;
|
|
lhs = lhs.as_index().Normalize();
|
|
return lhs.is_constant() && rhs.is_constant() &&
|
|
lhs.as_int64() >= rhs.as_int64();
|
|
}
|
|
|
|
// 2. For dynamic shape, as CINN is currently unable to prove equality with
|
|
// symbols, we just check their forms rather than values. We require that
|
|
// the extent be like `(<expr_a> / <expr_b>) + 1`, and the rhs of if_cond
|
|
// be `<expr_a>`.
|
|
auto* add_node = extent.As<ir::Add>();
|
|
if (!add_node) return false;
|
|
if (add_node->b() != ir::Expr(1)) return false;
|
|
auto* div_node = add_node->a().As<ir::Div>();
|
|
if (!div_node) return false;
|
|
return div_node->a() == if_cond->b();
|
|
}
|
|
|
|
bool CanEntailLoopCondition(const For& for_stmt) {
|
|
// 1. The loop extent must be >32 or be dynamic, otherwise the loop will be
|
|
// completely unrolled and we have nothing to do.
|
|
ir::Expr extent = for_stmt->extent();
|
|
if (extent.is_constant() && extent.as_int64() <= 32) return false;
|
|
|
|
// 2. The loop body must contains exactly one IfThenElse.
|
|
if (for_stmt->body()->stmts().size() != 1) return false;
|
|
StmtRef for_body = for_stmt->body()->stmts().front();
|
|
if (!for_body.isa<IfThenElse>()) return false;
|
|
|
|
// 3. The IfThenElse must be a leaf node, i.e., it only contains Schedule
|
|
// nodes.
|
|
IfThenElse if_stmt = for_body.as<IfThenElse>();
|
|
if (!if_stmt->condition().As<ir::LT>()) return false;
|
|
if (!if_stmt->false_case()->stmts().empty()) return false;
|
|
for (auto& inner_stmt : if_stmt->true_case()->stmts()) {
|
|
if (!inner_stmt.isa<Schedule>()) return false;
|
|
}
|
|
|
|
// 4. Check whether the if-condition actually entails the loop extent.
|
|
return CanProveEntailment(
|
|
for_stmt->loop_var(), extent, if_stmt->condition().As<ir::LT>());
|
|
}
|
|
|
|
struct CommonFactorExtractor : public ir::IRMutator<> {
|
|
explicit CommonFactorExtractor(const ir::Var& var) : var_(var) {}
|
|
|
|
void operator()(ir::Expr* expr) { ir::IRMutator<>::Visit(expr, expr); }
|
|
|
|
int64_t GetResult() { return common_factor_ > 0 ? common_factor_ : 1; }
|
|
|
|
private:
|
|
void Visit(const ir::Mul* op, ir::Expr* expr) override {
|
|
if (op->a() != var_ && op->b() != var_) {
|
|
auto* node = expr->As<ir::Mul>();
|
|
ir::IRMutator<>::Visit(&node->a(), &node->a());
|
|
ir::IRMutator<>::Visit(&node->b(), &node->b());
|
|
return;
|
|
}
|
|
|
|
int64_t factor = 1;
|
|
if (op->a() == var_ && op->b().is_constant()) {
|
|
factor = op->b().as_int64();
|
|
} else if (op->b() == var_ && op->a().is_constant()) {
|
|
factor = op->a().as_int64();
|
|
}
|
|
|
|
if (common_factor_ == 0) {
|
|
common_factor_ = factor;
|
|
} else {
|
|
common_factor_ = std::gcd(common_factor_, factor);
|
|
}
|
|
}
|
|
|
|
void Visit(const ir::_Var_* op, ir::Expr* expr) override {
|
|
// Note: if we visit a standalone `var` that is not an operand of a Mul op,
|
|
// the common factor must be 1.
|
|
if (*expr == var_) {
|
|
common_factor_ = 1;
|
|
}
|
|
}
|
|
|
|
ir::Expr var_;
|
|
int64_t common_factor_{0};
|
|
};
|
|
|
|
int64_t ExtractCommonFactorOfLoopVar(const For& for_stmt) {
|
|
CommonFactorExtractor extractor(for_stmt->loop_var());
|
|
|
|
const auto VisitFn = [&](const StmtRef& stmt) {
|
|
if (stmt.isa<IfThenElse>()) {
|
|
IfThenElse if_stmt = stmt.as<IfThenElse>();
|
|
ir::Expr condition = if_stmt->condition();
|
|
extractor(&condition);
|
|
} else if (stmt.isa<Store>()) {
|
|
Store store_stmt = stmt.as<Store>();
|
|
for (ir::Expr index : store_stmt->indices()) {
|
|
extractor(&index);
|
|
}
|
|
ir::Expr value = store_stmt->value();
|
|
extractor(&value);
|
|
}
|
|
};
|
|
|
|
ir::stmt::Visit(for_stmt->body(), VisitFn, [](auto) {});
|
|
return extractor.GetResult();
|
|
}
|
|
|
|
struct StridedLoopVarReplacer : public ir::IRMutator<> {
|
|
explicit StridedLoopVarReplacer(const ir::Var& var,
|
|
const ir::Var& new_var,
|
|
int64_t common_factor)
|
|
: var_(var), new_var_(new_var), common_factor_(common_factor) {}
|
|
|
|
void operator()(ir::Expr* expr) { ir::IRMutator<>::Visit(expr, expr); }
|
|
|
|
private:
|
|
void Visit(const ir::Mul* op, ir::Expr* expr) override {
|
|
auto* node = expr->As<ir::Mul>();
|
|
ir::IRMutator<>::Visit(&node->a(), &node->a());
|
|
ir::IRMutator<>::Visit(&node->b(), &node->b());
|
|
|
|
if (op->a() == var_) {
|
|
Replace(expr, &node->b());
|
|
} else if (op->b() == var_) {
|
|
Replace(expr, &node->a());
|
|
}
|
|
}
|
|
|
|
void Replace(ir::Expr* expr, ir::Expr* scale) {
|
|
if (*scale == ir::Expr(common_factor_)) {
|
|
*expr = new_var_;
|
|
} else {
|
|
scale->As<ir::IntImm>()->value /= common_factor_;
|
|
}
|
|
}
|
|
|
|
ir::Expr var_;
|
|
ir::Expr new_var_;
|
|
int64_t common_factor_;
|
|
};
|
|
|
|
void ReplaceWithStridedLoopVar(For for_stmt,
|
|
const ir::Var& strided_loop_var,
|
|
int64_t common_factor) {
|
|
ir::Var loop_var = for_stmt->loop_var();
|
|
StridedLoopVarReplacer replacer(loop_var, strided_loop_var, common_factor);
|
|
|
|
const auto VisitFn = [&](StmtRef stmt) {
|
|
if (stmt.isa<IfThenElse>()) {
|
|
IfThenElse if_stmt = stmt.as<IfThenElse>();
|
|
ir::Expr condition = if_stmt->condition();
|
|
replacer(&condition);
|
|
} else if (stmt.isa<Store>()) {
|
|
Store store_stmt = stmt.as<Store>();
|
|
for (ir::Expr index : store_stmt->indices()) {
|
|
replacer(&index);
|
|
}
|
|
ir::Expr value = store_stmt->value();
|
|
replacer(&value);
|
|
}
|
|
};
|
|
|
|
ir::stmt::Mutate(for_stmt->body(), VisitFn, [](auto) {});
|
|
}
|
|
|
|
void EntailLoopCondition(For for_stmt) {
|
|
ir::Var loop_var = for_stmt->loop_var();
|
|
IfThenElse if_stmt = for_stmt->body()->stmts().front().as<IfThenElse>();
|
|
std::vector<StmtRef> new_body_stmts;
|
|
|
|
// Step 1. Try extract common factor of loop_var.
|
|
// Note: type of common_factor should be strictly the same as loop_var. Don't
|
|
// mix int64 value in an int32 expression.
|
|
int64_t common_factor = ExtractCommonFactorOfLoopVar(for_stmt);
|
|
ir::Expr common_factor_expr =
|
|
common::make_const(loop_var->type(), common_factor);
|
|
|
|
// Step 2. If loop_var has a non-unit common factor, replace all occurrences
|
|
// of loop_var with strided_loop_var = loop_var * common_factor.
|
|
if (common_factor > 1) {
|
|
ir::Var strided_loop_var = loop_var->Copy();
|
|
strided_loop_var->name += "_strided";
|
|
new_body_stmts.push_back(
|
|
Let(strided_loop_var, ir::Mul::Make(loop_var, common_factor_expr)));
|
|
|
|
ReplaceWithStridedLoopVar(for_stmt, strided_loop_var, common_factor);
|
|
loop_var = strided_loop_var;
|
|
}
|
|
|
|
// Step 3. Declare the entailment relation on the loop condition.
|
|
ir::Expr entail_expr =
|
|
ir::Call::Make(Void(),
|
|
"CINN_ENTAIL_LOOP_CONDITION",
|
|
{loop_var, if_stmt->condition(), common_factor_expr},
|
|
{},
|
|
ir::CallType::Intrinsic);
|
|
new_body_stmts.push_back(Evaluate(entail_expr));
|
|
|
|
new_body_stmts.push_back(if_stmt);
|
|
for_stmt->set_body(BlockRef(new_body_stmts));
|
|
}
|
|
|
|
} // namespace
|
|
|
|
LogicalResult EntailLoopConditionPass::Run(ir::stmt::BlockRef block) {
|
|
std::vector<StmtRef> new_stmts;
|
|
for (auto stmt : block->stmts()) {
|
|
if (stmt.isa<For>()) {
|
|
For for_stmt = stmt.as<For>();
|
|
if (CanEntailLoopCondition(for_stmt)) {
|
|
EntailLoopCondition(for_stmt);
|
|
}
|
|
}
|
|
new_stmts.push_back(stmt);
|
|
}
|
|
block->set_stmts(new_stmts);
|
|
return LogicalResult::success();
|
|
}
|
|
|
|
std::unique_ptr<BlockPass> CreateEntailLoopConditionPass() {
|
|
return std::make_unique<EntailLoopConditionPass>();
|
|
}
|
|
|
|
} // namespace optim
|
|
} // namespace cinn
|