Make VulkanProgram a ThreadSafeResource

filament/backend/src/vulkan/VulkanAsyncHandles.cpp

@@ -16,6 +16,10 @@
 
 #include "VulkanAsyncHandles.h"
 
+#include "VulkanConstants.h"
+
+#include "vulkan/utils/Spirv.h"
+
 #include <backend/DriverEnums.h>
 
 #include <utils/debug.h>
@@ -29,6 +33,154 @@ using namespace bluevk;
 
 namespace filament::backend {
 
+namespace {
+
+inline VkShaderStageFlags getVkStage(backend::ShaderStage stage) {
+    switch(stage) {
+        case backend::ShaderStage::VERTEX:
+            return VK_SHADER_STAGE_VERTEX_BIT;
+        case backend::ShaderStage::FRAGMENT:
+            return VK_SHADER_STAGE_FRAGMENT_BIT;
+        case backend::ShaderStage::COMPUTE:
+            PANIC_POSTCONDITION("Unsupported stage");
+    }
+}
+
+} // namespace
+
+PushConstantDescription::PushConstantDescription(backend::Program const& program) {
+    mRangeCount = 0;
+    uint32_t offset = 0;
+
+    // The range is laid out so that the vertex constants are defined as the first set of bytes,
+    // followed by fragment and compute. This means we need to keep track of the offset for each
+    // stage. We do the bookeeping in mDescriptions.
+    for (auto stage: { ShaderStage::VERTEX, ShaderStage::FRAGMENT, ShaderStage::COMPUTE }) {
+        auto const& constants = program.getPushConstants(stage);
+        if (constants.empty()) {
+            continue;
+        }
+
+        auto& description = mDescriptions[(uint8_t) stage];
+        // We store the type of the constant for type-checking when writing.
+        description.types.reserve(constants.size());
+        std::for_each(constants.cbegin(), constants.cend(),
+                [&description](Program::PushConstant t) { description.types.push_back(t.type); });
+
+        uint32_t const constantsSize = (uint32_t) constants.size() * ENTRY_SIZE;
+        mRanges[mRangeCount++] = {
+            .stageFlags = getVkStage(stage),
+            .offset = offset,
+            .size = constantsSize,
+        };
+        description.offset = offset;
+        offset += constantsSize;
+    }
+}
+
+void PushConstantDescription::write(VkCommandBuffer cmdbuf, VkPipelineLayout layout,
+        backend::ShaderStage stage, uint8_t index, backend::PushConstantVariant const& value) {
+
+    uint32_t binaryValue = 0;
+    auto const& description = mDescriptions[(uint8_t) stage];
+    UTILS_UNUSED_IN_RELEASE auto const& types = description.types;
+    uint32_t const offset = description.offset;
+
+    if (std::holds_alternative<bool>(value)) {
+        assert_invariant(types[index] == ConstantType::BOOL);
+        bool const bval = std::get<bool>(value);
+        binaryValue = static_cast<uint32_t const>(bval ? VK_TRUE : VK_FALSE);
+    } else if (std::holds_alternative<float>(value)) {
+        assert_invariant(types[index] == ConstantType::FLOAT);
+        float const fval = std::get<float>(value);
+        binaryValue = *reinterpret_cast<uint32_t const*>(&fval);
+    } else {
+        assert_invariant(types[index] == ConstantType::INT);
+        int const ival = std::get<int>(value);
+        binaryValue = *reinterpret_cast<uint32_t const*>(&ival);
+    }
+
+    vkCmdPushConstants(cmdbuf, layout, getVkStage(stage), offset + index * ENTRY_SIZE, ENTRY_SIZE,
+            &binaryValue);
+}
+
+VulkanProgram::VulkanProgram(VkDevice device, Program const& builder) noexcept
+    : HwProgram(builder.getName()),
+      mInfo(new(std::nothrow) PipelineInfo(builder)),
+      mDevice(device) {
+
+    Program::ShaderSource const& blobs = builder.getShadersSource();
+    auto& modules = mInfo->shaders;
+    auto const& specializationConstants = builder.getSpecializationConstants();
+    std::vector<uint32_t> shader;
+
+    static_assert(static_cast<ShaderStage>(0) == ShaderStage::VERTEX &&
+            static_cast<ShaderStage>(1) == ShaderStage::FRAGMENT &&
+            MAX_SHADER_MODULES == 2);
+
+    for (size_t i = 0; i < MAX_SHADER_MODULES; i++) {
+        Program::ShaderBlob const& blob = blobs[i];
+
+        uint32_t* data = (uint32_t*) blob.data();
+        size_t dataSize = blob.size();
+
+        if (!specializationConstants.empty()) {
+            fvkutils::workaroundSpecConstant(blob, specializationConstants, shader);
+            data = (uint32_t*) shader.data();
+            dataSize = shader.size() * 4;
+        }
+
+        VkShaderModule& module = modules[i];
+        VkShaderModuleCreateInfo moduleInfo = {
+            .sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO,
+            .codeSize = dataSize,
+            .pCode = data,
+        };
+        VkResult result = vkCreateShaderModule(mDevice, &moduleInfo, VKALLOC, &module);
+        FILAMENT_CHECK_POSTCONDITION(result == VK_SUCCESS)
+                << "Unable to create shader module."
+                << " error=" << static_cast<int32_t>(result);
+
+#if FVK_ENABLED(FVK_DEBUG_DEBUG_UTILS)
+        utils::CString name{ builder.getName().c_str(), builder.getName().size() };
+        switch (static_cast<ShaderStage>(i)) {
+            case ShaderStage::VERTEX:
+                name += "_vs";
+                break;
+            case ShaderStage::FRAGMENT:
+                name += "_fs";
+                break;
+            default:
+                PANIC_POSTCONDITION("Unexpected stage");
+                break;
+        }
+        VulkanDriver::DebugUtils::setName(VK_OBJECT_TYPE_SHADER_MODULE,
+                reinterpret_cast<uint64_t>(module), name.c_str());
+#endif
+    }
+
+#if FVK_ENABLED(FVK_DEBUG_SHADER_MODULE)
+    FVK_LOGD << "Created VulkanProgram " << builder << ", shaders = (" << modules[0]
+             << ", " << modules[1] << ")";
+#endif
+}
+
+VulkanProgram::~VulkanProgram() {
+    for (auto shader: mInfo->shaders) {
+        vkDestroyShaderModule(mDevice, shader, VKALLOC);
+    }
+    delete mInfo;
+}
+
+void VulkanProgram::flushPushConstants(VkPipelineLayout layout) {
+    // At this point, we really ought to have a VkPipelineLayout.
+    assert_invariant(layout != VK_NULL_HANDLE);
+    for (const auto& c : mQueuedPushConstants) {
+        mInfo->pushConstantDescription.write(c.cmdbuf, layout, c.stage, c.index, c.value);
+    }
+    mQueuedPushConstants.clear();
+}
+
 std::shared_ptr<VulkanCmdFence> VulkanCmdFence::completed() noexcept {
     auto cmdFence = std::make_shared<VulkanCmdFence>(VK_NULL_HANDLE);
     cmdFence->mStatus = VK_SUCCESS;

filament/backend/src/vulkan/VulkanAsyncHandles.h

@@ -24,6 +24,9 @@
 #include "backend/Platform.h"
 
 #include "vulkan/memory/Resource.h"
+#include "vulkan/utils/StaticVector.h"
+
+#include <backend/Program.h>
 
 #include <chrono>
 #include <cstdint>
@@ -35,6 +38,123 @@
 
 namespace filament::backend {
 
+using PushConstantNameArray = utils::FixedCapacityVector<char const*>;
+using PushConstantNameByStage = std::array<PushConstantNameArray, Program::SHADER_TYPE_COUNT>;
+
+struct PushConstantDescription {
+    explicit PushConstantDescription(backend::Program const& program);
+
+    VkPushConstantRange const* getVkRanges() const noexcept { return mRanges; }
+    uint32_t getVkRangeCount() const noexcept { return mRangeCount; }
+    void write(VkCommandBuffer cmdbuf, VkPipelineLayout layout, backend::ShaderStage stage,
+            uint8_t index, backend::PushConstantVariant const& value);
+
+private:
+    static constexpr uint32_t ENTRY_SIZE = sizeof(uint32_t);
+
+    struct ConstantDescription {
+        utils::FixedCapacityVector<backend::ConstantType> types;
+        uint32_t offset = 0;
+    };
+
+    // Describes the constants in each shader stage.
+    ConstantDescription mDescriptions[Program::SHADER_TYPE_COUNT];
+    VkPushConstantRange mRanges[Program::SHADER_TYPE_COUNT];
+    uint32_t mRangeCount;
+};
+
+// Encapsulates a VkShaderModule. Note: this is a ThreadSafeResource
+// because we may use parallel compilation, in which case ref counts may
+// be modified on multiple backend threads. This is simply to avoid
+// destroying a VkShaderModule mid-compilation of a pipeline. It is
+// still assumed that all calls will originate from the backend.
+struct VulkanProgram : public HwProgram, fvkmemory::ThreadSafeResource {
+    using BindingList = fvkutils::StaticVector<uint16_t, MAX_SAMPLER_COUNT>;
+
+    VulkanProgram(VkDevice device, Program const& builder) noexcept;
+    ~VulkanProgram();
+
+    /**
+     * Cancels any parallel compilation jobs that have not yet run for this
+     * program.
+     */
+    inline void cancelParallelCompilation() {
+        mParallelCompilationCanceled.store(true, std::memory_order_release);
+    }
+
+    /**
+     * Writes out any queued push constants using the provided VkPipelineLayout.
+     *
+     * @param layout The layout that is to be used along with these push constants,
+     *               in the next draw call.
+     */
+    void flushPushConstants(VkPipelineLayout layout);
+
+    inline VkShaderModule getVertexShader() const {
+        return mInfo->shaders[0];
+    }
+
+    inline VkShaderModule getFragmentShader() const { return mInfo->shaders[1]; }
+
+    inline uint32_t getPushConstantRangeCount() const {
+        return mInfo->pushConstantDescription.getVkRangeCount();
+    }
+
+    inline VkPushConstantRange const* getPushConstantRanges() const {
+        return mInfo->pushConstantDescription.getVkRanges();
+    }
+
+    /**
+     * Returns true if parallel compilation is canceled, false if not. Parallel
+     * compilation will be canceled if this program is destroyed before relevant
+     * pipelines are created.
+     *
+     * @return true if parallel compilation should run for this program, false if not
+     */
+    inline bool isParallelCompilationCanceled() const {
+        return mParallelCompilationCanceled.load(std::memory_order_acquire);
+    }
+
+    inline void writePushConstant(VkCommandBuffer cmdbuf, VkPipelineLayout layout,
+            backend::ShaderStage stage, uint8_t index, backend::PushConstantVariant const& value) {
+        // It's possible that we don't have the layout yet. When external samplers are used, bindPipeline()
+        // in VulkanDriver returns early, without binding a layout. If that happens, the layout is not
+        // set until draw time. Any push constants that are written during that time should be saved for
+        // later, and flushed when the layout is set.
+        if (layout != VK_NULL_HANDLE) {
+            mInfo->pushConstantDescription.write(cmdbuf, layout, stage, index, value);
+        } else {
+            mQueuedPushConstants.push_back({cmdbuf, stage, index, value});
+        }
+    }
+
+    // TODO: handle compute shaders.
+    // The expected order of shaders - from frontend to backend - is vertex, fragment, compute.
+    static constexpr uint8_t const MAX_SHADER_MODULES = 2;
+
+private:
+    struct PipelineInfo {
+        explicit PipelineInfo(backend::Program const& program) noexcept
+            : pushConstantDescription(program)
+            {}
+
+        VkShaderModule shaders[MAX_SHADER_MODULES] = { VK_NULL_HANDLE };
+        PushConstantDescription pushConstantDescription;
+    };
+
+    struct PushConstantInfo {
+        VkCommandBuffer cmdbuf;
+        backend::ShaderStage stage;
+        uint8_t index;
+        backend::PushConstantVariant value;
+    };
+
+    PipelineInfo* mInfo;
+    VkDevice mDevice = VK_NULL_HANDLE;
+    std::atomic<bool> mParallelCompilationCanceled { false };
+    std::vector<PushConstantInfo> mQueuedPushConstants;
+};
+
 // Wrapper to enable use of shared_ptr for implementing shared ownership of low-level Vulkan fences.
 struct VulkanCmdFence {
     explicit VulkanCmdFence(VkFence fence) : mFence(fence) { }

filament/backend/src/vulkan/VulkanCommands.h

@@ -34,10 +34,11 @@
 #include <utils/Mutex.h>
 
 #include <atomic>
-
 #include <chrono>
 #include <list>
+#include <type_traits>
 #include <utility>
+#include <variant>
 
 namespace filament::backend {
 
@@ -72,7 +73,11 @@ struct VulkanCommandBuffer {
 
     ~VulkanCommandBuffer();
 
-    inline void acquire(fvkmemory::resource_ptr<fvkmemory::Resource> resource) {
+    template <typename T,
+              typename = std::enable_if_t<
+                std::is_same_v<T, fvkmemory::Resource> ||
+                std::is_same_v<T, fvkmemory::ThreadSafeResource>>>
+    inline void acquire(fvkmemory::resource_ptr<T> resource) {
         mResources.push_back(resource);
     }
 
@@ -115,6 +120,9 @@ struct VulkanCommandBuffer {
     }
 
 private:
+    using HeldResource = std::variant<fvkmemory::resource_ptr<Resource>,
+                                      fvkmemory::resource_ptr<ThreadSafeResource>>;
+
     static uint32_t sAgeCounter;
 
     VulkanContext const& mContext;
@@ -129,7 +137,7 @@ struct VulkanCommandBuffer {
     fvkmemory::resource_ptr<VulkanSemaphore> mSubmission;
     VkFence mFence;
     std::shared_ptr<VulkanCmdFence> mFenceStatus;
-    std::vector<fvkmemory::resource_ptr<Resource>> mResources;
+    std::vector<HeldResource> mResources;
     uint32_t mAge;
 };