Add push constants (#7817)

 - Push constants is a small set of bytes that can be recorded
   directly on the command buffer.
 - Implemented it for the vulkan/gl backend.

filament/backend/include/backend/DriverEnums.h

@@ -31,6 +31,7 @@
 
 #include <array>        // FIXME: STL headers are not allowed in public headers
 #include <type_traits>  // FIXME: STL headers are not allowed in public headers
+#include <variant>      // FIXME: STL headers are not allowed in public headers
 
 #include <stddef.h>
 #include <stdint.h>
@@ -91,12 +92,15 @@ static constexpr uint64_t SWAP_CHAIN_HAS_STENCIL_BUFFER         = SWAP_CHAIN_CON
  */
 static constexpr uint64_t SWAP_CHAIN_CONFIG_PROTECTED_CONTENT   = 0x40;
 
-
 static constexpr size_t MAX_VERTEX_ATTRIBUTE_COUNT  = 16;   // This is guaranteed by OpenGL ES.
 static constexpr size_t MAX_SAMPLER_COUNT           = 62;   // Maximum needed at feature level 3.
 static constexpr size_t MAX_VERTEX_BUFFER_COUNT     = 16;   // Max number of bound buffer objects.
 static constexpr size_t MAX_SSBO_COUNT              = 4;    // This is guaranteed by OpenGL ES.
 
+static constexpr size_t MAX_PUSH_CONSTANT_COUNT     = 32;   // Vulkan 1.1 spec allows for 128-byte
+                                                            // of push constant (we assume 4-byte
+                                                            // types).
+
 // Per feature level caps
 // Use (int)FeatureLevel to index this array
 static constexpr struct {
@@ -332,7 +336,7 @@ enum class UniformType : uint8_t {
 /**
  * Supported constant parameter types
  */
- enum class ConstantType : uint8_t {
+enum class ConstantType : uint8_t {
   INT,
   FLOAT,
   BOOL
@@ -1219,6 +1223,8 @@ struct StencilState {
     uint8_t padding = 0;
 };
 
+using PushConstantVariant = std::variant<int32_t, float, bool>;
+
 static_assert(sizeof(StencilState::StencilOperations) == 5u,
         "StencilOperations size not what was intended");
 

filament/backend/include/backend/Program.h

@@ -117,6 +117,14 @@ class Program {
     Program& specializationConstants(
             utils::FixedCapacityVector<SpecializationConstant> specConstants) noexcept;
 
+    struct PushConstant {
+        utils::CString name;
+        ConstantType type;
+    };
+
+    Program& pushConstants(ShaderStage stage,
+            utils::FixedCapacityVector<PushConstant> constants) noexcept;
+
     Program& cacheId(uint64_t cacheId) noexcept;
 
     Program& multiview(bool multiview) noexcept;
@@ -148,6 +156,15 @@ class Program {
         return mSpecializationConstants;
     }
 
+    utils::FixedCapacityVector<PushConstant> const& getPushConstants(
+            ShaderStage stage) const noexcept {
+        return mPushConstants[static_cast<uint8_t>(stage)];
+    }
+
+    utils::FixedCapacityVector<PushConstant>& getPushConstants(ShaderStage stage) noexcept {
+        return mPushConstants[static_cast<uint8_t>(stage)];
+    }
+
     uint64_t getCacheId() const noexcept { return mCacheId; }
 
     bool isMultiview() const noexcept { return mMultiview; }
@@ -165,6 +182,7 @@ class Program {
     uint64_t mCacheId{};
     utils::Invocable<utils::io::ostream&(utils::io::ostream& out)> mLogger;
     utils::FixedCapacityVector<SpecializationConstant> mSpecializationConstants;
+    std::array<utils::FixedCapacityVector<PushConstant>, SHADER_TYPE_COUNT> mPushConstants;
     utils::FixedCapacityVector<std::pair<utils::CString, uint8_t>> mAttributes;
     std::array<UniformInfo, Program::UNIFORM_BINDING_COUNT> mBindingUniformInfo;
     CompilerPriorityQueue mPriorityQueue = CompilerPriorityQueue::HIGH;

filament/backend/include/private/backend/DriverAPI.inc

@@ -434,6 +434,11 @@ DECL_DRIVER_API_N(bindSamplers,
         uint32_t, index,
         backend::SamplerGroupHandle, sbh)
 
+DECL_DRIVER_API_N(setPushConstant,
+        backend::ShaderStage, stage,
+        uint8_t, index,
+        backend::PushConstantVariant, value)
+
 DECL_DRIVER_API_N(insertEventMarker,
         const char*, string,
         uint32_t, len = 0)

filament/backend/src/Program.cpp

@@ -91,6 +91,12 @@ Program& Program::specializationConstants(
     return *this;
 }
 
+Program& Program::pushConstants(ShaderStage stage,
+        utils::FixedCapacityVector<PushConstant> constants) noexcept {
+    mPushConstants[static_cast<uint8_t>(stage)] = std::move(constants);
+    return *this;
+}
+
 Program& Program::cacheId(uint64_t cacheId) noexcept {
     mCacheId = cacheId;
     return *this;

filament/backend/src/metal/MetalDriver.mm

@@ -1252,6 +1252,9 @@
     mContext->samplerBindings[index] = sb;
 }
 
+void MetalDriver::setPushConstant(backend::ShaderStage stage, uint8_t index,
+        backend::PushConstantVariant value) {}
+
 void MetalDriver::insertEventMarker(const char* string, uint32_t len) {
 
 }

filament/backend/src/noop/NoopDriver.cpp

@@ -312,6 +312,10 @@ void NoopDriver::unbindBuffer(BufferObjectBinding bindingType, uint32_t index) {
 void NoopDriver::bindSamplers(uint32_t index, Handle<HwSamplerGroup> sbh) {
 }
 
+void NoopDriver::setPushConstant(backend::ShaderStage stage, uint8_t index,
+        backend::PushConstantVariant value) {
+}
+
 void NoopDriver::insertEventMarker(char const* string, uint32_t len) {
 }
 

filament/backend/src/opengl/OpenGLDriver.cpp

@@ -269,6 +269,10 @@ void OpenGLDriver::terminate() {
     assert_invariant(mGpuCommandCompleteOps.empty());
 #endif
 
+    if (mCurrentPushConstants) {
+        delete mCurrentPushConstants;
+    }
+
     mContext.terminate();
 
     mPlatform.terminate();
@@ -289,6 +293,42 @@ void OpenGLDriver::bindSampler(GLuint unit, GLuint sampler) noexcept {
     mContext.bindSampler(unit, sampler);
 }
 
+void OpenGLDriver::setPushConstant(backend::ShaderStage stage, uint8_t index,
+        backend::PushConstantVariant value) {
+    assert_invariant(mCurrentPushConstants &&
+                     "Calling setPushConstant() before binding a pipeline");
+
+    assert_invariant(stage == ShaderStage::VERTEX || stage == ShaderStage::FRAGMENT);
+    utils::Slice<std::pair<GLint, ConstantType>> constants;
+    if (stage == ShaderStage::VERTEX) {
+        constants = mCurrentPushConstants->vertexConstants;
+    } else if (stage == ShaderStage::FRAGMENT) {
+        constants = mCurrentPushConstants->fragmentConstants;
+    }
+
+    assert_invariant(index < constants.size());
+    auto const& [location, type] = constants[index];
+
+    // This push constant wasn't found in the shader. It's ok to return without error-ing here.
+    if (location < 0) {
+        return;
+    }
+
+    if (std::holds_alternative<bool>(value)) {
+        assert_invariant(type == ConstantType::BOOL);
+        bool const bval = std::get<bool>(value);
+        glUniform1i(location, bval ? 1 : 0);
+    } else if (std::holds_alternative<float>(value)) {
+        assert_invariant(type == ConstantType::FLOAT);
+        float const fval = std::get<float>(value);
+        glUniform1f(location, fval);
+    } else {
+        assert_invariant(type == ConstantType::INT);
+        int const ival = std::get<int>(value);
+        glUniform1i(location, ival);
+    }
+}
+
 void OpenGLDriver::bindTexture(GLuint unit, GLTexture const* t) noexcept {
     assert_invariant(t != nullptr);
     mContext.bindTexture(unit, t->gl.target, t->gl.id);
@@ -3808,6 +3848,12 @@ void OpenGLDriver::bindPipeline(PipelineState state) {
     gl.polygonOffset(state.polygonOffset.slope, state.polygonOffset.constant);
     OpenGLProgram* const p = handle_cast<OpenGLProgram*>(state.program);
     mValidProgram = useProgram(p);
+
+    if (!mCurrentPushConstants) {
+        mCurrentPushConstants = new (std::nothrow) PushConstantBundle{p->getPushConstants()};
+    } else {
+        (*mCurrentPushConstants) = p->getPushConstants();
+    }
 }
 
 void OpenGLDriver::bindRenderPrimitive(Handle<HwRenderPrimitive> rph) {

filament/backend/src/opengl/OpenGLDriver.h

@@ -66,9 +66,9 @@ namespace filament::backend {
 class OpenGLPlatform;
 class PixelBufferDescriptor;
 struct TargetBufferInfo;
-
 class OpenGLProgram;
 class TimerQueryFactoryInterface;
+struct PushConstantBundle;
 
 class OpenGLDriver final : public DriverBase {
     inline explicit OpenGLDriver(OpenGLPlatform* platform,
@@ -375,6 +375,8 @@ class OpenGLDriver final : public DriverBase {
     // for ES2 sRGB support
     GLSwapChain* mCurrentDrawSwapChain = nullptr;
     bool mRec709OutputColorspace = false;
+
+    PushConstantBundle* mCurrentPushConstants = nullptr;
 };
 
 // ------------------------------------------------------------------------------------------------

filament/backend/src/opengl/OpenGLProgram.cpp

@@ -46,21 +46,24 @@ struct OpenGLProgram::LazyInitializationData {
     Program::UniformBlockInfo uniformBlockInfo;
     Program::SamplerGroupInfo samplerGroupInfo;
     std::array<Program::UniformInfo, Program::UNIFORM_BINDING_COUNT> bindingUniformInfo;
+    utils::FixedCapacityVector<Program::PushConstant> vertexPushConstants;
+    utils::FixedCapacityVector<Program::PushConstant> fragmentPushConstants;
 };
 
 
 OpenGLProgram::OpenGLProgram() noexcept = default;
 
 OpenGLProgram::OpenGLProgram(OpenGLDriver& gld, Program&& program) noexcept
         : HwProgram(std::move(program.getName())) {
-
     auto* const lazyInitializationData = new(std::nothrow) LazyInitializationData();
     lazyInitializationData->samplerGroupInfo = std::move(program.getSamplerGroupInfo());
     if (UTILS_UNLIKELY(gld.getContext().isES2())) {
         lazyInitializationData->bindingUniformInfo = std::move(program.getBindingUniformInfo());
     } else {
         lazyInitializationData->uniformBlockInfo = std::move(program.getUniformBlockBindings());
     }
+    lazyInitializationData->vertexPushConstants = std::move(program.getPushConstants(ShaderStage::VERTEX));
+    lazyInitializationData->fragmentPushConstants = std::move(program.getPushConstants(ShaderStage::FRAGMENT));
 
     ShaderCompilerService& compiler = gld.getShaderCompilerService();
     mToken = compiler.createProgram(name, std::move(program));
@@ -203,6 +206,21 @@ void OpenGLProgram::initializeProgramState(OpenGLContext& context, GLuint progra
         }
     }
     mUsedBindingsCount = usedBindingCount;
+
+    auto& vertexConstants = lazyInitializationData.vertexPushConstants;
+    auto& fragmentConstants = lazyInitializationData.fragmentPushConstants;
+
+    size_t const totalConstantCount = vertexConstants.size() + fragmentConstants.size();
+    if (totalConstantCount > 0) {
+        mPushConstants.reserve(totalConstantCount);
+        mPushConstantFragmentStageOffset = vertexConstants.size();
+        auto const transformAndAdd = [&](Program::PushConstant const& constant) {
+            GLint const loc = glGetUniformLocation(program, constant.name.c_str());
+            mPushConstants.push_back({loc, constant.type});
+        };
+        std::for_each(vertexConstants.cbegin(), vertexConstants.cend(), transformAndAdd);
+        std::for_each(fragmentConstants.cbegin(), fragmentConstants.cend(), transformAndAdd);
+    }
 }
 
 void OpenGLProgram::updateSamplers(OpenGLDriver* const gld) const noexcept {

filament/backend/src/opengl/OpenGLProgram.h

@@ -27,6 +27,7 @@
 
 #include <utils/compiler.h>
 #include <utils/FixedCapacityVector.h>
+#include <utils/Slice.h>
 
 #include <array>
 #include <limits>
@@ -38,6 +39,11 @@ namespace filament::backend {
 
 class OpenGLDriver;
 
+struct PushConstantBundle {
+    utils::Slice<std::pair<GLint, ConstantType>> vertexConstants;
+    utils::Slice<std::pair<GLint, ConstantType>> fragmentConstants;
+};
+
 class OpenGLProgram : public HwProgram {
 public:
 
@@ -78,6 +84,14 @@ class OpenGLProgram : public HwProgram {
         GLuint program = 0;
     } gl;                                               // 4 bytes
 
+    PushConstantBundle getPushConstants() {
+        auto fragBegin = mPushConstants.begin() + mPushConstantFragmentStageOffset;
+        return {
+            .vertexConstants = utils::Slice(mPushConstants.begin(), fragBegin),
+            .fragmentConstants = utils::Slice(fragBegin, mPushConstants.end()),
+        };
+    }
+
 private:
     // keep these away from of other class attributes
     struct LazyInitializationData;
@@ -95,23 +109,30 @@ class OpenGLProgram : public HwProgram {
     ShaderCompilerService::program_token_t mToken{};    // 16 bytes
 
     uint8_t mUsedBindingsCount = 0u;                    // 1 byte
-    UTILS_UNUSED uint8_t padding[3] = {};               // 3 bytes
+    UTILS_UNUSED uint8_t padding[2] = {};               // 2 byte
 
+    // Push constant array offset for fragment stage constants.
+    uint8_t mPushConstantFragmentStageOffset = 0u;      // 1 byte
 
     // only needed for ES2
-    GLint mRec709Location = -1; // 4 bytes
+    GLint mRec709Location = -1;                         // 4 bytes
+
     using LocationInfo = utils::FixedCapacityVector<GLint>;
     struct UniformsRecord {
         Program::UniformInfo uniforms;
         LocationInfo locations;
         mutable GLuint id = 0;
         mutable uint16_t age = std::numeric_limits<uint16_t>::max();
     };
-    UniformsRecord const* mUniformsRecords = nullptr;
+    UniformsRecord const* mUniformsRecords = nullptr;               // 8 bytes
+
+    // Note that this can be replaced with a raw pointer and an uint8_t (for size) to reduce the
+    // size of the container to 9 bytes if there is a need in the future.
+    utils::FixedCapacityVector<std::pair<GLint, ConstantType>> mPushConstants;// 16 bytes
 };
 
 // if OpenGLProgram is larger tha 64 bytes, it'll fall in a larger Handle bucket.
-static_assert(sizeof(OpenGLProgram) <= 64); // currently 48 bytes
+static_assert(sizeof(OpenGLProgram) <= 64); // currently 64 bytes
 
 } // namespace filament::backend
 

filament/backend/src/vulkan/VulkanDriver.cpp

@@ -259,8 +259,8 @@ VulkanDriver::VulkanDriver(VulkanPlatform* platform, VulkanContext const& contex
     mDescriptorSetManager.setPlaceHolders(mSamplerCache.getSampler({}), mEmptyTexture,
             mEmptyBufferObject);
 
-    mGetPipelineFunction = [this](VulkanDescriptorSetLayoutList const& layouts) {
-        return mPipelineLayoutCache.getLayout(layouts);
+    mGetPipelineFunction = [this](VulkanDescriptorSetLayoutList const& layouts, VulkanProgram* program) {
+        return mPipelineLayoutCache.getLayout(layouts, program);
     };
 }
 
@@ -1572,6 +1572,14 @@ void VulkanDriver::bindSamplers(uint32_t index, Handle<HwSamplerGroup> sbh) {
     mSamplerBindings[index] = hwsb;
 }
 
+void VulkanDriver::setPushConstant(backend::ShaderStage stage, uint8_t index,
+        backend::PushConstantVariant value) {
+    assert_invariant(mBoundPipeline.program && "Expect a program when writing to push constants");
+    VulkanCommands* commands = &mCommands;
+    mBoundPipeline.program->writePushConstant(commands, mBoundPipeline.pipelineLayout, stage, index,
+            value);
+}
+
 void VulkanDriver::insertEventMarker(char const* string, uint32_t len) {
 #if FVK_ENABLED(FVK_DEBUG_GROUP_MARKERS)
     mCommands.insertEventMarker(string, len);
@@ -1857,7 +1865,13 @@ void VulkanDriver::bindPipeline(PipelineState pipelineState) {
         mDescriptorSetManager.updateSampler({}, binding, texture, vksampler);
     }
 
-    mPipelineCache.bindLayout(mDescriptorSetManager.bind(commands, program, mGetPipelineFunction));
+    auto const pipelineLayout = mDescriptorSetManager.bind(commands, program, mGetPipelineFunction);
+    mBoundPipeline = {
+        .program = program,
+        .pipelineLayout = pipelineLayout,
+    };
+
+    mPipelineCache.bindLayout(pipelineLayout);
     mPipelineCache.bindPipeline(commands);
     FVK_SYSTRACE_END();
 }

filament/backend/src/vulkan/VulkanDriver.h

@@ -160,6 +160,13 @@ class VulkanDriver final : public DriverBase {
 
     VulkanDescriptorSetManager::GetPipelineLayoutFunction mGetPipelineFunction;
 
+    // This is necessary for us to write to push constants after binding a pipeline.
+    struct BoundPipeline {
+        VulkanProgram* program;
+        VkPipelineLayout pipelineLayout;
+    };
+    BoundPipeline mBoundPipeline = {};
+
     RenderPassFboBundle mRenderPassFboInfo;
 
     bool const mIsSRGBSwapChainSupported;