[genetic search] Switch to (mu,lamda) selection

Previously each generation had mu candidates and generated mu
children which all survided. This meant that really bad candidates
that were randomly generated would survive across generations.

With this change, lambda (typically larger thatn mu) children are
generated and the best mu survive. The previous behaviour is a
special case in which lambda = mu.

Author: Theodoros Theodoridis

examples/tensordot.cc

@@ -96,7 +96,6 @@ int main(int argc, char** argv) {
   ::gflags::ParseCommandLineFlags(&argc, &argv, true);
   ::google::InitGoogleLogging(argv[0]);
   setAtenSeed(tc::initRandomSeed(), at::Backend::CUDA);
-  tc::FLAGS_tuner_gen_number_elites = FLAGS_number_elites;
   tc::FLAGS_tuner_gen_generations = FLAGS_generations;
   tc::FLAGS_tuner_gen_pop_size = FLAGS_pop_size;
   tc::FLAGS_tuner_threads = FLAGS_threads;

include/tc/autotuner/genetic_search.h

@@ -70,7 +70,8 @@ class GeneticSearch {
       size_t n,
       uint8_t crossOverRate,
       uint8_t mutationRate,
-      size_t numberElites);
+      size_t matingPoolSize,
+      size_t selectionPoolSize);
 
   /**
    * confs are used to seed the first generation, the rest of the population is
@@ -92,15 +93,22 @@ class GeneticSearch {
       size_t n,
       uint8_t crossOverRate,
       uint8_t mutationRate,
-      size_t numberElites);
+      size_t matingPoolSize,
+      size_t selectionPoolSize);
 
-  void updateParameters();
+  void generateSelectionPool();
+  void selectSurvivors();
 
  private:
   std::vector<TuningConfiguration> stochasticUniversalSampling(
       const std::vector<double>& fitness) const;
+
   void breed();
 
+  void updateBestCandidate(const TuningConfiguration& c);
+
+  void resetPopulationIfNotEnoughCandidates();
+
   TuningConfiguration crossover(
       TuningConfiguration&,
       TuningConfiguration&,
@@ -113,12 +121,13 @@ class GeneticSearch {
   using Population = std::vector<std::unique_ptr<CandidateConfiguration>>;
 
   Population population;
+  Population selectionPool;
   TuningConfiguration lastBestConf;
   const size_t kMaxPopulationSize;
   const size_t kMatingPoolSize;
+  const size_t kSelectionPoolSize;
   const uint8_t kCrossOverRate;
   const uint8_t kMutationRate;
-  const size_t kNumberElites;
 
   /*
    * c++11 seeding is (apparently) not of the highest quality:

include/tc/autotuner/genetic_tuning_harness.h

@@ -38,7 +38,8 @@ class GeneticTunerHarness {
       size_t n,
       uint8_t crossoverRate,
       uint8_t mutationRate,
-      size_t numberElites,
+      size_t matingPoolSize,
+      size_t selectionPoolSize,
       lang::TreeRef tc,
       std::string kernelName,
       const std::unordered_map<size_t, std::vector<const DLTensor*>>& inputs,
@@ -66,12 +67,16 @@ class GeneticTunerHarness {
       size_t bestTimeSoFar);
 
   /// Helper function to delegate compiling on the cpu to different threads
-  template <typename ExecutorType>
-  void doCompile(ExecutorType& engine);
+  template <typename ExecutorType, typename Population>
+  void doCompile(ExecutorType& engine, Population& population);
 
   /// Helper function to delegate running on the gpu to different threads
-  template <typename ExecutorType>
-  void doGpuWork(size_t gpu, ExecutorType& engine, Printer& printer);
+  template <typename ExecutorType, typename Population>
+  void doGpuWork(
+      size_t gpu,
+      ExecutorType& engine,
+      Population& population,
+      Printer& printer);
 
   /// Make options from conf
   tc::CudaMappingOptions makeOptions(const CandidateConfiguration& conf);
@@ -90,7 +95,8 @@ class GeneticTunerHarness {
   const size_t kMaxPopulationSize;
   const uint8_t kCrossOverRate;
   const uint8_t kMutationRate;
-  const size_t kNumberElites;
+  const size_t kMatingPoolSize;
+  const size_t kSelectionPoolSize;
 
   TuningConfiguration configuration;
 

include/tc/autotuner/utils/printer.h

@@ -33,7 +33,7 @@ namespace autotune {
 class Printer {
  public:
   Printer(
-      size_t generation,
+      std::string prefix,
       size_t total,
       const std::atomic_size_t& currentCompilationJob,
       const std::atomic_size_t& numEvaluations);
@@ -47,7 +47,7 @@ class Printer {
  private:
   void printLoop();
 
-  size_t generation_;
+  std::string prefix_;
   std::vector<Duration> runtimes_;
   mutable std::mutex runtimesMtx_;
 

include/tc/core/flags.h

@@ -40,10 +40,11 @@ DECLARE_uint32(benchmark_iterations);
 
 // Used in autotuning
 DECLARE_uint32(tuner_gen_pop_size);
+DECLARE_uint32(tuner_gen_mating_pool_size);
+DECLARE_uint32(tuner_gen_selection_pool_size);
 DECLARE_uint32(tuner_gen_crossover_rate);
 DECLARE_uint32(tuner_gen_mutation_rate);
 DECLARE_uint32(tuner_gen_generations);
-DECLARE_uint32(tuner_gen_number_elites);
 DECLARE_uint32(tuner_threads);
 DECLARE_string(tuner_gpus);
 DECLARE_bool(tuner_print_best);

src/autotuner/genetic_autotuner.cc

@@ -118,7 +118,8 @@ llvm::Optional<CudaMappingOptions> GeneticAutotuner::tune(
       FLAGS_tuner_gen_pop_size,
       FLAGS_tuner_gen_crossover_rate,
       FLAGS_tuner_gen_mutation_rate,
-      FLAGS_tuner_gen_number_elites,
+      FLAGS_tuner_gen_mating_pool_size,
+      FLAGS_tuner_gen_selection_pool_size,
       tcNameMap_.at(tcName),
       tcName,
       inputs,

src/autotuner/genetic_search.cc

@@ -162,7 +162,8 @@ void dropInvalidConfigurations(GeneticSearch::Population& population) {
 } // namespace
 
 #define VALIDATE()                                     \
-  CHECK_LT(kNumberElites, kMaxPopulationSize);         \
+  CHECK_LT(kMaxPopulationSize, kMatingPoolSize);       \
+  CHECK_LT(kMaxPopulationSize, kSelectionPoolSize);    \
   CHECK(kMutationRate >= 0 and kMutationRate <= 100)   \
       << "the mutation rate (" << kMutationRate        \
       << ") should be in the [0,100] interval";        \
@@ -189,14 +190,15 @@ GeneticSearch::GeneticSearch(
     size_t n,
     uint8_t crossOverRate,
     uint8_t mutationRate,
-    size_t numberElites)
+    size_t matingPoolSize,
+    size_t selectionPoolSize)
     : population(),
       lastBestConf(confs[0]),
       kMaxPopulationSize(n),
-      kMatingPoolSize(n * 3),
+      kMatingPoolSize(matingPoolSize),
+      kSelectionPoolSize(selectionPoolSize),
       kCrossOverRate(crossOverRate),
       kMutationRate(mutationRate),
-      kNumberElites(numberElites),
       rng{std::random_device{}()} {
   restoreRngState(rng);
   VALIDATE();
@@ -222,14 +224,15 @@ GeneticSearch::GeneticSearch(
     size_t n,
     uint8_t crossOverRate,
     uint8_t mutationRate,
-    size_t numberElites)
+    size_t matingPoolSize,
+    size_t selectionPoolSize)
     : population(),
       lastBestConf(conf),
       kMaxPopulationSize(n),
-      kMatingPoolSize(n * 3),
+      kMatingPoolSize(matingPoolSize),
+      kSelectionPoolSize(selectionPoolSize),
       kCrossOverRate(crossOverRate),
       kMutationRate(mutationRate),
-      kNumberElites(numberElites),
       rng{std::random_device{}()} {
   restoreRngState(rng);
   VALIDATE();
@@ -301,13 +304,6 @@ void GeneticSearch::breed() {
   auto matingPool =
       stochasticUniversalSampling(computeAccumulatedFitness(population));
 
-  Population new_population;
-  new_population.reserve(kMatingPoolSize);
-  for (size_t c = 0; c < kNumberElites; ++c) {
-    new_population.push_back(
-        make_unique<CandidateConfiguration>(population.at(c)->configuration));
-  }
-
   auto select = [&]() -> TuningConfiguration& {
     auto idx = std::uniform_int_distribution<size_t>{
         size_t(0), matingPool.size() - 1}(rng);
@@ -323,45 +319,20 @@ void GeneticSearch::breed() {
     return dist(rng);
   };
 
-  while (new_population.size() < kMaxPopulationSize) {
+  while (selectionPool.size() < kSelectionPoolSize) {
     if (shouldCrossOver()) {
       auto parent1 = select();
       auto parent2 = select();
       auto parent3 = select();
-      new_population.emplace_back(make_unique<CandidateConfiguration>(
+      selectionPool.emplace_back(make_unique<CandidateConfiguration>(
           crossover(parent1, parent2, parent3)));
     } else {
-      new_population.emplace_back(
-          make_unique<CandidateConfiguration>(select()));
+      selectionPool.emplace_back(make_unique<CandidateConfiguration>(select()));
     }
   }
-  population = std::move(new_population);
 }
 
-void GeneticSearch::updateParameters() {
-  dropInvalidConfigurations(population);
-
-  // Sort population before taking any decision
-  std::sort(
-      population.begin(),
-      population.end(),
-      [](const std::unique_ptr<CandidateConfiguration>& a,
-         const std::unique_ptr<CandidateConfiguration>& b) {
-        checkRuntimeRecorded(a->runtime);
-        checkRuntimeRecorded(b->runtime);
-        return a->runtime < b->runtime;
-      });
-
-  // Update failsafe lastBestConf
-  lastBestConf =
-      population.size() > 0 ? population.front()->configuration : lastBestConf;
-  if (FLAGS_tuner_print_best) {
-    CudaMappingOptions options(
-        CudaMappingOptions::makeSingleThreadCudaMappingOptions());
-    lastBestConf.applyToCudaMappingOptions(options);
-    LOG(INFO) << "Best so far:\n" << options;
-  }
-
+void GeneticSearch::resetPopulationIfNotEnoughCandidates() {
   if (population.size() < kMinCandidatesForBreeding) {
     LOG_IF(ERROR, FLAGS_debug_tuner)
         << population.size() << " out of " << kMaxPopulationSize
@@ -380,12 +351,70 @@ void GeneticSearch::updateParameters() {
     // Don't lose the first one which was the best from before
     CHECK_LT(0, population.size());
     randomizePopulation(population.begin() + 1, population.end(), rng);
-    return;
   }
+}
 
+namespace {
+void sortByRuntime(GeneticSearch::Population& population) {
+  std::sort(
+      population.begin(),
+      population.end(),
+      [](const std::unique_ptr<CandidateConfiguration>& a,
+         const std::unique_ptr<CandidateConfiguration>& b) {
+        checkRuntimeRecorded(a->runtime);
+        checkRuntimeRecorded(b->runtime);
+        return a->runtime < b->runtime;
+      });
+}
+} // namespace
+
+void GeneticSearch::updateBestCandidate(const TuningConfiguration& c) {
+  lastBestConf = c;
+  if (FLAGS_tuner_print_best) {
+    CudaMappingOptions options(
+        CudaMappingOptions::makeSingleThreadCudaMappingOptions());
+    lastBestConf.applyToCudaMappingOptions(options);
+    LOG(INFO) << "Best so far:\n" << options;
+  }
+}
+
+void GeneticSearch::generateSelectionPool() {
+  dropInvalidConfigurations(population);
+  sortByRuntime(population);
+  updateBestCandidate(
+      population.size() > 0 ? population.front()->configuration : lastBestConf);
+  resetPopulationIfNotEnoughCandidates();
   breed();
-  for (int i = kNumberElites; i < population.size(); ++i) {
-    mutate(*population[i], kMutationRate, kMutateIterations, rng);
+  selectionPool.clear();
+  selectionPool.emplace_back(make_unique<CandidateConfiguration>(lastBestConf));
+  breed();
+  for (size_t i = 1; i < selectionPool.size(); ++i) {
+    mutate(*selectionPool[i], kMutationRate, kMutateIterations, rng);
+  }
+}
+
+void GeneticSearch::selectSurvivors() {
+  dropInvalidConfigurations(selectionPool);
+  sortByRuntime(selectionPool);
+  population.clear();
+  std::transform(
+      selectionPool.begin(),
+      selectionPool.begin() +
+          std::min(selectionPool.size(), kMaxPopulationSize),
+      std::back_inserter(population),
+      [](const std::unique_ptr<CandidateConfiguration>& c) {
+        return make_unique<CandidateConfiguration>(c->configuration);
+      });
+
+  if (selectionPool.size() < kMaxPopulationSize) {
+    auto numberMissing = kMaxPopulationSize - selectionPool.size();
+
+    for (size_t i = 0; i < numberMissing; ++i) {
+      selectionPool.emplace_back(
+          make_unique<CandidateConfiguration>(lastBestConf));
+    }
+    randomizePopulation(
+        selectionPool.end() - numberMissing, selectionPool.end(), rng);
   }
 }