Skip to content

[PyTorch] Replace cached thread_locals with stack allocation in AOTI - #112116

Closed
swolchok wants to merge 27 commits into
gh/swolchok/595/basefrom
gh/swolchok/595/head
Closed

[PyTorch] Replace cached thread_locals with stack allocation in AOTI#112116
swolchok wants to merge 27 commits into
gh/swolchok/595/basefrom
gh/swolchok/595/head

Conversation

@swolchok

@swolchok swolchok commented Oct 26, 2023 •
edited
Loading

Copy link
Copy Markdown
Contributor

Stack from ghstack (oldest at bottom):

This changes cached thread_local tensors to stack-allocated buffers. Since we were incidentally caching output in a thread_local, I had to add manual thread_local caching of outputs, which I implemented by caching a buffer and a Tensor whose storage is that buffer and then just memcpying the result into the cached buffer every time. Ideally, memory planning would be able to identify allocations that are the backing storage for outputs, but this should be good enough in the absence of planning.

Differential Revision: D50416438

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler

@swolchok
[PyTorch] Replace cached thread_locals with stack allocation in AOTI
6fc0267 
This changes cached thread_local tensors to stack-allocated buffers. Since we were incidentally caching output in a thread_local, I had to add manual thread_local caching of outputs, which I implemented by caching a buffer and a Tensor whose storage is that buffer and then just memcpying the result into the cached buffer every time. Ideally, memory planning would be able to identify allocations that are the backing storage for outputs, but this should be good enough in the absence of planning.

