[SVE] Fix crashes with inline assembly

[david-arm]

Summary

All the crashes found compiling inline assembly are fixed in this
patch by changing AArch64TargetLowering::getRegForInlineAsmConstraint
to be more resilient to mismatched value and register types. For
example, it makes no sense to request a predicate register for
a nxv2i64 type and so on.

Tests have been added here:

test/CodeGen/AArch64/inline-asm-constraints-bad-sve.ll

[sdesmalen]

llvm/test/CodeGen/AArch64/inline-asm-constraints-bad-sve.ll

33

How is foo4 different from foo1?

[david-arm]

llvm/test/CodeGen/AArch64/inline-asm-constraints-bad-sve.ll

33

Fair point. The original C tests were different, but I hadn't spotted they'd boiled down to the same IR. I'll remove it.

[sdesmalen]

LGTM

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

7528

nit: is it worth rewriting this to:

TypeSize VTSize = VT.getSizeInBits();
if (VTSize == TypeSize::getScalable(1))
  return ...
if (VTSize == TypeSize::getFixed(16))
  return ...
if (VTSize == TypeSize::getFixed(32))
  return ...

if (VT.isScalableVector())
  return std::make_pair(0U, nullptr);
break;

7561

nit:

if (!VT.isScalableVector() || VT.getVectorElementType() != MVT::i1)
  return std::make_pair(0U, nullptr);
bool restricted = (PC == PredicateConstraint::Upl);
return restricted ? std::make_pair(0U, &AArch64::PPR_3bRegClass)
                  : std::make_pair(0U, &AArch64::PPRRegClass);

(that saves having to indent the block)

[david-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

7528

Hi Sander, I tried rewriting this and then I realised we can't check "VTSize == TypeSize::getScalable(1)" because this is looking for a size of "vscale x 1 byte", which isn't the same as a predicate. I felt it was more obvious to check the element type explicitly. Hope that's ok!

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

 }
 
 std::pair<unsigned, const TargetRegisterClass *>
 AArch64TargetLowering::getRegForInlineAsmConstraint(
     const TargetRegisterInfo *TRI, StringRef Constraint, MVT VT) const {
   if (Constraint.size() == 1) {
     switch (Constraint[0]) {
     case 'r':
-      if (VT.getSizeInBits() == 64)
+      if (VT.isScalableVector())
+        return std::make_pair(0U, nullptr);
+      if (VT.getFixedSizeInBits() == 64)
         return std::make_pair(0U, &AArch64::GPR64commonRegClass);
       return std::make_pair(0U, &AArch64::GPR32commonRegClass);
-    case 'w':
+    case 'w': {
       if (!Subtarget->hasFPARMv8())
         break;
-      if (VT.isScalableVector())
+      if (VT.isScalableVector()) {
+        if (VT.getVectorElementType() != MVT::i1)
           return std::make_pair(0U, &AArch64::ZPRRegClass);
-      if (VT.getSizeInBits() == 16)
+        return std::make_pair(0U, nullptr);
+      }
+      uint64_t VTSize = VT.getFixedSizeInBits();
+      if (VTSize == 16)
         return std::make_pair(0U, &AArch64::FPR16RegClass);
-      if (VT.getSizeInBits() == 32)
+      if (VTSize == 32)
         return std::make_pair(0U, &AArch64::FPR32RegClass);
-      if (VT.getSizeInBits() == 64)
+      if (VTSize == 64)
         return std::make_pair(0U, &AArch64::FPR64RegClass);
-      if (VT.getSizeInBits() == 128)
+      if (VTSize == 128)
         return std::make_pair(0U, &AArch64::FPR128RegClass);
       break;
+    }
     // The instructions that this constraint is designed for can
     // only take 128-bit registers so just use that regclass.
     case 'x':
       if (!Subtarget->hasFPARMv8())
         break;
       if (VT.isScalableVector())
         return std::make_pair(0U, &AArch64::ZPR_4bRegClass);
       if (VT.getSizeInBits() == 128)
         return std::make_pair(0U, &AArch64::FPR128_loRegClass);
       break;
     case 'y':
       if (!Subtarget->hasFPARMv8())
         break;
       if (VT.isScalableVector())
         return std::make_pair(0U, &AArch64::ZPR_3bRegClass);
       break;
     }
   } else {
     PredicateConstraint PC = parsePredicateConstraint(Constraint);
     if (PC != PredicateConstraint::Invalid) {
-      assert(VT.isScalableVector());
+      if (!VT.isScalableVector() || VT.getVectorElementType() != MVT::i1)
+        return std::make_pair(0U, nullptr);
       bool restricted = (PC == PredicateConstraint::Upl);
       return restricted ? std::make_pair(0U, &AArch64::PPR_3bRegClass)
-                          : std::make_pair(0U, &AArch64::PPRRegClass);
     }
   }
   if (StringRef("{cc}").equals_lower(Constraint))
     return std::make_pair(unsigned(AArch64::NZCV), &AArch64::CCRRegClass);
 
   // Use the default implementation in TargetLowering to convert the register
   // constraint into a member of a register class.
   std::pair<unsigned, const TargetRegisterClass *> Res;

llvm/test/CodeGen/AArch64/inline-asm-constraints-bad-sve.ll

+; RUN: not llc -mtriple=aarch64-none-linux-gnu -mattr=+sve -o - %s 2>&1 | FileCheck %s
+
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64-unknown-linux-gnu"
+
+; CHECK: error: couldn't allocate input reg for constraint 'Upa'
+; CHECK: error: couldn't allocate input reg for constraint 'r'
+; CHECK: error: couldn't allocate output register for constraint 'w'
+
+define <vscale x 16 x i1> @foo1(i32 *%in) {
+entry:
+  %0 = load i32, i32* %in, align 4
+  %1 = call <vscale x 16 x i1> asm sideeffect "mov $0.b, $1.b \0A", "=@3Upa,@3Upa"(i32 %0)
+  ret <vscale x 16 x i1> %1
+}
+
+define <vscale x 4 x float> @foo2(<vscale x 4 x i32> *%in) {
+entry:
+  %0 = load <vscale x 4 x i32>, <vscale x 4 x i32>* %in, align 16
+  %1 = call <vscale x 4 x float> asm sideeffect "ptrue p0.s, #1 \0Afabs $0.s, p0/m, $1.s \0A", "=w,r"(<vscale x 4 x i32> %0)
+  ret <vscale x 4 x float> %1
+}
+
+define <vscale x 16 x i1> @foo3(<vscale x 16 x i1> *%in) {
+entry:
+  %0 = load <vscale x 16 x i1>, <vscale x 16 x i1>* %in, align 2
+  %1 = call <vscale x 16 x i1> asm sideeffect "mov $0.b, $1.b \0A", "=&w,w"(<vscale x 16 x i1> %0)
+  ret <vscale x 16 x i1> %1
+}