mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
[inductor][ez] ExternChoice with maybe_append_choice (#161336)
# why - make the API for ExternChoice the same as KernelTemplate - make it possible to use the same retrieval point as templates # what - add a maybe_append_choice to ExternChoice that under the hood invokes self.bind This pr does not actuate the new path, but just exposes it # testing ``` python3 -bb -m pytest test/inductor/test_max_autotune.py ``` Differential Revision: [D81520578](https://our.internmc.facebook.com/intern/diff/D81520578) Pull Request resolved: https://github.com/pytorch/pytorch/pull/161336 Approved by: https://github.com/jansel ghstack dependencies: #161123, #161124, #161125, #161126
This commit is contained in:
committed by
PyTorch MergeBot
parent
6cb13dd3cc
commit
c31dee6fa5
@ -1907,6 +1907,19 @@ class ExternKernelChoice:
|
||||
self, input_nodes, layout, kwargs, has_out_variant=self.has_out_variant
|
||||
)
|
||||
|
||||
def maybe_append_choice(
|
||||
self, choices: list[Any], **kwargs: Any
|
||||
) -> Optional[NotImplementedError]:
|
||||
# convenience function to match the Template interface, so that
|
||||
# templates and ExternKernelChoice can be treated the same when
|
||||
# generating choice callers
|
||||
assert "input_nodes" in kwargs, "input_nodes argument required"
|
||||
assert "layout" in kwargs, "layout argument required"
|
||||
input_nodes = kwargs.pop("input_nodes")
|
||||
layout = kwargs.pop("layout")
|
||||
choices.append(self.bind(input_nodes=input_nodes, layout=layout, **kwargs))
|
||||
return None
|
||||
|
||||
|
||||
class TritonTemplateCaller(ir.TritonTemplateCallerBase):
|
||||
def __init__(
|
||||
|
Reference in New Issue
Block a user