Improved hashing performance and simplified compatible_auto_opaque

Author: DoeringChristian

src/python/freeze.cpp

@@ -30,9 +30,9 @@ struct ProfilerPhase {
     std::string m_message;
     ProfilerPhase(const char *message) : m_message(message) {
         jit_log(LogLevel::Debug, "profiler start: %s", message);
-#if defined(DRJIT_ENABLE_NVTX)
+// #if defined(DRJIT_ENABLE_NVTX)
         jit_profile_range_push(message);
-#endif
+// #endif
     }
 
     ProfilerPhase(const drjit::TraversableBase *traversable) {
@@ -47,9 +47,9 @@ struct ProfilerPhase {
     }
 
     ~ProfilerPhase() {
-#if defined(DRJIT_ENABLE_NVTX)
+// #if defined(DRJIT_ENABLE_NVTX)
         jit_profile_range_pop();
-#endif
+// #endif
         jit_log(LogLevel::Debug, "profiler end: %s", m_message.c_str());
     }
 };
@@ -208,53 +208,13 @@ uint32_t FlatVariables::add_jit_index(uint32_t index) {
  * consists of changes in literal values.
  */
 bool compatible_auto_opaque(FlatVariables &cur, FlatVariables &prev){
+    // NOTE: We only test the size of the layout, as a full test is somewhat
+    // expensive, and the worst case is that we make too many variables opaque,
+    // which does not impact correctness. If this causes problems, more
+    // extensive tests might have to be reintroduced.
     if (cur.layout.size() != prev.layout.size()) {
         return false;
     }
-    for (uint32_t i = 0; i < cur.layout.size(); i++) {
-        Layout &cur_layout  = cur.layout[i];
-        Layout &prev_layout = prev.layout[i];
-
-        if (((bool) cur_layout.type != (bool) prev_layout.type) ||
-            !(cur_layout.type.equal(prev_layout.type))){
-            // jit_log(LogLevel::Warn, "type");
-            return false;
-        }
-
-        if (cur_layout.num != prev_layout.num){
-            // jit_log(LogLevel::Warn, "num");
-            return false;
-        }
-
-        if (cur_layout.fields.size() != prev_layout.fields.size()){
-            // jit_log(LogLevel::Warn, "fields.size()");
-            return false;
-        }
-
-        for (uint32_t i = 0; i < cur_layout.fields.size(); ++i) {
-            if (!(cur_layout.fields[i].equal(prev_layout.fields[i]))){
-                // jit_log(LogLevel::Warn, "fields[%u]", i);
-                return false;
-            }
-        }
-
-        // if (cur_layout.index != prev_layout.index)
-        //     return false;
-
-        // if (cur_layout.flags != prev_layout.flags)
-        //     return false;
-
-        if (cur_layout.vt != prev_layout.vt){
-            // jit_log(LogLevel::Warn, "vt");
-            return false;
-        }
-
-        if (((bool) cur_layout.py_object != (bool) prev_layout.py_object) ||
-            !cur_layout.py_object.equal(prev_layout.py_object)){
-            // jit_log(LogLevel::Warn, "py_object");
-            return false;
-        }
-    }
     return true;
 }
 
@@ -1477,12 +1437,12 @@ bool log_diff(LogLevel level, const FlatVariables &curr,
 
     return true;
 }
-inline void hash_combine(size_t &seed, size_t value) {
+inline void hash_combine(uint64_t &seed, uint64_t value) {
     /// From CityHash (https://github.com/google/cityhash)
-    const size_t mult = 0x9ddfea08eb382d69ull;
-    size_t a          = (value ^ seed) * mult;
+    const uint64_t mult = 0x9ddfea08eb382d69ull;
+    uint64_t a          = (value ^ seed) * mult;
     a ^= (a >> 47);
-    size_t b = (seed ^ a) * mult;
+    uint64_t b = (seed ^ a) * mult;
     b ^= (b >> 47);
     seed = b * mult;
 }
@@ -1491,46 +1451,39 @@ size_t
 FlatVariablesHasher::operator()(const std::shared_ptr<FlatVariables> &key) const {
     ProfilerPhase profiler("hash");
     // Hash the layout
-    // NOTE: string hashing seems to be less efficient
-    size_t hash = key->layout.size();
-    for (const Layout &layout : key->layout) {
-        hash_combine(hash, layout.num);
-        hash_combine(hash, layout.fields.size());
-        hash_combine(hash, (size_t) layout.flags);
-        hash_combine(hash, (size_t) layout.index);
-        hash_combine(hash, (size_t) layout.literal);
-        hash_combine(hash, (size_t) layout.vt);
-        if (layout.type)
-            hash_combine(hash, nb::hash(layout.type));
-        if (layout.py_object){
-            PyObject *ptr = layout.py_object.ptr();
-            uint32_t object_hash;
-            Py_hash_t rv = PyObject_Hash(ptr);
-
-            // Try to hash the object, and otherwise fallback to ``id()``
-            if (rv == -1 && PyErr_Occurred()) {
-                PyErr_Clear();
-                object_hash = (uintptr_t) ptr;
-            } else {
-                object_hash = rv;
-            }
 
-            hash_combine(hash, object_hash);
-        }
-        for (auto &field : layout.fields) {
-            hash_combine(hash, nb::hash(field));
-        }
+    // TODO: Maybe we can use xxh by first collecting in vector<uint64_t>?
+
+    uint64_t hash = (uint64_t) (key->layout.size() << 32) |
+                    (uint64_t) (key->var_layout.size() << 2);
+
+    for (const Layout &layout : key->layout) {
+        // if layout.fields is not 0 then layout.num == layout.fields.size()
+        // therefore we can omit layout.fields.size().
+        hash_combine(hash,
+                     ((uint64_t) layout.num << 32) | ((uint64_t) layout.index));
+        hash_combine(hash, layout.literal);
+        hash_combine(hash,
+                     ((uint64_t) layout.flags << 32) | ((uint64_t) layout.vt));
+        hash_combine(
+            hash,
+            ((uint64_t) (uint32_t) (uintptr_t) layout.type.ptr() << 32) |
+                ((uint64_t) (uint32_t) (uintptr_t) layout.py_object.ptr()));
+        hash_combine(hash, (uint64_t) layout.vt);
+        for (auto &field : layout.fields)
+             hash_combine(hash, (uintptr_t)field.ptr());
     }
 
     for (const VarLayout &layout : key->var_layout) {
-        hash_combine(hash, (size_t) layout.vt);
-        hash_combine(hash, (size_t) layout.vs);
-        hash_combine(hash, (size_t) layout.flags);
-        hash_combine(hash, (size_t) layout.size_index);
+        // layout.vt: 4
+        // layout.vs: 4
+        // layout.flags: 6
+        hash_combine(hash, ((uint64_t) layout.size_index << 32) |
+                               ((uint64_t) layout.flags << 8) |
+                               ((uint64_t) layout.vs << 4) |
+                               ((uint64_t) layout.vt));
     }
 
-    hash_combine(hash, (size_t) key->flags);
-
     return hash;
 }
 

src/python/freeze.h

@@ -62,6 +62,8 @@ struct Layout {
     uint64_t literal = 0;
     /// Optional drjit type of the variable
     VarType vt = VarType::Void;
+    /// Variable index of literal. Instead of constructing a literal every time,
+    /// we keep a reference to it.
     uint32_t literal_index = 0;
 
     /// If a non drjit type is passed as function arguments or result, we simply