GlobalISel/Utils: Refactor constant splat match functions

[Petar.Avramovic]

Summary

Add generic helper function that matches constant splat. It also has an option
to match constant splat with undef (some elements can be undef but not all).
Add util function and matcher for G_FCONSTANT splat.

[foad]

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

362

Would a return type of Optional<ValueAndVReg> (like getConstantVRegValWithLookThrough) be more useful? Actually I'm not sure if returning the register is useful, maybe just Optional<APInt>?

llvm/lib/CodeGen/GlobalISel/Utils.cpp

1022

Isn't there an existing utility you can call to do this?

[Petar.Avramovic]

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

362

ValueAndVReg would be fine but in later patch APFloat would be useful. I did it with reg so that once it is matched caller knows if it matched int or float constant and can get APInt or APFloat and check what it needs.

llvm/lib/CodeGen/GlobalISel/Utils.cpp

1022

There is code that does similar thing with look through in a few places but not a utility function.

[paquette]

llvm/lib/CodeGen/GlobalISel/Utils.cpp

1019

This looks really similar to getBuildVectorConstantSplat. Would it make sense to improve that function instead?

1022

getDefIgnoringCopies doesn't work here?

[Petar.Avramovic]

Make use of the more generic helper function in existing splat matcher functions.

[paquette]

llvm/lib/CodeGen/GlobalISel/Utils.cpp

374

This will need a nullptr check.

I'm not sure about how I feel about adding functions for generic opcode checks, since I think it could end up being confusing if we don't add them for every opcode. :/

That being said, these do look nicer than using MI->getOpcode() everywhere.

(Maybe someone else can give a stronger opinion here.)

[aemerson]

llvm/lib/CodeGen/GlobalISel/Utils.cpp

374

I don't think alone this helper has much benefit. I think it would be nice to have some wrapper classes around MachineInstr to do some of that while also proving some abstractions over operand accessors etc. I haven't got around to it but was on my todo list.

[Petar.Avramovic]

llvm/lib/CodeGen/GlobalISel/Utils.cpp

374

I meant to add it for a number of opcodes (~50).
I did grep of '==' and '!=' compares with 'TargetOpcode::G_' Here are some stats:
there are 44 different opcodes with 1 or 2 mentions, in total 58 compares.
there are 49 different opcodes with more that 3 mentions (G_CONSTANT has the most; 30). In total 450 compares
190 are asserts.
Would result in nicer code formatting.
I will directly check opcode for this patch.

374

I like this better than util function.

[Petar.Avramovic]

Remove isImplicitDef.

[arsenm]

llvm/lib/CodeGen/GlobalISel/Utils.cpp

948

Use mi_match?

957

Typo contant

958

• instead of ->Value?

[Petar.Avramovic]

llvm/lib/CodeGen/GlobalISel/Utils.cpp

948

There is nothing to match for G_IMPLICIT_DEF only opcode, but there is this new GenericMachineInstr wrapper around MachineInstr which does opcode checks.

958

I dind't get this. What should be the change again?

[Petar.Avramovic]

Ping.

[foad]

Generally looks OK to me. Any more comments @paquette?

llvm/lib/CodeGen/GlobalISel/Utils.cpp

988

What does "Allows nan splat matching" mean?

[Petar.Avramovic]

llvm/lib/CodeGen/GlobalISel/Utils.cpp

988

Splat of NaN value, if we get value as APFloat all comparisons with it are false and we cant check for splat. We check for splat using APInt and check if constant is g_fconstant with getFConstantVRegValWithLookThrough and get APFloat.

[foad]

llvm/lib/CodeGen/GlobalISel/Utils.cpp

988

OK, I just assumed that any floating point splat detection code would use bitwiseIsEqual instead of operator== anyway.

[Petar.Avramovic]

Remove comment about constant splat with NaN value. Seems unnecessary to mention it since we match all constants.

[paquette]

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

111

Could this function be a one-liner?

return getFConstantSplat(Reg, MRI) || getFConstantVRegValWithLookThrough(Reg, MRI);

llvm/lib/CodeGen/GlobalISel/Utils.cpp

933

check that this is not nullptr?

939

Maybe use

