[AArch64] Recognize legal add immediates for addvl

Author: huntergr-arm

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -16595,6 +16595,21 @@ bool AArch64TargetLowering::isLegalAddImmediate(int64_t Immed) const {
   return IsLegal;
 }
 
+bool AArch64TargetLowering::isLegalAddScalableImmediate(int64_t Imm) const {
+  // Scalable immediates require SVE support.
+  if (!Subtarget->hasSVE())
+    return false;
+
+  // addvl's immediates are in terms of the number of bytes in a register.
+  // Since there are 16 in the base supported size (128bits), we need to
+  // divide the immediate by that much to give us a useful immediate to
+  // multiply by vscale. We can't have a remainder as a result of this.
+  if (Imm % 16 != 0)
+    return false;
+
+  return isInt<6>(Imm / 16);
+}
+
 // Return false to prevent folding
 // (mul (add x, c1), c2) -> (add (mul x, c2), c2*c1) in DAGCombine,
 // if the folding leads to worse code.

llvm/lib/Target/AArch64/AArch64ISelLowering.h

@@ -689,6 +689,7 @@ class AArch64TargetLowering : public TargetLowering {
                                        StoreInst *SI) const override;
 
   bool isLegalAddImmediate(int64_t) const override;
+  bool isLegalAddScalableImmediate(int64_t) const override;
   bool isLegalICmpImmediate(int64_t) const override;
 
   bool isMulAddWithConstProfitable(SDValue AddNode,

llvm/unittests/Target/AArch64/Immediates.cpp

@@ -20,19 +20,19 @@ struct TestCase {
 const std::initializer_list<TestCase> Tests = {
     // ScalableImm, Result
     // No change, easily 'supported'
-    {0, false},
+    {0, true},
 
     // addvl increments by whole registers, range [-32,31]
     // +(16 * vscale), one register's worth
-    {16, false},
+    {16, true},
     // +(8 * vscale), half a register's worth
     {8, false},
     // -(32 * 16 * vscale)
-    {-512, false},
+    {-512, true},
     // -(33 * 16 * vscale)
     {-528, false},
     // +(31 * 16 * vscale)
-    {496, false},
+    {496, true},
     // +(32 * 16 * vscale)
     {512, false},
 };