Add caching for source modules

slangpy/tests/device/test_module_cache.py

@@ -25,6 +25,12 @@ def test_module_cache(device_type: spy.DeviceType, tmpdir: str):
     kernel = device.create_compute_kernel(program)
     kernel.dispatch(thread_count=[1, 1, 1])
     assert device.flush_print_to_string().strip() == "Hello module cache!"
+
+    # Check that cached binary module files were created.
+    cache_root = Path(cache_dir)
+    slang_module_files = list(cache_root.rglob("*.slang-module"))
+    assert len(slang_module_files) > 0, "Expected cached .slang-module files"
+
     # Close device.
     device.close()
 
@@ -47,5 +53,70 @@ def test_module_cache(device_type: spy.DeviceType, tmpdir: str):
     device.close()
 
 
+@pytest.mark.parametrize("device_type", helpers.DEFAULT_DEVICE_TYPES)
+def test_source_module_cache(device_type: spy.DeviceType, tmpdir: str):
+    cache_dir = tmpdir
+    source = """
+RWStructuredBuffer<float> output;
+
+[shader("compute")]
+[numthreads(1, 1, 1)]
+void compute_main() {
+    output[0] = 42.0;
+}
+"""
+    # Create device with a module cache.
+    device = spy.Device(
+        type=device_type,
+        module_cache_path=cache_dir,
+        label=f"source-module-cache-1-{device_type.name}",
+    )
+    # Load module from source, link program, and dispatch.
+    module = device.load_module_from_source("test_source_cache", source)
+    ep = module.entry_point("compute_main")
+    program = device.link_program([module], [ep])
+    kernel = device.create_compute_kernel(program)
+    output = device.create_buffer(
+        element_count=1,
+        struct_size=4,
+        usage=spy.BufferUsage.shader_resource | spy.BufferUsage.unordered_access,
+    )
+    kernel.dispatch(thread_count=[1, 1, 1], vars={"output": output})
+    import numpy as np
+
+    result = np.frombuffer(output.to_numpy(), dtype=np.float32)
+    assert result[0] == 42.0
+
+    # Check that cache files were created under source_modules directory.
+    cache_root = Path(cache_dir)
+    slang_module_files = list(cache_root.rglob("source_modules/*.slang-module"))
+    assert len(slang_module_files) > 0, "Expected cached .slang-module files for source module"
+
+    # Close device.
+    device.close()
+
+    # Re-create device using same module cache location.
+    device = spy.Device(
+        type=device_type,
+        module_cache_path=cache_dir,
+        label=f"source-module-cache-2-{device_type.name}",
+    )
+    # Load same source again - should load from cache.
+    module = device.load_module_from_source("test_source_cache", source)
+    ep = module.entry_point("compute_main")
+    program = device.link_program([module], [ep])
+    kernel = device.create_compute_kernel(program)
+    output = device.create_buffer(
+        element_count=1,
+        struct_size=4,
+        usage=spy.BufferUsage.shader_resource | spy.BufferUsage.unordered_access,
+    )
+    kernel.dispatch(thread_count=[1, 1, 1], vars={"output": output})
+    result = np.frombuffer(output.to_numpy(), dtype=np.float32)
+    assert result[0] == 42.0
+    # Close device.
+    device.close()
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v", "-s"])

src/sgl/device/shader.cpp

@@ -476,6 +476,11 @@ void SlangSession::create_session(SlangSessionBuild& build)
             session_options.insert_cache_include(data->cache_include_paths[i], i);
         }
 
+        // Create cache directory for source-loaded modules and register as include path.
+        data->source_cache_path = data->cache_path / "source_modules";
+        std::filesystem::create_directories(data->source_cache_path);
+        session_options.add_include(data->source_cache_path);
+
         // Update session descriptor with the patched include paths.
         slang_session_option_entries = session_options.slang_entries();
         session_desc.compilerOptionEntries = slang_session_option_entries.data();
@@ -537,6 +542,7 @@ ref<SlangModule> SlangSession::load_module_from_source(
     desc.module_name = module_name;
     desc.source = source;
     desc.path = path;
+    desc.source_digest = digest;
     return create_module(std::move(desc));
 }
 
