Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add CanProveDivisible for symbolic calculation #60572

Merged
merged 3 commits into from
Jan 6, 2024
Merged

Add CanProveDivisible for symbolic calculation #60572

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
93d18ec
add CanProveDivisible for symbolic calculation
Courtesy-Xs Jan 4, 2024
11a877c
delete extra cout for debug
Courtesy-Xs Jan 5, 2024
9709f44
fix according to some comments
Courtesy-Xs Jan 5, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
91 changes: 87 additions & 4 deletions paddle/cinn/common/integer_set.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,8 @@
// limitations under the License.

#include "paddle/cinn/common/integer_set.h"

#include "paddle/cinn/common/arithmatic.h"
#include "paddle/cinn/ir/ir_mutator.h"
#include "paddle/cinn/ir/op/ir_operators.h"
#include "paddle/cinn/ir/utils/ir_copy.h"
Expand Down Expand Up @@ -164,11 +166,85 @@ std::optional<bool> SymbolicExprAnalyzer::ProveLT(const ir::Expr& lhs,
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 {
CHECK(rhs.is_var()) << "Rhs in ProveDivisible must be a var temporarily!\n";
CHECK(lhs.defined());
CHECK(rhs.defined());
CHECK(cinn::common::IsPureMath(lhs));

ir::Expr lhs_copy = ir::ir_utils::IRCopy(lhs);
if (cinn::common::is_zero(lhs_copy)) return true;

std::vector<ir::Expr> ops{};
bool res = false;
ir::Expr zero(0);
bool is_positive = ProveGT(lhs, zero).value_or(false);
lhs_copy = (is_positive ? 1 : -1) * lhs;
bool is_ge = ProveGE(lhs_copy, rhs).value_or(false);

switch (lhs.node_type()) {
case cinn::ir::IrNodeTy::_Var_:
return ProveEQ(lhs, rhs);
case cinn::ir::IrNodeTy::IntImm:
return false;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

判断不了或暂时没有实现的情况返回nullopt

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done

case cinn::ir::IrNodeTy::Sum:
res = true;
ops = lhs.As<ir::Sum>()->operands();
CHECK(!ops.empty());
std::for_each(ops.begin(), ops.end(), [&](const ir::Expr& expr) {
res = res && this->ProveDivisible(expr, rhs).value_or(false);
});
res = res && is_ge;
return res;
case cinn::ir::IrNodeTy::Product:
res = false;
ops = lhs.As<ir::Product>()->operands();
CHECK(!ops.empty());
std::for_each(ops.begin(), ops.end(), [&](const ir::Expr& expr) {
res = res || this->ProveDivisible(expr, rhs).value_or(false);
});
res = res && is_ge;
return res;
case cinn::ir::IrNodeTy::FracOp:
return ProveDivisible(lhs.As<ir::FracOp>()->a(), rhs).value_or(false) &&
is_ge;
case cinn::ir::IrNodeTy::FloatImm:
return false;
case cinn::ir::IrNodeTy::Add:
return ProveDivisible(lhs.As<ir::Add>()->a(), rhs).value_or(false) &&
ProveDivisible(lhs.As<ir::Add>()->b(), rhs).value_or(false) &&
is_ge;
case cinn::ir::IrNodeTy::Sub:
return ProveDivisible(lhs.As<ir::Sub>()->a(), rhs).value_or(false) &&
ProveDivisible(lhs.As<ir::Sub>()->b(), rhs).value_or(false) &&
is_ge;
case cinn::ir::IrNodeTy::Div:
return ProveDivisible(lhs.As<ir::Div>()->a(), rhs).value_or(false) &&
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这种情况不一定成立吧
lhs: a/b rhs: c
a%c == 0，a/b%c不一定为0
比如a = 10, b = 2, c = 2
a%c = 0, a/b%c != 0

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done

is_ge;
case cinn::ir::IrNodeTy::Mul:
return (ProveDivisible(lhs.As<ir::Mul>()->a(), rhs).value_or(false) ||
ProveDivisible(lhs.As<ir::Mul>()->b(), rhs).value_or(false)) &&
is_ge;
case cinn::ir::IrNodeTy::Mod:
return false;
case cinn::ir::IrNodeTy::Minus:
return ProveDivisible(lhs.As<ir::Minus>()->v(), rhs).value_or(false);
default:
LOG(FATAL) << "Not supported yet!";
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

我一直对这些地方的完备性担忧。

break;
}
}

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_intervals_(var_intervals),
sign_(is_lower_bound),
var_visited_({}) {}

void operator()(ir::Expr* expr) { IRMutator::Visit(expr, expr); }

Expand All @@ -183,10 +259,16 @@ class BoundReplacer : public ir::IRMutator<> {
upper_bound =
interval.e_r.defined() ? interval.e_r : ir::Expr(interval.r);
}
if (sign_) {
*op = ir::ir_utils::IRCopy(lower_bound);
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(upper_bound);
*op = ir::ir_utils::IRCopy(var_visited_.at(var->name));
}
}