Differential Revision: [D50416438](https://our.internmc.facebook.com/intern/diff/D50416438/)

[ghstack-poisoned]
@pytorch-bot

pytorch-bot Bot commented Oct 26, 2023 •
edited
Loading

Copy link
Copy Markdown

🔗 Helpful Links

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

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

✅ No Failures

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

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

@swolchok swolchok mentioned this pull request Oct 26, 2023
Closed
@swolchok
Update on "[PyTorch] Replace cached thread_locals with stack allocati…
d76fcf2 
…on in AOTI"

This changes cached thread_local tensors to stack-allocated buffers. Since we were incidentally caching output in a thread_local, I had to add manual thread_local caching of outputs, which I implemented by caching a buffer and a Tensor whose storage is that buffer and then just memcpying the result into the cached buffer every time. Ideally, memory planning would be able to identify allocations that are the backing storage for outputs, but this should be good enough in the absence of planning.

Differential Revision: [D50416438](https://our.internmc.facebook.com/intern/diff/D50416438/)

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler

[ghstack-poisoned]
@swolchok swolchok mentioned this pull request Oct 26, 2023
Closed
jansel
jansel requested changes Oct 27, 2023
View reviewed changes

@jansel jansel left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Test failures?

@chenyang78

chenyang78 commented Oct 27, 2023

Copy link
Copy Markdown

LGTM overall. Left a minor comment. There are some test failures that might be related to this change?

chenyang78
chenyang78 reviewed Oct 27, 2023
View reviewed changes
Comment thread torch/_inductor/codegen/wrapper.py
self.expr_printer = pexpr
self.cached_thread_locals = set()
self.user_defined_kernel_count = 0
self.allow_stack_allocation = None

@chenyang78 chenyang78 Oct 27, 2023

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

Do we initialize allow_stack_allocation with None on purpose? If so, wondering what the difference would be between None and False for this flag.

@swolchok swolchok Oct 27, 2023

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

None is "not decided yet", False is "decided false"

@swolchok

swolchok commented Oct 27, 2023

Copy link
Copy Markdown
Contributor Author

Test failures?

the broken test is disabled in fbcode. I'll try a blind fix but if that doesn't work I have no way to debug it.

@swolchok
Update on "[PyTorch] Replace cached thread_locals with stack allocati…
24751c3 
…on in AOTI"

This changes cached thread_local tensors to stack-allocated buffers. Since we were incidentally caching output in a thread_local, I had to add manual thread_local caching of outputs, which I implemented by caching a buffer and a Tensor whose storage is that buffer and then just memcpying the result into the cached buffer every time. Ideally, memory planning would be able to identify allocations that are the backing storage for outputs, but this should be good enough in the absence of planning.

Differential Revision: [D50416438](https://our.internmc.facebook.com/intern/diff/D50416438/)

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler

[ghstack-poisoned]
@swolchok swolchok mentioned this pull request Oct 30, 2023
Closed
@swolchok
Update on "[PyTorch] Replace cached thread_locals with stack allocati…
70734df 
…on in AOTI"

This changes cached thread_local tensors to stack-allocated buffers. Since we were incidentally caching output in a thread_local, I had to add manual thread_local caching of outputs, which I implemented by caching a buffer and a Tensor whose storage is that buffer and then just memcpying the result into the cached buffer every time. Ideally, memory planning would be able to identify allocations that are the backing storage for outputs, but this should be good enough in the absence of planning.

Differential Revision: [D50416438](https://our.internmc.facebook.com/intern/diff/D50416438/)

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler

[ghstack-poisoned]
@swolchok
Update on "[PyTorch] Replace cached thread_locals with stack allocati…
dbf57cc 
…on in AOTI"

This changes cached thread_local tensors to stack-allocated buffers. Since we were incidentally caching output in a thread_local, I had to add manual thread_local caching of outputs, which I implemented by caching a buffer and a Tensor whose storage is that buffer and then just memcpying the result into the cached buffer every time. Ideally, memory planning would be able to identify allocations that are the backing storage for outputs, but this should be good enough in the absence of planning.

Differential Revision: [D50416438](https://our.internmc.facebook.com/intern/diff/D50416438/)

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler

[ghstack-poisoned]
@swolchok
swolchok requested review from chenyang78 and jansel October 30, 2023 21:27
@swolchok

swolchok commented Oct 30, 2023

Copy link
Copy Markdown
Contributor Author

lintrunner failure is clearly not related to this diff

@swolchok swolchok added the topic: not user facing topic category label Oct 30, 2023
@swolchok
Update on "[PyTorch] Replace cached thread_locals with stack allocati…
3db082c 
…on in AOTI"

This changes cached thread_local tensors to stack-allocated buffers. Since we were incidentally caching output in a thread_local, I had to add manual thread_local caching of outputs, which I implemented by caching a buffer and a Tensor whose storage is that buffer and then just memcpying the result into the cached buffer every time. Ideally, memory planning would be able to identify allocations that are the backing storage for outputs, but this should be good enough in the absence of planning.

Differential Revision: [D50416438](https://our.internmc.facebook.com/intern/diff/D50416438/)

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler

[ghstack-poisoned]
swolchok added a commit that referenced this pull request Oct 30, 2023
@swolchok
[PyTorch] Replace cached thread_locals with stack allocation in AOTI
9acbc7f 
Pull Request resolved: #112116

This changes cached thread_local tensors to stack-allocated buffers. Since we were incidentally caching output in a thread_local, I had to add manual thread_local caching of outputs, which I implemented by caching a buffer and a Tensor whose storage is that buffer and then just memcpying the result into the cached buffer every time. Ideally, memory planning would be able to identify allocations that are the backing storage for outputs, but this should be good enough in the absence of planning.
ghstack-source-id: 205864117
@exported-using-ghexport

Differential Revision: [D50416438](https://our.internmc.facebook.com/intern/diff/D50416438/)
@jansel

jansel commented Oct 31, 2023

Copy link
Copy Markdown
Contributor

lintrunner failure is clearly not related to this diff

Lint failure looks related to me.

>>> Lint for torch/_inductor/ir.py:
2023-10-31T00:00:01.7590302Z 
2023-10-31T00:00:01.7590494Z   Warning (UFMT) format
2023-10-31T00:00:01.7591001Z     Run `lintrunner -a` to apply this patch.
2023-10-31T00:00:01.7591306Z 
2023-10-31T00:00:01.7591590Z     You can run `lintrunner -a` to apply this patch.
2023-10-31T00:00:01.7591963Z 
2023-10-31T00:00:01.7592288Z     5232  5232 |         algorithm,
2023-10-31T00:00:01.7592834Z     5233  5233 |     ):
2023-10-31T00:00:01.7593433Z     5234  5234 |         transposed = True
2023-10-31T00:00:01.7594388Z     5234       |-        (inputs, constant_args, kernel_layout, _,) = _prepare_convolution_fusion_create(
2023-10-31T00:00:01.7595248Z           5235 |+        (
2023-10-31T00:00:01.7595820Z           5236 |+            inputs,
2023-10-31T00:00:01.7596438Z           5237 |+            constant_args,
2023-10-31T00:00:01.7597088Z           5238 |+            kernel_layout,
2023-10-31T00:00:01.7597685Z           5239 |+            _,
2023-10-31T00:00:01.7598389Z           5240 |+        ) = _prepare_convolution_fusion_create(
2023-10-31T00:00:01.7599088Z     5236  5241 |             cls,
2023-10-31T00:00:01.7599658Z     5237  5242 |             x,
2023-10-31T00:00:01.7600245Z     5238  5243 |             weight,

desertfire
desertfire reviewed Nov 1, 2023
View reviewed changes

@desertfire desertfire left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

The coding LGTM. Left some minor comment. A more general concern I have is this adds more complexity to the wrapper codegen. With the memory planning PR #112178 merged (reverted with some ROCM test failures to be fixed/skiped), should we give that a try to see it is sufficient enough for your use case?

Comment thread torch/_inductor/codegen/wrapper.py Outdated

static_shape = self.wrapper.static_shape_for_buffer_or_none(self.node)
if static_shape is not None:
state.total_allocated_buffer_size += int(

@desertfire desertfire Nov 1, 2023

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

In a case where there are both CPU and GPU buffers, I think we are counting both, which is less ideal.

@desertfire

desertfire commented Nov 1, 2023

Copy link
Copy Markdown
Contributor

Also I left some comments in D50416438 for fixing lint issues.

pytorchmergebot pushed a commit that referenced this pull request Dec 12, 2023
@swolchok @pytorchmergebot
[PyTorch] AOTI: Emit static constexpr int array vars when possible (#…
24c67fe 
…112174)

No need to populate a stack-based array for a shape/stride array when it's statically known.

Differential Revision: [D50699889](https://our.internmc.facebook.com/intern/diff/D50699889/)

Pull Request resolved: #112174
Approved by: https://github.com/chenyang78, https://github.com/desertfire, https://github.com/jansel
ghstack dependencies: #112116
pytorchmergebot pushed a commit that referenced this pull request Dec 12, 2023
@swolchok @pytorchmergebot
[PyTorch] AOTI: Avoid aoti_torch_data_ptr calls for constants at infe…
ca52195 
…rence time (#112405)

Cache aoti_torch_get_data_ptr at constants update time.

Differential Revision: [D50708982](https://our.internmc.facebook.com/intern/diff/D50708982/)

Pull Request resolved: #112405
Approved by: https://github.com/chenyang78, https://github.com/desertfire, https://github.com/khabinov
ghstack dependencies: #112116, #112174
pytorchmergebot pushed a commit that referenced this pull request Dec 12, 2023
@swolchok @pytorchmergebot
[PyTorch] AOTI: Use static_cast, not dynamic_cast (#112798)
2b323e6 
dynamic_cast is for when we aren't certain about the type. We are certain (and will crash anyway if we're wrong).

Differential Revision: [D50812978](https://our.internmc.facebook.com/intern/diff/D50812978/)

Pull Request resolved: #112798
Approved by: https://github.com/chenyang78, https://github.com/desertfire, https://github.com/jansel, https://github.com/khabinov
ghstack dependencies: #112116, #112174, #112405
pytorchmergebot pushed a commit that referenced this pull request Dec 12, 2023
@swolchok @pytorchmergebot
[PyTorch] AOTI: add more basic aoti_torch getters (#112799)
3de2596 
Lot of simple information about tensors we couldn't get. In
particular, we didn't know the lengths of the arrays returned by sizes
and strides.

Differential Revision: [D50949929](https://our.internmc.facebook.com/intern/diff/D50949929/)
Pull Request resolved: #112799
Approved by: https://github.com/desertfire, https://github.com/aakhundov
ghstack dependencies: #112116, #112174, #112405, #112798
pytorchmergebot pushed a commit that referenced this pull request Dec 12, 2023
@swolchok @pytorchmergebot
[PyTorch] AOTI: add more basic aoti_torch getters (#112799)
de4b2e5 
Lot of simple information about tensors we couldn't get. In
particular, we didn't know the lengths of the arrays returned by sizes
and strides.

Differential Revision: [D50949929](https://our.internmc.facebook.com/intern/diff/D50949929/)
Pull Request resolved: #112799
Approved by: https://github.com/desertfire, https://github.com/aakhundov
ghstack dependencies: #112116, #112174, #112405, #112798
@facebook-github-bot
facebook-github-bot deleted the gh/swolchok/595/head branch December 15, 2023 15:27
guilhermeleobas pushed a commit to guilhermeleobas/pytorch that referenced this pull request Dec 18, 2023
@swolchok @guilhermeleobas
[PyTorch] Replace cached thread_locals with stack allocation in AOTI (p…
1b01834 
…ytorch#112116)

This changes cached thread_local tensors to stack-allocated buffers. Since we were incidentally caching output in a thread_local, I had to add manual thread_local caching of outputs, which I implemented by caching a buffer and a Tensor whose storage is that buffer and then just memcpying the result into the cached buffer every time. Ideally, memory planning would be able to identify allocations that are the backing storage for outputs, but this should be good enough in the absence of planning.

Differential Revision: [D50416438](https://our.internmc.facebook.com/intern/diff/D50416438/)

Pull Request resolved: pytorch#112116
Approved by: https://github.com/jansel, https://github.com/desertfire
guilhermeleobas pushed a commit to guilhermeleobas/pytorch that referenced this pull request Dec 18, 2023
@swolchok @guilhermeleobas
[PyTorch] AOTI: Emit static constexpr int array vars when possible (p…
86b7f8c 
…ytorch#112174)

No need to populate a stack-based array for a shape/stride array when it's statically known.

Differential Revision: [D50699889](https://our.internmc.facebook.com/intern/diff/D50699889/)

Pull Request resolved: pytorch#112174
Approved by: https://github.com/chenyang78, https://github.com/desertfire, https://github.com/jansel
ghstack dependencies: pytorch#112116
guilhermeleobas pushed a commit to guilhermeleobas/pytorch that referenced this pull request Dec 18, 2023
@swolchok @guilhermeleobas
[PyTorch] AOTI: Avoid aoti_torch_data_ptr calls for constants at infe…
06e280f 
…rence time (pytorch#112405)

Cache aoti_torch_get_data_ptr at constants update time.

Differential Revision: [D50708982](https://our.internmc.facebook.com/intern/diff/D50708982/)

Pull Request resolved: pytorch#112405
Approved by: https://github.com/chenyang78, https://github.com/desertfire, https://github.com/khabinov
ghstack dependencies: pytorch#112116, pytorch#112174
guilhermeleobas pushed a commit to guilhermeleobas/pytorch that referenced this pull request Dec 18, 2023
@swolchok @guilhermeleobas
[PyTorch] AOTI: Use static_cast, not dynamic_cast (pytorch#112798)
c589ad7 
dynamic_cast is for when we aren't certain about the type. We are certain (and will crash anyway if we're wrong).

Differential Revision: [D50812978](https://our.internmc.facebook.com/intern/diff/D50812978/)

Pull Request resolved: pytorch#112798
Approved by: https://github.com/chenyang78, https://github.com/desertfire, https://github.com/jansel, https://github.com/khabinov
ghstack dependencies: pytorch#112116, pytorch#112174, pytorch#112405
guilhermeleobas pushed a commit to guilhermeleobas/pytorch that referenced this pull request Dec 18, 2023
@swolchok @guilhermeleobas
[PyTorch] AOTI: add more basic aoti_torch getters (pytorch#112799)
17a5a2d 
Lot of simple information about tensors we couldn't get. In
particular, we didn't know the lengths of the arrays returned by sizes
and strides.

Differential Revision: [D50949929](https://our.internmc.facebook.com/intern/diff/D50949929/)
Pull Request resolved: pytorch#112799
Approved by: https://github.com/desertfire, https://github.com/aakhundov
ghstack dependencies: pytorch#112116, pytorch#112174, pytorch#112405, pytorch#112798
guilhermeleobas pushed a commit to guilhermeleobas/pytorch that referenced this pull request Dec 18, 2023
@swolchok @guilhermeleobas
[PyTorch] AOTI: add more basic aoti_torch getters (pytorch#112799)
dac7055 
Lot of simple information about tensors we couldn't get. In
particular, we didn't know the lengths of the arrays returned by sizes
and strides.

Differential Revision: [D50949929](https://our.internmc.facebook.com/intern/diff/D50949929/)
Pull Request resolved: pytorch#112799
Approved by: https://github.com/desertfire, https://github.com/aakhundov
ghstack dependencies: pytorch#112116, pytorch#112174, pytorch#112405, pytorch#112798
dmenig pushed a commit to dmenig/pytorch that referenced this pull request Dec 21, 2023
@swolchok @dmenig
[PyTorch] Replace cached thread_locals with stack allocation in AOTI (p…
e613f4c 
…ytorch#112116)

This changes cached thread_local tensors to stack-allocated buffers. Since we were incidentally caching output in a thread_local, I had to add manual thread_local caching of outputs, which I implemented by caching a buffer and a Tensor whose storage is that buffer and then just memcpying the result into the cached buffer every time. Ideally, memory planning would be able to identify allocations that are the backing storage for outputs, but this should be good enough in the absence of planning.

Differential Revision: [D50416438](https://our.internmc.facebook.com/intern/diff/D50416438/)

Pull Request resolved: pytorch#112116
Approved by: https://github.com/jansel, https://github.com/desertfire
dmenig pushed a commit to dmenig/pytorch that referenced this pull request Dec 21, 2023
@swolchok @dmenig
[PyTorch] AOTI: Emit static constexpr int array vars when possible (p…
9510605 
…ytorch#112174)

No need to populate a stack-based array for a shape/stride array when it's statically known.

Differential Revision: [D50699889](https://our.internmc.facebook.com/intern/diff/D50699889/)

Pull Request resolved: pytorch#112174
Approved by: https://github.com/chenyang78, https://github.com/desertfire, https://github.com/jansel
ghstack dependencies: pytorch#112116
dmenig pushed a commit to dmenig/pytorch that referenced this pull request Dec 21, 2023
@swolchok @dmenig
[PyTorch] AOTI: Avoid aoti_torch_data_ptr calls for constants at infe…
ed663ea 
…rence time (pytorch#112405)

Cache aoti_torch_get_data_ptr at constants update time.

Differential Revision: [D50708982](https://our.internmc.facebook.com/intern/diff/D50708982/)

Pull Request resolved: pytorch#112405
Approved by: https://github.com/chenyang78, https://github.com/desertfire, https://github.com/khabinov
ghstack dependencies: pytorch#112116, pytorch#112174
dmenig pushed a commit to dmenig/pytorch that referenced this pull request Dec 21, 2023
@swolchok @dmenig
[PyTorch] AOTI: Use static_cast, not dynamic_cast (pytorch#112798)
0584dd2 
dynamic_cast is for when we aren't certain about the type. We are certain (and will crash anyway if we're wrong).

Differential Revision: [D50812978](https://our.internmc.facebook.com/intern/diff/D50812978/)

Pull Request resolved: pytorch#112798
Approved by: https://github.com/chenyang78, https://github.com/desertfire, https://github.com/jansel, https://github.com/khabinov
ghstack dependencies: pytorch#112116, pytorch#112174, pytorch#112405
dmenig pushed a commit to dmenig/pytorch that referenced this pull request Dec 21, 2023
@swolchok @dmenig
[PyTorch] AOTI: add more basic aoti_torch getters (pytorch#112799)
576f5fd 
Lot of simple information about tensors we couldn't get. In
particular, we didn't know the lengths of the arrays returned by sizes
and strides.

Differential Revision: [D50949929](https://our.internmc.facebook.com/intern/diff/D50949929/)
Pull Request resolved: pytorch#112799
Approved by: https://github.com/desertfire, https://github.com/aakhundov
ghstack dependencies: pytorch#112116, pytorch#112174, pytorch#112405, pytorch#112798
dmenig pushed a commit to dmenig/pytorch that referenced this pull request Dec 21, 2023
@swolchok @dmenig
[PyTorch] AOTI: add more basic aoti_torch getters (pytorch#112799)
f17732e 
Lot of simple information about tensors we couldn't get. In
particular, we didn't know the lengths of the arrays returned by sizes
and strides.

Differential Revision: [D50949929](https://our.internmc.facebook.com/intern/diff/D50949929/)
Pull Request resolved: pytorch#112799
Approved by: https://github.com/desertfire, https://github.com/aakhundov
ghstack dependencies: pytorch#112116, pytorch#112174, pytorch#112405, pytorch#112798
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@ezyang ezyang ezyang left review comments
@jansel jansel jansel approved these changes
@desertfire desertfire desertfire approved these changes
@chenyang78 chenyang78 Awaiting requested review from chenyang78

Assignees

No one assigned

Labels

ciflow/inductor ciflow/trunk Trigger trunk jobs on your pull request Merged module: inductor topic: not user facing topic category

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

6 participants

@swolchok @chenyang78 @jansel @desertfire @ezyang @pytorchmergebot