AOTInductor dynamic shape #109012

chenyang78 · 2023-09-11T08:59:20Z

Summary: This PR adds dynamic-shape support for AOTInductor

On the runtime/interface side, we added two structs, StaticDimInfo
and DynamicDimInfo, to hold values for static and dynamic dimensions,
respectively. Dynamic dimensions are tracked by an unordered map field
defined in AOTInductorModelBase. At inference time, the inference run
method will assign the current real dimensional value to each dynamic
dimension before executing any kernel.
On the CUDA wrapper codegen side, we generate dynamic symbols
appropriately for shape computations. We simulate kernel launch grids
in the C++ land by re-using the grid functions from the Python world.
The returned grid configs, which may contain symbolic expressions,
are printed out in their C++ forms via the CppPrinter. Note that
when dynamic shapes are involved, we have to compute grid configs
for each kernel at runtime in the same way as we do for launching
the corresponding Triton kernel. Otherwise, we may end up with
memory-access failures or mis-computations caused by invalid indices
for fetching or storing data in device memory.

Differential Revision: D49100472

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

pytorch-bot · 2023-09-11T08:59:24Z

🔗 Helpful Links

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

📄 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.

✅ No Failures

As of commit 0f6853b with merge base 025d1a1 ():
💚 Looks good so far! There are no failures yet. 💚

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

facebook-github-bot · 2023-09-11T08:59:38Z

This pull request was exported from Phabricator. Differential Revision: D49100472

facebook-github-bot · 2023-09-11T21:52:20Z

This pull request was exported from Phabricator. Differential Revision: D49100472

Summary: This PR adds dynamic-shape support for AOTInductor Differential Revision: D49100472

facebook-github-bot · 2023-09-11T21:52:36Z

This pull request was exported from Phabricator. Differential Revision: D49100472

facebook-github-bot · 2023-09-12T08:52:18Z

This pull request was exported from Phabricator. Differential Revision: D49100472

Summary: This PR adds dynamic-shape support for AOTInductor Differential Revision: D49100472

facebook-github-bot · 2023-09-12T20:27:43Z

This pull request was exported from Phabricator. Differential Revision: D49100472

facebook-github-bot · 2023-09-12T20:29:02Z

This pull request was exported from Phabricator. Differential Revision: D49100472

Summary: This PR adds dynamic-shape support for AOTInductor Differential Revision: D49100472

ipiszy

Thanks @chenyang78 ! Left some comments, mainly some clarification questions. It would be great if you could help summarize your changes in the PR description.

ipiszy · 2023-09-13T00:29:34Z

torch/_inductor/codegen/wrapper.py


        return ", ".join(new_args)

+    def generate_default_grid(self, name, grid, cuda=True):


Is this cuda arg used anywhere?

Seems unnecessary since the cpu backend doesn't need to generate grid

ipiszy · 2023-09-13T00:32:01Z

torch/_inductor/codegen/wrapper.py

+            return grid
+        assert isinstance(grid, list), f"expected {grid=} to be a list"
+        grid = [e.inner_expr if isinstance(e, SymbolicCallArg) else e for e in grid]
+        grid_fn = default_grid(*grid)


What is default_grid? I cannot find it in the pytorch repo.

It's from from ..triton_heuristics import grid as default_grid

ipiszy · 2023-09-13T00:33:59Z

torch/_inductor/codegen/wrapper.py

        )
        stack.enter_context(self.wrapper_call.indent())

+    def generate_default_grid(self, name, grid_args):


nit: Add more comments to describe what this function does, and the args. Also add type hints.

Will do. Thanks.

ipiszy · 2023-09-13T00:46:59Z

torch/_inductor/codegen/wrapper.py

 @dataclasses.dataclass
 class SymbolicCallArg:
    inner: Any
+    inner_expr: sympy.Expr


Add some comments?

desertfire

Left some nits. LGTM overall.

Currently AOTInductor on the OSS benchmarks are not a part of the CI yet. I have a pending PR, #108419, but it won't be merged soon given the recent CI infra instability. So can you manually trigger a run at https://github.com/pytorch/pytorch/actions/workflows/inductor-perf-test-nightly.yml?

You will need to select the options like in the screenshot,

Unfortunately the Branch selection does not work with a PR from your forked repo. You can create a mirror PR on a pytorch branch in order to do the testing. Let me know if you need more help on the instructions.

desertfire · 2023-09-13T01:54:45Z

test/inductor/test_aot_inductor.py

+        for example_inputs, example_outputs in zip(
+            list_example_inputs, list_example_outputs
+        ):
+            output_tensors = [torch.empty_like(output) for output in example_outputs]


