[inductor] Refactor MutableBox to make IRNode typing easier #140895

jansel · 2024-11-17T03:04:56Z

Stack from ghstack (oldest at bottom):

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-17T03:05:00Z

🔗 Helpful Links

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

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

[DomainsOnly] Jobs fail with GLIBC version not found

✅ No Failures

As of commit 85dbca2 with merge base 260d1dc ():
💚 Looks good so far! There are no failures yet. 💚

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

[ghstack-poisoned]

ghstack-source-id: f288bdd Pull Request resolved: #140895

[ghstack-poisoned]

jansel · 2024-11-18T03:41:22Z

torch/_inductor/ir.py

-    def __getattr__(self, name):  # type: ignore[no-untyped-def]
-        fn = getattr(self.data, name)
-        if callable(fn):
-            return fn
-        raise AttributeError(f"{type(self.data).__name__}.{name} not callable")


@ezyang @aorenste I believe this function was the main reason typing ir.py was hard. This little trick does save ~100 lines of boilerplate, but once I got rid of this (and a few hasattr -> isinstance changes in this PR) typing the rest of things seem straightforward.

ezyang

Thankssss

ezyang · 2024-11-18T04:30:33Z

torch/_inductor/dependencies.py

-        size = input_node.get_size()
-        reduction_size = input_node.get_reduction_size()
+        size = [*input_node.get_size()]
+        reduction_size = [*input_node.get_reduction_size()]


Curious, is this equivalent to list(...)?

Yeah curious what the motivation here is for this. It should actually be marginally less efficient than the ctor list method, shouldn't it?

It is faster (for small lists) and looks cleaner IMO:

>>> timeit.timeit("list(range(10))") 0.12816928839311004 >>> timeit.timeit("[*range(10)]") 0.11478627705946565 >>>

list gets mapped to a LOAD_GLOBAL bytecode (since you could do list = tuple in global scope), while [*x] gets mapped to list-specific bytecodes without touching global scope.

@ezyang @Skylion007 maybe we should add a rule to automate this change?

ezyang · 2024-11-18T04:30:46Z

torch/_inductor/dependencies.py

                continue
            op = buffer.get_defining_op()
-            if op is None:
+            if op is None or isinstance(op, ExternKernel):


ezyang · 2024-11-18T04:37:30Z

torch/_inductor/lowering.py

 ):
-    def inner(*inputs: List[TensorBox], alpha=None):
-        if triton_fallback is not None and any(map(is_triton, inputs)):
+    def inner(*inputs: TensorBox, alpha=None):


ezyang · 2024-11-18T04:37:59Z

torch/_inductor/lowering.py

-        if triton_fallback is not None and any(map(is_triton, inputs)):
+    def inner(*inputs: TensorBox, alpha=None):
+        if triton_fallback is not None and any(
+            isinstance(inp, IRNode) and is_triton(inp) for inp in inputs


Huh, when is inp not an IRNode?

Usually it is an int/float/None/etc.

ezyang · 2024-11-18T04:41:27Z

torch/_inductor/scheduler.py

    def codegen(self) -> None:
-        assert isinstance(self.node, ir.ComputedBuffer), f"{type(self.node)=}"
-        self.node.get_store_function()(self.node.make_loader()())
+        raise NotImplementedError


Was unused code (with incorrect num args to make_loader())

torch/_inductor/ir.py

ezyang · 2024-11-18T04:43:01Z

torch/_inductor/ir.py

+    elif isinstance(x, IRNode):
+        try:
+            return x.get_device().type
+        except NoDevice:


ezyang · 2024-11-18T04:44:02Z

torch/_inductor/ir.py

+        raise NoDevice(f"{self.__class__.__name__} does not implement get_device()")
+
+    def has_exceeded_max_reads(self) -> bool:
+        raise NotImplementedError(type(self).__name__)


I don't think commit should change this, but would it be better for these to return NotImplemented instead, so that a catch isn't required from call sites? Not sure...

Yeah, I think we could also move some more stuff to the base class (and other stuff down to only on subclasses) -- but that would be for another PR.

ezyang · 2024-11-18T04:45:38Z

torch/_inductor/ir.py

-    def __getattr__(self, name):  # type: ignore[no-untyped-def]
-        fn = getattr(self.data, name)
-        if callable(fn):
-            return fn
-        raise AttributeError(f"{type(self.data).__name__}.{name} not callable")


[ghstack-poisoned]

jansel · 2024-11-19T00:16:41Z

@pytorchbot merge

pytorchmergebot · 2024-11-19T00:18:31Z

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

[ghstack-poisoned]

jansel · 2024-11-20T19:43:18Z

@pytorchbot merge

pytorchmergebot · 2024-11-20T19:44:57Z

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

This separate the concepts of a Layout (size/stride/etc) and an OutputSpec (which includes multiple outputs). Which should make typing easier. Pull Request resolved: #140910 Approved by: https://github.com/ezyang ghstack dependencies: #140895

Pull Request resolved: #140912 Approved by: https://github.com/aorenste ghstack dependencies: #140895, #140910

…140895) Pull Request resolved: pytorch#140895 Approved by: https://github.com/ezyang, https://github.com/Skylion007

