[inductor] Support autotune restore_value for user-defined Triton kernels by aakhundov · Pull Request #139851 · pytorch/pytorch

aakhundov · 2024-11-06T04:53:44Z

Stack from ghstack (oldest at bottom):

-> [inductor] Support autotune restore_value for user-defined Triton kernels #139851

This PR adds support for the restore_value argument of the
@triton.autotune for the user-defined Triton kernels in PT2.

The kernel.restore_idx are extracted in the
ir.UserDefinedTritonKernel and the corresponding arg names are
placed into the triton_meta["restore_value"]. From there, those
are added to the existing mutated_arg_names in the caching autotuner
infra which already exists and leads to the listed argss being cloned.
This achieves the equivalent effect to the native restore_value.

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

[ghstack-poisoned]

pytorch-bot · 2024-11-06T04:53:47Z

🔗 Helpful Links

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

📄 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 Cancelled Job

As of commit 8191aa2 with merge base 157c18a ():

CANCELLED JOB - The following job was cancelled. Please retry:

trunk / macos-py3-arm64 / build (gh)
##[error]The operation was canceled.

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

[ghstack-poisoned]

This PR adds support for the `restore_value` argument of the `triton.autotune` for the user-defined Triton kernels in PT2. The `kernel.restore_idx` are extracted in the `ir.UserDefinedTritonKernel` and the corresponding arg names are placed into the `triton_meta["restore_value"]`. From there, those are added to the existing `mutated_arg_names` in the caching autotuner infra which already exists and leads to the listed argss being cloned. This achieves the equivalent effect to the native `restore_value`. ghstack-source-id: a027031 Pull Request resolved: #139851

[ghstack-poisoned]

This PR adds support for the `restore_value` argument of the `triton.autotune` for the user-defined Triton kernels in PT2. The `kernel.restore_idx` are extracted in the `ir.UserDefinedTritonKernel` and the corresponding arg names are placed into the `triton_meta["restore_value"]`. From there, those are added to the existing `mutated_arg_names` in the caching autotuner infra which already exists and leads to the listed argss being cloned. This achieves the equivalent effect to the native `restore_value`. ghstack-source-id: 03ea337 Pull Request resolved: #139851

oulgen · 2024-11-07T06:13:39Z

torch/_higher_order_ops/triton_kernel_wrap.py

+    # move as many positional arguments from dicts to args as we
+    # can to circumvent the bug with the kwargs and pre_/post_hook:
+    # https://github.com/triton-lang/triton/issues/5082
+    # TODO: remove this when the Triton issue above is fixed
+    args = []


this feels really bad, and confusing. Maybe only do this when pre/post hook is used?

At this stage of the Triton code, we can't tell when pre/post hook is used. Just submitted a PR so that we'll be able to tell: triton-lang/triton#5092. Can update then in a follow-up. Although, there is another PR fixing the mentioned issue that should make this whole thing unnecessary: triton-lang/triton#5083.

Could you elaborate why it's bad, btw? Is this incorrect? I'm just moving the args around, but not skipping or duplicating any passed to the kernel? It may not always help the issue, but in most cases it should (i.e., when the args used in the hook are not coming from @triton.heuristic or smth).

oulgen · 2024-11-07T06:15:21Z

torch/_higher_order_ops/triton_kernel_wrap.py

+    kwargs = kwargs.copy()
+    constant_args = constant_args.copy()
+    for name in kernel.arg_names:
+        if name in kwargs:
+            args.append(kwargs.pop(name))
+        elif name in constant_args:
+            args.append(constant_args.pop(name))
+        else:
+            break


you can write this for loop better, you dont need copy/pop, just read indices?

I need to pop, because I'm passing the kwargs and constant_args together with the collected args, no? (That's also why I'm copying kwargs and constant_args, as I mutate the dicts). Otherwise, there will be duplicate argument passing in the line following this?

[ghstack-poisoned]

This PR adds support for the `restore_value` argument of the `triton.autotune` for the user-defined Triton kernels in PT2. The `kernel.restore_idx` are extracted in the `ir.UserDefinedTritonKernel` and the corresponding arg names are placed into the `triton_meta["restore_value"]`. From there, those are added to the existing `mutated_arg_names` in the caching autotuner infra which already exists and leads to the listed argss being cloned. This achieves the equivalent effect to the native `restore_value`. ghstack-source-id: 9f806fb Pull Request resolved: #139851

aakhundov · 2024-11-07T23:05:35Z

@pytorchbot merge

pytorchmergebot · 2024-11-07T23:07:17Z

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

