[AArch64][GlobalISel] Fix assertion fail in C++ selection for vector zext of <4 x s8>

We bailed out of dealing with vectors only after the assertion, move it before.

Fixes PR43794

2 files changed, 12 insertions, 2 deletions

diff --git a/llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp b/llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp
index 961f38cad1e..c9d4654ade0 100644
--- a/llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp
+++ b/llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp
@@ -2065,14 +2065,15 @@ bool AArch64InstructionSelector::select(MachineInstr &I) {
     unsigned DstSize = DstTy.getSizeInBits();
     unsigned SrcSize = SrcTy.getSizeInBits();
 
+    if (DstTy.isVector())
+      return false; // Should be handled by imported patterns.
+
     assert((*RBI.getRegBank(DefReg, MRI, TRI)).getID() ==
                AArch64::GPRRegBankID &&
            "Unexpected ext regbank");
 
     MachineIRBuilder MIB(I);
     MachineInstr *ExtI;
-    if (DstTy.isVector())
-      return false; // Should be handled by imported patterns.
 
     // First check if we're extending the result of a load which has a dest type
     // smaller than 32 bits, then this zext is redundant. GPR32 is the smallest
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/arm64-fallback.ll b/llvm/test/CodeGen/AArch64/GlobalISel/arm64-fallback.ll
index cd247718eba..dab44b84355 100644
--- a/llvm/test/CodeGen/AArch64/GlobalISel/arm64-fallback.ll
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/arm64-fallback.ll
@@ -200,3 +200,12 @@ entry:
   call void @use_s128(i128 %p2, i128 %p6)
   ret i32 0
 }
+
+; FALLBACK-WITH-REPORT-ERR: remark: <unknown>:0:0: cannot select: %2:fpr(<4 x s16>) = G_ZEXT %0:fpr(<4 x s8>) (in function: zext_v4s8)
+; FALLBACK-WITH-REPORT-ERR: warning: Instruction selection used fallback path for zext_v4s8
+; FALLBACK-WITH-REPORT-OUT-LABEL: zext_v4s8
+define <4 x i16> @zext_v4s8(<4 x i8> %in) {
+  %ext = zext <4 x i8> %in to <4 x i16>
+  ret <4 x i16> %ext
+}
+