-
Notifications
You must be signed in to change notification settings - Fork 374
Propagate onnx.dim_params in shape inference #3160
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
base: main
Are you sure you want to change the base?
Conversation
Signed-off-by: Chen Tong <[email protected]>
Signed-off-by: Chen Tong <[email protected]>
Signed-off-by: Chen Tong <[email protected]>
Signed-off-by: Chen Tong <[email protected]>
| if (dimParamsStr == "") | ||
| return; | ||
| StringAttr dimParamsAttr = StringAttr::get(op->getContext(), dimParamsStr); | ||
| op->setAttr("onnx.dim_params_" + std::to_string(resultIndex), |
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.
Suggestion: how about using onnx.out_dim_params_? I thought onnx.dim_params_0 was for the first input, but actually it was for the first output.
tungld
left a comment
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.
LGTM with a minor change in naming.
The propagation looks effective. Let's see how we can utilize this information in future PRs.
For the changes to IndexExpr, I would like to hear comments from @AlexandreEichenberger as well.
Signed-off-by: Chen Tong <[email protected]>
Signed-off-by: Chen Tong <[email protected]>
Signed-off-by: Chen Tong <[email protected]>
Signed-off-by: Chen Tong <[email protected]>
Signed-off-by: Chen Tong <[email protected]>
Signed-off-by: Chen Tong <[email protected]>
Signed-off-by: Chen Tong <[email protected]>
Signed-off-by: Chen Tong <[email protected]>
Signed-off-by: Chen Tong <[email protected]>
Signed-off-by: Chen Tong <[email protected]>
Signed-off-by: Tong Chen <[email protected]>
ONNX model provides symbolic name for dynamic dimension of input tensor with argument attribute "onnx.dim_params". This PR tries to propagate this symbolic name for dynamic dimension during shape inference with extra attribute to operation "onnx.dim_params_[n]", where n is the index of the output tensor.
The major change is on IndexExpr to carry such info, if it exists. Other changes include the use of the dim_param info in shape inference and test case.
So far, the usu of dim_param info is far from complete. The PR is kind of experimental.
At least one benefit of this approach is that the readability of the output of onnx-mlir can be improved with the symbolic names.
Another potential benefit is the reuse the existing shape inference code to do symbolic dim propagation. For example, without touching the shape inference for onnx.MatMul, we can get the result of the following test case:
Result from
onnx-mlir-opt --shape-inference