pytorchmergebot · 2024-11-08T05:05:58Z

The merge job was canceled or timed out. This most often happen if two merge requests were issued for the same PR, or if merge job was waiting for more than 6 hours for tests to finish. In later case, please do not hesitate to reissue the merge command
For more information see pytorch-bot wiki.

bertmaher · 2024-11-08T14:34:56Z

@pytorchbot merge

pytorchmergebot · 2024-11-08T14:36:37Z

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

pytorchmergebot · 2024-11-08T14:36:54Z

Merge failed

Reason: 1 jobs have failed, first few of them are: trunk / macos-py3-arm64 / build

Details for Dev Infra team

Raised by workflow job

bertmaher · 2024-11-08T14:57:20Z

@pytorchbot merge -f 'macos-py3-arm64 seems to be throwing spurious cancellations and cant possibly be a failure on this diff'

pytorchmergebot · 2024-11-08T14:58:49Z

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use -f as last resort and instead consider -i/--ignore-current to continue the merge ignoring current failures. This will allow currently pending tests to finish and report signal before the merge.

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

…nels (pytorch#139851) This PR adds support for the `restore_value` argument of the `@triton.autotune` for the user-defined Triton kernels in PT2. The `kernel.restore_idx` are extracted in the `ir.UserDefinedTritonKernel` and the corresponding arg names are placed into the `triton_meta["restore_value"]`. From there, those are added to the existing `mutated_arg_names` in the caching autotuner infra which already exists and leads to the listed argss being cloned. This achieves the equivalent effect to the native `restore_value`. Pull Request resolved: pytorch#139851 Approved by: https://github.com/oulgen

Update

cb4b9fb

[ghstack-poisoned]

aakhundov requested a review from zou3519 as a code owner November 6, 2024 04:53

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

aakhundov requested review from davidberard98 and oulgen November 6, 2024 04:55

aakhundov added the topic: not user facing topic category label Nov 6, 2024

aakhundov changed the title ~~[inductor] Add restore_value to user-defined Triton kernel support~~ [inductor] Support autotune restore_value for user-defined Triton kernels in PT2 Nov 6, 2024

aakhundov changed the title ~~[inductor] Support autotune restore_value for user-defined Triton kernels in PT2~~ [inductor] Support autotune restore_value for user-defined Triton kernels Nov 6, 2024

Update

70586a8

[ghstack-poisoned]

Update

a1c5703

[ghstack-poisoned]

aakhundov requested a review from eellison November 7, 2024 05:31

Update

b9b8b21

[ghstack-poisoned]

oulgen reviewed Nov 7, 2024

View reviewed changes

oulgen approved these changes Nov 7, 2024

View reviewed changes

Update

8191aa2

[ghstack-poisoned]

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

pytorchmergebot added the merging label Nov 7, 2024

pytorchmergebot removed the merging label Nov 8, 2024

pytorchmergebot added the merging label Nov 8, 2024

pytorchmergebot closed this in 838958d Nov 8, 2024

pytorchmergebot added Merged and removed merging labels Nov 8, 2024

github-actions bot deleted the gh/aakhundov/16/head branch December 9, 2024 02:14

Conversation

aakhundov commented Nov 6, 2024 • edited by pytorch-bot bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Nov 6, 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/139851

❌ 1 Cancelled Job

Uh oh!

oulgen Nov 7, 2024

Choose a reason for hiding this comment

Uh oh!

aakhundov Nov 7, 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!

aakhundov Nov 7, 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!

oulgen Nov 7, 2024

Choose a reason for hiding this comment

Uh oh!

aakhundov Nov 7, 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!

aakhundov commented Nov 7, 2024

Uh oh!

pytorchmergebot commented Nov 7, 2024

Merge started

Uh oh!

pytorchmergebot commented Nov 8, 2024

Uh oh!

bertmaher commented Nov 8, 2024

Uh oh!

pytorchmergebot commented Nov 8, 2024

Merge started

Uh oh!

pytorchmergebot commented Nov 8, 2024

Merge failed

Uh oh!

bertmaher commented Nov 8, 2024

Uh oh!

pytorchmergebot commented Nov 8, 2024

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

aakhundov commented Nov 6, 2024 •

edited by pytorch-bot bot

Loading

pytorch-bot bot commented Nov 6, 2024 •

edited

Loading

aakhundov Nov 7, 2024 •

edited

Loading

aakhundov Nov 7, 2024 •

edited

Loading

aakhundov Nov 7, 2024 •

edited

Loading