Add push constants support in the vulkan backend

src/anki/gr/CommandBuffer.h

@@ -261,6 +261,9 @@ public:
 	/// Bind texture buffer.
 	void bindTextureBuffer(U32 set, U32 binding, BufferPtr buff, PtrSize offset, PtrSize range, Format fmt);
 
+	/// Set push constants.
+	void setPushConstants(const void* data, U32 dataSize);
+
 	/// Bind a program.
 	void bindShaderProgram(ShaderProgramPtr prog);
 

src/anki/gr/ShaderCompiler.cpp

@@ -41,6 +41,7 @@ static const char* SHADER_HEADER = R"(#version 450 core
 #	define ANKI_TEX_BINDING(set_, binding_) binding = set_ * %u + binding_
 #	define ANKI_IMAGE_BINDING(set_, binding_) binding = set_ * %u + binding_
 #	define ANKI_SPEC_CONST(binding_, type_, name_) const type_ name_ = _anki_spec_const_ ## binding_
+#	define ANKI_PUSH_CONSTANTS(struct_, name_) layout(location = 0) uniform struct_ name_
 #else
 #	define gl_VertexID gl_VertexIndex
 #	define gl_InstanceID gl_InstanceIndex
@@ -49,6 +50,7 @@ static const char* SHADER_HEADER = R"(#version 450 core
 #	define ANKI_SS_BINDING(set_, binding_) set = set_, binding = %u + binding_
 #	define ANKI_IMAGE_BINDING(set_, binding_) set = set_, binding = %u + binding_
 #	define ANKI_SPEC_CONST(binding_, type_, name_) layout(constant_id = binding_) const type_ name_ = type_(0)
+#	define ANKI_PUSH_CONSTANTS(struct_, name_) layout(push_constant) uniform pushConst_ {struct_ name_;}
 #endif
 
 #if %u

src/anki/gr/gl/CommandBuffer.cpp

@@ -1486,4 +1486,9 @@ void CommandBuffer::writeOcclusionQueryResultToBuffer(OcclusionQueryPtr query, P
 	self.pushBackNewCommand<WriteOcclResultToBuff>(query, offset, buff);
 }
 
+void CommandBuffer::setPushConstants(const void* data, U32 dataSize)
+{
+	ANKI_ASSERT(!"TODO");
+}
+
 } // end namespace anki

src/anki/gr/vulkan/CommandBuffer.cpp

@@ -358,4 +358,10 @@ Bool CommandBuffer::isEmpty() const
 	return self.isEmpty();
 }
 
+void CommandBuffer::setPushConstants(const void* data, U32 dataSize)
+{
+	ANKI_VK_SELF(CommandBufferImpl);
+	self.setPushConstants(data, dataSize);
+}
+
 } // end namespace anki

src/anki/gr/vulkan/CommandBufferImpl.h

@@ -329,6 +329,8 @@ public:
 
 	void copyBufferToBuffer(BufferPtr& src, PtrSize srcOffset, BufferPtr& dst, PtrSize dstOffset, PtrSize range);
 
+	void setPushConstants(const void* data, U32 dataSize);
+
 private:
 	StackAllocator<U8> m_alloc;
 
@@ -342,6 +344,7 @@ private:
 	ThreadId m_tid = ~ThreadId(0);
 #if ANKI_EXTRA_CHECKS
 	U32 m_commandCount = 0;
+	U32 m_setPushConstantsSize = 0;
 #endif
 
 	U m_rpCommandCount = 0; ///< Number of drawcalls or pushed cmdbs in rp.

src/anki/gr/vulkan/CommandBufferImpl.inl.h