for (MachineOperand &Op : MI->uses()) {
   ...

[Petar.Avramovic]

Addressed review comments.

[paquette]

LGTM

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

public:
   Register getBaseReg() const { return getReg(1); }
   Register getOffsetReg() const { return getReg(2); }
 
   static bool classof(const MachineInstr *MI) {
     return MI->getOpcode() == TargetOpcode::G_PTR_ADD;
   }
 };
 
+/// Represents a G_IMPLICIT_DEF.
+class GImplicitDef : public GenericMachineInstr {
+public:
+  static bool classof(const MachineInstr *MI) {
+    return MI->getOpcode() == TargetOpcode::G_IMPLICIT_DEF;
+  }
+};
+
 } // namespace llvm
 
 #endif // LLVM_CODEGEN_GLOBALISEL_GENERICMACHINEINSTRS_H
 No newline at end of file

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

bool match(const MachineRegisterInfo &MRI, Register Reg) {
     return FPValReg ? true : false;
   }
 };
 
 inline GFCstAndRegMatch m_GFCst(Optional<FPValueAndVReg> &FPValReg) {
   return GFCstAndRegMatch(FPValReg);
 }
 
+struct GFCstOrSplatGFCstMatch {
+  Optional<FPValueAndVReg> &FPValReg;
+  GFCstOrSplatGFCstMatch(Optional<FPValueAndVReg> &FPValReg)
+      : FPValReg(FPValReg) {}
+  bool match(const MachineRegisterInfo &MRI, Register Reg) {
+    return (FPValReg = getFConstantSplat(Reg, MRI)) ||
+           (FPValReg = getFConstantVRegValWithLookThrough(Reg, MRI));
+  };
+};
+
+inline GFCstOrSplatGFCstMatch
+m_GFCstOrSplat(Optional<FPValueAndVReg> &FPValReg) {
+  return GFCstOrSplatGFCstMatch(FPValReg);
+}
+
 /// Matcher for a specific constant value.
 struct SpecificConstantMatch {
   int64_t RequestedVal;
   SpecificConstantMatch(int64_t RequestedVal) : RequestedVal(RequestedVal) {}
   bool match(const MachineRegisterInfo &MRI, Register Reg) {
     int64_t MatchedVal;
     return mi_match(Reg, MRI, m_ICst(MatchedVal)) && MatchedVal == RequestedVal;
   }

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

 /// \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 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 instruction is a G_BUILD_VECTOR or
 /// G_BUILD_VECTOR_TRUNC where all of the elements are 0 or undef.
 bool isBuildVectorAllZeros(const MachineInstr &MI,
-                           const MachineRegisterInfo &MRI);
+                           const MachineRegisterInfo &MRI,
+                           bool AllowUndef = false);
 
 /// Return true if the specified instruction is a G_BUILD_VECTOR or
 /// G_BUILD_VECTOR_TRUNC where all of the elements are ~0 or undef.
 bool isBuildVectorAllOnes(const MachineInstr &MI,
-                          const MachineRegisterInfo &MRI);
+                          const MachineRegisterInfo &MRI,
+                          bool AllowUndef = false);
 
 /// \returns a value when \p MI is a vector splat. The splat can be either a
 /// Register or a constant.
 ///
 /// Examples:
 ///
 /// \code
 ///   %reg = COPY $physreg

llvm/lib/CodeGen/GlobalISel/Utils.cpp

 //===----------------------------------------------------------------------===//
 
 #include "llvm/CodeGen/GlobalISel/Utils.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/CodeGen/GlobalISel/GISelChangeObserver.h"
 #include "llvm/CodeGen/GlobalISel/GISelKnownBits.h"
