Converts all Adaptive Pooling Ops to Linalg #2808

zjgarvey · 2024-01-25T21:43:24Z

The previous conversions for AtenAdaptiveAvgPool1dOp and AtenAdaptiveMaxPool2dOp are refactored into a general templated conversion that works for all of the AtenAdaptive...PoolNdOp's.

New support is added for the following ops:

AtenAdaptiveMaxPool1d
AtenAdaptiveMaxPool3d
AtenAdaptiveAvgPool3d

Support is also provided for passing inputs without batch dimensions. For example, applying adaptive_avg_pool2d to an input tensor of rank 3.

After pytorch #118162 gets down to torch-mlir, I'll add a test for AdaptiveMaxPool1d with return_indices (which will pass with that upstream fix).

Also adds support for: - aten.adaptive_max_pool1d - aten.adaptive_max_pool3d - using inputs of rank n+1 instead of n+2 (i.e., no batches)

All adaptive pooling ops will now use the same templated conversion pattern.

zjgarvey · 2024-01-26T06:03:30Z

The AdaptiveAvgPool2d tests are failing on dynamo with the expected error message:

Exception: Unsupported: missing default value for argument 0 in schema for aten.div.Tensor_mode

So I know to add those to the appropriate place in the xfail sets.

However, the dynamo failures for AdaptiveAvgPool3d are more cryptic and say:

error: unknown: failed to legalize operation 'torch.constant.int'

I'll do some digging to figure out what is going on there. I would like to assume it is secretly encountering the same issue as the other tests, but I'd like to figure out why it isn't giving the same error message.

amended for a typo

newling

Looks like a nice improvement! Just a few minor suggestions.

lib/Conversion/TorchToLinalg/Pooling.cpp

also stores isMaxPool as part of the same struct as Dim

Seeing if it's possible to use polymorphism to clean up conversion

newling

This looks good to me. Thanks for making the suggested changes (the template to determine dim and isMax is very nice). @zjgarvey feel free to get the opinion of a more experienced torch-mlir-er if you like.

This one off-loads compile-time cf to polymorphic helper functions

zjgarvey · 2024-02-05T21:45:43Z

The helper functions require quite a few arguments, but overall I think this polymorphism approach is probably better. Let me know what you think, @newling.

newling · 2024-02-05T23:51:49Z

The helper functions require quite a few arguments, but overall I think this polymorphism approach is probably better. Let me know what you think, @newling.

Nice! I think it makes the main function easier to follow. I have made a PR zjgarvey#1 against this which may/mayn't be a better design.

zjgarvey · 2024-02-06T00:57:01Z

Nice! I think it makes the main function easier to follow. I have made a PR zjgarvey#1 against this which may/mayn't be a better design.

Yeah, I can see the advantage there especially if I can move all of the helper functions to the external structure and not need to define derived classes for the conversion. I'll check this out, thanks for the prototyping @newling .

zjgarvey · 2024-02-06T18:39:29Z

I'm going to make a few more changes right now to move the rest of the methods to the external helper.

zjgarvey · 2024-03-14T23:33:47Z

@rsuderman I just updated this PR so that it should be fine to merge without problems.

Would you mind taking a look?

renxida

I'm by no means familiar with this codebase and most of my comments are "i wonder" type comments. I would so love it if @zjgarvey could explain the thought processes behind some of the decisions. Maybe we should @ somebody more experienced to comment on it afterwards.

lib/Conversion/TorchToLinalg/Pooling.cpp

projects/pt1/python/torch_mlir_e2e_test/test_suite/pooling.py

lib/Conversion/TorchToLinalg/Pooling.cpp

newling

Apologies for the delay in reviewing. This looks good to me!

zjgarvey · 2024-03-21T22:30:45Z

@renxida do you think the comments I added in the last commit are sufficient to clear up the readability concerns you had earlier?

The previous conversions for AtenAdaptiveAvgPool1dOp and AtenAdaptiveMaxPool2dOp are refactored into a general templated conversion that works for all of the AtenAdaptive...PoolNdOp's. New support is added for the following ops: 1. AtenAdaptiveMaxPool1d 2. AtenAdaptiveMaxPool3d 3. AtenAdaptiveAvgPool3d Support is also provided for passing inputs without batch dimensions. For example, applying adaptive_avg_pool2d to an input tensor of rank 3. After [pytorch #118162](pytorch/pytorch#118162) gets down to torch-mlir, I'll add a test for AdaptiveMaxPool1d with return_indices (which will pass with that upstream fix). --------- Co-authored-by: James Newling <[email protected]>

zjgarvey added 7 commits January 24, 2024 20:58

Refactors Adaptive Pooling Conversion

0e1772b

