Code motion CompiledFxGraph to a dedicated file #141654

ezyang · 2024-11-27T04:07:51Z

Stack from ghstack (oldest at bottom):

Signed-off-by: Edward Z. Yang [email protected]

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

[ghstack-poisoned]

pytorch-bot · 2024-11-27T04:07:54Z

🔗 Helpful Links

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

📄 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! (2 Unrelated Failures)

As of commit 197cc26 with merge base 4c1f50a ():

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

inductor / cuda12.1-py3.10-gcc9-sm86 / test (inductor_timm, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh) (#141703)
convnext_base
inductor / cuda12.4-py3.10-gcc9-sm86 / test (inductor_timm, 1, 2, linux.g5.4xlarge.nvidia.gpu, unstable) (gh) (#141498)
convnext_base

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

Signed-off-by: Edward Z. Yang <[email protected]> ghstack-source-id: 3bcb4b1 Pull Request resolved: #141654

[ghstack-poisoned]

ezyang · 2024-11-27T05:08:04Z

torch/_inductor/codecache.py

+# python test/inductor/test_codecache.py
+# TestFxGraphCache.test_constant_handling_device_cpu
+# TODO: Why are we monkeypatching it......
+import torch._inductor.output_code as output_code


@masnesral you added the test in question, is there really no way to do it without monkeypatching?

I can get the testing I want more directly from another test in that same file: #141898

ezyang · 2024-11-27T05:08:31Z

torch/_inductor/output_code.py

@@ -0,0 +1,179 @@
+"""


This stuff is new and explains the motivation for code motion

ezyang · 2024-11-27T05:08:48Z

torch/_inductor/output_code.py

+class OutputCode(Protocol):
+    def __call__(self, inputs: Sequence[Any]) -> Any:
+        ...
+


Everything below is code movement only

ezyang · 2024-11-27T05:09:03Z

torch/_inductor/output_code.py

+
+
+def _typecheck_CompiledFxGraph(h: CompiledFxGraph) -> OutputCode:
+    return h


Except this, this is new

Maybe put this in an if TYPE_CHECKING block along w/ a comment that it's just used to verify that CompiledFxGraph implements the OutputCode protocol.

Oh, there's an interesting idea. I'll do it in follow up (there are a few other places I did this pattern too)

[ghstack-poisoned]

Signed-off-by: Edward Z. Yang <[email protected]> ghstack-source-id: c5c09cb Pull Request resolved: #141654

aorenste · 2024-11-27T15:28:39Z

torch/_inductor/output_code.py

+_StrideExprStr: TypeAlias = str
+
+
+def has_frozen_params(gm: torch.fx.GraphModule) -> bool:


Why isn't this just a method on GraphModule?

It probably shouldn't live straight on GraphModule, because it's a very Inductor specific concept. But the real reason is I was just shoving code around lol.

aorenste · 2024-11-27T15:39:02Z

torch/_inductor/output_code.py

+
+
+def _typecheck_CompiledFxGraph(h: CompiledFxGraph) -> OutputCode:
+    return h


Maybe put this in an if TYPE_CHECKING block along w/ a comment that it's just used to verify that CompiledFxGraph implements the OutputCode protocol.

jamesjwu · 2024-11-27T15:44:43Z

torch/_inductor/output_code.py

+    from .triton_bundler import TritonKernelArtifacts
+
+
+class OutputCode(Protocol):


What's a protocol?

Structural subtyping (duck typing). You're basically saying "I take/return something that looks like this" without having to use inheritance to force the things to actually be related.

Signed-off-by: Edward Z. Yang <[email protected]> ghstack-source-id: c5c09cb Pull Request resolved: #141654

ezyang · 2024-11-27T18:28:02Z

@pytorchbot merge -i

pytorchmergebot · 2024-11-27T18:30:05Z

Merge started

Your change will be merged while ignoring the following 2 checks: inductor / cuda12.1-py3.10-gcc9-sm86 / test (inductor_timm, 1, 2, linux.g5.4xlarge.nvidia.gpu), inductor / cuda12.4-py3.10-gcc9-sm86 / test (inductor_timm, 1, 2, linux.g5.4xlarge.nvidia.gpu, unstable)

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

Signed-off-by: Edward Z. Yang <[email protected]> Pull Request resolved: pytorch#141654 Approved by: https://github.com/aorenste, https://github.com/jansel ghstack dependencies: pytorch#141491, pytorch#141492, pytorch#141574

Signed-off-by: Edward Z. Yang <[email protected]> ghstack-source-id: c5c09cb Pull Request resolved: pytorch/pytorch#141654

Update

38d2ce8

[ghstack-poisoned]

ezyang mentioned this pull request Nov 27, 2024

Modest code motion in compile_fx #141574

Closed

pytorch-bot bot added ciflow/inductor module: inductor labels Nov 27, 2024

ezyang added a commit that referenced this pull request Nov 27, 2024

Code motion CompiledFxGraph to a dedicated file

78bc785

Signed-off-by: Edward Z. Yang <[email protected]> ghstack-source-id: 3bcb4b1 Pull Request resolved: #141654

github-actions bot requested review from SherlockNoMad, albanD, antoniojkim, bdhirsh and miladm November 27, 2024 04:08

ezyang requested review from jamesjwu, jansel and masnesral November 27, 2024 04:49

ezyang added the topic: not user facing topic category label Nov 27, 2024

Update

a38ef51

[ghstack-poisoned]

ezyang mentioned this pull request Nov 27, 2024

Move post compile steps into post_compile1/post_compile2 method #141656

Closed

ezyang requested a review from aorenste November 27, 2024 05:06

ezyang commented Nov 27, 2024

View reviewed changes

torch/_inductor/output_code.py

@@ -0,0 +1,179 @@

"""

Copy link

Contributor Author

ezyang Nov 27, 2024

Choose a reason for hiding this comment

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

This stuff is new and explains the motivation for code motion

ezyang commented Nov 27, 2024

View reviewed changes

Update

197cc26

[ghstack-poisoned]

pytorch-bot bot added the module: dynamo label Nov 27, 2024

This was referenced Nov 27, 2024

Inline FxGraphCache.load into its sole call site #141681

Closed

Hoist set_feature_use out of conditional, rename some variables #141683

Closed

Unify cache disable and cache bypass paths #141685

Closed

ezyang added a commit that referenced this pull request Nov 27, 2024

Code motion CompiledFxGraph to a dedicated file

dcbb5bc

Signed-off-by: Edward Z. Yang <[email protected]> ghstack-source-id: c5c09cb Pull Request resolved: #141654

ezyang added a commit that referenced this pull request Nov 27, 2024

Code motion CompiledFxGraph to a dedicated file

fba9cb9

Signed-off-by: Edward Z. Yang <[email protected]> ghstack-source-id: c5c09cb Pull Request resolved: #141654

ezyang added a commit that referenced this pull request Nov 27, 2024

Code motion CompiledFxGraph to a dedicated file

861189b

Signed-off-by: Edward Z. Yang <[email protected]> ghstack-source-id: c5c09cb Pull Request resolved: #141654

ezyang added a commit that referenced this pull request Nov 27, 2024

Code motion CompiledFxGraph to a dedicated file

ad524c1

Signed-off-by: Edward Z. Yang <[email protected]> ghstack-source-id: c5c09cb Pull Request resolved: #141654

ezyang added a commit that referenced this pull request Nov 27, 2024

Code motion CompiledFxGraph to a dedicated file

cb77016

Signed-off-by: Edward Z. Yang <[email protected]> ghstack-source-id: c5c09cb Pull Request resolved: #141654

ezyang added a commit that referenced this pull request Nov 27, 2024

Code motion CompiledFxGraph to a dedicated file

754286f

Signed-off-by: Edward Z. Yang <[email protected]> ghstack-source-id: c5c09cb Pull Request resolved: #141654

aorenste approved these changes Nov 27, 2024

View reviewed changes

jamesjwu reviewed Nov 27, 2024

View reviewed changes

ezyang added a commit that referenced this pull request Nov 27, 2024

Code motion CompiledFxGraph to a dedicated file

f8827ee

Signed-off-by: Edward Z. Yang <[email protected]> ghstack-source-id: c5c09cb Pull Request resolved: #141654

jansel approved these changes Nov 27, 2024

View reviewed changes

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Nov 27, 2024

pytorchmergebot added the merging label Nov 27, 2024

pytorchmergebot added the Merged label Nov 27, 2024

pytorchmergebot closed this in dbbebee Nov 27, 2024

pytorchmergebot removed the merging label Nov 27, 2024

Esquains pushed a commit to Esquains/study1 that referenced this pull request Dec 15, 2024

Code motion CompiledFxGraph to a dedicated file

3eeba0e

Signed-off-by: Edward Z. Yang <[email protected]> ghstack-source-id: c5c09cb Pull Request resolved: pytorch/pytorch#141654

github-actions bot deleted the gh/ezyang/3021/head branch January 2, 2025 02:04



		def _typecheck_CompiledFxGraph(h: CompiledFxGraph) -> OutputCode:
		return h

		_StrideExprStr: TypeAlias = str


		def has_frozen_params(gm: torch.fx.GraphModule) -> bool:

		from .triton_bundler import TritonKernelArtifacts


		class OutputCode(Protocol):

Code motion CompiledFxGraph to a dedicated file #141654

Code motion CompiledFxGraph to a dedicated file #141654

Uh oh!

Conversation

ezyang commented Nov 27, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Nov 27, 2024 • 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/141654

✅ You can merge normally! (2 Unrelated Failures)

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!

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!

ezyang commented Nov 27, 2024

Uh oh!

pytorchmergebot commented Nov 27, 2024

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

ezyang commented Nov 27, 2024 •

edited

Loading

pytorch-bot bot commented Nov 27, 2024 •

edited

Loading