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

[aemerson]

Summary

Adds a new utility function: isConstantOrConstantVector().

[aemerson]

llvm/lib/CodeGen/GlobalISel/Utils.cpp

1019–1030

Actually this doesn't have the same meaning as the DAG counterpart. This one only checks for constant scalars or splats of scalars. I'll rename this to make it clear.

[aemerson]

Rename to isConstantOrConstantSplatVector().

[foad]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

2247

New code should use isZero instead.

[foad]

llvm/lib/CodeGen/GlobalISel/Utils.cpp

1028

64 seems a bit imprecise here. Can't we handle the scalar and vector cases consistently, either using APInt throughout or int64_t throughout?

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

3

Unintended change?

[aemerson]

I've changed it to use the vector element size for the APInt.

[foad]

It still seems a bit messy that we throw away the constant's width in getBuildVectorConstantSplat and then put it back again in isConstantOrConstantSplatVector, but I guess that needn't hold up this patch.

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

 /// Represents a G_IMPLICIT_DEF.
 class GImplicitDef : public GenericMachineInstr {
 public:
   static bool classof(const MachineInstr *MI) {
     return MI->getOpcode() == TargetOpcode::G_IMPLICIT_DEF;
   }
 };
 
+/// 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
 No newline at end of file

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

 ///   %cst = G_CONSTANT iN 4
 ///   %constant_splat = G_BUILD_VECTOR %cst, %cst, ..., %cst
 /// \endcode
 ///
 /// In the above case, this will return a RegOrConstant containing 4.
 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.
 bool matchUnaryPredicate(const MachineRegisterInfo &MRI, Register Reg,
                          std::function<bool(const Constant *ConstVal)> Match,
                          bool AllowUndefs = false);
 
 /// Returns true if given the TargetLowering's boolean contents information,

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

 
 bool CombinerHelper::matchUndefSelectCmp(MachineInstr &MI) {
   assert(MI.getOpcode() == TargetOpcode::G_SELECT);
   return getOpcodeDef(TargetOpcode::G_IMPLICIT_DEF, MI.getOperand(1).getReg(),
                       MRI);
 }
 
 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;
-  }
+  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) {
   MI.eraseFromParent();
   return true;
 }
 
 bool CombinerHelper::matchEqualDefs(const MachineOperand &MOP1,

llvm/lib/CodeGen/GlobalISel/Utils.cpp

if (auto Splat = getBuildVectorConstantSplat(MI, MRI))
     return RegOrConstant(*Splat);
   auto Reg = MI.getOperand(1).getReg();
   if (any_of(make_range(MI.operands_begin() + 2, MI.operands_end()),
              [&Reg](const MachineOperand &Op) { return Op.getReg() != Reg; }))
     return None;
   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 RegOrCst = getBuildVectorConstantSplat(MI, MRI);
+  if (!RegOrCst)
+    return None;
+  return APInt(64, *RegOrCst, true);
+}
+
 bool llvm::matchUnaryPredicate(
     const MachineRegisterInfo &MRI, Register Reg,
     std::function<bool(const Constant *ConstVal)> Match, bool AllowUndefs) {
 
   const MachineInstr *Def = getDefIgnoringCopies(Reg, MRI);
   if (AllowUndefs && Def->getOpcode() == TargetOpcode::G_IMPLICIT_DEF)
     return Match(nullptr);
 

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

 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
 # RUN: llc -run-pass=aarch64-prelegalizer-combiner -verify-machineinstrs -mtriple aarch64-unknown-unknown %s -o - | FileCheck %s
-# RUN: llc -debugify-and-strip-all-safe -run-pass=aarch64-prelegalizer-combiner -verify-machineinstrs -mtriple aarch64-unknown-unknown %s -o - | FileCheck %s
+# R UN: llc -debugify-and-strip-all-safe -run-pass=aarch64-prelegalizer-combiner -verify-machineinstrs -mtriple aarch64-unknown-unknown %s -o - | FileCheck %s
 ---
 # select (c, x, x) -> x
 name:            test_combine_select_same_res
 body:             |
   bb.1:
     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
     $x0 = COPY %2(s64)
 ...
 ---
 # select (undef, x, y) -> y
 name:            test_combine_select_undef_res0_res1
 body:             |
   bb.1:
     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
     %3:_(s64) = G_SELECT %2, %0, %1
     $x0 = COPY %3(s64)
 ...
 ---
 # select (false, x, y) -> y
 name:            test_combine_select_false_res0_res1
 body:             |
   bb.1:
     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>)
+...