fix shader source dump for cuda and metal

LuisaGroup · Apr 9, 2024 · 3309143 · 3309143
1 parent b1d7fd4
commit 3309143
Show file tree

Hide file tree

Showing 7 changed files with 40 additions and 4 deletions.
diff --git a/include/luisa/core/binary_io.h b/include/luisa/core/binary_io.h
@@ -74,17 +74,19 @@ class BinaryStream {
     virtual ~BinaryStream() noexcept = default;
 };
 
-class BinaryIO {
+class LC_CORE_API BinaryIO {
 
 public:
     virtual ~BinaryIO() noexcept = default;
     virtual void clear_shader_cache() const noexcept = 0;
     [[nodiscard]] virtual luisa::unique_ptr<BinaryStream> read_shader_bytecode(luisa::string_view name) const noexcept = 0;
     [[nodiscard]] virtual luisa::unique_ptr<BinaryStream> read_shader_cache(luisa::string_view name) const noexcept = 0;
+    [[nodiscard]] virtual luisa::unique_ptr<BinaryStream> read_shader_source(luisa::string_view name) const noexcept;
     [[nodiscard]] virtual luisa::unique_ptr<BinaryStream> read_internal_shader(luisa::string_view name) const noexcept = 0;
     // returns the path of the written file (if stored on disk, otherwise returns empty path)
     [[nodiscard]] virtual luisa::filesystem::path write_shader_bytecode(luisa::string_view name, luisa::span<std::byte const> data) const noexcept = 0;
     [[nodiscard]] virtual luisa::filesystem::path write_shader_cache(luisa::string_view name, luisa::span<std::byte const> data) const noexcept = 0;
+    [[nodiscard]] virtual luisa::filesystem::path write_shader_source(luisa::string_view name, luisa::span<std::byte const> data) const noexcept;
     [[nodiscard]] virtual luisa::filesystem::path write_internal_shader(luisa::string_view name, luisa::span<std::byte const> data) const noexcept = 0;
 };
 

diff --git a/src/backends/common/default_binary_io.cpp b/src/backends/common/default_binary_io.cpp
@@ -1,3 +1,4 @@
+#include <new>
 #include <luisa/core/stl/filesystem.h>
 #include <luisa/core/logging.h>
 #include <luisa/core/binary_file_stream.h>
@@ -147,6 +148,12 @@ luisa::unique_ptr<BinaryStream> DefaultBinaryIO::read_internal_shader(luisa::str
     return luisa::make_unique<LMDBBinaryStream>(r.data(), r.size());
 }
 
+luisa::unique_ptr<BinaryStream> DefaultBinaryIO::read_shader_source(luisa::string_view name) const noexcept {
+    std::filesystem::path local_path{name};
+    if (local_path.is_absolute()) { return _read(luisa::to_string(name)); }
+    return _read(luisa::to_string(_cache_dir / name));
+}
+
 luisa::filesystem::path DefaultBinaryIO::write_shader_bytecode(luisa::string_view name, luisa::span<std::byte const> data) const noexcept {
     std::filesystem::path local_path{name};
     if (local_path.is_absolute()) {
@@ -158,6 +165,17 @@ luisa::filesystem::path DefaultBinaryIO::write_shader_bytecode(luisa::string_vie
     return file_path;
 }
 
+luisa::filesystem::path DefaultBinaryIO::write_shader_source(luisa::string_view name, luisa::span<std::byte const> data) const noexcept {
+    std::filesystem::path local_path{name};
+    if (local_path.is_absolute()) {
+        _write(luisa::to_string(name), data);
+        return local_path;
+    }
+    auto file_path = _cache_dir / name;
+    _write(luisa::to_string(file_path), data);
+    return file_path;
+}
+
 luisa::filesystem::path DefaultBinaryIO::write_shader_cache(luisa::string_view name, luisa::span<std::byte const> data) const noexcept {
     _cache_lmdb.write(name, data);
     return _cache_dir / name;
@@ -167,6 +185,7 @@ luisa::filesystem::path DefaultBinaryIO::write_internal_shader(luisa::string_vie
     _data_lmdb.write(name, data);
     return _data_dir / name;
 }
+
 void DefaultBinaryIO::clear_shader_cache() const noexcept {
     vstd::destruct(std::addressof(_cache_lmdb));
     std::error_code ec;

diff --git a/src/backends/common/default_binary_io.h b/src/backends/common/default_binary_io.h
@@ -44,6 +44,8 @@ class DefaultBinaryIO final : public BinaryIO {
     luisa::filesystem::path write_shader_bytecode(luisa::string_view name, luisa::span<std::byte const> data) const noexcept override;
     luisa::filesystem::path write_shader_cache(luisa::string_view name, luisa::span<std::byte const> data) const noexcept override;
     luisa::filesystem::path write_internal_shader(luisa::string_view name, luisa::span<std::byte const> data) const noexcept override;
+    luisa::unique_ptr<BinaryStream> read_shader_source(luisa::string_view name) const noexcept override;
+    luisa::filesystem::path write_shader_source(luisa::string_view name, luisa::span<std::byte const> data) const noexcept override;
     void clear_shader_cache() const noexcept override;
 };
 

diff --git a/src/backends/cuda/cuda_device.cpp b/src/backends/cuda/cuda_device.cpp
@@ -546,7 +546,7 @@ ShaderCreationInfo CUDADevice::_create_shader(luisa::string name,
             if (uses_user_path) {
                 src_dump_path = _io->write_shader_bytecode(src_name, src_data);
             } else if (option.enable_cache) {
-                src_dump_path = _io->write_shader_cache(src_name, src_data);
+                src_dump_path = _io->write_shader_source(src_name, src_data);
             }
         }
         luisa::string src_filename{src_dump_path.string()};

diff --git a/src/backends/metal/metal_compiler.cpp b/src/backends/metal/metal_compiler.cpp
@@ -333,15 +333,14 @@ MetalShaderHandle MetalCompiler::compile(luisa::string_view src,
 
         auto is_aot = !option.name.empty();
         auto uses_cache = is_aot || option.enable_cache;
-
         if (option.enable_debug_info || detail::get_bool_env("LUISA_DUMP_SOURCE")) {
             auto src_dump_name = luisa::format("{}.metal", name);
             luisa::span src_dump{reinterpret_cast<const std::byte *>(src.data()), src.size()};
             luisa::filesystem::path src_dump_path;
             if (is_aot) {
                 src_dump_path = _device->io()->write_shader_bytecode(src_dump_name, src_dump);
             } else if (option.enable_cache) {
-                src_dump_path = _device->io()->write_shader_cache(src_dump_name, src_dump);
+                src_dump_path = _device->io()->write_shader_source(src_dump_name, src_dump);
             }
             // TODO: attach shader source to Metal shader archive for debugging.
             //       Is it possible without using the command line?

diff --git a/src/core/CMakeLists.txt b/src/core/CMakeLists.txt
@@ -7,6 +7,7 @@ set(LUISA_COMPUTE_CORE_SOURCES
         basic_types.cpp
         binary_buffer.cpp
         binary_file_stream.cpp
+        binary_io.cpp
         dynamic_module.cpp
         first_fit.cpp
         logging.cpp

diff --git a/src/core/binary_io.cpp b/src/core/binary_io.cpp
@@ -0,0 +1,13 @@
+#include <luisa/core/binary_io.h>
+
+namespace luisa {
+
+luisa::unique_ptr<BinaryStream> BinaryIO::read_shader_source(luisa::string_view name) const noexcept {
+    return this->read_shader_cache(name);
+}
+
+luisa::filesystem::path BinaryIO::write_shader_source(luisa::string_view name, luisa::span<const std::byte> data) const noexcept {
+    return this->write_shader_cache(name, data);
+}
+
+}