[AArch64][GlobalISel] Implement selection of <2 x float> vector splat.

[aemerson]

Summary

[AArch64][GlobalISel] Implement selection of <2 x float> vector splat.

Also requires making G_IMPLICIT_DEF of v2s32 legal.

[paquette]

llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp

3718

Would it make sense to assert that there are 4 elements?

3731

Same here?

[aemerson]

Add asserts.

[paquette]

LGTM

llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp

bool AArch64InstructionSelector::tryOptVectorDup(MachineInstr &I) const {
   // The shuffle's second operand doesn't matter if the mask is all zero.
   const Constant *Mask = I.getOperand(3).getShuffleMask();
   if (!isa<ConstantAggregateZero>(Mask))
     return false;
 
   // We're done, now find out what kind of splat we need.
   LLT VecTy = MRI.getType(I.getOperand(0).getReg());
   LLT EltTy = VecTy.getElementType();
-  if (VecTy.getSizeInBits() != 128 || EltTy.getSizeInBits() < 32) {
-    LLVM_DEBUG(dbgs() << "Could not optimize splat pattern < 128b yet");
+  if (EltTy.getSizeInBits() < 32) {
+    LLVM_DEBUG(dbgs() << "Could not optimize splat pattern < 32b elts yet");
     return false;
   }
   bool IsFP = ScalarRB->getID() == AArch64::FPRRegBankID;
-  static const unsigned OpcTable[2][2] = {
-      {AArch64::DUPv4i32gpr, AArch64::DUPv2i64gpr},
-      {AArch64::DUPv4i32lane, AArch64::DUPv2i64lane}};
-  unsigned Opc = OpcTable[IsFP][EltTy.getSizeInBits() == 64];
+  unsigned Opc = 0;
+  if (IsFP) {
+    switch (EltTy.getSizeInBits()) {
+    case 32:
+      if (VecTy.getNumElements() == 2) {
+        Opc = AArch64::DUPv2i32lane;
+      } else {
+        Opc = AArch64::DUPv4i32lane;
+        assert(VecTy.getNumElements() == 4);
+      }
+      break;
+    case 64:
+      assert(VecTy.getNumElements() == 2 && "Unexpected num elts");
+      Opc = AArch64::DUPv2i64lane;
+      break;
+    }
+  } else {
+    switch (EltTy.getSizeInBits()) {
+    case 32:
+      if (VecTy.getNumElements() == 2) {
+        Opc = AArch64::DUPv2i32gpr;
+      } else {
+        Opc = AArch64::DUPv4i32gpr;
+        assert(VecTy.getNumElements() == 4);
+      }
+      break;
+    case 64:
+      assert(VecTy.getNumElements() == 2 && "Unexpected num elts");
+      Opc = AArch64::DUPv2i64gpr;
+      break;
+    }
+  }
+  assert(Opc && "Did not compute an opcode for a dup");
 
   // For FP splats, we need to widen the scalar reg via undef too.
   if (IsFP) {
     MachineInstr *Widen = emitScalarToVector(
         EltTy.getSizeInBits(), &AArch64::FPR128RegClass, ScalarReg, MIB);
     if (!Widen)
       return false;
     ScalarReg = Widen->getOperand(0).getReg();

llvm/lib/Target/AArch64/AArch64LegalizerInfo.cpp

AArch64LegalizerInfo::AArch64LegalizerInfo(const AArch64Subtarget &ST) {
 
   // FIXME: support subtargets which have neon/fp-armv8 disabled.
   if (!ST.hasNEON() || !ST.hasFPARMv8()) {
     computeTables();
     return;
   }
 
   getActionDefinitionsBuilder(G_IMPLICIT_DEF)
-    .legalFor({p0, s1, s8, s16, s32, s64, v4s32, v2s64})
+    .legalFor({p0, s1, s8, s16, s32, s64, v2s32, v4s32, v2s64})
     .clampScalar(0, s1, s64)
     .widenScalarToNextPow2(0, 8)
     .fewerElementsIf(
       [=](const LegalityQuery &Query) {
         return Query.Types[0].isVector() &&
           (Query.Types[0].getElementType() != s64 ||
            Query.Types[0].getNumElements() != 2);
       },

llvm/test/CodeGen/AArch64/GlobalISel/legalize-undef.mir

bb.0:
     %0:_(s3) = G_IMPLICIT_DEF
     %1:_(s64) = G_SEXT %0
     $x0 = COPY %1(s64)
 ...
 
 # FIXME: s2 not correctly handled
 
 ---
-name: test_implicit_def_v2s32
+name: test_implicit_def_v4s32
 body: |
   bb.0:
 
-    ; CHECK-LABEL: name: test_implicit_def_v2s32
+    ; CHECK-LABEL: name: test_implicit_def_v4s32
     ; CHECK: [[DEF:%[0-9]+]]:_(<4 x s32>) = G_IMPLICIT_DEF
     ; CHECK: [[UV:%[0-9]+]]:_(<2 x s32>), [[UV1:%[0-9]+]]:_(<2 x s32>) = G_UNMERGE_VALUES [[DEF]](<4 x s32>)
     ; CHECK: $x0 = COPY [[UV]](<2 x s32>)
     ; CHECK: $x1 = COPY [[UV1]](<2 x s32>)
     %0:_(<4 x s32>) = G_IMPLICIT_DEF
     %1:_(<2 x s32> ), %2:_(<2 x s32>) = G_UNMERGE_VALUES %0
     $x0 = COPY %1
     $x1 = COPY %2
bb.0:
     ; CHECK: [[DEF1:%[0-9]+]]:_(<2 x s64>) = G_IMPLICIT_DEF
     ; CHECK: $q0 = COPY [[DEF]](<2 x s64>)
     ; CHECK: $q1 = COPY [[DEF1]](<2 x s64>)
     %0:_(<4 x s64>) = G_IMPLICIT_DEF
     %1:_(<2 x s64> ), %2:_(<2 x s64>) = G_UNMERGE_VALUES %0
     $q0 = COPY %1
     $q1 = COPY %2
 ...
+---
+name: test_implicit_def_v2s32
+body: |
+  bb.0:
+
+    ; CHECK-LABEL: name: test_implicit_def_v2s32
+    ; CHECK: [[DEF:%[0-9]+]]:_(<2 x s32>) = G_IMPLICIT_DEF
+    ; CHECK: [[UV:%[0-9]+]]:_(s32), [[UV1:%[0-9]+]]:_(s32) = G_UNMERGE_VALUES [[DEF]](<2 x s32>)
+    ; CHECK: $w0 = COPY [[UV]](s32)
+    ; CHECK: $w1 = COPY [[UV1]](s32)
+    %0:_(<2 x s32>) = G_IMPLICIT_DEF
+    %1:_(s32), %2:_(s32) = G_UNMERGE_VALUES %0
+    $w0 = COPY %1
+    $w1 = COPY %2
+...

llvm/test/CodeGen/AArch64/GlobalISel/opt-shuffle-splat.mir

bb.1.entry:
     %3:gpr(s32) = G_CONSTANT i32 0
     %1:fpr(<2 x s64>) = G_INSERT_VECTOR_ELT %2, %0(s64), %3(s32)
     %4:fpr(<2 x s64>) = G_SHUFFLE_VECTOR %1(<2 x s64>), %2, shufflemask(0, 0)
     $q0 = COPY %4(<2 x s64>)
     RET_ReallyLR implicit $q0
 
 ...
 ---
+name:            splat_2xi32
+alignment:       4
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.1.entry:
+    liveins: $w0
+
+    ; CHECK-LABEL: name: splat_2xi32
+    ; CHECK: liveins: $w0
+    ; CHECK: [[COPY:%[0-9]+]]:gpr32 = COPY $w0
+    ; CHECK: [[DUPv2i32gpr:%[0-9]+]]:fpr64 = DUPv2i32gpr [[COPY]]
+    ; CHECK: $d0 = COPY [[DUPv2i32gpr]]
+    ; CHECK: RET_ReallyLR implicit $d0
+    %0:gpr(s32) = COPY $w0
+    %2:fpr(<2 x s32>) = G_IMPLICIT_DEF
+    %3:gpr(s32) = G_CONSTANT i32 0
+    %1:fpr(<2 x s32>) = G_INSERT_VECTOR_ELT %2, %0(s32), %3(s32)
+    %4:fpr(<2 x s32>) = G_SHUFFLE_VECTOR %1(<2 x s32>), %2, shufflemask(0, 0)
+    $d0 = COPY %4(<2 x s32>)
+    RET_ReallyLR implicit $d0
+
+...
+---
 name:            splat_4xf32
 alignment:       4
 legalized:       true
 regBankSelected: true
 tracksRegLiveness: true
 body:             |
   bb.1.entry:
     liveins: $s0
bb.1.entry:
     %3:gpr(s32) = G_CONSTANT i32 0
     %1:fpr(<2 x s64>) = G_INSERT_VECTOR_ELT %2, %0(s64), %3(s32)
     %4:fpr(<2 x s64>) = G_SHUFFLE_VECTOR %1(<2 x s64>), %2, shufflemask(0, 0)
     $q0 = COPY %4(<2 x s64>)
     RET_ReallyLR implicit $q0
 
 ...
 ---
+name:            splat_2xf32
+alignment:       4
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.1.entry:
+    liveins: $s0
+
+    ; CHECK-LABEL: name: splat_2xf32
+    ; CHECK: liveins: $s0
+    ; CHECK: [[COPY:%[0-9]+]]:fpr32 = COPY $s0
+    ; CHECK: [[DEF:%[0-9]+]]:fpr128 = IMPLICIT_DEF
+    ; CHECK: [[INSERT_SUBREG:%[0-9]+]]:fpr128 = INSERT_SUBREG [[DEF]], [[COPY]], %subreg.ssub
+    ; CHECK: [[DUPv2i32lane:%[0-9]+]]:fpr64 = DUPv2i32lane [[INSERT_SUBREG]], 0
+    ; CHECK: $d0 = COPY [[DUPv2i32lane]]
+    ; CHECK: RET_ReallyLR implicit $d0
+    %0:fpr(s32) = COPY $s0
+    %2:fpr(<2 x s32>) = G_IMPLICIT_DEF
+    %3:gpr(s32) = G_CONSTANT i32 0
+    %1:fpr(<2 x s32>) = G_INSERT_VECTOR_ELT %2, %0(s32), %3(s32)
+    %4:fpr(<2 x s32>) = G_SHUFFLE_VECTOR %1(<2 x s32>), %2, shufflemask(0, 0)
+    $d0 = COPY %4(<2 x s32>)
+    RET_ReallyLR implicit $d0
+
+...
+---
 name:            splat_2xf64_copies
 alignment:       4
 legalized:       true
 regBankSelected: true
 tracksRegLiveness: true
 body:             |
   bb.1.entry:
     liveins: $d0

llvm/test/CodeGen/AArch64/arm64-rev.ll

 ; CHECK-NEXT:    rev64.2s v0, v0
 ; CHECK-NEXT:    ret
 ;
 ; FALLBACK-LABEL: test_vrev64D32:
 ; FALLBACK:       // %bb.0:
 ; FALLBACK-NEXT:    ldr d0, [x0]
 ; FALLBACK-NEXT:    adrp x8, .LCPI13_0
 ; FALLBACK-NEXT:    ldr d1, [x8, :lo12:.LCPI13_0]
-; FALLBACK-NEXT:    mov.s v2[1], w8
-; FALLBACK-NEXT:    mov.d v0[1], v2[0]
+; FALLBACK-NEXT:    mov.d v0[1], v0[0]
 ; FALLBACK-NEXT:    tbl.16b v0, { v0 }, v1
 ; FALLBACK-NEXT:    // kill: def $d0 killed $d0 killed $q0
 ; FALLBACK-NEXT:    ret
   %tmp1 = load <2 x i32>, <2 x i32>* %A
   %tmp2 = shufflevector <2 x i32> %tmp1, <2 x i32> undef, <2 x i32> <i32 1, i32 0>
   ret <2 x i32> %tmp2
 }
 
 define <2 x float> @test_vrev64Df(<2 x float>* %A) nounwind {
 ; CHECK-LABEL: test_vrev64Df:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ldr d0, [x0]
 ; CHECK-NEXT:    rev64.2s v0, v0
 ; CHECK-NEXT:    ret
 ;
 ; FALLBACK-LABEL: test_vrev64Df:
 ; FALLBACK:       // %bb.0:
 ; FALLBACK-NEXT:    ldr d0, [x0]
 ; FALLBACK-NEXT:    adrp x8, .LCPI14_0
 ; FALLBACK-NEXT:    ldr d1, [x8, :lo12:.LCPI14_0]
-; FALLBACK-NEXT:    mov.s v2[1], w8
-; FALLBACK-NEXT:    mov.d v0[1], v2[0]
+; FALLBACK-NEXT:    mov.d v0[1], v0[0]
 ; FALLBACK-NEXT:    tbl.16b v0, { v0 }, v1
 ; FALLBACK-NEXT:    // kill: def $d0 killed $d0 killed $q0
 ; FALLBACK-NEXT:    ret
   %tmp1 = load <2 x float>, <2 x float>* %A
   %tmp2 = shufflevector <2 x float> %tmp1, <2 x float> undef, <2 x i32> <i32 1, i32 0>
   ret <2 x float> %tmp2
 }