Expand Down Expand Up @@ -248,6 +330,7 @@ class BoundReplacer : public ir::IRMutator<> {

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_;
Expand Down
2 changes: 2 additions & 0 deletions paddle/cinn/common/integer_set.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,8 @@ class SymbolicExprAnalyzer {
std::optional<bool> ProveLE(const ir::Expr& lhs, const ir::Expr& rhs) const;
std::optional<bool> ProveGT(const ir::Expr& lhs, const ir::Expr& rhs) const;
std::optional<bool> ProveLT(const ir::Expr& lhs, const ir::Expr& rhs) const;
std::optional<bool> ProveDivisible(const ir::Expr& lhs,
const ir::Expr& rhs) const;

ir::Expr LowerBound(const ir::Expr& expr) const;
ir::Expr UpperBound(const ir::Expr& expr) const;
Expand Down
39 changes: 39 additions & 0 deletions paddle/cinn/common/integer_set_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,45 @@ TEST_F(TestSymbolicExprAnalyzer, compare) {
analyzer.Prove(e3 < e4).value());
}

TEST_F(TestSymbolicExprAnalyzer, Divisible) {
auto x = ir::Var(ir::Expr(1), ir::Expr(7), "x");
auto y = ir::Var(ir::Expr(1), ir::Expr(15), "y");
auto S = ir::Var(ir::Expr(16), ir::Expr(256), "S");

cas_intervals_t divisible_var_intervals = {
{"x", CasInterval(x->lower_bound, x->upper_bound)},
{"y", CasInterval(y->lower_bound, y->upper_bound)},
{"S", CasInterval(S->lower_bound, S->upper_bound)},
};
SymbolicExprAnalyzer divisible_analyzer{divisible_var_intervals};

// case 1
ir::Expr e1 = 4 * x + 2 * y * x;
ir::Expr e2 = x;
ir::Expr e3 = y;

EXPECT_TRUE(divisible_analyzer.ProveDivisible(e1, e2).value_or(false));
EXPECT_FALSE(divisible_analyzer.ProveDivisible(e1, e3).value_or(false));

// case 2
ir::Expr e4 = y + y * x + 4 * y - x * y;

EXPECT_TRUE(divisible_analyzer.ProveDivisible(e4, e3).value_or(false));
EXPECT_FALSE(divisible_analyzer.ProveDivisible(e4, e2).value_or(false));

// case 3
ir::Expr e5 = x / y + x + y;

EXPECT_FALSE(divisible_analyzer.ProveDivisible(e5, e3).value_or(false));
EXPECT_FALSE(divisible_analyzer.ProveDivisible(e5, e2).value_or(false));

// case 4
ir::Expr e6 = S * x / 4 + x * y;

EXPECT_TRUE(divisible_analyzer.ProveDivisible(e6, e2).value_or(false));
EXPECT_FALSE(divisible_analyzer.ProveDivisible(e6, e3).value_or(false));
}

TEST(SingleIntervalIntSet, constant) {
SingleIntervalIntSet empty_set(ir::Expr(0), ir::Expr(-1));
SingleIntervalIntSet all_set(SymbolicExprLimit::negative_inf,
Expand Down