wehub-resource-sync
202 lines
6.4 KiB
C++
202 lines
6.4 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.
|
|
|
|
#pragma once
|
|
|
|
#include <map>
|
|
#include <memory>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
#include "paddle/cinn/common/bfloat16.h"
|
|
#include "paddle/cinn/common/float16.h"
|
|
#include "paddle/cinn/common/float8e4m3.h"
|
|
#include "paddle/cinn/common/integer_set.h"
|
|
#include "paddle/cinn/ir/ir.h"
|
|
#include "paddle/utils/flat_hash_map.h"
|
|
|
|
namespace cinn {
|
|
namespace common {
|
|
|
|
Expr IndiceToAbsOffset(const std::vector<Expr> &shape,
|
|
const std::vector<Expr> &indices);
|
|
Expr IndiceToAbsOffset(const std::vector<int> &shape,
|
|
const std::vector<Expr> &indices);
|
|
|
|
Expr PrecedingAxisToAbsOffset(const std::vector<Expr> &shape,
|
|
int preceding_n_axis);
|
|
|
|
Expr CastIfNeeded(Expr body, Type type);
|
|
|
|
//! Substitute vars to other expressions.
|
|
//! @param expr The expression to do modification.
|
|
//! @param var_map The map from variables to the target expressions.
|
|
void Substitute(Expr *expr, const std::map<const ir::_Var_ *, Expr> &var_map);
|
|
|
|
//! Get a stack of forloops(For and PolyFor nodes) to a Store node target to \p
|
|
//! tensor_name
|
|
std::vector<Expr *> GetForloopStackToStore(Expr *expr,
|
|
const std::string &tensor_name);
|
|
|
|
// make const
|
|
// @{
|
|
inline Expr make_const(int32_t x) { return Expr(static_cast<int32_t>(x)); }
|
|
inline Expr make_const(int64_t x) { return Expr(static_cast<int64_t>(x)); }
|
|
inline Expr make_const(bfloat16 x) { return Expr(static_cast<bfloat16>(x)); }
|
|
inline Expr make_const(float8e4m3 x) {
|
|
return Expr(static_cast<float8e4m3>(x));
|
|
}
|
|
inline Expr make_const(float16 x) { return Expr(static_cast<float16>(x)); }
|
|
inline Expr make_const(float x) { return Expr(static_cast<float>(x)); }
|
|
inline Expr make_const(double x) { return Expr(static_cast<double>(x)); }
|
|
inline Expr make_const(bool x) { return Expr(static_cast<bool>(x)); }
|
|
// @}
|
|
|
|
//! maker for some general consts.
|
|
// @{
|
|
template <typename T = int32_t>
|
|
inline Expr make_zero() {
|
|
return make_const(static_cast<T>(0));
|
|
}
|
|
template <typename T = int32_t>
|
|
inline Expr make_one() {
|
|
return make_const(static_cast<T>(1));
|
|
}
|
|
inline Expr make_bool(bool x) {
|
|
return cinn::common::make_shared<ir::UIntImm>(Bool(), x);
|
|
}
|
|
inline Expr make_bool(bool x, int lanes) {
|
|
return cinn::common::make_shared<ir::UIntImm>(Bool(lanes), x);
|
|
}
|
|
// @}
|
|
|
|
/**
|
|
* \brief Check all the tensors are unique in an expression.
|
|
*/
|
|
void CheckTensorUniqueInExpr(Expr expr);
|
|
|
|
std::vector<std::string> GatherItersToTensorProducer(
|
|
const std::string &target_tensor_name, Expr *expr);
|
|
|
|
bool is_zero(Expr v);
|
|
|
|
Expr NormalizeUpperBound(Expr upper_bound, bool minus_one = true);
|
|
|
|
bool MathEqual(const Expr &a, const Expr &b);
|
|
|
|
//! helper function to get a ir::Select node.
|
|
Expr select(Expr cond, Expr true_value, Expr false_value);
|
|
|
|
//! helper function to get the And of all the conditions.
|
|
Expr and_all(const std::vector<Expr> &conds);
|
|
|
|
//! helper function to get the Or of all the conditions.
|
|
Expr or_any(const std::vector<Expr> &conds);
|
|
|
|
//! Cast the expression \p e to type \type.
|
|
Expr cast(Expr e, Type type);
|
|
|
|
Expr max(Expr a, Expr b);
|
|
|
|
Expr min(Expr a, Expr b);
|
|
|
|
template <typename T>
|
|
Expr make_const(Type t, T v) {
|
|
if (t.is_vector()) {
|
|
if (t.is_int()) {
|
|
return ir::Broadcast::Make(
|
|
make_shared<ir::IntImm>(t.ElementOf(), static_cast<int64_t>(v)),
|
|
t.lanes());
|
|
} else if (t.is_uint()) {
|
|
return ir::Broadcast::Make(
|
|
make_shared<ir::UIntImm>(t.ElementOf(), static_cast<uint64_t>(v)),
|
|
t.lanes());
|
|
} else if (t.is_float()) {
|
|
return ir::Broadcast::Make(
|
|
make_shared<ir::FloatImm>(t.ElementOf(), static_cast<double>(v)),
|
|
t.lanes());
|
|
} else if (t.is_bool()) {
|
|
return ir::Broadcast::Make(
|
|
make_shared<ir::UIntImm>(t.ElementOf(), static_cast<bool>(v)),
|
|
t.lanes());
|
|
} else {
|
|
CINN_NOT_IMPLEMENTED
|
|
}
|
|
} else {
|
|
if (t.is_int()) {
|
|
return make_shared<ir::IntImm>(t, static_cast<int64_t>(v));
|
|
} else if (t.is_uint()) {
|
|
return make_shared<ir::UIntImm>(t, static_cast<uint64_t>(v));
|
|
} else if (t.is_float()) {
|
|
return make_shared<ir::FloatImm>(t, static_cast<double>(v));
|
|
} else if (t.is_bool()) {
|
|
return make_shared<ir::UIntImm>(t, static_cast<bool>(v));
|
|
} else {
|
|
CINN_NOT_IMPLEMENTED
|
|
}
|
|
}
|
|
return Expr();
|
|
}
|
|
|
|
template <typename FuncOp>
|
|
Expr FoldExpr(FuncOp func_op, const std::vector<Expr> &values) {
|
|
Expr init_value;
|
|
for (const Expr &val : values) {
|
|
if (!init_value.defined()) {
|
|
init_value = val;
|
|
} else {
|
|
init_value = func_op(val, init_value);
|
|
}
|
|
}
|
|
return init_value;
|
|
}
|
|
|
|
inline bool IsIterExpr(const Expr &a, const Expr &b) {
|
|
return a.As<ir::IterSplit>() || a.As<ir::IterSum>() ||
|
|
b.As<ir::IterSplit>() || b.As<ir::IterSum>();
|
|
}
|
|
|
|
inline bool IsOne(const Expr &expr) {
|
|
if (expr.is_constant() && expr.get_constant() == 1) {
|
|
return true;
|
|
}
|
|
return false;
|
|
}
|
|
inline bool IsZero(const Expr &expr) {
|
|
if (expr.is_constant() && expr.get_constant() == 0) {
|
|
return true;
|
|
}
|
|
return false;
|
|
}
|
|
|
|
// Promote int32 to int64 type if needed.
|
|
void OpDataTypePromote(ir::Expr *expr);
|
|
void OpDataTypePromote(ir::Module *module);
|
|
void OpDataTypePromote(ir::LoweredFunc *func);
|
|
|
|
// only process ir::Min and ir::Max where the operands 1. contains dynamic shape
|
|
// symbols. 2. the operands are both int types and both are 32/64 bits. Returns
|
|
// the number of bits for unifying operands (by casting). The bool flag
|
|
// indicates whether both sides has different dynamic shape symbols, since if
|
|
// true (like min(S0, S1))), we should not make a ir::Cast but a ir::Call
|
|
// (coercion)
|
|
std::pair<int, bool> UnifiedOperandTypeBits(
|
|
const std::unordered_map<std::string, common::Type> *search_map,
|
|
const ir::Min *op);
|
|
std::pair<int, bool> UnifiedOperandTypeBits(
|
|
const std::unordered_map<std::string, common::Type> *search_map,
|
|
const ir::Max *op);
|
|
} // namespace common
|
|
} // namespace cinn
|