Decomposition of aten.pixel_shuffle with static input shape

[newling]

For static tests (that is when the shape is know) for example:

@annotate_args([None, ([3, 18, 2, 2], torch.float32, True)])

The e2e passes. But only if the replacement op's return type is set as undefined (optional shape and type must be explicitly made unset), otherwise there's a error about the function return type.

For dynamic cases, for example if the above is replaced with

@annotate_args([None, ([-1, -1, -1, -1], torch.float32, True)])

There is a failure to lower to linalg from torch ("view op explicitly labelled as illegal"). This seems to be because the support for lowering from torch to linalg with dynamic shapes is limited.

[newling]

@vivekkhandelwal1 I have an e2e test case working, but only for static shapes (see my comment above). I am not sure what to do about the dynamic shape case, is my thought that torch->linalg for aten.view is limited, and that's why I'm having problems?

More generally, I'd like to understand dynamic shapes end-to-end in torch-mlir, any pointers?

[vivekkhandelwal1]

@vivekkhandelwal1 I have an e2e test case working, but only for static shapes (see my comment above). I am not sure what to do about the dynamic shape case, is my thought that torch->linalg for aten.view is limited, and that's why I'm having problems?

More generally, I'd like to understand dynamic shapes end-to-end in torch-mlir, any pointers?

Hi @newling, the support for dynamic shapes in Torch-MLIR is limited, especially for view-like ops. I think, for now, you can follow either of the below two paths:
1.) Add a condition in your decomp restricting the support for only static shapes (Easy and fast path, but limited support).
2.) Directly add a Torch -> Linalg lowering for this op, if it seems feasible (Hard and slow path, but you will get complete support for the op).

For the Torch -> Linalg support you can refer to this: https://github.com/pytorch/pytorch/blob/fa9045a8725214c05ae4dcec5a855820b861155e/aten/src/ATen/native/cpu/PixelShuffleKernel.cpp#L16-L53

Also, unfortunately, there is no documentation for the dynamic shapes in Torch-MLIR. It's just that ops in Torch-MLIR broadly fall into one or a combination of the below four categories:
1.) Full-Dynamic Support
2.) Partial-Dynamic Support
3.) Full-Static Support
4.) Partial-Static Support.

You can go through some of the op lowerings for view-like ops to understand the functioning or dynamic-shape support in Torch-MLIR.

[newling]

Thanks for these information gems @vivekkhandelwal1

I am tempted to go with option 1 (just support static) for now, and leave comments in the code about what the issues currently are with dynamic shapes, and what possible remedies there are (like go directly to linalg, or maybe extend lowering of reshape).

Hopefully the use cases that @qedawkins has in mind are static.

[vivekkhandelwal1]

Thanks for these information gems @vivekkhandelwal1

I am tempted to go with option 1 (just support static) for now, and leave comments in the code about what the issues currently are with dynamic shapes, and what possible remedies there are (like go directly to linalg, or maybe extend lowering of reshape).

Hopefully the use cases that @qedawkins has in mind are static.

If it's required for Llama-2 then the use case would probably be dynamic. But @qedawkins can comment more on this,

[qedawkins]

Thanks for these information gems @vivekkhandelwal1
I am tempted to go with option 1 (just support static) for now, and leave comments in the code about what the issues currently are with dynamic shapes, and what possible remedies there are (like go directly to linalg, or maybe extend lowering of reshape).
Hopefully the use cases that @qedawkins has in mind are static.

If it's required for Llama-2 then the use case would probably be dynamic. But @qedawkins can comment more on this,

Ah no, this mostly shows up in vision models, and I would guess most of the time we would see static dims for the ones that are actually shuffled. There might be some dynamic outer/batch like dims though.

[newling]

Ok, so I've got it supporting just static shapes. I've left a comment in the code about why, and possible next steps.

I don't seem to be able to add reviewers, and I also need approval to start the workflows. How I get these permissions?

[qedawkins]

Ok, so I've got it supporting just static shapes. I've left a comment in the code about why, and possible next steps.

I don't seem to be able to add reviewers, and I also need approval to start the workflows. How I get these permissions?

Good question. IIRC there is a process to be added to the repo after you have landed a commit, but I don't know if the process has changed.

[qedawkins]

Mostly looks good, thanks for the patch! It has been a while since I've worked in depth with these decompositions and there might be some helpers already in place to help simplify code in a few places (looking for another reviewer who might know).

[qedawkins]

I left the comments that still need addressing as unresolved. Mainly some comment cleanup and missing shape/dtype refinement functions (sorry for the segmented review).

[qedawkins]

LGTM! Thanks for pushing this through

[qedawkins]

I think some linting check wants you to run

./build_tools/update_abstract_interp_lib.sh

[newling]

Thanks for the help @qedawkins!

I think some linting check wants you to run

./build_tools/update_abstract_interp_lib.sh

Yeah, that was a very useful diagnostic message from CI 😍

[newling]

@qedawkins would you mind kickstarting CI? It was failing with an "unexpected pass" failure, so I added the 2 tests to projects/pt1/e2e_testing/xfail_sets.py (a bit of a shot in the dark)

General question: Is there a simple way to run all the tests locally, with a single command?

[qedawkins]

There are some notes on testing here: https://github.com/llvm/torch-mlir/blob/main/docs/development.md#testing
(I don't remember off the top of my head what the command you're looking for is)

[newling]

There isn't much material about updating xfail_sets.py... but I'm optimistic I've got it right now (added to LTC_XFAIL_SET). Can you please restart the workflow @qedawkins ?

[qedawkins]

There isn't much material about updating xfail_sets.py... but I'm optimistic I've got it right now (added to LTC_XFAIL_SET). Can you please restart the workflow @qedawkins ?

Sorry about needing approval each time to run the CI :(
Let's get this in and get you direct access

[ramiro050]

Some small requests

[qedawkins]

(Merging because it looks like Ramiro's concerns were addressed and there is no PR blocker in place)