[GlobalISel] Handle constant splat in funnel shift combine

[abinavpp]

Summary

This change adds the constant splat versions of m_ICst() (by using
getBuildVectorConstantSplat()) and uses it in
matchOrShiftToFunnelShift(). The getBuildVectorConstantSplat() name is
shortened to getIConstantSplatVal() so that the *SExtVal() version would
have a more compact name.

[abinavpp]

Moved matchConstantSplat() closer to struct ICstOrSplatMatch

[foad]

Looks reasonable to me.

llvm/test/CodeGen/AMDGPU/GlobalISel/combine-fsh.mir

135

Precommit this?

llvm/test/CodeGen/AMDGPU/GlobalISel/combine-rot.mir

125

Precommit this?

[abinavpp]

Pre-committed tests

[foad]

LGTM. Maybe wait a day in case anyone else has opinions about the names you used for the new matchers.

llvm/include/llvm/CodeGen/GlobalISel/MIPatternMatch.h

121

No else after return.

[Petar.Avramovic]

Name changes look good to me. Wait a little if @paquette or @aemerson have comments.

[paquette]

Aside from @foad's one comment, I think this looks fine.

[abinavpp]

Rebased; Addressed review comment

llvm/include/llvm/CodeGen/GlobalISel/MIPatternMatch.h

 
 inline ConstantMatch<APInt> m_ICst(APInt &Cst) {
   return ConstantMatch<APInt>(Cst);
 }
 inline ConstantMatch<int64_t> m_ICst(int64_t &Cst) {
   return ConstantMatch<int64_t>(Cst);
 }
 
+template <typename ConstT>
+inline Optional<ConstT> matchConstantSplat(Register,
+                                           const MachineRegisterInfo &);
+
+template <>
+inline Optional<APInt> matchConstantSplat(Register Reg,
+                                          const MachineRegisterInfo &MRI) {
+  return getIConstantSplatVal(Reg, MRI);
+}
+
+template <>
+inline Optional<int64_t> matchConstantSplat(Register Reg,
+                                            const MachineRegisterInfo &MRI) {
+  return getIConstantSplatSExtVal(Reg, MRI);
+}
+
+template <typename ConstT> struct ICstOrSplatMatch {
+  ConstT &CR;
+  ICstOrSplatMatch(ConstT &C) : CR(C) {}
+  bool match(const MachineRegisterInfo &MRI, Register Reg) {
+    if (auto MaybeCst = matchConstant<ConstT>(Reg, MRI)) {
+      CR = *MaybeCst;
+      return true;
+    }
+
+    if (auto MaybeCstSplat = matchConstantSplat<ConstT>(Reg, MRI)) {
+      CR = *MaybeCstSplat;
+      return true;
+    }
+
+    return false;
+  };
+};
+
+inline ICstOrSplatMatch<APInt> m_ICstOrSplat(APInt &Cst) {
+  return ICstOrSplatMatch<APInt>(Cst);
+}
+
+inline ICstOrSplatMatch<int64_t> m_ICstOrSplat(int64_t &Cst) {
+  return ICstOrSplatMatch<int64_t>(Cst);
+}
+
 struct GCstAndRegMatch {
   Optional<ValueAndVReg> &ValReg;
   GCstAndRegMatch(Optional<ValueAndVReg> &ValReg) : ValReg(ValReg) {}
   bool match(const MachineRegisterInfo &MRI, Register Reg) {
     ValReg = getIConstantVRegValWithLookThrough(Reg, MRI);
     return ValReg ? true : false;
   }
 };

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

int64_t getCst() const {
     return Cst;
   }
 };
 
 /// \returns The splat index of a G_SHUFFLE_VECTOR \p MI when \p MI is a splat.
 /// If \p MI is not a splat, returns None.
 Optional<int> getSplatIndex(MachineInstr &MI);
 