In the prod scenario, do we allocate a list of tensors with the max batch size and reuse those? If yes, we need to test that here.

No, we don't re-use output tensors. Instead, we allocate output tensors for each inference run and return these tensors to the caller of the forward function.

desertfire · 2023-09-13T02:07:20Z

torch/_inductor/codegen/wrapper.py

+            return grid
+        assert isinstance(grid, list), f"expected {grid=} to be a list"
+        grid = [e.inner_expr if isinstance(e, SymbolicCallArg) else e for e in grid]
+        grid_fn = default_grid(*grid)


It's from from ..triton_heuristics import grid as default_grid

desertfire · 2023-09-13T02:09:19Z

torch/_inductor/codegen/wrapper.py


        return ", ".join(new_args)

+    def generate_default_grid(self, name, grid, cuda=True):


Seems unnecessary since the cpu backend doesn't need to generate grid

desertfire · 2023-09-13T02:22:31Z

torch/csrc/inductor/aot_runtime/interface.h

    size_t num_inputs,
    AOTInductorTensorHandle outputs_handle,
    size_t num_outputs,
+    AOTInductorParamShape* output_shapes,


Needs to revisit for ABI compatibility, but ok for now.

facebook-github-bot · 2023-09-13T08:30:13Z

This pull request was exported from Phabricator. Differential Revision: D49100472

Summary: This PR adds dynamic-shape support for AOTInductor Reviewed By: khabinov Differential Revision: D49100472

facebook-github-bot · 2023-09-13T08:31:15Z

This pull request was exported from Phabricator. Differential Revision: D49100472

chenyang78 · 2023-09-13T08:42:34Z

Left some nits. LGTM overall.

Currently AOTInductor on the OSS benchmarks are not a part of the CI yet. I have a pending PR, #108419, but it won't be merged soon given the recent CI infra instability. So can you manually trigger a run at https://github.com/pytorch/pytorch/actions/workflows/inductor-perf-test-nightly.yml?

You will need to select the options like in the screenshot, Unfortunately the Branch selection does not work with a PR from your forked repo. You can create a mirror PR on a pytorch branch in order to do the testing. Let me know if you need more help on the instructions.

Running it: https://github.com/pytorch/pytorch/actions/runs/6170063555

Hopefully I was doing it right. Thanks.

desertfire · 2023-09-13T21:12:54Z

Left some nits. LGTM overall.
Currently AOTInductor on the OSS benchmarks are not a part of the CI yet. I have a pending PR, #108419, but it won't be merged soon given the recent CI infra instability. So can you manually trigger a run at https://github.com/pytorch/pytorch/actions/workflows/inductor-perf-test-nightly.yml?
You will need to select the options like in the screenshot, Unfortunately the Branch selection does not work with a PR from your forked repo. You can create a mirror PR on a pytorch branch in order to do the testing. Let me know if you need more help on the instructions.

Running it: https://github.com/pytorch/pytorch/actions/runs/6170063555

Hopefully I was doing it right. Thanks.

Given the long task queueing time, I ran this PR locally to verify. I haven't checked every model, but the results look good to me. I am ok with landing it.

Summary: This PR adds dynamic-shape support for AOTInductor Reviewed By: frank-wei, khabinov Differential Revision: D49100472

facebook-github-bot · 2023-09-14T01:16:10Z

This pull request was exported from Phabricator. Differential Revision: D49100472

facebook-github-bot · 2023-09-14T01:16:36Z

This pull request was exported from Phabricator. Differential Revision: D49100472

chenyang78 · 2023-09-14T06:00:20Z

@pytorchbot merge

pytorchmergebot · 2023-09-14T06:04:12Z

Merge failed

Reason: This PR has internal changes and must be landed via Phabricator

Details for Dev Infra team

Raised by workflow job

facebook-github-bot · 2023-09-14T07:58:44Z

@pytorchbot merge

(Initiating merge automatically since Phabricator Diff has merged)

pytorchmergebot · 2023-09-14T08:00:23Z

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

pytorch-bot bot added the release notes: foreach_frontend release notes category label Sep 11, 2023

github-actions bot added module: inductor ciflow/inductor labels Sep 11, 2023

facebook-github-bot added the fb-exported label Sep 11, 2023

chenyang78 requested review from desertfire and jansel September 11, 2023 09:17

davidradl mentioned this pull request Sep 11, 2023

pr build failures in inductor dynamic shape test for operation tests with simple tensors. Side effect of current test framework #109016

Closed

chenyang78 added the topic: not user facing topic category label Sep 11, 2023

