wehub-resource-sync
303 lines
10 KiB
C++
303 lines
10 KiB
C++
// Copyright (c) 2023 PaddlePaddle 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 <memory>
|
|
#include <optional>
|
|
#include "paddle/pir/include/core/builtin_attribute.h"
|
|
#include "paddle/pir/include/core/builtin_op.h"
|
|
#include "paddle/pir/include/core/builtin_type_interfaces.h"
|
|
#include "paddle/pir/include/core/dll_decl.h"
|
|
#include "paddle/pir/include/core/utils.h"
|
|
#include "paddle/pir/include/dialect/shape/ir/shape_op.h"
|
|
#include "paddle/pir/include/dialect/shape/utils/constraints_manager.h"
|
|
#include "paddle/pir/include/dialect/shape/utils/dim_expr_builder.h"
|
|
#include "paddle/pir/include/dialect/shape/utils/shape_or_data_expr.h"
|
|
|
|
namespace pir {
|
|
using InferSymbolicShapeCacheValue = std::vector<symbol::ShapeOrDataDimExprs>;
|
|
|
|
enum TransLayoutType { NCHW2NHWC, NHWC2NCHW, INVALID };
|
|
|
|
/**
|
|
* This class represents information needed to determine the output
|
|
* shape of an operator, which includes the operator's name, input shapes, and
|
|
* attributes.
|
|
*/
|
|
class IR_API InferSymbolicShapeCacheKey {
|
|
public:
|
|
InferSymbolicShapeCacheKey(
|
|
const std::string& op_name,
|
|
const std::vector<symbol::ShapeOrDataDimExprs>& input_shape_or_datas,
|
|
const std::map<std::string, Attribute>& attributes)
|
|
: op_name_(op_name),
|
|
input_shape_or_datas_(input_shape_or_datas),
|
|
attributes_(attributes) {}
|
|
bool operator==(const InferSymbolicShapeCacheKey& other) const;
|
|
std::size_t GetHashValue() const;
|
|
friend std::ostream& operator<<(std::ostream& os,
|
|
const InferSymbolicShapeCacheKey& info);
|
|
friend class InferSymbolicShapeContext;
|
|
|
|
private:
|
|
std::string op_name_;
|
|
std::vector<symbol::ShapeOrDataDimExprs> input_shape_or_datas_;
|
|
std::map<std::string, Attribute> attributes_;
|
|
const std::vector<symbol::ShapeOrDataDimExprs>& GetInputShapeOrDatas() const;
|
|
void SetInputShapeOrDatas(
|
|
const std::vector<symbol::ShapeOrDataDimExprs>& input_shape_or_datas);
|
|
};
|
|
|
|
struct InputDynamicDimSpec {
|
|
std::string dim_name;
|
|
// input_bind = [(input_name, dim_index)]
|
|
std::vector<std::pair<std::string, int>> input_bind;
|
|
symbol::ConstraintsManager::Range range;
|
|
};
|
|
} // namespace pir
|
|
|
|
namespace std {
|
|
template <>
|
|
struct hash<pir::InferSymbolicShapeCacheKey> {
|
|
std::size_t operator()(const pir::InferSymbolicShapeCacheKey& obj) const {
|
|
return obj.GetHashValue();
|
|
}
|
|
};
|
|
} // namespace std
|
|
|
|
namespace pir {
|
|
void InferSymExprForAllValues(ModuleOp module_op);
|
|
|
|
class IR_API InferSymbolicShapeContext {
|
|
public:
|
|
InferSymbolicShapeContext() = default;
|
|
InferSymbolicShapeContext(const InferSymbolicShapeContext&) = delete;
|
|
InferSymbolicShapeContext(InferSymbolicShapeContext&&) = delete;
|
|
void Init(const std::vector<InputDynamicDimSpec>& input_dynamic_dim_spec);
|
|
|
|
// Note: Only initialize the symbol info, the value info is not update.
|
|
void RegisterSymbolConstraintFromContext(
|
|
const InferSymbolicShapeContext& other);
|
|
|
|
const std::string GetNextSymName();
|
|
|
|
bool HasShapeOrDataForValue(Value val) const;
|
|
|
|
const symbol::ShapeOrDataDimExprs& GetShapeOrDataForValue(Value val) const;
|
|
|
|
void SetSymbolForValueByStaticShape(Value val);
|
|
|
|
void SetShapeOrDataForValue(Value val,
|
|
const symbol::ShapeOrDataDimExprs& shape_or_data);
|
|
|
|
void AddEqualCstr(const symbol::DimExpr& lhs, const symbol::DimExpr& rhs);
|
|
|
|
// Add equal constraints for each dim in lhs and rhs.
|
|
void AddEqualCstr(const std::vector<symbol::DimExpr>& lhs,
|
|
const std::vector<symbol::DimExpr>& rhs);
|
|
|
|
bool IsEqual(const symbol::DimExpr& lhs, const symbol::DimExpr& rhs) const;
|
|
|
|
// Returns true if:
|
|
// lhs[i] == rhs[i] for all i
|
|
bool IsEqual(const std::vector<symbol::DimExpr>& lhs,
|
|
const std::vector<symbol::DimExpr>& rhs) const;
|
|
|
|
void AddGreatThanOneCstr(const symbol::DimExpr& dim_expr);
|
|
|
|
bool IsGreatThanOne(const symbol::DimExpr& dim_expr) const;
|
|
|
|
void AddBroadcastableCstr(const symbol::DimExpr& lhs,
|
|
const symbol::DimExpr& rhs);
|
|
|
|
bool IsBroadcastable(const symbol::DimExpr& lhs,
|
|
const symbol::DimExpr& rhs) const;
|
|
|
|
bool HasPredefinedRange(const symbol::DimExpr& dim_expr) const;
|
|
|
|
void PrintShapeOrDatas() const;
|
|
|
|
void SetOpInferSymbolicShapeCache(
|
|
const InferSymbolicShapeCacheKey& op_infer_cache_key,
|
|
InferSymbolicShapeCacheValue result_shape);
|
|
|
|
std::optional<InferSymbolicShapeCacheValue> GetOpInferSymbolicShapeCache(
|
|
const InferSymbolicShapeCacheKey& op_infer_cache_key) const;
|
|
|
|
void ClearOpInferSymbolicShapeCache();
|
|
|
|
const symbol::ConstraintsManager& constraints_manager() const {
|
|
return constraints_manager_;
|
|
}
|
|
|
|
struct DimIndexAndExpr {
|
|
int index;
|
|
symbol::DimExpr dim_expr;
|
|
DimIndexAndExpr(int index_val, const symbol::DimExpr& dim_expr_val)
|
|
: index(index_val), dim_expr(dim_expr_val) {}
|
|
};
|
|
|
|
bool HasPredefinedDimExprForInputName(const std::string& input_name) const;
|
|
|
|
const std::vector<DimIndexAndExpr> GetPredefinedDimExprForInputName(
|
|
const std::string& input_name) const;
|
|
|
|
private:
|
|
symbol::ShapeOrDataDimExprs SimplifyBroadcastForShapeOrData(
|
|
const symbol::ShapeOrDataDimExprs& shape_or_data);
|
|
|
|
void SubstituteDimExpr(const symbol::DimExpr& origin,
|
|
const symbol::DimExpr& substituted);
|
|
|
|
int64_t sym_idx_begin_ = 0;
|
|
int64_t next_sym_idx_ = 0;
|
|
|
|
std::unordered_map<uint64_t, symbol::ShapeOrDataDimExprs>
|
|
value_id_to_shape_or_data_;
|
|
|
|
symbol::ConstraintsManager constraints_manager_;
|
|
|
|
using DimExprSubstitutionPattern =
|
|
std::unordered_map<symbol::DimExpr, symbol::DimExpr>;
|
|
DimExprSubstitutionPattern substitution_pattern_;
|
|
|
|
std::unordered_map<InferSymbolicShapeCacheKey, InferSymbolicShapeCacheValue>
|
|
infer_symbolic_shape_cache_;
|
|
|
|
std::unordered_map<std::string, std::vector<DimIndexAndExpr>>
|
|
predefined_dimexpr_map_for_inputs_;
|
|
|
|
std::unordered_map<std::string, symbol::DimExpr>
|
|
input_dynamic_dim_name_spec_to_dimexpr_map_;
|
|
};
|
|
|
|
class IR_API ShapeConstraintIRAnalysis final
|
|
: public std::enable_shared_from_this<ShapeConstraintIRAnalysis> {
|
|
public:
|
|
ShapeConstraintIRAnalysis() = default;
|
|
ShapeConstraintIRAnalysis(const ShapeConstraintIRAnalysis&) = delete;
|
|
ShapeConstraintIRAnalysis(ShapeConstraintIRAnalysis&&) = delete;
|
|
void InitInferContext();
|
|
|
|
void RegisterSymbolConstraintFromShapeAnalysis(
|
|
const ShapeConstraintIRAnalysis& other);
|
|
|
|
const std::string GetNextSymName();
|
|
|
|
const symbol::ShapeOrDataDimExprs& GetShapeOrDataForValue(Value val);
|
|
|
|
void SetShapeOrDataForValue(Value val,
|
|
const symbol::ShapeOrDataDimExprs& shape_or_data);
|
|
|
|
// Set ShapeOrData of `to` value by ShapeOrData of `from` value.
|
|
void ShareShapeOrData(Value from, Value to);
|
|
|
|
// Update Symbol Shape for value by layout transformation.
|
|
void UpdateShapeOrDataByTransLayout(Value val,
|
|
TransLayoutType trans_layout_type);
|
|
|
|
void AddEqualCstr(const symbol::DimExpr& lhs, const symbol::DimExpr& rhs);
|
|
|
|
bool IsEqual(const symbol::DimExpr& lhs, const symbol::DimExpr& rhs) const;
|
|
|
|
bool IsGreatThanOne(const symbol::DimExpr& dim_expr) const;
|
|
|
|
bool IsBroadcastable(const symbol::DimExpr& lhs,
|
|
const symbol::DimExpr& rhs) const;
|
|
|
|
// Used to debug
|
|
void PrintShapeOrDatas() const;
|
|
|
|
// Returns true if the two value have the same symbolic shape.
|
|
bool IsShapeEqual(Value lhs, Value rhs);
|
|
|
|
// Suppose:
|
|
// lhs_dim_idxs = {ld0, ld1, ...}
|
|
// rhs_dim_idxs = {rd0, rd1, ...}
|
|
// Returns true if:
|
|
// lhs.shape[ld0] * lhs.shape[ld1] * ... ==
|
|
// rhs.shape[rd0] * rhs.shape[rd1] * ...
|
|
bool IsProductEqual(Value lhs,
|
|
const std::vector<int>& lhs_dim_idxs,
|
|
Value rhs,
|
|
const std::vector<int>& rhs_dim_idxs);
|
|
|
|
// Returns true if:
|
|
// lhs.shape[lhs_from] * ... lhs.shape[lhs_to-1] ==
|
|
// rhs.shape[rhs_from] * ... rhs.shape[rhs_to-1]
|
|
bool IsProductEqual(
|
|
Value lhs, int lhs_from, int lhs_to, Value rhs, int rhs_from, int rhs_to);
|
|
|
|
// Returns true if the two value have the same number elements.
|
|
bool IsSameNumel(Value lhs, Value rhs);
|
|
|
|
pir::PrintHooks PrintHook();
|
|
|
|
symbol::DimExpr GetProductDimExpr(Value lhs,
|
|
const std::vector<int>& lhs_dim_idxs);
|
|
|
|
const symbol::ConstraintsManager& constraints_manager() const {
|
|
return context_.constraints_manager();
|
|
}
|
|
|
|
void ClearOpInferSymbolicShapeCache() {
|
|
context_.ClearOpInferSymbolicShapeCache();
|
|
}
|
|
|
|
void SetInputDynamicDimSpec(
|
|
const std::vector<InputDynamicDimSpec>& input_dynamic_dim_spec);
|
|
|
|
void AppendInputDynamicDimSpec(
|
|
const std::vector<InputDynamicDimSpec>& input_dynamic_dim_spec);
|
|
|
|
private:
|
|
InferSymbolicShapeContext* MutInferSymbolicShapeContext() {
|
|
return &context_;
|
|
}
|
|
|
|
friend void InferSymExprForAllValues(ModuleOp module_op);
|
|
|
|
void SetSymbolForValueByStaticShape(Value val);
|
|
|
|
void InferShapeOrDataForValue(Value val);
|
|
|
|
private:
|
|
InferSymbolicShapeContext context_;
|
|
std::vector<InputDynamicDimSpec> input_dynamic_dim_spec_;
|
|
};
|
|
|
|
class IR_API ShapeAnalysisManager {
|
|
public:
|
|
static ShapeAnalysisManager& Instance();
|
|
ShapeConstraintIRAnalysis& Get(const pir::Program* program);
|
|
|
|
ShapeAnalysisManager(const ShapeAnalysisManager&) = delete;
|
|
ShapeAnalysisManager(ShapeAnalysisManager&&) = delete;
|
|
ShapeAnalysisManager& operator=(const ShapeAnalysisManager&) = delete;
|
|
|
|
private:
|
|
ShapeAnalysisManager() {}
|
|
std::unordered_map<uint64_t, std::shared_ptr<ShapeConstraintIRAnalysis>>
|
|
tables_;
|
|
};
|
|
|
|
#define OP_DECLARE_INFER_SYMBOLIC_SHAPE(name) \
|
|
bool name##OpInferSymbolicShape( \
|
|
pir::Operation* op, pir::InferSymbolicShapeContext* infer_context);
|
|
|
|
bool IsStaticShape(const Value& value);
|
|
|
|
} // namespace pir
|