[C++ API] Consolidate autotuner/utils

This changeset consolidates autotuner/utils/* by moving the remaining
functionality into autotuner/utils.cc|h. This allows preparing for the
final step of refactoring and reduce the delta in the last commit.

This also helps centralize some current dependencies of autotuner/utils on
cuda_mapping_options.h, cuda_compilation_cache.h, dlpack.h and canonicalize.h.
These dependencies will go away in the last refactoring commit.

Lastly this commit also removes some dead code related to compilation and
options caches but cannot remove all such code because it is still needed
until we switch to the new API (in particular getOptionsAndMedianRuntimes that
will go away in a later commit).

Author: nicolasvasilache

tc/autotuner/CMakeLists.txt

@@ -9,8 +9,7 @@ if (WITH_CUDA)
     genetic_search.cc
     genetic_tuning_harness.cc
     parameters.cc
-    utils/printer.cc
-    utils/utils.cc)
+    utils.cc)
 
   target_include_directories(tc_autotuner PUBLIC ${PROJECT_SOURCE_DIR}/include)
   target_link_libraries(tc_autotuner PUBLIC ${ATEN_LIBRARIES} tc_cuda tc_proto)

tc/autotuner/genetic_autotuner.cc

@@ -20,6 +20,7 @@
 #include <csignal>
 #include <thread>
 
+#include "tc/autotuner/utils.h"
 #include "tc/core/cuda/cuda_compilation_cache.h"
 #include "tc/core/cuda/cuda_tc_executor.h"
 #include "tc/core/flags.h"

tc/autotuner/genetic_autotuner.h

@@ -20,7 +20,7 @@
 #include <vector>
 
 #include "tc/autotuner/genetic_tuning_harness.h"
-#include "tc/autotuner/utils/utils.h"
+#include "tc/autotuner/utils.h"
 #include "tc/core/execution_engine.h"
 #include "tc/lang/parser.h"
 

tc/autotuner/genetic_autotuner_aten.h

@@ -22,7 +22,7 @@
 #include "tc/aten/aten_compiler.h"
 #include "tc/autotuner/genetic_autotuner.h"
 #include "tc/autotuner/genetic_tuning_harness.h"
-#include "tc/autotuner/utils/utils.h"
+#include "tc/autotuner/utils.h"
 #include "tc/lang/parser.h"
 
 namespace tc {

tc/autotuner/genetic_tuning_harness.cc

@@ -24,8 +24,7 @@
 #include <cuda_runtime_api.h>
 #include <glog/stl_logging.h>
 
-#include "tc/autotuner/utils/printer.h"
-#include "tc/autotuner/utils/utils.h"
+#include "tc/autotuner/utils.h"
 #include "tc/core/cuda/cuda.h"
 #include "tc/core/cuda/cuda_mapping_options_cpp_printer.h"
 #include "tc/core/cuda/cuda_tc_executor.h"

tc/autotuner/genetic_tuning_harness.h

@@ -24,7 +24,7 @@
 
 #include "tc/autotuner/genetic_search.h"
 #include "tc/autotuner/parameters.h"
-#include "tc/autotuner/utils/printer.h"
+#include "tc/autotuner/utils.h"
 #include "tc/core/utils/dlpack.h"
 #include "tc/lang/parser.h"
 

tc/autotuner/utils/utils-inl.h renamed to tc/autotuner/utils-inl.h

@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 #pragma once
+
 #include <algorithm>
 #include <iterator>
 #include <type_traits>

tc/autotuner/utils/utils.cc renamed to tc/autotuner/utils.cc

@@ -13,33 +13,36 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+#include "tc/autotuner/utils.h"
+
 #include <algorithm>
 #include <cmath>
+#include <sstream>
+
+#include <glog/stl_logging.h>
 
-#include "tc/aten/aten_compiler.h"
-#include "tc/autotuner/utils/utils.h"
 #include "tc/core/cuda/cuda_compilation_cache.h"
+#include "tc/core/cuda/cuda_mapping_options.h"
+#include "tc/core/flags.h"
+#include "tc/core/utils/dlpack.h"
 #include "tc/core/utils/math.h"
+#include "tc/core/utils/time.h"
 #include "tc/lang/canonicalize.h"
 
 namespace tc {
 namespace autotune {
 
-namespace {
-std::vector<std::size_t> firstPowers2(std::size_t n) {
-  std::vector<std::size_t> powers(n + 1);
+std::vector<std::size_t> powers2andCeilDivisors(std::size_t val) {
+  auto numPowers = static_cast<std::size_t>(std::ceil(std::log2(val)));
+  // 1. generate `numPowers' powers of 2
+  std::vector<std::size_t> res(numPowers + 1);
   std::size_t p = 1;
-  std::generate(powers.begin(), powers.end(), [p]() mutable {
+  std::generate(res.begin(), res.end(), [p]() mutable {
     auto old_p = p;
     p *= 2;
     return old_p;
   });
-  return powers;
-}
-} // namespace
-
-std::vector<std::size_t> powers2andCeilDivisors(std::size_t val) {
-  auto res = firstPowers2(static_cast<std::size_t>(std::ceil(std::log2(val))));
+  // 2. additionally insert ceil(val / powers2)
   res.reserve(res.size() * 2);
   for (std::size_t i = 0, s = res.size(); i < s; ++i) {
     if (res[i] > val) {
@@ -52,6 +55,11 @@ std::vector<std::size_t> powers2andCeilDivisors(std::size_t val) {
   return res;
 }
 
+struct OptionsWithMedianTime {
+  CudaMappingOptions options;
+  Duration medianRuntime;
+};
+
 std::vector<OptionsWithMedianTime> getOptionsAndMedianRuntimes(
     const lang::CanonicalTcString& id,
     const std::vector<const DLTensor*>& inputs,
@@ -109,5 +117,92 @@ llvm::Optional<CudaMappingOptions> getBestOptions(
   return llvm::Optional<CudaMappingOptions>{};
 }
 
+void Printer::record(Duration runtime) {
+  std::lock_guard<std::mutex> lock(runtimesMtx_);
+  runtimes_.push_back(runtime);
+}
+
+namespace {
+uint64_t toMicroseconds(const Duration& d) {
+  return std::chrono::duration_cast<std::chrono::microseconds>(d).count();
+}
+} // namespace
+
+void Printer::printLoop() {
+  while (true) {
+    std::this_thread::sleep_for(std::chrono::seconds(1));
+
+    std::stringstream ss;
+    ss << "Iteration " << iteration_;
+    ss << "\tJobs(Compiled, Evaluated)/total  ("
+       << std::min(total_, currentCompilationJob_.load()) << ", "
+       << std::min(total_, numEvaluations_.load()) << ")/" << total_;
+
+    {
+      std::lock_guard<std::mutex> lock(runtimesMtx_);
+      if (not runtimes_.empty()) {
+        std::sort(runtimes_.begin(), runtimes_.end());
+        auto best = toMicroseconds(runtimes_.front());
+        auto median = toMicroseconds(runtimes_.at(runtimes_.size() / 2));
+        auto worst = toMicroseconds(runtimes_.back());
+        ss << "   (best/median/worst)us: " << best << '/' << median << '/'
+           << worst;
+      }
+    }
+    // XXX: platform specific erase current line and move cursor to begining
+    // of line. Currently works with python/C++ both.
+    std::cout << "\u001b[2K\r" << ss.str() << std::flush;
+    LOG_IF(INFO, FLAGS_debug_tuner) << "\u001b[2K\r" << ss.str() << std::endl;
+
+    if (stopPrinting_.load()) {
+      // Print one more time to flush
+      // XXX: platform specific erase current line and move cursor to begining
+      // of line. Currently works with python/C++ both.
+      std::cout << "\u001b[2K\r" << ss.str() << std::flush;
+      LOG_IF(INFO, FLAGS_debug_tuner) << "\u001b[2K\r" << ss.str() << std::endl;
+      // commit line so it does not get erased at the next iteration
+      std::cerr << std::endl;
+      return;
+    }
+  }
+}
+
+Printer::Printer(
+    size_t iteration,
+    size_t total,
+    const std::atomic_size_t& currentCompilationJob,
+    const std::atomic_size_t& numEvaluations)
+    : iteration_(iteration),
+      printerThread_([this]() { printLoop(); }),
+      total_(total),
+      currentCompilationJob_(currentCompilationJob),
+      numEvaluations_(numEvaluations) {}
+
+Printer::~Printer() {
+  stop();
+  printerThread_.join();
+}
+
+void Printer::stop() {
+  stopPrinting_.store(true);
+}
+
+void Printer::printAll() {
+  auto runtimes = [this]() {
+    std::lock_guard<std::mutex> lock(runtimesMtx_);
+    std::sort(runtimes_.begin(), runtimes_.end());
+    std::vector<uint64_t> runtimes;
+    runtimes.reserve(runtimes_.size());
+    std::transform(
+        runtimes_.begin(),
+        runtimes_.end(),
+        std::back_inserter(runtimes),
+        toMicroseconds);
+    return runtimes;
+  }();
+  LOG_IF(INFO, FLAGS_debug_tuner)
+      << "\n [TUNER][ITERATION LOG] median times of each candidate (in us) "
+      << runtimes << std::endl;
+}
 } // namespace autotune
 } // namespace tc

tc/autotuner/utils/utils.h renamed to tc/autotuner/utils.h

@@ -14,15 +14,21 @@
  * limitations under the License.
  */
 #pragma once
