mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
d91eecc9a5
# why - unnecessary as we only ever need to know the dtype and maybe the device - we already take in the kernel inputs which have the device - enable us to specify the layout after finding all the configs but before generating the ChoiceCallers # what - replace all calls in template_heuristics that used to take Layout with now just taking out_dtype # testing ci Differential Revision: [D81820115](https://our.internmc.facebook.com/intern/diff/D81820115) Pull Request resolved: https://github.com/pytorch/pytorch/pull/162238 Approved by: https://github.com/eellison ghstack dependencies: #161347, #161348, #161349
19 lines
540 B
Python
19 lines
540 B
Python
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING
|
|
|
|
from .. import config as inductor_config
|
|
from .base import TemplateConfigHeuristics
|
|
|
|
|
|
if TYPE_CHECKING:
|
|
from ..kernel_inputs import KernelInputs
|
|
|
|
|
|
class GemmMaxAutotuneTemplateConfigHeuristics(TemplateConfigHeuristics):
|
|
def should_run(self, inputs: KernelInputs) -> bool:
|
|
"""
|
|
simple base override for GEMM family templates that run only in max-autotune
|
|
"""
|
|
return inductor_config.max_autotune or inductor_config.max_autotune_gemm
|