chenyang78 force-pushed the export-D49100472 branch from bb538c0 to e2083bc Compare September 11, 2023 21:52

chenyang78 added a commit to chenyang78/pytorch that referenced this pull request Sep 11, 2023

AOTInductor dynamic shape (pytorch#109012)

bfb55cf

Summary: This PR adds dynamic-shape support for AOTInductor Differential Revision: D49100472

chenyang78 force-pushed the export-D49100472 branch from e2083bc to bfb55cf Compare September 11, 2023 21:52

chenyang78 force-pushed the export-D49100472 branch from bfb55cf to 64287a7 Compare September 12, 2023 08:52

chenyang78 added a commit to chenyang78/pytorch that referenced this pull request Sep 12, 2023

AOTInductor dynamic shape (pytorch#109012)

64287a7

Summary: This PR adds dynamic-shape support for AOTInductor Differential Revision: D49100472

chenyang78 force-pushed the export-D49100472 branch from 64287a7 to 1003ba8 Compare September 12, 2023 20:27

chenyang78 added a commit to chenyang78/pytorch that referenced this pull request Sep 12, 2023

AOTInductor dynamic shape (pytorch#109012)

89223bb

Summary: This PR adds dynamic-shape support for AOTInductor Differential Revision: D49100472

chenyang78 force-pushed the export-D49100472 branch from 1003ba8 to 89223bb Compare September 12, 2023 20:29

khabinov approved these changes Sep 12, 2023

View reviewed changes

ipiszy reviewed Sep 13, 2023

View reviewed changes

desertfire approved these changes Sep 13, 2023

View reviewed changes

chenyang78 force-pushed the export-D49100472 branch from 89223bb to a7158ba Compare September 13, 2023 08:29

chenyang78 added a commit to chenyang78/pytorch that referenced this pull request Sep 13, 2023

AOTInductor dynamic shape (pytorch#109012)

9cf331c

Summary: This PR adds dynamic-shape support for AOTInductor Reviewed By: khabinov Differential Revision: D49100472

chenyang78 force-pushed the export-D49100472 branch from a7158ba to 9cf331c Compare September 13, 2023 08:31

AOTInductor dynamic shape (pytorch#109012)

0f6853b

Summary: This PR adds dynamic-shape support for AOTInductor Reviewed By: frank-wei, khabinov Differential Revision: D49100472

chenyang78 force-pushed the export-D49100472 branch from 9cf331c to 768090f Compare September 14, 2023 01:15

chenyang78 force-pushed the export-D49100472 branch from 768090f to 0f6853b Compare September 14, 2023 01:16

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Sep 14, 2023

pytorchmergebot added the merging label Sep 14, 2023

pytorchmergebot removed the merging label Sep 14, 2023

hl475 approved these changes Sep 14, 2023

View reviewed changes

pytorchmergebot added the merging label Sep 14, 2023

pytorchmergebot added Merged and removed merging labels Sep 14, 2023

pytorchmergebot closed this in 9cd4548 Sep 14, 2023


		return ", ".join(new_args)

		def generate_default_grid(self, name, grid, cuda=True):

AOTInductor dynamic shape #109012

AOTInductor dynamic shape #109012

Uh oh!

Conversation

chenyang78 commented Sep 11, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Sep 11, 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/109012

✅ No Failures

Uh oh!

facebook-github-bot commented Sep 11, 2023

Uh oh!

facebook-github-bot commented Sep 11, 2023

Uh oh!

facebook-github-bot commented Sep 11, 2023

Uh oh!

facebook-github-bot commented Sep 12, 2023

Uh oh!

facebook-github-bot commented Sep 12, 2023

Uh oh!

facebook-github-bot commented Sep 12, 2023

Uh oh!

ipiszy left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

desertfire left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot commented Sep 13, 2023

Uh oh!

facebook-github-bot commented Sep 13, 2023

Uh oh!

chenyang78 commented Sep 13, 2023

Uh oh!

desertfire commented Sep 13, 2023

Uh oh!

facebook-github-bot commented Sep 14, 2023

Uh oh!

facebook-github-bot commented Sep 14, 2023

Uh oh!

chenyang78 commented Sep 14, 2023

Uh oh!

pytorchmergebot commented Sep 14, 2023

Merge failed

Uh oh!

facebook-github-bot commented Sep 14, 2023

Uh oh!

pytorchmergebot commented Sep 14, 2023

Merge started

Uh oh!

Reviewers

chenyang78 commented Sep 11, 2023 •

edited

Loading

pytorch-bot bot commented Sep 11, 2023 •

edited

Loading