[AOTI] Support non auto-tuned triton kernels in aoti #113090

oulgen · 2023-11-06T22:55:31Z

Stack from ghstack (oldest at bottom):

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

[ghstack-poisoned]

pytorch-bot · 2023-11-06T22:55:33Z

🔗 Helpful Links

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

📄 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

As of commit 036cbc8 with merge base b7acd37 ():
💚 Looks good so far! There are no failures yet. 💚

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

oulgen · 2023-11-06T22:56:28Z

torch/_inductor/ir.py

            for arg_name in self.ordered_kwargs_for_cpp_kernel:
                v = self.get_kwargs_value(arg_name)
-                kwargs.append(V.graph.wrapper_code.val_to_arg_str(v))
+                if isinstance(v, sympy.Expr):


I don't like this but not sure if there's a better solution..

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler [ghstack-poisoned]

ghstack-source-id: 0f7105a Pull Request resolved: #113090

jansel · 2023-11-07T04:20:09Z

@desertfire @chenyang78 mind reviewing this one?

chenyang78 · 2023-11-07T08:19:49Z

test/inductor/test_aot_inductor.py

+    @requires_cuda()
+    @skipIfRocm


Given that we have if self.device != "cuda" check below, are these two redundant?

triton is not installed on all machines, so in the past, if i didn't include the both of these, i ran into some symbol problems. I can try deleting them and seeing if anything fails

chenyang78 · 2023-11-07T09:03:05Z

torch/_inductor/ir.py

+            if self.kwargs and not self.ordered_kwargs_for_cpp_kernel:
                assert (
                    self.ordered_kwargs_for_cpp_kernel
                ), "ordered_kwargs_for_cpp_kernel is missing"


Looks like the assert condition is always False? Probably we could just thrown an exception here?

chenyang78 · 2023-11-07T09:33:44Z

torch/_inductor/ir.py

            for arg_name in self.ordered_kwargs_for_cpp_kernel:
                v = self.get_kwargs_value(arg_name)
-                kwargs.append(V.graph.wrapper_code.val_to_arg_str(v))
+                if isinstance(v, sympy.Expr):


Could we add a clause to handle sympy.Expr from CppWrapperCodeGen's val_to_arg_str so that we could make kwargs hold only strings.

But I think the point here is to use sympy.Expr as kwarg to codegen it at a later stage? cc @oulgen.

Yes, if we converted to a string, then codegen will not be able to tell it was a sympy expr and emit it like a tensor. @chenyang78 The problem is codegen tries to figure out things based on their string representation but for sympy expr, it keeps them as is.. which is not great.

I see. Thanks for the clarification. This is quite unfortunate...

chenyang78 · 2023-11-07T09:48:48Z

torch/_inductor/ir.py

        self.grid = grid

+        kernel, _ = self.get_kernel_and_configs()
+        self.ordered_kwargs_for_cpp_kernel = kernel.arg_names


It seems to be pretty hacky to rely on ordered_kwargs_for_cpp_kernel to generate kernel args indirectly. Wondering if we could implement a codegen_args method for UserDefinedTritonKernel, where we return arg strings.

I'm not following. It is not possible to return all strings because then sympy expr can no longer be distinguished. Look at

pytorch/torch/_inductor/codegen/wrapper.py

Lines 2468 to 2501 in 562c4ae

def generate_args_decl(self, call_args):

dynamic_symbols = V.graph.sizevars.free_symbols()

# TODO: only works for constant now, need type info

new_args = []

for arg in call_args:

var_name = f"var_{next(self.arg_var_id)}"

if isinstance(

arg,

(

sympy.Integer,

sympy.Symbol,

SymbolicCallArg,

),

):

self.writeline(f"auto {var_name} = {arg};")

elif is_int(arg):

self.writeline(f"int {var_name} = {arg};")

elif is_float(arg):

self.writeline(f"float {var_name} = {arg};")

elif any(str(arg) == s.name for s in dynamic_symbols):

self.writeline(f"auto {var_name} = {arg};")

else:

if config.aot_inductor.abi_compatible:

self.writeline(f"CUdeviceptr {var_name};")

