diff --git a/slangpy/tests/device/test_module_cache.py b/slangpy/tests/device/test_module_cache.py index ad4edc44e..973267a0b 100644 --- a/slangpy/tests/device/test_module_cache.py +++ b/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 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"]) diff --git a/src/sgl/device/shader.cpp b/src/sgl/device/shader.cpp index e412becb2..fbc58f000 100644 --- a/src/sgl/device/shader.cpp +++ b/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 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 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 + ); + } } } diff --git a/src/sgl/device/shader.h b/src/sgl/device/shader.h index 11db396e6..eeabfbe69 100644 --- a/src/sgl/device/shader.h +++ b/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 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 _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 m_device; @@ -374,6 +386,9 @@ struct SlangModuleDesc { /// If source specified, additional path for compilation. std::optional path; + /// SHA1 digest of source (set when loading from source, for caching). + std::optional source_digest; + /// Source modules that are composed together to form this module (for composed modules only). std::vector> source_modules;