Skip to content

[Inductor] support masked vectorization for the tail_loop for dynamic shapes#131745

Closed
jiayisunx wants to merge 27 commits intogh/jiayisunx/16/basefrom
gh/jiayisunx/16/head
Closed

[Inductor] support masked vectorization for the tail_loop for dynamic shapes#131745
jiayisunx wants to merge 27 commits intogh/jiayisunx/16/basefrom
gh/jiayisunx/16/head

Conversation

@jiayisunx
Copy link
Copy Markdown
Collaborator

@jiayisunx jiayisunx commented Jul 25, 2024 •
edited
Loading

@jiayisunx
Update
55766ac 
[ghstack-poisoned]
@pytorch-bot
Copy link
Copy Markdown

pytorch-bot Bot commented Jul 25, 2024 •
edited
Loading

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/131745

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit fcfebe5 with merge base 1754850 (image):
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@jiayisunx jiayisunx mentioned this pull request Jul 25, 2024
Closed
@pytorch-bot pytorch-bot Bot added ciflow/inductor module: cpu CPU specific problem (e.g., perf, algorithm) module: inductor labels Jul 25, 2024
This was referenced Jul 25, 2024
Closed
Closed
Closed
@jiayisunx jiayisunx marked this pull request as draft July 25, 2024 06:11
jiayisunx added 19 commits July 24, 2024 23:31
@jiayisunx
Update
a72a81b 
[ghstack-poisoned]
@jiayisunx
Update
9b2ee93 
[ghstack-poisoned]
@jiayisunx
Update
1e37bcd 
[ghstack-poisoned]
@jiayisunx
Update
c64dcc6 
[ghstack-poisoned]
@jiayisunx
Update
89b6b27 
[ghstack-poisoned]
@jiayisunx
Update
423f94e 
[ghstack-poisoned]
@jiayisunx
Update
bc376f3 
[ghstack-poisoned]
@jiayisunx
Update
e6ec444 
[ghstack-poisoned]
@jiayisunx
Update
d2f2ddb 
[ghstack-poisoned]
@jiayisunx
Update
a563c1c 
[ghstack-poisoned]
@jiayisunx
Update
80e03d6 
[ghstack-poisoned]
@jiayisunx
Update
ea8a871 
[ghstack-poisoned]
@jiayisunx
Update
3db8d56 
[ghstack-poisoned]
@jiayisunx
Update
50ae6db 
[ghstack-poisoned]
@jiayisunx
Update
9c0dc58 
[ghstack-poisoned]
@jiayisunx
Update
64ad42a 
[ghstack-poisoned]
@jiayisunx
Update
9d6c37b 
[ghstack-poisoned]
@jiayisunx
Update
e0e95cd 
[ghstack-poisoned]
@jiayisunx
Update
f7299d0 
[ghstack-poisoned]
jgong5
jgong5 requested changes Aug 29, 2024
View reviewed changes
Comment thread aten/src/ATen/cpu/vec/vec512/vec512_bfloat16.h
template <typename T, int M, int N,
typename std::enable_if_t<std::is_same<T, BFloat16>::value && ((M < 32 && M != 16) || (N < 32 && N != 16)), int> = 0>
inline void transpose_mxn(const BFloat16* src, int64_t ld_src, BFloat16* dst, int64_t ld_dst) {
inline void transpose_mxn<BFloat16>(const BFloat16* src, int64_t ld_src, BFloat16* dst, int64_t ld_dst, int M, int N) {
Copy link
Copy Markdown
Collaborator

@jgong5 jgong5 Aug 29, 2024

Choose a reason for hiding this comment

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

I'm concerned about the perf impact. Originally, M and N are template args and compile-time constants.

Copy link
Copy Markdown
Collaborator Author

@jiayisunx jiayisunx Aug 29, 2024

Choose a reason for hiding this comment

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

I will check it with TAS.

Copy link
Copy Markdown
Collaborator Author

@jiayisunx jiayisunx Sep 4, 2024

Choose a reason for hiding this comment

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

No performance regression in TAS, and we compared assembly code before and after PR, it doesn't seem to have performance impact.

Copy link
Copy Markdown
Collaborator

@jgong5 jgong5 Sep 5, 2024

Choose a reason for hiding this comment

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

Thanks! So perhaps we don't have to keep the version with M and N as template args.

leslie-fang-intel
leslie-fang-intel approved these changes Aug 29, 2024
View reviewed changes
Copy link
Copy Markdown
Collaborator

@leslie-fang-intel leslie-fang-intel left a comment

Choose a reason for hiding this comment

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

LGTM, please kindly address Jiong's comment.

@jiayisunx jiayisunx requested a review from jansel September 3, 2024 07:23
@jiayisunx
Update
c58b2ad 
[ghstack-poisoned]
@jiayisunx
Update
c088cad 
[ghstack-poisoned]
@jiayisunx jiayisunx added the topic: not user facing topic category label Sep 4, 2024
jansel
jansel requested changes Sep 4, 2024
View reviewed changes
Copy link
Copy Markdown
Contributor

@jansel jansel left a comment

Choose a reason for hiding this comment

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

Failing tests?

@jiayisunx
Update
d046d55 
[ghstack-poisoned]
@jiayisunx
Update
839b37d 
[ghstack-poisoned]
@jiayisunx jiayisunx requested a review from jgong5 September 4, 2024 07:57
@jiayisunx
Copy link
Copy Markdown
Collaborator Author

jiayisunx commented Sep 4, 2024

Failing tests?

A bug introduced by rebase, I have fixed it, please help to review this PR again, thanks!

@jiayisunx jiayisunx requested a review from jansel September 4, 2024 12:07
jgong5
jgong5 approved these changes Sep 5, 2024
View reviewed changes
Comment thread torch/_inductor/codegen/cpp.py
Comment on lines +4705 to +4707
steps_str = (
f"{self.var}+=({cexpr_index(self.steps)} == 0 ? "
f"1 : {cexpr_index(self.steps)})"
Copy link
Copy Markdown
Collaborator

@jgong5 jgong5 Sep 5, 2024

Choose a reason for hiding this comment

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

Can you add comment on why we need this trick here?

Copy link
Copy Markdown
Collaborator Author

@jiayisunx jiayisunx Sep 5, 2024

Choose a reason for hiding this comment

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

Added, thanks!

Comment thread aten/src/ATen/cpu/vec/vec512/vec512_bfloat16.h
template <typename T, int M, int N,
typename std::enable_if_t<std::is_same<T, BFloat16>::value && ((M < 32 && M != 16) || (N < 32 && N != 16)), int> = 0>
inline void transpose_mxn(const BFloat16* src, int64_t ld_src, BFloat16* dst, int64_t ld_dst) {
inline void transpose_mxn<BFloat16>(const BFloat16* src, int64_t ld_src, BFloat16* dst, int64_t ld_dst, int M, int N) {
Copy link
Copy Markdown
Collaborator

@jgong5 jgong5 Sep 5, 2024

Choose a reason for hiding this comment

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

Thanks! So perhaps we don't have to keep the version with M and N as template args.

@jiayisunx
Update
fcfebe5 
[ghstack-poisoned]
@jiayisunx
Copy link
Copy Markdown
Collaborator Author

jiayisunx commented Sep 5, 2024

@pytorchbot merge

@pytorchmergebot
Copy link
Copy Markdown
Collaborator

pytorchmergebot commented Sep 5, 2024

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging
Check the merge workflow status
here

@xuhancn xuhancn mentioned this pull request Sep 5, 2024
Closed
@jiayisunx jiayisunx mentioned this pull request Sep 6, 2024
Closed
pytorchmergebot pushed a commit that referenced this pull request Sep 6, 2024
@jiayisunx @pytorchmergebot
[inductor] Fix gen_transposed_tile_load_store (#135307)
67c7924 
Recent PR: #131745 bring new VLA logical in cpp codegen. And it will raise build fail error on MSVC and error code is `Compiler Error C2131`: https://learn.microsoft.com/en-us/cpp/error-messages/compiler-errors-1/compiler-error-c2131?view=msvc-170

reproduce UT:
```cmd
pytest test\inductor\test_torchinductor_dynamic_shapes.py -v -k test_large_block_sizes_dynamic_shapes_cpu
```

Original generated code:
```c++
alignas(16) float tmp1[static_cast<int64_t>(((-256LL)*(c10::div_floor_integer(static_cast<int64_t>(ks1), static_cast<int64_t>(16LL)))) + (16LL*ks1))];
```

Changes:
allocate a large-enough fixed-sized buffer.

New genarated code:
```c++
alignas(16) float tmp1[16*16];
```

Pull Request resolved: #135307
Approved by: https://github.com/jgong5, https://github.com/jansel
Chao1Han pushed a commit to Chao1Han/pytorch that referenced this pull request Sep 20, 2024
@jiayisunx @Chao1Han
[Inductor] support masked vectorization for the tail_loop for dynamic…
5e55d56 
… shapes (pytorch#131745)

Pull Request resolved: pytorch#131745
Approved by: https://github.com/jgong5, https://github.com/leslie-fang-intel, https://github.com/jansel
Chao1Han pushed a commit to Chao1Han/pytorch that referenced this pull request Sep 20, 2024
@jiayisunx @Chao1Han
[inductor] Fix gen_transposed_tile_load_store (pytorch#135307)
3538a8d 
Recent PR: pytorch#131745 bring new VLA logical in cpp codegen. And it will raise build fail error on MSVC and error code is `Compiler Error C2131`: https://learn.microsoft.com/en-us/cpp/error-messages/compiler-errors-1/compiler-error-c2131?view=msvc-170

reproduce UT:
```cmd
pytest test\inductor\test_torchinductor_dynamic_shapes.py -v -k test_large_block_sizes_dynamic_shapes_cpu
```

Original generated code:
```c++
alignas(16) float tmp1[static_cast<int64_t>(((-256LL)*(c10::div_floor_integer(static_cast<int64_t>(ks1), static_cast<int64_t>(16LL)))) + (16LL*ks1))];
```

Changes:
allocate a large-enough fixed-sized buffer.

New genarated code:
```c++
alignas(16) float tmp1[16*16];
```

Pull Request resolved: pytorch#135307
Approved by: https://github.com/jgong5, https://github.com/jansel
@github-actions github-actions Bot deleted the gh/jiayisunx/16/head branch October 6, 2024 02:08
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@jansel jansel jansel approved these changes

@jgong5 jgong5 jgong5 approved these changes

@leslie-fang-intel leslie-fang-intel leslie-fang-intel approved these changes

Assignees

No one assigned

Labels

ciflow/inductor ciflow/trunk Trigger trunk jobs on your pull request Merged module: cpu CPU specific problem (e.g., perf, algorithm) module: inductor open source topic: not user facing topic category

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

6 participants

@jiayisunx @pytorchmergebot @jansel @jgong5 @leslie-fang-intel @pytorchbot