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.

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;

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.

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