self.writeline(

f"AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_data_ptr({arg}, reinterpret_cast<void**>(&{var_name})));"

)

else:

self.writeline(

f"CUdeviceptr {var_name} = reinterpret_cast<CUdeviceptr>({arg}.data_ptr());"

)

new_args.append(f"&{var_name}")

return ", ".join(new_args)

test/inductor/test_aot_inductor.py

torch/_inductor/codegen/wrapper.py

aakhundov · 2023-11-07T14:15:10Z

torch/_inductor/codegen/wrapper.py

                arg,
                (
                    sympy.Integer,
+                    sympy.Expr,


I'm wondering if we can simply use {arg} in the f-string below if the arg is sympy.Expr. From what I remember, in C++ codegen we need to use self.expr_printer for expressions?

Good point!

aakhundov · 2023-11-07T14:16:40Z

torch/_inductor/ir.py

            for arg_name in self.ordered_kwargs_for_cpp_kernel:
                v = self.get_kwargs_value(arg_name)
-                kwargs.append(V.graph.wrapper_code.val_to_arg_str(v))
+                if isinstance(v, sympy.Expr):


But I think the point here is to use sympy.Expr as kwarg to codegen it at a later stage? cc @oulgen.

aakhundov

LGTM, with a few nits. Thanks!

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler [ghstack-poisoned]

chenyang78 · 2023-11-07T23:01:23Z

torch/_inductor/ir.py

            for arg_name in self.ordered_kwargs_for_cpp_kernel:
                v = self.get_kwargs_value(arg_name)
-                kwargs.append(V.graph.wrapper_code.val_to_arg_str(v))
+                if isinstance(v, sympy.Expr):


I see. Thanks for the clarification. This is quite unfortunate...

oulgen · 2023-11-08T02:26:48Z

@pytorchbot merge

pytorchmergebot · 2023-11-08T02:28:36Z

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 · 2023-11-08T02:28:55Z

Merge failed

Reason: Command git -C /home/runner/work/pytorch/pytorch cherry-pick -x 8b2992137f78e28a48c9fa782d7124697a828436 returned non-zero exit code 1

Auto-merging test/inductor/test_aot_inductor.py
Auto-merging torch/_inductor/ir.py
CONFLICT (content): Merge conflict in torch/_inductor/ir.py
error: could not apply 8b2992137f7... [AOTI] Support non auto-tuned triton kernels in aoti
hint: After resolving the conflicts, mark them with
hint: "git add/rm <pathspec>", then run
hint: "git cherry-pick --continue".
hint: You can instead skip this commit with "git cherry-pick --skip".
hint: To abort and get back to the state before "git cherry-pick",
hint: run "git cherry-pick --abort".

Details for Dev Infra team

Raised by workflow job

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler [ghstack-poisoned]

oulgen · 2023-11-08T06:22:54Z

@pytorchbot merge

pytorchmergebot · 2023-11-08T06:24:58Z

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

Pull Request resolved: pytorch#113090 Approved by: https://github.com/aakhundov, https://github.com/chenyang78, https://github.com/desertfire

[AOTI] Support non auto-tuned triton kernels in aoti

9c85aa4

[ghstack-poisoned]

oulgen mentioned this pull request Nov 6, 2023

[Inductor] Cache generated user defined triton kernels on tensor dtype and non tensor parameters #112752

Closed

This was referenced Nov 6, 2023

Move all triton related testing utils into shared file #113008

Closed

[Inductor] Improve reinplace_scatters pass #112801

Closed

[Inductor] Allow None values to be passed in as arguments to triton kernels #113056

Closed

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

oulgen commented Nov 6, 2023

View reviewed changes

oulgen requested review from Chillee, aakhundov, jansel and zou3519 November 6, 2023 22:59

oulgen added ciflow/trunk Trigger trunk jobs on your pull request rocm This tag is for PRs from ROCm team ciflow/rocm Trigger "default" config CI on ROCm and removed rocm This tag is for PRs from ROCm team labels Nov 6, 2023

Update on "[AOTI] Support non auto-tuned triton kernels in aoti"

12a14fb

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler [ghstack-poisoned]

Update on "[AOTI] Support non auto-tuned triton kernels in aoti"

557e54d

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler [ghstack-poisoned]

Update on "[AOTI] Support non auto-tuned triton kernels in aoti"

142a204

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler [ghstack-poisoned]

oulgen added a commit that referenced this pull request Nov 7, 2023

[AOTI] Support non auto-tuned triton kernels in aoti

f345b47

ghstack-source-id: 0f7105a Pull Request resolved: #113090

jansel requested review from chenyang78 and desertfire November 7, 2023 04:18

jansel removed their request for review November 7, 2023 04:20

chenyang78 reviewed Nov 7, 2023

View reviewed changes

aakhundov reviewed Nov 7, 2023

View reviewed changes

aakhundov approved these changes Nov 7, 2023

View reviewed changes

Update on "[AOTI] Support non auto-tuned triton kernels in aoti"

c367dee

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler [ghstack-poisoned]

chenyang78 approved these changes Nov 7, 2023

View reviewed changes

oulgen added the release notes: inductor label Nov 8, 2023

pytorchmergebot added the merging label Nov 8, 2023

pytorchmergebot removed the merging label Nov 8, 2023

Update on "[AOTI] Support non auto-tuned triton kernels in aoti"

673d34e

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler [ghstack-poisoned]

oulgen mentioned this pull request Nov 8, 2023

[AOTI] Implement support for user defined kernels that use triton.autotune #113229

Closed

Update on "[AOTI] Support non auto-tuned triton kernels in aoti"

036cbc8

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler [ghstack-poisoned]

desertfire approved these changes Nov 8, 2023

View reviewed changes

pytorchmergebot added the merging label Nov 8, 2023

pytorchmergebot added Merged and removed merging labels Nov 8, 2023

pytorchmergebot closed this in 8ba11bf Nov 8, 2023

facebook-github-bot deleted the gh/oulgen/29/head branch November 11, 2023 15:24

	def generate_args_decl(self, call_args):
	dynamic_symbols = V.graph.sizevars.free_symbols()
	# TODO: only works for constant now, need type info
	new_args = []
	for arg in call_args:
	var_name = f"var_{next(self.arg_var_id)}"
	if isinstance(
	arg,
	(
	sympy.Integer,
	sympy.Symbol,
	SymbolicCallArg,
	),
	):
	self.writeline(f"auto {var_name} = {arg};")
	elif is_int(arg):
	self.writeline(f"int {var_name} = {arg};")
	elif is_float(arg):
	self.writeline(f"float {var_name} = {arg};")
	elif any(str(arg) == s.name for s in dynamic_symbols):
	self.writeline(f"auto {var_name} = {arg};")
	else:
	if config.aot_inductor.abi_compatible:
	self.writeline(f"CUdeviceptr {var_name};")
	self.writeline(
	f"AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_data_ptr({arg}, reinterpret_cast<void**>(&{var_name})));"
	)
	else:
	self.writeline(
	f"CUdeviceptr {var_name} = reinterpret_cast<CUdeviceptr>({arg}.data_ptr());"
	)
	new_args.append(f"&{var_name}")

	return ", ".join(new_args)

[AOTI] Support non auto-tuned triton kernels in aoti #113090

[AOTI] Support non auto-tuned triton kernels in aoti #113090

Uh oh!

Conversation

oulgen commented Nov 6, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

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

✅ No Failures

Uh oh!

oulgen Nov 6, 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!

jansel commented Nov 7, 2023

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!

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!

aakhundov left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

oulgen commented Nov 8, 2023

Uh oh!

pytorchmergebot commented Nov 8, 2023

Merge started

Uh oh!

pytorchmergebot commented Nov 8, 2023

Merge failed

Uh oh!

oulgen commented Nov 8, 2023

Uh oh!

pytorchmergebot commented Nov 8, 2023

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

oulgen commented Nov 6, 2023 •

edited

Loading

pytorch-bot bot commented Nov 6, 2023 •

edited

Loading

oulgen Nov 6, 2023 •

edited

Loading