Shader cache index implementation (#2452)

7 files changed, 332 insertions, 40 deletions

diff --git a/tools/gfx/persistent-shader-cache.cpp b/tools/gfx/persistent-shader-cache.cpp
new file mode 100644
index 000000000..a378e0f4f
--- /dev/null
+++ b/tools/gfx/persistent-shader-cache.cpp
@@ -0,0 +1,182 @@
+// slang-shader-cache-index.cpp
+#include "persistent-shader-cache.h"
+
+#include "../../source/core/slang-digest-util.h"
+#include "../../source/core/slang-io.h"
+#include "../../source/core/slang-string-util.h"
+#include "../../source/core/slang-file-system.h"
+
+namespace Slang
+{
+
+PersistentShaderCache::PersistentShaderCache(const Desc& inDesc)
+{
+    desc = inDesc;
+
+    if (!desc.shaderCachePath.getBuffer())
+    {
+        if (!desc.shaderCacheFileSystem)
+        {
+            // Only a path was provided, so we get a mutable file system
+            // using OSFileSystem::getMutableSingleton.
+            desc.shaderCacheFileSystem = OSFileSystem::getMutableSingleton();
+        }
+        desc.shaderCacheFileSystem = new RelativeFileSystem(desc.shaderCacheFileSystem, desc.shaderCachePath);
+    }
+
+    // If our shader cache has an underlying file system, check if it's mutable. If so, store a pointer
+    // to the mutable version in order to save new entries later.
+    if (desc.shaderCacheFileSystem)
+    {
+        desc.shaderCacheFileSystem->queryInterface(ISlangMutableFileSystem::getTypeGuid(), (void**)mutableShaderCacheFileSystem.writeRef());
+    }
+
+    loadCacheFromFile();
+}
+
+// Load a previous cache index saved to disk. If not found, create a new cache index
+// and save it to disk as filename.
+void PersistentShaderCache::loadCacheFromFile()
+{
+    ComPtr<ISlangBlob> indexBlob;
+    if (SLANG_FAILED(desc.shaderCacheFileSystem->loadFile(desc.cacheFilename.getBuffer(), indexBlob.writeRef())))
+    {
+        // Cache index not found, so we'll create and save a new one.
+        if (mutableShaderCacheFileSystem)
+        {
+            mutableShaderCacheFileSystem->saveFile(desc.cacheFilename.getBuffer(), nullptr, 0);
+            return;
+        }
+        // Cache index not found and we can't save a new one due to the file system being immutable.
+        return;
+    }
+
+    String indexString;
+    File::readAllText(desc.cacheFilename, indexString);
+
+    List<UnownedStringSlice> lines;
+    StringUtil::calcLines(indexString.getUnownedSlice(), lines);
+    for (auto line : lines)
+    {
+        List<UnownedStringSlice> digests;
+        StringUtil::split(line, ' ', digests); // This will return our two hashes as two elements in digests.
+        auto dependencyDigest = DigestUtil::fromString(digests[0]);
+        auto astDigest = DigestUtil::fromString(digests[1]);
+
+        ShaderCacheEntry entry = { dependencyDigest, astDigest };
+        auto entryNode = entries.AddLast(entry);
+        keyToEntry.Add(dependencyDigest, entryNode);
+    }
+}
+
+LinkedNode<ShaderCacheEntry>* PersistentShaderCache::findEntry(const slang::Digest& key, ISlangBlob** outCompiledCode)
+{
+    LinkedNode<ShaderCacheEntry>* entryNode;
+    if (!keyToEntry.TryGetValue(key, entryNode))
+    {
+        // The key was not found in the cache, so we return nullptr.
+        *outCompiledCode = nullptr;
+        return nullptr;
+    }
+
+    // If the key is found, load the stored contents from disk. We then move the corresponding
+    // entry to the front of the linked list and update the cache file on disk
+    desc.shaderCacheFileSystem->loadFile(DigestUtil::toString(key).getBuffer(), outCompiledCode);
+    if (entries.FirstNode() != entryNode)
+    {
+        entries.RemoveFromList(entryNode);
+        entries.AddFirst(entryNode);
+        if (mutableShaderCacheFileSystem)
+        {
+            saveCacheToFile();
+        }
+    }
+    return entryNode;
+}
+
+void PersistentShaderCache::addEntry(const slang::Digest& dependencyDigest, const slang::Digest& astDigest, ISlangBlob* compiledCode)
+{
+    if (!mutableShaderCacheFileSystem)
+    {
+        // Should not save new entries if the underlying file system isn't mutable.
+        return;
+    }
+    
+    // Check that we do not exceed the cache's size limit by adding another entry. If so,
+    // remove the least recently used entry first.
+    //
+    // In theory, this could loop infinitely since deleteLRUEntry() immediately returns if
+    // mutableShaderCacheFileSystem is not set. However, this situation is functionally impossible
+    // because we check immediately before this as well.
+    while (desc.entryCountLimit > 0 && entries.Count() >= desc.entryCountLimit)
+    {
+        deleteLRUEntry();
+    }
+
+    ShaderCacheEntry entry = { dependencyDigest, astDigest };
+    auto entryNode = entries.AddFirst(entry);
+    keyToEntry.Add(dependencyDigest, entryNode);
+
+    mutableShaderCacheFileSystem->saveFileBlob(DigestUtil::toString(dependencyDigest).getBuffer(), compiledCode);
+
+    saveCacheToFile();
+}
+
+void PersistentShaderCache::updateEntry(
+    LinkedNode<ShaderCacheEntry>* entryNode,
+    const slang::Digest& dependencyDigest,
+    const slang::Digest& astDigest,
+    ISlangBlob* updatedCode)
+{
+    if (!mutableShaderCacheFileSystem)
+    {
+        // Updating entries requires saving to disk in order to overwrite the old shader file
+        // on disk, so we return if the underlying file system isn't mutable.
+        return;
+    }
+
+    entryNode->Value.astBasedDigest = astDigest;
+    mutableShaderCacheFileSystem->saveFileBlob(DigestUtil::toString(dependencyDigest).getBuffer(), updatedCode);
+
+    saveCacheToFile();
+}
+
+void PersistentShaderCache::saveCacheToFile()
+{
+    if (!mutableShaderCacheFileSystem)
+    {
+        // Cannot save the index to disk if the underlying file system isn't mutable.
+        return;
+    }
+
+    StringBuilder indexSb;
+    for (auto& entry : entries)
+    {
+        indexSb << entry.dependencyBasedDigest;
+        indexSb << " ";
+        indexSb << entry.astBasedDigest;
+        indexSb << "\n";
+    }
+
+    mutableShaderCacheFileSystem->saveFile(desc.cacheFilename.getBuffer(), indexSb.getBuffer(), indexSb.getLength());
+}
+
+void PersistentShaderCache::deleteLRUEntry()
+{
+    if (!mutableShaderCacheFileSystem)
+    {
+        // This is here as a safety precaution but should never be hit as
+        // addEntry() is the only function that should call this.
+        return;
+    }
+
+    auto lruEntry = entries.LastNode();
+    auto shaderKey = lruEntry->Value.dependencyBasedDigest;
+
+    keyToEntry.Remove(shaderKey);
+    mutableShaderCacheFileSystem->remove(DigestUtil::toString(shaderKey).getBuffer());
+
+    entries.Delete(lruEntry);
+}
+
+}
diff --git a/tools/gfx/persistent-shader-cache.h b/tools/gfx/persistent-shader-cache.h
new file mode 100644
index 000000000..55738a9af
--- /dev/null
+++ b/tools/gfx/persistent-shader-cache.h
@@ -0,0 +1,79 @@
+// slang-shader-cache-index.h
+#pragma once
+#include "../../slang.h"
+#include "../../slang-com-ptr.h"
+
+#include "../../source/core/slang-string.h"
+#include "../../source/core/slang-dictionary.h"
+#include "../../source/core/slang-linked-list.h"
+
+namespace Slang
+{
+
+struct ShaderCacheEntry
+{
+    slang::Digest dependencyBasedDigest;
+    slang::Digest astBasedDigest;
+};
+
+class PersistentShaderCache
+{
+public:
+    // TODO: Remove in integration PR in favor of new ShaderCacheDesc in slang-gfx.h
+    struct Desc
+    {
+        String cacheFilename;
+        String shaderCachePath;
+        SlangInt entryCountLimit = 1000;
+        ISlangFileSystem* shaderCacheFileSystem = nullptr;
+    };
+
+    PersistentShaderCache(const Desc& inDesc);
+
+    // Fetch the cache entry corresponding to the provided key. If found, move the entry to
+    // the front of entries and return the entry and the corresponding compiled code in
+    // outCompiledCode. Else, return nullptr.
+    LinkedNode<ShaderCacheEntry>* findEntry(const slang::Digest& key, ISlangBlob** outCompiledCode);
+
+    // Add an entry to the cache with the provided key and contents hashes. If
+    // adding an entry causes the cache to exceed size limitations, this will also
+    // delete the least recently used entry.
+    void addEntry(const slang::Digest& dependencyDigest, const slang::Digest& astDigest, ISlangBlob* compiledCode);
+
+    // Update the contents hash for the specified entry in the cache and update the
+    // corresponding file on disk.
+    void updateEntry(
+        LinkedNode<ShaderCacheEntry>* entryNode,
+        const slang::Digest& dependencyDigest,
+        const slang::Digest& astDigest,
+        ISlangBlob* updatedCode);
+
+private:
+    // Load a previous cache index saved to disk. If not found, create a new cache index
+    // and save it to disk as filename.
+    void loadCacheFromFile();
+
+    // Update the cache index on disk. This should be called any time an entry changes.
+    void saveCacheToFile();
+
+    // Delete the last entry (the least recently used) from entries, remove its key/value pair
+    // from keyToEntry, and remove the corresponding file on disk. This should only be called
+    // by addEntry() when the cache reaches maximum capacity.
+    void deleteLRUEntry();
+
+    // The shader cache's description.
+    Desc desc;
+
+    // Dictionary mapping each shader's key to its corresponding node (entry) in the list
+    // of entries.
+    Dictionary<slang::Digest, LinkedNode<ShaderCacheEntry>*> keyToEntry;
+
+    // Linked list containing the entries stored in the shader cache in order
+    // of most to least recently used.
+    LinkedList<ShaderCacheEntry> entries;
+
+    // The underlying file system used for the shader cache.
+    ComPtr<ISlangMutableFileSystem> mutableShaderCacheFileSystem = nullptr;
+};
+
+}
diff --git a/tools/gfx/renderer-shared.cpp b/tools/gfx/renderer-shared.cpp
index f76fa2325..48cd6907f 100644
--- a/tools/gfx/renderer-shared.cpp
+++ b/tools/gfx/renderer-shared.cpp
@@ -6,7 +6,7 @@
 #include "../../source/core/slang-file-system.h"
 
 #include "../../slang.h"
-#include "../../source/slang/slang-hash-utils.h"
+#include "../../source/core/slang-digest-util.h"
 
 using namespace Slang;
 
@@ -368,7 +368,7 @@ Result RendererBase::getEntryPointCodeFromShaderCache(
     slang::Digest shaderKeyHash;
     program->computeDependencyBasedHash(entryPointIndex, targetIndex, &shaderKeyHash);
 
-    StringBuilder shaderKey = hashToString(shaderKeyHash);
+    String shaderKey = DigestUtil::toString(shaderKeyHash);
 
     // Produce a hash using the AST for this program - This is needed to check whether a cache entry is effectively dirty,
     // or to save along with the compiled code into an entry so the entry can be checked if fetched later on.
@@ -463,6 +463,9 @@ IDevice* gfx::RendererBase::getInterface(const Guid& guid)
 
 SLANG_NO_THROW Result SLANG_MCALL RendererBase::initialize(const Desc& desc)
 {
+    // TODO: This logic for initalizing the shader cache has been replicated inside
+    // the constructor for ShaderCacheIndex. Remove when ShaderCacheIndex is integrated in.
+
     // If a shader cache file system was provided, use the provided system.
     if (desc.shaderCacheFileSystem)
     {
diff --git a/tools/gfx/renderer-shared.h b/tools/gfx/renderer-shared.h
index b140e1f40..fc1c35d20 100644
--- a/tools/gfx/renderer-shared.h
+++ b/tools/gfx/renderer-shared.h
@@ -1372,6 +1372,8 @@ public:
     SlangContext slangContext;
     ShaderCache shaderCache;
 
+    // TODO: These should be removed when ShaderCacheIndex is ready to be integrated. ShaderCacheIndex
+    // will be responsible for keeping track of the underlying filesystem for the cache.
     ISlangFileSystem* shaderCacheFileSystem = nullptr;
     ComPtr<ISlangMutableFileSystem> mutableShaderCacheFileSystem = nullptr;
 
diff --git a/tools/slang-unit-test/unit-test-checksum.cpp b/tools/slang-unit-test/unit-test-checksum.cpp
deleted file mode 100644
index 90f426f94..000000000
--- a/tools/slang-unit-test/unit-test-checksum.cpp
+++ /dev/null
@@ -1,32 +0,0 @@
-// unit-test-checksum.cpp
-
-#include "tools/unit-test/slang-unit-test.h"
-
-#include "../../source/slang/slang-hash-utils.h"
-
-using namespace Slang;
-
-SLANG_UNIT_TEST(checksum)
-{
-    {
-        slang::Digest testA;
-        testA.values[0] = 1;
-        testA.values[1] = 2;
-        testA.values[2] = 3;
-        testA.values[3] = 4;
-
-        String testAString = hashToString(testA);
-        SLANG_CHECK(testAString.equals(String("01000000020000000300000004000000")));
-    }
-
-    {
-        slang::Digest testC;
-        testC.values[0] = 0x11111111;
-        testC.values[1] = 0x22222222;
-        testC.values[2] = 0x33333333;
-        testC.values[3] = 0x44444444;
-
-        String testCString = hashToString(testC);
-        SLANG_CHECK(testCString.equals(String("11111111222222223333333344444444")));
-    }
-}
diff --git a/tools/slang-unit-test/unit-test-digest-utils.cpp b/tools/slang-unit-test/unit-test-digest-utils.cpp
new file mode 100644
index 000000000..a463fe07f
--- /dev/null
+++ b/tools/slang-unit-test/unit-test-digest-utils.cpp
@@ -0,0 +1,58 @@
+// unit-test-digest-utils.cpp
+
+#include "tools/unit-test/slang-unit-test.h"
+
+#include "../../source/core/slang-digest-util.h"
+
+using namespace Slang;
+
+SLANG_UNIT_TEST(digestUtils)
+{
+    {
+        slang::Digest testA;
+        testA.values[0] = 1;
+        testA.values[1] = 2;
+        testA.values[2] = 3;
+        testA.values[3] = 4;
+
+        String testAString = DigestUtil::toString(testA);
+        SLANG_CHECK(testAString.equals(String("01000000020000000300000004000000")));
+    }
+
+    {
+        slang::Digest testC;
+        testC.values[0] = 0x11111111;
+        testC.values[1] = 0x22222222;
+        testC.values[2] = 0x33333333;
+        testC.values[3] = 0x44444444;
+
+        String testCString = DigestUtil::toString(testC);
+        SLANG_CHECK(testCString.equals(String("11111111222222223333333344444444")));
+    }
+
+    {
+        auto digestString = UnownedStringSlice("5D6CC58E1824A4DFD0CF57395B603316");
+        slang::Digest digest = DigestUtil::fromString(digestString);
+        auto resultString = DigestUtil::toString(digest);
+        SLANG_CHECK(resultString == digestString);
+    }
+
+    {
+        auto digestString = UnownedStringSlice("01000000020000000300000004000000");
+        slang::Digest digest = DigestUtil::fromString(digestString);
+        auto resultString = DigestUtil::toString(digest);
+        SLANG_CHECK(resultString == digestString);
+    }
+
+    {
+        slang::Digest testD;
+        testD.values[0] = 1;
+        testD.values[1] = 2;
+        testD.values[2] = 3;
+        testD.values[3] = 4;
+
+        StringBuilder testDSb;
+        testDSb << testD;
+        SLANG_CHECK(testDSb.equals(String("01000000020000000300000004000000")));
+    }
+}
diff --git a/tools/slang-unit-test/unit-test-md5.cpp b/tools/slang-unit-test/unit-test-md5.cpp
index 0e7fef57c..41f3d6cf6 100644
--- a/tools/slang-unit-test/unit-test-md5.cpp
+++ b/tools/slang-unit-test/unit-test-md5.cpp
@@ -3,7 +3,7 @@
 
 #include "../../source/core/slang-md5.h"
 #include "../../source/core/slang-string.h"
-#include "../../source/slang/slang-hash-utils.h"
+#include "../../source/core/slang-digest-util.h"
 
 using namespace Slang;
 
@@ -23,7 +23,7 @@ SLANG_UNIT_TEST(md5hash)
         slang::Digest testA;
         testHashGen.finalize(&testCtx, &testA);
 
-        String testAString = hashToString(testA);
+        String testAString = DigestUtil::toString(testA);
         SLANG_CHECK(testAString.equals(String("5BA171E20898BDD205639013746F2679")));
     }
 
@@ -43,7 +43,7 @@ SLANG_UNIT_TEST(md5hash)
         slang::Digest testB;
         testHashGen.finalize(&testCtx, &testB);
 
-        String testBString = hashToString(testB);
+        String testBString = DigestUtil::toString(testB);
         SLANG_CHECK(testBString.equals(String("4352D88A78AA39750BF70CD6F27BCAA5")));
     }
 
@@ -59,7 +59,7 @@ SLANG_UNIT_TEST(md5hash)
         slang::Digest testC;
         testHashGen.finalize(&testCtx, &testC);
 
-        String testCString = hashToString(testC);
+        String testCString = DigestUtil::toString(testC);
         SLANG_CHECK(testCString.equals(String("5D6CC58E1824A4DFD0CF57395B603316")));
     }
 
@@ -75,7 +75,7 @@ SLANG_UNIT_TEST(md5hash)
         slang::Digest testD;
         testHashGen.finalize(&testCtx, &testD);
 
-        String testDString = hashToString(testD);
+        String testDString = DigestUtil::toString(testD);
         SLANG_CHECK(testDString.equals(String("DF5A79CC2170C7401CF0A506CEB0CE24")));
     }
 
@@ -91,7 +91,7 @@ SLANG_UNIT_TEST(md5hash)
         slang::Digest testE;
         testHashGen.finalize(&testCtx, &testE);
 
-        String testEString = hashToString(testE);
+        String testEString = DigestUtil::toString(testE);
         SLANG_CHECK(testEString.equals(String("4AE71336E44BF9BF79D2752E234818A5")));
     }
 }