@@ -306,6 +306,9 @@ inline void CommandBufferImpl::dispatchCompute(U32 groupCountX, U32 groupCountY,
 {
 	ANKI_ASSERT(m_computeProg);
 	ANKI_ASSERT(!!(m_flags & CommandBufferFlag::COMPUTE_WORK));
+	ANKI_ASSERT(m_state.tryGetBoundShaderProgram()->getReflectionInfo().m_pushConstantsSize == m_setPushConstantsSize
+				&& "Forgot to set pushConstants");
+
 	commandCommon();
 
 	// Bind descriptors
@@ -463,6 +466,9 @@ inline void CommandBufferImpl::drawcallCommon()
 	commandCommon();
 	ANKI_ASSERT(insideRenderPass() || secondLevel());
 	ANKI_ASSERT(m_subpassContents == VK_SUBPASS_CONTENTS_MAX_ENUM || m_subpassContents == VK_SUBPASS_CONTENTS_INLINE);
+	ANKI_ASSERT(m_state.tryGetBoundShaderProgram()->getReflectionInfo().m_pushConstantsSize == m_setPushConstantsSize
+				&& "Forgot to set pushConstants");
+
 	m_subpassContents = VK_SUBPASS_CONTENTS_INLINE;
 
 	if(ANKI_UNLIKELY(m_rpCommandCount == 0) && !secondLevel())
@@ -705,6 +711,10 @@ inline void CommandBufferImpl::bindShaderProgram(ShaderProgramPtr& prog)
 	}
 
 	m_microCmdb->pushObjectRef(prog);
+
+#if ANKI_EXTRA_CHECKS
+	m_setPushConstantsSize = 0;
+#endif
 }
 
 inline void CommandBufferImpl::copyBufferToBuffer(
@@ -734,4 +744,20 @@ inline Bool CommandBufferImpl::flipViewport() const
 		   && !!(getGrManagerImpl().getExtensions() & VulkanExtensions::KHR_MAINENANCE1);
 }
 
+inline void CommandBufferImpl::setPushConstants(const void* data, U32 dataSize)
+{
+	const ShaderProgramImpl* prog = m_state.tryGetBoundShaderProgram();
+	ANKI_ASSERT(prog && "Need have bound the ShaderProgram first");
+	ANKI_ASSERT(prog->getReflectionInfo().m_pushConstantsSize == dataSize
+				&& "The bound program should have push constants equal to the \"dataSize\" parameter");
+
+	ANKI_CMD(
+		vkCmdPushConstants(m_handle, prog->getPipelineLayout().getHandle(), VK_SHADER_STAGE_ALL, 0, dataSize, data),
+		ANY_OTHER_COMMAND);
+
+#if ANKI_EXTRA_CHECKS
+	m_setPushConstantsSize = dataSize;
+#endif
+}
+
 } // end namespace anki

src/anki/gr/vulkan/Pipeline.h

@@ -364,6 +364,18 @@ public:
 		}
 	}
 
