Typesafe autotuner cache interaction

This PR introduces a simple explicit CanonicalTcString type that specializes std::string.
All functions in tc::autotune::utils use this type.
Update tests and benchmarks to comply to the modified API.

Author: nicolasvasilache

benchmarks/benchmark_fixture.h

@@ -36,6 +36,7 @@
 #include "tc/core/cuda/cuda_tc_executor.h"
 #include "tc/core/flags.h"
 #include "tc/core/scope_guard.h"
+#include "tc/lang/canonicalize.h"
 
 #include <cublas_v2.h> // Must be the same as Caffe2
 #include <cuda_runtime_api.h>
@@ -270,7 +271,8 @@ struct Benchmark : public ::testing::Test {
       auto inputsPair = tc::toConstDlpackTensors(inputs);
       auto outputs = atCompl.inferOutputTensorInfo(name, inputs);
       tc::ScopeGuard g([&]() { tc::deleteDlmTensors(inputsPair.second); });
-      return tc::autotune::restoreCandidates(name, inputsPair.first, outputs);
+      return tc::autotune::restoreCandidates(
+          lang::canonicalTc(tc), inputsPair.first, outputs);
     }();
     auto handle = atCompl.compile(name, inputs, mappingOptions[0]);
     std::vector<at::Tensor> outputs;

include/tc/autotuner/utils/utils.h

@@ -21,18 +21,14 @@
 #include "tc/core/cuda/cuda.h"
 #include "tc/core/cuda/cuda_mapping_options.h"
 #include "tc/core/utils/dlpack.h"
+#include "tc/lang/canonicalize.h"
 #include "tc/lang/tree.h"
 
 #include <llvm/ADT/Optional.h>
 
 namespace tc {
 namespace autotune {
 
-struct OptionsWithMedianTime {
-  CudaMappingOptions options;
-  Duration medianRuntime;
-};
-
 /// Returns all the powers of 2 up to the first one that is larger than val
 /// and the result of ceil(val/pow2) for each of those powers of 2 (except for
 /// the larger one)
@@ -41,26 +37,28 @@ std::vector<std::size_t> powers2andCeilDivisors(std::size_t val);
 template <typename Vector, typename... Vectors>
 Vector mergeVectors(Vector&& v, Vectors&&... vs);
 
-std::vector<OptionsWithMedianTime> getOptionsAndMedianRuntimes(
-    const std::string& id,
-    const std::vector<const DLTensor*>& inputs);
-
+/// 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 std::string& id,
-    const std::vector<const DLTensor*>& inputs,
-    const std::vector<const DLTensor*>& outputs);
-
-std::vector<CudaMappingOptions> restoreCandidates(
-    const lang::TreeRef& tc,
+    const lang::CanonicalTcString& tc,
     const std::vector<const DLTensor*>& inputs,
     const std::vector<const DLTensor*>& outputs);
 
 llvm::Optional<CudaMappingOptions> getBestOptions(
-    const std::string& id,
+    const lang::CanonicalTcString& id,
     const std::vector<const DLTensor*>& inputs,
     const std::vector<const DLTensor*>& outputs);
 
-std::string canonicalTC(const lang::TreeRef& tc);
+struct OptionsWithMedianTime {
+  CudaMappingOptions options;
+  Duration medianRuntime;
+};
+
+std::vector<OptionsWithMedianTime> getOptionsAndMedianRuntimes(
+    const lang::CanonicalTcString& id,
+    const std::vector<const DLTensor*>& inputs);
 
 } // namespace autotune
 } // namespace tc

include/tc/core/tc_executor.h

@@ -25,7 +25,7 @@
 #include "tc/core/utils/dlpack.h"
 #include "tc/core/utils/time.h"
 
