Add AnyKernelDescriptor class to replace KernelRegistry::CacheKey

Author: stijnh

include/kernel_launcher/registry.h

@@ -19,82 +19,91 @@ struct KernelDescriptor {
     }
 };
 
-struct KernelRegistry {
+struct AnyKernelDescriptor {
+    AnyKernelDescriptor(AnyKernelDescriptor&&) noexcept = default;
+    AnyKernelDescriptor(const AnyKernelDescriptor&) = default;
+
+    template<typename D>
+    AnyKernelDescriptor(D&& descriptor) {
+        using T = typename std::decay<D>::type;
+        descriptor_type_ = type_of<T>();
+        descriptor_ = std::make_unique<T>(std::forward<D>(descriptor));
+        hash_ = hash_fields(descriptor_type_, descriptor_->hash());
+    }
+
+    const KernelDescriptor& descriptor() const {
+        return *descriptor_;
+    }
+
+    hash_t hash() const {
+        return hash_;
+    }
+
+    bool operator==(const AnyKernelDescriptor& that) const {
+        return that.hash_ == hash_ && that.descriptor_type_ == descriptor_type_
+            && that.descriptor_->equals(*descriptor_);
+    }
+
+    bool operator!=(const AnyKernelDescriptor& that) const {
+        return !(*this == that);
+    }
+
   private:
-    struct CacheKey {
-        struct hasher {
-            size_t operator()(const CacheKey& key) const;
-        };
-
-        struct equals {
-            bool operator()(const CacheKey& lhs, const CacheKey& rhs) const;
-        };
-
-        CacheKey(CacheKey&&) noexcept = default;
-
-        template<typename D>
-        explicit CacheKey(D&& descriptor) {
-            using T = typename std::decay<D>::type;
-            descriptor_type_ = type_of<T>();
-            descriptor_ = std::make_unique<T>(std::forward<D>(descriptor));
-            hash_ = hash_fields(descriptor_type_, descriptor_->hash());
-        }
-
-        const KernelDescriptor& descriptor() const {
-            return *descriptor_;
-        }
-
-      private:
-        hash_t hash_;
-        TypeInfo descriptor_type_;
-        std::unique_ptr<KernelDescriptor> descriptor_;
-    };
-
-  public:
+    hash_t hash_;
+    TypeInfo descriptor_type_;
+    std::shared_ptr<KernelDescriptor> descriptor_;
+};
+}
+
+namespace std {
+template <>
+struct hash<kernel_launcher::AnyKernelDescriptor> {
+    size_t operator()(const kernel_launcher::AnyKernelDescriptor& d) const {
+        return d.hash();
+    }
+};
+}
+
+namespace kernel_launcher {
+struct KernelRegistry {
     explicit KernelRegistry(
         Compiler compiler = default_compiler(),
         WisdomSettings settings = default_wisdom_settings()) :
         compiler_(std::move(compiler)),
         settings_(std::move(settings)) {}
 
-    template<typename D>
-    WisdomKernel& lookup(D&& descriptor) const {
-        return lookup_internal(CacheKey(std::forward<D>(descriptor)));
+    WisdomKernel& lookup(AnyKernelDescriptor descriptor) const {
+        return lookup_internal(std::move(descriptor));
     }
 
-    template<typename D>
-    WisdomKernelLaunch
-    instantiate(D&& descriptor, cudaStream_t stream, ProblemSize problem_size)
-        const {
-        return lookup(std::forward<D>(descriptor))
-            .instantiate(stream, problem_size);
+    WisdomKernelLaunch instantiate(
+        AnyKernelDescriptor descriptor,
+        cudaStream_t stream,
+        ProblemSize problem_size) const {
+        return lookup(std::move(descriptor)).instantiate(stream, problem_size);
     }
 
-    template<typename D>
-    WisdomKernelLaunch
-    operator()(D&& descriptor, cudaStream_t stream, ProblemSize problem_size)
-        const {
-        return instantiate(std::forward<D>(descriptor), stream, problem_size);
+    WisdomKernelLaunch operator()(
+        AnyKernelDescriptor descriptor,
+        cudaStream_t stream,
+        ProblemSize problem_size) const {
+        return instantiate(std::move(descriptor), stream, problem_size);
     }
 
-    template<typename D>
     WisdomKernelLaunch
-    operator()(D&& descriptor, ProblemSize problem_size) const {
-        return instantiate(std::forward<D>(descriptor), nullptr, problem_size);
+    operator()(AnyKernelDescriptor descriptor, ProblemSize problem_size) const {
+        return instantiate(std::move(descriptor), nullptr, problem_size);
     }
 
   private:
-    WisdomKernel& lookup_internal(CacheKey key) const;
+    WisdomKernel& lookup_internal(AnyKernelDescriptor key) const;
 
     Compiler compiler_;
     WisdomSettings settings_;
     mutable std::mutex mutex_;
-    mutable std::unordered_map<
-        CacheKey,
-        std::unique_ptr<WisdomKernel>,
-        CacheKey::hasher,
-        CacheKey::equals>
-        cache_ = {};
+    mutable std::
+        unordered_map<AnyKernelDescriptor, std::unique_ptr<WisdomKernel>>
+            cache_ = {};
 };
 
 const KernelRegistry& default_registry();
@@ -113,4 +122,5 @@ WisdomKernelLaunch launch(D&& descriptor, ProblemSize size) {
 
 }  // namespace kernel_launcher
 
+
 #endif  //KERNEL_LAUNCHER_CACHE_H

src/registry.cpp

@@ -11,19 +11,7 @@ const KernelRegistry& default_registry() {
     return *global_default_registry;
 }
 
-size_t KernelRegistry::CacheKey::hasher::operator()(const CacheKey& key) const {
-    return key.hash_;
-}
-
-bool KernelRegistry::CacheKey::equals::operator()(
-    const CacheKey& lhs,
-    const CacheKey& rhs) const {
-    return lhs.hash_ == rhs.hash_
-        && lhs.descriptor_type_ == rhs.descriptor_type_
-        && lhs.descriptor_->equals(*rhs.descriptor_);
-}
-
-WisdomKernel& KernelRegistry::lookup_internal(CacheKey key) const {
+WisdomKernel& KernelRegistry::lookup_internal(AnyKernelDescriptor key) const {
     std::lock_guard<std::mutex> guard(mutex_);
 
     auto it = cache_.find(key);