// 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 #include #include #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(); 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(); 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(); auto *b_ramp = b.As(); auto *a_broadcast = a.As(); auto *b_broadcast = b.As(); 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(); auto *b_ramp = b.As(); auto *a_broadcast = a.As(); auto *b_broadcast = b.As(); 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 &shape, const std::vector &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 &shape, const std::vector &indices) { std::vector shape_; for (int v : shape) shape_.push_back(Expr(v)); return IndiceToAbsOffset(shape, indices); } Expr PrecedingAxisToAbsOffset(const std::vector &shape, int preceding_n_axis) { std::vector indices; for (int i = 0; i < preceding_n_axis; i++) indices.push_back(shape[i]); return IndiceToAbsOffset(shape, indices); } namespace { class SubstituteMutator : ir::IRMutator { public: explicit SubstituteMutator(const std::map &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 var_map_; }; } // namespace void Substitute(Expr *expr, const std::map &var_map) { SubstituteMutator mutator(var_map); mutator(expr); } bool is_zero(Expr v) { v = optim::ArithSimplify(v); auto *int_n = v.As(); auto *float_n = v.As(); 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 &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 &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 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(e.get_constant())); } else if (type.is_int(8)) { return Expr(static_cast(e.get_constant())); } else if (type.is_int(16)) { return Expr(static_cast(e.get_constant())); } else if (type.is_int(32)) { return Expr(static_cast(e.get_constant())); } else if (type.is_int(64)) { return Expr(static_cast(e.get_constant())); } else if (type.is_uint(8)) { return Expr(static_cast(e.get_constant())); } else if (type.is_uint(16)) { return Expr(static_cast(e.get_constant())); } else if (type.is_uint(32)) { return Expr(static_cast(e.get_constant())); } else if (type.is_uint(64)) { return Expr(static_cast(e.get_constant())); } else if (type.is_float(32)) { return Expr(static_cast(e.get_constant())); } else if (type.is_float(64)) { return Expr(static_cast(e.get_constant())); } else if (type.is_bfloat16()) { return Expr(static_cast(e.get_constant())); } else if (type.is_float16()) { return Expr(static_cast(e.get_constant())); } else { CINN_NOT_IMPLEMENTED } } return ir::Cast::Make(type, e); } std::vector GatherItersToTensorProducer( const std::string &target_tensor_name, Expr *expr) { struct Visitor : public ir::IRMutator<> { std::vector iters; const std::string &target_tensor_name; explicit Visitor(const std::string &target_tensor_name) : target_tensor_name(target_tensor_name) {} std::vector 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(); auto *polyfor_n = e->As(); 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 for_stack; }; return Visitor(target_tensor_name)(expr); } std::vector 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 forloop_stack; bool found{false}; std::string tensor_name; explicit Mutator(const std::string &tensor_name) : tensor_name(tensor_name) {} std::vector operator()(Expr *expr) { ir::IRMutator<>::Visit(expr, expr); return forloop_stack; } void Visit(const ir::For *op, Expr *expr) { auto *node = expr->As(); 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(); 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(); 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::TryElevateInt32ToInt64_(node->indices); IRMutator::Visit(op, expr); } void Visit(const ir::Store *op, ir::Expr *expr) { auto node = expr->As(); ir::TryElevateInt32ToInt64_(node->indices); IRMutator::Visit(op, expr); } void Visit(const ir::Let *op, ir::Expr *expr) { auto node = expr->As(); // 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(); 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(); 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 *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 *search_map; int dyn_symbol_bit = 0; }; #define VISIT_OP(NodeType) \ std::pair UnifiedOperandTypeBits( \ const std::unordered_map *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