[AIE2][AIE2P] Allow single splat in symmetric build vector combine

Author: abhinay-anubola

Diff for: llvm/lib/Target/AIE/AIECombinerHelper.cpp

@@ -1228,6 +1228,40 @@ void llvm::applyExtractVecEltAndExt(
   MatchMI->eraseFromParent();
 }
 
+static std::optional<Register>
+getSplatVectorSrcReg(const MachineInstr &MI, const MachineRegisterInfo &MRI,
+                     std::pair<unsigned, unsigned> Range) {
+  auto IsUndef = [&](const MachineOperand &Op) {
+    const MachineInstr *Undef = MRI.getVRegDef(Op.getReg());
+    return Undef && Undef->getOpcode() == TargetOpcode::G_IMPLICIT_DEF;
+  };
+  const unsigned Start = Range.first;
+  const unsigned End = Range.second;
+  // First non-undef operand.
+  Register SrcReg = 0;
+  bool FoundSrc = false;
+  bool AllUndef = true;
+
+  // Find the first non-undef operand as the reference.
+  for (unsigned I = Start; I < End; I++) {
+    const MachineOperand &Op = MI.getOperand(I);
+    if (!IsUndef(Op)) {
+      if (!FoundSrc) {
+        SrcReg = Op.getReg();
+        FoundSrc = true;
+      } else if (Op.getReg() != SrcReg) {
+        return std::nullopt;
+      }
+      AllUndef = false;
+    }
+  }
+
+  if (AllUndef)
+    SrcReg = MI.getOperand(1).getReg();
+
+  return SrcReg;
+}
+
 // Match something like:
 // %0:_(<32 x s16>) = G_BUILD_VECTOR %1:_(s16), ... x32
 //
@@ -1255,34 +1289,12 @@ bool llvm::matchSplatVector(MachineInstr &MI, MachineRegisterInfo &MRI,
     return false;
   }
 
-  auto IsUndef = [&](const MachineOperand &Op) {
-    const MachineInstr *Undef = MRI.getVRegDef(Op.getReg());
-    return Undef && Undef->getOpcode() == TargetOpcode::G_IMPLICIT_DEF;
-  };
   const unsigned NumOps = MI.getNumOperands();
-  // First non-undef operand.
-  unsigned SrcReg = 0;
-  bool FoundSrc = false;
-  bool AllUndef = true;
-
-  // Find the first non-undef operand as the reference.
-  for (unsigned I = 1; I < NumOps; I++) {
-    const MachineOperand &Op = MI.getOperand(I);
-    if (!IsUndef(Op)) {
-      if (!FoundSrc) {
-        SrcReg = Op.getReg();
-        FoundSrc = true;
-      } else if (Op.getReg() != SrcReg) {
-        return false;
-      }
-      AllUndef = false;
-    }
-  }
-
-  if (AllUndef)
-    SrcReg = MI.getOperand(1).getReg();
+  auto SrcReg = getSplatVectorSrcReg(MI, MRI, std::make_pair(1, NumOps));
+  if (!SrcReg)
+    return false;
 
-  MatchInfo = {DstVecReg, SrcReg};
+  MatchInfo = {DstVecReg, *SrcReg};
   return true;
 }
 
