Implement mesh viewer support for task/mesh shaders

renderdoc/data/glsl/glsl_globals.h

@@ -172,10 +172,16 @@ vec3 CalcCubeCoord(vec2 uv, int face)
 
 #endif
 
+// first few match SolidShade enum
 #define MESHDISPLAY_SOLID 0x1
 #define MESHDISPLAY_FACELIT 0x2
 #define MESHDISPLAY_SECONDARY 0x3
-#define MESHDISPLAY_SECONDARY_ALPHA 0x4
+#define MESHDISPLAY_MESHLET 0x4
+
+// extra values below
+#define MESHDISPLAY_SECONDARY_ALPHA 0x5
+
+#define MAX_NUM_MESHLETS (512 * 1024)
 
 #define TEXDISPLAY_TYPEMASK 0xF
 #define TEXDISPLAY_UINT_TEX 0x10

renderdoc/data/glsl/glsl_ubos.h

@@ -41,11 +41,24 @@ BINDING(0) uniform MeshUBOData
   uint rawoutput;
   uint flipY;
   vec2 padding;
+  uvec4 meshletColours[12];
 }
 INST_NAME(Mesh);
 
 #endif    // defined(MESH_UBO) || defined(__cplusplus)
 
+#if defined(MESH_UBO) && defined(VULKAN)
+
+layout(binding = 1, std140) readonly buffer meshlet_data
+{
+  uint meshletCount;
+  uint meshletOffset;
+  uvec4 data[];
+}
+meshlet;
+
+#endif    // defined(MESH_UBO)
+
 #if defined(CHECKER_UBO) || defined(__cplusplus)
 
 BINDING(0) uniform CheckerboardUBOData

renderdoc/data/glsl/mesh.frag

@@ -46,7 +46,7 @@ void main(void)
 {
   int type = Mesh.displayFormat;
 
-  if(type == MESHDISPLAY_SECONDARY)
+  if(type == MESHDISPLAY_SECONDARY || type == MESHDISPLAY_MESHLET)
   {
     color_out = vec4(SECONDARY_NAME.xyz, 1);
   }

renderdoc/data/glsl/mesh.vert

@@ -42,6 +42,69 @@ IO_LOCATION(1) in SECONDARY_TYPE vsin_secondary;
 IO_LOCATION(0) out vec4 vsout_secondary;
 IO_LOCATION(1) out vec4 vsout_norm;
 
+#ifdef VULKAN
+
+uint getMeshletCountAt(uint m)
+{
+  uint vecIdx = m / 4;
+
+  if((m % 4) == 0)
+    return meshlet.data[vecIdx].x;
+  else if((m % 4) == 1)
+    return meshlet.data[vecIdx].y;
+  else if((m % 4) == 2)
+    return meshlet.data[vecIdx].z;
+  else if((m % 4) == 3)
+    return meshlet.data[vecIdx].w;
+}
+
+vec4 getMeshletColor()
+{
+  uint searchIdx = VERTEX_ID;
+
+  // array of prefix summed counts accessible via getMeshletCountAt [x, x+y, x+y+z, ...] we do a
+  // binary search to find which meshlet this index corresponds to
+
+  uint first = 0, last = meshlet.meshletCount - 1;
+  uint count = last - first;
+
+  while(count > 0)
+  {
+    uint halfrange = count / 2;
+    uint mid = first + halfrange;
+
+    if(searchIdx < getMeshletCountAt(mid))
+    {
+      count = halfrange;
+    }
+    else
+    {
+      first = mid + 1;
+      count -= halfrange + 1;
+    }
+  }
+
+  uint meshletIndex = first;
+
+  if(VERTEX_ID < getMeshletCountAt(meshletIndex))
+  {
+    meshletIndex += meshlet.meshletOffset;
+    meshletIndex %= 48;
+    uvec4 meshletColor = Mesh.meshletColours[meshletIndex / 4];
+    if((meshletIndex % 4) == 0)
+      return unpackUnorm4x8(meshletColor.x);
+    else if((meshletIndex % 4) == 1)
+      return unpackUnorm4x8(meshletColor.y);
+    else if((meshletIndex % 4) == 2)
+      return unpackUnorm4x8(meshletColor.z);
+    else if((meshletIndex % 4) == 3)
+      return unpackUnorm4x8(meshletColor.w);
+  }
+
+  return vec4(0, 0, 0, 1);
+}
+#endif
+
 void main(void)
 {
   vec2 psprite[4] =
@@ -65,6 +128,9 @@ void main(void)
   vsout_norm = vec4(0, 0, 1, 1);
 
 #ifdef VULKAN
+  if(Mesh.displayFormat == MESHDISPLAY_MESHLET)
+    vsout_secondary = getMeshletColor();
+
   // GL->VK conventions
   gl_Position.y = -gl_Position.y;
   if(Mesh.flipY == 1)

renderdoc/data/hlsl/hlsl_cbuffers.h

@@ -106,8 +106,15 @@ cbuffer MeshVertexCBuffer REG(b0)
   row_major float4x4 ModelViewProj;
 
   float2 SpriteSize;
-
   uint homogenousInput;
+  uint vertMeshDisplayFormat;
+
+  uint meshletOffset;
+  uint meshletCount;
+  uint padding1;
+  uint padding2;
+
+  uint4 meshletColours[12];
 };
 
 cbuffer MeshGeometryCBuffer REG(b0)
@@ -243,10 +250,16 @@ cbuffer DebugSampleOperation REG(b0)
 #define RESTYPE_DEPTH_STENCIL_MS 0x7
 #define RESTYPE_TEX2D_MS 0x9
 
+// first few match SolidShade enum
 #define MESHDISPLAY_SOLID 0x1
 #define MESHDISPLAY_FACELIT 0x2
 #define MESHDISPLAY_SECONDARY 0x3
-#define MESHDISPLAY_SECONDARY_ALPHA 0x4
+#define MESHDISPLAY_MESHLET 0x4
+
+// extra values below
+#define MESHDISPLAY_SECONDARY_ALPHA 0x5
+
+#define MAX_NUM_MESHLETS (512 * 1024)
 
 #define TEXDISPLAY_TYPEMASK 0xF
 #define TEXDISPLAY_NANS 0x0100

renderdoc/data/hlsl/mesh.hlsl

@@ -37,6 +37,79 @@ struct meshA2V
   float4 secondary : sec;
 };
 
