opt: Mark InterpolateAt* argument as live for DCE (#5824)

The GLSL 450 InterpolateAt* instructions should be treated as a load by dead code elimination. This is part of https://github.com/microsoft/DirectXShaderCompiler/issues/3649.
diff --git a/source/opt/aggressive_dead_code_elim_pass.cpp b/source/opt/aggressive_dead_code_elim_pass.cpp
index 953a7f5..7965422 100644
--- a/source/opt/aggressive_dead_code_elim_pass.cpp
+++ b/source/opt/aggressive_dead_code_elim_pass.cpp
@@ -40,6 +40,9 @@
 constexpr uint32_t kLoadSourceAddrInIdx = 0;
 constexpr uint32_t kDebugDeclareOperandVariableIndex = 5;
 constexpr uint32_t kGlobalVariableVariableIndex = 12;
+constexpr uint32_t kExtInstSetInIdx = 0;
+constexpr uint32_t kExtInstOpInIdx = 1;
+constexpr uint32_t kInterpolantInIdx = 2;
 
 // Sorting functor to present annotation instructions in an easy-to-process
 // order. The functor orders by opcode first and falls back on unique id
@@ -422,6 +425,19 @@
     case spv::Op::OpCopyMemorySized:
       return GetVariableId(
           inst->GetSingleWordInOperand(kCopyMemorySourceAddrInIdx));
+    case spv::Op::OpExtInst: {
+      if (inst->GetSingleWordInOperand(kExtInstSetInIdx) ==
+          context()->get_feature_mgr()->GetExtInstImportId_GLSLstd450()) {
+        auto ext_inst = inst->GetSingleWordInOperand(kExtInstOpInIdx);
+        switch (ext_inst) {
+          case GLSLstd450InterpolateAtCentroid:
+          case GLSLstd450InterpolateAtOffset:
+          case GLSLstd450InterpolateAtSample:
+            return inst->GetSingleWordInOperand(kInterpolantInIdx);
+        }
+      }
+      break;
+    }
     default:
       break;
   }
diff --git a/test/opt/aggressive_dead_code_elim_test.cpp b/test/opt/aggressive_dead_code_elim_test.cpp
index a4353b1..5b88117 100644
--- a/test/opt/aggressive_dead_code_elim_test.cpp
+++ b/test/opt/aggressive_dead_code_elim_test.cpp
@@ -8106,6 +8106,133 @@
   SinglePassRunAndCheck<AggressiveDCEPass>(text, text, true, true);
 }
 
