Don't add non-integer Triton kernel arg 1 to equal_to_1 #123886

aakhundov · 2024-04-12T00:50:15Z

Stack from ghstack (oldest at bottom):

Summary: Triton compiler adds constnat argument 1 to equal_to_1 only when it's an int. Here we restrict Inductor's equal_to_1 in the same way.

Test Plan:

$ python test/inductor/test_triton_kernels.py -k test_triton_kernel_equal_to_1_float_arg
...
----------------------------------------------------------------------
Ran 1 test in 6.528s

OK

$ python test/inductor/test_triton_kernels.py -k test_triton_kernel_equal_to_1_arg
...
----------------------------------------------------------------------
Ran 2 tests in 10.142s

OK

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

Summary: Triton compiler adds constnat argument 1 to `equal_to_1` [only when it's an int](https://github.com/openai/triton/blob/8c5e33c77ef83e0cb99c744e58842930e602df31/python/triton/runtime/jit.py#L275). Here we restrict Inductor's `equal_to_1` in the same way. Test Plan: ``` $ python test/inductor/test_triton_kernels.py -k test_triton_kernel_equal_to_1_float_arg ... ---------------------------------------------------------------------- Ran 1 test in 6.528s OK $ python test/inductor/test_triton_kernels.py -k test_triton_kernel_equal_to_1_arg ... ---------------------------------------------------------------------- Ran 2 tests in 10.142s OK ``` [ghstack-poisoned]

pytorch-bot · 2024-04-12T00:50:20Z

🔗 Helpful Links

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

📄 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, 1 Pending

As of commit da76a63 with merge base 4f244cf ():
💚 Looks good so far! There are no failures yet. 💚

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

Summary: Triton compiler adds constnat argument 1 to `equal_to_1` [only when it's an int](https://github.com/openai/triton/blob/8c5e33c77ef83e0cb99c744e58842930e602df31/python/triton/runtime/jit.py#L275). Here we restrict Inductor's `equal_to_1` in the same way. Test Plan: ``` $ python test/inductor/test_triton_kernels.py -k test_triton_kernel_equal_to_1_float_arg ... ---------------------------------------------------------------------- Ran 1 test in 6.528s OK $ python test/inductor/test_triton_kernels.py -k test_triton_kernel_equal_to_1_arg ... ---------------------------------------------------------------------- Ran 2 tests in 10.142s OK ``` ghstack-source-id: d02b427 Pull Request resolved: #123886

oulgen · 2024-04-12T00:55:07Z

torch/_inductor/codegen/triton_utils.py

        for i, arg in zip(indices, args)
        if isinstance(arg, SizeArg)
        and arg.expr is not None
+        and isinstance(arg.expr, (int, sympy.Expr))


Why did you also include sympy.Expr? not symint?

unit test does not cover this case

If I don't include sympy.Expr, symbols statically know to be equal to 1 won't be considered. I assume that the next line

and V.graph.sizevars.statically_known_equals(arg.expr, 1)

will filter out non-1 sympy.Exprs (including 1.0). @ezyang could you confirm if this is a valid assumption? Or should I replace with sympy.Integer?

So, I agree that we won't have SymInts at this point in time, in Inductor generally you're only ever dealing with int/sympy.Expr

That being said, I don't really know what the type invariants on SizeArg are, so I don't really know how to review this. Maybe I'll be satisfied with a comment here explaining the situation.

Added a comment.

@ezyang does this look good to you now? Thanks!

Summary: Triton compiler adds constnat argument 1 to `equal_to_1` [only when it's an int](https://github.com/openai/triton/blob/8c5e33c77ef83e0cb99c744e58842930e602df31/python/triton/runtime/jit.py#L275). Here we restrict Inductor's `equal_to_1` in the same way. Test Plan: ``` $ python test/inductor/test_triton_kernels.py -k test_triton_kernel_equal_to_1_float_arg ... ---------------------------------------------------------------------- Ran 1 test in 6.528s OK $ python test/inductor/test_triton_kernels.py -k test_triton_kernel_equal_to_1_arg ... ---------------------------------------------------------------------- Ran 2 tests in 10.142s OK ``` cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler amjames desertfire chauhang [ghstack-poisoned]

Summary: Triton compiler adds constnat argument 1 to `equal_to_1` [only when it's an int](https://github.com/openai/triton/blob/8c5e33c77ef83e0cb99c744e58842930e602df31/python/triton/runtime/jit.py#L275). Here we restrict Inductor's `equal_to_1` in the same way. Test Plan: ``` $ python test/inductor/test_triton_kernels.py -k test_triton_kernel_equal_to_1_float_arg ... ---------------------------------------------------------------------- Ran 1 test in 6.528s OK $ python test/inductor/test_triton_kernels.py -k test_triton_kernel_equal_to_1_arg ... ---------------------------------------------------------------------- Ran 2 tests in 10.142s OK ``` ghstack-source-id: ca041c8 Pull Request resolved: #123886

Summary: Triton compiler adds constnat argument 1 to `equal_to_1` [only when it's an int](https://github.com/openai/triton/blob/8c5e33c77ef83e0cb99c744e58842930e602df31/python/triton/runtime/jit.py#L275). Here we restrict Inductor's `equal_to_1` in the same way. Test Plan: ``` $ python test/inductor/test_triton_kernels.py -k test_triton_kernel_equal_to_1_float_arg ... ---------------------------------------------------------------------- Ran 1 test in 6.528s OK $ python test/inductor/test_triton_kernels.py -k test_triton_kernel_equal_to_1_arg ... ---------------------------------------------------------------------- Ran 2 tests in 10.142s OK ``` cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler amjames desertfire chauhang [ghstack-poisoned]

Summary: Triton compiler adds constnat argument 1 to `equal_to_1` [only when it's an int](https://github.com/openai/triton/blob/8c5e33c77ef83e0cb99c744e58842930e602df31/python/triton/runtime/jit.py#L275). Here we restrict Inductor's `equal_to_1` in the same way. Test Plan: ``` $ python test/inductor/test_triton_kernels.py -k test_triton_kernel_equal_to_1_float_arg ... ---------------------------------------------------------------------- Ran 1 test in 6.528s OK $ python test/inductor/test_triton_kernels.py -k test_triton_kernel_equal_to_1_arg ... ---------------------------------------------------------------------- Ran 2 tests in 10.142s OK ``` ghstack-source-id: a0ab100 Pull Request resolved: #123886

Summary: Triton compiler adds constnat argument 1 to `equal_to_1` [only when it's an int](https://github.com/openai/triton/blob/8c5e33c77ef83e0cb99c744e58842930e602df31/python/triton/runtime/jit.py#L275). Here we restrict Inductor's `equal_to_1` in the same way. Test Plan: ``` $ python test/inductor/test_triton_kernels.py -k test_triton_kernel_equal_to_1_float_arg ... ---------------------------------------------------------------------- Ran 1 test in 6.528s OK $ python test/inductor/test_triton_kernels.py -k test_triton_kernel_equal_to_1_arg ... ---------------------------------------------------------------------- Ran 2 tests in 10.142s OK ``` cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler amjames desertfire chauhang [ghstack-poisoned]

aakhundov · 2024-04-13T02:31:22Z

We discussed this offline with @oulgen and, due to statically_known_equals claiming sympy.Float(1.0) == sympy.Integer(1), decided to restrict possible types to sympy.Integer instead of the broader sympy.Expr. Unit test is added to verify that sympy.Float(1.0) is not added to equal_to_1.

torch/_inductor/codegen/wrapper.py

torch/_inductor/codegen/triton_utils.py

test/inductor/test_aot_inductor.py

Summary: Triton compiler adds constnat argument 1 to `equal_to_1` [only when it's an int](https://github.com/openai/triton/blob/8c5e33c77ef83e0cb99c744e58842930e602df31/python/triton/runtime/jit.py#L275). Here we restrict Inductor's `equal_to_1` in the same way. Test Plan: ``` $ python test/inductor/test_triton_kernels.py -k test_triton_kernel_equal_to_1_float_arg ... ---------------------------------------------------------------------- Ran 1 test in 6.528s OK $ python test/inductor/test_triton_kernels.py -k test_triton_kernel_equal_to_1_arg ... ---------------------------------------------------------------------- Ran 2 tests in 10.142s OK ``` cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler amjames desertfire chauhang [ghstack-poisoned]

Summary: Triton compiler adds constnat argument 1 to `equal_to_1` [only when it's an int](https://github.com/openai/triton/blob/8c5e33c77ef83e0cb99c744e58842930e602df31/python/triton/runtime/jit.py#L275). Here we restrict Inductor's `equal_to_1` in the same way. Test Plan: ``` $ python test/inductor/test_triton_kernels.py -k test_triton_kernel_equal_to_1 ... ---------------------------------------------------------------------- Ran 4 tests in 16.044s OK $ python test/inductor/test_aot_inductor.py -k test_triton_kernel_equal_to_1 ... ---------------------------------------------------------------------- Ran 12 tests in 71.502s OK (skipped=8) ``` ghstack-source-id: db374f9 Pull Request resolved: #123886

aakhundov · 2024-04-14T20:32:08Z

@pytorchbot merge

pytorchmergebot · 2024-04-14T20:33:51Z

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

Summary: Triton compiler adds constnat argument 1 to `equal_to_1` [only when it's an int](https://github.com/openai/triton/blob/8c5e33c77ef83e0cb99c744e58842930e602df31/python/triton/runtime/jit.py#L275). Here we restrict Inductor's `equal_to_1` in the same way. Test Plan: ``` $ python test/inductor/test_triton_kernels.py -k test_triton_kernel_equal_to_1_float_arg ... ---------------------------------------------------------------------- Ran 1 test in 6.528s OK $ python test/inductor/test_triton_kernels.py -k test_triton_kernel_equal_to_1_arg ... ---------------------------------------------------------------------- Ran 2 tests in 10.142s OK ``` Pull Request resolved: pytorch#123886 Approved by: https://github.com/oulgen ghstack dependencies: pytorch#123703

pytorch-bot bot added ciflow/inductor module: inductor labels Apr 12, 2024

aakhundov requested review from ezyang and oulgen April 12, 2024 00:51

oulgen reviewed Apr 12, 2024

View reviewed changes

aakhundov added topic: not user facing topic category ciflow/trunk Trigger trunk jobs on your pull request labels Apr 12, 2024

aakhundov requested a review from oulgen April 13, 2024 02:30

oulgen approved these changes Apr 13, 2024

View reviewed changes

torch/_inductor/codegen/wrapper.py Outdated Show resolved Hide resolved

torch/_inductor/codegen/triton_utils.py Outdated Show resolved Hide resolved

test/inductor/test_aot_inductor.py Outdated Show resolved Hide resolved

pytorchmergebot added the merging label Apr 14, 2024

pytorchmergebot closed this in 03a05e7 Apr 14, 2024

pytorchmergebot added Merged and removed merging labels Apr 14, 2024

github-actions bot deleted the gh/aakhundov/24/head branch May 15, 2024 01:54

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Don't add non-integer Triton kernel arg 1 to equal_to_1 #123886

Don't add non-integer Triton kernel arg 1 to equal_to_1 #123886

Uh oh!

aakhundov commented Apr 12, 2024 •

edited by pytorch-bot bot

Loading

Uh oh!

pytorch-bot bot commented Apr 12, 2024 •

edited

Loading

Uh oh!

oulgen Apr 12, 2024

Uh oh!

aakhundov Apr 12, 2024

Uh oh!

ezyang Apr 12, 2024

Uh oh!

aakhundov Apr 12, 2024

Uh oh!

aakhundov Apr 12, 2024

Uh oh!

aakhundov commented Apr 13, 2024 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

aakhundov commented Apr 14, 2024

Uh oh!

pytorchmergebot commented Apr 14, 2024

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Don't add non-integer Triton kernel arg 1 to equal_to_1 #123886

Don't add non-integer Triton kernel arg 1 to equal_to_1 #123886

Uh oh!

Conversation

aakhundov commented Apr 12, 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 Apr 12, 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/123886

⏳ No Failures, 1 Pending

Uh oh!

oulgen Apr 12, 2024

Choose a reason for hiding this comment

Uh oh!

aakhundov Apr 12, 2024

Choose a reason for hiding this comment

Uh oh!

ezyang Apr 12, 2024

Choose a reason for hiding this comment

Uh oh!

aakhundov Apr 12, 2024

Choose a reason for hiding this comment

Uh oh!

aakhundov Apr 12, 2024

Choose a reason for hiding this comment

Uh oh!

aakhundov commented Apr 13, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

aakhundov commented Apr 14, 2024

Uh oh!

pytorchmergebot commented Apr 14, 2024

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

aakhundov commented Apr 12, 2024 •

edited by pytorch-bot bot

Loading

pytorch-bot bot commented Apr 12, 2024 •

edited

Loading

aakhundov commented Apr 13, 2024 •

edited

Loading