@@ -1466,6 +1478,7 @@ bool llvm::applySingleDiffLaneBuildVector(
 // %0:_(<32 x s16>) = G_CONCAT_VECTORS %3:_(<16 x s16>), %4:_(<16 x s16>)
 // These sub-G_BUILD_VECTOR instructions may later be combined into broadcast
 // instructions by combine_splat_vector.
+// TODO: Remove the original splat vector match and implement the same here.
 bool llvm::matchSymmetricBuildVector(MachineInstr &MI, MachineRegisterInfo &MRI,
                                      GISelChangeObserver &Observer,
                                      BuildFnTy &MatchInfo) {
@@ -1487,30 +1500,16 @@ bool llvm::matchSymmetricBuildVector(MachineInstr &MI, MachineRegisterInfo &MRI,
     return false;
   }
 
-  auto getSrcOperand = [&](unsigned I) { return MI.getOperand(I + 1); };
-
-  const unsigned NumElts = MI.getNumOperands() - 1;
-  const unsigned HalfNumElts = NumElts / 2;
-  const MachineOperand FirstOp = getSrcOperand(0);
-  const MachineOperand SecondOp = getSrcOperand(HalfNumElts);
-
-  // Ensures that each operand in the first half matches FirstOp, and each
-  // operand in the second half matches SecondOp.
-  for (unsigned i = 0; i < HalfNumElts; i++) {
-    if (!getSrcOperand(i).isIdenticalTo(FirstOp)) {
-      return false;
-    }
-    if (!getSrcOperand(HalfNumElts + i).isIdenticalTo(SecondOp)) {
-      return false;
-    }
-  }
-
-  // If both halves are the same register, it's effectively a splat, and the
-  // splat vector combine already handles that case.
-  if (FirstOp.isIdenticalTo(SecondOp))
-    return false;
+  // TODO: Split the G_BUILD_VECTOR either into 3/4 and 1/4 parts,
+  // or 1/4 and 3/4 parts, and then check if any part qualifies as a splat.
+  const unsigned NumOps = MI.getNumOperands();
+  const unsigned HalfNumElts = NumOps / 2 + 1;
+  auto FirstHalfSrcReg =
+      getSplatVectorSrcReg(MI, MRI, std::make_pair(1, HalfNumElts));
+  auto SecondHalfSrcReg =
+      getSplatVectorSrcReg(MI, MRI, std::make_pair(HalfNumElts, NumOps));
 
-  MatchInfo = [&MI, &Observer, &MRI, DstVecTy](MachineIRBuilder &B) {
+  MatchInfo = [&MI, &Observer, DstVecTy](MachineIRBuilder &B) {
     B.setInstrAndDebugLoc(MI);
     LegalizerHelper Helper(B.getMF(), Observer, B);
     // Splits the G_BUILD_VECTOR into two half-sized G_BUILD_VECTOR operations
@@ -1520,7 +1519,8 @@ bool llvm::matchSymmetricBuildVector(MachineInstr &MI, MachineRegisterInfo &MRI,
         DstVecTy.changeElementCount(
             DstVecTy.getElementCount().divideCoefficientBy(2)));
   };
-  return true;
+
+  return (FirstHalfSrcReg.has_value() || SecondHalfSrcReg.has_value());
 }
 
 // Match something like:

Diff for: llvm/test/CodeGen/AIE/GlobalISel/prelegalizercombiner-pad-vector.mir

@@ -50,8 +50,11 @@ body:             |
     ; CHECK: [[COPY:%[0-9]+]]:_(<4 x s32>) = COPY $q0
     ; CHECK-NEXT: [[UV:%[0-9]+]]:_(s32), [[UV1:%[0-9]+]]:_(s32), [[UV2:%[0-9]+]]:_(s32), [[UV3:%[0-9]+]]:_(s32) = G_UNMERGE_VALUES [[COPY]](<4 x s32>)
     ; CHECK-NEXT: [[DEF:%[0-9]+]]:_(s32) = G_IMPLICIT_DEF
-    ; CHECK-NEXT: [[BUILD_VECTOR:%[0-9]+]]:_(<8 x s32>) = G_BUILD_VECTOR [[UV]](s32), [[UV2]](s32), [[UV1]](s32), [[UV3]](s32), [[DEF]](s32), [[DEF]](s32), [[DEF]](s32), [[DEF]](s32)
-    ; CHECK-NEXT: $wl0 = COPY [[BUILD_VECTOR]](<8 x s32>)
+    ; CHECK-NEXT: [[BUILD_VECTOR:%[0-9]+]]:_(<4 x s32>) = G_BUILD_VECTOR [[UV]](s32), [[UV2]](s32), [[UV1]](s32), [[UV3]](s32)
+    ; CHECK-NEXT: [[AIE_BROADCAST_VECTOR:%[0-9]+]]:_(<16 x s32>) = G_AIE_BROADCAST_VECTOR [[DEF]](s32)
+    ; CHECK-NEXT: [[AIE_UNPAD_VECTOR:%[0-9]+]]:_(<4 x s32>) = G_AIE_UNPAD_VECTOR [[AIE_BROADCAST_VECTOR]](<16 x s32>)
+    ; CHECK-NEXT: [[CONCAT_VECTORS:%[0-9]+]]:_(<8 x s32>) = G_CONCAT_VECTORS [[BUILD_VECTOR]](<4 x s32>), [[AIE_UNPAD_VECTOR]](<4 x s32>)
+    ; CHECK-NEXT: $wl0 = COPY [[CONCAT_VECTORS]](<8 x s32>)
       %10:_(<4 x s32>) = COPY $q0
       %0:_(s32), %1:_(s32), %2:_(s32), %3:_(s32) = G_UNMERGE_VALUES %10
       %4:_(s32) = G_IMPLICIT_DEF
@@ -68,8 +71,11 @@ body:             |
     ; CHECK: [[COPY:%[0-9]+]]:_(<4 x s32>) = COPY $q0
     ; CHECK-NEXT: [[UV:%[0-9]+]]:_(s32), [[UV1:%[0-9]+]]:_(s32), [[UV2:%[0-9]+]]:_(s32), [[UV3:%[0-9]+]]:_(s32) = G_UNMERGE_VALUES [[COPY]](<4 x s32>)
     ; CHECK-NEXT: [[DEF:%[0-9]+]]:_(s32) = G_IMPLICIT_DEF
-    ; CHECK-NEXT: [[BUILD_VECTOR:%[0-9]+]]:_(<8 x s32>) = G_BUILD_VECTOR [[UV]](s32), [[UV1]](s32), [[UV2]](s32), [[DEF]](s32), [[DEF]](s32), [[DEF]](s32), [[DEF]](s32), [[DEF]](s32)
-    ; CHECK-NEXT: $wl0 = COPY [[BUILD_VECTOR]](<8 x s32>)
+    ; CHECK-NEXT: [[BUILD_VECTOR:%[0-9]+]]:_(<4 x s32>) = G_BUILD_VECTOR [[UV]](s32), [[UV1]](s32), [[UV2]](s32), [[DEF]](s32)
+    ; CHECK-NEXT: [[AIE_BROADCAST_VECTOR:%[0-9]+]]:_(<16 x s32>) = G_AIE_BROADCAST_VECTOR [[DEF]](s32)
+    ; CHECK-NEXT: [[AIE_UNPAD_VECTOR:%[0-9]+]]:_(<4 x s32>) = G_AIE_UNPAD_VECTOR [[AIE_BROADCAST_VECTOR]](<16 x s32>)
+    ; CHECK-NEXT: [[CONCAT_VECTORS:%[0-9]+]]:_(<8 x s32>) = G_CONCAT_VECTORS [[BUILD_VECTOR]](<4 x s32>), [[AIE_UNPAD_VECTOR]](<4 x s32>)
+    ; CHECK-NEXT: $wl0 = COPY [[CONCAT_VECTORS]](<8 x s32>)
       %10:_(<4 x s32>) = COPY $q0
       %0:_(s32), %1:_(s32), %2:_(s32), %3:_(s32) = G_UNMERGE_VALUES %10
       %4:_(s32) = G_IMPLICIT_DEF
@@ -86,8 +92,11 @@ body:             |
     ; CHECK: [[COPY:%[0-9]+]]:_(<4 x s32>) = COPY $q0
     ; CHECK-NEXT: [[UV:%[0-9]+]]:_(s32), [[UV1:%[0-9]+]]:_(s32), [[UV2:%[0-9]+]]:_(s32), [[UV3:%[0-9]+]]:_(s32) = G_UNMERGE_VALUES [[COPY]](<4 x s32>)
     ; CHECK-NEXT: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 1
-    ; CHECK-NEXT: [[BUILD_VECTOR:%[0-9]+]]:_(<8 x s32>) = G_BUILD_VECTOR [[UV]](s32), [[UV2]](s32), [[UV1]](s32), [[UV3]](s32), [[C]](s32), [[C]](s32), [[C]](s32), [[C]](s32)
-    ; CHECK-NEXT: $wl0 = COPY [[BUILD_VECTOR]](<8 x s32>)
+    ; CHECK-NEXT: [[BUILD_VECTOR:%[0-9]+]]:_(<4 x s32>) = G_BUILD_VECTOR [[UV]](s32), [[UV2]](s32), [[UV1]](s32), [[UV3]](s32)
+    ; CHECK-NEXT: [[AIE_BROADCAST_VECTOR:%[0-9]+]]:_(<16 x s32>) = G_AIE_BROADCAST_VECTOR [[C]](s32)
+    ; CHECK-NEXT: [[AIE_UNPAD_VECTOR:%[0-9]+]]:_(<4 x s32>) = G_AIE_UNPAD_VECTOR [[AIE_BROADCAST_VECTOR]](<16 x s32>)
+    ; CHECK-NEXT: [[CONCAT_VECTORS:%[0-9]+]]:_(<8 x s32>) = G_CONCAT_VECTORS [[BUILD_VECTOR]](<4 x s32>), [[AIE_UNPAD_VECTOR]](<4 x s32>)
+    ; CHECK-NEXT: $wl0 = COPY [[CONCAT_VECTORS]](<8 x s32>)
       %10:_(<4 x s32>) = COPY $q0
       %0:_(s32), %1:_(s32), %2:_(s32), %3:_(s32) = G_UNMERGE_VALUES %10
       %4:_(s32) = G_CONSTANT i32 1
@@ -104,8 +113,11 @@ body:             |
     ; CHECK: [[COPY:%[0-9]+]]:_(<8 x s32>) = COPY $wh0
     ; CHECK-NEXT: [[UV:%[0-9]+]]:_(s32), [[UV1:%[0-9]+]]:_(s32), [[UV2:%[0-9]+]]:_(s32), [[UV3:%[0-9]+]]:_(s32), [[UV4:%[0-9]+]]:_(s32), [[UV5:%[0-9]+]]:_(s32), [[UV6:%[0-9]+]]:_(s32), [[UV7:%[0-9]+]]:_(s32) = G_UNMERGE_VALUES [[COPY]](<8 x s32>)
     ; CHECK-NEXT: [[DEF:%[0-9]+]]:_(s32) = G_IMPLICIT_DEF
-    ; CHECK-NEXT: [[BUILD_VECTOR:%[0-9]+]]:_(<16 x s32>) = G_BUILD_VECTOR [[UV]](s32), [[UV2]](s32), [[UV1]](s32), [[UV3]](s32), [[UV4]](s32), [[UV5]](s32), [[UV6]](s32), [[UV7]](s32), [[DEF]](s32), [[DEF]](s32), [[DEF]](s32), [[DEF]](s32), [[DEF]](s32), [[DEF]](s32), [[DEF]](s32), [[DEF]](s32)
-    ; CHECK-NEXT: $x0 = COPY [[BUILD_VECTOR]](<16 x s32>)
+    ; CHECK-NEXT: [[BUILD_VECTOR:%[0-9]+]]:_(<8 x s32>) = G_BUILD_VECTOR [[UV]](s32), [[UV2]](s32), [[UV1]](s32), [[UV3]](s32), [[UV4]](s32), [[UV5]](s32), [[UV6]](s32), [[UV7]](s32)
+    ; CHECK-NEXT: [[AIE_BROADCAST_VECTOR:%[0-9]+]]:_(<16 x s32>) = G_AIE_BROADCAST_VECTOR [[DEF]](s32)
+    ; CHECK-NEXT: [[AIE_UNPAD_VECTOR:%[0-9]+]]:_(<8 x s32>) = G_AIE_UNPAD_VECTOR [[AIE_BROADCAST_VECTOR]](<16 x s32>)
+    ; CHECK-NEXT: [[CONCAT_VECTORS:%[0-9]+]]:_(<16 x s32>) = G_CONCAT_VECTORS [[BUILD_VECTOR]](<8 x s32>), [[AIE_UNPAD_VECTOR]](<8 x s32>)
+    ; CHECK-NEXT: $x0 = COPY [[CONCAT_VECTORS]](<16 x s32>)
       %10:_(<8 x s32>) = COPY $wh0
       %0:_(s32), %1:_(s32), %2:_(s32), %3:_(s32), %4:_(s32), %5:_(s32), %6:_(s32), %7:_(s32) = G_UNMERGE_VALUES %10
       %8:_(s32) = G_IMPLICIT_DEF

Diff for: llvm/test/CodeGen/AIE/GlobalISel/prelegalizercombiner-splat-vector.mir

@@ -156,11 +156,11 @@ body:             |
     ; CHECK-LABEL: name: test_build_vector_256_32bit_scl_invalid
     ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY $r0
     ; CHECK-NEXT: [[COPY1:%[0-9]+]]:_(s32) = COPY $r1
-    ; CHECK-NEXT: [[BUILD_VECTOR:%[0-9]+]]:_(<8 x s32>) = G_BUILD_VECTOR [[COPY1]](s32), [[COPY]](s32), [[COPY]](s32), [[COPY1]](s32), [[COPY]](s32), [[COPY]](s32), [[COPY]](s32), [[COPY]](s32)
+    ; CHECK-NEXT: [[BUILD_VECTOR:%[0-9]+]]:_(<8 x s32>) = G_BUILD_VECTOR [[COPY1]](s32), [[COPY]](s32), [[COPY]](s32), [[COPY]](s32), [[COPY]](s32), [[COPY]](s32), [[COPY1]](s32), [[COPY]](s32)
     ; CHECK-NEXT: $wl0 = COPY [[BUILD_VECTOR]](<8 x s32>)
     %1:_(s32) = COPY $r0
     %2:_(s32) = COPY $r1
-    %3:_(<8 x s32>) = G_BUILD_VECTOR %2:_(s32), %1:_(s32), %1:_(s32), %2:_(s32), %1:_(s32), %1:_(s32), %1:_(s32), %1:_(s32)
+    %3:_(<8 x s32>) = G_BUILD_VECTOR %2:_(s32), %1:_(s32), %1:_(s32), %1:_(s32), %1:_(s32), %1:_(s32), %2:_(s32), %1:_(s32)
     $wl0 = COPY %3:_(<8 x s32>)
 ...
 
@@ -229,11 +229,11 @@ body:             |
     ; CHECK-LABEL: name: test_build_vector_512_32bit_scl_invalid
     ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY $r0
     ; CHECK-NEXT: [[COPY1:%[0-9]+]]:_(s32) = COPY $r1
-    ; CHECK-NEXT: [[BUILD_VECTOR:%[0-9]+]]:_(<16 x s32>) = G_BUILD_VECTOR [[COPY]](s32), [[COPY1]](s32), [[COPY]](s32), [[COPY1]](s32), [[COPY]](s32), [[COPY]](s32), [[COPY]](s32), [[COPY]](s32), [[COPY]](s32), [[COPY]](s32), [[COPY]](s32), [[COPY]](s32), [[COPY]](s32), [[COPY]](s32), [[COPY]](s32), [[COPY]](s32)
+    ; CHECK-NEXT: [[BUILD_VECTOR:%[0-9]+]]:_(<16 x s32>) = G_BUILD_VECTOR [[COPY]](s32), [[COPY]](s32), [[COPY]](s32), [[COPY1]](s32), [[COPY]](s32), [[COPY]](s32), [[COPY]](s32), [[COPY]](s32), [[COPY]](s32), [[COPY]](s32), [[COPY]](s32), [[COPY]](s32), [[COPY]](s32), [[COPY1]](s32), [[COPY]](s32), [[COPY]](s32)
     ; CHECK-NEXT: $x0 = COPY [[BUILD_VECTOR]](<16 x s32>)
     %1:_(s32) = COPY $r0
     %2:_(s32) = COPY $r1
-    %3:_(<16 x s32>) = G_BUILD_VECTOR %1:_(s32), %2:_(s32), %1:_(s32), %2:_(s32), %1:_(s32), %1:_(s32), %1:_(s32), %1:_(s32), %1:_(s32), %1:_(s32), %1:_(s32), %1:_(s32), %1:_(s32), %1:_(s32), %1:_(s32), %1:_(s32)
+    %3:_(<16 x s32>) = G_BUILD_VECTOR %1:_(s32), %1:_(s32), %1:_(s32), %2:_(s32), %1:_(s32), %1:_(s32), %1:_(s32), %1:_(s32), %1:_(s32), %1:_(s32), %1:_(s32), %1:_(s32), %1:_(s32), %2:_(s32), %1:_(s32), %1:_(s32)
     $x0 = COPY %3:_(<16 x s32>)
 ...