Skip to content

Commit c448d28

Browse files
zhangboSJTUzhangboSJTU
authored
[PIR][DynamicShape] Add nullary_infer_sym and binary nullary_infer_sym (#62383)
* add nullary_infer_sym * add infer
1 parent 2c34d76 commit c448d28

17 files changed

+692
-73
lines changed

‎paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/cinn_op_infer_sym.cc‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,7 @@ bool ConcatOpInferSymbolicShape(
7676
out_dims[axis] = out_dims[axis] + operand_shape_or_data.shape()[axis];
7777
}
7878

79-
for (size_t i = 1; i < rank; ++i) {
79+
for (size_t i = 0; i < rank; ++i) {
8080
if (i == static_cast<size_t>(axis)) continue;
8181
paddle::dialect::details::BuildCstrEqForTensorListAlongAxis(
8282
shape_analysis, input_values, i);
@@ -85,6 +85,9 @@ bool ConcatOpInferSymbolicShape(
8585
return out_dims;
8686
};
8787

88+
VLOG(3) << "constraints size:"
89+
<< shape_analysis->CreateDimExprBuilder().constraints().size();
90+
8891
symbol::ShapeOrDataDimExprs shape_data{
8992
symbol::TensorShapeOrDataDimExprs(GetOutDimExprs())};
9093

‎paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/cinn_op_infer_sym.h‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@
1313
// limitations under the License.
1414

1515
#pragma once
16-
#include "paddle/fluid/pir/dialect/operator/utils/utils.h"
1716
#include "paddle/pir/include/dialect/shape/utils/shape_analysis.h"
1817

1918
namespace cinn::dialect {

‎paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_element_wise_binary.h‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@
1414

1515
#pragma once
1616

17-
#include "paddle/fluid/pir/dialect/operator/utils/utils.h"
1817
#include "paddle/pir/include/dialect/shape/utils/shape_analysis.h"
1918

2019
namespace paddle::dialect {

‎paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_symbolic_shape.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
#include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/cinn_op_infer_sym.h"
1818
#include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_element_wise_binary.h"
19+
#include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/nullary_infer_sym.h"
1920
#include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/paddle_op_infer_sym.h"
2021
#include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_and_result.h"
2122
#include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.h"

‎paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/nullary_infer_sym.cc‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
1+
// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
#include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/nullary_infer_sym.h"
16+
#include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_utils.h"
17+
18+
namespace paddle::dialect {
19+
20+
bool EmptyOpInferSymbolicShape(pir::Operation *op,
21+
pir::ShapeConstraintIRAnalysis *shape_analysis) {
22+
const auto &shape_gen_op = op->operand_source(0).defining_op();
23+
if (shape_gen_op->isa<paddle::dialect::FullIntArrayOp>()) {
24+
std::vector<int64_t> shape = details::GetVectorAttr(
25+
shape_gen_op->dyn_cast<paddle::dialect::FullIntArrayOp>(), "value");
26+
std::vector<symbol::DimExpr> sym_dims;
27+
sym_dims.reserve(shape.size());
28+
for (const int64_t &dim : shape) {
29+
sym_dims.emplace_back(symbol::DimExpr(dim));
30+
}
31+
32+
symbol::ShapeOrDataDimExprs shape_data{
33+
symbol::TensorShapeOrDataDimExprs(sym_dims)};
34+
shape_analysis->SetShapeOrDataForValue(op->result(0), shape_data);
35+
return true;
36+
37+
} else {
38+
pir::Value operand_source = op->operand_source(0);
39+
const symbol::ShapeOrDataDimExprs &operand_shape_or_data =
40+
shape_analysis->GetShapeOrDataForValue(operand_source);
41+
42+
shape_analysis->SetShapeOrDataForValue(op->result(0),
43+
operand_shape_or_data);
44+
return true;
45+
}
46+
}
47+
48+
bool GaussianOpInferSymbolicShape(
49+
pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
50+
const auto &shape_gen_op = op->operand_source(0).defining_op();
51+
52+
if (shape_gen_op->isa<paddle::dialect::FullIntArrayOp>()) {
53+
std::vector<int64_t> shape = details::GetVectorAttr(
54+
shape_gen_op->dyn_cast<paddle::dialect::FullIntArrayOp>(), "value");
55+
std::vector<symbol::DimExpr> sym_dims;
56+
sym_dims.reserve(shape.size());
57+
for (const int64_t &dim : shape) {
58+
sym_dims.emplace_back(symbol::DimExpr(dim));
59+
}
60+
61+
symbol::ShapeOrDataDimExprs shape_data{
62+
symbol::TensorShapeOrDataDimExprs(sym_dims)};
63+
shape_analysis->SetShapeOrDataForValue(op->result(0), shape_data);
64+
return true;
65+
66+
} else {
67+
PADDLE_THROW(phi::errors::Unimplemented(
68+
op->name() +
69+
" 's InferSymbolicShape interface is NOT implemented now."));
70+
return true;
71+
}
72+
}
73+
74+
} // namespace paddle::dialect

‎paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/nullary_infer_sym.h‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
#pragma once
16+
17+
#include "paddle/pir/include/dialect/shape/utils/shape_analysis.h"
18+
19+
namespace paddle::dialect {
20+
OP_DECLARE_INFER_SYMBOLIC_SHAPE(Empty)
21+
OP_DECLARE_INFER_SYMBOLIC_SHAPE(Gaussian)
22+
} // namespace paddle::dialect

‎paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/paddle_op_infer_sym.cc‎

Lines changed: 24 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,6 @@ bool StackOpInferSymbolicShape(pir::Operation *op,
9797
static_cast<std::int64_t>(shape_data_list.size()));
9898
} else {
9999
for (int i = 0; i < rank; ++i) {
100-
if (i == axis) continue;
101100
details::BuildCstrEqForTensorListAlongAxis(
102101
shape_analysis, shape_data_list, i);
103102
}
@@ -931,26 +930,6 @@ bool SplitOpInferSymbolicShape(pir::Operation *op,
931930
}
932931

933932
// Not Implemented Ops.
934-
935-
bool DiagEmbedOpInferSymbolicShape(
936-
pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
937-
PADDLE_THROW(phi::errors::Unimplemented(
938-
op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
939-
return true;
940-
}
941-
bool DiagonalOpInferSymbolicShape(
942-
pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
943-
PADDLE_THROW(phi::errors::Unimplemented(
944-
op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
945-
return true;
946-
}
947-
bool DirichletOpInferSymbolicShape(
948-
pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
949-
PADDLE_THROW(phi::errors::Unimplemented(
950-
op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
951-
return true;
952-
}
953-
954933
bool GatherOpInferSymbolicShape(
955934
pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
956935
const auto &input_shape_or_data =
@@ -1020,17 +999,33 @@ bool GatherOpInferSymbolicShape(
1020999

10211000
bool KronOpInferSymbolicShape(pir::Operation *op,
10221001
pir::ShapeConstraintIRAnalysis *shape_analysis) {
1023-
PADDLE_THROW(phi::errors::Unimplemented(
1024-
op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
1025-
return true;
1026-
}
1027-
bool KthvalueOpInferSymbolicShape(
1028-
pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
1029-
PADDLE_THROW(phi::errors::Unimplemented(
1030-
op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
1002+
const auto &x_shape_or_data =
1003+
shape_analysis->GetShapeOrDataForValue(op->operand_source(0)).shape();
1004+
const auto &y_shape_or_data =
1005+
shape_analysis->GetShapeOrDataForValue(op->operand_source(1)).shape();
1006+
const int rank_x = x_shape_or_data.size();
1007+
const int rank_y = y_shape_or_data.size();
1008+
const int rank = (rank_x > rank_y) ? rank_x : rank_y;
1009+
1010+
std::vector<symbol::DimExpr> dim_out;
1011+
dim_out.reserve(rank);
1012+
const auto one = symbol::DimExpr{1};
1013+
const auto minus_one = symbol::DimExpr{-1};
1014+
for (int i = 0; i < rank; i++) {
1015+
symbol::DimExpr dim_xi =
1016+
(i < rank - rank_x) ? one : x_shape_or_data.at(i - (rank - rank_x));
1017+
symbol::DimExpr dim_yi =
1018+
(i < rank - rank_y) ? one : y_shape_or_data.at(i - (rank - rank_y));
1019+
dim_out.push_back(dim_xi * dim_yi);
1020+
}
1021+
symbol::ShapeOrDataDimExprs shape_data{
1022+
symbol::TensorShapeOrDataDimExprs(dim_out)};
1023+
pir::Value res = op->result(0);
1024+
shape_analysis->SetShapeOrDataForValue(res, shape_data);
10311025
return true;
10321026
}
10331027

1028+
// Not Impelmented Ops.
10341029
bool LogcumsumexpOpInferSymbolicShape(
10351030
pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
10361031
PADDLE_THROW(phi::errors::Unimplemented(
@@ -1095,32 +1090,6 @@ bool UniqueConsecutiveOpInferSymbolicShape(
10951090
op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
10961091
return true;
10971092
}
1098-
1099-
bool EinsumOpInferSymbolicShape(
1100-
pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
1101-
PADDLE_THROW(phi::errors::Unimplemented(
1102-
op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
1103-
return true;
1104-
}
1105-
bool EmptyOpInferSymbolicShape(pir::Operation *op,
1106-
pir::ShapeConstraintIRAnalysis *shape_analysis) {
1107-
PADDLE_THROW(phi::errors::Unimplemented(
1108-
op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
1109-
return true;
1110-
}
1111-
bool Exponential_OpInferSymbolicShape(
1112-
pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
1113-
PADDLE_THROW(phi::errors::Unimplemented(
1114-
op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
1115-
return true;
1116-
}
1117-
bool GaussianOpInferSymbolicShape(
1118-
pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
1119-
PADDLE_THROW(phi::errors::Unimplemented(
1120-
op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
1121-
return true;
1122-
}
1123-
11241093
bool LinspaceOpInferSymbolicShape(
11251094
pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
11261095
PADDLE_THROW(phi::errors::Unimplemented(

‎paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/paddle_op_infer_sym.h‎

Lines changed: 0 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@
1414

1515
#pragma once
1616

17-
#include "paddle/fluid/pir/dialect/operator/utils/utils.h"
1817
#include "paddle/pir/include/dialect/shape/utils/shape_analysis.h"
1918

2019
namespace paddle::dialect {
@@ -51,12 +50,8 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Split)
5150

5251
// Not Impelmented Ops.
5352

54-
OP_DECLARE_INFER_SYMBOLIC_SHAPE(DiagEmbed)
55-
OP_DECLARE_INFER_SYMBOLIC_SHAPE(Diagonal)
56-
OP_DECLARE_INFER_SYMBOLIC_SHAPE(Dirichlet)
5753
OP_DECLARE_INFER_SYMBOLIC_SHAPE(Gather)
5854
OP_DECLARE_INFER_SYMBOLIC_SHAPE(Kron)
59-
OP_DECLARE_INFER_SYMBOLIC_SHAPE(Kthvalue)
6055
OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logcumsumexp)
6156
OP_DECLARE_INFER_SYMBOLIC_SHAPE(MaskedSelect)
6257
OP_DECLARE_INFER_SYMBOLIC_SHAPE(Poisson)
@@ -67,10 +62,6 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(TakeAlongAxis)
6762
OP_DECLARE_INFER_SYMBOLIC_SHAPE(Topk)
6863
OP_DECLARE_INFER_SYMBOLIC_SHAPE(Unbind)
6964
OP_DECLARE_INFER_SYMBOLIC_SHAPE(UniqueConsecutive)
70-
OP_DECLARE_INFER_SYMBOLIC_SHAPE(Einsum)
71-
OP_DECLARE_INFER_SYMBOLIC_SHAPE(Empty)
72-
OP_DECLARE_INFER_SYMBOLIC_SHAPE(Exponential_)
73-
OP_DECLARE_INFER_SYMBOLIC_SHAPE(Gaussian)
7465
OP_DECLARE_INFER_SYMBOLIC_SHAPE(Linspace)
7566
OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logspace)
7667
OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logsumexp)

‎paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_and_result.cc‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -154,6 +154,10 @@ bool Digamma_OpInferSymbolicShape(
154154
pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
155155
return SameOperandsAndResultShape(op, shape_analysis);
156156
}
157+
bool DirichletOpInferSymbolicShape(
158+
pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
159+
return SameOperandsAndResultShape(op, shape_analysis);
160+
}
157161
bool EqualOpInferSymbolicShape(pir::Operation *op,
158162
pir::ShapeConstraintIRAnalysis *shape_analysis) {
159163
return SameOperandsAndResultShape(op, shape_analysis);
@@ -194,6 +198,10 @@ bool Expm1_OpInferSymbolicShape(
194198
pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
195199
return SameOperandsAndResultShape(op, shape_analysis);
196200
}
201+
bool Exponential_OpInferSymbolicShape(
202+
pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
203+
return SameOperandsAndResultShape(op, shape_analysis);
204+
}
197205
bool FetchOpInferSymbolicShape(pir::Operation *op,
198206
pir::ShapeConstraintIRAnalysis *shape_analysis) {
199207
return SameOperandsAndResultShape(op, shape_analysis);

‎paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_and_result.h‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@
1414

1515
#pragma once
1616

17-
#include "paddle/fluid/pir/dialect/operator/utils/utils.h"
1817
#include "paddle/pir/include/dialect/shape/utils/shape_analysis.h"
1918

2019
namespace paddle::dialect {
@@ -50,6 +49,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cosh)
5049
OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cosh_)
5150
OP_DECLARE_INFER_SYMBOLIC_SHAPE(Digamma)
5251
OP_DECLARE_INFER_SYMBOLIC_SHAPE(Digamma_)
52+
OP_DECLARE_INFER_SYMBOLIC_SHAPE(Dirichlet)
5353
OP_DECLARE_INFER_SYMBOLIC_SHAPE(Equal)
5454
OP_DECLARE_INFER_SYMBOLIC_SHAPE(Equal_)
5555
OP_DECLARE_INFER_SYMBOLIC_SHAPE(Erf)
@@ -60,6 +60,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Exp)
6060
OP_DECLARE_INFER_SYMBOLIC_SHAPE(Exp_)
6161
OP_DECLARE_INFER_SYMBOLIC_SHAPE(Expm1)
6262
OP_DECLARE_INFER_SYMBOLIC_SHAPE(Expm1_)
63+
OP_DECLARE_INFER_SYMBOLIC_SHAPE(Exponential_)
6364
OP_DECLARE_INFER_SYMBOLIC_SHAPE(Fetch)
6465
OP_DECLARE_INFER_SYMBOLIC_SHAPE(Flip)
6566
OP_DECLARE_INFER_SYMBOLIC_SHAPE(Floor)

0 commit comments

Comments
 (0)