[GlobalISel][AArch64] Add partial selection support for G_INSERT_VECTOR_ELT

[paquette]

Summary

This adds support for inserting elements into packed vectors. It also adds two tests: one for selection, and one for regbank select.

Unpacked vectors will come in a follow-up.

[paquette]

rebase against changes to emitLaneInsert and include change to the legalizer info test I forgot to include

[paquette]

Rebase against changes to emitLaneInsert

[aemerson]

llvm/lib/Target/AArch64/AArch64LegalizerInfo.cpp

455

Type 0 should always be a vector, else it (should be) a verifier error so I don't think it's related to legality.

Likewise with the "Types[0].getElementType() == Query.Types[1]", that should be true otherwise it's an invalid instruction.

Here we should probably specify legality by using a whitelist of vector types that Type 0 could take, or some other way or narrowing down the legal types.

[paquette]

Fix up legalizer info stuff

[aemerson]

llvm/lib/Target/AArch64/AArch64LegalizerInfo.cpp

457

Tests for < 128b vectors?

[paquette]

Restrict legal types to exactly v4s32 and v2s64, add TODOs for other types.

llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp

private:
   ///
   /// The lane inserted into is defined by \p LaneIdx. The vector source
   /// register is given by \p SrcReg. The register containing the element is
   /// given by \p EltReg.
   MachineInstr *emitLaneInsert(Optional<unsigned> DstReg, unsigned SrcReg,
                                unsigned EltReg, unsigned LaneIdx,
                                const RegisterBank &RB,
                                MachineIRBuilder &MIRBuilder) const;
+  bool selectInsertElt(MachineInstr &I, MachineRegisterInfo &MRI) const;
   bool selectBuildVector(MachineInstr &I, MachineRegisterInfo &MRI) const;
   bool selectMergeValues(MachineInstr &I, MachineRegisterInfo &MRI) const;
   bool selectUnmergeValues(MachineInstr &I, MachineRegisterInfo &MRI) const;
 
   void collectShuffleMaskIndices(MachineInstr &I, MachineRegisterInfo &MRI,
                                  SmallVectorImpl<int> &Idxs) const;
   bool selectShuffleVector(MachineInstr &I, MachineRegisterInfo &MRI) const;
   bool selectExtractElt(MachineInstr &I, MachineRegisterInfo &MRI) const;
#endif
   case TargetOpcode::G_MERGE_VALUES:
     return selectMergeValues(I, MRI);
   case TargetOpcode::G_UNMERGE_VALUES:
     return selectUnmergeValues(I, MRI);
   case TargetOpcode::G_SHUFFLE_VECTOR:
     return selectShuffleVector(I, MRI);
   case TargetOpcode::G_EXTRACT_VECTOR_ELT:
     return selectExtractElt(I, MRI);
+  case TargetOpcode::G_INSERT_VECTOR_ELT:
+    return selectInsertElt(I, MRI);
   }
 
   return false;
 }
 
 MachineInstr *AArch64InstructionSelector::emitScalarToVector(
     unsigned EltSize, const TargetRegisterClass *DstRC, unsigned Scalar,
     MachineIRBuilder &MIRBuilder) const {
InsElt = MIRBuilder.buildInstr(Opc, {*DstReg}, {SrcReg})
                  .addImm(LaneIdx)
                  .addUse(EltReg);
   }
 
   constrainSelectedInstRegOperands(*InsElt, TII, TRI, RBI);
   return InsElt;
 }
 