+#include "llvm/CodeGen/GlobalISel/GenericMachineInstrs.h"
 #include "llvm/CodeGen/GlobalISel/MIPatternMatch.h"
 #include "llvm/CodeGen/GlobalISel/RegisterBankInfo.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineOptimizationRemarkEmitter.h"
 #include "llvm/CodeGen/MachineSizeOpts.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/StackProtector.h"

 
 bool isAnyConstant(const MachineInstr *MI) {
   if (!MI)
     return false;
   unsigned Opc = MI->getOpcode();
   return Opc == TargetOpcode::G_CONSTANT || Opc == TargetOpcode::G_FCONSTANT;
 }
 
 Optional<APInt> getCImmAsAPInt(const MachineInstr *MI) {
   const MachineOperand &CstVal = MI->getOperand(1);
   if (CstVal.isCImm())
     return CstVal.getCImm()->getValue();
   return None;
 }
 
 Optional<APInt> getCImmOrFPImmAsAPInt(const MachineInstr *MI) {
   const MachineOperand &CstVal = MI->getOperand(1);
Optional<int> llvm::getSplatIndex(MachineInstr &MI) {
   return SplatValue;
 }
 
 static bool isBuildVectorOp(unsigned Opcode) {
   return Opcode == TargetOpcode::G_BUILD_VECTOR ||
          Opcode == TargetOpcode::G_BUILD_VECTOR_TRUNC;
 }
 
-// TODO: Handle mixed undef elements.
-static bool isBuildVectorConstantSplat(const MachineInstr &MI,
+namespace {
+
+Optional<ValueAndVReg> getAnyConstantSplat(Register VReg,
                                            const MachineRegisterInfo &MRI,
-                                       int64_t SplatValue) {
-  if (!isBuildVectorOp(MI.getOpcode()))
-    return false;
+                                           bool AllowUndef) {
+  MachineInstr *MI = getDefIgnoringCopies(VReg, MRI);
+  if (!MI)
+    return None;
 
-  const unsigned NumOps = MI.getNumOperands();
-  for (unsigned I = 1; I != NumOps; ++I) {
-    Register Element = MI.getOperand(I).getReg();
-    if (!mi_match(Element, MRI, m_SpecificICst(SplatValue)))
-      return false;
+  if (!isBuildVectorOp(MI->getOpcode()))
+    return None;
+
+  Optional<ValueAndVReg> SplatValAndReg = None;
+  for (MachineOperand &Op : MI->uses()) {
+    Register Element = Op.getReg();
+    auto ElementValAndReg =
+        getAnyConstantVRegValWithLookThrough(Element, MRI, true, true);
+
+    // If AllowUndef, treat undef as value that will result in a constant splat.
+    if (!ElementValAndReg) {
+      if (AllowUndef && isa<GImplicitDef>(MRI.getVRegDef(Element)))
+        continue;
+      return None;
     }
 
-  return true;
+    // Record splat value
+    if (!SplatValAndReg)
+      SplatValAndReg = ElementValAndReg;
+
+    // Different constant then the one already recorded, not a constant splat.
+    if (SplatValAndReg->Value != ElementValAndReg->Value)
+      return None;
+  }
+
+  return SplatValAndReg;
 }
 
+bool isBuildVectorConstantSplat(const MachineInstr &MI,
+                                const MachineRegisterInfo &MRI,
+                                int64_t SplatValue, bool AllowUndef) {
+  if (auto SplatValAndReg =
+          getAnyConstantSplat(MI.getOperand(0).getReg(), MRI, AllowUndef))
+    return mi_match(SplatValAndReg->VReg, MRI, m_SpecificICst(SplatValue));
+  return false;
+}
+
+} // end anonymous namespace
+
 Optional<int64_t>
 llvm::getBuildVectorConstantSplat(const MachineInstr &MI,
                                   const MachineRegisterInfo &MRI) {
-  if (!isBuildVectorOp(MI.getOpcode()))
-    return None;
-
-  const unsigned NumOps = MI.getNumOperands();
-  Optional<int64_t> Scalar;
-  for (unsigned I = 1; I != NumOps; ++I) {
-    Register Element = MI.getOperand(I).getReg();
-    int64_t ElementValue;
-    if (!mi_match(Element, MRI, m_ICst(ElementValue)))
-      return None;
-    if (!Scalar)
-      Scalar = ElementValue;
-    else if (*Scalar != ElementValue)
-      return None;
-  }
+  if (auto SplatValAndReg =
+          getAnyConstantSplat(MI.getOperand(0).getReg(), MRI, false))
+    return getIConstantVRegSExtVal(SplatValAndReg->VReg, MRI);
 
-  return Scalar;
+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::isBuildVectorAllZeros(const MachineInstr &MI,
-                                 const MachineRegisterInfo &MRI) {
-  return isBuildVectorConstantSplat(MI, MRI, 0);
+                                 const MachineRegisterInfo &MRI,
+                                 bool AllowUndef) {
+  return isBuildVectorConstantSplat(MI, MRI, 0, AllowUndef);
 }
 
 bool llvm::isBuildVectorAllOnes(const MachineInstr &MI,
-                                const MachineRegisterInfo &MRI) {
-  return isBuildVectorConstantSplat(MI, MRI, -1);
+                                const MachineRegisterInfo &MRI,
+                                bool AllowUndef) {
+  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))
     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);
 }
 
 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);
 
   // TODO: Also handle fconstant
   if (Def->getOpcode() == TargetOpcode::G_CONSTANT)
     return Match(Def->getOperand(1).getCImm());
 