+StructuredBuffer<uint4> meshletSizesBuf : register(t0);
+
+float4 unpackUnorm4x8(uint value)
+{
+  uint4 shifted = uint4(value & 0xff, (value >> 8) & 0xff, (value >> 16) & 0xff, value >> 24);
+  return float4(shifted) / 255.0;
+}
+
+uint getMeshletCountAt(uint m)
+{
+  uint vecIdx = m / 4;
+
+  uint count = 0;
+  if((m % 4) == 0)
+    count = meshletSizesBuf[vecIdx].x;
+  else if((m % 4) == 1)
+    count = meshletSizesBuf[vecIdx].y;
+  else if((m % 4) == 2)
+    count = meshletSizesBuf[vecIdx].z;
+  else if((m % 4) == 3)
+    count = meshletSizesBuf[vecIdx].w;
+
+  return count;
+}
+
+float4 getMeshletColor(uint vid)
+{
+  uint searchIdx = vid;
+
+  // array of prefix summed counts accessible via getMeshletCountAt [x, x+y, x+y+z, ...] we do a
+  // binary search to find which meshlet this index corresponds to
+
+  uint first = 0, last = meshletCount - 1;
+  uint count = last - first;
+
+  while(count > 0)
+  {
+    uint halfrange = count / 2;
+    uint mid = first + halfrange;
+
+    if(searchIdx < getMeshletCountAt(mid))
+    {
+      count = halfrange;
+    }
+    else
+    {
+      first = mid + 1;
+      count -= halfrange + 1;
+    }
+  }
+
+  uint meshletIndex = first;
+
+  float4 col = float4(0, 0, 0, 1);
+
+  if(vid < getMeshletCountAt(meshletIndex))
+  {
+    meshletIndex += meshletOffset;
+    meshletIndex %= 48;
+    uint4 meshletColor = meshletColours[meshletIndex / 4];
+    if((meshletIndex % 4) == 0)
+      col = unpackUnorm4x8(meshletColor.x);
+    else if((meshletIndex % 4) == 1)
+      col = unpackUnorm4x8(meshletColor.y);
+    else if((meshletIndex % 4) == 2)
+      col = unpackUnorm4x8(meshletColor.z);
+    else if((meshletIndex % 4) == 3)
+      col = unpackUnorm4x8(meshletColor.w);
+  }
+
+  return col;
+}
+
 meshV2F RENDERDOC_MeshVS(meshA2V IN, uint vid : SV_VertexID)
 {
   meshV2F OUT = (meshV2F)0;
@@ -56,6 +129,9 @@ meshV2F RENDERDOC_MeshVS(meshA2V IN, uint vid : SV_VertexID)
   OUT.norm = float3(0, 0, 1);
   OUT.secondary = IN.secondary;
 
+  if(vertMeshDisplayFormat == MESHDISPLAY_MESHLET)
+    OUT.secondary = getMeshletColor(vid);
+
   return OUT;
 }
 
@@ -128,7 +204,7 @@ float4 RENDERDOC_MeshPS(meshV2F IN)
 {
   uint type = MeshDisplayFormat;
 
-  if(type == MESHDISPLAY_SECONDARY)
+  if(type == MESHDISPLAY_SECONDARY || type == MESHDISPLAY_MESHLET)
     return float4(IN.secondary.xyz, 1);
   else if(type == MESHDISPLAY_SECONDARY_ALPHA)
     return float4(IN.secondary.www, 1);