From ec9b1c7a64f20950a3296f3bba9aa4c5ecd71e44 Mon Sep 17 00:00:00 2001
From: David Beckingsale <davidbeckingsale@gmail.com>
Date: Mon, 23 Dec 2024 10:07:14 -0800
Subject: [PATCH] Allow variadic args in hash function

---
 lib/JitCache.hpp        | 7 ++++---
 lib/JitEngineDevice.hpp | 6 ++++--
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/lib/JitCache.hpp b/lib/JitCache.hpp
index 26361b38..ddb46194 100644
--- a/lib/JitCache.hpp
+++ b/lib/JitCache.hpp
@@ -34,10 +34,11 @@ inline hash_code hash_value(const proteus::RuntimeConstant &RC) {
 
 template <typename Function_t> class JitCache {
 public:
+  template<typename... Ts>
   uint64_t hash(StringRef ModuleUniqueId, StringRef FnName,
-                const RuntimeConstant *RC, int NumRuntimeConstants) const {
+                const RuntimeConstant *RC, int NumRuntimeConstants, Ts... args) const {
     ArrayRef<RuntimeConstant> Data(RC, NumRuntimeConstants);
-    auto HashValue = hash_combine(ExePath, ModuleUniqueId, FnName, Data);
+    auto HashValue = hash_combine(ExePath, ModuleUniqueId, FnName, Data, args...);
     return HashValue;
   }
 
@@ -122,4 +123,4 @@ template <typename Function_t> class JitCache {
 
 } // namespace proteus
 
-#endif
\ No newline at end of file
+#endif
diff --git a/lib/JitEngineDevice.hpp b/lib/JitEngineDevice.hpp
index aa253171..8ed93d67 100644
--- a/lib/JitEngineDevice.hpp
+++ b/lib/JitEngineDevice.hpp
@@ -439,8 +439,10 @@ JitEngineDevice<ImplT>::compileAndRun(
 
   getRuntimeConstantValues(KernelArgs, RCIndices, RCTypes, RCsVec);
 
-  uint64_t HashValue = CodeCache.hash(ModuleUniqueId, KernelName, RCsVec.data(),
-                                      NumRuntimeConstants);
+  uint64_t HashValue = CodeCache.hash(ModuleUniqueId,
+    KernelName, RCsVec.data(), NumRuntimeConstants, GridDim.x, GridDim.y,
+    GridDim.z, BlockDim.x, BlockDim.y, BlockDim.z);
+
   typename DeviceTraits<ImplT>::KernelFunction_t KernelFunc =
       CodeCache.lookup(HashValue);
   if (KernelFunc)