ALMIOPEN-2598: Reuse KernDb connections via cached-database mechanism

projects/miopen/src/include/miopen/kern_db.hpp

@@ -86,6 +86,7 @@ class KernDb : public SQLiteBase<KernDb>
     template <typename T>
     bool RemoveRecordUnsafe(const T& problem_config)
     {
+        const std::lock_guard<std::recursive_mutex> lock{instance_mutex};
         if(filename.empty())
             return true;
         auto del_query =
@@ -108,52 +109,66 @@ class KernDb : public SQLiteBase<KernDb>
     {
         if(filename.empty())
             return {};
-        // Where clause with inserted values defeats the purpose of a prepraed statement
-        auto select_query = "SELECT kernel_blob, kernel_hash, uncompressed_size FROM " +
-                            T::table_name() + " WHERE " + problem_config.Where() + ";";
-        auto stmt = SQLite::Statement{sql, select_query};
-        // only one result field
-        // assert one row
-        auto rc = stmt.Step(sql);
-        if(rc == SQLITE_ROW)
+
+        // Lock only around the SQLite query so that bz2 decompression and MD5
+        // checking can run concurrently across threads.
+        std::vector<char> compressed_blob;
+        std::string md5_hash;
+        int64_t uncompressed_size = 0;
         {
-            auto compressed_blob                 = stmt.ColumnBlob(0);
-            auto md5_hash                        = stmt.ColumnText(1);
-            auto uncompressed_size               = stmt.ColumnInt64(2);
-            std::vector<char>& decompressed_blob = compressed_blob;
-            if(uncompressed_size != 0)
+            const std::lock_guard<std::recursive_mutex> lock{instance_mutex};
+            // Where clause with inserted values defeats the purpose of a prepared statement
+            auto select_query = "SELECT kernel_blob, kernel_hash, uncompressed_size FROM " +
+                                T::table_name() + " WHERE " + problem_config.Where() + ";";
+            auto stmt = SQLite::Statement{sql, select_query};
+            auto rc   = stmt.Step(sql);
+            if(rc == SQLITE_ROW)
             {
-                decompressed_blob = decompress_fn(compressed_blob, uncompressed_size);
+                compressed_blob   = stmt.ColumnBlob(0);
+                md5_hash          = stmt.ColumnText(1);
+                uncompressed_size = stmt.ColumnInt64(2);
+            }
+            else if(rc == SQLITE_DONE)
+            {
+                return {};
+            }
+            else
+            {
+                MIOPEN_THROW(miopenStatusInternalError, sql.ErrorMessage());
+                return {};
             }
-            auto new_md5 = md5(decompressed_blob);
-            if(new_md5 != md5_hash)
-                MIOPEN_THROW(miopenStatusInternalError, "Possible database corruption");
-            return decompressed_blob;
-        }
-        else if(rc == SQLITE_DONE)
-        {
-            return {};
         }
-        else
+
+        // Decompression and MD5 verification run outside the lock.
+        std::vector<char>& decompressed_blob = compressed_blob;
+        if(uncompressed_size != 0)
         {
-            MIOPEN_THROW(miopenStatusInternalError, sql.ErrorMessage());
+            decompressed_blob = decompress_fn(compressed_blob, uncompressed_size);
         }
-        return {};
+        auto new_md5 = md5(decompressed_blob);
+        if(new_md5 != md5_hash)
+            MIOPEN_THROW(miopenStatusInternalError, "Possible database corruption");
+        return decompressed_blob;
     }
 
     template <typename T>
     bool StoreRecordUnsafe(const T& problem_config)
     {
         if(filename.empty())
             return false;
-        auto insert_query = "INSERT OR REPLACE INTO " + T::table_name() +
-                            "(kernel_name, kernel_args, kernel_blob, kernel_hash, "
-                            "uncompressed_size) VALUES(?, ?, ?, ?, ?);";
+
+        // Compression and MD5 run outside the lock.
         auto md5_sum           = md5(problem_config.kernel_blob);
         auto uncompressed_size = problem_config.kernel_blob.size();
         bool success           = false;
         auto compressed_blob   = compress_fn(problem_config.kernel_blob, &success);
-        auto stmt              = SQLite::Statement{sql, insert_query};
+
+        // Lock only around the SQLite insert.
+        const std::lock_guard<std::recursive_mutex> lock{instance_mutex};
+        auto insert_query = "INSERT OR REPLACE INTO " + T::table_name() +
+                            "(kernel_name, kernel_args, kernel_blob, kernel_hash, "
+                            "uncompressed_size) VALUES(?, ?, ?, ?, ?);";
+        auto stmt = SQLite::Statement{sql, insert_query};
         stmt.BindPath(1, problem_config.kernel_name);
         stmt.BindText(2, problem_config.kernel_args);
         if(!success)

projects/miopen/src/include/miopen/sqlite_db.hpp

@@ -20,6 +20,8 @@
 #include "sqlite3.h"
 #include <mutex>
 
+#include <map>
+#include <memory>
 #include <string>
 #include <optional>
 #include <unordered_map>
@@ -265,7 +267,8 @@ class SQLiteBase
         }
     }
 
