[GlobalISel] Extend G_SELECT of known condition combine to vectors.

Adds a new utility function: isConstantOrConstantSplatVector().

Differential Revision: https://reviews.llvm.org/D110786

Author: aemerson

llvm/include/llvm/CodeGen/GlobalISel/GenericMachineInstrs.h

@@ -214,6 +214,18 @@ class GImplicitDef : public GenericMachineInstr {
   }
 };
 
+/// Represents a G_SELECT.
+class GSelect : public GenericMachineInstr {
+public:
+  Register getCondReg() const { return getReg(1); }
+  Register getTrueReg() const { return getReg(2); }
+  Register getFalseReg() const { return getReg(3); }
+
+  static bool classof(const MachineInstr *MI) {
+    return MI->getOpcode() == TargetOpcode::G_SELECT;
+  }
+};
+
 } // namespace llvm
 
 #endif // LLVM_CODEGEN_GLOBALISEL_GENERICMACHINEINSTRS_H

llvm/include/llvm/CodeGen/GlobalISel/Utils.h

@@ -397,6 +397,12 @@ bool isBuildVectorAllOnes(const MachineInstr &MI,
 Optional<RegOrConstant> getVectorSplat(const MachineInstr &MI,
                                        const MachineRegisterInfo &MRI);
 
+/// Determines if \p MI defines a constant integer or a splat vector of
+/// constant integers.
+/// \returns the scalar constant or None.
+Optional<APInt> isConstantOrConstantSplatVector(MachineInstr &MI,
+                                                const MachineRegisterInfo &MRI);
+
 /// Attempt to match a unary predicate against a scalar/splat constant or every
 /// element of a constant G_BUILD_VECTOR. If \p ConstVal is null, the source
 /// value was undef.

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

@@ -2239,13 +2239,13 @@ bool CombinerHelper::matchUndefSelectCmp(MachineInstr &MI) {
 }
 
 bool CombinerHelper::matchConstantSelectCmp(MachineInstr &MI, unsigned &OpIdx) {
-  assert(MI.getOpcode() == TargetOpcode::G_SELECT);
-  if (auto MaybeCstCmp =
-          getIConstantVRegValWithLookThrough(MI.getOperand(1).getReg(), MRI)) {
-    OpIdx = MaybeCstCmp->Value.isNullValue() ? 3 : 2;
-    return true;
-  }
-  return false;
+  GSelect &SelMI = cast<GSelect>(MI);
+  auto Cst =
+      isConstantOrConstantSplatVector(*MRI.getVRegDef(SelMI.getCondReg()), MRI);
+  if (!Cst)
+    return false;
+  OpIdx = Cst->isZero() ? 3 : 2;
+  return true;
 }
 
 bool CombinerHelper::eraseInst(MachineInstr &MI) {

llvm/lib/CodeGen/GlobalISel/Utils.cpp

@@ -1016,6 +1016,19 @@ Optional<RegOrConstant> llvm::getVectorSplat(const MachineInstr &MI,
   return RegOrConstant(Reg);
 }
 
+Optional<APInt>
+llvm::isConstantOrConstantSplatVector(MachineInstr &MI,
+                                      const MachineRegisterInfo &MRI) {
+  Register Def = MI.getOperand(0).getReg();
+  if (auto C = getIConstantVRegValWithLookThrough(Def, MRI))
+    return C->Value;
+  auto MaybeCst = getBuildVectorConstantSplat(MI, MRI);
+  if (!MaybeCst)
+    return None;
+  const unsigned ScalarSize = MRI.getType(Def).getScalarSizeInBits();
+  return APInt(ScalarSize, *MaybeCst, true);
+}
+
 bool llvm::matchUnaryPredicate(
     const MachineRegisterInfo &MRI, Register Reg,
     std::function<bool(const Constant *ConstVal)> Match, bool AllowUndefs) {

llvm/test/CodeGen/AArch64/GlobalISel/combine-select.mir

@@ -9,7 +9,7 @@ body:             |
     liveins: $x0, $x1
     ; CHECK-LABEL: name: test_combine_select_same_res
     ; CHECK: [[COPY:%[0-9]+]]:_(s64) = COPY $x0
-    ; CHECK: $x0 = COPY [[COPY]](s64)
+    ; CHECK-NEXT: $x0 = COPY [[COPY]](s64)
     %0:_(s64) = COPY $x0
     %1:_(s1) = G_TRUNC %0
     %2:_(s64) = G_SELECT %1, %0, %0
@@ -23,7 +23,7 @@ body:             |
     liveins: $x0, $x1
     ; CHECK-LABEL: name: test_combine_select_undef_res0_res1
     ; CHECK: [[COPY:%[0-9]+]]:_(s64) = COPY $x0
-    ; CHECK: $x0 = COPY [[COPY]](s64)
+    ; CHECK-NEXT: $x0 = COPY [[COPY]](s64)
     %0:_(s64) = COPY $x0
     %1:_(s64) = COPY $x1
     %2:_(s1) = G_IMPLICIT_DEF
@@ -38,25 +38,57 @@ body:             |
     liveins: $x0, $x1
     ; CHECK-LABEL: name: test_combine_select_false_res0_res1
     ; CHECK: [[COPY:%[0-9]+]]:_(s64) = COPY $x1
-    ; CHECK: $x0 = COPY [[COPY]](s64)
+    ; CHECK-NEXT: $x0 = COPY [[COPY]](s64)
     %0:_(s64) = COPY $x0
     %1:_(s64) = COPY $x1
     %2:_(s1) = G_CONSTANT i1 false
     %3:_(s64) = G_SELECT %2, %0, %1
     $x0 = COPY %3(s64)
 ...
 ---
+# vector select (false, x, y) -> y
+name:            test_combine_vector_select_false_res0_res1
+body:             |
+  bb.1:
+    liveins: $q0, $q1
+    ; CHECK-LABEL: name: test_combine_vector_select_false_res0_res1
+    ; CHECK: [[COPY:%[0-9]+]]:_(<4 x s32>) = COPY $q1
+    ; CHECK-NEXT: $q0 = COPY [[COPY]](<4 x s32>)
+    %0:_(<4 x s32>) = COPY $q0
+    %1:_(<4 x s32>) = COPY $q1
+    %2:_(s1) = G_CONSTANT i1 false
+    %condvec:_(<4 x s1>) = G_BUILD_VECTOR %2, %2, %2, %2
+    %3:_(<4 x s32>) = G_SELECT %condvec, %0, %1
+    $q0 = COPY %3(<4 x s32>)
+...
+---
 # select (true, x, y) -> x
 name:            test_combine_select_true_res0_res1
 body:             |
   bb.1:
     liveins: $x0, $x1
     ; CHECK-LABEL: name: test_combine_select_true_res0_res1
     ; CHECK: [[COPY:%[0-9]+]]:_(s64) = COPY $x0
-    ; CHECK: $x0 = COPY [[COPY]](s64)
+    ; CHECK-NEXT: $x0 = COPY [[COPY]](s64)
     %0:_(s64) = COPY $x0
     %1:_(s64) = COPY $x1
     %2:_(s1) = G_CONSTANT i1 true
     %3:_(s64) = G_SELECT %2, %0, %1
     $x0 = COPY %3(s64)
 ...
+---
+# vector select (true, x, y) -> x
+name:            test_combine_vector_select_true_res0_res1
+body:             |
+  bb.1:
+    liveins: $q0, $q1
+    ; CHECK-LABEL: name: test_combine_vector_select_true_res0_res1
+    ; CHECK: [[COPY:%[0-9]+]]:_(<4 x s32>) = COPY $q0
+    ; CHECK-NEXT: $q0 = COPY [[COPY]](<4 x s32>)
+    %0:_(<4 x s32>) = COPY $q0
+    %1:_(<4 x s32>) = COPY $q1
+    %2:_(s1) = G_CONSTANT i1 true
+    %condvec:_(<4 x s1>) = G_BUILD_VECTOR %2, %2, %2, %2
+    %3:_(<4 x s32>) = G_SELECT %condvec, %0, %1
+    $q0 = COPY %3(<4 x s32>)
+...