[AArch64][GlobalISel] Port buildvector -> dup pattern from AArch64ISelLowering

[paquette]

Summary

Given this:

%x:_(<n x sK>) = G_BUILD_VECTOR %lane, ...
...
%y:_(<n x sK>) = G_SHUFFLE_VECTOR %x(<n x sK>), %foo, shufflemask(0, 0, ...)

We can produce:

%y:_(<n x sK) = G_DUP %lane(sK)

Doesn't seem to be too common, but AArch64ISelLowering attempts to do this before trying to produce a DUPLANE. Might as well port it.

Factor out the other G_DUP pattern into a helper function, and create a helper for this pattern.

Also make it so that when the splat has an undef mask, we try setting it to 0. SDAG does this, and it makes sure that when we get the build vector operand, we actually get a source operand.

[aemerson]

What's the reason for avoiding matching a constant lane build_vector?

llvm/lib/Target/AArch64/GISel/AArch64PostLegalizerCombiner.cpp

290

misleading whitespace change.

[paquette]

In D81979#2097404, @aemerson wrote:

What's the reason for avoiding matching a constant lane build_vector?

I think that it's probably this comment in arm64-dup.ll:

; We used to spot this as a BUILD_VECTOR implementable by dup, but assume that
; the single value needed was of the same type as the vector. This is false if
; the scalar corresponding to the vector type is illegal (e.g. a <4 x i16>
; BUILD_VECTOR will have an i32 as its source). In that case, the operation is
; not a simple "dup vD.4h, vN.h[idx]" after all, and we crashed.
;
; *However*, it is a dup vD.4h, vN.h[2*idx].

This is wrt the following IR:

define <4 x i16> @test_build_illegal(<4 x i32> %in) {
  %val = extractelement <4 x i32> %in, i32 3
  %smallval = trunc i32 %val to i16
  %vec = insertelement <4x i16> undef, i16 %smallval, i32 3
  ret <4 x i16> %vec
}

[aemerson]

So is this still applicable to GISel? If not can we remove that restriction.

[paquette]

• Remove restriction on constants. I don't think that we can have any problems here.
• Fix whitespace in comment

[aemerson]

LGTM.

llvm/lib/Target/AArch64/GISel/AArch64PostLegalizerCombiner.cpp

if (!isZipMask(ShuffleMask, NumElts, WhichResult))
     return false;
   unsigned Opc = (WhichResult == 0) ? AArch64::G_ZIP1 : AArch64::G_ZIP2;
   Register V1 = MI.getOperand(1).getReg();
   Register V2 = MI.getOperand(2).getReg();
   MatchInfo = ShuffleVectorPseudo(Opc, Dst, {V1, V2});
   return true;
 }
 