-/// Returns a scalar constant of a G_BUILD_VECTOR splat if it exists.
-Optional<int64_t> getBuildVectorConstantSplat(const MachineInstr &MI,
-                                              const MachineRegisterInfo &MRI);
+/// \returns the scalar integral splat value of \p Reg if possible.
+Optional<APInt> getIConstantSplatVal(const Register Reg,
+                                     const MachineRegisterInfo &MRI);
+
+/// \returns the scalar integral splat value defined by \p MI if possible.
+Optional<APInt> getIConstantSplatVal(const MachineInstr &MI,
+                                     const MachineRegisterInfo &MRI);
+
+/// \returns the scalar sign extended integral splat value of \p Reg if
+/// possible.
+Optional<int64_t> getIConstantSplatSExtVal(const Register Reg,
+                                           const MachineRegisterInfo &MRI);
+
+/// \returns the scalar sign extended integral splat value defined by \p MI if
+/// possible.
+Optional<int64_t> getIConstantSplatSExtVal(const MachineInstr &MI,
 
 /// Returns a floating point scalar constant of a build vector splat if it
 /// exists. When \p AllowUndef == true some elements can be undef but not all.
 Optional<FPValueAndVReg> getFConstantSplat(Register VReg,
                                            const MachineRegisterInfo &MRI,
                                            bool AllowUndef = true);
 
 /// Return true if the specified register is defined by G_BUILD_VECTOR or

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

for (unsigned I = 0; I < RegsToNegate.size(); ++I) {
     }
   }
 
   // Now we know whether the comparisons are integer or floating point, check
   // the constant in the xor.
   int64_t Cst;
   if (Ty.isVector()) {
     MachineInstr *CstDef = MRI.getVRegDef(CstReg);
-    auto MaybeCst = getBuildVectorConstantSplat(*CstDef, MRI);
+    auto MaybeCst = getIConstantSplatSExtVal(*CstDef, MRI);
     if (!MaybeCst)
       return false;
     if (!isConstValidTrue(TLI, Ty.getScalarSizeInBits(), *MaybeCst, true, IsFP))
       return false;
   } else {
     if (!mi_match(CstReg, MRI, m_ICst(Cst)))
       return false;
     if (!isConstValidTrue(TLI, Ty.getSizeInBits(), Cst, false, IsFP))
bool CombinerHelper::matchOrShiftToFunnelShift(MachineInstr &MI,
   if (!mi_match(Dst, MRI,
                 // m_GOr() handles the commuted version as well.
                 m_GOr(m_GShl(m_Reg(ShlSrc), m_Reg(ShlAmt)),
                       m_GLShr(m_Reg(LShrSrc), m_Reg(LShrAmt)))))
     return false;
 
   // Given constants C0 and C1 such that C0 + C1 is bit-width:
   // (or (shl x, C0), (lshr y, C1)) -> (fshl x, y, C0) or (fshr x, y, C1)
-  // TODO: Match constant splat.
   int64_t CstShlAmt, CstLShrAmt;
-  if (mi_match(ShlAmt, MRI, m_ICst(CstShlAmt)) &&
-      mi_match(LShrAmt, MRI, m_ICst(CstLShrAmt)) &&
+  if (mi_match(ShlAmt, MRI, m_ICstOrSplat(CstShlAmt)) &&
+      mi_match(LShrAmt, MRI, m_ICstOrSplat(CstLShrAmt)) &&
       CstShlAmt + CstLShrAmt == BitWidth) {
     FshOpc = TargetOpcode::G_FSHR;
     Amt = LShrAmt;
 
   } else if (mi_match(LShrAmt, MRI,
                       m_GSub(m_SpecificICstOrSplat(BitWidth), m_Reg(Amt))) &&
              ShlAmt == Amt) {
     // (or (shl x, amt), (lshr y, (sub bw, amt))) -> (fshl x, y, amt)

llvm/lib/CodeGen/GlobalISel/Utils.cpp

 
 bool llvm::isBuildVectorConstantSplat(const MachineInstr &MI,
                                       const MachineRegisterInfo &MRI,
                                       int64_t SplatValue, bool AllowUndef) {
   return isBuildVectorConstantSplat(MI.getOperand(0).getReg(), MRI, SplatValue,
                                     AllowUndef);
 }
 
+Optional<APInt> llvm::getIConstantSplatVal(const Register Reg,
+                                           const MachineRegisterInfo &MRI) {
+  if (auto SplatValAndReg =
+          getAnyConstantSplat(Reg, MRI, /* AllowUndef */ false)) {
+    Optional<ValueAndVReg> ValAndVReg =
+        getIConstantVRegValWithLookThrough(SplatValAndReg->VReg, MRI);
+    return ValAndVReg->Value;
+  }
+
+  return None;
+}
+
+Optional<APInt> getIConstantSplatVal(const MachineInstr &MI,
+                                     const MachineRegisterInfo &MRI) {
+  return getIConstantSplatVal(MI.getOperand(0).getReg(), MRI);
+}
+
 Optional<int64_t>
-llvm::getBuildVectorConstantSplat(const MachineInstr &MI,
-                                  const MachineRegisterInfo &MRI) {
+llvm::getIConstantSplatSExtVal(const Register Reg,
   if (auto SplatValAndReg =
-          getAnyConstantSplat(MI.getOperand(0).getReg(), MRI, false))
+          getAnyConstantSplat(Reg, MRI, /* AllowUndef */ false))
     return getIConstantVRegSExtVal(SplatValAndReg->VReg, MRI);
   return None;
 }
 
+Optional<int64_t>
+llvm::getIConstantSplatSExtVal(const MachineInstr &MI,
+                               const MachineRegisterInfo &MRI) {
+  return getIConstantSplatSExtVal(MI.getOperand(0).getReg(), MRI);
+}
+
 Optional<FPValueAndVReg> llvm::getFConstantSplat(Register VReg,
                                                  const MachineRegisterInfo &MRI,
                                                  bool AllowUndef) {
   if (auto SplatValAndReg = getAnyConstantSplat(VReg, MRI, AllowUndef))
     return getFConstantVRegValWithLookThrough(SplatValAndReg->VReg, MRI);
   return None;
 }
 
bool llvm::isBuildVectorAllOnes(const MachineInstr &MI,
   return isBuildVectorConstantSplat(MI, MRI, -1, AllowUndef);
 }
 
 Optional<RegOrConstant> llvm::getVectorSplat(const MachineInstr &MI,
                                              const MachineRegisterInfo &MRI) {
   unsigned Opc = MI.getOpcode();
   if (!isBuildVectorOp(Opc))
     return None;
-  if (auto Splat = getBuildVectorConstantSplat(MI, MRI))
+  if (auto Splat = getIConstantSplatSExtVal(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 MaybeCst = getBuildVectorConstantSplat(MI, MRI);
+  auto MaybeCst = getIConstantSplatSExtVal(MI, MRI);
   if (!MaybeCst)
     return None;
   const unsigned ScalarSize = MRI.getType(Def).getScalarSizeInBits();
   return APInt(ScalarSize, *MaybeCst, true);
 }
 
 bool llvm::isNullOrNullSplat(const MachineInstr &MI,
                              const MachineRegisterInfo &MRI, bool AllowUndefs) {

llvm/test/CodeGen/AMDGPU/GlobalISel/combine-fsh.mir

bb.0:
     %amt1:_(s32) = G_CONSTANT i32 12
     %shl:_(s32) = G_SHL %a, %amt0
     %lshr:_(s32) = G_LSHR %b, %amt1
     %or:_(s32) = G_OR %shl, %lshr
     $vgpr2 = COPY %or
 ...
 
 ---
 name: fsh_v2i32_const
 tracksRegLiveness: true
 body: |
   bb.0:
     liveins: $vgpr0_vgpr1, $vgpr2_vgpr3, $vgpr4_vgpr5
 
     ; CHECK-LABEL: name: fsh_v2i32_const
     ; CHECK: liveins: $vgpr0_vgpr1, $vgpr2_vgpr3, $vgpr4_vgpr5
     ; CHECK-NEXT: {{  $}}
     ; CHECK-NEXT: %a:_(<2 x s32>) = COPY $vgpr0_vgpr1
     ; CHECK-NEXT: %b:_(<2 x s32>) = COPY $vgpr2_vgpr3
-    ; CHECK-NEXT: %scalar_amt0:_(s32) = G_CONSTANT i32 20
-    ; CHECK-NEXT: %amt0:_(<2 x s32>) = G_BUILD_VECTOR %scalar_amt0(s32), %scalar_amt0(s32)
     ; CHECK-NEXT: %scalar_amt1:_(s32) = G_CONSTANT i32 12
     ; CHECK-NEXT: %amt1:_(<2 x s32>) = G_BUILD_VECTOR %scalar_amt1(s32), %scalar_amt1(s32)
-    ; CHECK-NEXT: %shl:_(<2 x s32>) = G_SHL %a, %amt0(<2 x s32>)
-    ; CHECK-NEXT: %lshr:_(<2 x s32>) = G_LSHR %b, %amt1(<2 x s32>)
-    ; CHECK-NEXT: %or:_(<2 x s32>) = G_OR %shl, %lshr
+    ; CHECK-NEXT: %or:_(<2 x s32>) = G_FSHR %a, %b, %amt1(<2 x s32>)
     ; CHECK-NEXT: $vgpr4_vgpr5 = COPY %or(<2 x s32>)
     %a:_(<2 x s32>) = COPY $vgpr0_vgpr1
     %b:_(<2 x s32>) = COPY $vgpr2_vgpr3
     %scalar_amt0:_(s32) = G_CONSTANT i32 20
     %amt0:_(<2 x s32>) = G_BUILD_VECTOR %scalar_amt0, %scalar_amt0
     %scalar_amt1:_(s32) = G_CONSTANT i32 12
     %amt1:_(<2 x s32>) = G_BUILD_VECTOR %scalar_amt1, %scalar_amt1
     %shl:_(<2 x s32>) = G_SHL %a, %amt0

llvm/test/CodeGen/AMDGPU/GlobalISel/combine-rot.mir

bb.0:
     %amt1:_(s32) = G_CONSTANT i32 12
     %shl:_(s32) = G_SHL %a, %amt0
     %lshr:_(s32) = G_LSHR %a, %amt1
     %or:_(s32) = G_OR %shl, %lshr
     $vgpr1 = COPY %or
 ...
 
 ---
 name: rot_v2i32_const
 tracksRegLiveness: true
 body: |
   bb.0:
     liveins: $vgpr0_vgpr1, $vgpr2_vgpr3
 
     ; CHECK-LABEL: name: rot_v2i32_const
     ; CHECK: liveins: $vgpr0_vgpr1, $vgpr2_vgpr3
     ; CHECK-NEXT: {{  $}}
     ; CHECK-NEXT: %a:_(<2 x s32>) = COPY $vgpr0_vgpr1
-    ; CHECK-NEXT: %scalar_amt0:_(s32) = G_CONSTANT i32 20
-    ; CHECK-NEXT: %amt0:_(<2 x s32>) = G_BUILD_VECTOR %scalar_amt0(s32), %scalar_amt0(s32)
     ; CHECK-NEXT: %scalar_amt1:_(s32) = G_CONSTANT i32 12
     ; CHECK-NEXT: %amt1:_(<2 x s32>) = G_BUILD_VECTOR %scalar_amt1(s32), %scalar_amt1(s32)
-    ; CHECK-NEXT: %shl:_(<2 x s32>) = G_SHL %a, %amt0(<2 x s32>)
-    ; CHECK-NEXT: %lshr:_(<2 x s32>) = G_LSHR %a, %amt1(<2 x s32>)
-    ; CHECK-NEXT: %or:_(<2 x s32>) = G_OR %shl, %lshr
+    ; CHECK-NEXT: %or:_(<2 x s32>) = G_ROTR %a, %amt1(<2 x s32>)
     ; CHECK-NEXT: $vgpr2_vgpr3 = COPY %or(<2 x s32>)
     %a:_(<2 x s32>) = COPY $vgpr0_vgpr1
     %scalar_amt0:_(s32) = G_CONSTANT i32 20
     %amt0:_(<2 x s32>) = G_BUILD_VECTOR %scalar_amt0, %scalar_amt0
     %scalar_amt1:_(s32) = G_CONSTANT i32 12
     %amt1:_(<2 x s32>) = G_BUILD_VECTOR %scalar_amt1, %scalar_amt1
     %shl:_(<2 x s32>) = G_SHL %a, %amt0
     %lshr:_(<2 x s32>) = G_LSHR %a, %amt1

llvm/unittests/CodeGen/GlobalISel/PatternMatchTest.cpp

if (!TM)
     return;
   auto MIBCst = B.buildConstant(LLT::scalar(64), 42);
   Optional<ValueAndVReg> Src0;
   bool match = mi_match(MIBCst.getReg(0), *MRI, m_GCst(Src0));
   EXPECT_TRUE(match);
   EXPECT_EQ(Src0->VReg, MIBCst.getReg(0));
 }
 
+TEST_F(AArch64GISelMITest, MatchIntConstantSplat) {
+  setUp();
+  if (!TM)
+    return;
+
+  LLT s64 = LLT::scalar(64);
+  LLT v4s64 = LLT::fixed_vector(4, s64);
+
+  MachineInstrBuilder FortyTwoSplat =
+      B.buildSplatVector(v4s64, B.buildConstant(s64, 42));
+  int64_t Cst;
+  EXPECT_TRUE(mi_match(FortyTwoSplat.getReg(0), *MRI, m_ICstOrSplat(Cst)));
+  EXPECT_EQ(Cst, 42);
+
+  MachineInstrBuilder NonConstantSplat =
+      B.buildBuildVector(v4s64, {Copies[0], Copies[0], Copies[0], Copies[0]});
+  EXPECT_FALSE(mi_match(NonConstantSplat.getReg(0), *MRI, m_ICstOrSplat(Cst)));
+}
+
 TEST_F(AArch64GISelMITest, MachineInstrPtrBind) {
   setUp();
   if (!TM)
     return;
   auto MIBAdd = B.buildAdd(LLT::scalar(64), Copies[0], Copies[1]);
   // Test 'MachineInstr *' bind.
   // Default mi_match.
   MachineInstr *MIPtr = MIBAdd.getInstr();