Add the code of the task and mesh shaders in the GBufferGeneric shader

AnKi/Resource/ShaderProgramResource.cpp

@@ -158,8 +158,8 @@ Error ShaderProgramResource::load(const ResourceFilename& filename, [[maybe_unus
 		if(m_shaderStages != (ShaderTypeBit::kAnyHit | ShaderTypeBit::kClosestHit) && m_shaderStages != ShaderTypeBit::kMiss
 		   && m_shaderStages != ShaderTypeBit::kRayGen)
 		{
-			ANKI_RESOURCE_LOGE("Any and closest hit shaders shouldn't coexist with other stages. Miss can't coexist "
-							   "with other stages. Raygen can't coexist with other stages as well");
+			ANKI_RESOURCE_LOGE("Any and closest hit shaders shouldn't coexist with other stages. Miss can't coexist with other stages. Raygen can't "
+							   "coexist with other stages as well");
 			return Error::kUserData;
 		}
 	}
@@ -200,7 +200,7 @@ void ShaderProgramResource::getOrCreateVariant(const ShaderProgramResourceVarian
 	ANKI_ASSERT(info.m_setConstants.getSetBitCount() == m_consts.getSize());
 
 	// Compute variant hash
-	U64 hash = 0;
+	U64 hash = info.m_meshShaders * 0xBAD;
 	if(m_mutators.getSize())
 	{
 		hash = computeHash(info.m_mutation.getBegin(), m_mutators.getSize() * sizeof(info.m_mutation[0]));
@@ -354,6 +354,9 @@ ShaderProgramResourceVariant* ShaderProgramResource::createNewVariant(const Shad
 	// Time to init the shaders
 	if(!!(m_shaderStages & (ShaderTypeBit::kAllGraphics | ShaderTypeBit::kCompute)))
 	{
+		const ShaderTypeBit stages =
+			m_shaderStages & ((info.m_meshShaders) ? ~ShaderTypeBit::kAllLegacyGeometry : ~ShaderTypeBit::kAllModernGeometry);
+
 		// Create the program name
 		String progName;
 		getFilepathFilename(getFilename(), progName);
@@ -365,7 +368,7 @@ ShaderProgramResourceVariant* ShaderProgramResource::createNewVariant(const Shad
 
 		ShaderProgramInitInfo progInf(cprogName);
 		Array<ShaderPtr, U32(ShaderType::kCount)> shaderRefs; // Just for refcounting
-		for(ShaderType shaderType : EnumBitsIterable<ShaderType, ShaderTypeBit>(m_shaderStages))
+		for(ShaderType shaderType : EnumBitsIterable<ShaderType, ShaderTypeBit>(stages))
 		{
 			ShaderInitInfo inf(cprogName);
 			inf.m_shaderType = shaderType;

AnKi/Resource/ShaderProgramResource.h

@@ -166,6 +166,8 @@ public:
 
 	ShaderProgramResourceVariantInitInfo& addMutation(CString name, MutatorValue t);
 
+	void requestMeshShaders(Bool request);
+
 private:
 	static constexpr U32 kMaxConstants = 32;
 	static constexpr U32 kMaxMutators = 32;
@@ -177,6 +179,8 @@ private:
 
 	Array<MutatorValue, kMaxMutators> m_mutation; ///< The order of storing the values is important. It will be hashed.
 	BitSet<kMaxMutators> m_setMutators = {false};
+
+	Bool m_meshShaders = false;
 };
 
 /// Shader program resource. It loads special AnKi programs.
@@ -306,6 +310,12 @@ inline ShaderProgramResourceVariantInitInfo& ShaderProgramResourceVariantInitInf
 	m_setMutators.set(mutatorIdx);
 	return *this;
 }
+
+inline void ShaderProgramResourceVariantInitInfo::requestMeshShaders(Bool request)
+{
+	ANKI_ASSERT(!!(m_ptr->getStages() & ShaderTypeBit::kAllModernGeometry));
+	m_meshShaders = request;
+}
 /// @}
 
 } // end namespace anki

AnKi/Scene/RenderStateBucket.h

@@ -21,6 +21,7 @@ public:
 	ShaderProgramPtr m_program;
 	PrimitiveTopology m_primitiveTopology = PrimitiveTopology::kTriangles;
 	Bool m_indexedDrawcall = true;
+	Bool m_meshShaders = false;
 };
 
 class RenderStateBucketIndex

AnKi/Shaders/GBufferGeneric.ankiprog

@@ -99,7 +99,15 @@ struct FragOut
 };
 #endif
 
-#pragma anki start vert
+struct MeshShaderPayload
+{
+	U32 m_meshletIndices[ANKI_TASK_SHADER_THREADGROUP_SIZE];
+	U32 m_worldTransformsOffset;
+	U32 m_constantsOffset;
+	U32 m_boneTransformsOrParticleEmitterOffset;
+	Vec3 m_positionTranslation;
+	F32 m_positionScale;
+};
 
 struct Mat3x4_2
 {
@@ -107,26 +115,25 @@ struct Mat3x4_2
 	Mat3x4 m_b;
 };
 
-Mat3x4_2 loadBoneTransforms(UnpackedMeshVertex vert, GpuSceneRenderableVertex renderable, U32 index)
+Mat3x4_2 loadBoneTransforms(UnpackedMeshVertex vert, U32 boneTransformsOffset, U32 index)
 {
 	const U32 boneIdx = vert.m_boneIndices[index];
-	U32 byteOffset = renderable.m_boneTransformsOrParticleEmitterOffset;
+	U32 byteOffset = boneTransformsOffset;
 	byteOffset += boneIdx * sizeof(Mat3x4) * 2;
 	return g_gpuScene.Load<Mat3x4_2>(byteOffset);
 }
 
 #if ANKI_BONES
-void skinning(UnpackedMeshVertex vert, GpuSceneRenderableVertex renderable, inout Vec3 pos, inout Vec3 prevPos, inout RVec3 normal,
-			  inout RVec4 tangent)
+void skinning(UnpackedMeshVertex vert, U32 boneTransformsOffset, inout Vec3 pos, inout Vec3 prevPos, inout RVec3 normal, inout RVec4 tangent)
 {
-	Mat3x4_2 mats = loadBoneTransforms(vert, renderable, 0);
+	Mat3x4_2 mats = loadBoneTransforms(vert, boneTransformsOffset, 0);
 
 	Mat3x4 skinMat = mats.m_a * vert.m_boneWeights[0];
 	Mat3x4 prevSkinMat = mats.m_b * vert.m_boneWeights[0];
 
 	[unroll] for(U32 i = 1u; i < 4u; ++i)
 	{
-		mats = loadBoneTransforms(vert, renderable, i);
+		mats = loadBoneTransforms(vert, boneTransformsOffset, i);
 
 		skinMat = skinMat + mats.m_a * vert.m_boneWeights[i];
 		prevSkinMat = prevSkinMat + mats.m_b * vert.m_boneWeights[i];
@@ -167,6 +174,8 @@ void velocity(Mat3x4 worldTransform, Mat3x4 prevWorldTransform, Vec3 prevLocalPo
 }
 #endif
 
+#pragma anki start vert
+
 VertOut main(VertIn input)
 {
 	VertOut output;
@@ -188,7 +197,7 @@ VertOut main(VertIn input)
 
 	// Do stuff
 #if ANKI_BONES
-	skinning(vert, renderable, vert.m_position, prevPos, vert.m_normal, vert.m_tangent);
+	skinning(vert, renderable.m_boneTransformsOrParticleEmitterOffset, vert.m_position, prevPos, vert.m_normal, vert.m_tangent);
 #endif
 
 	output.m_position = Vec4(mul(worldTransform, Vec4(vert.m_position, 1.0)), 1.0);
@@ -209,6 +218,95 @@ VertOut main(VertIn input)
 
 #pragma anki end vert
 
+#pragma anki start task
+
+groupshared MeshShaderPayload s_payload;
+
+[numthreads(ANKI_TASK_SHADER_THREADGROUP_SIZE, 1, 1)] void main(U32 svGroupId : SV_GROUPID, U32 svGroupIndex : SV_GROUPINDEX)
+{
+	const GpuSceneTaskShaderPayload inPayload = g_taskShaderPayloads[svGroupId];
+
+	const U32 meshletCount = (inPayload.m_firstMeshlet_26bit_meshletCountMinusOne_6bit & 63u) + 1u;
+	const U32 firstMeshlet = inPayload.m_firstMeshlet_26bit_meshletCountMinusOne_6bit >> 6u;
+
+	if(svGroupIndex < meshletCount)
+	{
+		s_payload.m_meshletIndices[svGroupIndex] = firstMeshlet + svGroupIndex;
+		s_payload.m_worldTransformsOffset = inPayload.m_worldTransformsOffset;
+		s_payload.m_constantsOffset = inPayload.m_constantsOffset;
+		s_payload.m_boneTransformsOrParticleEmitterOffset = inPayload.m_boneTransformsOrParticleEmitterOffset;
+		s_payload.m_positionScale = inPayload.m_positionScale;
+		s_payload.m_positionTranslation = inPayload.m_positionTranslation;
+	}
+
+	GroupMemoryBarrierWithGroupSync();
+	DispatchMesh(meshletCount, 1, 1, s_payload);
+}
+
+#pragma anki end task
+
+#pragma anki start mesh
+
+constexpr U32 g_dummy = 0; // The formater is getting confused so add this
+
+[numthreads(ANKI_MESH_SHADER_THREADGROUP_SIZE, 1, 1)] [outputtopology("triangle")] void
+main(in payload MeshShaderPayload payload, out vertices VertOut verts[kMaxVerticesPerMeshlet], out indices UVec3 indices[kMaxPrimitivesPerMeshlet],
+	 U32 svGroupId : SV_GROUPID, U32 svGroupIndex : SV_GROUPINDEX)
+{
+	const Meshlet meshlet = g_meshlets[payload.m_meshletIndices[svGroupId]];
+	const U32 primCount = meshlet.m_primitiveCount_R16_Uint_vertexCount_R16_Uint >> 16u;
+	const U32 vertCount = meshlet.m_primitiveCount_R16_Uint_vertexCount_R16_Uint & 0xFFFFu;
+
+	SetMeshOutputCounts(vertCount, primCount);
+
+	// Write the verts
+	if(svGroupIndex < vertCount)
+	{
+		VertOut output;
+
+		UnpackedMeshVertex vert = loadVertex(meshlet, svGroupIndex, ANKI_BONES, payload.m_positionScale, payload.m_positionTranslation);
+
+		const Mat3x4 worldTransform = g_gpuScene.Load<Mat3x4>(payload.m_worldTransformsOffset);
+		const Mat3x4 prevWorldTransform = g_gpuScene.Load<Mat3x4>(payload.m_worldTransformsOffset + sizeof(Mat3x4));
+		ANKI_MAYBE_UNUSED(prevWorldTransform);
+
+#if UVS
+		output.m_uv = vert.m_uv;
+#endif
+		Vec3 prevPos = vert.m_position;
+		ANKI_MAYBE_UNUSED(prevPos);
+		output.m_constantsOffset = payload.m_constantsOffset;
+
+		// Do stuff
+#if ANKI_BONES
+		skinning(vert, payload.m_boneTransformsOrParticleEmitterOffset, vert.m_position, prevPos, vert.m_normal, vert.m_tangent);
+#endif
+
+		output.m_position = Vec4(mul(worldTransform, Vec4(vert.m_position, 1.0)), 1.0);
+		output.m_position = mul(g_globalConstants.m_viewProjectionMatrix, output.m_position);
+
+#if ANKI_TECHNIQUE == ANKI_RENDERING_TECHNIQUE_GBUFFER
+		output.m_normal = mul(worldTransform, Vec4(vert.m_normal, 0.0));
+		output.m_tangent = mul(worldTransform, Vec4(vert.m_tangent.xyz, 0.0));
+		output.m_bitangent = cross(output.m_normal, output.m_tangent) * vert.m_tangent.w;
+#endif
+
+#if REALLY_VELOCITY
+		velocity(worldTransform, prevWorldTransform, prevPos, output);
+#endif
+
+		verts[svGroupIndex] = output;
+	}
+
+	// Write the indices
+	if(svGroupIndex < primCount)
+	{
+		indices[svGroupIndex] = g_unifiedGeom_R8G8B8A8_Uint[meshlet.m_firstPrimitive + svGroupIndex].xyz;
+	}
+}
+
+#pragma anki end mesh
+
 #pragma anki start frag
 
 void doAlphaTest(RF32 alpha)

AnKi/Shaders/Include/Common.h

@@ -767,6 +767,9 @@ constexpr U32 kMaxMipsSinglePassDownsamplerCanProduce = 12u;
 
 constexpr U32 kMaxPrimitivesPerMeshlet = 64;
 constexpr U32 kMaxVerticesPerMeshlet = 64;
+#define ANKI_TASK_SHADER_THREADGROUP_SIZE 64u
+#define ANKI_MESH_SHADER_THREADGROUP_SIZE 64u
+static_assert(ANKI_MESH_SHADER_THREADGROUP_SIZE == max(kMaxPrimitivesPerMeshlet, kMaxVerticesPerMeshlet));
 
 struct DrawIndirectArgs
 {

AnKi/Shaders/Include/GpuSceneTypes.h

@@ -38,6 +38,19 @@ struct GpuSceneRenderableVertex
 };
 static_assert(sizeof(GpuSceneRenderableVertex) == sizeof(UVec4));
 
+/// Input to a single task shader threadgroup. Something similar to GpuSceneRenderableVertex but for mesh shading.
+struct GpuSceneTaskShaderPayload
+{
+	U32 m_firstMeshlet_26bit_meshletCountMinusOne_6bit;
+	U32 m_worldTransformsOffset;
+	U32 m_constantsOffset;
+	U32 m_boneTransformsOrParticleEmitterOffset;
+
+	Vec3 m_positionTranslation;
+	F32 m_positionScale;
+};
+static_assert(ANKI_TASK_SHADER_THREADGROUP_SIZE == 2u << (6u - 1u)); // Need to fit to 6bit
+
 /// Used in visibility testing.
 struct GpuSceneRenderableBoundingVolume
 {
@@ -58,7 +71,7 @@ struct GpuSceneMeshLod
 
 	UVec2 m_padding1;
 	U32 m_firstMeshlet; // In sizeof(Meshlet)
-	U32 m_meshletCount;
+	U32 m_meshletCount; // Can be zero if the mesh doesn't support mesh shading (or mesh shading is off)
 
 	Vec3 m_positionTranslation;
 	F32 m_positionScale;

AnKi/Shaders/Include/MaterialTypes.h

@@ -37,6 +37,9 @@ enum class MaterialBinding : U32
 #define ANKI_UNIFIED_GEOM_FORMAT(fmt, shaderType) kUnifiedGeometry_##fmt,
 #include <AnKi/Shaders/Include/UnifiedGeometryTypes.defs.h>
 
+	kMeshlets, // Pointing to the unified geom buffer
+	kTaskShaderPayloads,
+
 	// For FW shading:
 	kLinearClampSampler,
 	kShadowSampler,

AnKi/Shaders/Intellisense.hlsl

@@ -16,6 +16,7 @@
 #define SV_POSITION
 #define SV_INSTANCEID
 #define numthreads(x, y, z) [nodiscard]
+#define outputtopology(x) [nodiscard]
 #define unroll [nodiscard]
 #define loop [nodiscard]
 #define out
@@ -235,3 +236,12 @@ bool WaveIsFirstLane();
 unsigned WaveActiveCountBits(bool bit);
 
 unsigned WaveGetLaneCount();
+
+// Other
+
+void GroupMemoryBarrierWithGroupSync();
+
+template<typename T>
+void DispatchMesh(U32 groupSizeX, U32 groupSizeY, U32 groupSizeZ, T payload);
+
+void SetMeshOutputCounts(U32 vertexCount, U32 primitiveCount);

AnKi/Shaders/MaterialShadersCommon.hlsl

@@ -22,6 +22,9 @@ ANKI_BINDLESS_SET(MaterialSet::kBindless)
 	[[vk::binding(MaterialBinding::kUnifiedGeometry_##fmt, MaterialSet::kGlobal)]] Buffer<shaderType> g_unifiedGeom_##fmt;
 #include <AnKi/Shaders/Include/UnifiedGeometryTypes.defs.h>
 
+[[vk::binding(MaterialBinding::kMeshlets, MaterialSet::kGlobal)]] StructuredBuffer<Meshlet> g_meshlets;
+[[vk::binding(MaterialBinding::kTaskShaderPayloads, MaterialSet::kGlobal)]] StructuredBuffer<GpuSceneTaskShaderPayload> g_taskShaderPayloads;
+
 // FW shading specific
 #if defined(FORWARD_SHADING)
 #	include <AnKi/Shaders/ClusteredShadingFunctions.hlsl>
@@ -56,3 +59,22 @@ UnpackedMeshVertex loadVertex(GpuSceneMeshLod mlod, U32 svVertexId, Bool bones)
 
 	return v;
 }
+
+UnpackedMeshVertex loadVertex(Meshlet meshlet, U32 vertexIndex, Bool bones, F32 positionScale, Vec3 positionTranslation)
+{
+	UnpackedMeshVertex v;
+	v.m_position = g_unifiedGeom_R16G16B16A16_Unorm[meshlet.m_vertexOffsets[(U32)VertexStreamId::kPosition] + vertexIndex];
+	v.m_position = v.m_position * positionScale + positionTranslation;
+
+	v.m_normal = g_unifiedGeom_R8G8B8A8_Snorm[meshlet.m_vertexOffsets[(U32)VertexStreamId::kNormal] + vertexIndex].xyz;
+	v.m_tangent = g_unifiedGeom_R8G8B8A8_Snorm[meshlet.m_vertexOffsets[(U32)VertexStreamId::kTangent] + vertexIndex];
+	v.m_uv = g_unifiedGeom_R32G32_Sfloat[meshlet.m_vertexOffsets[(U32)VertexStreamId::kUv] + vertexIndex];
+
+	if(bones)
+	{
+		v.m_boneIndices = g_unifiedGeom_R8G8B8A8_Uint[meshlet.m_vertexOffsets[(U32)VertexStreamId::kBoneIds] + vertexIndex];
+		v.m_boneWeights = g_unifiedGeom_R8G8B8A8_Snorm[meshlet.m_vertexOffsets[(U32)VertexStreamId::kBoneWeights] + vertexIndex];
+	}
+
+	return v;
+}

Tools/FormatSource.py

@@ -19,9 +19,9 @@ hlsl_semantics = ["TEXCOORD", "SV_POSITION", "SV_TARGET0", "SV_TARGET1", "SV_TAR
                   "SV_TARGET5", "SV_TARGET6", "SV_TARGET7", "SV_DISPATCHTHREADID", "SV_GROUPINDEX", "SV_GROUPID",
                   "SV_GROUPTHREADID"]
 hlsl_attribs = ["[shader(\"closesthit\")]", "[shader(\"anyhit\")]", "[shader(\"raygeneration\")]", "[shader(\"miss\")]",
-                "[raypayload]"]
+                "[raypayload]", "[outputtopology(\"triangle\")]"]
 hlsl_attribs_fake = ["______shaderclosesthit", "______shaderanyhit", "______shaderraygeneration", "______shadermiss",
-                     "[[raypaylo]]"]
+                     "[[raypaylo]]", "______outputtopology_triangle"]
 
 
 def thread_callback(tid):