Allow using a user-defined key in the compilation cache

Author: Theodoros Theodoridis

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==;

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);

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;