[AArch64] Support scalable offsets with isLegalAddressingMode

Given a base register and a scalable offset (multiplied by vscale),
return true if the offset corresponds to the valid range for the
size of the vector type in memory; e.g. `[X0, #1, mul vl]`

Author: huntergr-arm

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -16671,6 +16671,16 @@ bool AArch64TargetLowering::isLegalAddressingMode(const DataLayout &DL,
 
   if (Ty->isScalableTy()) {
     if (isa<ScalableVectorType>(Ty)) {
+      // See if we have a foldable vscale-based offset, for vector types which
+      // are either legal or smaller than the minimum; more work will be
+      // required if we need to consider addressing for types which need
+      // legalization by splitting.
+      uint64_t VecNumBytes = DL.getTypeSizeInBits(Ty).getKnownMinValue() / 8;
+      if (AM.HasBaseReg && !AM.BaseOffs && AM.ScalableOffset && !AM.Scale &&
+          (AM.ScalableOffset % VecNumBytes == 0) && VecNumBytes <= 16 &&
+          isPowerOf2_64(VecNumBytes))
+        return isInt<4>(AM.ScalableOffset / (int64_t)VecNumBytes);
+
       uint64_t VecElemNumBytes =
           DL.getTypeSizeInBits(cast<VectorType>(Ty)->getElementType()) / 8;
       return AM.HasBaseReg && !AM.BaseOffs && !AM.ScalableOffset &&
@@ -16680,6 +16690,10 @@ bool AArch64TargetLowering::isLegalAddressingMode(const DataLayout &DL,
     return AM.HasBaseReg && !AM.BaseOffs && !AM.ScalableOffset && !AM.Scale;
   }
 
+  // No scalable offsets allowed for non-scalable types.
+  if (AM.ScalableOffset)
+    return false;
+
   // check reg + imm case:
   // i.e., reg + 0, reg + imm9, reg + SIZE_IN_BYTES * uimm12
   uint64_t NumBytes = 0;

llvm/unittests/Target/AArch64/AddressingModes.cpp

@@ -167,20 +167,20 @@ const std::initializer_list<SVETestCase> SVETests = {
     // {BaseGV, BaseOffs, HasBaseReg, Scale, SOffs}, EltBits, Count, Result
     // Test immediate range -- [-8,7] vector's worth.
     // <vscale x 16 x i8>, increment by one vector
-    {{nullptr, 0, true, 0, 16}, 8, 16, false},
+    {{nullptr, 0, true, 0, 16}, 8, 16, true},
     // <vscale x 4 x i32>, increment by eight vectors
     {{nullptr, 0, true, 0, 128}, 32, 4, false},
     // <vscale x 8 x i16>, increment by seven vectors
-    {{nullptr, 0, true, 0, 112}, 16, 8, false},
+    {{nullptr, 0, true, 0, 112}, 16, 8, true},
     // <vscale x 2 x i64>, decrement by eight vectors
-    {{nullptr, 0, true, 0, -128}, 64, 2, false},
+    {{nullptr, 0, true, 0, -128}, 64, 2, true},
     // <vscale x 16 x i8>, decrement by nine vectors
     {{nullptr, 0, true, 0, -144}, 8, 16, false},
 
     // Half the size of a vector register, but allowable with extending
     // loads and truncating stores
     // <vscale x 8 x i8>, increment by three vectors
-    {{nullptr, 0, true, 0, 24}, 8, 8, false},
+    {{nullptr, 0, true, 0, 24}, 8, 8, true},
 
     // Test invalid types or offsets
     // <vscale x 5 x i32>, increment by one vector (base size > 16B)