Added option to control number of kernel options displayed (pytorch#138788)

Pull Request resolved: pytorch#138788
Approved by: https://github.com/drisspg

Author: Chillee

torch/_inductor/config.py

@@ -298,6 +298,9 @@ def bundle_triton_into_fx_graph_cache_default() -> Optional[bool]:
 # enable slow autotuning passes to select gemm algorithms
 max_autotune_gemm = os.environ.get("TORCHINDUCTOR_MAX_AUTOTUNE_GEMM") == "1"
 
+# Modifies the number of autotuning choices displayed, set to None for all
+autotune_num_choices_displayed = 10
+
 # force cublas and triton to use the same precision; cublas supports TF32 for matmul operations
 # when m, n, k are multiples of 16, 16, 8, whereas triton supports TF32 for matmul operations
 # for any combinations of m, n, k, regardless of their alignment. setting this flag will ensure

torch/_inductor/select_algorithm.py

@@ -1730,9 +1730,15 @@ def log_results(
                 for n in input_nodes
             ]
         )
+        if config.autotune_num_choices_displayed == 0:
+            return
+        elif config.autotune_num_choices_displayed is None:
+            n = -1
+        else:
+            n = config.autotune_num_choices_displayed
 
-        n = None if log.getEffectiveLevel() == logging.DEBUG else 10
         top_k = sorted(timings, key=timings.__getitem__)[:n]
+
         best = top_k[0]
 
         def get_choice_info(choice):