Add push constants

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

@@ -40,6 +40,8 @@ class Program {
     static constexpr size_t UNIFORM_BINDING_COUNT = CONFIG_UNIFORM_BINDING_COUNT;
     static constexpr size_t SAMPLER_BINDING_COUNT = CONFIG_SAMPLER_BINDING_COUNT;
 
+    static constexpr char const* PUSH_CONSTANT_STRUCT_VAR_NAME = "pushConstants";
+
     struct Sampler {
         utils::CString name = {};   // name of the sampler in the shader
         uint32_t binding = 0;       // binding point of the sampler in the shader
@@ -117,6 +119,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 +158,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 +184,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

@@ -1241,6 +1241,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

@@ -285,6 +285,32 @@ void OpenGLDriver::bindSampler(GLuint unit, GLuint sampler) noexcept {
     mContext.bindSampler(unit, sampler);
 }
 
+void OpenGLDriver::setPushConstant(backend::ShaderStage, uint8_t index,
+        backend::PushConstantVariant value) {
+    assert_invariant(mCurrentPushConstants &&
+                     "Calling setPushConstant() when program does not define push constants");
+
+    // Note that stage is not applicable to the GL backend.
+
+    auto const& constants = *mCurrentPushConstants;
+    ASSERT_PRECONDITION(index < constants.size(), "Push constant index=%d is out-of-bounds", index);
+    auto const& constant = constants[index];
+
+    if (std::holds_alternative<bool>(value)) {
+        assert_invariant(constant.type == ConstantType::BOOL);
+        bool const bval = std::get<bool>(value);
+        glUniform1i(constant.location, bval ? 1 : 0);
+    } else if (std::holds_alternative<float>(value)) {
+        assert_invariant(constant.type == ConstantType::FLOAT);
+        float const fval = std::get<float>(value);
+        glUniform1f(constant.location, fval);
+    } else {
+        assert_invariant(constant.type == ConstantType::INT);
+        int const ival = std::get<int>(value);
+        glUniform1i(constant.location, ival);
+    }
+}
+
 void OpenGLDriver::bindTexture(GLuint unit, GLTexture const* t) noexcept {
     assert_invariant(t != nullptr);
     mContext.bindTexture(unit, t->gl.target, t->gl.id);
@@ -3802,6 +3828,7 @@ void OpenGLDriver::bindPipeline(PipelineState state) {
     gl.polygonOffset(state.polygonOffset.slope, state.polygonOffset.constant);
     OpenGLProgram* const p = handle_cast<OpenGLProgram*>(state.program);
     mValidProgram = useProgram(p);
+    mCurrentPushConstants = p->getPushConstants();
 }
 
 void OpenGLDriver::bindRenderPrimitive(Handle<HwRenderPrimitive> rph) {

filament/backend/src/opengl/OpenGLDriver.h

@@ -70,6 +70,8 @@ struct TargetBufferInfo;
 class OpenGLProgram;
 class TimerQueryFactoryInterface;
 
+struct PushConstantBundle;
+
 class OpenGLDriver final : public DriverBase {
     inline explicit OpenGLDriver(OpenGLPlatform* platform,
             const Platform::DriverConfig& driverConfig) noexcept;
@@ -375,6 +377,9 @@ class OpenGLDriver final : public DriverBase {
     // for ES2 sRGB support
     GLSwapChain* mCurrentDrawSwapChain = nullptr;
     bool mRec709OutputColorspace = false;
+
+    utils::FixedCapacityVector<PushConstantBundle> const* mCurrentPushConstants =
+            nullptr;
 };
 
 // ------------------------------------------------------------------------------------------------

filament/backend/src/opengl/OpenGLProgram.cpp

@@ -30,6 +30,7 @@
 #include <utils/Systrace.h>
 
 #include <array>
+#include <sstream>
 #include <string_view>
 #include <utility>
 #include <new>
@@ -46,6 +47,7 @@ struct OpenGLProgram::LazyInitializationData {
     Program::UniformBlockInfo uniformBlockInfo;
     Program::SamplerGroupInfo samplerGroupInfo;
     std::array<Program::UniformInfo, Program::UNIFORM_BINDING_COUNT> bindingUniformInfo;
+    utils::FixedCapacityVector<Program::PushConstant> pushConstants;
 };
 
 
@@ -62,6 +64,11 @@ OpenGLProgram::OpenGLProgram(OpenGLDriver& gld, Program&& program) noexcept
         lazyInitializationData->uniformBlockInfo = std::move(program.getUniformBlockBindings());
     }
 
+    // We only keep the push constants for vertex stage because we'd like to keep this class a
+    // certain size.
+    lazyInitializationData->pushConstants =
+            std::move(program.getPushConstants(ShaderStage::VERTEX));
+
     ShaderCompilerService& compiler = gld.getShaderCompilerService();
     mToken = compiler.createProgram(name, std::move(program));
 
@@ -203,6 +210,29 @@ void OpenGLProgram::initializeProgramState(OpenGLContext& context, GLuint progra
         }
     }
     mUsedBindingsCount = usedBindingCount;
+
+    // Push constant initialization
+    auto& constants = lazyInitializationData.pushConstants;
+    if (!constants.empty()) {
+        mVertexPushConstants.reserve(constants.size());
+        mVertexPushConstants.resize(constants.size());
+        // The constants are defined in a struct of name PUSH_CONSTANT_STRUCT_VAR_NAME. We prepend
+        // the prefix here to avoid string manipulation in a the draw loop.
+        std::stringbuf constantPrefix;
+        constantPrefix.sputn(Program::PUSH_CONSTANT_STRUCT_VAR_NAME,
+                strlen(Program::PUSH_CONSTANT_STRUCT_VAR_NAME));
+        constantPrefix.sputn(".", 1);
+
+        std::transform(constants.begin(), constants.end(), mVertexPushConstants.begin(),
+                [&constantPrefix, program](Program::PushConstant& constant) -> PushConstantBundle {
+                    constant.name.insert(0, utils::CString(constantPrefix.str().c_str()));
+                    auto const loc = glGetUniformLocation(program, constant.name.c_str());
+                    return {
+                        .location = loc,
+                        .type = constant.type,
+                    };
+                });
+    }
 }
 
 void OpenGLProgram::updateSamplers(OpenGLDriver* const gld) const noexcept {

filament/backend/src/opengl/OpenGLProgram.h

@@ -38,6 +38,13 @@ namespace filament::backend {
 
 class OpenGLDriver;
 
+// Holds information needed to write to push constants. This struct cannot nest within OpenGLProgram
+// due to the need for forward declaration.
+struct PushConstantBundle {
+    GLint location = -1;
+    ConstantType type;
+};
+
 class OpenGLProgram : public HwProgram {
 public:
 
@@ -78,6 +85,13 @@ class OpenGLProgram : public HwProgram {
         GLuint program = 0;
     } gl;                                               // 4 bytes
 
+    utils::FixedCapacityVector<PushConstantBundle> const* getPushConstants() {
+        if (mVertexPushConstants.empty()) {
+            return nullptr;
+        }
+        return &mVertexPushConstants;
+    }
+
 private:
     // keep these away from of other class attributes
     struct LazyInitializationData;
@@ -97,7 +111,6 @@ class OpenGLProgram : public HwProgram {
     uint8_t mUsedBindingsCount = 0u;                    // 1 byte
     UTILS_UNUSED uint8_t padding[3] = {};               // 3 bytes
 
-
     // only needed for ES2
     GLint mRec709Location = -1; // 4 bytes
     using LocationInfo = utils::FixedCapacityVector<GLint>;
@@ -108,10 +121,12 @@ class OpenGLProgram : public HwProgram {
         mutable uint16_t age = std::numeric_limits<uint16_t>::max();
     };
     UniformsRecord const* mUniformsRecords = nullptr;
+
+    utils::FixedCapacityVector<PushConstantBundle> mVertexPushConstants;
 };
 
 // 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 54 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;