-    static Derived& GetCached(const fs::path& path, bool is_system);
+    static Derived& GetCached(DbKinds db_kind, const fs::path& path, bool is_system);
+    static void EvictCached(const fs::path& path, bool is_system);
     // TODO: Fix this for the overhead of having fields per record
 
     inline auto CheckTableColumns(const std::string& tableName,
@@ -350,24 +353,64 @@ class SQLiteBase
     /// behind IsUserDbDisabled() is flipped between constructions. Always false for the system
     /// databases, which the user-db switch does not govern.
     const bool disable_file_io;
+
+    /// Per-instance mutex that guards SQLite operations when the instance is shared via GetCached.
+    /// Recursive so that call chains (e.g. LoadUnsafe → FindRecordUnsafe) can re-lock safely.
+    mutable std::recursive_mutex instance_mutex;
 };
 
+/// Cached instances are stored for process lifetime (static storage).
+/// Thread-safe: a static mutex guards the map, and a per-instance mutex (see instance_mutex)
+/// guards individual operations once instances are shared across threads.
+/// Keyed by (path, is_system) so that read-only and read-write connections are not conflated.
+
+namespace detail {
 template <typename Derived>
-Derived& SQLiteBase<Derived>::GetCached(const fs::path& path, bool is_system)
+struct CacheStorage
 {
+    using Key = std::pair<fs::path, bool>;
     // NOLINTNEXTLINE (cppcoreguidelines-avoid-non-const-global-variables)
     static std::mutex mutex;
-    const std::lock_guard<std::mutex> lock{mutex};
-
     // NOLINTNEXTLINE (cppcoreguidelines-avoid-non-const-global-variables)
-    static auto instances = std::map<fs::path, Derived>{};
-    const auto it         = instances.find(path);
+    static std::map<Key, std::unique_ptr<Derived>> instances;
+};
+template <typename Derived>
+// NOLINTNEXTLINE (cppcoreguidelines-avoid-non-const-global-variables)
+std::mutex CacheStorage<Derived>::mutex;
+template <typename Derived>
+// NOLINTNEXTLINE (cppcoreguidelines-avoid-non-const-global-variables)
+std::map<typename CacheStorage<Derived>::Key, std::unique_ptr<Derived>>
+    CacheStorage<Derived>::instances;
+} // namespace detail
+
+template <typename Derived>
+Derived& SQLiteBase<Derived>::GetCached(DbKinds db_kind, const fs::path& path, bool is_system)
+{
+    using Storage = detail::CacheStorage<Derived>;
+    using Key     = typename Storage::Key;
+
+    const std::lock_guard<std::mutex> lock{Storage::mutex};
+
+    const Key key{path, is_system};
+    const auto it = Storage::instances.find(key);
+
+    if(it != Storage::instances.end())
+        return *it->second;
 
-    if(it != instances.end())
-        return it->second;
+    return *Storage::instances.emplace(key, std::make_unique<Derived>(db_kind, path, is_system))
+                .first->second;
+}
+
+/// Evict a cached instance. Closes the underlying connection so that temporary files
+/// can be deleted on platforms where open files cannot be removed (Windows).
+template <typename Derived>
+void SQLiteBase<Derived>::EvictCached(const fs::path& path, bool is_system)
+{
+    using Storage = detail::CacheStorage<Derived>;
+    using Key     = typename Storage::Key;
 
-    instances.emplace(path, Derived{path, is_system});
-    return instances.at(path);
+    const std::lock_guard<std::mutex> lock{Storage::mutex};
+    Storage::instances.erase(Key{path, is_system});
 }
 
 class SQLitePerfDb : public SQLiteBase<SQLitePerfDb>
