spirv-fuzz: Use component-wise selectors when flattening conditional branches

source/fuzz/fuzzer_pass_flatten_conditional_branches.cpp

@@ -72,6 +72,107 @@ void FuzzerPassFlattenConditionalBranches::Apply() {
         continue;
       }
 
+      uint32_t convergence_block_id =
+          TransformationFlattenConditionalBranch::FindConvergenceBlock(
+              GetIRContext(), *header);
+
+      // If the SPIR-V version is restricted so that OpSelect can only work on
+      // scalar, pointer and vector types then we cannot apply this
+      // transformation to a header whose convergence block features OpPhi
+      // instructions on different types, as we cannot convert such instructions
+      // to OpSelect instructions.
+      if (TransformationFlattenConditionalBranch::
+              OpSelectArgumentsAreRestricted(GetIRContext())) {
+        if (!GetIRContext()
+                 ->cfg()
+                 ->block(convergence_block_id)
+                 ->WhileEachPhiInst(
+                     [this](opt::Instruction* phi_instruction) -> bool {
+                       switch (GetIRContext()
+                                   ->get_def_use_mgr()
+                                   ->GetDef(phi_instruction->type_id())
+                                   ->opcode()) {
+                         case SpvOpTypeBool:
+                         case SpvOpTypeInt:
+                         case SpvOpTypeFloat:
+                         case SpvOpTypePointer:
+                         case SpvOpTypeVector:
+                           return true;
+                         default:
+                           return false;
+                       }
+                     })) {
+          // An OpPhi is performed on a type not supported by OpSelect; we
+          // cannot flatten this selection.
+          continue;
+        }
+      }
+
+      // If the construct's convergence block features OpPhi instructions with
+      // vector result types then we may be *forced*, by the SPIR-V version, to
+      // turn these into component-wise OpSelect instructions, or we might wish
+      // to do so anyway.  The following booleans capture whether we will opt
+      // to use a component-wise select even if we don't have to.
+      bool use_component_wise_2d_select_even_if_optional =
+          GetFuzzerContext()->ChooseEven();
+      bool use_component_wise_3d_select_even_if_optional =
+          GetFuzzerContext()->ChooseEven();
+      bool use_component_wise_4d_select_even_if_optional =
+          GetFuzzerContext()->ChooseEven();
+
+      // If we do need to perform any component-wise selections, we will need a
+      // fresh id for a boolean vector representing the selection's condition
+      // repeated N times, where N is the vector dimension.
+      uint32_t fresh_id_for_bvec2_selector = 0;
+      uint32_t fresh_id_for_bvec3_selector = 0;
+      uint32_t fresh_id_for_bvec4_selector = 0;
+
+      GetIRContext()
+          ->cfg()
+          ->block(convergence_block_id)
+          ->ForEachPhiInst([this, &fresh_id_for_bvec2_selector,
+                            &fresh_id_for_bvec3_selector,
+                            &fresh_id_for_bvec4_selector,
+                            use_component_wise_2d_select_even_if_optional,
+                            use_component_wise_3d_select_even_if_optional,
+                            use_component_wise_4d_select_even_if_optional](
+                               opt::Instruction* phi_instruction) {
+            opt::Instruction* type_instruction =
+                GetIRContext()->get_def_use_mgr()->GetDef(
+                    phi_instruction->type_id());
+            switch (type_instruction->opcode()) {
+              case SpvOpTypeVector: {
+                uint32_t dimension =
+                    type_instruction->GetSingleWordInOperand(1);
+                switch (dimension) {
+                  case 2:
+                    PrepareForOpPhiOnVectors(
+                        dimension,
+                        use_component_wise_2d_select_even_if_optional,
+                        &fresh_id_for_bvec2_selector);
+                    break;
+                  case 3:
+                    PrepareForOpPhiOnVectors(
+                        dimension,
+                        use_component_wise_3d_select_even_if_optional,
+                        &fresh_id_for_bvec3_selector);
+                    break;
+                  case 4:
+                    PrepareForOpPhiOnVectors(
+                        dimension,
+                        use_component_wise_4d_select_even_if_optional,
+                        &fresh_id_for_bvec4_selector);
+                    break;
+                  default:
+                    assert(false && "Invalid vector dimension.");
+                }
+                break;
+              }
+              default:
+                break;
+            }
+          });
+
       // Some instructions will require to be enclosed inside conditionals
       // because they have side effects (for example, loads and stores). Some of
       // this have no result id, so we require instruction descriptors to