+	const ShaderProgramImpl* tryGetBoundShaderProgram()
+	{
+		if(m_state.m_prog)
+		{
+			return static_cast<const ShaderProgramImpl*>(m_state.m_prog.get());
+		}
+		else
+		{
+			return nullptr;
+		}
+	}
+
 	void beginRenderPass(const FramebufferPtr& fb)
 	{
 		ANKI_ASSERT(m_rpass == VK_NULL_HANDLE);

src/anki/gr/vulkan/PipelineLayout.cpp

@@ -21,9 +21,9 @@ void PipelineLayoutFactory::destroy()
 }
 
 Error PipelineLayoutFactory::newPipelineLayout(
-	const WeakArray<DescriptorSetLayout>& dsetLayouts, PipelineLayout& layout)
+	const WeakArray<DescriptorSetLayout>& dsetLayouts, U32 pushConstantsSize, PipelineLayout& layout)
 {
-	U64 hash = 1;
+	U64 hash = computeHash(&pushConstantsSize, sizeof(pushConstantsSize));
 	Array<VkDescriptorSetLayout, MAX_DESCRIPTOR_SETS> vkDsetLayouts;
 	U dsetLayoutCount = 0;
 	for(const DescriptorSetLayout& dl : dsetLayouts)
@@ -33,7 +33,7 @@ Error PipelineLayoutFactory::newPipelineLayout(
 
 	if(dsetLayoutCount > 0)
 	{
-		hash = computeHash(&vkDsetLayouts[0], sizeof(vkDsetLayouts[0]) * dsetLayoutCount);
+		hash = appendHash(&vkDsetLayouts[0], sizeof(vkDsetLayouts[0]) * dsetLayoutCount, hash);
 	}
 
 	LockGuard<Mutex> lock(m_layoutsMtx);
@@ -54,6 +54,16 @@ Error PipelineLayoutFactory::newPipelineLayout(
 		ci.pSetLayouts = &vkDsetLayouts[0];
 		ci.setLayoutCount = dsetLayoutCount;
 
+		VkPushConstantRange pushConstantRange;
+		if(pushConstantsSize > 0)
+		{
+			pushConstantRange.offset = 0;
+			pushConstantRange.size = pushConstantsSize;
+			pushConstantRange.stageFlags = VK_SHADER_STAGE_ALL;
+			ci.pushConstantRangeCount = 1;
+			ci.pPushConstantRanges = &pushConstantRange;
+		}
+
 		VkPipelineLayout pplineLayHandle;
 		ANKI_VK_CHECK(vkCreatePipelineLayout(m_dev, &ci, nullptr, &pplineLayHandle));
 

src/anki/gr/vulkan/PipelineLayout.h

@@ -45,7 +45,8 @@ public:
 	void destroy();
 
 	/// @note It's thread-safe.
-	ANKI_USE_RESULT Error newPipelineLayout(const WeakArray<DescriptorSetLayout>& dsetLayouts, PipelineLayout& layout);
+	ANKI_USE_RESULT Error newPipelineLayout(
+		const WeakArray<DescriptorSetLayout>& dsetLayouts, U32 pushConstantsSize, PipelineLayout& layout);
 
 private:
 	GrAllocator<U8> m_alloc;

src/anki/gr/vulkan/ShaderImpl.cpp

@@ -192,6 +192,13 @@ void ShaderImpl::doReflection(ConstWeakArray<U8> spirv, SpecConstsVector& specCo
 
 	// Spec consts
 	specConstIds.m_vec = spvc.get_specialization_constants();
+
+	// Push consts
+	if(rsrc.push_constant_buffers.size() == 1)
+	{
+		U32 blockSize = spvc.get_declared_struct_size(spvc.get_type(rsrcActive.push_constant_buffers[0].base_type_id));
+		m_pushConstantsSize = blockSize;
+	}
 }
 
 } // end namespace anki

src/anki/gr/vulkan/ShaderImpl.h

@@ -27,6 +27,7 @@ public:
 	BitSet<MAX_VERTEX_ATTRIBUTES, U8> m_attributeMask = {false};
 	BitSet<MAX_DESCRIPTOR_SETS, U8> m_descriptorSetMask = {false};
 	Array<BitSet<MAX_BINDINGS_PER_DESCRIPTOR_SET, U8>, MAX_DESCRIPTOR_SETS> m_activeBindingMask = {{{false}, {false}}};
+	U32 m_pushConstantsSize = 0;
 
 	ShaderImpl(GrManager* manager, CString name)
 		: Shader(manager, name)

src/anki/gr/vulkan/ShaderProgramImpl.cpp

@@ -78,6 +78,8 @@ Error ShaderProgramImpl::init(const ShaderProgramInitInfo& inf)
 					bindings[set][counts[set]++] = simpl.m_bindings[set][i];
 				}
 			}
+
+			m_refl.m_pushConstantsSize = max(m_refl.m_pushConstantsSize, simpl.m_pushConstantsSize);
 		}
 
 		if(counts[set])
@@ -101,7 +103,8 @@ Error ShaderProgramImpl::init(const ShaderProgramInitInfo& inf)
 	//
 	WeakArray<DescriptorSetLayout> dsetLayouts(
 		(descriptorSetCount) ? &m_descriptorSetLayouts[0] : nullptr, descriptorSetCount);
-	ANKI_CHECK(getGrManagerImpl().getPipelineLayoutFactory().newPipelineLayout(dsetLayouts, m_pplineLayout));
+	ANKI_CHECK(getGrManagerImpl().getPipelineLayoutFactory().newPipelineLayout(
+		dsetLayouts, m_refl.m_pushConstantsSize, m_pplineLayout));
 
 	// Get some masks
 	//

src/anki/gr/vulkan/ShaderProgramImpl.h

@@ -25,6 +25,7 @@ public:
 	BitSet<MAX_VERTEX_ATTRIBUTES, U8> m_attributeMask = {false};
 	BitSet<MAX_DESCRIPTOR_SETS, U8> m_descriptorSetMask = {false};
 	Array<BitSet<MAX_BINDINGS_PER_DESCRIPTOR_SET, U8>, MAX_DESCRIPTOR_SETS> m_activeBindingMask = {{{false}, {false}}};
+	U32 m_pushConstantsSize = 0;
 };
 
 /// Shader program implementation.

tests/gr/Gr.cpp

@@ -176,8 +176,6 @@ layout(location = 0) in vec2 in_uv;
 layout(ANKI_TEX_BINDING(0, 0)) uniform sampler2D u_tex0;
 layout(ANKI_TEX_BINDING(0, 1)) uniform sampler2D u_tex1;
 
