[inductor] Add size, stride, storage_offset to RAIIAtenTensorHandle #110764

aakhundov · 2023-10-06T21:07:30Z

Stack from ghstack (oldest at bottom):

Summary: For unbacked SymInts, the C++ wrapper codegen can generate expressions like buf123.size() or .stride() or .storage_offset():

pytorch/torch/_inductor/ir.py

Lines 2504 to 2520 in 7cc0020

    
           for i, s in enumerate(self.get_size()): 
        
               if s in symbols_to_define: 
        
                   wrapper.writeline( 
        
                       f"{wrapper.declare}{s} = {self.get_name()}.size({i}){wrapper.ending}" 
        
                   ) 
        
                   symbols_to_define.remove(s) 
        
           for i, s in enumerate(self.get_stride()): 
        
               if s in symbols_to_define: 
        
                   wrapper.writeline( 
        
                       f"{wrapper.declare}{s} = {self.get_name()}.stride({i}){wrapper.ending}" 
        
                   ) 
        
                   symbols_to_define.remove(s) 
        
           if (s := self.get_offset()) in symbols_to_define: 
        
               wrapper.writeline( 
        
                   f"{wrapper.declare}{s} = {self.get_name()}.storage_offset(){wrapper.ending}" 
        
               ) 
        
               symbols_to_define.remove(s)

Here we add corresponding methods to the RAIIAtenTensorHandle class so that the above codegen works in the ABI compatibility mode.

Test Plan: CI + the following PR

Reviewers:

Subscribers:

Tasks:

Tags:

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

Summary: For unbacked SymInts, the C++ wrapper codegen can generate expressions like `buf123.size()` or `.stride()` or `.storage_offset()`: https://github.com/pytorch/pytorch/blob/7cc0020a80527207a1725e6d21ce7c326668fe0d/torch/_inductor/ir.py#L2504-L2520 Here we add corresponding methods to the `RAIIAtenTensorHandle` class so that the above codegen works in the ABI compatibility mode. Test Plan: CI + the following PR Reviewers: Subscribers: Tasks: Tags: [ghstack-poisoned]

pytorch-bot · 2023-10-06T21:07:32Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

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

✅ You can merge normally! (3 Unrelated Failures)

As of commit 24f17e4 with merge base 371d8ba ():

UNSTABLE - The following jobs failed but were likely due to flakiness present on trunk and has been marked as unstable:

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

…sorHandle" Summary: For unbacked SymInts, the C++ wrapper codegen can generate expressions like `buf123.size()` or `.stride()` or `.storage_offset()`: https://github.com/pytorch/pytorch/blob/7cc0020a80527207a1725e6d21ce7c326668fe0d/torch/_inductor/ir.py#L2504-L2520 Here we add corresponding methods to the `RAIIAtenTensorHandle` class so that the above codegen works in the ABI compatibility mode. Test Plan: CI + the following PR Reviewers: Subscribers: Tasks: Tags: [ghstack-poisoned]

Summary: For unbacked SymInts, the C++ wrapper codegen can generate expressions like `buf123.size()` or `.stride()` or `.storage_offset()`: https://github.com/pytorch/pytorch/blob/7cc0020a80527207a1725e6d21ce7c326668fe0d/torch/_inductor/ir.py#L2504-L2520 Here we add corresponding methods to the `RAIIAtenTensorHandle` class so that the above codegen works in the ABI compatibility mode. Test Plan: CI + the following PR Reviewers: Subscribers: Tasks: Tags: ghstack-source-id: 56f1bd2 Pull Request resolved: #110764

chenyang78 · 2023-10-06T21:54:13Z

I have a high-level question. We already have API calls like aoti_torch_get_sizes that can be used to implement aoti_torch_get_size, e.g.

int64_t *sizes;
aoti_torch_get_sizes(..., &sizes, ...);
int64_t sz = sizes[d];

So, instead of adding new APIs, I am wondering if we could just code-gen a similar code like above by leveraging the existing APIs. One benefit would be that we end up maintaining a minimal set of interfaces.

aakhundov · 2023-10-06T22:01:35Z

So, instead of adding new APIs, I am wondering if we could just code-gen a similar code like above by leveraging the existing APIs. One benefit would be that we end up maintaining a minimal set of interfaces.

@chenyang78 I agree that reducing the extent of the API is generally good. However, the problem with relying on the existing functions is that it's impossible to do range checking, as they simply provide a pointer to the list of sizes. We've discussed this briefly offline with @desertfire and decided to delegate the range checking (as well as proper error messaging) to ATen ops. And that's the reason for the new functions relying on tensor->size(d) and friends instead of tensor->sizes() as in the pre-existing API. Does this make sense?

chenyang78 · 2023-10-06T22:08:16Z

So, instead of adding new APIs, I am wondering if we could just code-gen a similar code like above by leveraging the existing APIs. One benefit would be that we end up maintaining a minimal set of interfaces.

@chenyang78 I agree that reducing the extent of the API is generally good. However, the problem with relying on the existing functions is that it's impossible to do range checking, as they simply provide a pointer to the list of sizes. We've discussed this briefly offline with @desertfire and decided to delegate the range checking (as well as proper error messaging) to ATen ops. And that's the reason for the new functions relying on tensor->size(d) and friends instead of tensor->sizes() as in the pre-existing API. Does this make sense?

I see. Make sense. Thanks for the clarification.

…sorHandle" Summary: For unbacked SymInts, the C++ wrapper codegen can generate expressions like `buf123.size()` or `.stride()` or `.storage_offset()`: https://github.com/pytorch/pytorch/blob/7cc0020a80527207a1725e6d21ce7c326668fe0d/torch/_inductor/ir.py#L2504-L2520 Here we add corresponding methods to the `RAIIAtenTensorHandle` class so that the above codegen works in the ABI compatibility mode. Test Plan: CI + the following PR Reviewers: Subscribers: Tasks: Tags: cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler [ghstack-poisoned]

