[AOTInductor] ProxyExecutor supports List[Tensor] return type #110182
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/110182
Note: Links to docs will display an error until the docs builds have been completed. ⏳ No Failures, 1 PendingAs of commit 7f77a31 with merge base 81da6db ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
This pull request was exported from Phabricator. Differential Revision: D49694691 |
This PR needs a
|
|
Does this PR overlap with #110140? |
|
This pull request was exported from Phabricator. Differential Revision: D49694691 |
…h#110182) Summary: Pull Request resolved: pytorch#110182 Support custom ops returns List[Tensor] type, like `"fn_with_list_output(Tensor[] tensors, int i) -> Tensor[]"` As an example `out5, out6 = torch.ops.fb.fn_with_list_output([out3, out4], 1)` got compiled into ``` AtenTensorHandle buf8_handle; // output buffer AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_new_uninitialized_tensor(&buf8_handle)); RAIIAtenTensorHandle buf8(buf8_handle); AtenTensorHandle buf9_handle; // output buffer AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_new_uninitialized_tensor(&buf9_handle)); RAIIAtenTensorHandle buf9(buf9_handle); AtenTensorHandle tensor_args_var_5[] = {buf5.get(), buf6.get(), buf8.get(), buf9.get()}; int64_t int_args_var_6[] = {1}; aoti_torch_proxy_executor_call_function(proxy_executor, 2, 1, int_args_var_6, 4, tensor_args_var_5); ``` Test Plan: Test Differential Revision: D49694691 fbshipit-source-id: 4be9fe4c4786f7099710e8cbe4ce01cd5a3d70b8
SherlockNoMad
force-pushed
the
export-D49694691
branch
from
69d2aec
to
b44b4cd
Compare
…h#110182) Summary: Support custom ops returns List[Tensor] type, like `"fn_with_list_output(Tensor[] tensors, int i) -> Tensor[]"` As an example `out5, out6 = torch.ops.fb.fn_with_list_output([out3, out4], 1)` got compiled into ``` AtenTensorHandle buf8_handle; // output buffer AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_new_uninitialized_tensor(&buf8_handle)); RAIIAtenTensorHandle buf8(buf8_handle); AtenTensorHandle buf9_handle; // output buffer AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_new_uninitialized_tensor(&buf9_handle)); RAIIAtenTensorHandle buf9(buf9_handle); AtenTensorHandle tensor_args_var_5[] = {buf5.get(), buf6.get(), buf8.get(), buf9.get()}; int64_t int_args_var_6[] = {1}; aoti_torch_proxy_executor_call_function(proxy_executor, 2, 1, int_args_var_6, 4, tensor_args_var_5); ``` Test Plan: Test Differential Revision: D49694691
SherlockNoMad
force-pushed
the
export-D49694691
branch
from
b44b4cd
to
7f77a31
Compare
|
This pull request was exported from Phabricator. Differential Revision: D49694691 |
| output_arguments = [ | ||
| export_schema.Argument.create( | ||
| as_tensor=export_schema.TensorArgument(name=output.get_name()) | ||
| ) | ||
| for output in self.outputs | ||
| ] | ||
| elif isinstance(self.outputs, list): |
There was a problem hiding this comment.
Wondering why we are differentiating list from tuple above. Particularly, why do we pass as_tensors for list but as_tensor for tuple?
There was a problem hiding this comment.
This is by design of the PT2 IR.
TensorList is a first class data type, whereas tuple is used only when there are multiple returns.
Consider following two cases foo(...) -> (Tensor, Tensor) and bar(...) -> Tensor[]
In the first case, self.output is a python tuple; The output would be serialized as [Arguemnt(asTensor="buf3"), Argument(asTensor = "buf4")]
In the second case, self.output is a python list. The output would be serialized as [Argument(asTensors=["buf3", "buf4"])].
For more details, see description in #110187
|
@pytorchbot merge (Initiating merge automatically since Phabricator Diff has merged) |
Merge startedYour 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 |
Summary:
Support custom ops returns List[Tensor] type, like
"fn_with_list_output(Tensor[] tensors, int i) -> Tensor[]"As an example
out5, out6 = torch.ops.fb.fn_with_list_output([out3, out4], 1)got compiled into
Test Plan: Test
Differential Revision: D49694691
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler