Allow ray tracing shaders in inst bindle check pass. (#2733)

Adds the ray tracing stages (ray gen, intersection, any hit, closest hit,
miss, and callable) to the allowed stages in pass instrumentation and add
debug records for these stages to output the global launch id.

More information for ray tracing shaders:
- https://github.com/KhronosGroup/GLSL/blob/master/extensions/nv/GLSL_NV_ray_tracing.txt

include/spirv-tools/instrument.hpp

@@ -117,6 +117,11 @@ static const int kInstGeomOutPrimitiveId = kInstCommonOutCnt;
 static const int kInstGeomOutInvocationId = kInstCommonOutCnt + 1;
 static const int kInstGeomOutUnused = kInstCommonOutCnt + 2;
 
+// Ray Tracing Shader Output Record Offsets
+static const int kInstRayTracingOutLaunchIdX = kInstCommonOutCnt;
+static const int kInstRayTracingOutLaunchIdY = kInstCommonOutCnt + 1;
+static const int kInstRayTracingOutLaunchIdZ = kInstCommonOutCnt + 2;
+
 // Size of Common and Stage-specific Members
 static const int kInstStageOutCnt = kInstCommonOutCnt + 2;
 static const int kInst2StageOutCnt = kInstCommonOutCnt + 3;

source/opt/instrument_pass.cpp

@@ -264,6 +264,28 @@ void InstrumentPass::GenStageStreamWriteCode(uint32_t stage_idx,
         GenFragCoordEltDebugOutputCode(
             base_offset_id, uint_frag_coord_inst->result_id(), u, builder);
     } break;
+    case SpvExecutionModelRayGenerationNV:
+    case SpvExecutionModelIntersectionNV:
+    case SpvExecutionModelAnyHitNV:
+    case SpvExecutionModelClosestHitNV:
+    case SpvExecutionModelMissNV:
+    case SpvExecutionModelCallableNV: {
+      // Load and store LaunchIdNV.
+      uint32_t launch_id = GenVarLoad(
+          context()->GetBuiltinInputVarId(SpvBuiltInLaunchIdNV), builder);
+      Instruction* x_launch_inst = builder->AddIdLiteralOp(
+          GetUintId(), SpvOpCompositeExtract, launch_id, 0);
+      Instruction* y_launch_inst = builder->AddIdLiteralOp(
+          GetUintId(), SpvOpCompositeExtract, launch_id, 1);
+      Instruction* z_launch_inst = builder->AddIdLiteralOp(
+          GetUintId(), SpvOpCompositeExtract, launch_id, 2);
+      GenDebugOutputFieldCode(base_offset_id, kInstRayTracingOutLaunchIdX,
+                              x_launch_inst->result_id(), builder);
+      GenDebugOutputFieldCode(base_offset_id, kInstRayTracingOutLaunchIdY,
+                              y_launch_inst->result_id(), builder);
+      GenDebugOutputFieldCode(base_offset_id, kInstRayTracingOutLaunchIdZ,
+                              z_launch_inst->result_id(), builder);
+    } break;
     default: { assert(false && "unsupported stage"); } break;
   }
 }
@@ -843,7 +865,12 @@ bool InstrumentPass::InstProcessEntryPointCallTree(InstProcessFunction& pfn) {
       stage != SpvExecutionModelGeometry &&
       stage != SpvExecutionModelGLCompute &&
       stage != SpvExecutionModelTessellationControl &&
-      stage != SpvExecutionModelTessellationEvaluation)
+      stage != SpvExecutionModelTessellationEvaluation &&
+      stage != SpvExecutionModelRayGenerationNV &&
+      stage != SpvExecutionModelIntersectionNV &&
+      stage != SpvExecutionModelAnyHitNV &&
+      stage != SpvExecutionModelClosestHitNV &&
+      stage != SpvExecutionModelMissNV && stage != SpvExecutionModelCallableNV)
     return false;
   // Add together the roots of all entry points
   std::queue<uint32_t> roots;

source/opt/ir_context.cpp

@@ -683,7 +683,8 @@ uint32_t IRContext::GetBuiltinInputVarId(uint32_t builtin) {
         reg_type = type_mgr->GetRegisteredType(&uint_ty);
         break;
       }
-      case SpvBuiltInGlobalInvocationId: {
+      case SpvBuiltInGlobalInvocationId:
+      case SpvBuiltInLaunchIdNV: {
         analysis::Integer uint_ty(32, false);
         analysis::Type* reg_uint_ty = type_mgr->GetRegisteredType(&uint_ty);
         analysis::Vector v3uint_ty(reg_uint_ty, 3);