[llvm] [AArch64] Support scalable offsets with isLegalAddressingMode (PR #83255)

[David Green via llvm-commits]

================
@@ -16374,15 +16374,31 @@ 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)) {
+        int64_t Idx = AM.ScalableOffset / (int64_t)VecNumBytes;
+        return Idx >= -8 && Idx <= 7;
+      }
+
       uint64_t VecElemNumBytes =
           DL.getTypeSizeInBits(cast<VectorType>(Ty)->getElementType()) / 8;
-      return AM.HasBaseReg && !AM.BaseOffs &&
+      return AM.HasBaseReg && !AM.BaseOffs && !AM.ScalableOffset &&
              (AM.Scale == 0 || (uint64_t)AM.Scale == VecElemNumBytes);
     }
 
     return AM.HasBaseReg && !AM.BaseOffs && !AM.Scale;
----------------
davemgreen wrote:

This could add a check for !AM.ScalableOffset too, in case it is a struct type.

https://github.com/llvm/llvm-project/pull/83255