Handle out-of-bounds accesses in VDCE (#4518)

It is possible that other optimization will propagate a value into an OpCompositeExtract or OpVectorShuffle instruction that is larger than the vector size. Vector DCE has to be able to handle it. Fixes https://github.com/KhronosGroup/SPIRV-Tools/issues/4513.
2024-11-27 17:40:28 +00:00 · 2021-09-13 09:57:44 -04:00 · 2021-09-13 09:57:44 -04:00 · 8865b20295
commit 8865b20295
parent 4f4f76037c
3 changed files with 94 additions and 11 deletions
--- a/source/opt/vector_dce.cpp
+++ b/source/opt/vector_dce.cpp
@ -110,7 +110,11 @@ void VectorDCE::MarkExtractUseAsLive(const Instruction* current_inst,
    if (current_inst->NumInOperands() < 2) {
      new_item.components = live_elements;
    } else {
-      new_item.components.Set(current_inst->GetSingleWordInOperand(1));
+      uint32_t element_index = current_inst->GetSingleWordInOperand(1);
+      uint32_t item_size = GetVectorComponentCount(operand_inst->type_id());
+      if (element_index < item_size) {
+        new_item.components.Set(element_index);
+      }
    }
    AddItemToWorkListIfNeeded(new_item, live_components, work_list);
  }
@ -176,10 +180,10 @@ void VectorDCE::MarkVectorShuffleUsesAsLive(
  second_operand.instruction =
      def_use_mgr->GetDef(current_item.instruction->GetSingleWordInOperand(1));

-  analysis::TypeManager* type_mgr = context()->get_type_mgr();
-  analysis::Vector* first_type =
-      type_mgr->GetType(first_operand.instruction->type_id())->AsVector();
-  uint32_t size_of_first_operand = first_type->element_count();
+  uint32_t size_of_first_operand =
+      GetVectorComponentCount(first_operand.instruction->type_id());
+  uint32_t size_of_second_operand =
+      GetVectorComponentCount(second_operand.instruction->type_id());

  for (uint32_t in_op = 2; in_op < current_item.instruction->NumInOperands();
       ++in_op) {
@ -187,7 +191,7 @@ void VectorDCE::MarkVectorShuffleUsesAsLive(
    if (current_item.components.Get(in_op - 2)) {
      if (index < size_of_first_operand) {
        first_operand.components.Set(index);
-      } else {
+      } else if (index - size_of_first_operand < size_of_second_operand) {
        second_operand.components.Set(index - size_of_first_operand);
      }
    }
@ -202,7 +206,6 @@ void VectorDCE::MarkCompositeContructUsesAsLive(
    VectorDCE::LiveComponentMap* live_components,
    std::vector<VectorDCE::WorkListItem>* work_list) {
  analysis::DefUseManager* def_use_mgr = context()->get_def_use_mgr();
-  analysis::TypeManager* type_mgr = context()->get_type_mgr();

  uint32_t current_component = 0;
  Instruction* current_inst = work_item.instruction;
@ -223,8 +226,7 @@ void VectorDCE::MarkCompositeContructUsesAsLive(
      assert(HasVectorResult(op_inst));
      WorkListItem new_work_item;
      new_work_item.instruction = op_inst;
-      uint32_t op_vector_size =
-          type_mgr->GetType(op_inst->type_id())->AsVector()->element_count();
+      uint32_t op_vector_size = GetVectorComponentCount(op_inst->type_id());

      for (uint32_t op_vector_idx = 0; op_vector_idx < op_vector_size;
           op_vector_idx++, current_component++) {
@ -297,6 +299,18 @@ bool VectorDCE::HasScalarResult(const Instruction* inst) const {
  }
 }

+uint32_t VectorDCE::GetVectorComponentCount(uint32_t type_id) {
+  assert(type_id != 0 &&
+         "Trying to get the vector element count, but the type id is 0");
+  analysis::TypeManager* type_mgr = context()->get_type_mgr();
+  const analysis::Type* type = type_mgr->GetType(type_id);
+  const analysis::Vector* vector_type = type->AsVector();
+  assert(
+      vector_type &&
+      "Trying to get the vector element count, but the type is not a vector");
+  return vector_type->element_count();
+}
+
 bool VectorDCE::RewriteInstructions(
    Function* function, const VectorDCE::LiveComponentMap& live_components) {
  bool modified = false;
--- a/source/opt/vector_dce.h
+++ b/source/opt/vector_dce.h
@ -94,12 +94,15 @@ class VectorDCE : public MemPass {
  // Returns true if the result of |inst| is a vector or a scalar.
  bool HasVectorOrScalarResult(const Instruction* inst) const;

-  // Returns true if the result of |inst| is a scalar.
+  // Returns true if the result of |inst| is a vector.
  bool HasVectorResult(const Instruction* inst) const;

-  // Returns true if the result of |inst| is a vector.
+  // Returns true if the result of |inst| is a scalar.
  bool HasScalarResult(const Instruction* inst) const;

+  // Returns the number of elements in the vector type with id |type_id|.
+  uint32_t GetVectorComponentCount(uint32_t type_id);
+
  // Adds |work_item| to |work_list| if it is not already live according to
  // |live_components|.  |live_components| is updated to indicate that
  // |work_item| is now live.
--- a/test/opt/vector_dce_test.cpp
+++ b/test/opt/vector_dce_test.cpp
@ -1351,6 +1351,72 @@ OpFunctionEnd
  SinglePassRunAndMatch<VectorDCE>(text, true);
 }

+TEST_F(VectorDCETest, OutOfBoundsExtract) {
+  // It tests that the vector DCE pass is able to handle an extract with an
+  // index that is out of bounds.
+  const std::string text = R"(
+; CHECK: [[undef:%\w+]] = OpUndef %v4float
+; CHECK: OpCompositeExtract %float [[undef]] 8
+                     OpCapability Shader
+                     OpMemoryModel Logical GLSL450
+                     OpEntryPoint Fragment %main "main" %OutColor
+                     OpExecutionMode %main OriginUpperLeft
+                     OpDecorate %OutColor Location 0
+             %void = OpTypeVoid
+               %10 = OpTypeFunction %void
+            %float = OpTypeFloat 32
+          %v4float = OpTypeVector %float 4
+%_ptr_Output_float = OpTypePointer Output %float
+         %OutColor = OpVariable %_ptr_Output_float Output
+             %null = OpConstantNull %v4float
+          %float_1 = OpConstant %float 1
+             %main = OpFunction %void None %10
+               %28 = OpLabel
+               %33 = OpCompositeInsert %v4float %float_1 %null 1
+          %extract = OpCompositeExtract %float %33 8
+                     OpStore %OutColor %extract
+                     OpReturn
+                     OpFunctionEnd
+)";
+
+  SetDisassembleOptions(SPV_BINARY_TO_TEXT_OPTION_NO_HEADER |
+                        SPV_BINARY_TO_TEXT_OPTION_FRIENDLY_NAMES);
+  SinglePassRunAndMatch<VectorDCE>(text, false);
+}
+
+TEST_F(VectorDCETest, OutOfBoundsShuffle) {
+  // It tests that the vector DCE pass is able to handle a shuffle with an
+  // index that is out of bounds.
+  const std::string text = R"(
+; CHECK: [[undef:%\w+]] = OpUndef %v4float
+; CHECK: OpVectorShuffle %v4float [[undef]] [[undef]] 9 10 11 12
+                     OpCapability Shader
+                     OpMemoryModel Logical GLSL450
+                     OpEntryPoint Fragment %main "main" %OutColor
+                     OpExecutionMode %main OriginUpperLeft
+                     OpDecorate %OutColor Location 0
+             %void = OpTypeVoid
+               %10 = OpTypeFunction %void
+            %float = OpTypeFloat 32
+          %v4float = OpTypeVector %float 4
+%_ptr_Output_v4float = OpTypePointer Output %v4float
+         %OutColor = OpVariable %_ptr_Output_v4float Output
+             %null = OpConstantNull %v4float
+          %float_1 = OpConstant %float 1
+             %main = OpFunction %void None %10
+               %28 = OpLabel
+               %33 = OpCompositeInsert %v4float %float_1 %null 1
+          %shuffle = OpVectorShuffle %v4float %33 %33 9 10 11 12
+                     OpStore %OutColor %shuffle
+                     OpReturn
+                     OpFunctionEnd
+)";
+
+  SetDisassembleOptions(SPV_BINARY_TO_TEXT_OPTION_NO_HEADER |
+                        SPV_BINARY_TO_TEXT_OPTION_FRIENDLY_NAMES);
+  SinglePassRunAndMatch<VectorDCE>(text, false);
+}
+
 }  // namespace
 }  // namespace opt
 }  // namespace spvtools