From 282392dda2e9349f1972942bbd535886793dd539 Mon Sep 17 00:00:00 2001
From: greg-lunarg <greg@lunarg.com>
Date: Tue, 14 Jul 2020 09:55:24 -0600
Subject: [PATCH] Add support to GPU-AV instrumentation for Task and Mesh
 shaders (#3512)

---
 include/spirv-tools/instrument.hpp | 10 ++++++++++
 source/opt/instrument_pass.cpp     |  5 ++++-
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/include/spirv-tools/instrument.hpp b/include/spirv-tools/instrument.hpp
index ef5136a8c..b4f33554d 100644
--- a/include/spirv-tools/instrument.hpp
+++ b/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;
 
diff --git a/source/opt/instrument_pass.cpp b/source/opt/instrument_pass.cpp
index 4210ad5d7..e6a55a8a4 100644
--- a/source/opt/instrument_pass.cpp
+++ b/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 &&