+
+#include <atomic>
+#include <iostream>
+#include <mutex>
+#include <numeric>
+#include <thread>
 #include <vector>
 
-#include "tc/core/cuda/cuda.h"
+#include <llvm/ADT/Optional.h>
+
+#include "tc/core/cuda/cuda_compilation_cache.h"
 #include "tc/core/cuda/cuda_mapping_options.h"
 #include "tc/core/utils/dlpack.h"
+#include "tc/core/utils/time.h"
 #include "tc/lang/canonicalize.h"
-#include "tc/lang/tree.h"
-
-#include <llvm/ADT/Optional.h>
 
 namespace tc {
 namespace autotune {
@@ -35,10 +41,6 @@ std::vector<std::size_t> powers2andCeilDivisors(std::size_t val);
 template <typename Vector, typename... Vectors>
 Vector mergeVectors(Vector&& v, Vectors&&... vs);
 
-/// The following API allows interacting with the autotuner caches.
-/// Caches generally take arbitrary strings for keys.
-/// The autotuner uses a canonicalized TC expression to load / store into
-/// caches. Add a layer of type safety to interact with these.
 std::vector<CudaMappingOptions> restoreCandidates(
     const lang::CanonicalTcString& tc,
     const std::vector<const DLTensor*>& inputs,
@@ -49,16 +51,38 @@ llvm::Optional<CudaMappingOptions> getBestOptions(
     const std::vector<const DLTensor*>& inputs,
     const std::vector<const DLTensor*>& outputs);
 
-struct OptionsWithMedianTime {
-  CudaMappingOptions options;
-  Duration medianRuntime;
-};
+/**
+ * Helper class to pretty print autotuning progress
+ */
+class Printer {
+ public:
+  Printer(
+      size_t iteration,
+      size_t total,
+      const std::atomic_size_t& currentCompilationJob,
+      const std::atomic_size_t& numEvaluations);
+  ~Printer();
 
-std::vector<OptionsWithMedianTime> getOptionsAndMedianRuntimes(
-    const lang::CanonicalTcString& id,
-    const std::vector<const DLTensor*>& inputs);
+  void record(Duration runtime);
+  void stop();
+
+  void printAll();
+
+ private:
+  void printLoop();
 
+  size_t iteration_;
+  std::vector<Duration> runtimes_;
+  mutable std::mutex runtimesMtx_;
+
+  std::atomic_bool stopPrinting_{false};
+  std::thread printerThread_;
+
+  const size_t total_;
+  const std::atomic_size_t& currentCompilationJob_;
+  const std::atomic_size_t& numEvaluations_;
+};
 } // namespace autotune
 } // namespace tc
 
-#include "tc/autotuner/utils/utils-inl.h"
+#include "tc/autotuner/utils-inl.h"