Skip to content

[CINN]fix CompileBroadcastTreeToConditionBlock #64659

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zyfncg merged 2 commits into from
May 29, 2024
Merged

[CINN]fix CompileBroadcastTreeToConditionBlock #64659

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
716afd3
[CINN]fix CompileBroadcastTreeToConditionBlock by using local ShapeOr…
Hongqing-work May 28, 2024
77addfa
fix lambda return type
Hongqing-work May 28, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 11 additions & 11 deletions paddle/cinn/hlir/dialect/operator/transforms/lowering_pass/broadcast_with_cf.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -202,14 +202,11 @@ void ReplaceExpandWithBroadcast(pir::IrContext* ir_context,

std::tuple<pir::Value, pir::Value, pir::Value> BroadcastableToCondValue(
const symbol::Broadcastable<symbol::DimExpr>& broadcastable_condition,
pir::ShapeConstraintIRAnalysis& shape_analysis, // NOLINT
const ShapeOrDataDimExprs4ValueT& ShapeOrDataDimExprs4Value,
const std::vector<pir::Value>& group_inputs,
pir::Builder& builder) { // NOLINT
const auto& lhs_expr = broadcastable_condition->lhs;
const auto& rhs_expr = broadcastable_condition->rhs;
auto ShapeOrDataDimExprs4Value = [&shape_analysis](pir::Value value) {
return shape_analysis.GetShapeOrDataForValue(value);
};

std::vector<pir::Value> lhs_minimal_inputs;
std::vector<pir::Attribute> lhs_output_dim_expr_attrs;
Expand Down Expand Up @@ -322,7 +319,7 @@ void InsertYieldOpForCondBlock(pir::Operation* cond_op,
pir::Operation* CreateConditionBlock(
const cinn::common::BroadcastTree& broadcast_tree,
const OpLoweringGroupPtr& origin_group,
pir::ShapeConstraintIRAnalysis& shape_analysis, // NOLINT
const ShapeOrDataDimExprs4ValueT& ShapeOrDataDimExprs4Value,
const std::unordered_map<pir::Value, size_t>& value_to_dim_expr_idx,
const std::vector<pir::Value>& group_inputs,
const std::vector<pir::Type>& output_types,
Expand All @@ -345,7 +342,7 @@ pir::Operation* CreateConditionBlock(
.Get<cinn::common::BroadcastBranch<cinn::common::BroadcastTree>>();
const auto& [lhs_eq_rhs_cond, lhs_eq_one_cond, rhs_eq_one_cond] =
BroadcastableToCondValue(
branch.Get<0>(), shape_analysis, group_inputs, builder);
branch.Get<0>(), ShapeOrDataDimExprs4Value, group_inputs, builder);

// lhs == rhs
auto lhs_eq_rhs_cond_op = builder.Build<paddle::dialect::IfOp>(
Expand All @@ -354,7 +351,7 @@ pir::Operation* CreateConditionBlock(
builder.SetInsertionPointToBlockEnd(&lhs_eq_rhs_block);
auto* lhs_eq_rhs_block_op = CreateConditionBlock(branch.Get<1>(),
origin_group,
shape_analysis,
ShapeOrDataDimExprs4Value,
value_to_dim_expr_idx,
group_inputs,
output_types,
Expand All @@ -373,7 +370,7 @@ pir::Operation* CreateConditionBlock(
builder.SetInsertionPointToBlockEnd(&lhs_eq_one_block);
auto* lhs_eq_one_block_op = CreateConditionBlock(branch.Get<2>(),
origin_group,
shape_analysis,
ShapeOrDataDimExprs4Value,
value_to_dim_expr_idx,
group_inputs,
output_types,
Expand All @@ -387,7 +384,7 @@ pir::Operation* CreateConditionBlock(
builder.SetInsertionPointToBlockEnd(&rhs_eq_one_block);
auto* rhs_eq_one_block_op = CreateConditionBlock(branch.Get<3>(),
origin_group,
shape_analysis,
ShapeOrDataDimExprs4Value,
value_to_dim_expr_idx,
group_inputs,
output_types,
Expand Down Expand Up @@ -487,17 +484,20 @@ bool NeedBroadcastWithCF(const cinn::common::BroadcastLeaf& leaves) {
pir::Operation* CompileBroadcastTreeToConditionBlock(
const OpLoweringGroupPtr& group,
const BroadcastTree& broadcast_tree,
pir::ShapeConstraintIRAnalysis& shape_analysis, // NOLINT
const std::unordered_map<pir::Value, size_t>& value_to_dim_expr_idx,
const std::vector<pir::Value>& group_inputs,
const std::vector<pir::Type>& output_types,
pir::PatternRewriter& rewriter) { // NOLINT
auto ShapeOrDataDimExprs4Value =
[&group](pir::Value value) -> const symbol::ShapeOrDataDimExprs& {
return group->GetShapeOrDataExprs(value);
};
// 1. broadcast tree to condition op
VLOG(4) << "broadcast tree to condition op";
std::unordered_map<pir::Block*, OpLoweringGroupPtr> group_map;
pir::Operation* cond_op = CreateConditionBlock(broadcast_tree,
group,
shape_analysis,
ShapeOrDataDimExprs4Value,
value_to_dim_expr_idx,
group_inputs,
output_types,
Expand Down
1 change: 0 additions & 1 deletion paddle/cinn/hlir/dialect/operator/transforms/lowering_pass/broadcast_with_cf.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,6 @@ GroupDimExprInfo GetGroupDimExprInfo(const OpLoweringGroupPtr& group);
pir::Operation* CompileBroadcastTreeToConditionBlock(
const OpLoweringGroupPtr& group,
const BroadcastTree& broadcast_tree,
pir::ShapeConstraintIRAnalysis& shape_analysis, // NOLINT
const std::unordered_map<pir::Value, size_t>& value_to_dim_expr_idx,
const std::vector<pir::Value>& group_inputs,
const std::vector<pir::Type>& output_types,
Expand Down
17 changes: 6 additions & 11 deletions paddle/cinn/hlir/dialect/operator/transforms/lowering_pass/lower_cinn_fusion_op_pass.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,10 +30,8 @@

namespace cinn::dialect::ir::details {

pir::Operation* ProcessDyShapeGroup(
const OpLoweringGroupPtr& group,
pir::ShapeConstraintIRAnalysis& shape_analysis, // NOLINT
pir::PatternRewriter& rewriter) { // NOLINT
pir::Operation* ProcessDyShapeGroup(const OpLoweringGroupPtr& group,
pir::PatternRewriter& rewriter) { // NOLINT
// NOTE(dev): Need UpdateShapeOrDataExprs firstly and the logic
// will be migated into BucketLower later.
UpdateGroupShapeOrDataExprs(const_cast<OpLoweringGroupPtr&>(group));
Expand All @@ -53,7 +51,6 @@ pir::Operation* ProcessDyShapeGroup(
}
return CompileBroadcastTreeToConditionBlock(group,
*broadcast_tree,
shape_analysis,
value_to_dim_expr_idx,
group_inputs,
output_types,
Expand Down Expand Up @@ -85,7 +82,7 @@ class FusionOpPattern : public pir::OpRewritePattern<cinn::dialect::FusionOp> {

// TODO(zhangyuqin1998): Replace pir::Group with a new structure
OpLoweringGroupPtr group = GetGroup(fusion_op);
pir::Operation* compiled_op = ProcessGroup(group, shape_analysis, rewriter);
pir::Operation* compiled_op = ProcessGroup(group, rewriter);

for (size_t i = 0; i < fusion_op.num_results(); ++i) {
rewriter.ReplaceAllUsesWith(fusion_op.result(i), compiled_op->result(i));
Expand All @@ -104,8 +101,7 @@ class FusionOpPattern : public pir::OpRewritePattern<cinn::dialect::FusionOp> {

virtual pir::Operation* ProcessGroup(
const OpLoweringGroupPtr& group,
pir::ShapeConstraintIRAnalysis& shape_analysis, // NOLINT
pir::PatternRewriter& rewriter) const { // NOLINT
pir::PatternRewriter& rewriter) const { // NOLINT
auto group_inputs = GetBlockOutsideInput(group->ops());
// compile group to jit_kernel_op
std::vector<pir::Type> output_types;
Expand Down Expand Up @@ -156,9 +152,8 @@ class DyShapeFusionOpPattern : public FusionOpPattern {
protected:
virtual pir::Operation* ProcessGroup(
const OpLoweringGroupPtr& group,
pir::ShapeConstraintIRAnalysis& shape_analysis, // NOLINT
pir::PatternRewriter& rewriter) const { // NOLINT
return ProcessDyShapeGroup(group, shape_analysis, rewriter);
pir::PatternRewriter& rewriter) const { // NOLINT
return ProcessDyShapeGroup(group, rewriter);
}
};

Expand Down