[MLIR][ONNX] Add OnnxToTorch support for GlobalMaxPool Op

[NeverRaR]

nod-ai/SHARK-Turbine#658

[renxida]

Hello! Welcome!

Nice PR & much appreciated! Left some comments for code quality. I hope you find them helpful.

[vivekkhandelwal1]

Hi @NeverRaR, I see many things done wrongly in this lowering. Please correct them.

• First of all, you should not be lowering this op to AdaptiveMaxPooling, instead use "AtenMaxPooldOp".
• Your lowering returns 2 results while the original Onnx op expects 1 result.
• The logic for the computation of output shape is also wrong, we can't just assume any dynamic dim to be 1.

[NeverRaR]

Hi @NeverRaR, I see many things done wrongly in this lowering. Please correct them.

• First of all, you should not be lowering this op to AdaptiveMaxPooling, instead use "AtenMaxPooldOp".
• Your lowering returns 2 results while the original Onnx op expects 1 result.
• The logic for the computation of output shape is also wrong, we can't just assume any dynamic dim to be 1.

@vivekkhandelwal1

• The reason for using AdaptiveMaxPooling is that the current torch-mlir does not support AtenAdaptiveMaxPool1d, and my model happens to require this situation.
• Directly replacing onnx.GlobalMaxPool with AdaptiveMaxPool will still use the result it needs (the first result of AtenAdaptiveMaxPool)
• And why it is assumed that all dynamic dims are 1. Actually, according to the onnx specification:

Output data tensor from pooling across the input tensor The output tensor has the same rank as the input The first two dimensions of output shape are the same as the input (N x C), while the other dimensions are all 1

[vivekkhandelwal1]

• The reason for using AdaptiveMaxPooling is that the current torch-mlir does not support AtenAdaptiveMaxPool1d, and my model happens to require this situation.

If that's the case, you could have added a lowering for AtenMaxPool1dOp.

• Directly replacing onnx.GlobalMaxPool with AdaptiveMaxPool will still use the result it needs (the first result of AtenAdaptiveMaxPool)

No, that's not true. You can't replace an op with one result with an op with different number of results. See the crash in the CI for the same.

  torch-mlir-opt: /_work/torch-mlir/torch-mlir/externals/llvm-project/mlir/lib/Transforms/Utils/DialectConversion.cpp:1645: virtual void mlir::ConversionPatternRewriter::replaceOp(Operation *, ValueRange): Assertion `op->getNumResults() == newValues.size() && "incorrect # of replacement values"' failed.

• And why it is assumed that all dynamic dims are 1. Actually, according to the onnx specification:

Output data tensor from pooling across the input tensor The output tensor has the same rank as the input The first two dimensions of output shape are the same as the input (N x C), while the other dimensions are all 1

If that's the case then there's no need to put the check, just assign value 1.

[NeverRaR]

• The reason for using AdaptiveMaxPooling is that the current torch-mlir does not support AtenAdaptiveMaxPool1d, and my model happens to require this situation.

If that's the case, you could have added a lowering for AtenMaxPool1dOp.

• Directly replacing onnx.GlobalMaxPool with AdaptiveMaxPool will still use the result it needs (the first result of AtenAdaptiveMaxPool)

No, that's not true. You can't replace an op with one result with an op with different number of results. See the crash in the CI for the same.

  torch-mlir-opt: /_work/torch-mlir/torch-mlir/externals/llvm-project/mlir/lib/Transforms/Utils/DialectConversion.cpp:1645: virtual void mlir::ConversionPatternRewriter::replaceOp(Operation *, ValueRange): Assertion `op->getNumResults() == newValues.size() && "incorrect # of replacement values"' failed.

• And why it is assumed that all dynamic dims are 1. Actually, according to the onnx specification:

Output data tensor from pooling across the input tensor The output tensor has the same rank as the input The first two dimensions of output shape are the same as the input (N x C), while the other dimensions are all 1

If that's the case then there's no need to put the check, just assign value 1.

Okay, I think I understand now. All modifications have been completed.I will added a lowering for AtenMaxPool1dOp in next pr.

[NeverRaR]

@vivekkhandelwal1 Does the failure of CI seem unrelated to my code?

[vivekkhandelwal1]

@vivekkhandelwal1 Does the failure of CI seem unrelated to my code?

Yes

[vivekkhandelwal1]

LGTM

[vivekkhandelwal1]

@NeverRaR, you need approval from other reviewers as well to merge this.