-ANKI_USING_FRAG_COORD(768)
-
 void main()
 {
 	if(anki_fragCoord.x < 1024 / 2)
@@ -2054,4 +2052,114 @@ void main()
 	COMMON_END()
 }
 
+ANKI_TEST(Gr, PushConsts)
+{
+	COMMON_BEGIN()
+
+	static const char* VERT_SRC = R"(
+struct PC
+{
+	vec4 color;
+	vec4 color2;
+};
+ANKI_PUSH_CONSTANTS(PC, regs);
+	
+out gl_PerVertex
+{
+	vec4 gl_Position;
+};
+	
+layout(location = 0) out vec4 out_color;
+
+void main()
+{
+	vec2 uv = vec2(gl_VertexID & 1, gl_VertexID >> 1) * 2.0;
+	vec2 pos = uv * 2.0 - 1.0;
+	gl_Position = vec4(pos, 0.0, 1.0);
+	
+	out_color = regs.color;
+}
+)";
+
+	static const char* FRAG_SRC = R"(
+struct PC
+{
+	vec4 color;
+	vec4 color2;
+};
+ANKI_PUSH_CONSTANTS(PC, regs);
+	
+layout(location = 0) in vec4 in_color;
+layout(location = 0) out vec4 out_color;
+
+layout(ANKI_SS_BINDING(0, 0)) buffer s_
+{
+	uvec4 u_result;
+};
+
+void main()
+{
+	out_color = vec4(1.0);
+
+	if(gl_FragCoord.x == 0.5 && gl_FragCoord.y == 0.5)
+	{
+		if(in_color != vec4(1.0, 0.0, 1.0, 0.0) || regs.color2 != vec4(0.0, 1.0, 0.25, 0.5))
+		{
+			u_result = uvec4(1u);
+		}
+		else
+		{
+			u_result = uvec4(2u);
+		}
+	}
+}
+)";
+
+	ShaderProgramPtr prog = createProgram(VERT_SRC, FRAG_SRC, *gr);
+
+	// Create the result buffer
+	BufferPtr resultBuff = gr->newBuffer(
+		BufferInitInfo(sizeof(UVec4), BufferUsageBit::STORAGE_ALL | BufferUsageBit::FILL, BufferMapAccessBit::READ));
+
+	// Draw
+	gr->beginFrame();
+
+	CommandBufferInitInfo cinit;
+	cinit.m_flags = CommandBufferFlag::GRAPHICS_WORK;
+	CommandBufferPtr cmdb = gr->newCommandBuffer(cinit);
+
+	cmdb->fillBuffer(resultBuff, 0, resultBuff->getSize(), 0);
+	cmdb->setBufferBarrier(
+		resultBuff, BufferUsageBit::FILL, BufferUsageBit::STORAGE_FRAGMENT_WRITE, 0, resultBuff->getSize());
+
+	cmdb->setViewport(0, 0, WIDTH, HEIGHT);
+	cmdb->bindShaderProgram(prog);
+
+	struct PushConstants
+	{
+		Vec4 m_color = Vec4(1.0, 0.0, 1.0, 0.0);
+		Vec4 m_color1 = Vec4(0.0, 1.0, 0.25, 0.5);
+	} pc;
+	cmdb->setPushConstants(&pc, sizeof(pc));
+
+	cmdb->bindStorageBuffer(0, 0, resultBuff, 0, resultBuff->getSize());
+	cmdb->beginRenderPass(createDefaultFb(*gr), {}, {});
+	cmdb->drawArrays(PrimitiveTopology::TRIANGLES, 3);
+	cmdb->endRenderPass();
+	cmdb->flush();
+
+	gr->swapBuffers();
+	gr->finish();
+
+	// Get the result
+	UVec4* result = static_cast<UVec4*>(resultBuff->map(0, resultBuff->getSize(), BufferMapAccessBit::READ));
+	ANKI_TEST_EXPECT_EQ(result->x(), 2);
+	ANKI_TEST_EXPECT_EQ(result->y(), 2);
+	ANKI_TEST_EXPECT_EQ(result->z(), 2);
+	ANKI_TEST_EXPECT_EQ(result->w(), 2);
+	resultBuff->unmap();
+
+	COMMON_END()
+}
+
 } // end namespace anki