+bool AArch64InstructionSelector::selectInsertElt(
+    MachineInstr &I, MachineRegisterInfo &MRI) const {
+  assert(I.getOpcode() == TargetOpcode::G_INSERT_VECTOR_ELT);
+
+  // Get information on the destination.
+  unsigned DstReg = I.getOperand(0).getReg();
+  const LLT DstTy = MRI.getType(DstReg);
+  if (DstTy.getSizeInBits() < 128) {
+    // TODO: Handle unpacked vectors.
+    LLVM_DEBUG(dbgs() << "Unpacked vectors not supported yet!");
+    return false;
+  }
+
+  // Get information on the element we want to insert into the destination.
+  unsigned EltReg = I.getOperand(2).getReg();
+  const LLT EltTy = MRI.getType(EltReg);
+  unsigned EltSize = EltTy.getSizeInBits();
+  if (EltSize < 16 || EltSize > 64)
+    return false; // Don't support all element types yet.
+
+  // Find the definition of the index. Bail out if it's not defined by a
+  // G_CONSTANT.
+  unsigned IdxReg = I.getOperand(3).getReg();
+  unsigned LaneIdx = 0;
+  if (!getConstantValueForReg(IdxReg, MRI, LaneIdx))
+    return false;
+
+  // Perform the lane insert.
+  unsigned SrcReg = I.getOperand(1).getReg();
+  const RegisterBank &EltRB = *RBI.getRegBank(EltReg, MRI, TRI);
+  MachineIRBuilder MIRBuilder(I);
+  emitLaneInsert(DstReg, SrcReg, EltReg, LaneIdx, EltRB, MIRBuilder);
+  I.eraseFromParent();
+  return true;
+}
+
 bool AArch64InstructionSelector::selectBuildVector(
     MachineInstr &I, MachineRegisterInfo &MRI) const {
   assert(I.getOpcode() == TargetOpcode::G_BUILD_VECTOR);
   // Until we port more of the optimized selections, for now just use a vector
   // insert sequence.
   const LLT DstTy = MRI.getType(I.getOperand(0).getReg());
   const LLT EltTy = MRI.getType(I.getOperand(1).getReg());
   unsigned EltSize = EltTy.getSizeInBits();

llvm/lib/Target/AArch64/AArch64LegalizerInfo.cpp

getActionDefinitionsBuilder(G_EXTRACT_VECTOR_ELT)
       })
       .minScalar(2, s64)
       .legalIf([=](const LegalityQuery &Query) {
         const LLT &VecTy = Query.Types[1];
         return VecTy == v2s16 || VecTy == v4s16 || VecTy == v4s32 ||
                VecTy == v2s64 || VecTy == v2s32;
       });
 
+  getActionDefinitionsBuilder(G_INSERT_VECTOR_ELT)
+      .legalIf([=](const LegalityQuery &Query) {
+        const LLT &VecTy = Query.Types[0];
+        // TODO: Support destination sizes of < 128 bits.
+        // TODO: Support s8 and s16
+        return VecTy == v4s32 || VecTy == v2s64;
+      });
+
   getActionDefinitionsBuilder(G_BUILD_VECTOR)
       .legalFor({{v4s16, s16},
                  {v8s16, s16},
                  {v2s32, s32},
                  {v4s32, s32},
                  {v2s64, s64}})
       .clampNumElements(0, v4s32, v4s32)
       .clampNumElements(0, v2s64, v2s64)

llvm/lib/Target/AArch64/AArch64RegisterBankInfo.cpp

