Skip to content

[TOSA] Add legalization for avg_pool with count_include_pad=True #4273

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

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

Lallapallooza
Copy link

Before this patch, the avg_pool2d and avg_pool1d legalizations lacked support for pooling with count_include_pad=True. This patch introduces that support.

@sjarus sjarus requested a review from sahas3 July 16, 2025 06:41
Copy link
Collaborator

@sjarus sjarus left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please add an e2e test under projects/pt1/python/torch_mlir_e2e_test/test_suite/pooling.py as well.

// CHECK: %[[VAL_11:.*]] = tosa.const_shape {values = dense<[0, 0, 0, 0, 1, 1, 1, 1]> : tensor<8xindex>} : () -> !tosa.shape<8>
// CHECK: %[[VAL_12:.*]] = "tosa.const"() <{values = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor<1xf32>
// CHECK: %[[VAL_13:.*]] = tosa.pad %[[VAL_1]], %[[VAL_11]], %[[VAL_12]] : (tensor<1x192x35x35xf32>, !tosa.shape<8>, tensor<1xf32>) -> tensor<1x192x37x37xf32>
// CHECK: %[[VAL_14:.*]] = tosa.transpose %[[VAL_13]] {perms = array<i32: 0, 2, 3, 1>} : (tensor<1x192x37x37xf32>) -> tensor<1x37x37x192xf32>
Copy link
Member

@sahas3 sahas3 Jul 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the change.

I think it'll be better to pad after transposing the input data -- the transpose operations for back-to-back pool (or conv -> pool pattern) ops will be optimized out leading to better performance.

@Lallapallooza Lallapallooza marked this pull request as draft July 17, 2025 20:58
@Lallapallooza Lallapallooza force-pushed the tosa_legalize_avgpool branch from 85a9c86 to 2c0cc57 Compare July 17, 2025 21:02
Before this patch, the `avg_pool2d` and `avg_pool1d` legalizations lacked
support for pooling with count_include_pad=True. This patch introduces that support.

Signed-off-by: Vitalii Shutov <[email protected]>
Change-Id: I73fa26a58379e2c021929ade81c983ff91c59667
@Lallapallooza Lallapallooza force-pushed the tosa_legalize_avgpool branch from 2c0cc57 to 79612ae Compare July 17, 2025 21:27
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants