[AArch64][GlobalISel] Fix a crash when selecting shufflevectors with an undef mask element.

If a shufflevector's mask vector has an element with "undef" then the generic
instruction defining that element register is a G_IMPLICT_DEF instead of G_CONSTANT.
This fixes the selector to handle this case, and for now assumes that undef just means
zero. In future we'll optimize this case properly.

llvm-svn: 358312

2 files changed, 68 insertions, 7 deletions

diff --git a/llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp b/llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp
index a756d6b19da..2a7f2181951 100644
--- a/llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp
+++ b/llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp
@@ -93,7 +93,7 @@ private:
   bool selectUnmergeValues(MachineInstr &I, MachineRegisterInfo &MRI) const;
 
   void collectShuffleMaskIndices(MachineInstr &I, MachineRegisterInfo &MRI,
-                                 SmallVectorImpl<int> &Idxs) const;
+                                 SmallVectorImpl<Optional<int>> &Idxs) const;
   bool selectShuffleVector(MachineInstr &I, MachineRegisterInfo &MRI) const;
   bool selectExtractElt(MachineInstr &I, MachineRegisterInfo &MRI) const;
   bool selectConcatVectors(MachineInstr &I, MachineRegisterInfo &MRI) const;
@@ -2430,7 +2430,7 @@ bool AArch64InstructionSelector::selectConcatVectors(
 
 void AArch64InstructionSelector::collectShuffleMaskIndices(
     MachineInstr &I, MachineRegisterInfo &MRI,
-    SmallVectorImpl<int> &Idxs) const {
+    SmallVectorImpl<Optional<int>> &Idxs) const {
   MachineInstr *MaskDef = MRI.getVRegDef(I.getOperand(3).getReg());
   assert(
       MaskDef->getOpcode() == TargetOpcode::G_BUILD_VECTOR &&
@@ -2444,8 +2444,13 @@ void AArch64InstructionSelector::collectShuffleMaskIndices(
       ScalarDef = MRI.getVRegDef(ScalarDef->getOperand(1).getReg());
       assert(ScalarDef && "Could not find def of copy operand");
     }
-    assert(ScalarDef->getOpcode() == TargetOpcode::G_CONSTANT);
-    Idxs.push_back(ScalarDef->getOperand(1).getCImm()->getSExtValue());
+    if (ScalarDef->getOpcode() != TargetOpcode::G_CONSTANT) {
+      // This be an undef if not a constant.
+      assert(ScalarDef->getOpcode() == TargetOpcode::G_IMPLICIT_DEF);
+      Idxs.push_back(None);
+    } else {
+      Idxs.push_back(ScalarDef->getOperand(1).getCImm()->getSExtValue());
+    }
   }
 }
 
@@ -2692,8 +2697,10 @@ bool AArch64InstructionSelector::selectShuffleVector(
 
   // G_SHUFFLE_VECTOR doesn't really have a strictly enforced constant mask
   // operand, it comes in as a normal vector value which we have to analyze to
-  // find the mask indices.
-  SmallVector<int, 8> Mask;
+  // find the mask indices. If the mask element is undef, then
+  // collectShuffleMaskIndices() will add a None entry for that index into
+  // the list.
+  SmallVector<Optional<int>, 8> Mask;
   collectShuffleMaskIndices(I, MRI, Mask);
   assert(!Mask.empty() && "Expected to find mask indices");
 
@@ -2708,7 +2715,10 @@ bool AArch64InstructionSelector::selectShuffleVector(
   unsigned BytesPerElt = DstTy.getElementType().getSizeInBits() / 8;
 
   SmallVector<Constant *, 64> CstIdxs;
-  for (int Val : Mask) {
+  for (auto &MaybeVal : Mask) {
+    // For now, any undef indexes we'll just assume to be 0. This should be
+    // optimized in future, e.g. to select DUP etc.
+    int Val = MaybeVal.hasValue() ? *MaybeVal : 0;
     for (unsigned Byte = 0; Byte < BytesPerElt; ++Byte) {
       unsigned Offset = Byte + Val * BytesPerElt;
       CstIdxs.emplace_back(ConstantInt::get(Type::getInt8Ty(Ctx), Offset));
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/select-shufflevec-undef-mask-elt.mir b/llvm/test/CodeGen/AArch64/GlobalISel/select-shufflevec-undef-mask-elt.mir
new file mode 100644
index 00000000000..7a973b8804e
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/select-shufflevec-undef-mask-elt.mir
@@ -0,0 +1,51 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple=aarch64-- -O0 -run-pass=instruction-select -verify-machineinstrs %s -global-isel-abort=1 -o - | FileCheck %s
+
+# This test checks that a shuffle mask with an undef value, instead of a constant,
+# doesn't crash. The code generated definitely isn't optimal.
+...
+---
+name:            shuffle_undef_mask_elt
+alignment:       2
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+machineFunctionInfo: {}
+body:             |
+  bb.1:
+    liveins: $d0
+
+    ; CHECK-LABEL: name: shuffle_undef_mask_elt
+    ; CHECK: liveins: $d0
+    ; CHECK: [[COPY:%[0-9]+]]:fpr64 = COPY $d0
+    ; CHECK: [[DEF:%[0-9]+]]:gpr32 = IMPLICIT_DEF
+    ; CHECK: [[DEF1:%[0-9]+]]:gpr32 = IMPLICIT_DEF
+    ; CHECK: [[DEF2:%[0-9]+]]:fpr128 = IMPLICIT_DEF
+    ; CHECK: [[INSERT_SUBREG:%[0-9]+]]:fpr128 = INSERT_SUBREG [[DEF2]], [[DEF]], %subreg.ssub
+    ; CHECK: [[INSvi32gpr:%[0-9]+]]:fpr128 = INSvi32gpr [[INSERT_SUBREG]], 1, [[DEF1]]
+    ; CHECK: [[COPY1:%[0-9]+]]:fpr64 = COPY [[INSvi32gpr]].dsub
+    ; CHECK: [[ADRP:%[0-9]+]]:gpr64common = ADRP target-flags(aarch64-page) %const.0
+    ; CHECK: [[LDRDui:%[0-9]+]]:fpr64 = LDRDui [[ADRP]], target-flags(aarch64-pageoff, aarch64-nc) %const.0
+    ; CHECK: [[DEF3:%[0-9]+]]:fpr128 = IMPLICIT_DEF
+    ; CHECK: [[INSERT_SUBREG1:%[0-9]+]]:fpr128 = INSERT_SUBREG [[DEF3]], [[COPY]], %subreg.dsub
+    ; CHECK: [[DEF4:%[0-9]+]]:fpr128 = IMPLICIT_DEF
+    ; CHECK: [[INSERT_SUBREG2:%[0-9]+]]:fpr128 = INSERT_SUBREG [[DEF4]], [[COPY1]], %subreg.dsub
+    ; CHECK: [[INSvi64lane:%[0-9]+]]:fpr128 = INSvi64lane [[INSERT_SUBREG1]], 1, [[INSERT_SUBREG2]], 0
+    ; CHECK: [[DEF5:%[0-9]+]]:fpr128 = IMPLICIT_DEF
+    ; CHECK: [[INSERT_SUBREG3:%[0-9]+]]:fpr128 = INSERT_SUBREG [[DEF5]], [[LDRDui]], %subreg.dsub
+    ; CHECK: [[TBLv16i8One:%[0-9]+]]:fpr128 = TBLv16i8One [[INSvi64lane]], [[INSERT_SUBREG3]]
+    ; CHECK: [[COPY2:%[0-9]+]]:fpr64 = COPY [[TBLv16i8One]].dsub
+    ; CHECK: $d0 = COPY [[COPY2]]
+    ; CHECK: RET_ReallyLR implicit $d0
+    %0:fpr(<2 x s32>) = COPY $d0
+    %6:gpr(s32) = G_IMPLICIT_DEF
+    %7:gpr(s32) = G_IMPLICIT_DEF
+    %2:fpr(<2 x s32>) = G_BUILD_VECTOR %6(s32), %7(s32)
+    %4:gpr(s32) = G_CONSTANT i32 1
+    %5:gpr(s32) = G_IMPLICIT_DEF
+    %3:fpr(<2 x s32>) = G_BUILD_VECTOR %4(s32), %5(s32)
+    %1:fpr(<2 x s32>) = G_SHUFFLE_VECTOR %0(<2 x s32>), %2, %3(<2 x s32>)
+    $d0 = COPY %1(<2 x s32>)
+    RET_ReallyLR implicit $d0
+
+...