[wrapper] Parse shader stage from SPIR-V

include/inexor/vulkan-renderer/wrapper/shader.hpp

@@ -2,6 +2,7 @@
 
 #include <vulkan/vulkan_core.h>
 
+#include <cstdint>
 #include <string>
 #include <vector>
 
@@ -13,30 +14,27 @@ class Device;
 class Shader {
     const Device &m_device;
     std::string m_name;
-    VkShaderStageFlagBits m_type;
-    VkShaderModule m_shader_module{VK_NULL_HANDLE};
+    VkShaderModule m_module{VK_NULL_HANDLE};
 
-public:
-    /// @brief Construct a shader module from a block of SPIR-V memory.
-    /// @param device The const reference to a device RAII wrapper instance.
-    /// @param type The shader type.
-    /// @param name The internal debug marker name of the VkShaderModule.
-    /// @param code The memory block of the SPIR-V shader.
-    /// @param entry_point The name of the entry point, "main" by default.
-    Shader(const Device &m_device, VkShaderStageFlagBits type, const std::string &name, const std::vector<char> &code);
+    VkShaderStageFlagBits m_stage;
 
+public:
     /// @brief Construct a shader module from a SPIR-V file.
     /// This constructor loads the file content and just calls the other constructor.
     /// @param device The const reference to a device RAII wrapper instance.
-    /// @param type The shader type.
     /// @param name The internal debug marker name of the VkShaderModule.
     /// @param file_name The name of the SPIR-V shader file to load.
     /// @param entry_point The name of the entry point, "main" by default.
-    Shader(const Device &m_device, VkShaderStageFlagBits type, const std::string &name, const std::string &file_name);
+    Shader(const Device &m_device, const std::string &name, const std::string &file_name);
 
+    /// @brief Construct a shader module from a block of SPIR-V memory.
+    /// @param device The const reference to a device RAII wrapper instance.
+    /// @param name The internal debug marker name of the VkShaderModule.
+    /// @param binary The memory block of the SPIR-V shader.
+    /// @param entry_point The name of the entry point, "main" by default.
+    Shader(const Device &m_device, const std::string &name, std::vector<char> &&binary);
     Shader(const Shader &) = delete;
     Shader(Shader &&) noexcept;
-
     ~Shader();
 
     Shader &operator=(const Shader &) = delete;
@@ -46,12 +44,12 @@ class Shader {
         return m_name;
     }
 
-    [[nodiscard]] VkShaderStageFlagBits type() const {
-        return m_type;
+    [[nodiscard]] VkShaderModule module() const {
+        return m_module;
     }
 
-    [[nodiscard]] VkShaderModule module() const {
-        return m_shader_module;
+    [[nodiscard]] VkShaderStageFlagBits stage() const {
+        return m_stage;
     }
 };
 

src/vulkan-renderer/application.cpp

@@ -163,7 +163,7 @@ void Application::load_shaders() {
         spdlog::debug("Loading vertex shader file {}.", vertex_shader_file);
 
         // Insert the new shader into the list of shaders.
-        m_shaders.emplace_back(*m_device, VK_SHADER_STAGE_VERTEX_BIT, "unnamed vertex shader", vertex_shader_file);
+        m_shaders.emplace_back(*m_device, "unnamed vertex shader", vertex_shader_file);
     }
 
     spdlog::debug("Loading fragment shaders.");
@@ -177,8 +177,7 @@ void Application::load_shaders() {
         spdlog::debug("Loading fragment shader file {}.", fragment_shader_file);
 
         // Insert the new shader into the list of shaders.
-        m_shaders.emplace_back(*m_device, VK_SHADER_STAGE_FRAGMENT_BIT, "unnamed fragment shader",
-                               fragment_shader_file);
+        m_shaders.emplace_back(*m_device, "unnamed fragment shader", fragment_shader_file);
     }
 
     spdlog::debug("Loading shaders finished.");

src/vulkan-renderer/imgui.cpp

@@ -38,10 +38,8 @@ ImGUIOverlay::ImGUIOverlay(const wrapper::Device &device, const wrapper::Swapcha
     io.FontGlobalScale = m_scale;
 
     spdlog::debug("Loading ImGUI shaders");
-    m_vertex_shader = std::make_unique<wrapper::Shader>(m_device, VK_SHADER_STAGE_VERTEX_BIT, "ImGUI vertex shader",
-                                                        "shaders/ui.vert.spv");
-    m_fragment_shader = std::make_unique<wrapper::Shader>(m_device, VK_SHADER_STAGE_FRAGMENT_BIT,
-                                                          "ImGUI fragment shader", "shaders/ui.frag.spv");
+    m_vertex_shader = std::make_unique<wrapper::Shader>(m_device, "ImGUI vertex shader", "shaders/ui.vert.spv");
+    m_fragment_shader = std::make_unique<wrapper::Shader>(m_device, "ImGUI fragment shader", "shaders/ui.frag.spv");
 
     // Load font texture
 

src/vulkan-renderer/render_graph.cpp

@@ -39,7 +39,7 @@ void GraphicsStage::bind_buffer(const BufferResource *buffer, const std::uint32_
 void GraphicsStage::uses_shader(const wrapper::Shader &shader) {
     auto create_info = wrapper::make_info<VkPipelineShaderStageCreateInfo>();
     create_info.module = shader.module();
-    create_info.stage = shader.type();
+    create_info.stage = shader.stage();
     create_info.pName = "main";
     m_shaders.push_back(create_info);
 }

src/vulkan-renderer/wrapper/shader.cpp

@@ -5,6 +5,7 @@
 #include "inexor/vulkan-renderer/wrapper/make_info.hpp"
 
 #include <spdlog/spdlog.h>
+#include <spirv/unified1/spirv.h>
 
 #include <cassert>
 #include <fstream>
@@ -30,47 +31,71 @@ std::vector<char> read_binary(const std::string &file_name) {
     return buffer;
 }
 
+VkShaderStageFlagBits shader_stage(SpvExecutionModel execution_model) {
+    switch (execution_model) {
+    case SpvExecutionModelVertex:
+        return VK_SHADER_STAGE_VERTEX_BIT;
+    case SpvExecutionModelFragment:
+        return VK_SHADER_STAGE_FRAGMENT_BIT;
+    case SpvExecutionModelGLCompute:
+        return VK_SHADER_STAGE_COMPUTE_BIT;
+    default:
+        assert(false);
+    }
+}
+
 } // namespace
 
 namespace inexor::vulkan_renderer::wrapper {
 
-Shader::Shader(const Device &device, const VkShaderStageFlagBits type, const std::string &name,
-               const std::string &file_name)
-    : Shader(device, type, name, read_binary(file_name)) {}
+Shader::Shader(const Device &device, const std::string &name, const std::string &file_name)
+    : Shader(device, name, read_binary(file_name)) {}
 
-Shader::Shader(const Device &device, const VkShaderStageFlagBits type, const std::string &name,
-               const std::vector<char> &code)
-    : m_device(device), m_type(type), m_name(name) {
-    assert(device.device());
+Shader::Shader(const Device &device, const std::string &name, std::vector<char> &&binary)
+    : m_device(device), m_name(name) {
     assert(!name.empty());
-    assert(!code.empty());
-
-    auto shader_module_ci = make_info<VkShaderModuleCreateInfo>();
-    shader_module_ci.codeSize = code.size();
+    assert(!binary.empty());
 
     // When you perform a cast like this, you also need to ensure that the data satisfies the alignment
     // requirements of std::uint32_t. Lucky for us, the data is stored in an std::vector where the default
     // allocator already ensures that the data satisfies the worst case alignment requirements.
-    shader_module_ci.pCode = reinterpret_cast<const std::uint32_t *>(code.data()); // NOLINT
+    const auto *code = reinterpret_cast<const std::uint32_t *>(binary.data()); // NOLINT
+    auto shader_module_ci = make_info<VkShaderModuleCreateInfo>();
+    shader_module_ci.codeSize = binary.size();
+    shader_module_ci.pCode = code;
 
-    spdlog::debug("Creating shader module {}.", name);
-    if (const auto result = vkCreateShaderModule(device.device(), &shader_module_ci, nullptr, &m_shader_module);
+    spdlog::debug("Creating shader module {}", name);
+    if (const auto result = vkCreateShaderModule(device.device(), &shader_module_ci, nullptr, &m_module);
         result != VK_SUCCESS) {
         throw VulkanException("Error: vkCreateShaderModule failed for shader " + name + "!", result);
     }
 
     // Assign an internal name using Vulkan debug markers.
-    m_device.set_debug_marker_name(m_shader_module, VK_DEBUG_REPORT_OBJECT_TYPE_SHADER_MODULE_EXT, name);
+    m_device.set_debug_marker_name(m_module, VK_DEBUG_REPORT_OBJECT_TYPE_SHADER_MODULE_EXT, name);
+
+    // Parse SPIR-V to extract the shader stage.
+    assert(code[0] == SpvMagicNumber);
+    const auto *inst = code + 5;
+    while (inst != code + (binary.size() / 4)) {
+        // Each instruction starts with a dword with the upper 16 bits holding the total number of words in the
+        // instruction and the lower 16 bits holding the opcode.
+        std::uint16_t opcode = (inst[0] >> 0u) & 0xffffu;
+        std::uint16_t word_count = (inst[0] >> 16u) & 0xffffu;
+        if (opcode == SpvOpEntryPoint) {
+            assert(word_count >= 2);
+            m_stage = shader_stage(static_cast<SpvExecutionModel>(inst[1]));
+        }
+        inst += word_count;
+    }
 }
 
-Shader::Shader(Shader &&other) noexcept : m_device(other.m_device) {
-    m_type = other.m_type;
+Shader::Shader(Shader &&other) noexcept : m_device(other.m_device), m_stage(other.m_stage) {
     m_name = std::move(other.m_name);
-    m_shader_module = std::exchange(other.m_shader_module, nullptr);
+    m_module = std::exchange(other.m_module, nullptr);
 }
 
 Shader::~Shader() {
-    vkDestroyShaderModule(m_device.device(), m_shader_module, nullptr);
+    vkDestroyShaderModule(m_device.device(), m_module, nullptr);
 }
 
 } // namespace inexor::vulkan_renderer::wrapper