@@ -380,6 +423,7 @@ class SQLitePerfDb : public SQLiteBase<SQLitePerfDb>
     template <class T>
     inline void InsertConfig(const T& prob_desc)
     {
+        const std::lock_guard<std::recursive_mutex> lock{instance_mutex};
         std::string clause;
         std::vector<std::string> vals;
         std::tie(clause, vals) = prob_desc.InsertQuery();
@@ -396,6 +440,7 @@ class SQLitePerfDb : public SQLiteBase<SQLitePerfDb>
     template <class T>
     inline std::string GetConfigIDs(const T& prob_desc)
     {
+        const std::lock_guard<std::recursive_mutex> lock{instance_mutex};
         std::string clause;
         std::vector<std::string> vals;
         std::tie(clause, vals) = prob_desc.WhereClause();
@@ -421,6 +466,7 @@ class SQLitePerfDb : public SQLiteBase<SQLitePerfDb>
     template <typename T>
     inline std::optional<DbRecord> FindRecordUnsafe(const T& problem_config)
     {
+        const std::lock_guard<std::recursive_mutex> lock{instance_mutex};
         if(dbInvalid)
             return {};
 
@@ -490,6 +536,7 @@ class SQLitePerfDb : public SQLiteBase<SQLitePerfDb>
     template <class T>
     inline bool RemoveUnsafe(const T& problem_config, const std::string& id)
     {
+        const std::lock_guard<std::recursive_mutex> lock{instance_mutex};
         if(dbInvalid)
             return false;
         std::string clause;
@@ -523,6 +570,7 @@ class SQLitePerfDb : public SQLiteBase<SQLitePerfDb>
     inline std::optional<DbRecord>
     UpdateUnsafe(const T& problem_config, const std::string& id, const V& values)
     {
+        const std::lock_guard<std::recursive_mutex> lock{instance_mutex};
         if(dbInvalid)
             return {};
         // UPSERT the value
@@ -576,6 +624,7 @@ class SQLitePerfDb : public SQLiteBase<SQLitePerfDb>
     template <class T, class V>
     inline bool StoreRecordUnsafe(const T& problem_config, const std::string& id, const V& values)
     {
+        const std::lock_guard<std::recursive_mutex> lock{instance_mutex};
         if(dbInvalid)
             return false;
         return bool(UpdateUnsafe(problem_config, id, values));
@@ -587,6 +636,7 @@ class SQLitePerfDb : public SQLiteBase<SQLitePerfDb>
     template <class T>
     inline bool ClearRecordUnsafe(const T& problem_config)
     {
+        const std::lock_guard<std::recursive_mutex> lock{instance_mutex};
         if(dbInvalid)
             return true;
         std::string clause;
@@ -619,6 +669,7 @@ class SQLitePerfDb : public SQLiteBase<SQLitePerfDb>
     template <class T, class V>
     inline bool LoadUnsafe(const T& problem_config, const std::string& id, V& values)
     {
+        const std::lock_guard<std::recursive_mutex> lock{instance_mutex};
         if(dbInvalid)
             return false;
         const auto record = FindRecordUnsafe(problem_config);

projects/miopen/test/gtest/cache.cpp

@@ -27,13 +27,16 @@
 
 #include <miopen/binary_cache.hpp>
 #include <miopen/bz2.hpp>
+#include <miopen/db.hpp>
 #include <miopen/kern_db.hpp>
 #include <miopen/temp_file.hpp>
 
 #include "test.hpp"
 #include "random.hpp"
 
 #include <algorithm>
+#include <thread>
+#include <type_traits>
 #include <vector>
 
 #if MIOPEN_ENABLE_SQLITE
@@ -134,6 +137,79 @@ TEST(CPU_Cache_NONE, check_kern_db)
         EXPECT_TRUE(err_db.RemoveRecordUnsafe(cfg0));
     }
 }
+
+// Compile-time check: GetDbInstance<KernDb> must resolve to the rank<1> caching overload
+// (returning a reference). If the SFINAE probe drifts again, this fires at build time.
+static_assert(std::is_lvalue_reference_v<decltype(miopen::GetDbInstance<miopen::KernDb>(
+                  miopen::DbKinds::KernelDb, {}, true))>,
+              "GetDbInstance<KernDb> should return an lvalue reference (cached overload)");
+
+TEST(CPU_Cache_NONE, check_kern_db_cached_reuse)
+{
+    miopen::TempFile temp_file("tmp-kerndb-cached");
+    auto& db1 = miopen::KernDb::GetCached(miopen::DbKinds::KernelDb, temp_file, false);
+    auto& db2 = miopen::KernDb::GetCached(miopen::DbKinds::KernelDb, temp_file, false);
+    EXPECT_EQ(&db1, &db2);
+    miopen::KernDb::EvictCached(temp_file, false);
+}
+
+TEST(CPU_Cache_NONE, check_kern_db_cached_distinct_paths)
+{
+    miopen::TempFile temp_file_a("tmp-kerndb-cached-a");
+    miopen::TempFile temp_file_b("tmp-kerndb-cached-b");
+    auto& db_a = miopen::KernDb::GetCached(miopen::DbKinds::KernelDb, temp_file_a, false);
+    auto& db_b = miopen::KernDb::GetCached(miopen::DbKinds::KernelDb, temp_file_b, false);
+    EXPECT_NE(&db_a, &db_b);
+    miopen::KernDb::EvictCached(temp_file_a, false);
+    miopen::KernDb::EvictCached(temp_file_b, false);
+}
+
+TEST(CPU_Cache_NONE, check_kern_db_cached_thread_safety)
+{
+    miopen::TempFile temp_file("tmp-kerndb-cached-mt");
+    auto& db = miopen::KernDb::GetCached(miopen::DbKinds::KernelDb, temp_file, false);
+
+    constexpr int kNumThreads = 8;
+
+    // Each thread gets its own config with a unique key and blob.
+    struct ThreadData
+    {
+        miopen::KernelConfig cfg;
+        std::vector<char> expected_blob;
+    };
+    std::vector<ThreadData> thread_data(kNumThreads);
+    for(int i = 0; i < kNumThreads; ++i)
+    {
+        thread_data[i].cfg.kernel_name = "kernel_mt_" + std::to_string(i);
+        thread_data[i].cfg.kernel_args = {random_bytes(512).data(), 512};
+        thread_data[i].cfg.kernel_blob = random_bytes(8192);
+        thread_data[i].expected_blob   = thread_data[i].cfg.kernel_blob;
+    }
+
+    // Use vector<int> instead of vector<bool> to avoid bit-packing data races.
+    std::vector<int> results(kNumThreads, 0);
+    std::vector<std::thread> threads;
+
+    for(int i = 0; i < kNumThreads; ++i)
+    {
+        threads.emplace_back([&db, &thread_data, &results, i]() {
+            // Each thread stores its own record, then reads it back.
+            if(!db.StoreRecord(thread_data[i].cfg))
+                return;
+            auto rec = db.FindRecord(thread_data[i].cfg);
+            if(rec && rec.value() == thread_data[i].expected_blob)
+                results[i] = 1;
+        });
+    }
+
+    for(auto& t : threads)
+        t.join();
+
+    for(int i = 0; i < kNumThreads; ++i)
+        EXPECT_EQ(results[i], 1) << "Thread " << i << " failed to store/find its own record";
+
+    miopen::KernDb::EvictCached(temp_file, false);
+}
 #endif
 
 TEST(CPU_Cache_NONE, check_cache_file)