spirv-opt: create OpDecorate for OpMemberDecorate in desc-sroa (#4617)

The scalar replacement of a resource array/struct variable must create
OpDecorate for elements if OpMemberDecorate instructions decorate
the elements.

source/opt/desc_sroa.cpp

@@ -19,6 +19,14 @@
 
 namespace spvtools {
 namespace opt {
+namespace {
+
+bool IsDecorationBinding(Instruction* inst) {
+  if (inst->opcode() != SpvOpDecorate) return false;
+  return inst->GetSingleWordInOperand(1u) == SpvDecorationBinding;
+}
+
+}  // namespace
 
 Pass::Status DescriptorScalarReplacement::Process() {
   bool modified = false;
@@ -157,6 +165,74 @@ uint32_t DescriptorScalarReplacement::GetReplacementVariable(Instruction* var,
   return replacement_vars->second[idx];
 }
 
+void DescriptorScalarReplacement::CopyDecorationsForNewVariable(
+    Instruction* old_var, uint32_t index, uint32_t new_var_id,
+    uint32_t new_var_ptr_type_id, const bool is_old_var_array,
+    const bool is_old_var_struct, Instruction* old_var_type) {
+  // Handle OpDecorate instructions.
+  for (auto old_decoration :
+       get_decoration_mgr()->GetDecorationsFor(old_var->result_id(), true)) {
+    uint32_t new_binding = 0;
+    if (IsDecorationBinding(old_decoration)) {
+      new_binding = GetNewBindingForElement(
+          old_decoration->GetSingleWordInOperand(2), index, new_var_ptr_type_id,
+          is_old_var_array, is_old_var_struct, old_var_type);
+    }
+    CreateNewDecorationForNewVariable(old_decoration, new_var_id, new_binding);
+  }
+
+  // Handle OpMemberDecorate instructions.
+  for (auto old_decoration : get_decoration_mgr()->GetDecorationsFor(
+           old_var_type->result_id(), true)) {
+    assert(old_decoration->opcode() == SpvOpMemberDecorate);
+    if (old_decoration->GetSingleWordInOperand(1u) != index) continue;
+    CreateNewDecorationForMemberDecorate(old_decoration, new_var_id);
+  }
+}
+
+uint32_t DescriptorScalarReplacement::GetNewBindingForElement(
+    uint32_t old_binding, uint32_t index, uint32_t new_var_ptr_type_id,
+    const bool is_old_var_array, const bool is_old_var_struct,
+    Instruction* old_var_type) {
+  if (is_old_var_array) {
+    return old_binding + index * GetNumBindingsUsedByType(new_var_ptr_type_id);
+  }
+  if (is_old_var_struct) {
+    // The binding offset that should be added is the sum of binding
+    // numbers used by previous members of the current struct.
+    uint32_t new_binding = old_binding;
+    for (uint32_t i = 0; i < index; ++i) {
+      new_binding +=
+          GetNumBindingsUsedByType(old_var_type->GetSingleWordInOperand(i));
+    }
+    return new_binding;
+  }
+  return old_binding;
+}
+
+void DescriptorScalarReplacement::CreateNewDecorationForNewVariable(
+    Instruction* old_decoration, uint32_t new_var_id, uint32_t new_binding) {
+  assert(old_decoration->opcode() == SpvOpDecorate);
+  std::unique_ptr<Instruction> new_decoration(old_decoration->Clone(context()));
+  new_decoration->SetInOperand(0, {new_var_id});
+
+  if (IsDecorationBinding(new_decoration.get())) {
+    new_decoration->SetInOperand(2, {new_binding});
+  }
+  context()->AddAnnotationInst(std::move(new_decoration));
+}
+
+void DescriptorScalarReplacement::CreateNewDecorationForMemberDecorate(
+    Instruction* old_member_decoration, uint32_t new_var_id) {
+  std::vector<Operand> operands(
+      {{spv_operand_type_t::SPV_OPERAND_TYPE_ID, {new_var_id}}});
+  auto new_decorate_operand_begin = old_member_decoration->begin() + 2u;
+  auto new_decorate_operand_end = old_member_decoration->end();
+  operands.insert(operands.end(), new_decorate_operand_begin,
+                  new_decorate_operand_end);
+  get_decoration_mgr()->AddDecoration(SpvOpDecorate, std::move(operands));
+}
+
 uint32_t DescriptorScalarReplacement::CreateReplacementVariable(
     Instruction* var, uint32_t idx) {
   // The storage class for the new variable is the same as the original.
@@ -192,33 +268,8 @@ uint32_t DescriptorScalarReplacement::CreateReplacementVariable(
                            {static_cast<uint32_t>(storage_class)}}}));
   context()->AddGlobalValue(std::move(variable));
 
-  // Copy all of the decorations to the new variable.  The only difference is
+  CopyDecorationsForNewVariable(var, idx, id, ptr_element_type_id, is_array,
-  // the Binding decoration needs to be adjusted.
+                                is_struct, pointee_type_inst);
-  for (auto old_decoration :
-       get_decoration_mgr()->GetDecorationsFor(var->result_id(), true)) {
-    assert(old_decoration->opcode() == SpvOpDecorate);
-    std::unique_ptr<Instruction> new_decoration(
-        old_decoration->Clone(context()));
-    new_decoration->SetInOperand(0, {id});
-
-    uint32_t decoration = new_decoration->GetSingleWordInOperand(1u);
-    if (decoration == SpvDecorationBinding) {
-      uint32_t new_binding = new_decoration->GetSingleWordInOperand(2);
-      if (is_array) {
-        new_binding += idx * GetNumBindingsUsedByType(ptr_element_type_id);
-      }
-      if (is_struct) {
-        // The binding offset that should be added is the sum of binding numbers
-        // used by previous members of the current struct.
-        for (uint32_t i = 0; i < idx; ++i) {
-          new_binding += GetNumBindingsUsedByType(
-              pointee_type_inst->GetSingleWordInOperand(i));
-        }
-      }
-      new_decoration->SetInOperand(2, {new_binding});
-    }
-    context()->AddAnnotationInst(std::move(new_decoration));
-  }
 
   // Create a new OpName for the replacement variable.
   std::vector<std::unique_ptr<Instruction>> names_to_add;

source/opt/desc_sroa.h

@@ -89,6 +89,46 @@ class DescriptorScalarReplacement : public Pass {
   // bindings used by its members.
   uint32_t GetNumBindingsUsedByType(uint32_t type_id);
 
+  // Copy all of the decorations of variable |old_var| and make them as
+  // decorations for the new variable whose id is |new_var_id|. The new variable
+  // is supposed to replace |index|th element of |old_var|.
+  // |new_var_ptr_type_id| is the id of the pointer to the type of the new
+  // variable. |is_old_var_array| is true if |old_var| has an array type.
+  // |is_old_var_struct| is true if |old_var| has a structure type.
+  // |old_var_type| is the pointee type of |old_var|.
+  void CopyDecorationsForNewVariable(Instruction* old_var, uint32_t index,
+                                     uint32_t new_var_id,
+                                     uint32_t new_var_ptr_type_id,
+                                     const bool is_old_var_array,
+                                     const bool is_old_var_struct,
+                                     Instruction* old_var_type);
+
+  // Get the new binding number for a new variable that will be replaced with an
+  // |index|th element of an old variable. The old variable has |old_binding|
+  // as its binding number. |ptr_elem_type_id| the id of the pointer to the
+  // element type. |is_old_var_array| is true if the old variable has an array
+  // type. |is_old_var_struct| is true if the old variable has a structure type.
+  // |old_var_type| is the pointee type of the old variable.
+  uint32_t GetNewBindingForElement(uint32_t old_binding, uint32_t index,
+                                   uint32_t ptr_elem_type_id,
+                                   const bool is_old_var_array,
+                                   const bool is_old_var_struct,
+                                   Instruction* old_var_type);
+
+  // Create a new OpDecorate instruction by cloning |old_decoration|. The new
+  // OpDecorate instruction will be used for a variable whose id is
+  // |new_var_ptr_type_id|. If |old_decoration| is a decoration for a binding,
+  // the new OpDecorate instruction will have |new_binding| as its binding.
+  void CreateNewDecorationForNewVariable(Instruction* old_decoration,
+                                         uint32_t new_var_id,
+                                         uint32_t new_binding);
+
+  // Create a new OpDecorate instruction whose operand is the same as an
+  // OpMemberDecorate instruction |old_member_decoration| except Target operand.
+  // The Target operand of the new OpDecorate instruction will be |new_var_id|.
+  void CreateNewDecorationForMemberDecorate(Instruction* old_decoration,
+                                            uint32_t new_var_id);
+
   // A map from an OpVariable instruction to the set of variables that will be
   // used to replace it. The entry |replacement_variables_[var][i]| is the id of
   // a variable that will be used in the place of the the ith element of the

test/opt/desc_sroa_test.cpp

@@ -770,6 +770,69 @@ TEST_F(DescriptorScalarReplacementTest, BindingForResourceArrayOfStructs) {
   SinglePassRunAndMatch<DescriptorScalarReplacement>(shader, true);
 }
 
+TEST_F(DescriptorScalarReplacementTest, MemberDecorationForResourceStruct) {
+  // Check that an OpMemberDecorate instruction is correctly converted to a
+  // OpDecorate instruction.
+
+  const std::string shader = R"(
+; CHECK: OpDecorate [[t:%\w+]] DescriptorSet 0
+; CHECK: OpDecorate [[t]] Binding 0
+; CHECK: OpDecorate [[t]] RelaxedPrecision
+; CHECK: OpDecorate [[s:%\w+]] DescriptorSet 0
+; CHECK: OpDecorate [[s]] Binding 1
+               OpCapability Shader
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint Fragment %PSMain "PSMain" %in_var_TEXCOORD %out_var_SV_Target
+               OpExecutionMode %PSMain OriginUpperLeft
+               OpSource HLSL 600
+               OpName %sampler2D_h "sampler2D_h"
+               OpMemberName %sampler2D_h 0 "t"
+               OpMemberName %sampler2D_h 1 "s"
+               OpName %type_2d_image "type.2d.image"
+               OpName %type_sampler "type.sampler"
+               OpName %_MainTex "_MainTex"
+               OpName %in_var_TEXCOORD "in.var.TEXCOORD"
+               OpName %out_var_SV_Target "out.var.SV_Target"
+               OpName %PSMain "PSMain"
+               OpName %type_sampled_image "type.sampled.image"
+               OpDecorate %in_var_TEXCOORD Location 0
+               OpDecorate %out_var_SV_Target Location 0
+               OpDecorate %_MainTex DescriptorSet 0
+               OpDecorate %_MainTex Binding 0
+               OpMemberDecorate %sampler2D_h 0 RelaxedPrecision
+               OpDecorate %out_var_SV_Target RelaxedPrecision
+               OpDecorate %69 RelaxedPrecision
+      %float = OpTypeFloat 32
+%type_2d_image = OpTypeImage %float 2D 2 0 0 1 Unknown
+%type_sampler = OpTypeSampler
+%sampler2D_h = OpTypeStruct %type_2d_image %type_sampler
+%_ptr_UniformConstant_sampler2D_h = OpTypePointer UniformConstant %sampler2D_h
+    %v2float = OpTypeVector %float 2
+%_ptr_Input_v2float = OpTypePointer Input %v2float
+    %v4float = OpTypeVector %float 4
+%_ptr_Output_v4float = OpTypePointer Output %v4float
+       %void = OpTypeVoid
+         %35 = OpTypeFunction %void
+%type_sampled_image = OpTypeSampledImage %type_2d_image
+   %_MainTex = OpVariable %_ptr_UniformConstant_sampler2D_h UniformConstant
+%in_var_TEXCOORD = OpVariable %_ptr_Input_v2float Input
+%out_var_SV_Target = OpVariable %_ptr_Output_v4float Output
+     %PSMain = OpFunction %void None %35
+         %43 = OpLabel
+         %44 = OpLoad %v2float %in_var_TEXCOORD
+         %57 = OpLoad %sampler2D_h %_MainTex
+         %72 = OpCompositeExtract %type_2d_image %57 0
+         %73 = OpCompositeExtract %type_sampler %57 1
+         %68 = OpSampledImage %type_sampled_image %72 %73
+         %69 = OpImageSampleImplicitLod %v4float %68 %44 None
+               OpStore %out_var_SV_Target %69
+               OpReturn
+               OpFunctionEnd
+)";
+
+  SinglePassRunAndMatch<DescriptorScalarReplacement>(shader, true);
+}
+
 }  // namespace
 }  // namespace opt
 }  // namespace spvtools