Add support to GPU-AV instrumentation for Task and Mesh shaders (#3512)

include/spirv-tools/instrument.hpp

@@ -110,6 +110,16 @@ static const int kInstRayTracingOutLaunchIdX = kInstCommonOutCnt;
 static const int kInstRayTracingOutLaunchIdY = kInstCommonOutCnt + 1;
 static const int kInstRayTracingOutLaunchIdZ = kInstCommonOutCnt + 2;
 
+// Mesh Shader Output Record Offsets
+static const int kInstMeshOutGlobalInvocationIdX = kInstCommonOutCnt;
+static const int kInstMeshOutGlobalInvocationIdY = kInstCommonOutCnt + 1;
+static const int kInstMeshOutGlobalInvocationIdZ = kInstCommonOutCnt + 2;
+
+// Task Shader Output Record Offsets
+static const int kInstTaskOutGlobalInvocationIdX = kInstCommonOutCnt;
+static const int kInstTaskOutGlobalInvocationIdY = kInstCommonOutCnt + 1;
+static const int kInstTaskOutGlobalInvocationIdZ = kInstCommonOutCnt + 2;
+
 // Size of Common and Stage-specific Members
 static const int kInstStageOutCnt = kInstCommonOutCnt + 3;
 

source/opt/instrument_pass.cpp

@@ -174,7 +174,9 @@ void InstrumentPass::GenStageStreamWriteCode(uint32_t stage_idx,
           context()->GetBuiltinInputVarId(SpvBuiltInInstanceIndex),
           kInstVertOutInstanceIndex, base_offset_id, builder);
     } break;
-    case SpvExecutionModelGLCompute: {
+    case SpvExecutionModelGLCompute:
+    case SpvExecutionModelTaskNV:
+    case SpvExecutionModelMeshNV: {
       // Load and store GlobalInvocationId.
       uint32_t load_id = GenVarLoad(
           context()->GetBuiltinInputVarId(SpvBuiltInGlobalInvocationId),
@@ -914,6 +916,7 @@ bool InstrumentPass::InstProcessEntryPointCallTree(InstProcessFunction& pfn) {
       stage != SpvExecutionModelGLCompute &&
       stage != SpvExecutionModelTessellationControl &&
       stage != SpvExecutionModelTessellationEvaluation &&
+      stage != SpvExecutionModelTaskNV && stage != SpvExecutionModelMeshNV &&
       stage != SpvExecutionModelRayGenerationNV &&
       stage != SpvExecutionModelIntersectionNV &&
       stage != SpvExecutionModelAnyHitNV &&