[GlobalISel] Detect splats built with G_CONCAT_VECTORS

[rovka]

Summary

Add support to the MI matching of vector splats for patterns that
consist of G_CONCAT_VECTORS of smaller splats with the same constant
value. With this, we would consider the following pseudo-MIR to be a splat:

%0 = G_[F]CONSTANT [...]
%1 = G_BUILD_VECTOR %0, %0, ..., %0
%2 = G_CONCAT_VECTORS %1, %1, ..., %1

Since it uses recursion for matching splats, it could match pretty
complicated patterns with all sorts of combinations of G_BUILD_VECTOR
and G_CONCAT_VECTORS (e.g. a G_CONCAT_VECTORS with
a G_BUILD_VECTOR_TRUNC and another G_CONCAT_VECTORS as operands),
and it should also look through copies etc.

This should make it easier to match complex immediates for certain
instructions on AMDGPU, where for instance a <8 x s16> will be split
before instruction selection into a G_CONCAT_VECTORS of <2 x s16>
splats.

[foad]

This should make it easier to match complex immediates for certain
instructions on AMDGPU, where for instance a <8 x s16> will be split
before instruction selection into a G_CONCAT_VECTORS of <2 x s16>
splats.

I don't quite understand when this would be useful. If <8 x s16> is split up into <2 x s16> pieces during legalization, because only the latter is legal, then why would we still have a G_CONCAT_VECTORS with result type <8 x s16>? What operations would use that result? Do you have an example in mind?

[rovka]

In D141902#4058303, @foad wrote:

This should make it easier to match complex immediates for certain
instructions on AMDGPU, where for instance a <8 x s16> will be split
before instruction selection into a G_CONCAT_VECTORS of <2 x s16>
splats.

I don't quite understand when this would be useful. If <8 x s16> is split up into <2 x s16> pieces during legalization, because only the latter is legal, then why would we still have a G_CONCAT_VECTORS with result type <8 x s16>? What operations would use that result? Do you have an example in mind?

There are some intrinsics that use large vector types, e.g. v16f16.

[arsenm]

This suggests to me we're missing some useful vector combines, but I don't see a problem having this handle this case also

llvm/unittests/CodeGen/GlobalISel/PatternMatchTest.cpp

810

Could use a case with some undefs in it (undef simple element and undef vector component)

[rovka]

llvm/unittests/CodeGen/GlobalISel/PatternMatchTest.cpp

810

Will do, thanks.

[rovka]

Add tests with undef

llvm/lib/CodeGen/GlobalISel/Utils.cpp

 
 std::optional<ValueAndVReg> getAnyConstantSplat(Register VReg,
                                                 const MachineRegisterInfo &MRI,
                                                 bool AllowUndef) {
   MachineInstr *MI = getDefIgnoringCopies(VReg, MRI);
   if (!MI)
     return std::nullopt;
 
-  if (!isBuildVectorOp(MI->getOpcode()))
+  bool isConcatVectorsOp = MI->getOpcode() == TargetOpcode::G_CONCAT_VECTORS;
+  if (!isBuildVectorOp(MI->getOpcode()) && !isConcatVectorsOp)
     return std::nullopt;
 
   std::optional<ValueAndVReg> SplatValAndReg;
   for (MachineOperand &Op : MI->uses()) {
     Register Element = Op.getReg();
+    // If we have a G_CONCAT_VECTOR, we recursively look into the
+    // vectors that we're concatenating to see if they're splats.
     auto ElementValAndReg =
-        getAnyConstantVRegValWithLookThrough(Element, MRI, true, true);
+        isConcatVectorsOp
+            ? getAnyConstantSplat(Element, MRI, AllowUndef)
+            : 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 std::nullopt;
     }
 

llvm/unittests/CodeGen/GlobalISel/PatternMatchTest.cpp

 }
 
 TEST_F(AArch64GISelMITest, MatchIntConstantSplat) {
   setUp();
   if (!TM)
     return;
 
   LLT s64 = LLT::scalar(64);
+  LLT v2s64 = LLT::fixed_vector(2, s64);
   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)));
+
+  auto ICst = B.buildConstant(s64, 15).getReg(0);
+  auto SmallSplat = B.buildBuildVector(v2s64, {ICst, ICst}).getReg(0);
+  auto LargeSplat = B.buildConcatVectors(v4s64, {SmallSplat, SmallSplat});
+  EXPECT_TRUE(mi_match(LargeSplat.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.

 }
 
 TEST_F(AArch64GISelMITest, MatchConstantSplat) {
   setUp();
   if (!TM)
     return;
 
   LLT s64 = LLT::scalar(64);
+  LLT v2s64 = LLT::fixed_vector(2, 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);
   std::optional<FPValueAndVReg> FValReg;
 
   // GFCstOrSplatGFCstMatch allows undef as part of splat. Undef often comes
TEST_F(AArch64GISelMITest, MatchConstantSplat) {
   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)));
+
+  // Look through G_CONCAT_VECTORS.
+  auto SmallZeroSplat = B.buildBuildVector(v2s64, {FPZero, FPZero}).getReg(0);
+  auto LargeZeroSplat =
+      B.buildConcatVectors(v4s64, {SmallZeroSplat, SmallZeroSplat});
+  EXPECT_TRUE(mi_match(LargeZeroSplat.getReg(0), *MRI,
+                       GFCstOrSplatGFCstMatch(FValReg)));
+
+  auto SmallZeroSplat2 = B.buildBuildVector(v2s64, {FPZero, FPZero}).getReg(0);
+  auto SmallZeroSplat3 = B.buildCopy(v2s64, SmallZeroSplat).getReg(0);
+  auto LargeZeroSplat2 =
+      B.buildConcatVectors(v4s64, {SmallZeroSplat2, SmallZeroSplat3});
+  EXPECT_TRUE(mi_match(LargeZeroSplat2.getReg(0), *MRI,
+                       GFCstOrSplatGFCstMatch(FValReg)));
+
+  // Not all G_CONCAT_VECTORS are splats.
+  auto SmallOneSplat = B.buildBuildVector(v2s64, {FPOne, FPOne}).getReg(0);
+  auto LargeMixedSplat =
+      B.buildConcatVectors(v4s64, {SmallZeroSplat, SmallOneSplat});
+  EXPECT_FALSE(mi_match(LargeMixedSplat.getReg(0), *MRI,
+                        GFCstOrSplatGFCstMatch(FValReg)));
+
+  auto SmallMixedSplat = B.buildBuildVector(v2s64, {FPOne, FPZero}).getReg(0);
+  auto LargeSplat =
+      B.buildConcatVectors(v4s64, {SmallMixedSplat, SmallMixedSplat});
+  EXPECT_FALSE(
+      mi_match(LargeSplat.getReg(0), *MRI, GFCstOrSplatGFCstMatch(FValReg)));
+
+  auto SmallUndefSplat = B.buildBuildVector(v2s64, {Undef, Undef}).getReg(0);
+  auto LargeUndefSplat =
+      B.buildConcatVectors(v4s64, {SmallUndefSplat, SmallUndefSplat});
+  EXPECT_FALSE(mi_match(LargeUndefSplat.getReg(0), *MRI,
+                        GFCstOrSplatGFCstMatch(FValReg)));
+
+  auto UndefVec = B.buildUndef(v2s64).getReg(0);
+  auto LargeUndefSplat2 = B.buildConcatVectors(v4s64, {UndefVec, UndefVec});
+  EXPECT_FALSE(mi_match(LargeUndefSplat2.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]);