wehub-resource-sync
798 lines
28 KiB
C++
798 lines
28 KiB
C++
// Copyright (c) 2021 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/optim/ir_simplify.h"
|
|
|
|
#include <ginac/ginac.h>
|
|
#include <glog/logging.h>
|
|
|
|
#include <map>
|
|
#include <string>
|
|
|
|
#include "paddle/cinn/common/ir_util.h"
|
|
#include "paddle/cinn/ir/ir_mutator.h"
|
|
#include "paddle/cinn/ir/ir_printer.h"
|
|
#include "paddle/cinn/ir/ir_visitor.h"
|
|
#include "paddle/cinn/ir/op/ir_operators.h"
|
|
#include "paddle/cinn/ir/tensor.h"
|
|
#include "paddle/cinn/ir/utils/ir_copy.h"
|
|
#include "paddle/cinn/optim/simplify_util.h"
|
|
#include "paddle/cinn/utils/string.h"
|
|
#include "paddle/utils/flat_hash_map.h"
|
|
|
|
namespace cinn {
|
|
namespace optim {
|
|
using namespace ir; // NOLINT
|
|
using cinn::common::bfloat16;
|
|
using cinn::common::float16;
|
|
using utils::GetStreamCnt;
|
|
using utils::Replace;
|
|
|
|
namespace {
|
|
|
|
//! Simplify the expression but Load.
|
|
struct SimplifyNoPureMathMutator : public ir::IRMutator<ir::Expr*> {
|
|
SimplifyNoPureMathMutator(
|
|
ir::IndexExpr::OptLevel opt_level = ir::IndexExpr::OptLevel::kLevel1)
|
|
: opt_level_(opt_level) {}
|
|
void operator()(Expr* x) { ir::IRMutator<ir::Expr*>::Visit(x, x); }
|
|
|
|
using ir::IRMutator<>::Visit;
|
|
|
|
#define __(op__) \
|
|
void Visit(const op__* op, Expr* expr) override { \
|
|
*expr = ArithSimplify(*expr, opt_level_); \
|
|
}
|
|
|
|
__(Add)
|
|
__(Mul)
|
|
__(Sub)
|
|
__(Div)
|
|
__(Mod)
|
|
__(Min)
|
|
__(Max)
|
|
#undef __
|
|
ir::IndexExpr::OptLevel opt_level_;
|
|
};
|
|
|
|
struct ReplaceFracWithDivMutator : public ir::IRMutator<> {
|
|
void operator()(Expr* x) { ir::IRMutator<>::Visit(x, x); }
|
|
|
|
void Visit(const FracOp* op, Expr* expr) override {
|
|
auto* node = expr->As<ir::FracOp>();
|
|
|
|
ir::IRMutator<>::Visit(&node->operand(0), &node->operand(0));
|
|
ir::IRMutator<>::Visit(&node->operand(1), &node->operand(1));
|
|
|
|
*expr = ir::Div::Make(node->operand(0), node->operand(1));
|
|
}
|
|
};
|
|
|
|
template <typename CastType, typename T>
|
|
CastType NormCastValue(T value) {
|
|
if (type_of<CastType>().is_uint() || type_of<T>().is_uint()) {
|
|
// not support uint
|
|
return static_cast<CastType>(value);
|
|
}
|
|
|
|
if (std::isinf(value)) {
|
|
if (CastType(value) == -std::numeric_limits<CastType>::infinity()) {
|
|
return -std::numeric_limits<CastType>::infinity();
|
|
}
|
|
return std::numeric_limits<CastType>::infinity();
|
|
} else if (std::isnan(value)) {
|
|
return std::numeric_limits<CastType>::signaling_NaN();
|
|
} else if (value >= static_cast<T>(std::numeric_limits<CastType>::max())) {
|
|
return std::numeric_limits<CastType>::max();
|
|
} else if (value <= static_cast<T>(std::numeric_limits<CastType>::lowest())) {
|
|
return std::numeric_limits<CastType>::lowest();
|
|
}
|
|
return static_cast<CastType>(value);
|
|
}
|
|
|
|
struct SimplifyCastMutator : public ir::IRMutator<> {
|
|
void operator()(Expr* expr) { ir::IRMutator<ir::Expr*>::Visit(expr, expr); }
|
|
|
|
void Visit(const ir::Cast* op, Expr* expr) {
|
|
auto* node = expr->As<ir::Cast>();
|
|
|
|
ir::IRMutator<ir::Expr*>::Visit(&node->v(), &node->v());
|
|
|
|
if (op->type() == op->v().type()) {
|
|
*expr = op->v();
|
|
return;
|
|
}
|
|
|
|
#define __CAST_TO_TYPE(type__) \
|
|
if (auto* i = op->v().As<ir::IntImm>()) { \
|
|
*expr = Expr(static_cast<type__>(i->value)); \
|
|
} else if (auto* f = op->v().As<ir::FloatImm>()) { \
|
|
*expr = Expr(static_cast<type__>(NormCastValue<type__>(f->value))); \
|
|
} else if (auto* u = op->v().As<ir::UIntImm>()) { \
|
|
*expr = Expr(static_cast<type__>(u->value)); \
|
|
} else { \
|
|
CINN_NOT_IMPLEMENTED \
|
|
}
|
|
|
|
if (op->v().is_constant()) {
|
|
if (op->type() == type_of<int8_t>()) {
|
|
__CAST_TO_TYPE(int8_t)
|
|
} else if (op->type() == type_of<int16_t>()) {
|
|
__CAST_TO_TYPE(int16_t)
|
|
} else if (op->type() == type_of<int32_t>()) {
|
|
__CAST_TO_TYPE(int32_t)
|
|
} else if (op->type() == type_of<int64_t>()) {
|
|
__CAST_TO_TYPE(int64_t)
|
|
} else if (op->type() == type_of<uint8_t>()) {
|
|
__CAST_TO_TYPE(uint8_t)
|
|
} else if (op->type() == type_of<uint16_t>()) {
|
|
__CAST_TO_TYPE(uint16_t)
|
|
} else if (op->type() == type_of<uint32_t>()) {
|
|
__CAST_TO_TYPE(uint32_t)
|
|
} else if (op->type() == type_of<uint64_t>()) {
|
|
__CAST_TO_TYPE(uint64_t)
|
|
} else if (op->type() == type_of<float>()) {
|
|
__CAST_TO_TYPE(float)
|
|
} else if (op->type() == type_of<double>()) {
|
|
__CAST_TO_TYPE(double)
|
|
} else if (op->type() == type_of<bool>()) {
|
|
__CAST_TO_TYPE(bool)
|
|
} else if (op->type() == type_of<bfloat16>()) {
|
|
// Cannot simplify!!! pass
|
|
__CAST_TO_TYPE(bfloat16)
|
|
} else if (op->type() == type_of<float16>()) {
|
|
// Cannot simplify!!! pass
|
|
__CAST_TO_TYPE(float16)
|
|
} else {
|
|
CINN_NOT_IMPLEMENTED
|
|
}
|
|
}
|
|
#undef __CAST_TO_TYPE
|
|
}
|
|
};
|
|
|
|
struct SimplifyRampMutator : public ir::IRMutator<Expr*> {
|
|
void operator()(Expr* x) { ir::IRMutator<ir::Expr*>::Visit(x, x); }
|
|
|
|
void Visit(const Ramp* op, Expr* expr) override {
|
|
auto* node = expr->As<ir::Ramp>();
|
|
|
|
PADDLE_ENFORCE_EQ(
|
|
IsPureMath(node->base),
|
|
true,
|
|
::common::errors::InvalidArgument("node->base is not a pure math!"));
|
|
PADDLE_ENFORCE_EQ(
|
|
IsPureMath(node->stride),
|
|
true,
|
|
::common::errors::InvalidArgument("node->stride is not a pure math!"));
|
|
node->base = ArithSimplify(node->base);
|
|
node->stride = ArithSimplify(node->stride);
|
|
}
|
|
// ramp + ramp
|
|
void Visit(const Add* op, Expr* expr) override {
|
|
auto* node = expr->As<ir::Add>();
|
|
Expr a = node->a();
|
|
Expr b = node->b();
|
|
auto a_ramp = a.As<ir::Ramp>();
|
|
auto b_ramp = b.As<ir::Ramp>();
|
|
|
|
if (a_ramp && b_ramp && a_ramp->lanes == b_ramp->lanes) {
|
|
Expr base_add = optim::ArithSimplify(a_ramp->base + b_ramp->base);
|
|
Expr stride_add = optim::ArithSimplify(a_ramp->stride + b_ramp->stride);
|
|
*expr = ir::Ramp::Make(base_add, stride_add, a_ramp->lanes);
|
|
}
|
|
}
|
|
};
|
|
|
|
struct SimplifyLoadStoreMutator : public ir::IRMutator<ir::Expr*> {
|
|
void operator()(Expr* x) { ir::IRMutator<ir::Expr*>::Visit(x, x); }
|
|
|
|
void Visit(const Load* expr, Expr* op) override {
|
|
auto* node = op->As<Load>();
|
|
for (auto& idx : node->indices) {
|
|
idx = ArithSimplify(idx);
|
|
}
|
|
}
|
|
|
|
void Visit(const Store* expr, Expr* op) override {
|
|
auto* node = op->As<Store>();
|
|
for (auto& idx : node->indices) {
|
|
idx = ArithSimplify(idx);
|
|
}
|
|
ir::IRMutator<ir::Expr*>::Visit(&node->value, &node->value);
|
|
}
|
|
};
|
|
|
|
struct SimplifyLogicalMutator : public ir::IRMutator<> {
|
|
void operator()(Expr* expr) { ir::IRMutator<>::Visit(expr, expr); }
|
|
|
|
#define DEFINE_VISIT_CMP_OP(OpType, Method) \
|
|
void Visit(const ir::OpType* op, Expr* expr) override { \
|
|
VLOG(7) << "Begin Visit Cmp op: " << *expr; \
|
|
auto* node = expr->As<ir::OpType>(); \
|
|
ir::IRMutator<>::Visit(&node->a(), &node->a()); \
|
|
ir::IRMutator<>::Visit(&node->b(), &node->b()); \
|
|
if (node->a().is_constant() && node->b().is_constant()) { \
|
|
if (node->a().get_constant() Method node->b().get_constant()) { \
|
|
*expr = Expr(true); \
|
|
} else { \
|
|
*expr = Expr(false); \
|
|
} \
|
|
} \
|
|
VLOG(7) << "End Visit Cmp op: " << *expr; \
|
|
}
|
|
DEFINE_VISIT_CMP_OP(LE, <=)
|
|
DEFINE_VISIT_CMP_OP(LT, <)
|
|
DEFINE_VISIT_CMP_OP(GE, >=)
|
|
DEFINE_VISIT_CMP_OP(GT, >)
|
|
DEFINE_VISIT_CMP_OP(EQ, ==)
|
|
DEFINE_VISIT_CMP_OP(NE, !=)
|
|
|
|
#undef DEFINE_VISIT_CMP_OP
|
|
|
|
void Visit(const ir::And* op, Expr* expr) override {
|
|
VLOG(7) << "Begin Visit And op: " << *expr;
|
|
auto* node = expr->As<ir::And>();
|
|
ir::IRMutator<>::Visit(&node->a(), &node->a());
|
|
if (common::IsZero(node->a())) {
|
|
*expr = Expr(false);
|
|
VLOG(7) << "End Visit And op: " << *expr;
|
|
return;
|
|
}
|
|
ir::IRMutator<>::Visit(&node->b(), &node->b());
|
|
if (common::IsZero(node->b())) {
|
|
VLOG(7) << "End Visit And op: " << *expr;
|
|
*expr = Expr(false);
|
|
return;
|
|
}
|
|
if (common::IsOne(node->a()) && common::IsOne(node->b())) {
|
|
*expr = Expr(true);
|
|
} else if (common::IsOne(node->a())) {
|
|
*expr = node->b();
|
|
} else if (common::IsOne(node->b())) {
|
|
*expr = node->a();
|
|
}
|
|
VLOG(7) << "End Visit And op: " << *expr;
|
|
}
|
|
|
|
void Visit(const ir::Or* op, Expr* expr) override {
|
|
VLOG(7) << "Begin Visit Or op: " << *expr;
|
|
auto* node = expr->As<ir::Or>();
|
|
ir::IRMutator<>::Visit(&node->a(), &node->a());
|
|
if (common::IsOne(node->a())) {
|
|
*expr = Expr(true);
|
|
VLOG(7) << "End visit Or op: " << *expr;
|
|
return;
|
|
}
|
|
ir::IRMutator<>::Visit(&node->b(), &node->b());
|
|
if (common::IsOne(node->b())) {
|
|
*expr = Expr(true);
|
|
VLOG(7) << "End visit Or op: " << *expr;
|
|
return;
|
|
}
|
|
if (common::IsZero(node->a()) && common::IsZero(node->b())) {
|
|
*expr = Expr(false);
|
|
} else if (common::IsZero(node->a())) {
|
|
*expr = node->b();
|
|
} else if (common::IsZero(node->b())) {
|
|
*expr = node->a();
|
|
}
|
|
VLOG(7) << "End visit Or op: " << *expr;
|
|
}
|
|
|
|
void Visit(const ir::Not* op, Expr* expr) override {
|
|
VLOG(7) << "Begin Visit Not op: " << *expr;
|
|
auto* node = expr->As<ir::Not>();
|
|
auto v = node->v();
|
|
ir::IRMutator<>::Visit(&v, &v);
|
|
switch (v.node_type()) {
|
|
case ir::IrNodeTy::IntImm:
|
|
case ir::IrNodeTy::UIntImm:
|
|
*expr = common::IsZero(v) ? Expr(true) : Expr(false);
|
|
return;
|
|
case ir::IrNodeTy::Not:
|
|
*expr = v.As<ir::Not>()->v();
|
|
return;
|
|
case ir::IrNodeTy::LE:
|
|
*expr = ir::GT::Make(v->operand(0), v->operand(1));
|
|
return;
|
|
case ir::IrNodeTy::LT:
|
|
*expr = ir::GE::Make(v->operand(0), v->operand(1));
|
|
return;
|
|
case ir::IrNodeTy::GE:
|
|
*expr = ir::LT::Make(v->operand(0), v->operand(1));
|
|
return;
|
|
case ir::IrNodeTy::GT:
|
|
*expr = ir::LE::Make(v->operand(0), v->operand(1));
|
|
return;
|
|
default:
|
|
VLOG(7) << "End Visit Not op: " << *expr;
|
|
return;
|
|
}
|
|
VLOG(7) << "End Visit Not op: " << *expr;
|
|
}
|
|
};
|
|
|
|
struct SimplifyIfThenElseMutator : public ir::ExprMutator<> {
|
|
void operator()(Expr* x) { ir::ExprMutator<>::Visit(x, x); }
|
|
|
|
using ir::ExprMutator<>::Visit;
|
|
|
|
void Visit(const IfThenElse* op, Expr* expr) override {
|
|
auto* node = expr->As<ir::IfThenElse>();
|
|
|
|
auto* condition_int = node->condition.As<ir::IntImm>();
|
|
auto* condition_uint = node->condition.As<ir::UIntImm>();
|
|
|
|
// 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;
|
|
}
|
|
|
|
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({});
|
|
}
|
|
}
|
|
};
|
|
|
|
struct SimplifySelectMutator : public ir::IRMutator<> {
|
|
void operator()(Expr* x) { ir::IRMutator<>::Visit(x, x); }
|
|
|
|
using ir::IRMutator<>::Visit;
|
|
|
|
void Visit(const Select* op, Expr* expr) override {
|
|
auto* node = expr->As<ir::Select>();
|
|
|
|
auto* condition_int = node->condition.As<ir::IntImm>();
|
|
auto* condition_uint = node->condition.As<ir::UIntImm>();
|
|
|
|
// not deterministic
|
|
if (!condition_int && !condition_uint) {
|
|
Visit(&node->true_value, &node->true_value);
|
|
Visit(&node->false_value, &node->false_value);
|
|
return;
|
|
}
|
|
|
|
bool value = condition_int ? condition_int->value : condition_uint->value;
|
|
if (value) {
|
|
*expr = op->true_value;
|
|
Visit(expr, expr);
|
|
} else {
|
|
*expr = op->false_value;
|
|
Visit(expr, expr);
|
|
}
|
|
}
|
|
};
|
|
|
|
/*
|
|
Example 1:
|
|
Select(a <= b, b, a) → max(a, b)
|
|
Example 2:
|
|
Select(a <= b, a, b) → min(a, b)
|
|
Example 3:
|
|
Select(a <= MAX, max(a, MIN), MAX) → min(max(a, MIN), MAX)
|
|
Select(a <= MAX, max(MIN, a), MAX) → min(max(a, MIN), MAX)
|
|
Example 4:
|
|
Select(MIN <= b, min(b, MAX), MIN) → max(min(b, MAX), MIN)
|
|
→ min(max(b, MIN), MAX)
|
|
Select(MIN <= b, min(MAX, b), MIN) → max(min(b, MAX), MIN)
|
|
→ min(max(b, MIN), MAX)
|
|
*/
|
|
struct SimplifySelect2MinMaxMutator : public ir::ExprMutator<> {
|
|
void operator()(Expr* x) { ir::ExprMutator<>::Visit(x, x); }
|
|
|
|
using ir::ExprMutator<>::Visit;
|
|
|
|
// Recursively optimize CompareOp operands
|
|
template <typename T>
|
|
void VisitCompare(T* op, Expr* expr) {
|
|
Expr a = op->a();
|
|
Expr b = op->b();
|
|
ir::ExprMutator<>::Visit(&a, &a);
|
|
ir::ExprMutator<>::Visit(&b, &b);
|
|
|
|
if (a.get() != op->a().get() || b.get() != op->b().get()) {
|
|
*expr = T::Make(a, b);
|
|
}
|
|
}
|
|
|
|
void Visit(const ir::GE* op, Expr* expr) override { VisitCompare(op, expr); }
|
|
void Visit(const ir::GT* op, Expr* expr) override { VisitCompare(op, expr); }
|
|
void Visit(const ir::LE* op, Expr* expr) override { VisitCompare(op, expr); }
|
|
void Visit(const ir::LT* op, Expr* expr) override { VisitCompare(op, expr); }
|
|
|
|
void Visit(const Select* op, Expr* expr) override {
|
|
auto* node = expr->As<ir::Select>();
|
|
|
|
// 1. Recursively optimize sub-expressions
|
|
Expr condition = node->condition;
|
|
Expr true_value = node->true_value;
|
|
Expr false_value = node->false_value;
|
|
|
|
ir::ExprMutator<>::Visit(&condition, &condition);
|
|
ir::ExprMutator<>::Visit(&true_value, &true_value);
|
|
ir::ExprMutator<>::Visit(&false_value, &false_value);
|
|
|
|
// 2. If sub-expressions are modified, rebuild the Select node
|
|
if (condition.get() != node->condition.get() ||
|
|
true_value.get() != node->true_value.get() ||
|
|
false_value.get() != node->false_value.get()) {
|
|
*expr = ir::Select::Make(condition, true_value, false_value);
|
|
node = expr->As<ir::Select>();
|
|
}
|
|
|
|
// 3. Function to optimize Select into Min/Max when possible
|
|
auto TryOptimizeSelect = [&](const Expr& a,
|
|
const Expr& b,
|
|
const Expr& x,
|
|
const Expr& y) -> Expr {
|
|
// Case 1: Select(a <= b, b, a) → max(a, b)
|
|
if (x == b && y == a) {
|
|
if (b.is_constant()) {
|
|
return ir::Max::Make(a, b);
|
|
} else {
|
|
return ir::Max::Make(b, a);
|
|
}
|
|
}
|
|
// Case 2: Select(a <= b, a, b) → min(a, b)
|
|
if (x == a && y == b) {
|
|
if (b.is_constant()) {
|
|
return ir::Min::Make(a, b);
|
|
} else {
|
|
return ir::Min::Make(b, a);
|
|
}
|
|
}
|
|
// Case 3: Select(a <= MAX, max(a, MIN), MAX) → min(max(a, MIN), MAX)
|
|
if (auto* max = x.As<ir::Max>()) {
|
|
if (max->a() == a) {
|
|
if (max->b().is_constant() && y.is_constant() && b.is_constant()) {
|
|
if (y.get_constant() == b.get_constant() &&
|
|
(max->b()).get_constant() <= y.get_constant()) {
|
|
return ir::Min::Make(ir::Max::Make(a, max->b()), b);
|
|
}
|
|
}
|
|
} else if (max->b() == a) {
|
|
// Select(a <= MAX, max(MIN, a), MAX) → min(max(a, MIN), MAX)
|
|
if (max->a().is_constant() && y.is_constant() && b.is_constant()) {
|
|
if (y.get_constant() == b.get_constant() &&
|
|
(max->a()).get_constant() <= y.get_constant()) {
|
|
return ir::Min::Make(ir::Max::Make(a, max->a()), b);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
// Case 4: Select(MIN <= b, min(b, Max), MIN) → max(min(b, MAX), MIN)
|
|
// → min(max(b, MIN), MAX)
|
|
if (auto* min = x.As<ir::Min>()) {
|
|
if (min->a() == b) {
|
|
if ((min->b()).is_constant() && y.is_constant() && a.is_constant()) {
|
|
if (y.get_constant() == a.get_constant() &&
|
|
y.get_constant() <= (min->b()).get_constant()) {
|
|
return ir::Min::Make(ir::Max::Make(b, a), min->b());
|
|
}
|
|
}
|
|
} else if (min->b() == b) {
|
|
// Select(MIN <= b, min(Max, b), MIN) → min(max(b, MIN), MAX)
|
|
if ((min->a()).is_constant() && y.is_constant() && a.is_constant()) {
|
|
if (y.get_constant() == a.get_constant() &&
|
|
y.get_constant() <= (min->a()).get_constant()) {
|
|
return ir::Min::Make(ir::Max::Make(b, a), min->a());
|
|
}
|
|
}
|
|
}
|
|
}
|
|
return Expr(nullptr);
|
|
};
|
|
|
|
// 4. Try to optimize different comparison conditions by converting them to
|
|
// <= logic
|
|
if (auto* ge = node->condition.As<ir::GE>()) {
|
|
// Select(a >= b, t, f) → Select(b <= a, t, f)
|
|
Expr optimized = TryOptimizeSelect(
|
|
ge->b(), ge->a(), node->true_value, node->false_value);
|
|
if (optimized.defined()) {
|
|
*expr = optimized;
|
|
return;
|
|
}
|
|
} else if (auto* gt = node->condition.As<ir::GT>()) {
|
|
// Select(a > b, t, f) → Select(a <= b, f, t)
|
|
Expr optimized = TryOptimizeSelect(
|
|
gt->a(), gt->b(), node->false_value, node->true_value);
|
|
if (optimized.defined()) {
|
|
*expr = optimized;
|
|
return;
|
|
}
|
|
} else if (auto* le = node->condition.As<ir::LE>()) {
|
|
// Select(a <= b, t, f) → Select(a <= b, t, f)
|
|
Expr optimized = TryOptimizeSelect(
|
|
le->a(), le->b(), node->true_value, node->false_value);
|
|
if (optimized.defined()) {
|
|
*expr = optimized;
|
|
return;
|
|
}
|
|
} else if (auto* lt = node->condition.As<ir::LT>()) {
|
|
// Select(a < b, t, f) → Select(b <= a, f, t)
|
|
Expr optimized = TryOptimizeSelect(
|
|
lt->b(), lt->a(), node->false_value, node->true_value);
|
|
if (optimized.defined()) {
|
|
*expr = optimized;
|
|
return;
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
// Optimizes pow(2.0f, ceil(log2(x))) pattern into more efficient bit
|
|
// manipulation:
|
|
// Original: pow(2.0f, ceil(log2(x)))
|
|
// Optimized: ldexpf(1.0f, exponent) where exponent is calculated via:
|
|
// 1. float_as_uint(x) - reinterpret float as uint32
|
|
// 2. right_shift(bits, 23) - extract exponent field
|
|
// 3. (exponent_raw & 0xFF) - 127 - adjust IEEE754 bias
|
|
// 4. +1 if mantissa is non-zero (for ceil behavior)
|
|
struct SimplifyPowerCeilLog2BitOpLdexpfMutator : public ir::ExprMutator<> {
|
|
void operator()(Expr* expr) { ir::ExprMutator<>::Visit(expr, expr); }
|
|
|
|
using ir::ExprMutator<>::Visit;
|
|
void Visit(const ir::Call* op, Expr* expr) override {
|
|
/// 1. First recursively process all sub-expressions
|
|
std::vector<Expr> new_args;
|
|
for (const auto& arg : op->read_args) {
|
|
Expr new_arg = arg;
|
|
Visit(&new_arg, &new_arg);
|
|
new_args.push_back(new_arg);
|
|
}
|
|
|
|
// 2. Match target pattern: pow(base, ceil(log2(x)))
|
|
if (op->name == "pow" && new_args.size() == 2) {
|
|
const Expr& base = new_args[0];
|
|
const Expr& exponent = new_args[1];
|
|
|
|
// Check if exponent is ceil(log2(x))
|
|
if (const ir::Call* ceil_call = exponent.As<ir::Call>()) {
|
|
if (ceil_call->name == "ceil" && ceil_call->read_args.size() == 1) {
|
|
if (const ir::Call* log2_call =
|
|
ceil_call->read_args[0].As<ir::Call>()) {
|
|
if (log2_call->name == "log2" && log2_call->read_args.size() == 1 &&
|
|
log2_call->read_args[0].type().is_float(32)) {
|
|
/// Verify base is 2.0f for optimization
|
|
bool is_base_two = false;
|
|
if (base.is_constant()) {
|
|
if (base.get_constant() == 2.0f) {
|
|
is_base_two = true;
|
|
}
|
|
}
|
|
if (is_base_two) {
|
|
// 3. Replace with bit operations + ldexpf
|
|
Expr x = log2_call->read_args[0]; // Extract log2's argument
|
|
|
|
// Create bit operations to compute ceil(log2(x))
|
|
// (1) Reinterpret float as 32-bit integer
|
|
Expr bits = ir::Call::Make(common::Int(32),
|
|
"__float_as_uint",
|
|
{x},
|
|
{},
|
|
ir::CallType::Extern,
|
|
ir::FunctionRef(),
|
|
0,
|
|
{});
|
|
|
|
std::vector<cinn::ir::Expr> shift_r_args = {bits, ir::Expr(23)};
|
|
Expr shift_r = ir::Call::Make(common::Int(32),
|
|
"right_shift",
|
|
shift_r_args,
|
|
{},
|
|
ir::CallType::Extern,
|
|
ir::FunctionRef(),
|
|
0,
|
|
{});
|
|
// (2) Extract exponent part: ((bits >> 23) & 0xFF) - 127
|
|
std::vector<cinn::ir::Expr> bitwise_and_exp_args = {
|
|
shift_r, ir::Expr(0xFF)};
|
|
Expr bitwise_and_exp = ir::Call::Make(common::Int(32),
|
|
"bitwise_and",
|
|
bitwise_and_exp_args,
|
|
{},
|
|
ir::CallType::Extern,
|
|
ir::FunctionRef(),
|
|
0,
|
|
{});
|
|
Expr exponent_raw =
|
|
ir::Sub::Make(bitwise_and_exp, ir::Expr(127));
|
|
// 3. Check if mantissa is non-zero (i.e., if exponent+1 is
|
|
// needed)
|
|
std::vector<cinn::ir::Expr> bitwise_and_tail_args = {
|
|
bits, ir::Expr(0x007FFFFF)};
|
|
Expr bitwise_and_tail = ir::Call::Make(common::Int(32),
|
|
"bitwise_and",
|
|
bitwise_and_tail_args,
|
|
{},
|
|
ir::CallType::Extern,
|
|
ir::FunctionRef(),
|
|
0,
|
|
{});
|
|
Expr mantissa_non_zero =
|
|
ir::NE::Make(bitwise_and_tail, ir::Expr(0));
|
|
// (4) Check if it's a normal number (exponent != -127)
|
|
Expr is_normal = ir::NE::Make(exponent_raw, ir::Expr(-127));
|
|
// (5) If needed, exponent += 1
|
|
Expr exponent_final = ir::Add::Make(
|
|
exponent_raw,
|
|
ir::Select::Make(
|
|
ir::And::Make(is_normal, mantissa_non_zero),
|
|
ir::Expr(1),
|
|
ir::Expr(0)));
|
|
// (6) Create final expression: ldexpf(1.0f, exponent_final)
|
|
Expr new_expr = ir::Call::Make(op->type(),
|
|
"ldexpf",
|
|
{ir::Expr(1.0f), exponent_final},
|
|
{},
|
|
ir::CallType::Extern,
|
|
ir::FunctionRef(),
|
|
0,
|
|
{});
|
|
*expr = new_expr;
|
|
return;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// For non-target patterns, reconstruct as-is
|
|
if (new_args != op->read_args) {
|
|
*expr = ir::Call::Make(op->type(),
|
|
op->name,
|
|
new_args,
|
|
op->write_args,
|
|
op->call_type,
|
|
op->func,
|
|
op->value_index,
|
|
op->attrs);
|
|
}
|
|
}
|
|
};
|
|
|
|
struct SimplifyUnitBlockMutator : public ir::ExprMutator<> {
|
|
void operator()(Expr* x) { ir::ExprMutator<ir::Expr*>::Visit(x, x); }
|
|
|
|
using ir::ExprMutator<>::Visit;
|
|
|
|
void Visit(const Block* op, Expr* expr) override {
|
|
auto* node = expr->As<ir::Block>();
|
|
|
|
if (node->stmts.size() == 1 && node->stmts[0].As<ir::Block>()) {
|
|
VLOG(6) << "Simplify size-1 ir::Block";
|
|
*expr = node->stmts[0];
|
|
Visit(expr, expr);
|
|
} else {
|
|
for (auto& s : node->stmts) {
|
|
Visit(&s, &s);
|
|
}
|
|
std::vector<Expr> stmts;
|
|
for (auto& s : node->stmts) {
|
|
if (s.As<ir::Block>()) {
|
|
VLOG(6) << "Simplify ir::Block inside ir::Block";
|
|
auto inner_block = s.As<ir::Block>();
|
|
for (auto inner_stmt : inner_block->stmts) {
|
|
stmts.push_back(inner_stmt);
|
|
}
|
|
} else {
|
|
stmts.push_back(s);
|
|
}
|
|
}
|
|
expr->As<ir::Block>()->stmts = stmts;
|
|
}
|
|
}
|
|
|
|
void Visit(const ScheduleBlock* op, Expr* expr) override {
|
|
auto* node = expr->As<ScheduleBlock>();
|
|
PADDLE_ENFORCE_NOT_NULL(node,
|
|
::common::errors::InvalidArgument(
|
|
"The node expr->As<ScheduleBlock>() is null"));
|
|
|
|
if (node->body.As<Block>()) {
|
|
if (node->body.As<Block>()->stmts.size() == 1) {
|
|
node->body = node->body.As<Block>()->stmts[0];
|
|
}
|
|
}
|
|
Visit(&(node->body), &(node->body));
|
|
}
|
|
};
|
|
|
|
struct SimplifyUnitLoopMutator : public ir::IRMutator<> {
|
|
paddle::flat_hash_map<std::string, Expr> var_mins;
|
|
void operator()(Expr* x) { ir::IRMutator<ir::Expr*>::Visit(x, x); }
|
|
|
|
using ir::IRMutator<>::Visit;
|
|
|
|
void Visit(const For* op, Expr* expr) override {
|
|
auto* node = expr->As<ir::For>();
|
|
Visit(&node->min, &node->min);
|
|
Visit(&node->extent, &node->extent);
|
|
auto* min_i = node->min.As<IntImm>();
|
|
auto* extent_i = node->extent.As<IntImm>();
|
|
if (min_i && extent_i && extent_i->value - min_i->value == 1) {
|
|
VLOG(6) << "Simplify current Unit For Loop";
|
|
std::string var_name = node->loop_var->name;
|
|
var_mins.emplace(var_name, node->min);
|
|
|
|
*expr = node->body;
|
|
|
|
Visit(expr, expr);
|
|
var_mins.erase(var_name);
|
|
} else {
|
|
Visit(&node->body, &node->body);
|
|
}
|
|
}
|
|
|
|
void Visit(const _Var_* op, Expr* expr) override {
|
|
auto* node = expr->As<ir::_Var_>();
|
|
|
|
if (var_mins.count(node->name)) {
|
|
*expr = var_mins.at(node->name);
|
|
}
|
|
}
|
|
};
|
|
} // namespace
|
|
|
|
void SimplifyCast(Expr* expr) { SimplifyCastMutator()(expr); }
|
|
void SimplifyUnitLoop(Expr* expr) { SimplifyUnitLoopMutator()(expr); }
|
|
void SimplifyUnitBlock(Expr* expr) { SimplifyUnitBlockMutator()(expr); }
|
|
|
|
void SimplifyLogical(Expr* expr) { SimplifyLogicalMutator()(expr); }
|
|
void SimplifyNoPureMath(Expr* expr, const ir::IndexExpr::OptLevel& opt_level) {
|
|
auto mutator = SimplifyNoPureMathMutator(opt_level);
|
|
mutator(expr);
|
|
}
|
|
|
|
Expr ArithSimplify(const Expr& u, const ir::IndexExpr::OptLevel& opt_level) {
|
|
VLOG(3) << "Begin ArithSimplify " << u;
|
|
if (!u.is_index()) return u;
|
|
auto copied = ir_utils::IRCopy(u);
|
|
auto res = copied.as_index().Normalize(opt_level);
|
|
VLOG(3) << "End ArithSimplify " << res;
|
|
return res;
|
|
}
|
|
|
|
void Simplify(Expr* expr) {
|
|
VLOG(6) << "Begin Simplify " << *expr;
|
|
ReplaceFracWithDivMutator()(expr);
|
|
SimplifyNoPureMathMutator()(expr);
|
|
SimplifyCastMutator()(expr);
|
|
SimplifyRampMutator()(expr);
|
|
SimplifyLoadStoreMutator()(expr);
|
|
SimplifyLogicalMutator()(expr);
|
|
SimplifyIfThenElseMutator()(expr);
|
|
SimplifySelectMutator()(expr);
|
|
SimplifySelect2MinMaxMutator()(expr);
|
|
SimplifyPowerCeilLog2BitOpLdexpfMutator()(expr);
|
|
SimplifyNoPureMathMutator()(expr);
|
|
VLOG(6) << "End Simplify " << *expr;
|
|
}
|
|
} // namespace optim
|
|
} // namespace cinn
|