Files
paddlepaddle--paddle/paddle/cinn/common/ir_util.cc
T
2026-07-13 12:40:42 +08:00

650 lines
22 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/common/ir_util.h"
#include <algorithm>
#include <stack>
#include <unordered_set>
#include "paddle/cinn/common/const_fold.h"
#include "paddle/cinn/common/simplify_special_pattern.h"
#include "paddle/cinn/ir/ir_mutator.h"
#include "paddle/cinn/ir/ir_printer.h"
#include "paddle/cinn/ir/op/ir_operators.h"
#include "paddle/cinn/ir/utils/ir_compare.h"
#include "paddle/cinn/ir/utils/ir_copy.h"
#include "paddle/cinn/optim/ir_simplify.h"
#include "paddle/cinn/utils/string.h"
#include "paddle/common/enforce.h"
namespace cinn {
namespace common {
namespace {
// ramp + scalar or broadcast
Expr RampRelatedMul(ir::Ramp *ramp, Expr other) {
PADDLE_ENFORCE_EQ(
other.type().ElementOf(),
Int(32),
::common::errors::InvalidArgument("The type of other should be int32."));
PADDLE_ENFORCE_EQ(ramp->base.type(),
Int(32),
::common::errors::InvalidArgument(
"The type of ramp->base should be int32."));
PADDLE_ENFORCE_EQ(ramp->stride.type(),
Int(32),
::common::errors::InvalidArgument(
"The type of ramp->stride should be int32."));
auto *other_broadcast = other.As<ir::Broadcast>();
if (other_broadcast) {
PADDLE_ENFORCE_EQ(ramp->lanes,
other_broadcast->lanes,
::common::errors::InvalidArgument(
"The lanes of ramp and other should be equal."));
other = other_broadcast->value;
}
return ir::Ramp::Make(ramp->base * other, ramp->stride * other, ramp->lanes);
}
Expr RampRelatedMul(ir::Broadcast *broadcast, Expr other) {
PADDLE_ENFORCE_EQ(
other.type().lanes(),
1,
::common::errors::InvalidArgument("The lanes of other should be 1."));
return ir::Broadcast::Make(broadcast->value * other, broadcast->lanes);
}
// ramp * ramp
Expr RampRelatedMul(ir::Ramp *ramp, ir::Ramp *other) {
CINN_NOT_IMPLEMENTED
return Expr();
}
// ramp + scalar
Expr RampRelatedAdd(ir::Ramp *ramp, Expr other) {
PADDLE_ENFORCE_EQ(
other.type().ElementOf(),
Int(32),
::common::errors::InvalidArgument("The type of other should be int32."));
auto *other_broadcast = other.As<ir::Broadcast>();
if (other_broadcast) {
PADDLE_ENFORCE_EQ(ramp->lanes,
other_broadcast->lanes,
::common::errors::InvalidArgument(
"The lanes of ramp and other should be equal."));
other = other_broadcast->value;
}
return ir::Ramp::Make(ramp->base + other, ramp->stride, ramp->lanes);
}
Expr RampRelatedAdd(ir::Broadcast *broadcast, Expr other) {
PADDLE_ENFORCE_EQ(
other.type().lanes(),
1,
::common::errors::InvalidArgument("The lanes of other should be 1."));
return ir::Broadcast::Make(broadcast->value + other, broadcast->lanes);
}
// ramp + ramp
Expr RampRelatedAdd(ir::Ramp *ramp, ir::Ramp *other) {
PADDLE_ENFORCE_NOT_NULL(
ramp,
::common::errors::InvalidArgument("Ramp pointer should not be null."));
PADDLE_ENFORCE_NOT_NULL(other,
::common::errors::InvalidArgument(
"Other ramp pointer should not be null."));
if (ramp->lanes == other->lanes) {
Expr base_add = optim::ArithSimplify(ramp->base + other->base);
Expr stride_add = optim::ArithSimplify(ramp->stride + other->stride);
VLOG(2) << base_add;
VLOG(2) << stride_add;
return ir::Ramp::Make(base_add, stride_add, ramp->lanes);
}
CINN_NOT_IMPLEMENTED
return Expr();
}
Expr RampRelatedAdd(Expr a, Expr b) {
auto *a_ramp = a.As<ir::Ramp>();
auto *b_ramp = b.As<ir::Ramp>();
auto *a_broadcast = a.As<ir::Broadcast>();
auto *b_broadcast = b.As<ir::Broadcast>();
if (a_ramp && !b_ramp && (b->type().lanes() == 1 || b_broadcast)) {
return RampRelatedAdd(a_ramp, b);
} else if (!a_ramp && b_ramp && (a->type().lanes() == 1 || a_broadcast)) {
return RampRelatedAdd(b_ramp, a);
} else if (!a_ramp && !b_ramp && !a->type().is_vector() &&
!b->type().is_vector()) {
return a + b;
} else if (a_ramp && b_ramp) { // a_ramp && b_ramp
return RampRelatedAdd(a_ramp, b_ramp);
} else if (a_broadcast && !b_broadcast) {
return RampRelatedAdd(a_broadcast, b);
} else if (!a_broadcast && b_broadcast) {
return RampRelatedAdd(b_broadcast, a);
} else if (a_broadcast && b_broadcast) {
PADDLE_ENFORCE_EQ(
a_broadcast->lanes,
b_broadcast->lanes,
::common::errors::InvalidArgument(
"The lanes of a_broadcast and b_broadcast should be equal."));
return ir::Broadcast::Make(a_broadcast->value + b_broadcast->value,
a_broadcast->lanes);
} else {
CINN_NOT_IMPLEMENTED
}
}
Expr RampRelatedMul(Expr a, Expr b) {
auto *a_ramp = a.As<ir::Ramp>();
auto *b_ramp = b.As<ir::Ramp>();
auto *a_broadcast = a.As<ir::Broadcast>();
auto *b_broadcast = b.As<ir::Broadcast>();
if (a_ramp && !b_ramp && (!b->type().is_vector() || b_broadcast)) {
return RampRelatedMul(a_ramp, b);
} else if (!a_ramp && b_ramp && (a->type().is_vector() || a_broadcast)) {
return RampRelatedMul(b_ramp, a);
} else if (!a_ramp && !b_ramp && !a->type().is_vector() &&
!b->type().is_vector()) {
return a * b;
} else if (a_ramp && b_ramp) { // a_ramp && b_ramp
return RampRelatedMul(a_ramp, b_ramp);
} else if (a_broadcast && !b_broadcast) {
return RampRelatedMul(a_broadcast, b);
} else if (!a_broadcast && b_broadcast) {
return RampRelatedMul(b_broadcast, a);
} else if (a_broadcast && b_broadcast) {
PADDLE_ENFORCE_EQ(
a_broadcast->lanes,
b_broadcast->lanes,
::common::errors::InvalidArgument(
"The lanes of a_broadcast and b_broadcast should be equal."));
return ir::Broadcast::Make(a_broadcast->value * b_broadcast->value,
a_broadcast->lanes);
} else {
VLOG(3) << "a,b: " << a << " " << b;
CINN_NOT_IMPLEMENTED
}
}
} // namespace
Expr IndiceToAbsOffset(const std::vector<Expr> &shape,
const std::vector<Expr> &indices) {
VLOG(3) << "Begin IndiceToAbsOffset";
VLOG(3) << "shape is : " << utils::Join(shape, ",");
VLOG(3) << "indices is : " << utils::Join(indices, ",");
PADDLE_ENFORCE_LE(shape.size(),
indices.size(),
::common::errors::InvalidArgument(
"The size of shape should be less than or "
"equal to the size of indices."));
Expr res(0);
for (int32_t i = 0; i < shape.size(); i++) {
PADDLE_ENFORCE_EQ(
shape[i].type() == Int(64) || shape[i].type() == Int(32),
true,
::common::errors::InvalidArgument(
"The shape data type currently supports only int32 or int64, but "
"the current data type of shape[%d] is %s",
i,
shape[i].type()));
Expr indice_cast = indices[i];
optim::SimplifyCast(&indice_cast);
res = RampRelatedAdd(RampRelatedMul(res, shape[i]), indice_cast);
if (res.is_index()) {
res = res.as_index().Normalize(ir::IndexExpr::OptLevel::kLevel2);
} else {
VLOG(8) << "**** expr is not index ****: " << res;
}
}
VLOG(3) << "End IndiceToAbsOffset";
return res;
}
Expr IndiceToAbsOffset(const std::vector<int> &shape,
const std::vector<Expr> &indices) {
std::vector<Expr> shape_;
for (int v : shape) shape_.push_back(Expr(v));
return IndiceToAbsOffset(shape, indices);
}
Expr PrecedingAxisToAbsOffset(const std::vector<Expr> &shape,
int preceding_n_axis) {
std::vector<Expr> indices;
for (int i = 0; i < preceding_n_axis; i++) indices.push_back(shape[i]);
return IndiceToAbsOffset(shape, indices);
}
namespace {
class SubstituteMutator : ir::IRMutator<ir::Expr *> {
public:
explicit SubstituteMutator(const std::map<const ir::_Var_ *, Expr> &var_map) {
for (auto &item : var_map) {
var_map_[item.first->name] = item.second;
}
}
void operator()(ir::Expr *expr) { Visit(expr); }
private:
void Visit(Expr *expr) { ir::IRMutator<>::Visit(expr, expr); }
void Visit(const ir::_Var_ *op, ir::Expr *expr) override {
auto it = var_map_.find(op->name);
if (it == var_map_.end()) return;
*expr = it->second;
}
Expr *expr_{};
std::map<std::string, Expr> var_map_;
};
} // namespace
void Substitute(Expr *expr, const std::map<const ir::_Var_ *, Expr> &var_map) {
SubstituteMutator mutator(var_map);
mutator(expr);
}
bool is_zero(Expr v) {
v = optim::ArithSimplify(v);
auto *int_n = v.As<ir::IntImm>();
auto *float_n = v.As<ir::FloatImm>();
if (int_n) return int_n->value == 0;
if (float_n) return float_n->value == 0.f;
return false;
}
Expr NormalizeUpperBound(Expr upper_bound, bool minus_one /* = true */) {
if (upper_bound == SymbolicExprLimit::positive_inf) {
return upper_bound;
}
if (minus_one) {
return upper_bound - ir::Expr(1); // [lower, upper) to [lower, upper]
}
return upper_bound + ir::Expr(1); // (lower, upper] to [lower, upper)
}
Expr CastIfNeeded(Expr body, Type type) {
if (body.type() == type) return body;
return ir::Cast::Make(type, body);
}
bool MathEqual(const Expr &a, const Expr &b) {
auto c = a - b;
c = optim::ArithSimplify(c);
return is_zero(c);
}
Expr select(Expr cond, Expr true_value, Expr false_value) {
return ir::Select::Make(cond, true_value, false_value);
}
Expr and_all(const std::vector<Expr> &conds) {
PADDLE_ENFORCE_NE(conds.empty(),
true,
::common::errors::InvalidArgument(
"The conditions vector should not be empty."));
Expr res = conds.front();
for (int i = 1; i < conds.size(); i++) {
res = ir::And::Make(res, conds[i]);
}
return res;
}
Expr or_all(const std::vector<Expr> &conds) {
PADDLE_ENFORCE_NE(conds.empty(),
true,
::common::errors::InvalidArgument(
"The conditions vector should not be empty."));
Expr res = conds.front();
for (int i = 1; i < conds.size(); i++) {
res = ir::Or::Make(res, conds[i]);
}
return res;
}
void CheckTensorUniqueInExpr(Expr expr) {
auto tensor_uniq = ir::ir_utils::CollectIRNodes(
expr, [](const Expr *x) { return x->as_tensor(); });
paddle::flat_hash_map<std::string, const ir::_Tensor_ *> tensor_names;
for (auto &t : tensor_uniq) {
auto *tp = t.as_tensor();
if (!tensor_names.count(tp->name)) {
tensor_names[tp->name] = tp;
} else {
PADDLE_ENFORCE_EQ(
tensor_names[tp->name],
tp,
::common::errors::InvalidArgument(
"Found tensor not unique, The original express is %d .", expr));
}
}
}
Expr cast(Expr e, Type type) {
if (e.is_constant()) {
if (type.is_bool()) {
return Expr(static_cast<bool>(e.get_constant()));
} else if (type.is_int(8)) {
return Expr(static_cast<int8_t>(e.get_constant()));
} else if (type.is_int(16)) {
return Expr(static_cast<int16_t>(e.get_constant()));
} else if (type.is_int(32)) {
return Expr(static_cast<int32_t>(e.get_constant()));
} else if (type.is_int(64)) {
return Expr(static_cast<int64_t>(e.get_constant()));
} else if (type.is_uint(8)) {
return Expr(static_cast<uint8_t>(e.get_constant()));
} else if (type.is_uint(16)) {
return Expr(static_cast<uint16_t>(e.get_constant()));
} else if (type.is_uint(32)) {
return Expr(static_cast<uint32_t>(e.get_constant()));
} else if (type.is_uint(64)) {
return Expr(static_cast<uint64_t>(e.get_constant()));
} else if (type.is_float(32)) {
return Expr(static_cast<float>(e.get_constant()));
} else if (type.is_float(64)) {
return Expr(static_cast<double>(e.get_constant()));
} else if (type.is_bfloat16()) {
return Expr(static_cast<cinn::common::bfloat16>(e.get_constant()));
} else if (type.is_float16()) {
return Expr(static_cast<cinn::common::float16>(e.get_constant()));
} else {
CINN_NOT_IMPLEMENTED
}
}
return ir::Cast::Make(type, e);
}
std::vector<std::string> GatherItersToTensorProducer(
const std::string &target_tensor_name, Expr *expr) {
struct Visitor : public ir::IRMutator<> {
std::vector<std::string> iters;
const std::string &target_tensor_name;
explicit Visitor(const std::string &target_tensor_name)
: target_tensor_name(target_tensor_name) {}
std::vector<std::string> operator()(Expr *expr) {
ir::IRMutator<>::Visit(expr, expr);
return iters;
}
void Visit(const ir::Store *op, Expr *expr) {
if (op->tensor.as_tensor()->name == target_tensor_name) {
PADDLE_ENFORCE_EQ(iters.empty(),
true,
::common::errors::InvalidArgument(
"The iterators vector should be empty."));
for (auto &e : for_stack) {
auto *for_n = e->As<ir::For>();
auto *polyfor_n = e->As<ir::PolyFor>();
if (for_n) {
iters.push_back(for_n->loop_var->name);
} else {
iters.push_back(polyfor_n->iterator->name);
}
}
}
}
void Visit(const ir::For *op, Expr *expr) {
for_stack.push_back(expr);
ir::IRMutator<>::Visit(op, expr);
for_stack.pop_back();
}
void Visit(const ir::PolyFor *op, Expr *expr) {
for_stack.push_back(expr);
ir::IRMutator<>::Visit(op, expr);
for_stack.pop_back();
}
std::vector<Expr *> for_stack;
};
return Visitor(target_tensor_name)(expr);
}
std::vector<Expr *> GetForloopStackToStore(Expr *expr,
const std::string &tensor_name) {
VLOG(4) << "search store " << tensor_name << " in expr:\n";
VLOG(4) << *expr;
struct Mutator : public ir::IRMutator<> {
std::vector<Expr *> forloop_stack;
bool found{false};
std::string tensor_name;
explicit Mutator(const std::string &tensor_name)
: tensor_name(tensor_name) {}
std::vector<Expr *> operator()(Expr *expr) {
ir::IRMutator<>::Visit(expr, expr);
return forloop_stack;
}
void Visit(const ir::For *op, Expr *expr) {
auto *node = expr->As<ir::For>();
forloop_stack.push_back(expr);
ir::IRMutator<>::Visit(&node->body, &node->body);
if (!found) forloop_stack.pop_back();
}
void Visit(const ir::PolyFor *op, Expr *expr) {
auto *node = expr->As<ir::PolyFor>();
forloop_stack.push_back(expr);
ir::IRMutator<>::Visit(&node->body, &node->body);
if (!found) forloop_stack.pop_back();
}
void Visit(const ir::Store *op, Expr *expr) {
found = op->tensor.as_tensor()->name == tensor_name;
}
};
return Mutator(tensor_name)(expr);
}
Expr max(Expr a, Expr b) {
PADDLE_ENFORCE_EQ(a.type(),
b.type(),
::common::errors::InvalidArgument(
"The type of a and b should be equal."));
return ir::Max::Make(a, b);
}
Expr min(Expr a, Expr b) {
PADDLE_ENFORCE_EQ(a.type(),
b.type(),
::common::errors::InvalidArgument(
"The type of a and b should be equal."));
return ir::Min::Make(a, b);
}
void OpDataTypePromote(Expr *expr) {
struct TypePromote : public ir::IRMutator<> {
void operator()(Expr *expr) { ir::IRMutator<>::Visit(expr, expr); }
// type promote for operand of binary op
#define __(op__) \
void Visit(const ir::op__ *op, ir::Expr *expr) override { \
ir::TryElevateInt32ToInt64_((*expr)->operands); \
IRMutator::Visit(op, expr); \
};
__(Sum)
__(Product)
NODETY_BINARY_OP_FOR_EACH(__)
#undef __
void Visit(const ir::Select *op, ir::Expr *expr) override {
auto node = expr->As<ir::Select>();
auto promote_args = std::move(
ir::TryElevateInt32ToInt64({node->true_value, node->false_value}));
node->true_value = promote_args.at(0);
node->false_value = promote_args.at(1);
IRMutator::Visit(op, expr);
}
void Visit(const ir::Load *op, ir::Expr *expr) {
auto node = expr->As<ir::Load>();
ir::TryElevateInt32ToInt64_(node->indices);
IRMutator::Visit(op, expr);
}
void Visit(const ir::Store *op, ir::Expr *expr) {
auto node = expr->As<ir::Store>();
ir::TryElevateInt32ToInt64_(node->indices);
IRMutator::Visit(op, expr);
}
void Visit(const ir::Let *op, ir::Expr *expr) {
auto node = expr->As<ir::Let>();
// For Symbol of LetOp, we need to insert a cast to convert its type, but
// inside LetOp, we should directly convert the Symbol type instead of
// inserting a cast.so we set the flag to false before the conversion and
// set it to true after the conversion, e.g.
// inside LetOp: type of v, v1 are int32.
// int32 v = v1 * 2 ==TypePromote==> int64 v = v1 * 2ll
// outside LetOp: type of v, v2 are int32 and v is defined by LetOp.
// v2 = v * 2 ==TypePromote==> v2 = (int64)v * 2ll
if (node->symbol.is_var()) {
node->symbol.as_var()->is_let_symbol = false;
}
auto promote_args =
std::move(ir::TryElevateInt32ToInt64({node->symbol, node->body}));
node->symbol = promote_args.at(0);
node->body = promote_args.at(1);
if (node->symbol.is_var()) {
node->symbol.as_var()->is_let_symbol = true;
}
IRMutator::Visit(op, expr);
}
};
TypePromote visitor;
visitor(expr);
}
void OpDataTypePromote(ir::Module *module) {
auto node = module->As<ir::_Module_>();
for (auto &func : node->functions) {
OpDataTypePromote(&func->body);
}
for (auto &buffer : node->buffers) {
OpDataTypePromote(&buffer);
}
for (auto &submodule : node->submodules) {
OpDataTypePromote(&submodule);
}
}
void OpDataTypePromote(ir::LoweredFunc *func) {
auto node = func->As<ir::_LoweredFunc_>();
OpDataTypePromote(&node->body);
}
struct DynamicSymbolExprBitTracker : public ir::IRVisitor {
DynamicSymbolExprBitTracker() = default;
#define TRAVERSE_EXPR_FIELDS(NodeType) \
void Visit(const ir::NodeType *x) override { \
if (dyn_symbol_bit == 64) return; \
for (auto expr : x->expr_fields()) { \
IRVisitor::Visit(expr); \
} \
}
NODETY_BINARY_OP_FOR_EACH(TRAVERSE_EXPR_FIELDS)
NODETY_UNARY_OP_FOR_EACH(TRAVERSE_EXPR_FIELDS)
NODETY_CONTROL_OP_FOR_INTRINSIC(TRAVERSE_EXPR_FIELDS)
TRAVERSE_EXPR_FIELDS(Cast)
TRAVERSE_EXPR_FIELDS(For)
TRAVERSE_EXPR_FIELDS(PolyFor)
TRAVERSE_EXPR_FIELDS(Select)
TRAVERSE_EXPR_FIELDS(IfThenElse)
TRAVERSE_EXPR_FIELDS(Block)
TRAVERSE_EXPR_FIELDS(Call)
TRAVERSE_EXPR_FIELDS(Load)
TRAVERSE_EXPR_FIELDS(Store)
TRAVERSE_EXPR_FIELDS(Alloc)
TRAVERSE_EXPR_FIELDS(Free)
TRAVERSE_EXPR_FIELDS(_Buffer_)
TRAVERSE_EXPR_FIELDS(_Tensor_)
TRAVERSE_EXPR_FIELDS(Let)
TRAVERSE_EXPR_FIELDS(Reduce)
TRAVERSE_EXPR_FIELDS(Ramp)
TRAVERSE_EXPR_FIELDS(Broadcast)
TRAVERSE_EXPR_FIELDS(FracOp)
TRAVERSE_EXPR_FIELDS(Product)
TRAVERSE_EXPR_FIELDS(Sum)
TRAVERSE_EXPR_FIELDS(PrimitiveNode)
TRAVERSE_EXPR_FIELDS(_BufferRange_)
TRAVERSE_EXPR_FIELDS(ScheduleBlock)
TRAVERSE_EXPR_FIELDS(ScheduleBlockRealize)
TRAVERSE_EXPR_FIELDS(_Dim_)
#undef TRAVERSE_EXPR_FIELDS
void Visit(const ir::_Var_ *x) override {
auto it = search_map->find(x->name);
if (it != search_map->cend()) {
int num_bits = it->second.bits();
if (num_bits == 32 || num_bits == 64) {
dyn_symbol_bit = std::max(dyn_symbol_bit, num_bits);
}
}
}
int operator()(const std::unordered_map<std::string, common::Type> *map_,
const Expr *e) {
// no need to continue if the result is already 64
if (dyn_symbol_bit == 64) return 64;
search_map = map_;
IRVisitor::Visit(e);
return dyn_symbol_bit;
}
const std::unordered_map<std::string, common::Type> *search_map;
int dyn_symbol_bit = 0;
};
#define VISIT_OP(NodeType) \
std::pair<int, bool> UnifiedOperandTypeBits( \
const std::unordered_map<std::string, common::Type> *search_map, \
const ir::NodeType *node) { \
if (search_map->empty()) return {0, false}; \
if (!node->a().type().is_int() || !node->b().type().is_int()) \
return {0, false}; \
int node_a_bits = node->a().type().bits(); \
int node_b_bits = node->b().type().bits(); \
if (node_a_bits < 32 || node_b_bits < 32) return {0, false}; \
DynamicSymbolExprBitTracker tracker; \
int b1 = tracker(search_map, &node->a()); \
tracker.dyn_symbol_bit = 0; \
int b2 = tracker(search_map, &node->b()); \
return std::make_pair(std::max(b1, b2), b1 > 0 && b2 > 0); \
}
VISIT_OP(Min)
VISIT_OP(Max)
#undef VISIT_OP
} // namespace common
} // namespace cinn