-#include "tc/lang/parser.h"
+#include "tc/lang/canonicalize.h"
 
 namespace tc {
 
@@ -124,7 +124,7 @@ class TcExecutor {
 
   tc2halide::HalideComponents halideComponents_;
   lang::TreeRef tcTree_;
-  std::string cacheKeyId;
+  lang::CanonicalTcString cacheKeyId;
 };
 
 // templating to match both const and non-const DLTensor pointers

include/tc/lang/canonicalize.h

@@ -17,14 +17,16 @@
 
 #include <string>
 
+#include "tc/lang/parser.h"
+#include "tc/lang/sema.h"
 #include "tc/lang/tree.h"
 #include "tc/lang/tree_views.h"
 
 namespace lang {
 
 // takes a tree after semantic analysis and create
 // a canonicalized version that is agnostic to the choice of identifiers
-TreeRef canonicalize(TreeRef tree) {
+inline TreeRef canonicalize(TreeRef tree) {
   struct Context {
     std::unordered_map<std::string, std::string> identMap;
     std::string rename(const std::string& name) {
@@ -53,4 +55,19 @@ TreeRef canonicalize(TreeRef tree) {
   Context ctx;
   return ctx.apply(tree);
 }
+
+struct CanonicalTcString : public std::string {
+  explicit CanonicalTcString(const std::string& s) : std::string(s) {}
+};
+
+inline CanonicalTcString canonicalTc(const lang::TreeRef& tc) {
+  std::stringstream ss;
+  // TODO: use tcFormat when more robust
+  ss << lang::canonicalize(lang::Sema().checkFunction(tc));
+  return CanonicalTcString(ss.str());
+}
+
+inline CanonicalTcString canonicalTc(const std::string& tc) {
+  return canonicalTc(lang::Parser(tc).parseFunction());
+}
 } // namespace lang

src/autotuner/genetic_autotuner.cc

@@ -81,7 +81,7 @@ std::vector<CudaMappingOptions> GeneticAutotuner::load(
   ee.define(tc_);
   auto outputs = ee.inferOutputTensorInfo(tcName, inputs);
   return tc::autotune::restoreCandidates(
-      tcNameMap_.at(tcName), inputs, outputs);
+      canonicalTc(tcNameMap_.at(tcName)), inputs, outputs);
 }
 
 namespace {
@@ -186,7 +186,7 @@ llvm::Optional<CudaMappingOptions> GeneticAutotuner::tune(
 
   CHECK_GT(inputs.size(), 0);
   return tc::autotune::getBestOptions(
-      canonicalTC(tcNameMap_.at(tcName)), inputs.begin()->second, outputPtrs);
+      canonicalTc(tcNameMap_.at(tcName)), inputs.begin()->second, outputPtrs);
 }
 
 } // namespace detail

src/autotuner/utils/utils.cc

@@ -21,8 +21,6 @@
 #include "tc/core/cuda/cuda_compilation_cache.h"
 #include "tc/core/utils/math.h"
 #include "tc/lang/canonicalize.h"
-#include "tc/lang/parser.h"
-#include "tc/lang/sema.h"
 
 namespace tc {
 namespace autotune {
@@ -55,7 +53,7 @@ std::vector<std::size_t> powers2andCeilDivisors(std::size_t val) {
 }
 
 std::vector<OptionsWithMedianTime> getOptionsAndMedianRuntimes(
-    const std::string& id,
+    const lang::CanonicalTcString& id,
     const std::vector<const DLTensor*>& inputs,
     const std::vector<const DLTensor*>& outputs) {
   auto candidates =
@@ -73,18 +71,11 @@ std::vector<OptionsWithMedianTime> getOptionsAndMedianRuntimes(
   return c;
 }
 
-std::string canonicalTC(const lang::TreeRef& tc) {
-  std::stringstream ss;
-  ss << lang::canonicalize(lang::Sema().checkFunction(tc));
-  return ss.str();
-}
-
 std::vector<CudaMappingOptions> restoreCandidates(
-    const lang::TreeRef& tc,
+    const lang::CanonicalTcString& tc,
     const std::vector<const DLTensor*>& inputs,
     const std::vector<const DLTensor*>& outputs) {
-  auto candidates =
-      getOptionsAndMedianRuntimes(canonicalTC(tc), inputs, outputs);
+  auto candidates = getOptionsAndMedianRuntimes(tc, inputs, outputs);
   LOG_IF(INFO, candidates.size() < FLAGS_tuner_gen_restore_number)
       << "Requested " << FLAGS_tuner_gen_restore_number
       << " candidates but there are only " << candidates.size() << " in cache.";
@@ -106,15 +97,8 @@ std::vector<CudaMappingOptions> restoreCandidates(
   return res;
 }
 
-std::vector<CudaMappingOptions> restoreCandidates(
-    const std::string& tc,
-    const std::vector<const DLTensor*>& inputs,
-    const std::vector<const DLTensor*>& outputs) {
-  return restoreCandidates(lang::Parser(tc).parseFunction(), inputs, outputs);
-}
-
 llvm::Optional<CudaMappingOptions> getBestOptions(
-    const std::string& id,
+    const lang::CanonicalTcString& id,
     const std::vector<const DLTensor*>& inputs,
     const std::vector<const DLTensor*>& outputs) {
   auto bestOptions =

src/core/tc_executor.cc

@@ -20,8 +20,6 @@
 
 #include "tc/core/utils/dlpack.h"
 #include "tc/lang/canonicalize.h"
-#include "tc/lang/parser.h"
-#include "tc/lang/sema.h"
 
 namespace tc {
 
@@ -32,13 +30,6 @@ int toTypeToken(DLDataType dtype) {
   return lang::TypeInfo(lang::TypeInfo::Code(dtype.code), dtype.bits)
       .toScalarToken();
 }
-
-std::string canonicalizedTc(const lang::TreeRef tcDefinition) {
-  std::stringstream ss;
-  ss << canonicalize(lang::Sema().checkFunction(tcDefinition));
-  return ss.str();
-}
-
 } // namespace
 
 TcExecutor::TcExecutor(
@@ -49,7 +40,8 @@ TcExecutor::TcExecutor(
     : identifier(id),
       inputsInfo(dlutils::makeDLTensorVector(inputsInfo)),
       options(options),
-      tcTree_(tcDefinition) {
+      tcTree_(tcDefinition),
+      cacheKeyId(lang::canonicalTc(tcDefinition)) {
   executionInfo_.kernelName = lang::Def(tcTree_).name().name();
   halideComponents_ =
       tc2halide::translate(isl::with_exceptions::globalIslCtx(), tcTree_);
@@ -58,7 +50,6 @@ TcExecutor::TcExecutor(
   // TODO: check if this is wrong, packed tensors may  have 0 strides stored
   executionInfo_.outputsInfo =
       tc::inferOutputTensorInfo(halideComponents_, inputsInfo);
-  cacheKeyId = canonicalizedTc(tcDefinition);
 }
 
 TcExecutor::~TcExecutor() {}

test/cuda/test_autotuner_utility.cc

@@ -20,6 +20,7 @@
 #include "tc/core/cuda/cuda_compilation_cache.h"
 #include "tc/core/cuda/cuda_tc_executor.h"
 #include "tc/core/scope_guard.h"
+#include "tc/lang/canonicalize.h"
 
 using namespace tc;
 using namespace autotune;
@@ -54,7 +55,7 @@ std::vector<CudaMappingOptions> restoreCandidates(
   });
 
   return tc::autotune::restoreCandidates(
-      tc, inputsPair.first, outputsPair.first);
+      lang::canonicalTc(tc), inputsPair.first, outputsPair.first);
 }
 
 TEST(RestoreCandidates, NoCache) {

test/test_lang.cc

@@ -147,12 +147,6 @@ TreeRef loadText(const std::string& text) {
   return Sema().checkFunction(Parser(text).parseFunction());
 }
 
-std::string canonicalText(const std::string& text) {
-  std::stringstream ss;
-  ss << canonicalize(loadText(text));
-  return ss.str();
-}
-
 void testTcFormat() {
   static std::ios_base::Init initIostreams;
   auto source = R"(def fun2(float(B, N, M) X, float(B, M, K) Y) -> (Q) {
@@ -334,7 +328,7 @@ int main(int argc, char** argv) {
             Q(b, ii, j) += X(b, ii, k) * Y(b, k, j)
     }
   )";
-  ASSERT(canonicalText(option_one) == canonicalText(option_two));
+  ASSERT(lang::canonicalTc(option_one) == lang::canonicalTc(option_two));
 
   testTcFormat();
 

test/test_tc_mapper_harness-inl.h

@@ -59,11 +59,8 @@ struct TcMapperTest : public ::testing::Test {
       tc::deleteDlmTensors(outputDLTensorsPair.second);
     });
     // Check that cache insertion worked properly (with canonicalized TC)
-    std::stringstream ss;
-    ss << lang::canonicalize(
-        lang::Sema().checkFunction(lang::Parser(tc).parseFunction()));
     auto cached = CacheType::getCache()->retrieveKernel(
-        ss.str(),
+        lang::canonicalTc(tc),
         mappingOptions,
         inputDLTensorsPair.first,
         outputDLTensorsPair.first);