Summary: For unbacked SymInts, the C++ wrapper codegen can generate expressions like `buf123.size()` or `.stride()` or `.storage_offset()`: https://github.com/pytorch/pytorch/blob/7cc0020a80527207a1725e6d21ce7c326668fe0d/torch/_inductor/ir.py#L2504-L2520 Here we add corresponding methods to the `RAIIAtenTensorHandle` class so that the above codegen works in the ABI compatibility mode. Test Plan: CI + the following PR Reviewers: Subscribers: Tasks: Tags: ghstack-source-id: 56f1bd2 Pull Request resolved: #110764

…sorHandle" Summary: For unbacked SymInts, the C++ wrapper codegen can generate expressions like `buf123.size()` or `.stride()` or `.storage_offset()`: https://github.com/pytorch/pytorch/blob/7cc0020a80527207a1725e6d21ce7c326668fe0d/torch/_inductor/ir.py#L2504-L2520 Here we add corresponding methods to the `RAIIAtenTensorHandle` class so that the above codegen works in the ABI compatibility mode. Test Plan: CI + the following PR Reviewers: Subscribers: Tasks: Tags: cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler [ghstack-poisoned]

aakhundov · 2023-10-07T08:23:43Z

@pytorchbot merge

pytorchmergebot · 2023-10-07T08:25:53Z

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

Summary: `torch.nonzero` doesn't have inductor lowering (yet). To invoke the operator in AOT Inductor's ABI compatibility mode we need a dedicated shim function. Test Plan: ``` $ python test/inductor/test_aot_inductor.py -k test_zero_grid_with_unbacked_symbols ... ---------------------------------------------------------------------- Ran 4 tests in 78.650s OK ``` Reviewers: Subscribers: Tasks: Tags: Pull Request resolved: #110766 Approved by: https://github.com/chenyang78 ghstack dependencies: #110713, #110745, #110764

This was referenced Oct 6, 2023

[inductor] Add aoti_torch_dtype_bool to AOTI ABI shim #110713

Closed

[inductor] Add AOTI ABI shim function for repeat_interleave.Tensor #110745

Closed

aakhundov requested review from SherlockNoMad, chenyang78 and desertfire October 6, 2023 21:08

aakhundov added the topic: not user facing topic category label Oct 6, 2023

aakhundov changed the title ~~Add size, stride and storage_offset methods to RAIIAtenTensorHandle~~ [inductor] Add size, stride, storage_offset to RAIIAtenTensorHandle Oct 6, 2023

aakhundov added ciflow/trunk Trigger trunk jobs on your pull request module: inductor ciflow/inductor labels Oct 6, 2023

aakhundov mentioned this pull request Oct 6, 2023

[inductor] Add AOTI ABI shim function for torch.nonzero #110766

Closed

chenyang78 approved these changes Oct 6, 2023

View reviewed changes

aakhundov force-pushed the gh/aakhundov/5/head branch from 9fb36c4 to eae46d9 Compare October 6, 2023 22:25

aakhundov added 2 commits October 6, 2023 15:46

pytorchmergebot added the merging label Oct 7, 2023

pytorchmergebot added Merged and removed merging labels Oct 7, 2023

pytorchmergebot closed this in 13a2f42 Oct 7, 2023

facebook-github-bot deleted the gh/aakhundov/5/head branch October 10, 2023 14:25

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[inductor] Add size, stride, storage_offset to RAIIAtenTensorHandle #110764

[inductor] Add size, stride, storage_offset to RAIIAtenTensorHandle #110764

Uh oh!

aakhundov commented Oct 6, 2023 •

edited

Loading

Uh oh!

pytorch-bot bot commented Oct 6, 2023 •

edited

Loading

Uh oh!

chenyang78 commented Oct 6, 2023

Uh oh!

aakhundov commented Oct 6, 2023

Uh oh!

chenyang78 commented Oct 6, 2023

Uh oh!

aakhundov commented Oct 7, 2023

Uh oh!

pytorchmergebot commented Oct 7, 2023

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

	for i, s in enumerate(self.get_size()):
	if s in symbols_to_define:
	wrapper.writeline(
	f"{wrapper.declare}{s} = {self.get_name()}.size({i}){wrapper.ending}"
	)
	symbols_to_define.remove(s)
	for i, s in enumerate(self.get_stride()):
	if s in symbols_to_define:
	wrapper.writeline(
	f"{wrapper.declare}{s} = {self.get_name()}.stride({i}){wrapper.ending}"
	)
	symbols_to_define.remove(s)
	if (s := self.get_offset()) in symbols_to_define:
	wrapper.writeline(
	f"{wrapper.declare}{s} = {self.get_name()}.storage_offset(){wrapper.ending}"
	)
	symbols_to_define.remove(s)

[inductor] Add size, stride, storage_offset to RAIIAtenTensorHandle #110764

[inductor] Add size, stride, storage_offset to RAIIAtenTensorHandle #110764

Uh oh!

Conversation

aakhundov commented Oct 6, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Oct 6, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

✅ You can merge normally! (3 Unrelated Failures)

Uh oh!

chenyang78 commented Oct 6, 2023

Uh oh!

aakhundov commented Oct 6, 2023

Uh oh!

chenyang78 commented Oct 6, 2023

Uh oh!

aakhundov commented Oct 7, 2023

Uh oh!

pytorchmergebot commented Oct 7, 2023

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

aakhundov commented Oct 6, 2023 •

edited

Loading

pytorch-bot bot commented Oct 6, 2023 •

edited

Loading