@@ -777,6 +783,87 @@ bool SlangSession::write_module_to_cache(slang::IModule* module)
     return true;
 }
 
+std::filesystem::path SlangSession::_get_source_module_cache_path(const SHA1::Digest& digest) const
+{
+    return m_data->source_cache_path / (string::hexlify(digest.data(), digest.size()) + ".slang");
+}
+
+bool SlangSession::_write_source_module_to_cache(
+    slang::IModule* module,
+    std::string_view module_name,
+    std::string_view source,
+    const SHA1::Digest& digest
+) const
+{
+    std::filesystem::path source_file = _get_source_module_cache_path(digest);
+    std::filesystem::path binary_file = source_file;
+    binary_file.replace_extension(".slang-module");
+
+    // Create directories to cache path.
+    std::error_code ec;
+    std::filesystem::create_directories(source_file.parent_path(), ec);
+    if (ec) {
+        log_warn(
+            "Failed to create directory \"{}\" for source module cache ({})",
+            source_file.parent_path(),
+            ec.message()
+        );
+        return false;
+    }
+
+    // Write source to a temporary file, then rename.
+    std::random_device rd;
+    {
+        std::filesystem::path tmp_path = source_file;
+        uint64_t uid = rd();
+        tmp_path.replace_extension(".slang-" + string::hexlify(&uid, sizeof(uid)));
+        if (std::filesystem::exists(tmp_path))
+            return false;
+        {
+            FileStream stream(tmp_path, FileStream::Mode::write);
+            stream.write(source.data(), source.size());
+        }
+        std::filesystem::rename(tmp_path, source_file, ec);
+        if (ec) {
+            log_warn("Failed to rename cached source file \"{}\" to \"{}\" ({})", tmp_path, source_file, ec.message());
+            std::filesystem::remove(tmp_path, ec);
+            return false;
+        }
+    }
+
+    // Write binary module to a temporary file, then rename.
+    {
+        std::filesystem::path tmp_path = binary_file;
+        uint64_t uid = rd();
+        tmp_path.replace_extension(".slang-module-" + string::hexlify(&uid, sizeof(uid)));
+        if (std::filesystem::exists(tmp_path)) {
+            std::filesystem::remove(source_file, ec);
+            return false;
+        }
+        if (!SLANG_SUCCEEDED(module->writeToFile(tmp_path.string().c_str()))) {
+            log_warn("Failed to write cached source module \"{}\" to \"{}\"", module_name, binary_file);
+            std::filesystem::remove(source_file, ec);
+            return false;
+        }
+        std::filesystem::rename(tmp_path, binary_file, ec);
+        if (ec) {
+            log_warn(
+                "Failed to rename cached source module \"{}\" to \"{}\" ({})",
+                tmp_path,
+                binary_file,
+                ec.message()
+            );
+            std::filesystem::remove(tmp_path, ec);
+            std::filesystem::remove(source_file, ec);
+            return false;
+        }
+    }
+
+    log_debug("Cached source module \"{}\" to \"{}\"", module_name, source_file);
+
+    return true;
+}
+
 std::string SlangSessionData::resolve_module_name(std::string_view module_name) const
 {
     // Return if module name is an absolute file path.
@@ -916,7 +1003,7 @@ void SlangModule::load(SlangSessionBuild& build_data) const
         // Regular module: load from file or source
         Timer timer;
         Slang::ComPtr<ISlangBlob> diagnostics;
-        slang::IModule* slang_module;
+        slang::IModule* slang_module = nullptr;
 
         // Load module either from resolved name or source depending on whether source specified
         if (!desc.source.has_value()) {
@@ -932,25 +1019,56 @@ void SlangModule::load(SlangSessionBuild& build_data) const
                 throw SlangCompileError(msg);
             }
         } else {
-            // TODO: This is a workaround until we use a Slang release with this fix:
-            // https://github.com/shader-slang/slang/pull/10996
-            // Once this is fixed on the Slang side, we can remove this.
-            std::string source_str = fmt::format("// {}\n{}", desc.module_name, desc.source.value());
+            // Try to load from source module cache if available.
+            bool loaded_from_cache = false;
+            if (session_data->cache_enabled && desc.source_digest.has_value()) {
+                std::filesystem::path cache_source_file = m_session->_get_source_module_cache_path(*desc.source_digest);
+                std::filesystem::path cache_binary_file = cache_source_file;
+                cache_binary_file.replace_extension(".slang-module");
+                if (std::filesystem::exists(cache_source_file) && std::filesystem::exists(cache_binary_file)) {
+                    std::string cache_filename = cache_source_file.filename().string();
+                    SGL_CATCH_INTERNAL_SLANG_ERROR(
+                        slang_module
+                        = session_data->slang_session->loadModule(cache_filename.c_str(), diagnostics.writeRef());
+                    );
+                    if (slang_module) {
+                        loaded_from_cache = true;
+                        log_debug("Loaded source module \"{}\" from cache", desc.module_name);
+                    }
+                }
+            }
 
-            SGL_CATCH_INTERNAL_SLANG_ERROR(
-                slang_module = session_data->slang_session->loadModuleFromSourceString(
-                    std::string{desc.module_name}.c_str(),
-                    desc.path ? desc.path->string().c_str() : nullptr,
-                    source_str.c_str(),
-                    diagnostics.writeRef()
-                )
-            );
-            if (!slang_module) {
-                std::string msg = append_diagnostics(
-                    fmt::format("Failed to load slang module \"{}\" from source", desc.module_name),
-                    diagnostics
+            if (!loaded_from_cache) {
+                // TODO: This is a workaround until we use a Slang release with this fix:
+                // https://github.com/shader-slang/slang/pull/10996
+                // Once this is fixed on the Slang side, we can remove this.
+                std::string source_str = fmt::format("// {}\n{}", desc.module_name, desc.source.value());
+
+                SGL_CATCH_INTERNAL_SLANG_ERROR(
+                    slang_module = session_data->slang_session->loadModuleFromSourceString(
+                        std::string{desc.module_name}.c_str(),
+                        desc.path ? desc.path->string().c_str() : nullptr,
+                        source_str.c_str(),
+                        diagnostics.writeRef()
+                    )
                 );
-                throw SlangCompileError(msg);
+                if (!slang_module) {
+                    std::string msg = append_diagnostics(
+                        fmt::format("Failed to load slang module \"{}\" from source", desc.module_name),
+                        diagnostics
+                    );
+                    throw SlangCompileError(msg);
+                }
+
+                // Write to source module cache.
+                if (session_data->cache_enabled && desc.source_digest.has_value()) {
+                    m_session->_write_source_module_to_cache(
+                        slang_module,
+                        desc.module_name,
+                        source_str,
+                        *desc.source_digest
+                    );
+                }
             }
         }
 

src/sgl/device/shader.h

@@ -266,6 +266,9 @@ struct SlangSessionData : Object {
     /// One cache path for each include path under the root cache path.
     std::vector<std::filesystem::path> cache_include_paths;
 
+    /// Cache path for source-loaded modules.
+    std::filesystem::path source_cache_path;
+
     /// Finds fully qualified module name by scanning the cache and include paths.
     std::string resolve_module_name(std::string_view module_name) const;
 };
@@ -334,6 +337,15 @@ class SGL_API SlangSession : public Object {
     // Internal access to the built session data.
     ref<SlangSessionData> _data() { return m_data; }
 
+    // Internal source module cache helpers.
+    std::filesystem::path _get_source_module_cache_path(const SHA1::Digest& digest) const;
+    bool _write_source_module_to_cache(
+        slang::IModule* module,
+        std::string_view module_name,
+        std::string_view source,
+        const SHA1::Digest& digest
+    ) const;
+
 private:
     ref<Device> m_device;
 
@@ -374,6 +386,9 @@ struct SlangModuleDesc {
     /// If source specified, additional path for compilation.
     std::optional<std::filesystem::path> path;
 
+    /// SHA1 digest of source (set when loading from source, for caching).
+    std::optional<SHA1::Digest> source_digest;
+
     /// Source modules that are composed together to form this module (for composed modules only).
     std::vector<ref<SlangModule>> source_modules;