wehub-resource-sync
127 lines
4.8 KiB
C++
127 lines
4.8 KiB
C++
// Copyright (c) 2025 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 <unordered_set>
|
|
#include <vector>
|
|
#include "paddle/cinn/common/dim_expr_converter.h"
|
|
#include "paddle/cinn/ir/ir.h"
|
|
#include "paddle/cinn/ir/ir_base.h"
|
|
#include "paddle/cinn/ir/utils/ir_compare.h"
|
|
#include "paddle/common/union_find_set.h"
|
|
#include "paddle/pir/include/dialect/shape/utils/dim_expr.h"
|
|
#include "paddle/pir/include/dialect/shape/utils/shape_analysis.h"
|
|
namespace cinn {
|
|
namespace common {
|
|
|
|
// In order to avoid an infinite loop when checking the IndexExpr in the
|
|
// UnionFindSet, We specify that IndexExpr directly compares IRNode pointers
|
|
// in the UnionFindSet
|
|
struct IndexExprDirectCompare {
|
|
bool operator()(const ir::IndexExpr& a, const ir::IndexExpr& b) const {
|
|
return a.get() == b.get();
|
|
}
|
|
};
|
|
|
|
/*
|
|
* ShapeConstraintManager is used to access and manage the shape constraints on
|
|
* the program generated by the frontend in the CINN backend.
|
|
* Note: The ShapeConstraintManager is a Thread Level Singleton class.
|
|
*/
|
|
class ShapeConstraintManager {
|
|
public:
|
|
// Returns a singleton object.
|
|
static ShapeConstraintManager& Instance();
|
|
// Initialize the ShapeConstraintManager with CINN frontend
|
|
// ConstraintsManager.
|
|
void Init(const symbol::ConstraintsManager& constraint);
|
|
// Initialize the ShapeConstraintManager with the DimExpr UnionFindSet.
|
|
void InitEqualExprs(
|
|
const ::common::UnionFindSet<symbol::DimExpr>& equal_dim_exprs);
|
|
// Initialize the ShapeConstraintManager with the IndexExpr UnionFindSet.
|
|
void InitEqualExprs(
|
|
const ::common::UnionFindSet<ir::IndexExpr, IndexExprDirectCompare>&
|
|
equal_exprs);
|
|
|
|
// Broadcast DimExpr: All parameters need to build edges between each other.
|
|
// Other DimExpr: Only build edges between them.
|
|
// e.g.
|
|
// 1. Broadcastable[ Broadcast(S0, S1), Broadcast(S2, S3) ]
|
|
// ===>
|
|
// {S0: {S1, S2, S3},
|
|
// S1: {S0, S2, S3},
|
|
// S2: {S0, S1, S3},
|
|
// S3: {S0, S1, S2}}
|
|
// 2. Broadcastable[ S0, S1 ]
|
|
// ===>
|
|
// {S0: {S1},
|
|
// S1: {S0}}
|
|
// 3. Broadcastable[ S0, Broadcast(S1, S2) ]
|
|
// ===>
|
|
// {S0: {S1, S2},
|
|
// S1: {S0, S2},
|
|
// S2: {S0, S1}}
|
|
void InitBroadcastableExprs(
|
|
const std::unordered_set<symbol::Broadcastable<symbol::DimExpr>>&
|
|
broadcastable_dim_exprs);
|
|
// Returns whether lhs and rhs are equal in the DimExpr UnionFindSet.
|
|
bool IsEqual(const ir::IndexExpr& lhs, const ir::IndexExpr& rhs);
|
|
bool IsBroadcastable(const ir::IndexExpr& lhs, const ir::IndexExpr& rhs);
|
|
bool IsBroadcastable(const std::vector<ir::IndexExpr>& vec);
|
|
|
|
friend std::ostream& operator<<(
|
|
std::ostream& os, const ShapeConstraintManager& constraints_manager);
|
|
|
|
private:
|
|
::common::UnionFindSet<symbol::DimExpr> equal_dim_exprs_;
|
|
::common::UnionFindSet<ir::IndexExpr, IndexExprDirectCompare> equal_exprs_;
|
|
|
|
// Since broadcast relationships are not conductive, `Bidirectional Adjacency
|
|
// Map` are used to store broadcast relationships. The edge between nodes
|
|
// indicates that the two nodes are broadcastable.
|
|
class BroadcastMap {
|
|
public:
|
|
void AddEdge(const ir::IndexExpr& lhs, const ir::IndexExpr& rhs) {
|
|
broadcastable_exprs_[lhs].insert(rhs);
|
|
broadcastable_exprs_[rhs].insert(lhs);
|
|
}
|
|
void AddEdge(const symbol::DimExpr& lhs, const symbol::DimExpr& rhs) {
|
|
DimExprConverter cvt;
|
|
auto lhs_ = cvt.ConvertToIrExpr(lhs);
|
|
auto rhs_ = cvt.ConvertToIrExpr(rhs);
|
|
broadcastable_exprs_[lhs_].insert(rhs_);
|
|
broadcastable_exprs_[rhs_].insert(lhs_);
|
|
}
|
|
bool HasEdge(const ir::IndexExpr& lhs, const ir::IndexExpr& rhs) {
|
|
return broadcastable_exprs_.count(lhs) &&
|
|
broadcastable_exprs_.at(lhs).count(rhs);
|
|
}
|
|
const std::unordered_map<ir::IndexExpr, std::unordered_set<ir::IndexExpr>>&
|
|
GetBroadcastableExprs() const {
|
|
return broadcastable_exprs_;
|
|
}
|
|
|
|
private:
|
|
std::unordered_map<ir::IndexExpr, std::unordered_set<ir::IndexExpr>>
|
|
broadcastable_exprs_;
|
|
};
|
|
BroadcastMap broadcastable_exprs_;
|
|
ShapeConstraintManager() = default;
|
|
~ShapeConstraintManager() = default;
|
|
ShapeConstraintManager(const ShapeConstraintManager&) = delete;
|
|
ShapeConstraintManager(ShapeConstraintManager&&) = delete;
|
|
ShapeConstraintManager& operator=(const ShapeConstraintManager&) = delete;
|
|
};
|
|
} // namespace common
|
|
} // namespace cinn
|