Merge pull request #31 from godlikepanos/push_constants_everywhere

Add push constant support in ShaderProgramResource.

src/anki/gr/ShaderCompiler.cpp

@@ -40,7 +40,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, row_major) uniform struct_ name_
+#	define ANKI_PUSH_CONSTANTS(struct_, name_) layout(location = %u, row_major) uniform struct_ name_
 #else
 #	define gl_VertexID gl_VertexIndex
 #	define gl_InstanceID gl_InstanceIndex
@@ -296,6 +296,7 @@ Error ShaderCompiler::compile(CString source, const ShaderCompilerOptions& optio
 		MAX_STORAGE_BUFFER_BINDINGS,
 		MAX_TEXTURE_BINDINGS,
 		MAX_IMAGE_BINDINGS,
+		MAX_TEXTURE_BINDINGS * MAX_DESCRIPTOR_SETS, // Push constant location
 		// VK bindings
 		0,
 		MAX_TEXTURE_BINDINGS,

src/anki/gr/gl/CommandBuffer.cpp

@@ -1509,24 +1509,33 @@ void CommandBuffer::setPushConstants(const void* data, U32 dataSize)
 
 			for(const ShaderProgramImplReflection::Uniform& uni : refl.m_uniforms)
 			{
-				const U idx = uni.m_location;
+				const U8* data = reinterpret_cast<const U8*>(&m_data[0]) + uni.m_pushConstantOffset;
 				const U count = uni.m_arrSize;
 				const GLint loc = uni.m_location;
 
 				switch(uni.m_type)
 				{
 				case ShaderVariableDataType::VEC4:
-					glUniform4fv(loc, count, reinterpret_cast<const GLfloat*>(&m_data[idx]));
+					glUniform4fv(loc, count, reinterpret_cast<const GLfloat*>(data));
 					break;
 				case ShaderVariableDataType::IVEC4:
-					glUniform4iv(loc, count, reinterpret_cast<const GLint*>(&m_data[idx]));
+					glUniform4iv(loc, count, reinterpret_cast<const GLint*>(data));
 					break;
 				case ShaderVariableDataType::UVEC4:
-					glUniform4uiv(loc, count, reinterpret_cast<const GLuint*>(&m_data[idx]));
+					glUniform4uiv(loc, count, reinterpret_cast<const GLuint*>(data));
 					break;
 				case ShaderVariableDataType::MAT4:
-					glUniformMatrix4fv(loc, count, false, reinterpret_cast<const GLfloat*>(&m_data[idx]));
+					glUniformMatrix4fv(loc, count, false, reinterpret_cast<const GLfloat*>(data));
 					break;
+				case ShaderVariableDataType::MAT3:
+				{
+					// Remove the padding
+					ANKI_ASSERT(count == 1 && "TODO");
+					const Mat3x4* m34 = reinterpret_cast<const Mat3x4*>(data);
+					Mat3 m3(m34->getRotationPart());
+					glUniformMatrix3fv(loc, count, false, reinterpret_cast<const GLfloat*>(&m3));
+					break;
+				}
 				default:
 					ANKI_ASSERT(!"TODO");
 				}

src/anki/gr/gl/ShaderProgramImpl.cpp

@@ -114,30 +114,39 @@ const ShaderProgramImplReflection& ShaderProgramImpl::getReflection()
 			GLint size;
 			Array<char, 128> name;
 			glGetActiveUniform(getGlName(), i, sizeof(name), &len, &size, &type, &name[0]);
-
 			name[len] = '\0';
+
+			if(CString(&name[0]).find("gl_") == 0)
+			{
+				// Builtin, skip
+				continue;
+			}
+
 			GLint location = glGetUniformLocation(getGlName(), &name[0]);
+			if(location < I(MAX_TEXTURE_BINDINGS * MAX_DESCRIPTOR_SETS))
+			{
+				// It must be a sampled image, skip it
+				continue;
+			}
 
 			// Store those info
 			ShaderVariableDataType akType = ShaderVariableDataType::NONE;
-			U32 dataSize = 0;
 			switch(type)
 			{
 			case GL_FLOAT_VEC4:
 				akType = ShaderVariableDataType::VEC4;
-				dataSize = 16;
 				break;
 			case GL_INT_VEC4:
 				akType = ShaderVariableDataType::IVEC4;
-				dataSize = 16;
 				break;
 			case GL_UNSIGNED_INT_VEC4:
 				akType = ShaderVariableDataType::UVEC4;
-				dataSize = 16;
 				break;
 			case GL_FLOAT_MAT4:
 				akType = ShaderVariableDataType::MAT4;
-				dataSize = 16 * 4;
+				break;
+			case GL_FLOAT_MAT3:
+				akType = ShaderVariableDataType::MAT3;
 				break;
 			default:
 				ANKI_ASSERT(!"Unsupported type");
@@ -149,6 +158,38 @@ const ShaderProgramImplReflection& ShaderProgramImpl::getReflection()
 			uni.m_arrSize = size;
 
 			m_refl.m_uniforms.emplaceBack(getAllocator(), uni);
+		}
+
+		// Sort the uniforms
+		std::sort(m_refl.m_uniforms.getBegin(),
+			m_refl.m_uniforms.getEnd(),
+			[](const ShaderProgramImplReflection::Uniform& a, const ShaderProgramImplReflection::Uniform& b) {
+				return a.m_location < b.m_location;
+			});
+
+		// Now calculate the offset inside the push constant buffer
+		m_refl.m_uniformDataSize = 0;
+		for(ShaderProgramImplReflection::Uniform& uni : m_refl.m_uniforms)
+		{
+			U32 dataSize = 0;
+			switch(uni.m_type)
+			{
+			case ShaderVariableDataType::VEC4:
+			case ShaderVariableDataType::IVEC4:
+			case ShaderVariableDataType::UVEC4:
+				dataSize = sizeof(F32) * 4;
+				break;
+			case ShaderVariableDataType::MAT4:
+				dataSize = sizeof(F32) * 16;
+				break;
+			case ShaderVariableDataType::MAT3:
+				dataSize = sizeof(F32) * 12;
+				break;
+			default:
+				ANKI_ASSERT(!"Unsupported type");
+			}
+
+			uni.m_pushConstantOffset = m_refl.m_uniformDataSize;
 			m_refl.m_uniformDataSize += dataSize * uni.m_arrSize;
 		}
 	}

src/anki/gr/gl/ShaderProgramImpl.h

@@ -20,6 +20,7 @@ public:
 	struct Uniform
 	{
 		I32 m_location;
+		U32 m_pushConstantOffset;
 		ShaderVariableDataType m_type;
 		U8 m_arrSize;
 	};

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

@@ -306,7 +306,7 @@ 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
+	ANKI_ASSERT(m_computeProg->getReflectionInfo().m_pushConstantsSize == m_setPushConstantsSize
 				&& "Forgot to set pushConstants");
 
 	commandCommon();
@@ -466,7 +466,7 @@ 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
+	ANKI_ASSERT(m_graphicsProg->getReflectionInfo().m_pushConstantsSize == m_setPushConstantsSize
 				&& "Forgot to set pushConstants");
 
 	m_subpassContents = VK_SUBPASS_CONTENTS_INLINE;
@@ -747,7 +747,7 @@ inline Bool CommandBufferImpl::flipViewport() const
 inline void CommandBufferImpl::setPushConstants(const void* data, U32 dataSize)
 {
 	ANKI_ASSERT(data && dataSize && dataSize % 16 == 0);
-	const ShaderProgramImpl* prog = m_state.tryGetBoundShaderProgram();
+	const ShaderProgramImpl* prog = (m_graphicsProg) ? m_graphicsProg : m_computeProg;
 	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");

src/anki/gr/vulkan/Pipeline.h

@@ -364,18 +364,6 @@ 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/resource/ShaderProgramResource.cpp

@@ -344,7 +344,7 @@ Error ShaderProgramResource::load(const ResourceFilename& filename, Bool async)
 	// <shader> again
 	inputVarCount = 0;
 	StringListAuto constSrcList(getTempAllocator());
-	StringListAuto blockSrcList(getTempAllocator());
+	StringListAuto uniStructList(getTempAllocator());
 	StringListAuto globalsSrcList(getTempAllocator());
 	StringListAuto definesSrcList(getTempAllocator());
 	ShaderTypeBit presentShaders = ShaderTypeBit::NONE;
@@ -384,7 +384,7 @@ Error ShaderProgramResource::load(const ResourceFilename& filename, Bool async)
 		if(inputsEl)
 		{
 			ANKI_CHECK(
-				parseInputs(inputsEl, inputVarCount, constSrcList, blockSrcList, globalsSrcList, definesSrcList));
+				parseInputs(inputsEl, inputVarCount, constSrcList, uniStructList, globalsSrcList, definesSrcList));
 		}
 
 		// <source>
@@ -398,7 +398,7 @@ Error ShaderProgramResource::load(const ResourceFilename& filename, Bool async)
 	// <inputs>
 	if(inputsEl)
 	{
-		ANKI_CHECK(parseInputs(inputsEl, inputVarCount, constSrcList, blockSrcList, globalsSrcList, definesSrcList));
+		ANKI_CHECK(parseInputs(inputsEl, inputVarCount, constSrcList, uniStructList, globalsSrcList, definesSrcList));
 	}
 
 	ANKI_ASSERT(inputVarCount == m_inputVars.getSize());
@@ -435,10 +435,25 @@ Error ShaderProgramResource::load(const ResourceFilename& filename, Bool async)
 	}
 
 	StringAuto backedUboSrc(getTempAllocator());
