[inductor] Pass None and skip constexpr in custom Triton kernel calls from C++ #114475

aakhundov · 2023-11-23T18:52:55Z

Stack from ghstack (oldest at bottom):

-> [inductor] Pass None and skip constexpr in custom Triton kernel calls from C++ #114475

Summary: None arguments are codegened as *i8 in the triton_meta of the generated or user-defined Triton kernels:

pytorch/torch/_inductor/codegen/triton_utils.py

Lines 33 to 36 in 85aa372

    
           if arg.expr is None: 
        
               # From triton/runtime/jit.py 
        
               # `None` is nullptr.  Implicitly convert to *i8. 
        
               return "*i8"

Due to this, in contrary to the conventional Triton, we actually should pass nullptr to the Triton kernels in C++ wrapper codegen instead of passing nothing (as normally None doesn't make it to the generated PTX parameters, just like tl.constexpr args).

This PR adds two things:

Proper C++ wrapper codegening (ABI and non-ABI) of nullptr and c10::nullopt, as the prior way codegening c10::nullopt as tensor breaks (also c10 breaks in the ABI mode).
Skipping tl.constexpr args when calling the loaded-from-cubin compiled Triton kernel in the C++ wrapper codegen. As a side effect, this also resolves an issue with string arguments: now they are simply omitted in the C++ wrapper codegen.

Test Plan:

$ python test/inductor/test_aot_inductor.py -k test_triton_kernel_with_none_input
...
----------------------------------------------------------------------
Ran 4 tests in 40.364s

OK (skipped=2)

Reviewers:

Subscribers:

Tasks:

Tags:

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

Summary: `None` arguments are codegened as `*i8` in the `triton_meta` of the generated or user-defined Triton kernels: https://github.com/pytorch/pytorch/blob/85aa3723749e0d06aa5fd34215b9b93529a60995/torch/_inductor/codegen/triton_utils.py#L33-L36 Due to this, in contrary to the conventional Triton, we actually should pass `nullptr` to the Triton kernels in C++ wrapper codegen instead of passing nothing (as normally `None` doesn't make it to the generated PTX parameters, just like `tl.constexpr` args). This PR adds two things: 1. Proper C++ wrapper codegening (ABI and non-ABI) of `nullptr` and `c10::nullopt`, as the prior way codegening `c10::nullopt` as tensor breaks (also `c10` breaks in the ABI mode). 2. Skipping `tl.constexpr` args when calling the loaded-from-cubin compiled Triton kernel in the C++ wrapper codegen. As a side effect, this also resolves an issue with string arguments: now they are simply omitted in the C++ wrapper codegen. Test Plan: ``` $ python test/inductor/test_aot_inductor.py -k test_triton_kernel_with_none_input ... ---------------------------------------------------------------------- Ran 4 tests in 40.364s OK (skipped=2) ``` Reviewers: Subscribers: Tasks: Tags: [ghstack-poisoned]

pytorch-bot · 2023-11-23T18:52:58Z

🔗 Helpful Links

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

📄 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! (1 Unrelated Failure)

As of commit 9a1f005 with merge base b27565a ():

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

pull / linux-focal-py3.8-clang10 / test (dynamo, 2, 2, linux.2xlarge, unstable) (gh)

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

oulgen · 2023-11-23T18:59:41Z

torch/_inductor/ir.py

+            from torch._higher_order_ops.triton_kernel_wrap import kernel_side_table
+
+            # in C++ wrapper, we don't pass constexpr args, as they don't
+            # get added asparameters to the PTX code compiled from the


Suggested change

# get added asparameters to the PTX code compiled from the

# get added as parameters to the PTX code compiled from the

test/inductor/test_aot_inductor.py

torch/_inductor/ir.py

…ernel calls from C++" Summary: `None` arguments are codegened as `*i8` in the `triton_meta` of the generated or user-defined Triton kernels: https://github.com/pytorch/pytorch/blob/85aa3723749e0d06aa5fd34215b9b93529a60995/torch/_inductor/codegen/triton_utils.py#L33-L36 Due to this, in contrary to the conventional Triton, we actually should pass `nullptr` to the Triton kernels in C++ wrapper codegen instead of passing nothing (as normally `None` doesn't make it to the generated PTX parameters, just like `tl.constexpr` args). This PR adds two things: 1. Proper C++ wrapper codegening (ABI and non-ABI) of `nullptr` and `c10::nullopt`, as the prior way codegening `c10::nullopt` as tensor breaks (also `c10` breaks in the ABI mode). 2. Skipping `tl.constexpr` args when calling the loaded-from-cubin compiled Triton kernel in the C++ wrapper codegen. As a side effect, this also resolves an issue with string arguments: now they are simply omitted in the C++ wrapper codegen. Test Plan: ``` $ python test/inductor/test_aot_inductor.py -k test_triton_kernel_with_none_input ... ---------------------------------------------------------------------- Ran 4 tests in 40.364s OK (skipped=2) ``` Reviewers: Subscribers: Tasks: Tags: cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler [ghstack-poisoned]

oulgen · 2023-11-23T21:16:10Z

torch/_inductor/ir.py

+            # in C++ wrapper, we don't pass constexpr args, as they don't
+            # get added as parameters to the PTX code compiled from the
+            # user-defined Triton kernel (only non-constexpr args do)
+            args = self._remove_constexpr_args(args)


The above

kernel, configs = self.get_kernel_and_configs()

already extracted the kernel for you, so just pass kernel to this function instead of duplicating that logic

oulgen · 2023-11-23T21:16:57Z

torch/_inductor/ir.py



 class UserDefinedTritonKernel(ExternKernel):
+    def _remove_constexpr_args(self, args):


take kernel here from below and make this a static function

you can in fact inline this to codegen since you dont call this function anywhere else

Actually after taking the extracted kernel in def codegen, this function is only one line now, please just inline it below instead, no need for new function

Few lines, but inlined.

…ernel calls from C++" Summary: `None` arguments are codegened as `*i8` in the `triton_meta` of the generated or user-defined Triton kernels: https://github.com/pytorch/pytorch/blob/85aa3723749e0d06aa5fd34215b9b93529a60995/torch/_inductor/codegen/triton_utils.py#L33-L36 Due to this, in contrary to the conventional Triton, we actually should pass `nullptr` to the Triton kernels in C++ wrapper codegen instead of passing nothing (as normally `None` doesn't make it to the generated PTX parameters, just like `tl.constexpr` args). This PR adds two things: 1. Proper C++ wrapper codegening (ABI and non-ABI) of `nullptr` and `c10::nullopt`, as the prior way codegening `c10::nullopt` as tensor breaks (also `c10` breaks in the ABI mode). 2. Skipping `tl.constexpr` args when calling the loaded-from-cubin compiled Triton kernel in the C++ wrapper codegen. As a side effect, this also resolves an issue with string arguments: now they are simply omitted in the C++ wrapper codegen. Test Plan: ``` $ python test/inductor/test_aot_inductor.py -k test_triton_kernel_with_none_input ... ---------------------------------------------------------------------- Ran 4 tests in 40.364s OK (skipped=2) ``` Reviewers: Subscribers: Tasks: Tags: cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler [ghstack-poisoned]

oulgen · 2023-11-23T21:53:30Z

torch/_inductor/ir.py

            # get added as parameters to the PTX code compiled from the
            # user-defined Triton kernel (only non-constexpr args do)
-            args = self._remove_constexpr_args(args)
+            from triton.runtime.jit import JITFunction


You dont need any of this, it is already a JITFunction

kernel, configs = self.get_kernel_and_configs()

took care of it for you

3839-3843 are not needed

Nope: autotuning tests fail. There kernel is an Autotuner instance which doesn't have constexpr. Need to unwrap. That's why I updated the PR.

…ernel calls from C++" Summary: `None` arguments are codegened as `*i8` in the `triton_meta` of the generated or user-defined Triton kernels: https://github.com/pytorch/pytorch/blob/85aa3723749e0d06aa5fd34215b9b93529a60995/torch/_inductor/codegen/triton_utils.py#L33-L36 Due to this, in contrary to the conventional Triton, we actually should pass `nullptr` to the Triton kernels in C++ wrapper codegen instead of passing nothing (as normally `None` doesn't make it to the generated PTX parameters, just like `tl.constexpr` args). This PR adds two things: 1. Proper C++ wrapper codegening (ABI and non-ABI) of `nullptr` and `c10::nullopt`, as the prior way codegening `c10::nullopt` as tensor breaks (also `c10` breaks in the ABI mode). 2. Skipping `tl.constexpr` args when calling the loaded-from-cubin compiled Triton kernel in the C++ wrapper codegen. As a side effect, this also resolves an issue with string arguments: now they are simply omitted in the C++ wrapper codegen. Test Plan: ``` $ python test/inductor/test_aot_inductor.py -k test_triton_kernel_with_none_input ... ---------------------------------------------------------------------- Ran 4 tests in 40.364s OK (skipped=2) ``` Reviewers: Subscribers: Tasks: Tags: cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler [ghstack-poisoned]

oulgen · 2023-11-24T00:44:27Z

torch/_inductor/ir.py

+            if isinstance(kernel, Autotuner):
+                kernel = kernel.fn


I'm still a bit confused about your comment here.

kernel, configs = self.get_kernel_and_configs()

is guaranteed to return a kernel such that

assert isinstance(kernel, JITFunction)

if you look at line 3820 above, it does kernel = kernel.fn when it is autotuner.

Yeah, I think, you're right. The autotuning test has failed before when I was fetching the kernel myself. Now that I'm reusing the kernel returned by get_kernel_and_configs, I can skip the unwrapping as it's already done.

…ernel calls from C++" Summary: `None` arguments are codegened as `*i8` in the `triton_meta` of the generated or user-defined Triton kernels: https://github.com/pytorch/pytorch/blob/85aa3723749e0d06aa5fd34215b9b93529a60995/torch/_inductor/codegen/triton_utils.py#L33-L36 Due to this, in contrary to the conventional Triton, we actually should pass `nullptr` to the Triton kernels in C++ wrapper codegen instead of passing nothing (as normally `None` doesn't make it to the generated PTX parameters, just like `tl.constexpr` args). This PR adds two things: 1. Proper C++ wrapper codegening (ABI and non-ABI) of `nullptr` and `c10::nullopt`, as the prior way codegening `c10::nullopt` as tensor breaks (also `c10` breaks in the ABI mode). 2. Skipping `tl.constexpr` args when calling the loaded-from-cubin compiled Triton kernel in the C++ wrapper codegen. As a side effect, this also resolves an issue with string arguments: now they are simply omitted in the C++ wrapper codegen. Test Plan: ``` $ python test/inductor/test_aot_inductor.py -k test_triton_kernel_with_none_input ... ---------------------------------------------------------------------- Ran 4 tests in 40.364s OK (skipped=2) ``` Reviewers: Subscribers: Tasks: Tags: cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler [ghstack-poisoned]

Summary: `None` arguments are codegened as `*i8` in the `triton_meta` of the generated or user-defined Triton kernels: https://github.com/pytorch/pytorch/blob/85aa3723749e0d06aa5fd34215b9b93529a60995/torch/_inductor/codegen/triton_utils.py#L33-L36 Due to this, in contrary to the conventional Triton, we actually should pass `nullptr` to the Triton kernels in C++ wrapper codegen instead of passing nothing (as normally `None` doesn't make it to the generated PTX parameters, just like `tl.constexpr` args). This PR adds two things: 1. Proper C++ wrapper codegening (ABI and non-ABI) of `nullptr` and `c10::nullopt`, as the prior way codegening `c10::nullopt` as tensor breaks (also `c10` breaks in the ABI mode). 2. Skipping `tl.constexpr` args when calling the loaded-from-cubin compiled Triton kernel in the C++ wrapper codegen. As a side effect, this also resolves an issue with string arguments: now they are simply omitted in the C++ wrapper codegen. Test Plan: ``` $ python test/inductor/test_aot_inductor.py -k test_triton_kernel ... ---------------------------------------------------------------------- Ran 108 tests in 1251.475s OK (skipped=54) ``` Reviewers: Subscribers: Tasks: Tags: ghstack-source-id: 419cc0e Pull Request resolved: #114475

aakhundov · 2023-11-24T12:48:49Z

@pytorchbot merge

pytorchmergebot · 2023-11-24T12:50: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

github-actions bot added module: inductor ciflow/inductor labels Nov 23, 2023

aakhundov requested review from chenyang78, desertfire and oulgen November 23, 2023 18:54

aakhundov added ciflow/trunk Trigger trunk jobs on your pull request topic: not user facing topic category labels Nov 23, 2023

aakhundov changed the title ~~Pass None and skip constexpr in custom Triton kernel calls from C++~~ [inductor] Pass None and skip constexpr in custom Triton kernel calls from C++ Nov 23, 2023

oulgen requested changes Nov 23, 2023

View reviewed changes

oulgen approved these changes Nov 23, 2023

View reviewed changes

oulgen reviewed Nov 23, 2023

View reviewed changes

oulgen reviewed Nov 24, 2023

View reviewed changes

pytorchmergebot added the merging label Nov 24, 2023

pytorchmergebot added the Merged label Nov 24, 2023

pytorchmergebot closed this in 0a063ad Nov 24, 2023

pytorchmergebot removed the merging label Nov 24, 2023

facebook-github-bot deleted the gh/aakhundov/4/head branch November 27, 2023 15:29

	if arg.expr is None:
	# From triton/runtime/jit.py
	# `None` is nullptr. Implicitly convert to *i8.
	return "*i8"

	# get added asparameters to the PTX code compiled from the
	# get added as parameters to the PTX code compiled from the



		class UserDefinedTritonKernel(ExternKernel):
		def _remove_constexpr_args(self, args):

[inductor] Pass None and skip constexpr in custom Triton kernel calls from C++ #114475

[inductor] Pass None and skip constexpr in custom Triton kernel calls from C++ #114475

Uh oh!

Conversation

aakhundov commented Nov 23, 2023 • 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 23, 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/114475

✅ You can merge normally! (1 Unrelated Failure)

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

oulgen Nov 23, 2023 • 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!

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!

oulgen Nov 24, 2023 • 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!

aakhundov commented Nov 24, 2023

Uh oh!

pytorchmergebot commented Nov 24, 2023

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

aakhundov commented Nov 23, 2023 •

edited by pytorch-bot bot

Loading

pytorch-bot bot commented Nov 23, 2023 •

edited

Loading

oulgen Nov 23, 2023 •

edited

Loading

oulgen Nov 24, 2023 •

edited

Loading