From a4e724513d726d1d47dd46813c939fd5b5aaaa81 Mon Sep 17 00:00:00 2001
From: Theodoros Theodoridis <theodoridisgr@gmail.com>
Date: Wed, 28 Mar 2018 18:58:31 +0200
Subject: [PATCH] Allow using a user-defined key in the compilation cache

---
 include/tc/core/compilation_cache-inl.h | 13 +++++++++++--
 include/tc/core/flags.h                 |  3 +++
 src/core/flags.cc                       |  5 +++++
 3 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/include/tc/core/compilation_cache-inl.h b/include/tc/core/compilation_cache-inl.h
index 44032409c..77c28a479 100644
--- a/include/tc/core/compilation_cache-inl.h
+++ b/include/tc/core/compilation_cache-inl.h
@@ -95,6 +95,15 @@ void Cache<CC>::clear() {
   static_cast<CC*>(this)->entries_.clear();
 }
 
+namespace {
+std::string gpuKey() {
+  if (not FLAGS_cache_custom_hw_key.empty()) {
+    return FLAGS_cache_custom_hw_key;
+  }
+  return CudaGPUInfo::GPUInfo().GetCudaDeviceStr();
+}
+} // namespace
+
 template <typename C, typename InputTy> // deduces whether C is const or
 // non-const
 auto CudaCache::searchKernelImpl(
@@ -104,7 +113,7 @@ auto CudaCache::searchKernelImpl(
     const std::vector<InputTy>& inputs,
     const std::vector<InputTy>& outputs)
     -> decltype(c.searchKernel(id, options, inputs, outputs)) {
-  auto gpuStr = CudaGPUInfo::GPUInfo().GetCudaDeviceStr();
+  auto gpuStr = gpuKey();
   auto it = std::find_if(
       c.entries_.begin(), c.entries_.end(), [&](const CachedEntry& c) {
         using tc::operator==;
@@ -134,7 +143,7 @@ auto OptionsCache::searchKernelImpl(
     const std::vector<const DLTensor*>& inputs,
     const std::vector<const DLTensor*>& outputs)
     -> decltype(c.searchKernel(id, inputs, outputs)) {
-  auto gpuStr = CudaGPUInfo::GPUInfo().GetCudaDeviceStr();
+  auto gpuStr = gpuKey();
   auto it = std::find_if(
       c.entries_.begin(), c.entries_.end(), [&](const CachedEntry& c) {
         using tc::operator==;
diff --git a/include/tc/core/flags.h b/include/tc/core/flags.h
index fdd662f26..42c9308ad 100644
--- a/include/tc/core/flags.h
+++ b/include/tc/core/flags.h
@@ -54,6 +54,9 @@ DECLARE_uint32(tuner_gen_restore_number);
 DECLARE_bool(tuner_gen_log_generations);
 DECLARE_uint64(tuner_min_launch_total_threads);
 
+
+DECLARE_string(cache_custom_hw_key);
+
 // Misc
 DECLARE_int64(random_seed);
 DECLARE_bool(schedule_tree_verbose_validation);
diff --git a/src/core/flags.cc b/src/core/flags.cc
index 0042e7180..dff25470b 100644
--- a/src/core/flags.cc
+++ b/src/core/flags.cc
@@ -111,6 +111,11 @@ DEFINE_int64(
     -1,
     "The number of best candidates to restore from the proto cache");
 
+DEFINE_string(
+    cache_custom_hw_key,
+    "",
+    "Replace the \"hardware description\" with a custom string when querying the cache.");
+
 uint64_t initRandomSeed() {
   static std::mutex mut;
   static bool inited = false;