Expose option for saving # best candidates in autotuning

Author: prigoyal

tc/autotuner/genetic_autotuner.cc

@@ -57,7 +57,8 @@ void GeneticAutotuner::storeCaches(const std::string& filename) {
   } else {
     std::cout << "Dumping cache to " << filename << ".cuda/options"
               << std::endl;
-    tc::OptionsCache::getCache()->keepOnlyBestCandidates(10);
+    tc::OptionsCache::getCache()->keepOnlyBestCandidates(
+        tc::FLAGS_tuner_save_best_candidates_count);
     tc::OptionsCache::dumpCacheToProtobuf(tc::makeOptionsFilename(filename));
 
     tc::OptionsCache::getCache()->keepOnlyBestCandidates(1);

tc/core/flags.cc

@@ -105,6 +105,10 @@ DEFINE_int64(
     random_seed,
     -1,
     "The number of best candidates to restore from the proto cache");
+DEFINE_uint32(
+    tuner_save_best_candidates_count,
+    10,
+    "Number of best candidates to save from autotuning");
 
 uint64_t initRandomSeed() {
   static std::mutex mut;

tc/core/flags.h

@@ -52,6 +52,7 @@ DECLARE_bool(tuner_gen_restore_from_proto);
 DECLARE_uint32(tuner_gen_restore_number);
 DECLARE_bool(tuner_gen_log_generations);
 DECLARE_uint64(tuner_min_launch_total_threads);
+DECLARE_uint32(tuner_save_best_candidates_count);
 
 // Misc
 DECLARE_int64(random_seed);

tensor_comprehensions/pybinds/pybind_autotuner.cc

@@ -110,6 +110,13 @@ PYBIND11_MODULE(autotuner, m) {
             tc::FLAGS_tuner_min_launch_total_threads =
                 tuner_min_launch_total_threads;
           })
+      .def(
+          "save_best_candidates_count",
+          [](tc::autotune::GeneticAutotunerATen& instance,
+             bool save_best_candidates_count) {
+            tc::FLAGS_tuner_save_best_candidates_count =
+                save_best_candidates_count;
+          })
       .def(
           "tune",
           [dlpack](

tensor_comprehensions/tc_unit.py

@@ -193,7 +193,7 @@ def __init__(self, tc_lang, **kwargs):
     def set_autotuner_parameters(
         self, pop_size=20, crossover_rate=80, mutation_rate=7, generations=10,
         number_elites=1, threads=8, gpus="0", restore_from_proto=False,
-        restore_number=10, log_generations=False,
+        restore_number=10, log_generations=False, save_best_candidates_count=10,
         tuner_min_launch_total_threads=64, **kwargs
     ):
         self.autotuner.pop_size(pop_size)
@@ -206,6 +206,7 @@ def set_autotuner_parameters(
         self.autotuner.restore_from_proto(restore_from_proto)
         self.autotuner.restore_number(restore_number)
         self.autotuner.log_generations(log_generations)
+        self.autotuner.save_best_candidates_count(save_best_candidates_count)
         self.autotuner.tuner_min_launch_total_threads(tuner_min_launch_total_threads)
 
     # We need to pass the inputs so that we can load the correct options from