-
Notifications
You must be signed in to change notification settings - Fork 5.5k
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.
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
【Infer Symbolic Shape No.120】[BUAA] affine_grid #67364
Changes from 4 commits
65f2395
9c66a79
e5d53c5
9083ae3
81cb0d6
0601169
679b590
7676cd7
9195723
7fd0275
5ff3f5c
9429d61
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -85,13 +85,6 @@ std::vector<symbol::DimExpr> GetRealPadding( | |
return real_padding; | ||
} | ||
|
||
// bool AffineGridOpInferSymbolicShape(pir::Operation *op, | ||
// pir::InferSymbolicShapeContext | ||
// *infer_context) { | ||
// // pass | ||
// return true; | ||
// } | ||
|
||
symbol::ShapeOrDataDimExprs Pool2dRawInferSymbolicShape( | ||
pir::Operation *op, | ||
const std::vector<symbol::DimExpr> &kernel_size, | ||
|
@@ -228,6 +221,42 @@ symbol::ShapeOrDataDimExprs Pool2dRawInferSymbolicShape( | |
namespace paddle::dialect { | ||
using paddle::dialect::details::CreateShapeOrDataForXShape; | ||
|
||
bool AffineGridOpInferSymbolicShape( | ||
pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) { | ||
const auto &input_shape_or_data = | ||
infer_context->GetShapeOrDataForValue(op->operand_source(0)); | ||
std::vector<symbol::DimExpr> input_dims = input_shape_or_data.shape(); | ||
|
||
const auto &output_shape_or_data = | ||
infer_context->GetShapeOrDataForValue(op->operand_source(1)); | ||
|
||
const int64_t output_shape_size = | ||
static_cast<int64_t>(output_shape_or_data.shape().size()); | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 非必要cast,为啥不直接用int接收 There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. done |
||
|
||
std::vector<symbol::DimExpr> output_dims; | ||
output_dims.push_back(input_dims[0]); // N | ||
|
||
if (output_shape_size == 4) { | ||
// N * H * W * 2 | ||
output_dims.push_back(infer_context->GetNextSymName()); // H | ||
output_dims.push_back(infer_context->GetNextSymName()); // W | ||
output_dims.push_back(symbol::DimExpr(2)); // 2 | ||
} else { | ||
// N * D * H * W * 3 | ||
output_dims.push_back(infer_context->GetNextSymName()); // D | ||
output_dims.push_back(infer_context->GetNextSymName()); // H | ||
output_dims.push_back(infer_context->GetNextSymName()); // W | ||
output_dims.push_back(symbol::DimExpr(3)); // 3 | ||
} | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. D, H, W可以从output_shape中得到,再确认下吧 There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. done |
||
|
||
infer_context->SetShapeOrDataForValue( | ||
op->result(0), | ||
symbol::ShapeOrDataDimExprs{ | ||
symbol::TensorShapeOrDataDimExprs(output_dims)}); | ||
|
||
return true; | ||
} | ||
|
||
bool AllOpInferSymbolicShape(pir::Operation *op, | ||
pir::InferSymbolicShapeContext *infer_context) { | ||
const auto &axis = details::GetVectorAttr(op, "axis"); | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
output_shape 是IntArray类型有可能从attribute传过来,加一个逻辑判断
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
done