llvm/unittests/CodeGen/GlobalISel/PatternMatchTest.cpp

TEST_F(AArch64GISelMITest, MatchFPOrIntConst) {
   EXPECT_EQ(IntOne, ValReg->VReg);
   EXPECT_FALSE(mi_match(IntOne, *MRI, m_GFCst(FValReg)));
 
   EXPECT_FALSE(mi_match(FPOne, *MRI, m_GCst(ValReg)));
   EXPECT_TRUE(mi_match(FPOne, *MRI, m_GFCst(FValReg)));
   EXPECT_EQ(FPOne, FValReg->VReg);
 }
 
+TEST_F(AArch64GISelMITest, MatchConstantSplat) {
+  setUp();
+  if (!TM)
+    return;
+
+  LLT s64 = LLT::scalar(64);
+  LLT v4s64 = LLT::fixed_vector(4, 64);
+
+  Register FPOne = B.buildFConstant(s64, 1.0).getReg(0);
+  Register FPZero = B.buildFConstant(s64, 0.0).getReg(0);
+  Register Undef = B.buildUndef(s64).getReg(0);
+  Optional<FPValueAndVReg> FValReg;
+
+  // GFCstOrSplatGFCstMatch allows undef as part of splat. Undef often comes
+  // from padding to legalize into available operation and then ignore added
+  // elements e.g. v3s64 to v4s64.
+
+  EXPECT_TRUE(mi_match(FPZero, *MRI, GFCstOrSplatGFCstMatch(FValReg)));
+  EXPECT_EQ(FPZero, FValReg->VReg);
+
+  EXPECT_FALSE(mi_match(Undef, *MRI, GFCstOrSplatGFCstMatch(FValReg)));
+
+  auto ZeroSplat = B.buildBuildVector(v4s64, {FPZero, FPZero, FPZero, FPZero});
+  EXPECT_TRUE(
+      mi_match(ZeroSplat.getReg(0), *MRI, GFCstOrSplatGFCstMatch(FValReg)));
+  EXPECT_EQ(FPZero, FValReg->VReg);
+
+  auto ZeroUndef = B.buildBuildVector(v4s64, {FPZero, FPZero, FPZero, Undef});
+  EXPECT_TRUE(
+      mi_match(ZeroUndef.getReg(0), *MRI, GFCstOrSplatGFCstMatch(FValReg)));
+  EXPECT_EQ(FPZero, FValReg->VReg);
+
+  // All undefs are not constant splat.
+  auto UndefSplat = B.buildBuildVector(v4s64, {Undef, Undef, Undef, Undef});
+  EXPECT_FALSE(
+      mi_match(UndefSplat.getReg(0), *MRI, GFCstOrSplatGFCstMatch(FValReg)));
+
+  auto ZeroOne = B.buildBuildVector(v4s64, {FPZero, FPZero, FPZero, FPOne});
+  EXPECT_FALSE(
+      mi_match(ZeroOne.getReg(0), *MRI, GFCstOrSplatGFCstMatch(FValReg)));
+
+  auto NonConstantSplat =
+      B.buildBuildVector(v4s64, {Copies[0], Copies[0], Copies[0], Copies[0]});
+  EXPECT_FALSE(mi_match(NonConstantSplat.getReg(0), *MRI,
+                        GFCstOrSplatGFCstMatch(FValReg)));
+
+  auto Mixed = B.buildBuildVector(v4s64, {FPZero, FPZero, FPZero, Copies[0]});
+  EXPECT_FALSE(
+      mi_match(Mixed.getReg(0), *MRI, GFCstOrSplatGFCstMatch(FValReg)));
+}
+
 TEST_F(AArch64GISelMITest, MatchNeg) {
   setUp();
   if (!TM)
     return;
 
   LLT s64 = LLT::scalar(64);
   auto Zero = B.buildConstant(LLT::scalar(64), 0);
   auto NegInst = B.buildSub(s64, Zero, Copies[0]);