Fix crash when optimizing shaders with DebugPrintf (#4280)

Fixes #4219

4 files changed, 97 insertions, 13 deletions

diff --git a/source/opt/eliminate_dead_functions_util.cpp b/source/opt/eliminate_dead_functions_util.cpp
index 6b5234bb..1379120f 100644
--- a/source/opt/eliminate_dead_functions_util.cpp
+++ b/source/opt/eliminate_dead_functions_util.cpp
@@ -23,9 +23,11 @@ Module::iterator EliminateFunction(IRContext* context,
                                    Module::iterator* func_iter) {
   bool first_func = *func_iter == context->module()->begin();
   bool seen_func_end = false;
+  std::unordered_set<Instruction*> to_kill;
   (*func_iter)
       ->ForEachInst(
-          [context, first_func, func_iter, &seen_func_end](Instruction* inst) {
+          [context, first_func, func_iter, &seen_func_end,
+           &to_kill](Instruction* inst) {
             if (inst->opcode() == SpvOpFunctionEnd) {
               seen_func_end = true;
             }
@@ -33,6 +35,7 @@ Module::iterator EliminateFunction(IRContext* context,
             // global values if this is the first function.
             if (seen_func_end && inst->opcode() == SpvOpExtInst) {
               assert(inst->IsNonSemanticInstruction());
+              if (to_kill.find(inst) != to_kill.end()) return;
               std::unique_ptr<Instruction> clone(inst->Clone(context));
               context->ForgetUses(inst);
               context->AnalyzeDefUse(clone.get());
@@ -44,12 +47,17 @@ Module::iterator EliminateFunction(IRContext* context,
                 prev_func_iter->AddNonSemanticInstruction(std::move(clone));
               }
               inst->ToNop();
-            } else {
-              context->KillNonSemanticInfo(inst);
+            } else if (to_kill.find(inst) == to_kill.end()) {
+              context->CollectNonSemanticTree(inst, &to_kill);
               context->KillInst(inst);
             }
           },
           true, true);
+
+  for (auto* dead : to_kill) {
+    context->KillInst(dead);
+  }
+
   return func_iter->Erase();
 }
 
diff --git a/source/opt/ir_context.cpp b/source/opt/ir_context.cpp
index 82107b5c..03afe6e8 100644
--- a/source/opt/ir_context.cpp
+++ b/source/opt/ir_context.cpp
@@ -214,10 +214,10 @@ Instruction* IRContext::KillInst(Instruction* inst) {
   return next_instruction;
 }
 
-void IRContext::KillNonSemanticInfo(Instruction* inst) {
+void IRContext::CollectNonSemanticTree(
+    Instruction* inst, std::unordered_set<Instruction*>* to_kill) {
   if (!inst->HasResultId()) return;
   std::vector<Instruction*> work_list;
-  std::vector<Instruction*> to_kill;
   std::unordered_set<Instruction*> seen;
   work_list.push_back(inst);
 
@@ -225,17 +225,13 @@ void IRContext::KillNonSemanticInfo(Instruction* inst) {
     auto* i = work_list.back();
     work_list.pop_back();
     get_def_use_mgr()->ForEachUser(
-        i, [&work_list, &to_kill, &seen](Instruction* user) {
+        i, [&work_list, to_kill, &seen](Instruction* user) {
           if (user->IsNonSemanticInstruction() && seen.insert(user).second) {
             work_list.push_back(user);
-            to_kill.push_back(user);
+            to_kill->insert(user);
           }
         });
   }
-
-  for (auto* dead : to_kill) {
-    KillInst(dead);
-  }
 }
 
 bool IRContext::KillDef(uint32_t id) {
diff --git a/source/opt/ir_context.h b/source/opt/ir_context.h
index 5aa25acd..aab35162 100644
--- a/source/opt/ir_context.h
+++ b/source/opt/ir_context.h
@@ -403,8 +403,10 @@ class IRContext {
   // instruction exists.
   Instruction* KillInst(Instruction* inst);
 
-  // Removes the non-semantic instruction tree that uses |inst|'s result id.
-  void KillNonSemanticInfo(Instruction* inst);
+  // Collects the non-semantic instruction tree that uses |inst|'s result id
+  // to be killed later.
+  void CollectNonSemanticTree(Instruction* inst,
+                              std::unordered_set<Instruction*>* to_kill);
 
   // Returns true if all of the given analyses are valid.
   bool AreAnalysesValid(Analysis set) { return (set & valid_analyses_) == set; }
diff --git a/test/opt/eliminate_dead_functions_test.cpp b/test/opt/eliminate_dead_functions_test.cpp
index 96ecdc60..96deb2a6 100644
--- a/test/opt/eliminate_dead_functions_test.cpp
+++ b/test/opt/eliminate_dead_functions_test.cpp
@@ -439,6 +439,84 @@ OpFunctionEnd
   SinglePassRunAndMatch<EliminateDeadFunctionsPass>(text, true);
 }
 
+TEST_F(EliminateDeadFunctionsBasicTest, NonSemanticInfoRemoveDebugPrintf) {
+  const std::string text = R"(
+; CHECK-NOT: %foo_ = OpFunction %void None % 3
+; CHECK-NOT: % 7 = OpLabel
+; CHECK-NOT: %c = OpVariable %_ptr_Function_v4float Function
+; CHECK-NOT: % 22 = OpAccessChain %_ptr_UniformConstant_13 %samplers %int_0
+; CHECK-NOT: % 23 = OpLoad % 13 % 22
+; CHECK-NOT: % 27 = OpImageSampleExplicitLod %v4float % 23 % 26 Lod %float_0
+; CHECK-NOT: OpStore %c % 27
+; CHECK-NOT: % 31 = OpAccessChain %_ptr_Function_float %c %uint_0
+; CHECK-NOT: % 32 = OpLoad %float %31
+; CHECK-NOT: % 34 = OpExtInst %void %33 1 % 28 % 32
+OpCapability RayTracingKHR
+OpExtension "SPV_KHR_non_semantic_info"
+OpExtension "SPV_KHR_ray_tracing"
+%1 = OpExtInstImport "GLSL.std.450"
+%33 = OpExtInstImport "NonSemantic.DebugPrintf"
+OpMemoryModel Logical GLSL450
+OpEntryPoint ClosestHitNV %main "main" %samplers
+%28 = OpString "%f"
+OpSource GLSL 460
+OpSourceExtension "GL_EXT_debug_printf"
+OpName %main "main"
+OpName %foo_ "foo("
+OpName %c "c"
+OpName %samplers "samplers"
+OpDecorate %samplers DescriptorSet 0
+OpDecorate %samplers Binding 0
+%void = OpTypeVoid
+%3 = OpTypeFunction %void
+%float = OpTypeFloat 32
+%v4float = OpTypeVector %float 4
+%_ptr_Function_v4float = OpTypePointer Function %v4float
+%12 = OpTypeImage %float 3D 0 0 0 1 Unknown
+%13 = OpTypeSampledImage %12
+%uint = OpTypeInt 32 0
+%uint_1 = OpConstant %uint 1
+%_arr_13_uint_1 = OpTypeArray %13 %uint_1
+%_ptr_UniformConstant__arr_13_uint_1 = OpTypePointer UniformConstant %_arr_13_uint_1
+%samplers = OpVariable %_ptr_UniformConstant__arr_13_uint_1 UniformConstant
+%int = OpTypeInt 32 1
+%int_0 = OpConstant %int 0
+%_ptr_UniformConstant_13 = OpTypePointer UniformConstant %13
+%v3float = OpTypeVector %float 3
+%float_0 = OpConstant %float 0
+%26 = OpConstantComposite %v3float %float_0 %float_0 %float_0
+%uint_0 = OpConstant %uint 0
+%_ptr_Function_float = OpTypePointer Function %float
+%main = OpFunction %void None %3
+%5 = OpLabel
+%36 = OpVariable %_ptr_Function_v4float Function
+%38 = OpAccessChain %_ptr_UniformConstant_13 %samplers %int_0
+%39 = OpLoad %13 %38
+%40 = OpImageSampleExplicitLod %v4float %39 %26 Lod %float_0
+OpStore %36 %40
+%41 = OpAccessChain %_ptr_Function_float %36 %uint_0
+%42 = OpLoad %float %41
+%43 = OpExtInst %void %33 1 %28 %42
+OpReturn
+OpFunctionEnd
+%foo_ = OpFunction %void None %3
+%7 = OpLabel
+%c = OpVariable %_ptr_Function_v4float Function
+%22 = OpAccessChain %_ptr_UniformConstant_13 %samplers %int_0
+%23 = OpLoad %13 %22
+%27 = OpImageSampleExplicitLod %v4float %23 %26 Lod %float_0
+OpStore %c %27
+%31 = OpAccessChain %_ptr_Function_float %c %uint_0
+%32 = OpLoad %float %31
+%34 = OpExtInst %void %33 1 %28 %32
+OpReturn
+OpFunctionEnd
+)";
+
+  SetTargetEnv(SPV_ENV_VULKAN_1_2);
+  SinglePassRunAndMatch<EliminateDeadFunctionsPass>(text, true);
+}
+
 }  // namespace
 }  // namespace opt
 }  // namespace spvtools