-	if(!blockSrcList.isEmpty())
+	if(!uniStructList.isEmpty())
 	{
-		blockSrcList.pushBack("};\n");
-		blockSrcList.join("", backedUboSrc);
+		// Create the uniform struct
+		uniStructList.pushFront("struct spr_Uniforms_ {\n");
+		uniStructList.pushBack("};\n");
+
+		// Create the uniform block
+		uniStructList.pushBack("#if USE_PUSH_CONSTANTS == 0\n");
+		uniStructList.pushBackSprintf(
+			"layout(ANKI_UBO_BINDING(%u, 0), std140, row_major) uniform spr_Block_ {spr_Uniforms_ spr_unis_;};\n",
+			U(m_descriptorSet));
+
+		// Create the push constants
+		uniStructList.pushBack("#else // if USE_PUSH_CONSTANTS == 0\n");
+		uniStructList.pushBackSprintf("ANKI_PUSH_CONSTANTS(spr_Uniforms_, spr_unis_);\n");
+		uniStructList.pushBack("#endif // if USE_PUSH_CONSTANTS == 0\n");
+
+		// Done
+		uniStructList.join("", backedUboSrc);
 	}
 
 	StringAuto backedGlobalsSrc(getTempAllocator());
@@ -490,7 +505,7 @@ Error ShaderProgramResource::load(const ResourceFilename& filename, Bool async)
 Error ShaderProgramResource::parseInputs(XmlElement& inputsEl,
 	U& inputVarCount,
 	StringListAuto& constsSrc,
-	StringListAuto& blockSrc,
+	StringListAuto& uniStruct,
 	StringListAuto& globalsSrc,
 	StringListAuto& definesSrc)
 {
@@ -682,35 +697,34 @@ Error ShaderProgramResource::parseInputs(XmlElement& inputsEl,
 		// Append to ubo source
 		if(var.inBlock())
 		{
-			if(blockSrc.isEmpty())
-			{
-				blockSrc.pushBackSprintf(
-					"layout(ANKI_UBO_BINDING(%u, 0), std140, row_major) uniform sprubo00_ {\n", U(m_descriptorSet));
-			}
-
-			blockSrc.pushBackSprintf("#if %s_DEFINED == 1\n", &name[0]);
+			uniStruct.pushBackSprintf("#if %s_DEFINED == 1\n", name.cstr());
+			globalsSrc.pushBackSprintf("#if %s_DEFINED == 1\n", name.cstr());
 
 			if(var.m_instanced)
 			{
-				blockSrc.pushBackSprintf("#if %s > 1\n", &m_instancingMutator->getName()[0]);
-				blockSrc.pushBackSprintf(
-					"%s %s_INSTARR[%s];\n", &typeTxt[0], &name[0], &m_instancingMutator->getName()[0]);
-				blockSrc.pushBack("#else\n");
-				blockSrc.pushBackSprintf("%s %s;\n", &typeTxt[0], &name[0]);
-				blockSrc.pushBack("#endif\n");
-
-				globalsSrc.pushBackSprintf("#if defined(ANKI_VERTEX_SHADER) && %s_DEFINED == 1 && %s > 1\n",
-					&name[0],
-					&m_instancingMutator->getName()[0]);
-				globalsSrc.pushBackSprintf("%s %s = %s_INSTARR[gl_InstanceID];\n", &typeTxt[0], &name[0], &name[0]);
-				globalsSrc.pushBack("#else\n// TODO\n#endif\n");
+				uniStruct.pushBackSprintf("#if %s > 1\n", &m_instancingMutator->getName()[0]);
+				uniStruct.pushBackSprintf("%s %s[%s];\n", &typeTxt[0], &name[0], &m_instancingMutator->getName()[0]);
+				uniStruct.pushBack("#else\n");
+				uniStruct.pushBackSprintf("%s %s;\n", &typeTxt[0], &name[0]);
+				uniStruct.pushBack("#endif\n");
+
+				globalsSrc.pushBackSprintf("#if %s > 1\n", m_instancingMutator->getName().cstr());
+				globalsSrc.pushBack("#if defined(ANKI_VERTEX_SHADER)\n");
+				globalsSrc.pushBackSprintf("%s %s = spr_unis_.%s[gl_InstanceID];\n", &typeTxt[0], &name[0], &name[0]);
+				globalsSrc.pushBack("#endif\n");
+				globalsSrc.pushBack("#else\n");
+				globalsSrc.pushBackSprintf("%s %s = spr_unis_.%s;\n", typeTxt.cstr(), name.cstr(), name.cstr());
+				globalsSrc.pushBack("#endif\n");
 			}
 			else
 			{
-				blockSrc.pushBackSprintf("%s %s;\n", &typeTxt[0], &name[0]);
+				uniStruct.pushBackSprintf("%s %s;\n", &typeTxt[0], &name[0]);
+
+				globalsSrc.pushBackSprintf("%s %s = spr_unis_.%s;\n", typeTxt.cstr(), name.cstr(), name.cstr());
 			}
 
-			blockSrc.pushBack("#endif\n");
+			uniStruct.pushBack("#endif\n");
+			globalsSrc.pushBack("#endif\n");
 		}
 
 		// Append the textures to global area
@@ -1053,8 +1067,14 @@ void ShaderProgramResource::initVariant(ConstWeakArray<ShaderProgramResourceMuta
 		}
 	}
 
+	// Check if we can use push constants
+	variant.m_usesPushConstants =
+		instanceCount == 1
+		&& variant.m_uniBlockSize <= getManager().getGrManager().getDeviceCapabilities().m_pushConstantsSize;
+
 	// Write the source header
 	StringListAuto shaderHeaderSrc(getTempAllocator());
+	shaderHeaderSrc.pushBackSprintf("#define USE_PUSH_CONSTANTS %d\n", I(variant.m_usesPushConstants));
 
 	for(const ShaderProgramResourceMutation& m : mutations)
 	{

src/anki/resource/ShaderProgramResource.h

@@ -198,6 +198,11 @@ public:
 		return m_uniBlockSize;
 	}
 
+	Bool usePushConstants() const
+	{
+		return m_usesPushConstants;
+	}
+
 	const ShaderVariableBlockInfo& getVariableBlockInfo(const ShaderProgramResourceInputVariable& var) const
 	{
 		ANKI_ASSERT(!var.isTexture() && variableActive(var));
@@ -234,6 +239,7 @@ private:
 	DynamicArray<ShaderVariableBlockInfo> m_blockInfos;
 	U32 m_uniBlockSize = 0;
 	DynamicArray<I16> m_texUnits;
+	Bool8 m_usesPushConstants = false;
 };
 
 /// The value of a constant.

src/anki/scene/components/RenderComponent.cpp

@@ -39,11 +39,25 @@ void RenderComponent::allocateAndSetupUniforms(
 	const ShaderProgramResourceVariant& progVariant = variant.getShaderProgramResourceVariant();
 
 	// Allocate uniform memory
-	StagingGpuMemoryToken token;
-	U8* uniforms =
-		static_cast<U8*>(alloc.allocateFrame(variant.getUniformBlockSize(), StagingGpuMemoryType::UNIFORM, token));
-	void* const uniformsBegin = uniforms;
-	const void* const uniformsEnd = uniforms + variant.getUniformBlockSize();
+	U8* uniforms;
+	void* uniformsBegin;
+	const void* uniformsEnd;
+	Array<U8, 256> pushConsts;
+	if(!progVariant.usePushConstants())
+	{
+		StagingGpuMemoryToken token;
+		uniforms =
+			static_cast<U8*>(alloc.allocateFrame(variant.getUniformBlockSize(), StagingGpuMemoryType::UNIFORM, token));
+
+		ctx.m_commandBuffer->bindUniformBuffer(set, 0, token.m_buffer, token.m_offset, token.m_range);
+	}
+	else
+	{
+		uniforms = &pushConsts[0];
+	}
+
+	uniformsBegin = uniforms;
+	uniformsEnd = uniforms + variant.getUniformBlockSize();
 
 	// Iterate variables
 	for(auto it = m_vars.getBegin(); it != m_vars.getEnd(); ++it)
@@ -236,7 +250,10 @@ void RenderComponent::allocateAndSetupUniforms(
 		} // end switch
 	}
 
-	ctx.m_commandBuffer->bindUniformBuffer(set, 0, token.m_buffer, token.m_offset, token.m_range);
+	if(progVariant.usePushConstants())
+	{
+		ctx.m_commandBuffer->setPushConstants(uniformsBegin, variant.getUniformBlockSize());
+	}
 }
 
 } // end namespace anki

src/anki/util/StringList.h

@@ -64,6 +64,17 @@ public:
 		Base::getBack() = std::move(str);
 	}
 
+	/// Push at the beginning of the list a formated string.
+	template<typename... TArgs>
+	void pushFrontSprintf(Allocator alloc, const TArgs&... args)
+	{
+		String str;
+		str.sprintf(alloc, args...);
+
+		Base::emplaceFront(alloc);
+		Base::getFront() = std::move(str);
+	}
+
 	/// Push back plain CString.
 	void pushBack(Allocator alloc, CString cstr)
 	{
@@ -134,6 +145,13 @@ public:
 		Base::pushBackSprintf(m_alloc, args...);
 	}
 
+	/// Push at the beginning of the list a formated string
+	template<typename... TArgs>
+	void pushFrontSprintf(const TArgs&... args)
+	{
+		Base::pushFrontSprintf(m_alloc, args...);
+	}
+
 	/// Push back plain CString.
 	void pushBack(CString cstr)
 	{