-static bool matchDup(MachineInstr &MI, MachineRegisterInfo &MRI,
+/// Helper function for matchDup.
+static bool matchDupFromInsertVectorElt(int Lane, MachineInstr &MI,
+                                        MachineRegisterInfo &MRI,
                                         ShuffleVectorPseudo &MatchInfo) {
-  assert(MI.getOpcode() == TargetOpcode::G_SHUFFLE_VECTOR);
-  auto Lane = getSplatIndex(MI);
-  if (!Lane || *Lane != 0)
+  if (Lane != 0)
     return false;
 
   // Try to match a vector splat operation into a dup instruction.
   // We're looking for this pattern:
   //
   // %scalar:gpr(s64) = COPY $x0
   // %undef:fpr(<2 x s64>) = G_IMPLICIT_DEF
   // %cst0:gpr(s32) = G_CONSTANT i32 0
   // %zerovec:fpr(<2 x s32>) = G_BUILD_VECTOR %cst0(s32), %cst0(s32)
   // %ins:fpr(<2 x s64>) = G_INSERT_VECTOR_ELT %undef, %scalar(s64), %cst0(s32)
   // %splat:fpr(<2 x s64>) = G_SHUFFLE_VECTOR %ins(<2 x s64>), %undef, %zerovec(<2 x s32>)
   //
   // ...into:
   // %splat = G_DUP %scalar
 
   // Begin matching the insert.
   auto *InsMI = getOpcodeDef(TargetOpcode::G_INSERT_VECTOR_ELT,
                              MI.getOperand(1).getReg(), MRI);
   if (!InsMI)
     return false;
-
   // Match the undef vector operand.
-  if (!getOpcodeDef(TargetOpcode::G_IMPLICIT_DEF,
-                               InsMI->getOperand(1).getReg(), MRI))
+  if (!getOpcodeDef(TargetOpcode::G_IMPLICIT_DEF, InsMI->getOperand(1).getReg(),
+                    MRI))
     return false;
 
   // Match the index constant 0.
   int64_t Index = 0;
   if (!mi_match(InsMI->getOperand(3).getReg(), MRI, m_ICst(Index)) || Index)
     return false;
 
-  Register Dst = MI.getOperand(0).getReg();
+  MatchInfo = ShuffleVectorPseudo(AArch64::G_DUP, MI.getOperand(0).getReg(),
+                                  {InsMI->getOperand(2).getReg()});
+  return true;
+}
+
+/// Helper function for matchDup.
+static bool matchDupFromBuildVector(int Lane, MachineInstr &MI,
+                                    MachineRegisterInfo &MRI,
+                                    ShuffleVectorPseudo &MatchInfo) {
+  assert(Lane >= 0 && "Expected positive lane?");
+  // Test if the LHS is a BUILD_VECTOR. If it is, then we can just reference the
+  // lane's definition directly.
+  auto *BuildVecMI = getOpcodeDef(TargetOpcode::G_BUILD_VECTOR,
+                                  MI.getOperand(1).getReg(), MRI);
+  if (!BuildVecMI)
+    return false;
+  Register Reg = BuildVecMI->getOperand(Lane + 1).getReg();
   MatchInfo =
-      ShuffleVectorPseudo(AArch64::G_DUP, Dst, {InsMI->getOperand(2).getReg()});
+      ShuffleVectorPseudo(AArch64::G_DUP, MI.getOperand(0).getReg(), {Reg});
   return true;
 }
 
+static bool matchDup(MachineInstr &MI, MachineRegisterInfo &MRI,
+                     ShuffleVectorPseudo &MatchInfo) {
+  assert(MI.getOpcode() == TargetOpcode::G_SHUFFLE_VECTOR);
+  auto MaybeLane = getSplatIndex(MI);
+  if (!MaybeLane)
+    return false;
+  int Lane = *MaybeLane;
+  // If this is undef splat, generate it via "just" vdup, if possible.
+  if (Lane < 0)
+    Lane = 0;
+  if (matchDupFromInsertVectorElt(Lane, MI, MRI, MatchInfo))
+    return true;
+  if (matchDupFromBuildVector(Lane, MI, MRI, MatchInfo))
+    return true;
+  return false;
+}
+
 static bool matchEXT(MachineInstr &MI, MachineRegisterInfo &MRI,
                      ShuffleVectorPseudo &MatchInfo) {
   assert(MI.getOpcode() == TargetOpcode::G_SHUFFLE_VECTOR);
   Register Dst = MI.getOperand(0).getReg();
   auto ExtInfo = getExtMask(MI.getOperand(3).getShuffleMask(),
                             MRI.getType(Dst).getNumElements());
   if (!ExtInfo)
     return false;

llvm/test/CodeGen/AArch64/GlobalISel/postlegalizer-combiner-shuffle-splat.mir

bb.1.entry:
     ; CHECK: RET_ReallyLR implicit $d0
     %copy:_(s32) = COPY $w0
     %undef:_(<8 x s8>) = G_IMPLICIT_DEF
     %cst:_(s32) = G_CONSTANT i32 0
     %ins:_(<8 x s8>) = G_INSERT_VECTOR_ELT %undef, %copy(s32), %cst(s32)
     %splat:_(<8 x s8>) = G_SHUFFLE_VECTOR %ins(<8 x s8>), %undef, shufflemask(0, 0, 0, 0, 0, 0, 0, 0)
     $d0 = COPY %splat(<8 x s8>)
     RET_ReallyLR implicit $d0
+
+...
+---
+name:            build_vector
+alignment:       4
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.1.entry:
+    liveins: $w0, $w1, $w2, $w3
+    ; The G_SHUFFLE_VECTOR is fed by a G_BUILD_VECTOR, and the 0th input
+    ; operand is not a constant. We should get a G_DUP.
+    ;
+    ; CHECK-LABEL: name: build_vector
+    ; CHECK: liveins: $w0, $w1, $w2, $w3
+    ; CHECK: %lane:_(s32) = COPY $w0
+    ; CHECK: %shuf:_(<4 x s32>) = G_DUP %lane(s32)
+    ; CHECK: $q0 = COPY %shuf(<4 x s32>)
+    ; CHECK: RET_ReallyLR implicit $q0
+    %lane:_(s32) = COPY $w0
+    %b:_(s32) = COPY $w1
+    %c:_(s32) = COPY $w2
+    %d:_(s32) = COPY $w3
+    %undef:_(<4 x s32>) = G_IMPLICIT_DEF
+    %buildvec:_(<4 x s32>) = G_BUILD_VECTOR %lane, %b, %c, %d
+    %shuf:_(<4 x s32>) = G_SHUFFLE_VECTOR %buildvec(<4 x s32>), %undef, shufflemask(0, 0, 0, 0)
+    $q0 = COPY %shuf(<4 x s32>)
+    RET_ReallyLR implicit $q0