opt: Update copy prop arrays to handle InterpolateAt* instructions (#5827)

cassiebeckley · web-flow · commit 48f9f4a99ed2 · 2024-10-02T09:38:32.000-04:00
Adds support for `InterpolateAtCentroid`, `InterpolateAtSample`, and `InterpolateAtOffset` to the copy propagate arrays pass, as well as propagating pointers with the `Input` storage class. Also handles situations where variables should be propagated in an order different than the order that they are declared in the source.
diff --git a/source/opt/copy_prop_arrays.cpp b/source/opt/copy_prop_arrays.cpp
@@ -28,6 +28,9 @@ constexpr uint32_t kStoreObjectInOperand = 1;
 constexpr uint32_t kCompositeExtractObjectInOperand = 0;
 constexpr uint32_t kTypePointerStorageClassInIdx = 0;
 constexpr uint32_t kTypePointerPointeeInIdx = 1;
+constexpr uint32_t kExtInstSetInIdx = 0;
+constexpr uint32_t kExtInstOpInIdx = 1;
+constexpr uint32_t kInterpolantInIdx = 2;
 
 bool IsDebugDeclareOrValue(Instruction* di) {
   auto dbg_opcode = di->GetCommonDebugOpcode();
@@ -74,28 +77,38 @@ Pass::Status CopyPropagateArrays::Process() {
 
     for (auto var_inst = entry_bb->begin();
          var_inst->opcode() == spv::Op::OpVariable; ++var_inst) {
-      if (!IsPointerToArrayType(var_inst->type_id())) {
-        continue;
-      }
+      worklist_.push(&*var_inst);
+    }
+  }
+
+  while (!worklist_.empty()) {
+    Instruction* var_inst = worklist_.front();
+    worklist_.pop();
 
-      // Find the only store to the entire memory location, if it exists.
-      Instruction* store_inst = FindStoreInstruction(&*var_inst);
+    // Find the only store to the entire memory location, if it exists.
+    Instruction* store_inst = FindStoreInstruction(&*var_inst);
 
-      if (!store_inst) {
+    if (!store_inst) {
+      continue;
+    }
+
+    std::unique_ptr<MemoryObject> source_object =
+        FindSourceObjectIfPossible(&*var_inst, store_inst);
+
+    if (source_object != nullptr) {
+      if (!IsPointerToArrayType(var_inst->type_id()) &&
+          source_object->GetStorageClass() != spv::StorageClass::Input) {
         continue;
       }
 
-      std::unique_ptr<MemoryObject> source_object =
-          FindSourceObjectIfPossible(&*var_inst, store_inst);
+      if (CanUpdateUses(&*var_inst, source_object->GetPointerTypeId(this))) {
+        modified = true;
 
-      if (source_object != nullptr) {
-        if (CanUpdateUses(&*var_inst, source_object->GetPointerTypeId(this))) {
-          modified = true;
-          PropagateObject(&*var_inst, source_object.get(), store_inst);
-        }
+        PropagateObject(&*var_inst, source_object.get(), store_inst);
       }
     }
   }
+
   return (modified ? Status::SuccessWithChange : Status::SuccessWithoutChange);
 }
 
@@ -204,6 +217,8 @@ bool CopyPropagateArrays::HasNoStores(Instruction* ptr_inst) {
       return true;
     } else if (use->opcode() == spv::Op::OpEntryPoint) {
       return true;
+    } else if (IsInterpolationInstruction(use)) {
+      return true;
     }
     // Some other instruction.  Be conservative.
     return false;
@@ -225,6 +240,13 @@ bool CopyPropagateArrays::HasValidReferencesOnly(Instruction* ptr_inst,
           // time to do the multiple traverses can add up.  Consider collecting
           // those loads and doing a single traversal.
           return dominator_analysis->Dominates(store_inst, use);
+        } else if (IsInterpolationInstruction(use)) {
+          // GLSL InterpolateAt* instructions work similarly to loads
+          uint32_t interpolant = use->GetSingleWordInOperand(kInterpolantInIdx);
+          if (interpolant !=
+              store_inst->GetSingleWordInOperand(kStorePointerInOperand))
+            return false;
+          return dominator_analysis->Dominates(store_inst, use);
         } else if (use->opcode() == spv::Op::OpAccessChain) {
           return HasValidReferencesOnly(use, store_inst);
         } else if (use->IsDecoration() || use->opcode() == spv::Op::OpName) {
@@ -489,6 +511,21 @@ bool CopyPropagateArrays::IsPointerToArrayType(uint32_t type_id) {
   return false;
 }
 
+bool CopyPropagateArrays::IsInterpolationInstruction(Instruction* inst) {
+  if (inst->opcode() == spv::Op::OpExtInst &&
+      inst->GetSingleWordInOperand(kExtInstSetInIdx) ==
+          context()->get_feature_mgr()->GetExtInstImportId_GLSLstd450()) {
+    uint32_t ext_inst = inst->GetSingleWordInOperand(kExtInstOpInIdx);
+    switch (ext_inst) {
+      case GLSLstd450InterpolateAtCentroid:
+      case GLSLstd450InterpolateAtOffset:
+      case GLSLstd450InterpolateAtSample:
+        return true;
+    }
+  }
+  return false;
+}
+
 bool CopyPropagateArrays::CanUpdateUses(Instruction* original_ptr_inst,
                                         uint32_t type_id) {
   analysis::TypeManager* type_mgr = context()->get_type_mgr();
@@ -522,6 +559,11 @@ bool CopyPropagateArrays::CanUpdateUses(Instruction* original_ptr_inst,
         }
         return true;
       }
+      case spv::Op::OpExtInst:
+        if (IsInterpolationInstruction(use)) {
+          return true;
+        }
+        return false;
       case spv::Op::OpAccessChain: {
         analysis::Pointer* pointer_type = type->AsPointer();
         const analysis::Type* pointee_type = pointer_type->pointee_type();
@@ -670,6 +712,18 @@ void CopyPropagateArrays::UpdateUses(Instruction* original_ptr_inst,
         } else {
           context()->AnalyzeUses(use);
         }
+
+        AddUsesToWorklist(use);
+      } break;
+      case spv::Op::OpExtInst: {
+        if (IsInterpolationInstruction(use)) {
+          // Replace the actual use.
+          context()->ForgetUses(use);
+          use->SetOperand(index, {new_ptr_inst->result_id()});
+          context()->AnalyzeUses(use);
+        } else {
+          assert(false && "Don't know how to rewrite instruction");
+        }
       } break;
       case spv::Op::OpAccessChain: {
         // Update the actual use.
@@ -799,6 +853,19 @@ uint32_t CopyPropagateArrays::GetMemberTypeId(
   return id;
 }
 
+void CopyPropagateArrays::AddUsesToWorklist(Instruction* inst) {
+  analysis::DefUseManager* def_use_mgr = context()->get_def_use_mgr();
+
+  def_use_mgr->ForEachUse(
+      inst, [this, def_use_mgr](Instruction* use, uint32_t) {
+        if (use->opcode() == spv::Op::OpStore) {
+          Instruction* target_pointer = def_use_mgr->GetDef(
+              use->GetSingleWordInOperand(kStorePointerInOperand));
+          worklist_.push(target_pointer);
+        }
+      });
+}
+
 void CopyPropagateArrays::MemoryObject::PushIndirection(
     const std::vector<AccessChainEntry>& access_chain) {
   access_chain_.insert(access_chain_.end(), access_chain.begin(),
diff --git a/source/opt/copy_prop_arrays.h b/source/opt/copy_prop_arrays.h
@@ -222,6 +222,10 @@ class CopyPropagateArrays : public MemPass {
   // Return true if |type_id| is a pointer type whose pointee type is an array.
   bool IsPointerToArrayType(uint32_t type_id);
 
+  // Return true if |inst| is one of the InterpolateAt* GLSL.std.450 extended
+  // instructions.
+  bool IsInterpolationInstruction(Instruction* inst);
+
   // Returns true if there are not stores using |ptr_inst| or something derived
   // from it.
   bool HasNoStores(Instruction* ptr_inst);
@@ -254,6 +258,14 @@ class CopyPropagateArrays : public MemPass {
   // same way the indexes are used in an |OpCompositeExtract| instruction.
   uint32_t GetMemberTypeId(uint32_t id,
                            const std::vector<uint32_t>& access_chain) const;
+
+  // If the result of inst is stored to a variable, add that variable to the
+  // worklist.
+  void AddUsesToWorklist(Instruction* inst);
+
+  // OpVariable worklist. An instruction is added to this list if we would like
+  // to run copy propagation on it.
+  std::queue<Instruction*> worklist_;
 };
 
 }  // namespace opt
diff --git a/test/opt/copy_prop_array_test.cpp b/test/opt/copy_prop_array_test.cpp
@@ -1977,6 +1977,127 @@ OpFunctionEnd
 
   SinglePassRunAndCheck<CopyPropagateArrays>(text, text, false);
 }
+
+TEST_F(CopyPropArrayPassTest, InterpolateFunctions) {
+  const std::string before = R"(OpCapability InterpolationFunction
+OpCapability Shader
+%1 = OpExtInstImport "GLSL.std.450"
+OpMemoryModel Logical GLSL450
+OpEntryPoint Fragment %main "main" %in_var_COLOR
+OpExecutionMode %main OriginUpperLeft
+OpSource HLSL 680
+OpName %in_var_COLOR "in.var.COLOR"
+OpName %main "main"
+OpName %offset "offset"
+OpDecorate %in_var_COLOR Location 0
+%int = OpTypeInt 32 1
+%int_0 = OpConstant %int 0
+%float = OpTypeFloat 32
+%float_0 = OpConstant %float 0
+%v2float = OpTypeVector %float 2
+%v4float = OpTypeVector %float 4
+%_ptr_Input_v4float = OpTypePointer Input %v4float
+%void = OpTypeVoid
+%19 = OpTypeFunction %void
+%_ptr_Function_v4float = OpTypePointer Function %v4float
+%in_var_COLOR = OpVariable %_ptr_Input_v4float Input
+%main = OpFunction %void None %19
+%20 = OpLabel
+%45 = OpVariable %_ptr_Function_v4float Function
+%25 = OpLoad %v4float %in_var_COLOR
+OpStore %45 %25
+; CHECK: OpExtInst %v4float %1 InterpolateAtCentroid %in_var_COLOR
+%52 = OpExtInst %v4float %1 InterpolateAtCentroid %45
+; CHECK: OpExtInst %v4float %1 InterpolateAtSample %in_var_COLOR %int_0
+%54 = OpExtInst %v4float %1 InterpolateAtSample %45 %int_0
+%offset = OpCompositeConstruct %v2float %float_0 %float_0
+; CHECK: OpExtInst %v4float %1 InterpolateAtOffset %in_var_COLOR %offset
+%56 = OpExtInst %v4float %1 InterpolateAtOffset %45 %offset
+OpReturn
+OpFunctionEnd
+)";
+
+  SetAssembleOptions(SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
+  SetDisassembleOptions(SPV_BINARY_TO_TEXT_OPTION_NO_HEADER |
+                        SPV_BINARY_TO_TEXT_OPTION_FRIENDLY_NAMES);
+  SinglePassRunAndMatch<CopyPropagateArrays>(before, false);
+}
+
+TEST_F(CopyPropArrayPassTest, InterpolateMultiPropagation) {
+  const std::string before = R"(OpCapability InterpolationFunction
+OpCapability Shader
+%1 = OpExtInstImport "GLSL.std.450"
+OpMemoryModel Logical GLSL450
+OpEntryPoint Fragment %main "main" %in_var_COLOR
+OpExecutionMode %main OriginUpperLeft
+OpSource HLSL 680
+OpName %in_var_COLOR "in.var.COLOR"
+OpName %main "main"
+OpName %param_var_color "param.var.color"
+OpDecorate %in_var_COLOR Location 0
+%float = OpTypeFloat 32
+%v4float = OpTypeVector %float 4
+%_ptr_Input_v4float = OpTypePointer Input %v4float
+%void = OpTypeVoid
+%19 = OpTypeFunction %void
+%_ptr_Function_v4float = OpTypePointer Function %v4float
+%in_var_COLOR = OpVariable %_ptr_Input_v4float Input
+%main = OpFunction %void None %19
+%20 = OpLabel
+%45 = OpVariable %_ptr_Function_v4float Function
+%param_var_color = OpVariable %_ptr_Function_v4float Function
+%25 = OpLoad %v4float %in_var_COLOR
+OpStore %param_var_color %25
+; CHECK: OpExtInst %v4float %1 InterpolateAtCentroid %in_var_COLOR
+%52 = OpExtInst %v4float %1 InterpolateAtCentroid %param_var_color
+%49 = OpLoad %v4float %param_var_color
+OpStore %45 %49
+; CHECK: OpExtInst %v4float %1 InterpolateAtCentroid %in_var_COLOR
+%54 = OpExtInst %v4float %1 InterpolateAtCentroid %45
+OpReturn
+OpFunctionEnd
+)";
+
+  SetAssembleOptions(SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
+  SetDisassembleOptions(SPV_BINARY_TO_TEXT_OPTION_NO_HEADER |
+                        SPV_BINARY_TO_TEXT_OPTION_FRIENDLY_NAMES);
+  SinglePassRunAndMatch<CopyPropagateArrays>(before, false);
+}
+
+TEST_F(CopyPropArrayPassTest, PropagateScalar) {
+  const std::string before = R"(OpCapability InterpolationFunction
+OpCapability Shader
+%1 = OpExtInstImport "GLSL.std.450"
+OpMemoryModel Logical GLSL450
+OpEntryPoint Fragment %main "main" %in_var_SV_InstanceID
+OpExecutionMode %main OriginUpperLeft
+OpSource HLSL 680
+OpName %in_var_SV_InstanceID "in.var.SV_InstanceID"
+OpName %main "main"
+OpDecorate %in_var_SV_InstanceID Location 0
+%float = OpTypeFloat 32
+%v4float = OpTypeVector %float 4
+%_ptr_Input_float = OpTypePointer Input %float
+%void = OpTypeVoid
+%19 = OpTypeFunction %void
+%_ptr_Function_float = OpTypePointer Function %float
+%in_var_SV_InstanceID = OpVariable %_ptr_Input_float Input
+%main = OpFunction %void None %19
+%20 = OpLabel
+%45 = OpVariable %_ptr_Function_float Function
+%25 = OpLoad %v4float %in_var_SV_InstanceID
+OpStore %45 %25
+; CHECK: OpExtInst %v4float %1 InterpolateAtCentroid %in_var_SV_InstanceID
+%52 = OpExtInst %v4float %1 InterpolateAtCentroid %45
+OpReturn
+OpFunctionEnd
+)";
+
+  SetAssembleOptions(SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
+  SetDisassembleOptions(SPV_BINARY_TO_TEXT_OPTION_NO_HEADER |
+                        SPV_BINARY_TO_TEXT_OPTION_FRIENDLY_NAMES);
+  SinglePassRunAndMatch<CopyPropagateArrays>(before, false);
+}
 }  // namespace
 }  // namespace opt
 }  // namespace spvtools
