Extend to 1D and 3D the torch to linalg lowering of the average pool operator with count_include_pad = false

[ivangarcia44]

Currently the avg_pool2d PyTorch operation supports the cases where count_include_pad is true and false, but the avg_pool1d and avg_pool3d only the true case is supported (which is simpler).

The count_include_pad = false support for avg_pool2d was added by @AmosLewis in this change (reviewed by @rsuderman and @nirvedhmeshram) : #3235

In this change I generalized the logic added above. I also did some refactoring to the original code to reduce the size of the functions and to avoid redundancy when possible.

@sahas3 @dixinzhou @rafaelubal

[ivangarcia44]

Hi @AmosLewis , @rsuderman ,

Could you please review this change when you get a chance?

The change is just to generalize the average pooling divisor computation when count_include_pad = false. This was implemented for 2 dimensions. This change makes it work for N dimensions, and does some refactoring to make the methods smaller.

Thanks,
Ivan

[ivangarcia44]

Link to old PR: #4010

The review of @sahas3 can be found there.

[ivangarcia44]

Including additional reviewers familiar with the Pooling.cpp file: @vivekkhandelwal1, @lingzhiz1998

In addition to @AmosLewis , @rsuderman, @sahas3, @dixinzhou, @rafaelubal

[ivangarcia44]

Including @rafaelubalmw in addition to the other reviewers: @vivekkhandelwal1, @lingzhiz1998, @AmosLewis , @rsuderman, @sahas3, @dixinzhou, @rafaelubal

[dixinzhou]

Thanks for making the changes @ivangarcia44. The changes look good to me.

[sahas3]

LGTM

[amd-vivekag]

Hi @ivangarcia44 ,

This change is causing a testcase failure: iree-test-suites/onnx_ops/onnx/node/generated/test_averagepool_2d_ceil/
Can you please fix it?

I've created an issue for this: #4079

Thanks Vivek

[ivangarcia44]

Hi @amd-vivekag ,

I will take a look at this.

Thanks,
Ivan