…ytorch#140895)" This reverts commit c79e78b. Reverted pytorch#140895 on behalf of https://github.com/huydhn due to Sorry for reverting your change but I think test_torchbind_inductor is failing in trunk after this lands ([comment](pytorch#140895 (comment)))

…140895) Pull Request resolved: pytorch#140895 Approved by: https://github.com/ezyang, https://github.com/Skylion007

This separate the concepts of a Layout (size/stride/etc) and an OutputSpec (which includes multiple outputs). Which should make typing easier. Pull Request resolved: pytorch#140910 Approved by: https://github.com/ezyang ghstack dependencies: pytorch#140895

Pull Request resolved: pytorch#140912 Approved by: https://github.com/aorenste ghstack dependencies: pytorch#140895, pytorch#140910

Update

279c272

[ghstack-poisoned]

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

Update

89f86ed

[ghstack-poisoned]

Update

f367659

[ghstack-poisoned]

jansel added a commit that referenced this pull request Nov 17, 2024

[inductor] Refactor MutableBox to make IRNode typing easier

8d65e22

ghstack-source-id: f288bdd Pull Request resolved: #140895

jansel added the topic: not user facing topic category label Nov 17, 2024

Update

e9100f5

[ghstack-poisoned]

Update

dcc2e46

[ghstack-poisoned]

Update

269279e

[ghstack-poisoned]

This was referenced Nov 17, 2024

[inductor] Refactor ir.Layout into ir.OutputSpec #140910

Closed

[inductor] Add typing to ir.py 1 #140912

Closed

Update

1b15f9a

[ghstack-poisoned]

jansel mentioned this pull request Nov 17, 2024

[inductor] Add typing to ir.py 2 #140915

Closed

jansel marked this pull request as ready for review November 18, 2024 03:35

jansel requested review from aorenste and ezyang November 18, 2024 03:37

jansel commented Nov 18, 2024

View reviewed changes

ezyang approved these changes Nov 18, 2024

View reviewed changes

Skylion007 approved these changes Nov 18, 2024

View reviewed changes

Update

ff3c8a9

[ghstack-poisoned]

jansel added the ciflow/trunk Trigger trunk jobs on your pull request label Nov 18, 2024

pytorchmergebot added the merging label Nov 19, 2024

pytorchmergebot added Reverted ci-no-td Do not run TD on this PR labels Nov 19, 2024

pytorchmergebot reopened this Nov 19, 2024

jansel added 3 commits November 18, 2024 21:19

Update

249bc6e

[ghstack-poisoned]

Update

1704b4c

[ghstack-poisoned]

Update

f245329

[ghstack-poisoned]

Update

85dbca2

[ghstack-poisoned]

pytorchmergebot added the merging label Nov 20, 2024

pytorchmergebot closed this in 808f0f6 Nov 20, 2024

pytorchmergebot removed the merging label Nov 20, 2024

github-actions bot deleted the gh/jansel/436/head branch December 21, 2024 02:05

[inductor] Refactor MutableBox to make IRNode typing easier #140895

[inductor] Refactor MutableBox to make IRNode typing easier #140895

Uh oh!

Conversation

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

Uh oh!

Uh oh!

pytorch-bot bot commented Nov 17, 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/140895

❗ 1 Active SEVs

✅ No Failures

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ezyang 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!

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

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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!

jansel commented Nov 19, 2024

Uh oh!

pytorchmergebot commented Nov 19, 2024

Merge started

Uh oh!

jansel commented Nov 20, 2024

Uh oh!

pytorchmergebot commented Nov 20, 2024

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

jansel commented Nov 17, 2024 •

edited

Loading

pytorch-bot bot commented Nov 17, 2024 •

edited

Loading

jansel Nov 18, 2024 •

edited

Loading

jansel Nov 18, 2024 •

edited

Loading