+TEST_F(AggressiveDCETest, MarkCentroidInterpolantLive) {
+  const std::string spirv =
+      R"(OpCapability InterpolationFunction
+OpCapability Shader
+%1 = OpExtInstImport "GLSL.std.450"
+OpMemoryModel Logical GLSL450
+OpEntryPoint Fragment %main "main" %in_var_COLOR %out_var_SV_Target
+OpExecutionMode %main OriginUpperLeft
+OpSource HLSL 680
+OpName %in_var_COLOR "in.var.COLOR"
+OpName %out_var_SV_Target "out.var.SV_Target"
+OpName %main "main"
+OpName %param_var_p1 "param.var.p1"
+OpDecorate %in_var_COLOR Location 0
+OpDecorate %out_var_SV_Target Location 0
+%float = OpTypeFloat 32
+%v4float = OpTypeVector %float 4
+%_ptr_Input_v4float = OpTypePointer Input %v4float
+%_ptr_Output_v4float = OpTypePointer Output %v4float
+%void = OpTypeVoid
+%11 = OpTypeFunction %void
+%_ptr_Function_v4float = OpTypePointer Function %v4float
+%in_var_COLOR = OpVariable %_ptr_Input_v4float Input
+%out_var_SV_Target = OpVariable %_ptr_Output_v4float Output
+%main = OpFunction %void None %11
+%13 = OpLabel
+%14 = OpVariable %_ptr_Function_v4float Function
+%param_var_p1 = OpVariable %_ptr_Function_v4float Function
+%15 = OpLoad %v4float %in_var_COLOR
+OpStore %param_var_p1 %15
+%16 = OpExtInst %v4float %1 InterpolateAtCentroid %param_var_p1
+OpStore %14 %16
+%17 = OpLoad %v4float %14
+OpStore %out_var_SV_Target %17
+OpReturn
+OpFunctionEnd
+)";
+
+  SinglePassRunAndCheck<AggressiveDCEPass>(spirv, spirv, true, false);
+}
+
+TEST_F(AggressiveDCETest, MarkSampleInterpolantLive) {
+  const std::string spirv =
+      R"(OpCapability InterpolationFunction
+OpCapability Shader
+%1 = OpExtInstImport "GLSL.std.450"
+OpMemoryModel Logical GLSL450
+OpEntryPoint Fragment %main "main" %in_var_COLOR %out_var_SV_Target
+OpExecutionMode %main OriginUpperLeft
+OpSource HLSL 680
+OpName %in_var_COLOR "in.var.COLOR"
+OpName %out_var_SV_Target "out.var.SV_Target"
+OpName %main "main"
+OpName %param_var_p1 "param.var.p1"
+OpDecorate %in_var_COLOR Location 0
+OpDecorate %out_var_SV_Target Location 0
+%float = OpTypeFloat 32
+%int = OpTypeInt 32 1
+%v4float = OpTypeVector %float 4
+%_ptr_Input_v4float = OpTypePointer Input %v4float
+%_ptr_Output_v4float = OpTypePointer Output %v4float
+%void = OpTypeVoid
+%12 = OpTypeFunction %void
+%_ptr_Function_v4float = OpTypePointer Function %v4float
+%in_var_COLOR = OpVariable %_ptr_Input_v4float Input
+%out_var_SV_Target = OpVariable %_ptr_Output_v4float Output
+%int_123 = OpConstant %int 123
+%main = OpFunction %void None %12
+%15 = OpLabel
+%16 = OpVariable %_ptr_Function_v4float Function
+%param_var_p1 = OpVariable %_ptr_Function_v4float Function
+%17 = OpLoad %v4float %in_var_COLOR
+OpStore %param_var_p1 %17
+%18 = OpExtInst %v4float %1 InterpolateAtSample %param_var_p1 %int_123
+OpStore %16 %18
+%19 = OpLoad %v4float %16
+OpStore %out_var_SV_Target %19
+OpReturn
+OpFunctionEnd
+)";
+
+  SinglePassRunAndCheck<AggressiveDCEPass>(spirv, spirv, true, false);
+}
+
+TEST_F(AggressiveDCETest, MarkOffsetInterpolantLive) {
+  const std::string spirv =
+      R"(OpCapability InterpolationFunction
+OpCapability Shader
+%1 = OpExtInstImport "GLSL.std.450"
+OpMemoryModel Logical GLSL450
+OpEntryPoint Fragment %main "main" %in_var_COLOR %out_var_SV_Target
+OpExecutionMode %main OriginUpperLeft
+OpSource HLSL 680
+OpName %in_var_COLOR "in.var.COLOR"
+OpName %out_var_SV_Target "out.var.SV_Target"
+OpName %main "main"
+OpName %param_var_p1 "param.var.p1"
+OpDecorate %in_var_COLOR Location 0
+OpDecorate %out_var_SV_Target Location 0
+%float = OpTypeFloat 32
+%int = OpTypeInt 32 1
+%v4float = OpTypeVector %float 4
+%_ptr_Input_v4float = OpTypePointer Input %v4float
+%_ptr_Output_v4float = OpTypePointer Output %v4float
+%void = OpTypeVoid
+%12 = OpTypeFunction %void
+%_ptr_Function_v4float = OpTypePointer Function %v4float
+%in_var_COLOR = OpVariable %_ptr_Input_v4float Input
+%out_var_SV_Target = OpVariable %_ptr_Output_v4float Output
+%int_123 = OpConstant %int 123
+%main = OpFunction %void None %12
+%15 = OpLabel
+%16 = OpVariable %_ptr_Function_v4float Function
+%param_var_p1 = OpVariable %_ptr_Function_v4float Function
+%17 = OpLoad %v4float %in_var_COLOR
+OpStore %param_var_p1 %17
+%18 = OpExtInst %v4float %1 InterpolateAtOffset %param_var_p1 %int_123
+OpStore %16 %18
+%19 = OpLoad %v4float %16
+OpStore %out_var_SV_Target %19
+OpReturn
+OpFunctionEnd
+)";
+
+  SinglePassRunAndCheck<AggressiveDCEPass>(spirv, spirv, true, false);
+}
+
 }  // namespace
 }  // namespace opt
 }  // namespace spvtools