Method __call__

torch/_inductor/select_algorithm.py:705–767

(
        self,
        name,
        choices: List[ChoiceCaller],
        input_nodes,
        layout,
        # optional dict mapping arg indices to the functions
        # generating a torch.Tensor for that input from the
        # corresponding ir.Buffer. if passed for a given
        # arg, the function will be called instead of
        # generating a random torch.Tensor for benchmarking.
        input_gen_fns: Optional[Dict[int, Callable[[ir.Buffer], torch.Tensor]]] = None,
    )

class AlgorithmSelectorCache(PersistentCache):
    def __call__(
        self,
        name,
        choices: List[ChoiceCaller],
        input_nodes,
        layout,
        # optional dict mapping arg indices to the functions
        # generating a torch.Tensor for that input from the
        # corresponding ir.Buffer. if passed for a given
        # arg, the function will be called instead of
        # generating a random torch.Tensor for benchmarking.
        input_gen_fns: Optional[Dict[int, Callable[[ir.Buffer], torch.Tensor]]] = None,
    ):
        from .codegen.cuda.cuda_kernel import CUDATemplateCaller

        # TODO(nmacchioni): remove once CI tests are fixed
        choices = [choice for choice in choices if choice is not None]
        if len(choices) == 0:
            raise RuntimeError(
                "No choices to select, please consider adding ATEN into max_autotune_gemm_backends "
                "config (defined in torch/_inductor/config.py) to allow at least one choice. "
            )
        log.info("Max autotune selects from %s choices.", str(len(choices)))

        if len(choices) == 1:
            if not isinstance(choices[0], CUDATemplateCaller):
                # CUDATemplateCaller still needs to go through autotuning process to retrieve workspace size.
                return choices[0].output_node()

        @functools.lru_cache(None)
        def make_benchmark_fn():
            return self.make_benchmark_fn(choices, input_nodes, layout, input_gen_fns)

        def autotune(choices):
            return make_benchmark_fn()(choices)

        if config.autotune_in_subproc:
            from .autotune_process import tuning_pool

            # do the optional warmup
            tuning_pool.initialize()

        autotune_start_ts = time.time()
        timings = self.lookup(
            choices,
            name,
            repr([self.key_of(x) for x in input_nodes]),
            autotune,
        )
        autotune_elapse = time.time() - autotune_start_ts
        if timings == {} or choices[0] not in timings:
            return choices[0].output_node()

        if make_benchmark_fn.cache_info().currsize:
            counters["inductor"]["select_algorithm_autotune"] += 1
        if (
            make_benchmark_fn.cache_info().currsize
            or log.getEffectiveLevel() == logging.DEBUG