AArch64RegisterBankInfo::getInstrMapping(const MachineInstr &MI) const {
   case TargetOpcode::G_EXTRACT_VECTOR_ELT:
     // Destination and source need to be FPRs.
     OpRegBankIdx[0] = PMI_FirstFPR;
     OpRegBankIdx[1] = PMI_FirstFPR;
 
     // Index needs to be a GPR.
     OpRegBankIdx[2] = PMI_FirstGPR;
     break;
+  case TargetOpcode::G_INSERT_VECTOR_ELT:
+    OpRegBankIdx[0] = PMI_FirstFPR;
+    OpRegBankIdx[1] = PMI_FirstFPR;
 
+    // The element may be either a GPR or FPR. Preserve that behaviour.
+    if (getRegBank(MI.getOperand(2).getReg(), MRI, TRI) == &AArch64::FPRRegBank)
+      OpRegBankIdx[2] = PMI_FirstFPR;
+    else
+      OpRegBankIdx[2] = PMI_FirstGPR;
+
+    // Index needs to be a GPR.
+    OpRegBankIdx[3] = PMI_FirstGPR;
+    break;
   case TargetOpcode::G_BUILD_VECTOR:
     // If the first source operand belongs to a FPR register bank, then make
     // sure that we preserve that.
     if (OpRegBankIdx[1] != PMI_FirstGPR)
       break;
     unsigned VReg = MI.getOperand(1).getReg();
     if (!VReg)
       break;

llvm/test/CodeGen/AArch64/GlobalISel/legalizer-info-validation.mir

 #
 # DEBUG-NEXT: G_PTR_MASK (opcode {{[0-9]+}}): 1 type index
 # DEBUG:      .. the first uncovered type index: 1, OK
 #
 # DEBUG-NEXT: G_BR (opcode {{[0-9]+}}): 0 type indices
 # DEBUG:      .. type index coverage check SKIPPED: no rules defined
 #
 # DEBUG-NEXT: G_INSERT_VECTOR_ELT (opcode {{[0-9]+}}): 3 type indices
-# DEBUG:      .. type index coverage check SKIPPED: no rules defined
+# DEBUG:      .. type index coverage check SKIPPED: user-defined predicate detected
 #
 # DEBUG-NEXT: G_EXTRACT_VECTOR_ELT (opcode {{[0-9]+}}): 3 type indices
 # DEBUG:      .. type index coverage check SKIPPED: user-defined predicate detected
 #
 # DEBUG-NEXT: G_SHUFFLE_VECTOR (opcode {{[0-9]+}}): 3 type indices
 # DEBUG:      .. type index coverage check SKIPPED: user-defined predicate detected
 #
 # DEBUG-NEXT: G_CTTZ (opcode {{[0-9]+}}): 2 type indices

llvm/test/CodeGen/AArch64/GlobalISel/regbank-insert-vector-elt.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -verify-machineinstrs -mtriple aarch64-unknown-unknown -run-pass=regbankselect %s -o - | FileCheck %s
+
+# The following should hold here:
+#
+# 1) The first and second operands of G_INSERT_VECTOR_ELT should be FPRs since
+#    they are vectors.
+#
+# 2) The third operand should be on the register bank given in the test name
+#    (e.g, v4s32_fpr). AArch64 supports native inserts of GPRs, so we need to
+#    preserve that.
+#
+# 3) The fourth operand should be a GPR, since it's a constant.
+
+name:            v4s32_fpr
+alignment:       2
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $q1, $s0
+
+    ; CHECK-LABEL: name: v4s32_fpr
+    ; CHECK: liveins: $q1, $s0
+    ; CHECK: [[COPY:%[0-9]+]]:fpr(s32) = COPY $s0
+    ; CHECK: [[COPY1:%[0-9]+]]:fpr(<4 x s32>) = COPY $q1
+    ; CHECK: [[C:%[0-9]+]]:gpr(s32) = G_CONSTANT i32 1
+    ; CHECK: [[IVEC:%[0-9]+]]:fpr(<4 x s32>) = G_INSERT_VECTOR_ELT [[COPY1]], [[COPY]](s32), [[C]](s32)
+    ; CHECK: $q0 = COPY [[IVEC]](<4 x s32>)
+    ; CHECK: RET_ReallyLR implicit $q0
+    %0:_(s32) = COPY $s0
+    %1:_(<4 x s32>) = COPY $q1
+    %3:_(s32) = G_CONSTANT i32 1
+    %2:_(<4 x s32>) = G_INSERT_VECTOR_ELT %1, %0(s32), %3(s32)
+    $q0 = COPY %2(<4 x s32>)
+    RET_ReallyLR implicit $q0
+
+...
+---
+name:            v4s32_gpr
+alignment:       2
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $q0, $w0
+
+    ; CHECK-LABEL: name: v4s32_gpr
+    ; CHECK: liveins: $q0, $w0
+    ; CHECK: [[COPY:%[0-9]+]]:gpr(s32) = COPY $w0
+    ; CHECK: [[COPY1:%[0-9]+]]:fpr(<4 x s32>) = COPY $q0
+    ; CHECK: [[C:%[0-9]+]]:gpr(s32) = G_CONSTANT i32 1
+    ; CHECK: [[IVEC:%[0-9]+]]:fpr(<4 x s32>) = G_INSERT_VECTOR_ELT [[COPY1]], [[COPY]](s32), [[C]](s32)
+    ; CHECK: $q0 = COPY [[IVEC]](<4 x s32>)
+    ; CHECK: RET_ReallyLR implicit $q0
+    %0:_(s32) = COPY $w0
+    %1:_(<4 x s32>) = COPY $q0
+    %3:_(s32) = G_CONSTANT i32 1
+    %2:_(<4 x s32>) = G_INSERT_VECTOR_ELT %1, %0(s32), %3(s32)
+    $q0 = COPY %2(<4 x s32>)
+    RET_ReallyLR implicit $q0
+
+...
+---
+name:            v2s64_fpr
+alignment:       2
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $d0, $q1
+
+    ; CHECK-LABEL: name: v2s64_fpr
+    ; CHECK: liveins: $d0, $q1
+    ; CHECK: [[COPY:%[0-9]+]]:fpr(s64) = COPY $d0
+    ; CHECK: [[COPY1:%[0-9]+]]:fpr(<2 x s64>) = COPY $q1
+    ; CHECK: [[C:%[0-9]+]]:gpr(s32) = G_CONSTANT i32 1
+    ; CHECK: [[IVEC:%[0-9]+]]:fpr(<2 x s64>) = G_INSERT_VECTOR_ELT [[COPY1]], [[COPY]](s64), [[C]](s32)
+    ; CHECK: $q0 = COPY [[IVEC]](<2 x s64>)
+    ; CHECK: RET_ReallyLR implicit $q0
+    %0:_(s64) = COPY $d0
+    %1:_(<2 x s64>) = COPY $q1
+    %3:_(s32) = G_CONSTANT i32 1
+    %2:_(<2 x s64>) = G_INSERT_VECTOR_ELT %1, %0(s64), %3(s32)
+    $q0 = COPY %2(<2 x s64>)
+    RET_ReallyLR implicit $q0
+
+...
+---
+name:            v2s64_gpr
+alignment:       2
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $q0, $x0
+
+    ; CHECK-LABEL: name: v2s64_gpr
+    ; CHECK: liveins: $q0, $x0
+    ; CHECK: [[COPY:%[0-9]+]]:gpr(s64) = COPY $x0
+    ; CHECK: [[COPY1:%[0-9]+]]:fpr(<2 x s64>) = COPY $q0
+    ; CHECK: [[C:%[0-9]+]]:gpr(s32) = G_CONSTANT i32 0
+    ; CHECK: [[IVEC:%[0-9]+]]:fpr(<2 x s64>) = G_INSERT_VECTOR_ELT [[COPY1]], [[COPY]](s64), [[C]](s32)
+    ; CHECK: $q0 = COPY [[IVEC]](<2 x s64>)
+    ; CHECK: RET_ReallyLR implicit $q0
+    %0:_(s64) = COPY $x0
+    %1:_(<2 x s64>) = COPY $q0
+    %3:_(s32) = G_CONSTANT i32 0
+    %2:_(<2 x s64>) = G_INSERT_VECTOR_ELT %1, %0(s64), %3(s32)
+    $q0 = COPY %2(<2 x s64>)
+    RET_ReallyLR implicit $q0
+
+...

llvm/test/CodeGen/AArch64/GlobalISel/select-insert-vector-elt.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -verify-machineinstrs -mtriple aarch64-unknown-unknown -run-pass=instruction-select %s -o - | FileCheck %s
+
+name:            v4s32_fpr
+alignment:       2
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $q1, $s0
+
+    ; CHECK-LABEL: name: v4s32_fpr
+    ; CHECK: liveins: $q1, $s0
+    ; CHECK: [[COPY:%[0-9]+]]:fpr32 = COPY $s0
+    ; CHECK: [[COPY1:%[0-9]+]]:fpr128 = COPY $q1
+    ; CHECK: [[DEF:%[0-9]+]]:fpr128 = IMPLICIT_DEF
+    ; CHECK: [[INSERT_SUBREG:%[0-9]+]]:fpr128 = INSERT_SUBREG [[DEF]], [[COPY]], %subreg.ssub
+    ; CHECK: [[INSvi32lane:%[0-9]+]]:fpr128 = INSvi32lane [[COPY1]], 1, [[INSERT_SUBREG]], 0
+    ; CHECK: $q0 = COPY [[INSvi32lane]]
+    ; CHECK: RET_ReallyLR implicit $q0
+    %0:fpr(s32) = COPY $s0
+    %1:fpr(<4 x s32>) = COPY $q1
+    %3:gpr(s32) = G_CONSTANT i32 1
+    %2:fpr(<4 x s32>) = G_INSERT_VECTOR_ELT %1, %0(s32), %3(s32)
+    $q0 = COPY %2(<4 x s32>)
+    RET_ReallyLR implicit $q0
+
+...
+---
+name:            v4s32_gpr
+alignment:       2
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $q0, $w0
+
+    ; CHECK-LABEL: name: v4s32_gpr
+    ; CHECK: liveins: $q0, $w0
+    ; CHECK: [[COPY:%[0-9]+]]:gpr32 = COPY $w0
+    ; CHECK: [[COPY1:%[0-9]+]]:fpr128 = COPY $q0
+    ; CHECK: [[INSvi32gpr:%[0-9]+]]:fpr128 = INSvi32gpr [[COPY1]], 1, [[COPY]]
+    ; CHECK: $q0 = COPY [[INSvi32gpr]]
+    ; CHECK: RET_ReallyLR implicit $q0
+    %0:gpr(s32) = COPY $w0
+    %1:fpr(<4 x s32>) = COPY $q0
+    %3:gpr(s32) = G_CONSTANT i32 1
+    %2:fpr(<4 x s32>) = G_INSERT_VECTOR_ELT %1, %0(s32), %3(s32)
+    $q0 = COPY %2(<4 x s32>)
+    RET_ReallyLR implicit $q0
+
+...
+---
+name:            v2s64_fpr
+alignment:       2
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $d0, $q1
+
+    ; CHECK-LABEL: name: v2s64_fpr
+    ; CHECK: liveins: $d0, $q1
+    ; CHECK: [[COPY:%[0-9]+]]:fpr64 = COPY $d0
+    ; CHECK: [[COPY1:%[0-9]+]]:fpr128 = COPY $q1
+    ; CHECK: [[DEF:%[0-9]+]]:fpr128 = IMPLICIT_DEF
+    ; CHECK: [[INSERT_SUBREG:%[0-9]+]]:fpr128 = INSERT_SUBREG [[DEF]], [[COPY]], %subreg.dsub
+    ; CHECK: [[INSvi64lane:%[0-9]+]]:fpr128 = INSvi64lane [[COPY1]], 1, [[INSERT_SUBREG]], 0
+    ; CHECK: $q0 = COPY [[INSvi64lane]]
+    ; CHECK: RET_ReallyLR implicit $q0
+    %0:fpr(s64) = COPY $d0
+    %1:fpr(<2 x s64>) = COPY $q1
+    %3:gpr(s32) = G_CONSTANT i32 1
+    %2:fpr(<2 x s64>) = G_INSERT_VECTOR_ELT %1, %0(s64), %3(s32)
+    $q0 = COPY %2(<2 x s64>)
+    RET_ReallyLR implicit $q0
+
+...
+---
+name:            v2s64_gpr
+alignment:       2
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $q0, $x0
+
+    ; CHECK-LABEL: name: v2s64_gpr
+    ; CHECK: liveins: $q0, $x0
+    ; CHECK: [[COPY:%[0-9]+]]:gpr64 = COPY $x0
+    ; CHECK: [[COPY1:%[0-9]+]]:fpr128 = COPY $q0
+    ; CHECK: [[INSvi64gpr:%[0-9]+]]:fpr128 = INSvi64gpr [[COPY1]], 0, [[COPY]]
+    ; CHECK: $q0 = COPY [[INSvi64gpr]]
+    ; CHECK: RET_ReallyLR implicit $q0
+    %0:gpr(s64) = COPY $x0
+    %1:fpr(<2 x s64>) = COPY $q0
+    %3:gpr(s32) = G_CONSTANT i32 0
+    %2:fpr(<2 x s64>) = G_INSERT_VECTOR_ELT %1, %0(s64), %3(s32)
+    $q0 = COPY %2(<2 x s64>)
+    RET_ReallyLR implicit $q0
+
+...