[opt] Fix constant folding of OpSelect with a scalar condition

source/opt/const_folding_rules.cpp

@@ -1276,6 +1276,67 @@ ConstantFoldingRule FoldInvariantSelect() {
   };
 }
 
+// Folds an OpSelect with a constant condition to the selected object.  A
+// scalar condition, which since SPIR-V 1.4 is allowed for any result type,
+// selects a whole object.  A vector condition selects each component
+// separately.
+ConstantFoldingRule FoldSelect() {
+  return [](IRContext* context, Instruction* inst,
+            const std::vector<const analysis::Constant*>& constants)
+             -> const analysis::Constant* {
+    assert(inst->opcode() == spv::Op::OpSelect);
+    const analysis::Constant* condition = constants[0];
+    const analysis::Constant* object_1 = constants[1];
+    const analysis::Constant* object_2 = constants[2];
+    if (!condition || !object_1 || !object_2) {
+      return nullptr;
+    }
+
+    if (condition->type()->AsVector() == nullptr) {
+      bool cond_is_true = false;
+      if (const analysis::BoolConstant* bool_condition =
+              condition->AsBoolConstant()) {
+        cond_is_true = bool_condition->value();
+      } else if (!condition->AsNullConstant()) {
+        return nullptr;
+      }
+      return cond_is_true ? object_1 : object_2;
+    }
+
+    analysis::ConstantManager* const_mgr = context->get_constant_mgr();
+    analysis::TypeManager* type_mgr = context->get_type_mgr();
+    const analysis::Vector* vector_type =
+        type_mgr->GetType(inst->type_id())->AsVector();
+    assert(vector_type != nullptr &&
+           "The result type of an OpSelect with a vector condition must be a "
+           "vector");
+
+    std::vector<const analysis::Constant*> cond_components =
+        condition->GetVectorComponents(const_mgr);
+    std::vector<const analysis::Constant*> object_1_components =
+        object_1->GetVectorComponents(const_mgr);
+    std::vector<const analysis::Constant*> object_2_components =
+        object_2->GetVectorComponents(const_mgr);
+
+    std::vector<uint32_t> ids;
+    for (uint32_t i = 0; i < cond_components.size(); ++i) {
+      const analysis::BoolConstant* bool_component =
+          cond_components[i]->AsBoolConstant();
+      const bool cond_is_true =
+          bool_component != nullptr && bool_component->value();
+      const analysis::Constant* component =
+          cond_is_true ? object_1_components[i] : object_2_components[i];
+      Instruction* component_inst =
+          const_mgr->GetDefiningInstruction(component);
+      if (component_inst == nullptr) {
+        return nullptr;
+      }
+      ids.push_back(component_inst->result_id());
+    }
+    return const_mgr->GetConstant(vector_type, ids);
+  };
+}
+
 // Folds an OpDot where all of the inputs are constants to a
 // constant.  A new constant is created if necessary.
 ConstantFoldingRule FoldOpDotWithConstants() {
@@ -2086,6 +2147,7 @@ void ConstantFoldingRules::AddFoldingRules() {
   rules_[spv::Op::OpFSub].push_back(FoldRedundantSub());
 
   rules_[spv::Op::OpSelect].push_back(FoldInvariantSelect());
+  rules_[spv::Op::OpSelect].push_back(FoldSelect());
 
   rules_[spv::Op::OpFOrdEqual].push_back(FoldFOrdEqual());
 

test/opt/fold_spec_const_op_composite_test.cpp

@@ -533,6 +533,163 @@ TEST_F(FoldSpecConstantOpAndCompositePassBasicTest,
   SinglePassRunAndMatch<FoldSpecConstantOpAndCompositePass>(test, false);
 }
 
+TEST_F(FoldSpecConstantOpAndCompositePassBasicTest,
+       SelectVectorsWithScalarCondition) {
+  SetTargetEnv(SPV_ENV_UNIVERSAL_1_4);
+  const std::string test =
+      R"(
+               OpCapability Shader
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint GLCompute %1 "main"
+               OpExecutionMode %1 LocalSize 1 1 1
+       %void = OpTypeVoid
+          %3 = OpTypeFunction %void
+       %bool = OpTypeBool
+        %int = OpTypeInt 32 1
+      %v2int = OpTypeVector %int 2
+       %true = OpConstantTrue %bool
+      %false = OpConstantFalse %bool
+      %int_1 = OpConstant %int 1
+      %int_2 = OpConstant %int 2
+      %int_3 = OpConstant %int 3
+      %int_4 = OpConstant %int 4
+
+; Since SPIR-V 1.4, the condition of OpSelect may be a scalar even when the
+; objects are vectors, in which case it selects the whole object.
+; CHECK: OpConstantComposite %v2int %int_1 %int_2
+; CHECK: OpConstantComposite %v2int %int_3 %int_4
+; CHECK: [[sel_t:%\w+]] = OpConstantComposite %v2int %int_1 %int_2
+; CHECK: [[sel_f:%\w+]] = OpConstantComposite %v2int %int_3 %int_4
+; CHECK-NOT: OpSpecConstantOp
+; CHECK: OpIAdd %v2int [[sel_t]] [[sel_f]]
+          %a = OpConstantComposite %v2int %int_1 %int_2
+          %b = OpConstantComposite %v2int %int_3 %int_4
+      %sel_t = OpSpecConstantOp %v2int Select %true %a %b
+      %sel_f = OpSpecConstantOp %v2int Select %false %a %b
+          %1 = OpFunction %void None %3
+      %label = OpLabel
+        %add = OpIAdd %v2int %sel_t %sel_f
+               OpReturn
+               OpFunctionEnd
+)";
+
+  SinglePassRunAndMatch<FoldSpecConstantOpAndCompositePass>(test, true);
+}
+
+TEST_F(FoldSpecConstantOpAndCompositePassBasicTest,
+       SelectVectorsWithVectorCondition) {
+  SetTargetEnv(SPV_ENV_UNIVERSAL_1_4);
+  const std::string test =
+      R"(
+               OpCapability Shader
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint GLCompute %1 "main"
+               OpExecutionMode %1 LocalSize 1 1 1
+       %void = OpTypeVoid
+          %3 = OpTypeFunction %void
+       %bool = OpTypeBool
+     %v2bool = OpTypeVector %bool 2
+        %int = OpTypeInt 32 1
+      %v2int = OpTypeVector %int 2
+       %true = OpConstantTrue %bool
+      %false = OpConstantFalse %bool
+      %int_1 = OpConstant %int 1
+      %int_2 = OpConstant %int 2
+      %int_3 = OpConstant %int 3
+      %int_4 = OpConstant %int 4
+       %cond = OpConstantComposite %v2bool %true %false
+          %a = OpConstantComposite %v2int %int_1 %int_2
+          %b = OpConstantComposite %v2int %int_3 %int_4
+
+; A vector condition mixes the two objects per component.
+; CHECK: [[mix:%\w+]] = OpConstantComposite %v2int %int_1 %int_4
+; CHECK-NOT: OpSpecConstantOp
+; CHECK: OpIAdd %v2int [[mix]] [[mix]]
+        %sel = OpSpecConstantOp %v2int Select %cond %a %b
+          %1 = OpFunction %void None %3
+      %label = OpLabel
+        %add = OpIAdd %v2int %sel %sel
+               OpReturn
+               OpFunctionEnd
+)";
+
+  SinglePassRunAndMatch<FoldSpecConstantOpAndCompositePass>(test, true);
+}
+
+TEST_F(FoldSpecConstantOpAndCompositePassBasicTest,
+       SelectVectorIdOverflowWhileMaterializingNullComponent) {
+  SetTargetEnv(SPV_ENV_UNIVERSAL_1_4);
+  const std::string test = R"(
+               OpCapability Shader
+               OpMemoryModel Logical GLSL450
+       %bool = OpTypeBool
+     %v2bool = OpTypeVector %bool 2
+        %int = OpTypeInt 32 1
+      %v2int = OpTypeVector %int 2
+       %true = OpConstantTrue %bool
+      %false = OpConstantFalse %bool
+       %cond = OpConstantComposite %v2bool %true %false
+      %int_1 = OpConstant %int 1
+      %int_2 = OpConstant %int 2
+  %non_zero = OpConstantComposite %v2int %int_1 %int_2
+       %null = OpConstantNull %v2int
+        %sel = OpSpecConstantOp %v2int Select %cond %null %non_zero
+)";
+
+  std::string diagnostic;
+  SetMessageConsumer(
+      [&diagnostic](spv_message_level_t, const char*, const spv_position_t&,
+                    const char* message) { diagnostic = message; });
+  std::unique_ptr<IRContext> context = AssembleModule(test);
+  ASSERT_NE(context, nullptr);
+  context->set_max_id_bound(context->module()->id_bound());
+
+  FoldSpecConstantOpAndCompositePass pass;
+  EXPECT_EQ(pass.Run(context.get()), Pass::Status::Failure);
+  EXPECT_EQ(diagnostic, "ID overflow. Try running compact-ids.");
+}
+
+TEST_F(FoldSpecConstantOpAndCompositePassBasicTest,
+       SelectStructsWithScalarCondition) {
+  SetTargetEnv(SPV_ENV_UNIVERSAL_1_4);
+  const std::string test =
+      R"(
+               OpCapability Shader
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint GLCompute %1 "main"
+               OpExecutionMode %1 LocalSize 1 1 1
+       %void = OpTypeVoid
+          %3 = OpTypeFunction %void
+       %bool = OpTypeBool
+        %int = OpTypeInt 32 1
+     %struct = OpTypeStruct %int %int
+       %true = OpConstantTrue %bool
+      %false = OpConstantFalse %bool
+      %int_1 = OpConstant %int 1
+      %int_2 = OpConstant %int 2
+      %int_3 = OpConstant %int 3
+      %int_4 = OpConstant %int 4
+
+; Since SPIR-V 1.4, the result type of OpSelect may be a non-vector
+; composite, with a scalar condition selecting the whole object.
+; CHECK: OpConstantComposite [[struct:%\w+]] %int_1 %int_2
+; CHECK: OpConstantComposite [[struct]] %int_3 %int_4
+; CHECK: OpConstantComposite [[struct]] %int_1 %int_2
+; CHECK: OpConstantComposite [[struct]] %int_3 %int_4
+; CHECK-NOT: OpSpecConstantOp
+          %a = OpConstantComposite %struct %int_1 %int_2
+          %b = OpConstantComposite %struct %int_3 %int_4
+      %sel_t = OpSpecConstantOp %struct Select %true %a %b
+      %sel_f = OpSpecConstantOp %struct Select %false %a %b
+          %1 = OpFunction %void None %3
+      %label = OpLabel
+               OpReturn
+               OpFunctionEnd
+)";
+
+  SinglePassRunAndMatch<FoldSpecConstantOpAndCompositePass>(test, true);
+}
+
 TEST_F(FoldSpecConstantOpAndCompositePassBasicTest,
        CompositeInsertVectorKeepNull) {
   const std::string test =