Also adds support for: - aten.adaptive_max_pool1d - aten.adaptive_max_pool3d - using inputs of rank n+1 instead of n+2 (i.e., no batches)

Some cleaning

6767eb9

formatting

6057388

Removes an Unused Argument

8467d56

Minor Edits

0ef27d0

Added support for AdaptiveAvgPool2+3d

d1925e7

Further Refactored Adaptive Pooling Conversions

dd984b6

All adaptive pooling ops will now use the same templated conversion pattern.

zjgarvey force-pushed the adaptpoolgen branch from 74cf105 to 68d404d Compare January 29, 2024 18:51

Added Aten_AdaptiveAvgPool3d to conversion

60e9346

zjgarvey force-pushed the adaptpoolgen branch 3 times, most recently from 8a4afeb to 60e9346 Compare January 29, 2024 19:04

zjgarvey and others added 2 commits January 29, 2024 13:11

Merge branch 'llvm:main' into adaptpoolgen

3ec0e5b

added two xfails for dynamo

2d6784c

amended for a typo

zjgarvey force-pushed the adaptpoolgen branch from 1762cba to 2d6784c Compare January 29, 2024 19:27

newling requested changes Jan 30, 2024

View reviewed changes

zjgarvey added 2 commits February 1, 2024 14:16

Adds Dim inference from OpTy

7a03acd

also stores isMaxPool as part of the same struct as Dim

Preliminary (modular) refactor for Adaptive Pool

0479f10

Seeing if it's possible to use polymorphism to clean up conversion

newling approved these changes Feb 1, 2024

View reviewed changes

Refactored Adaptive Conversion (again)

c43c107

This one off-loads compile-time cf to polymorphic helper functions

Merge branch 'main' into adaptpoolgen

e2a7efb

zjgarvey force-pushed the adaptpoolgen branch from 77e6a84 to 9d45ed1 Compare February 6, 2024 18:35

alternative design suggestion

ac7fb6c

Rolls other methods into helper classes

9e26534

zjgarvey force-pushed the adaptpoolgen branch from 9d45ed1 to 9e26534 Compare February 6, 2024 18:48

Merge remote-tracking branch 'origin' into adaptpoolgen

d55db0b

zjgarvey force-pushed the adaptpoolgen branch from ab48a2d to e325555 Compare March 14, 2024 22:18

Updated to match current repo

8660363

zjgarvey force-pushed the adaptpoolgen branch 2 times, most recently from 4e5753f to 8660363 Compare March 14, 2024 22:49

renxida reviewed Mar 18, 2024

View reviewed changes

renxida requested a review from newling March 19, 2024 20:33

zjgarvey force-pushed the adaptpoolgen branch from f1f878a to de73f80 Compare March 20, 2024 21:26

Removed uneccesary function definitions; added comments

82fb009

zjgarvey force-pushed the adaptpoolgen branch from de73f80 to 82fb009 Compare March 20, 2024 21:30

newling approved these changes Mar 20, 2024

View reviewed changes

zjgarvey requested a review from renxida March 21, 2024 22:31

zjgarvey mentioned this pull request Mar 21, 2024

torch.aten._adaptive_avg_pool3d nod-ai/SHARK-ModelDev#545

Closed

renxida approved these changes Mar 22, 2024

View reviewed changes

renxida merged commit 99b3a5f into llvm:main Mar 22, 2024
3 checks passed

josel-amd mentioned this pull request Jun 7, 2024

Fixes avg pool 3d Xilinx/torch-mlir#171

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Converts all Adaptive Pooling Ops to Linalg #2808

Converts all Adaptive Pooling Ops to Linalg #2808

zjgarvey commented Jan 25, 2024

zjgarvey commented Jan 26, 2024

newling left a comment

newling left a comment

zjgarvey commented Feb 5, 2024

newling commented Feb 5, 2024

zjgarvey commented Feb 6, 2024

zjgarvey commented Feb 6, 2024

zjgarvey commented Mar 14, 2024

renxida left a comment

newling left a comment

zjgarvey commented Mar 21, 2024

Converts all Adaptive Pooling Ops to Linalg #2808

Converts all Adaptive Pooling Ops to Linalg #2808

Conversation

zjgarvey commented Jan 25, 2024

zjgarvey commented Jan 26, 2024

newling left a comment

Choose a reason for hiding this comment

newling left a comment

Choose a reason for hiding this comment

zjgarvey commented Feb 5, 2024

newling commented Feb 5, 2024

zjgarvey commented Feb 6, 2024

zjgarvey commented Feb 6, 2024

zjgarvey commented Mar 14, 2024

renxida left a comment

Choose a reason for hiding this comment

newling left a comment

Choose a reason for hiding this comment

zjgarvey commented Mar 21, 2024