Extract out RetrievalResult in standalone class

Author: nicolasvasilache

tc/autotuner/utils/utils.cc

@@ -65,7 +65,7 @@ std::vector<OptionsWithMedianTime> getOptionsAndMedianRuntimes(
       candidates.begin(),
       candidates.end(),
       std::back_inserter(c),
-      [](const OptionsCache::RetrievalResult& rr) -> OptionsWithMedianTime {
+      [](const OptionsCacheRetrievalResult& rr) -> OptionsWithMedianTime {
         return {std::move(rr.options), median(rr.recordedRuntimes)};
       });
   return c;

tc/core/cuda/cuda_compilation_cache.cc

@@ -170,7 +170,7 @@ const CudaCachedEntry* CudaCache::searchKernel(
   return searchKernelImpl(*this, id, options, inputs, outputs);
 }
 
-std::unique_ptr<CudaCache::RetrievalResult> CudaCache::retrieveKernel(
+std::unique_ptr<CudaCacheRetrievalResult> CudaCache::retrieveKernel(
     const std::string& id,
     const CudaMappingOptions& options,
     const std::vector<const DLTensor*>& inputs,
@@ -182,12 +182,12 @@ std::unique_ptr<CudaCache::RetrievalResult> CudaCache::retrieveKernel(
     return nullptr;
   }
   ++numberSuccessfulRetrievals;
-  return std::unique_ptr<CudaCache::RetrievalResult>(
-      new CudaCache::RetrievalResult{entry->values.cudaSource,
-                                     entry->values.kernelSpecializedName,
-                                     entry->values.kernelParameters,
-                                     entry->values.grid,
-                                     entry->values.block});
+  return std::unique_ptr<CudaCacheRetrievalResult>(
+      new CudaCacheRetrievalResult{entry->values.cudaSource,
+                                   entry->values.kernelSpecializedName,
+                                   entry->values.kernelParameters,
+                                   entry->values.grid,
+                                   entry->values.block});
 }
 
 void CudaCache::removeEntriesNotInOptionsCache(const OptionsCache& oc) {
@@ -228,7 +228,7 @@ std::unique_ptr<CudaMappingOptions> OptionsCache::retrieveBestOptions(
       new CudaMappingOptions(ret.front()));
 }
 
-std::vector<OptionsCache::RetrievalResult>
+std::vector<OptionsCacheRetrievalResult>
 OptionsCache::retrieveOptionsAndRuntimes(
     const std::string& id,
     const std::vector<const DLTensor*>& inputs,
@@ -533,7 +533,7 @@ void removeFromCudaCacheEntriesNotInOptionsCache(
   cc.removeEntriesNotInOptionsCache(oc);
 }
 
-std::unique_ptr<CudaCache::RetrievalResult> ManualCudaCache::retrieveKernel(
+std::unique_ptr<CudaCacheRetrievalResult> ManualCudaCache::retrieveKernel(
     const std::string& id,
     const std::vector<const DLTensor*>& inputs,
     const std::vector<const DLTensor*>& outputs) const {
@@ -544,12 +544,12 @@ std::unique_ptr<CudaCache::RetrievalResult> ManualCudaCache::retrieveKernel(
     return nullptr;
   }
   ++numberSuccessfulRetrievals;
-  return std::unique_ptr<CudaCache::RetrievalResult>(
-      new CudaCache::RetrievalResult{entry->values.cudaSource,
-                                     entry->values.kernelSpecializedName,
-                                     entry->values.kernelParameters,
-                                     entry->values.grid,
-                                     entry->values.block});
+  return std::unique_ptr<CudaCacheRetrievalResult>(
+      new CudaCacheRetrievalResult{entry->values.cudaSource,
+                                   entry->values.kernelSpecializedName,
+                                   entry->values.kernelParameters,
+                                   entry->values.grid,
+                                   entry->values.block});
 }
 
 ManualCudaCachedEntry* ManualCudaCache::searchKernel(

tc/core/cuda/cuda_compilation_cache.h

@@ -72,6 +72,14 @@ struct CudaCachedEntry {
   Values values;
 };
 
+struct CudaCacheRetrievalResult {
+  std::string source;
+  std::string specializedName;
+  std::vector<int> parameters;
+  Grid grid;
+  Block block;
+};
+
 /**
  * CudaCache stores the Cuda source of optimized kernels
  * A CudaCache holds multiple CudaCachedEntry's.
@@ -92,16 +100,9 @@ class CudaCache : public Cache<CudaCache, CudaCachedEntry> {
  public:
   using ProtobufType = CudaCacheProto;
   using CachedEntry = CudaCachedEntry;
+  using RetrievalResult = CudaCacheRetrievalResult;
   static std::shared_ptr<CudaCache>& getGlobalSharedCache();
 
-  struct RetrievalResult {
-    std::string source;
-    std::string specializedName;
-    std::vector<int> parameters;
-    Grid grid;
-    Block block;
-  };
-
  private:
   /**
    * SearchKernel (through SearchKernelImpl) searches op in the cache
@@ -161,7 +162,7 @@ class CudaCache : public Cache<CudaCache, CudaCachedEntry> {
    * Returns the cache entry that matches op (id, isl options, target device)
    * and inputs' shapes.
    */
-  std::unique_ptr<RetrievalResult> retrieveKernel(
+  std::unique_ptr<CudaCacheRetrievalResult> retrieveKernel(
       const std::string& id,
       const CudaMappingOptions& options,
       const std::vector<const DLTensor*>& inputs,
@@ -219,10 +220,16 @@ struct OptionsCachedEntry {
   std::vector<Values> values;
 };
 
+struct OptionsCacheRetrievalResult {
+  CudaMappingOptions options;
+  std::vector<Duration> recordedRuntimes;
+};
+
 class OptionsCache : public Cache<OptionsCache, OptionsCachedEntry> {
  public:
   using ProtobufType = OptionsCacheProto;
   using CachedEntry = OptionsCachedEntry;
+  using RetrievalResult = OptionsCacheRetrievalResult;
   static std::shared_ptr<OptionsCache>& getGlobalSharedCache();
 
  private:
@@ -256,10 +263,6 @@ class OptionsCache : public Cache<OptionsCache, OptionsCachedEntry> {
   OptionsCache(const OptionsCacheProto& buf);
 
   OptionsCacheProto toProtobuf() const;
-  struct RetrievalResult {
-    CudaMappingOptions options;
-    std::vector<Duration> recordedRuntimes;
-  };
 
   // returns the sum of cache entry sizes (that is a single cache entry can have
   // multiple options and profiling information associated with it)
@@ -272,7 +275,7 @@ class OptionsCache : public Cache<OptionsCache, OptionsCachedEntry> {
       const std::vector<const DLTensor*>& outputs,
       Duration runtime);
 
-  std::vector<RetrievalResult> retrieveOptionsAndRuntimes(
+  std::vector<OptionsCacheRetrievalResult> retrieveOptionsAndRuntimes(
       const std::string& id,
       const std::vector<const DLTensor*>& inputs,
       const std::vector<const DLTensor*>& outputs) const;
@@ -339,13 +342,16 @@ struct ManualCudaCachedEntry {
   Values values;
 };
 
+typedef CudaCacheRetrievalResult ManualCudaCacheRetrievalResult;
+
 /*
  * ManualCudaCache stores the manually injected source of Cuda kernels
  */
 class ManualCudaCache : public Cache<ManualCudaCache, ManualCudaCachedEntry> {
  public:
   using ProtobufType = ManualCudaCacheProto;
   using CachedEntry = ManualCudaCachedEntry;
+  using RetrievalResult = ManualCudaCacheRetrievalResult;
   static std::shared_ptr<ManualCudaCache>& getGlobalSharedCache();
 
  private:
@@ -401,7 +407,7 @@ class ManualCudaCache : public Cache<ManualCudaCache, ManualCudaCachedEntry> {
    *Returns the cache entry that matches
    *op(id, target device) and inputs' shapes.
    */
-  std::unique_ptr<CudaCache::RetrievalResult> retrieveKernel(
+  std::unique_ptr<ManualCudaCacheRetrievalResult> retrieveKernel(
       const std::string& id,
       const std::vector<const DLTensor*>& inputs,
       const std::vector<const DLTensor*>& outputs) const;

tc/core/cuda/cuda_tc_executor.cc

@@ -53,7 +53,7 @@ void CudaTcExecutor::compile(const tc::CudaMappingOptions& options) {
   }
   executionInfo_.options = options.toProtobufSerializedString();
 
-  auto cachedOp = [&]() -> std::unique_ptr<CudaCache::RetrievalResult> {
+  auto cachedOp = [&]() -> std::unique_ptr<CudaCacheRetrievalResult> {
     if (ManualCudaCache::cacheEnabled()) {
       auto rr = ManualCudaCache::getCache()->retrieveKernel(
           cacheKeyId_,