wehub-resource-sync
812 lines
28 KiB
C++
812 lines
28 KiB
C++
// Copyright (c) 2023 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
|
|
//
|
|
// 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/common/integer_set.h"
|
|
|
|
#include "paddle/cinn/common/ir_util.h"
|
|
#include "paddle/cinn/ir/ir_mutator.h"
|
|
#include "paddle/cinn/ir/op/ir_operators.h"
|
|
#include "paddle/cinn/ir/utils/ir_copy.h"
|
|
#include "paddle/cinn/optim/replace_var_with_expr.h"
|
|
#include "paddle/cinn/optim/simplify_util.h"
|
|
#include "paddle/common/enforce.h"
|
|
|
|
namespace cinn {
|
|
namespace common {
|
|
CasInterval::CasInterval(ir::Expr expr_l, ir::Expr expr_r) {
|
|
VLOG(6) << "CasInterval is : [" << expr_l << ", " << expr_r << "].";
|
|
expr_r = ReplaceMinToConstant(expr_r);
|
|
expr_l = ReplaceMaxToConstant(expr_l);
|
|
expr_l = optim::ArithSimplify(expr_l);
|
|
expr_r = optim::ArithSimplify(expr_r);
|
|
VLOG(6) << "After simplify, CasInterval is : [" << expr_l << ", " << expr_r
|
|
<< "].";
|
|
|
|
if (expr_l.is_constant() && expr_r.is_constant()) {
|
|
PADDLE_ENFORCE_EQ(
|
|
expr_l->type().is_integer(),
|
|
true,
|
|
::common::errors::InvalidArgument("Expected expr_l to be an integer."));
|
|
PADDLE_ENFORCE_EQ(
|
|
expr_r->type().is_integer(),
|
|
true,
|
|
::common::errors::InvalidArgument("Expected expr_r to be an integer."));
|
|
l = expr_l.as_int64();
|
|
r = expr_r.as_int64();
|
|
return;
|
|
}
|
|
e_l = expr_l;
|
|
e_r = expr_r;
|
|
}
|
|
|
|
/**
|
|
* @brief Given an expr, visit it. If there is an ir::Min and its operands are 1
|
|
* constant value and 1 inconstant value, return the constant min value. For
|
|
* example, if a < min(5, b), then we get a < 5 and a < b. Using a < 5 to
|
|
* simplify the condition ensures correctness, though not sufficient.
|
|
*/
|
|
ir::Expr CasInterval::ReplaceMinToConstant(ir::Expr expr) {
|
|
ir::Expr copied = ir::ir_utils::IRCopy(expr);
|
|
struct Mutator : public ir::IRMutator<ir::Expr*> {
|
|
void operator()(ir::Expr* expr) { Visit(expr); }
|
|
void Visit(ir::Expr* expr) { ir::IRMutator<>::Visit(expr, expr); }
|
|
|
|
private:
|
|
void Visit(const ir::Min* op, ir::Expr* expr) override {
|
|
auto a = op->a();
|
|
auto b = op->b();
|
|
|
|
Visit(&a);
|
|
Visit(&b);
|
|
|
|
auto min_a = op->a();
|
|
auto min_b = op->b();
|
|
if (min_a.is_constant() && !min_b.is_constant()) {
|
|
PADDLE_ENFORCE_EQ(
|
|
min_a->type().is_integer(),
|
|
true,
|
|
::common::errors::InvalidArgument("Min a should be an integer."));
|
|
*expr = ir::ir_utils::IRCopy(min_a);
|
|
} else if (min_b.is_constant() && !min_a.is_constant()) {
|
|
PADDLE_ENFORCE_EQ(
|
|
min_b->type().is_integer(),
|
|
true,
|
|
::common::errors::InvalidArgument("Min b should be an integer."));
|
|
*expr = ir::ir_utils::IRCopy(min_b);
|
|
}
|
|
}
|
|
};
|
|
Mutator()(&copied);
|
|
return copied;
|
|
}
|
|
|
|
/**
|
|
* @brief Given an expr, visit it. If there is an ir::Max and its operands are 1
|
|
* constant value and 1 inconstant value, return the constant max value.
|
|
*/
|
|
ir::Expr CasInterval::ReplaceMaxToConstant(ir::Expr expr) {
|
|
ir::Expr copied = ir::ir_utils::IRCopy(expr);
|
|
struct Mutator : public ir::IRMutator<ir::Expr*> {
|
|
void operator()(ir::Expr* expr) { Visit(expr); }
|
|
void Visit(ir::Expr* expr) { ir::IRMutator<>::Visit(expr, expr); }
|
|
|
|
private:
|
|
void Visit(const ir::Max* op, ir::Expr* expr) override {
|
|
auto a = op->a();
|
|
auto b = op->b();
|
|
|
|
Visit(&a);
|
|
Visit(&b);
|
|
|
|
auto max_a = op->a();
|
|
auto max_b = op->b();
|
|
if (max_a.is_constant() && !max_b.is_constant()) {
|
|
PADDLE_ENFORCE_EQ(
|
|
max_a->type().is_integer(),
|
|
true,
|
|
::common::errors::InvalidArgument("Max a should be an integer."));
|
|
*expr = ir::ir_utils::IRCopy(max_a);
|
|
} else if (max_b.is_constant() && !max_a.is_constant()) {
|
|
PADDLE_ENFORCE_EQ(
|
|
max_b->type().is_integer(),
|
|
true,
|
|
::common::errors::InvalidArgument("Max b should be an integer."));
|
|
*expr = ir::ir_utils::IRCopy(max_b);
|
|
}
|
|
}
|
|
};
|
|
Mutator()(&copied);
|
|
return copied;
|
|
}
|
|
|
|
std::ostream& operator<<(std::ostream& os, const CasInterval& i) {
|
|
if (i.e_l.defined() && i.e_r.defined()) {
|
|
os << "Expr e_l Interval[" << i.e_l << ", " << i.e_r << "]";
|
|
} else {
|
|
os << "Int l Interval[" << i.l << ", " << i.r << "]";
|
|
}
|
|
return os;
|
|
}
|
|
|
|
ir::Expr SymbolicExprLimit::positive_inf =
|
|
ir::Expr(ir::Var("positive_infinity"));
|
|
ir::Expr SymbolicExprLimit::negative_inf =
|
|
ir::Expr(ir::Var("negative_infinity"));
|
|
|
|
cas_intervals_t CollectVarIntervalsOfExprs(const std::vector<ir::Expr>& exprs,
|
|
bool is_lower_bound_zero) {
|
|
cas_intervals_t var_intervals;
|
|
for (ir::Expr expr : exprs) {
|
|
ir::ir_utils::CollectIRNodes(expr, [&](const ir::Expr* x) {
|
|
if (const ir::_Var_* var = x->as_var()) {
|
|
ir::Expr lower_bound = is_lower_bound_zero
|
|
? ir::Expr(static_cast<int64_t>(1))
|
|
: SymbolicExprLimit::negative_inf;
|
|
ir::Expr upper_bound = SymbolicExprLimit::positive_inf;
|
|
if (var->lower_bound.defined()) {
|
|
lower_bound = var->lower_bound;
|
|
}
|
|
if (var->upper_bound.defined()) {
|
|
upper_bound = var->upper_bound;
|
|
}
|
|
if (var->is_symbolic_constant) {
|
|
lower_bound = ir::Expr(1);
|
|
}
|
|
var_intervals.insert(
|
|
{var->name,
|
|
CasInterval(lower_bound, NormalizeUpperBound(upper_bound))});
|
|
}
|
|
return false;
|
|
});
|
|
}
|
|
return var_intervals;
|
|
}
|
|
|
|
std::optional<bool> SymbolicExprAnalyzer::Prove(
|
|
const ir::Expr& condition) const {
|
|
try {
|
|
if (condition.As<ir::EQ>()) {
|
|
return ProveEQ(condition.As<ir::EQ>()->a(), condition.As<ir::EQ>()->b());
|
|
}
|
|
if (condition.As<ir::NE>()) {
|
|
return ProveNE(condition.As<ir::NE>()->a(), condition.As<ir::NE>()->b());
|
|
}
|
|
if (condition.As<ir::GE>()) {
|
|
return ProveGE(condition.As<ir::GE>()->a(), condition.As<ir::GE>()->b());
|
|
}
|
|
if (condition.As<ir::LE>()) {
|
|
return ProveLE(condition.As<ir::LE>()->a(), condition.As<ir::LE>()->b());
|
|
}
|
|
if (condition.As<ir::GT>()) {
|
|
return ProveGT(condition.As<ir::GT>()->a(), condition.As<ir::GT>()->b());
|
|
}
|
|
if (condition.As<ir::LT>()) {
|
|
return ProveLT(condition.As<ir::LT>()->a(), condition.As<ir::LT>()->b());
|
|
}
|
|
return std::nullopt;
|
|
} catch (const ::common::enforce::EnforceNotMet& e) {
|
|
LOG(WARNING) << "Error occurred during integer calculation: " << e.what()
|
|
<< ", so SymbolicExprAnalyzer cannot prove anything.";
|
|
return std::nullopt;
|
|
}
|
|
}
|
|
|
|
std::optional<bool> SymbolicExprAnalyzer::ProveEQ(const ir::Expr& lhs,
|
|
const ir::Expr& rhs) const {
|
|
try {
|
|
if (lhs == rhs) {
|
|
return true;
|
|
}
|
|
ir::Expr diff = optim::ArithSimplify(ir::Sub::Make(lhs, rhs));
|
|
if (diff.is_constant()) {
|
|
return diff.get_constant() == 0;
|
|
}
|
|
ir::Expr diff_lower_bound = LowerBound(diff);
|
|
VLOG(6) << "lower bound of " << diff << " = " << diff_lower_bound;
|
|
ir::Expr diff_upper_bound = UpperBound(diff);
|
|
VLOG(6) << "upper bound of " << diff << " = " << diff_upper_bound;
|
|
if (diff_lower_bound.is_constant() && diff_upper_bound.is_constant() &&
|
|
diff_lower_bound.get_constant() == diff_upper_bound.get_constant()) {
|
|
return diff_lower_bound.get_constant() == 0;
|
|
}
|
|
std::optional<bool> prove_gt = ProveGT(lhs, rhs);
|
|
if (prove_gt.has_value() && prove_gt.value()) {
|
|
return false;
|
|
}
|
|
std::optional<bool> prove_lt = ProveLT(lhs, rhs);
|
|
if (prove_lt.has_value() && prove_lt.value()) {
|
|
return false;
|
|
}
|
|
return std::nullopt;
|
|
} catch (const ::common::enforce::EnforceNotMet& e) {
|
|
LOG(WARNING) << "Error occurred during integer calculation: " << e.what()
|
|
<< ", so SymbolicExprAnalyzer cannot prove anything.";
|
|
return std::nullopt;
|
|
}
|
|
}
|
|
|
|
std::optional<bool> SymbolicExprAnalyzer::ProveNE(const ir::Expr& lhs,
|
|
const ir::Expr& rhs) const {
|
|
try {
|
|
std::optional<bool> prove_eq = ProveEQ(lhs, rhs);
|
|
if (!prove_eq.has_value()) {
|
|
return std::nullopt;
|
|
}
|
|
return !prove_eq.value();
|
|
} catch (const ::common::enforce::EnforceNotMet& e) {
|
|
LOG(WARNING) << "Error occurred during integer calculation: " << e.what()
|
|
<< ", so SymbolicExprAnalyzer cannot prove anything.";
|
|
return std::nullopt;
|
|
}
|
|
}
|
|
|
|
std::optional<bool> SymbolicExprAnalyzer::ProveGE(const ir::Expr& lhs,
|
|
const ir::Expr& rhs) const {
|
|
try {
|
|
if (lhs == rhs) {
|
|
return true;
|
|
}
|
|
if (rhs == SymbolicExprLimit::positive_inf ||
|
|
lhs == SymbolicExprLimit::negative_inf) {
|
|
return false;
|
|
}
|
|
if (lhs == SymbolicExprLimit::positive_inf ||
|
|
rhs == SymbolicExprLimit::negative_inf) {
|
|
return true;
|
|
}
|
|
ir::Expr diff = optim::ArithSimplify(ir::Sub::Make(lhs, rhs));
|
|
VLOG(6) << "diff of " << ir::Sub::Make(lhs, rhs) << " = " << diff;
|
|
if (diff.is_constant() && diff.get_constant() < 0) {
|
|
return false;
|
|
}
|
|
if (diff.is_constant() && diff.get_constant() >= 0) {
|
|
return true;
|
|
}
|
|
ir::Expr diff_upper_bound = UpperBound(diff);
|
|
VLOG(6) << "upper bound of " << diff << " = " << diff_upper_bound;
|
|
if (diff_upper_bound.is_constant() && diff_upper_bound.get_constant() < 0) {
|
|
return false;
|
|
}
|
|
ir::Expr diff_lower_bound = LowerBound(diff);
|
|
VLOG(6) << "lower bound of " << diff << " = " << diff_lower_bound;
|
|
if (diff_lower_bound.is_constant() &&
|
|
diff_lower_bound.get_constant() >= 0) {
|
|
return true;
|
|
}
|
|
return std::nullopt;
|
|
} catch (const ::common::enforce::EnforceNotMet& e) {
|
|
LOG(WARNING) << "Error occurred during integer calculation: " << e.what()
|
|
<< ", so SymbolicExprAnalyzer cannot prove anything.";
|
|
return std::nullopt;
|
|
}
|
|
}
|
|
|
|
std::optional<bool> SymbolicExprAnalyzer::ProveLE(const ir::Expr& lhs,
|
|
const ir::Expr& rhs) const {
|
|
return ProveGE(rhs, lhs);
|
|
}
|
|
|
|
std::optional<bool> SymbolicExprAnalyzer::ProveGT(const ir::Expr& lhs,
|
|
const ir::Expr& rhs) const {
|
|
try {
|
|
if (lhs == rhs) {
|
|
return false;
|
|
}
|
|
if (rhs == SymbolicExprLimit::positive_inf ||
|
|
lhs == SymbolicExprLimit::negative_inf) {
|
|
return false;
|
|
}
|
|
if (lhs == SymbolicExprLimit::positive_inf ||
|
|
rhs == SymbolicExprLimit::negative_inf) {
|
|
return true;
|
|
}
|
|
ir::Expr diff = optim::ArithSimplify(ir::Sub::Make(lhs, rhs));
|
|
VLOG(6) << "diff of " << ir::Sub::Make(lhs, rhs) << " = " << diff;
|
|
if (diff.is_constant() && diff.get_constant() <= 0) {
|
|
return false;
|
|
}
|
|
if (diff.is_constant() && diff.get_constant() > 0) {
|
|
return true;
|
|
}
|
|
ir::Expr diff_upper_bound = UpperBound(diff);
|
|
VLOG(6) << "upper bound of " << diff << " = " << diff_upper_bound;
|
|
if (diff_upper_bound.is_constant() &&
|
|
diff_upper_bound.get_constant() <= 0) {
|
|
return false;
|
|
}
|
|
ir::Expr diff_lower_bound = LowerBound(diff);
|
|
VLOG(6) << "lower bound of " << diff << " = " << diff_lower_bound;
|
|
if (diff_lower_bound.is_constant() && diff_lower_bound.get_constant() > 0) {
|
|
return true;
|
|
}
|
|
|
|
return std::nullopt;
|
|
} catch (const ::common::enforce::EnforceNotMet& e) {
|
|
LOG(WARNING) << "Error occurred during integer calculation: " << e.what()
|
|
<< ", so SymbolicExprAnalyzer cannot prove anything.";
|
|
return std::nullopt;
|
|
}
|
|
}
|
|
|
|
std::optional<bool> SymbolicExprAnalyzer::ProveLT(const ir::Expr& lhs,
|
|
const ir::Expr& rhs) const {
|
|
return ProveGT(rhs, lhs);
|
|
}
|
|
|
|
// Tell whether lhs can be divisible by rhs, lhs must be a pure math expression
|
|
// and rhs must be a var
|
|
std::optional<bool> SymbolicExprAnalyzer::ProveDivisible(
|
|
const ir::Expr& lhs, const ir::Expr& rhs) const {
|
|
PADDLE_ENFORCE_EQ(rhs.is_var(),
|
|
true,
|
|
::common::errors::InvalidArgument(
|
|
"Rhs in ProveDivisible must be a var temporarily!\n"));
|
|
PADDLE_ENFORCE_EQ(lhs.defined(),
|
|
true,
|
|
::common::errors::InvalidArgument(
|
|
"Lhs in ProveDivisible must be defined."));
|
|
PADDLE_ENFORCE_EQ(rhs.defined(),
|
|
true,
|
|
::common::errors::InvalidArgument(
|
|
"Rhs in ProveDivisible must be defined."));
|
|
PADDLE_ENFORCE_EQ(
|
|
cinn::optim::IsPureMath(lhs),
|
|
true,
|
|
::common::errors::InvalidArgument(
|
|
"Lhs in ProveDivisible must be a pure math expression."));
|
|
|
|
try {
|
|
ir::Expr lhs_copy = ir::ir_utils::IRCopy(lhs);
|
|
if (cinn::common::is_zero(lhs_copy)) return true;
|
|
|
|
auto OptionalAnd =
|
|
[](const std::optional<bool>& lhs,
|
|
const std::optional<bool>& rhs) -> std::optional<bool> {
|
|
if (lhs.has_value() && rhs.has_value()) {
|
|
return lhs.value() && rhs.value();
|
|
} else {
|
|
return std::nullopt;
|
|
}
|
|
};
|
|
auto OptionalOr =
|
|
[](const std::optional<bool>& lhs,
|
|
const std::optional<bool>& rhs) -> std::optional<bool> {
|
|
if (lhs.has_value() && rhs.has_value()) {
|
|
return lhs.value() || rhs.value();
|
|
} else if ((!lhs.has_value()) && (!rhs.has_value())) {
|
|
return std::nullopt;
|
|
} else if (lhs.has_value() && (!rhs.has_value())) {
|
|
return lhs.value() ? std::optional<bool>(lhs.value())
|
|
: std::optional<bool>(std::nullopt);
|
|
} else {
|
|
return rhs.value() ? std::optional<bool>(rhs.value())
|
|
: std::optional<bool>(std::nullopt);
|
|
}
|
|
};
|
|
|
|
std::vector<ir::Expr> ops{};
|
|
std::optional<bool> res = std::nullopt;
|
|
ir::Expr zero(0);
|
|
ir::Expr tmp_expr;
|
|
|
|
auto is_ge = ProveGE(lhs, rhs);
|
|
|
|
switch (lhs.node_type()) {
|
|
case cinn::ir::IrNodeTy::_Var_:
|
|
return ProveEQ(lhs, rhs);
|
|
case cinn::ir::IrNodeTy::IntImm:
|
|
return false;
|
|
case cinn::ir::IrNodeTy::Sum:
|
|
res = true;
|
|
ops = lhs.As<ir::Sum>()->operands();
|
|
PADDLE_ENFORCE_NE(ops.empty(),
|
|
true,
|
|
::common::errors::InvalidArgument(
|
|
"Operands in Sum node should not be empty."));
|
|
std::for_each(ops.begin(), ops.end(), [&](const ir::Expr& expr) {
|
|
res = OptionalAnd(res, this->ProveDivisible(expr, rhs));
|
|
});
|
|
res = OptionalAnd(res, is_ge);
|
|
return res;
|
|
case cinn::ir::IrNodeTy::Product:
|
|
res = false;
|
|
ops = lhs.As<ir::Product>()->operands();
|
|
PADDLE_ENFORCE_NE(ops.empty(),
|
|
true,
|
|
::common::errors::InvalidArgument(
|
|
"Operands in Sum node should not be empty."));
|
|
std::for_each(ops.begin(), ops.end(), [&](const ir::Expr& expr) {
|
|
res = OptionalOr(res, this->ProveDivisible(expr, rhs));
|
|
if (res.has_value() && res.value()) return;
|
|
});
|
|
res = OptionalAnd(res, is_ge);
|
|
return res;
|
|
case cinn::ir::IrNodeTy::FloatImm:
|
|
return false;
|
|
case cinn::ir::IrNodeTy::Add:
|
|
return OptionalAnd(
|
|
OptionalAnd(ProveDivisible(lhs.As<ir::Add>()->a(), rhs),
|
|
ProveDivisible(lhs.As<ir::Add>()->b(), rhs)),
|
|
is_ge);
|
|
case cinn::ir::IrNodeTy::Sub:
|
|
return OptionalAnd(
|
|
OptionalAnd(ProveDivisible(lhs.As<ir::Sub>()->a(), rhs),
|
|
ProveDivisible(lhs.As<ir::Sub>()->b(), rhs)),
|
|
is_ge);
|
|
case cinn::ir::IrNodeTy::Div:
|
|
tmp_expr = optim::ArithSimplify(lhs);
|
|
if (tmp_expr.node_type() == cinn::ir::IrNodeTy::Div)
|
|
return std::nullopt;
|
|
return OptionalAnd(ProveDivisible(tmp_expr, rhs), is_ge);
|
|
case cinn::ir::IrNodeTy::Mul:
|
|
return OptionalAnd(
|
|
OptionalOr(ProveDivisible(lhs.As<ir::Mul>()->a(), rhs),
|
|
ProveDivisible(lhs.As<ir::Mul>()->b(), rhs)),
|
|
is_ge);
|
|
case cinn::ir::IrNodeTy::Mod:
|
|
return false;
|
|
case cinn::ir::IrNodeTy::Minus:
|
|
return ProveDivisible(lhs.As<ir::Minus>()->v(), rhs);
|
|
default:
|
|
PADDLE_THROW(::common::errors::InvalidArgument("Not supported yet!"));
|
|
break;
|
|
}
|
|
} catch (const ::common::enforce::EnforceNotMet& e) {
|
|
LOG(WARNING) << "Error occurred during integer calculation: " << e.what()
|
|
<< ", so SymbolicExprAnalyzer cannot prove anything.";
|
|
return std::nullopt;
|
|
}
|
|
}
|
|
|
|
class BoundReplacer : public ir::IRMutator<> {
|
|
public:
|
|
explicit BoundReplacer(const cas_intervals_t& var_intervals,
|
|
bool is_lower_bound)
|
|
: var_intervals_(var_intervals),
|
|
sign_(is_lower_bound),
|
|
var_visited_({}) {}
|
|
|
|
void operator()(ir::Expr* expr) { IRMutator::Visit(expr, expr); }
|
|
|
|
private:
|
|
void Visit(const ir::_Var_* var, ir::Expr* op) override {
|
|
// if the variable is S0/S1..., do not replace it.
|
|
if (var->is_symbolic_constant) return;
|
|
|
|
ir::Expr lower_bound = SymbolicExprLimit::negative_inf;
|
|
ir::Expr upper_bound = SymbolicExprLimit::positive_inf;
|
|
if (var_intervals_.count(var->name) != 0) {
|
|
const CasInterval& interval = var_intervals_.at(var->name);
|
|
lower_bound =
|
|
interval.e_l.defined() ? interval.e_l : ir::Expr(interval.l);
|
|
upper_bound =
|
|
interval.e_r.defined() ? interval.e_r : ir::Expr(interval.r);
|
|
}
|
|
if (!var_visited_.count(var->name)) {
|
|
if (sign_) {
|
|
*op = ir::ir_utils::IRCopy(lower_bound);
|
|
var_visited_.insert({var->name, lower_bound});
|
|
} else {
|
|
*op = ir::ir_utils::IRCopy(upper_bound);
|
|
var_visited_.insert({var->name, upper_bound});
|
|
}
|
|
} else {
|
|
*op = ir::ir_utils::IRCopy(var_visited_.at(var->name));
|
|
}
|
|
}
|
|
|
|
void Visit(const ir::Add* expr, ir::Expr* op) override {
|
|
ir::Add* node = op->As<ir::Add>();
|
|
IRMutator::Visit(&node->a(), &node->a());
|
|
IRMutator::Visit(&node->b(), &node->b());
|
|
}
|
|
|
|
void Visit(const ir::Mul* expr, ir::Expr* op) override {
|
|
ir::Mul* node = op->As<ir::Mul>();
|
|
if (node->b().is_constant() && node->b().get_constant() < 0) {
|
|
sign_ = !sign_;
|
|
}
|
|
IRMutator::Visit(&node->a(), &node->a());
|
|
if (node->b().is_constant() && node->b().get_constant() < 0) {
|
|
sign_ = !sign_;
|
|
}
|
|
if (node->a().is_constant() && node->a().get_constant() < 0) {
|
|
sign_ = !sign_;
|
|
}
|
|
IRMutator::Visit(&node->b(), &node->b());
|
|
if (node->a().is_constant() && node->a().get_constant() < 0) {
|
|
sign_ = !sign_;
|
|
}
|
|
}
|
|
|
|
void Visit(const ir::Sub* expr, ir::Expr* op) override {
|
|
ir::Sub* node = op->As<ir::Sub>();
|
|
IRMutator::Visit(&node->a(), &node->a());
|
|
sign_ = !sign_;
|
|
IRMutator::Visit(&node->b(), &node->b());
|
|
sign_ = !sign_;
|
|
}
|
|
|
|
void Visit(const ir::Div* expr, ir::Expr* op) override {
|
|
ir::Div* node = op->As<ir::Div>();
|
|
if (node->b().is_constant() && node->b().get_constant() < 0) {
|
|
sign_ = !sign_;
|
|
}
|
|
IRMutator::Visit(&node->a(), &node->a());
|
|
if (node->b().is_constant() && node->b().get_constant() < 0) {
|
|
sign_ = !sign_;
|
|
}
|
|
sign_ = !sign_;
|
|
IRMutator::Visit(&node->b(), &node->b());
|
|
sign_ = !sign_;
|
|
}
|
|
|
|
void Visit(const ir::Mod* expr, ir::Expr* op) override {
|
|
ir::Mod* node = op->As<ir::Mod>();
|
|
if (sign_) {
|
|
*op = ir::Expr(0);
|
|
} else {
|
|
IRMutator::Visit(&node->b(), &node->b());
|
|
*op = node->b() - ir::Expr(1);
|
|
}
|
|
}
|
|
|
|
private:
|
|
const cas_intervals_t& var_intervals_;
|
|
std::unordered_map<std::string, ir::Expr> var_visited_;
|
|
// Determine replacing with upper or lower bound,
|
|
// True means lower bound and False means upper bound.
|
|
bool sign_;
|
|
};
|
|
|
|
ir::Expr SymbolicExprAnalyzer::LowerBound(const ir::Expr& expr) const {
|
|
BoundReplacer bound_replacer(var_intervals_, true);
|
|
ir::Expr bound = ir::ir_utils::IRCopy(expr);
|
|
if (bound.is_index()) {
|
|
bound = bound.as_index().Normalize(ir::IndexExpr::OptLevel::kLevel3);
|
|
}
|
|
bound_replacer(&bound);
|
|
return optim::ArithSimplify(bound);
|
|
}
|
|
|
|
ir::Expr SymbolicExprAnalyzer::UpperBound(const ir::Expr& expr) const {
|
|
BoundReplacer bound_replacer(var_intervals_, false);
|
|
ir::Expr bound = ir::ir_utils::IRCopy(expr);
|
|
if (bound.is_index()) {
|
|
bound = bound.as_index().Normalize(ir::IndexExpr::OptLevel::kLevel3);
|
|
}
|
|
bound_replacer(&bound);
|
|
|
|
return optim::ArithSimplify(bound);
|
|
}
|
|
|
|
std::optional<bool> ProveEQ(const SingleIntervalIntSet& lhs,
|
|
const SingleIntervalIntSet& rhs) {
|
|
cas_intervals_t merged_var_intervals = MergeVarIntervals(lhs, rhs);
|
|
SymbolicExprAnalyzer analyzer(merged_var_intervals);
|
|
std::optional<bool> prove_min_eq = analyzer.ProveEQ(lhs.min_, rhs.min_);
|
|
std::optional<bool> prove_max_eq = analyzer.ProveEQ(lhs.max_, rhs.max_);
|
|
if (!prove_min_eq.has_value() || !prove_max_eq.has_value()) {
|
|
return std::nullopt;
|
|
} else if (prove_min_eq.value() == true && prove_max_eq.value() == true) {
|
|
return true;
|
|
} else if (prove_min_eq.value() == false || prove_max_eq.value() == false) {
|
|
return false;
|
|
}
|
|
return std::nullopt;
|
|
}
|
|
|
|
std::optional<SingleIntervalIntSet> ProvedUnion(const SingleIntervalIntSet& a,
|
|
const SingleIntervalIntSet& b) {
|
|
bool is_a_empty = a.ProveEmpty().value_or(false);
|
|
bool is_a_all = a.ProveAll().value_or(false);
|
|
bool is_b_empty = b.ProveEmpty().value_or(false);
|
|
bool is_b_all = b.ProveAll().value_or(false);
|
|
if (is_a_empty || is_b_all) {
|
|
return b;
|
|
}
|
|
if (is_b_empty || is_a_all) {
|
|
return a;
|
|
}
|
|
|
|
// May be relaxed when (a.max < b.min - 1) or (b.max < a.min - 1)
|
|
cas_intervals_t merged_var_intervals = MergeVarIntervals(a, b);
|
|
SymbolicExprAnalyzer analyzer(merged_var_intervals);
|
|
ir::Expr min = SymbolicExprLimit::positive_inf;
|
|
ir::Expr max = SymbolicExprLimit::negative_inf;
|
|
|
|
std::optional<bool> prove_a_min_le_b_min = analyzer.ProveLE(a.Min(), b.Min());
|
|
if (!prove_a_min_le_b_min.has_value()) {
|
|
return std::nullopt;
|
|
} else if (prove_a_min_le_b_min.value() == true) {
|
|
min = a.Min();
|
|
} else if (prove_a_min_le_b_min.value() == false) {
|
|
min = b.Min();
|
|
}
|
|
|
|
std::optional<bool> prove_a_max_ge_b_max = analyzer.ProveGE(a.Max(), b.Max());
|
|
if (!prove_a_max_ge_b_max.has_value()) {
|
|
return std::nullopt;
|
|
} else if (prove_a_max_ge_b_max.value() == true) {
|
|
max = a.Max();
|
|
} else if (prove_a_max_ge_b_max.value() == false) {
|
|
max = b.Max();
|
|
}
|
|
|
|
return SingleIntervalIntSet(min, max, std::move(merged_var_intervals));
|
|
}
|
|
|
|
std::optional<SingleIntervalIntSet> ProvedIntersect(
|
|
const SingleIntervalIntSet& a, const SingleIntervalIntSet& b) {
|
|
bool is_a_empty = a.ProveEmpty().value_or(false);
|
|
bool is_a_all = a.ProveAll().value_or(false);
|
|
bool is_b_empty = b.ProveEmpty().value_or(false);
|
|
bool is_b_all = b.ProveAll().value_or(false);
|
|
if (is_a_all || is_b_empty) {
|
|
return b;
|
|
}
|
|
if (is_b_all || is_a_empty) {
|
|
return a;
|
|
}
|
|
cas_intervals_t merged_var_intervals = MergeVarIntervals(a, b);
|
|
SymbolicExprAnalyzer analyzer(merged_var_intervals);
|
|
ir::Expr min = SymbolicExprLimit::positive_inf;
|
|
ir::Expr max = SymbolicExprLimit::negative_inf;
|
|
|
|
std::optional<bool> prove_a_max_lt_b_min_sub1 =
|
|
analyzer.ProveLT(a.Max(), b.Min() - ir::Expr(1));
|
|
std::optional<bool> prove_b_max_lt_a_min_sub1 =
|
|
analyzer.ProveLT(b.Max(), a.Min() - ir::Expr(1));
|
|
if (prove_a_max_lt_b_min_sub1.has_value() &&
|
|
prove_a_max_lt_b_min_sub1.value() ||
|
|
prove_b_max_lt_a_min_sub1.has_value() &&
|
|
prove_b_max_lt_a_min_sub1.value()) {
|
|
return SingleIntervalIntSet(min, max, std::move(merged_var_intervals));
|
|
}
|
|
|
|
std::optional<bool> prove_a_min_ge_b_min = analyzer.ProveGE(a.Min(), b.Min());
|
|
if (!prove_a_min_ge_b_min.has_value()) {
|
|
return std::nullopt;
|
|
} else if (prove_a_min_ge_b_min.value() == true) {
|
|
min = a.Min();
|
|
} else if (prove_a_min_ge_b_min.value() == false) {
|
|
min = b.Min();
|
|
}
|
|
|
|
std::optional<bool> prove_a_max_le_b_max = analyzer.ProveLE(a.Max(), b.Max());
|
|
if (!prove_a_max_le_b_max.has_value()) {
|
|
return std::nullopt;
|
|
} else if (prove_a_max_le_b_max.value() == true) {
|
|
max = a.Max();
|
|
} else if (prove_a_max_le_b_max.value() == false) {
|
|
max = b.Max();
|
|
}
|
|
|
|
return SingleIntervalIntSet(min, max, std::move(merged_var_intervals));
|
|
}
|
|
|
|
cas_intervals_t MergeVarIntervals(const SingleIntervalIntSet& a,
|
|
const SingleIntervalIntSet& b) {
|
|
cas_intervals_t merged = a.var_intervals_;
|
|
merged.insert(b.var_intervals_.begin(), b.var_intervals_.end());
|
|
return merged;
|
|
}
|
|
|
|
SingleIntervalIntSet::SingleIntervalIntSet(const ir::Expr& min,
|
|
const ir::Expr& max,
|
|
cas_intervals_t var_intervals)
|
|
: min_(min), max_(max), var_intervals_(var_intervals) {
|
|
if (var_intervals_.empty()) {
|
|
auto insert_interval_func = [&](const ir::Expr* x) {
|
|
if (x->as_var()) {
|
|
ir::Expr lower_bound = x->as_var()->lower_bound.defined()
|
|
? x->as_var()->lower_bound
|
|
: SymbolicExprLimit::negative_inf;
|
|
ir::Expr upper_bound = x->as_var()->upper_bound.defined()
|
|
? x->as_var()->upper_bound
|
|
: SymbolicExprLimit::positive_inf;
|
|
var_intervals_.insert(
|
|
{x->as_var()->name,
|
|
CasInterval(lower_bound, NormalizeUpperBound(upper_bound))});
|
|
}
|
|
return false;
|
|
};
|
|
ir::ir_utils::CollectIRNodes(min_, insert_interval_func);
|
|
ir::ir_utils::CollectIRNodes(max_, insert_interval_func);
|
|
}
|
|
}
|
|
|
|
std::optional<bool> SingleIntervalIntSet::ProveEmpty() const {
|
|
if (min_ == SymbolicExprLimit::positive_inf ||
|
|
max_ == SymbolicExprLimit::negative_inf) {
|
|
return true;
|
|
}
|
|
SymbolicExprAnalyzer analyzer(var_intervals_);
|
|
return analyzer.ProveGT(min_, max_);
|
|
}
|
|
|
|
std::optional<bool> SingleIntervalIntSet::ProveAll() const {
|
|
return min_ == SymbolicExprLimit::negative_inf &&
|
|
max_ == SymbolicExprLimit::positive_inf;
|
|
}
|
|
|
|
std::optional<bool> SingleIntervalIntSet::ProvePoint() const {
|
|
SymbolicExprAnalyzer analyzer(var_intervals_);
|
|
return analyzer.ProveEQ(min_, max_);
|
|
}
|
|
|
|
std::optional<bool> SingleIntervalIntSet::ProveSubSet(
|
|
const SingleIntervalIntSet& other) const {
|
|
cas_intervals_t merged_var_intervals = MergeVarIntervals(*this, other);
|
|
SymbolicExprAnalyzer analyzer(merged_var_intervals);
|
|
std::optional<bool> prove_min_ge = analyzer.ProveGE(min_, other.Min());
|
|
std::optional<bool> prove_max_le = analyzer.ProveLE(max_, other.Max());
|
|
if (!prove_min_ge.has_value() || !prove_max_le.has_value()) {
|
|
return std::nullopt;
|
|
} else if (prove_min_ge.value() && prove_max_le.value()) {
|
|
return true;
|
|
} else {
|
|
return false;
|
|
}
|
|
return std::nullopt;
|
|
}
|
|
|
|
std::optional<bool> SingleIntervalIntSet::ProveSuperSet(
|
|
const SingleIntervalIntSet& other) const {
|
|
cas_intervals_t merged_var_intervals = MergeVarIntervals(*this, other);
|
|
SymbolicExprAnalyzer analyzer(merged_var_intervals);
|
|
std::optional<bool> prove_min_le = analyzer.ProveLE(min_, other.Min());
|
|
std::optional<bool> prove_max_ge = analyzer.ProveGE(max_, other.Max());
|
|
if (!prove_min_le.has_value() || !prove_max_ge.has_value()) {
|
|
return std::nullopt;
|
|
} else if (prove_min_le.value() && prove_max_ge.value()) {
|
|
return true;
|
|
} else {
|
|
return false;
|
|
}
|
|
return std::nullopt;
|
|
}
|
|
|
|
ir::Expr EnhancedSimplifyModExpr(
|
|
ir::Expr e,
|
|
const paddle::flat_hash_map<std::string, CasInterval>& var_intervals) {
|
|
struct Mutator : public ir::IRMutator<ir::Expr*> {
|
|
explicit Mutator(
|
|
const paddle::flat_hash_map<std::string, CasInterval>& var_intervals)
|
|
: var_intervals_(var_intervals), analyzer_(var_intervals_) {}
|
|
|
|
void operator()(ir::Expr* expr) { Visit(expr); }
|
|
void Visit(ir::Expr* expr) { ir::IRMutator<>::Visit(expr, expr); }
|
|
|
|
private:
|
|
void Visit(const ir::Mod* op, ir::Expr* expr) override {
|
|
std::optional<bool> prove_lt = analyzer_.ProveLT(op->a(), op->b());
|
|
if (prove_lt.has_value() && prove_lt.value()) {
|
|
*expr = op->a();
|
|
}
|
|
}
|
|
|
|
private:
|
|
const paddle::flat_hash_map<std::string, CasInterval>& var_intervals_;
|
|
SymbolicExprAnalyzer analyzer_;
|
|
};
|
|
|
|
Mutator mutator(var_intervals);
|
|
ir::Expr copied = ir::ir_utils::IRCopy(e);
|
|
mutator(&copied);
|
|
return copied;
|
|
}
|
|
} // namespace common
|
|
} // namespace cinn
|