Skip to content

Commit 3cad240

Browse files
coconutrubenpytorchmergebot
authored andcommitted
[inductor][choices] pass through annotations from KTC to ChoiceCaller (pytorch#163117)
# why - KTC might regenerate a choicecaller e.g. through FlexibleLayout optimization. This in turn would delete any annotations # what - provide an annotations dict inside KTC - forward that dict towards the ChoiceCaller's annotations - ChoiceCaller users e.g. in selectalgorithm now have access to the KTC and can register handlers do record/make decisions based on the KTC # testing n/a Differential Revision: [D82587631](https://our.internmc.facebook.com/intern/diff/D82587631) Pull Request resolved: pytorch#163117 Approved by: https://github.com/nmacchioni
1 parent e63476b commit 3cad240

File tree

1 file changed

+3
-0
lines changed

1 file changed

+3
-0
lines changed

torch/_inductor/kernel_template_choice.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ def __init__(
3333
self.extra_kwargs = extra_kwargs
3434
self.layout = layout
3535
self.inputs = inputs
36+
self.annotations: dict[str, Any] = {"ktc": self}
3637

3738
@property
3839
def choice(self) -> Optional[ChoiceCaller]:
@@ -54,6 +55,8 @@ def choice(self) -> Optional[ChoiceCaller]:
5455
input_nodes=self.inputs.nodes(),
5556
**self.extra_kwargs,
5657
)
58+
if self._choice is not None:
59+
self._choice.annotations = self.annotations
5760
return self._choice
5861

5962

0 commit comments

Comments
 (0)