@@ -116,10 +217,31 @@ void FuzzerPassFlattenConditionalBranches::Apply() {
       // Apply the transformation, evenly choosing whether to lay out the true
       // branch or the false branch first.
       ApplyTransformation(TransformationFlattenConditionalBranch(
-          header->id(), GetFuzzerContext()->ChooseEven(), wrappers_info));
+          header->id(), GetFuzzerContext()->ChooseEven(),
+          fresh_id_for_bvec2_selector, fresh_id_for_bvec3_selector,
+          fresh_id_for_bvec4_selector, wrappers_info));
     }
   }
 }
 
+void FuzzerPassFlattenConditionalBranches::PrepareForOpPhiOnVectors(
+    uint32_t vector_dimension, bool use_vector_select_if_optional,
+    uint32_t* fresh_id_for_bvec_selector) {
+  if (*fresh_id_for_bvec_selector != 0) {
+    // We already have a fresh id for a component-wise OpSelect of this
+    // dimension
+    return;
+  }
+  if (TransformationFlattenConditionalBranch::OpSelectArgumentsAreRestricted(
+          GetIRContext()) ||
+      use_vector_select_if_optional) {
+    // We either have to, or have chosen to, perform a component-wise select, so
+    // we ensure that the right boolean vector type is available, and grab a
+    // fresh id.
+    FindOrCreateVectorType(FindOrCreateBoolType(), vector_dimension);
+    *fresh_id_for_bvec_selector = GetFuzzerContext()->GetFreshId();
+  }
+}
+
 }  // namespace fuzz
 }  // namespace spvtools

source/fuzz/fuzzer_pass_flatten_conditional_branches.h

@@ -30,6 +30,16 @@ class FuzzerPassFlattenConditionalBranches : public FuzzerPass {
   ~FuzzerPassFlattenConditionalBranches() override;
 
   void Apply() override;
+
+ private:
+  // If the SPIR-V version requires vector OpSelects to be component-wise, or
+  // if |use_vector_select_if_optional| holds, |fresh_id_for_bvec_selector| is
+  // populated with a fresh id if it is currently zero, and a
+  // |vector_dimension|-dimensional boolean vector type is added to the module
+  // if not already present.
+  void PrepareForOpPhiOnVectors(uint32_t vector_dimension,
+                                bool use_vector_select_if_optional,
+                                uint32_t* fresh_id_for_bvec_selector);
 };
 }  // namespace fuzz
 }  // namespace spvtools

source/fuzz/fuzzer_pass_replace_opselects_with_conditional_branches.cpp

@@ -65,6 +65,16 @@ void FuzzerPassReplaceOpSelectsWithConditionalBranches::Apply() {
           continue;
         }
 
+        // If the selector does not have scalar boolean type (i.e., it is a
+        // boolean vector) then ignore this OpSelect.
+        if (GetIRContext()
+                ->get_def_use_mgr()
+                ->GetDef(fuzzerutil::GetTypeId(
+                    GetIRContext(), instruction.GetSingleWordInOperand(0)))
+                ->opcode() != SpvOpTypeBool) {
+          continue;
+        }
+
         // If the block is a loop header and we need to split it, the
         // transformation cannot be applied because loop headers cannot be
         // split. We can break out of this loop because the transformation can

source/fuzz/protobufs/spvtoolsfuzz.proto

@@ -1468,10 +1468,23 @@ message TransformationFlattenConditionalBranch {
   // field is true.
   bool true_branch_first = 2;
 
+  // If the convergence block contains an OpPhi with bvec2 result type, it may
+  // be necessary to introduce a bvec2 with the selection construct's condition
+  // in both components in order to turn the OpPhi into an OpSelect.  This
+  // this field provides a fresh id for an OpCompositeConstruct instruction for
+  // this purpose.  It should be set to 0 if no such instruction is required.
+  uint32 fresh_id_for_bvec2_selector = 3;
+
+  // The same as |fresh_id_for_bvec2_selector| but for the bvec3 case.
+  uint32 fresh_id_for_bvec3_selector = 4;
+
+  // The same as |fresh_id_for_bvec2_selector| but for the bvec4 case.
+  uint32 fresh_id_for_bvec4_selector = 5;
+
   // A list of instructions with side effects, which must be enclosed
   // inside smaller conditionals before flattening the main one, and
   // the corresponding fresh ids and module ids needed.
-  repeated SideEffectWrapperInfo side_effect_wrapper_info = 3;
+  repeated SideEffectWrapperInfo side_effect_wrapper_info = 6;
 }
 
 message TransformationFunctionCall {