[AArch64] Support optional constant offset for constraint "S" (#80255)

Modify the initial implementation (https://reviews.llvm.org/D46745) to
support a constant offset so that the following code will compile:
```
int a[2][2];
void foo() { asm("// %0" :: "S"(&a[1][1])); }
```

We use the generic code path for "s". In GCC's aarch64 port, "S" is
supported for PIC while "s" isn't, making "s" less useful. We implement
"S" but not "s".

Similar to #80201 for RISC-V.

3 files changed, 12 insertions, 13 deletions

diff --git a/llvm/docs/LangRef.rst b/llvm/docs/LangRef.rst
index 3648ea2..96bdba5 100644
--- a/llvm/docs/LangRef.rst
+++ b/llvm/docs/LangRef.rst
@@ -5107,6 +5107,8 @@ AArch64:
   offsets). (However, LLVM currently does this for the ``m`` constraint as
   well.)
 - ``r``: A 32 or 64-bit integer register (W* or X*).
+- ``S``: A symbol or label reference with a constant offset. The generic ``s``
+  is not supported.
 - ``Uci``: Like r, but restricted to registers 8 to 11 inclusive.
 - ``Ucj``: Like r, but restricted to registers 12 to 15 inclusive.
 - ``w``: A 32, 64, or 128-bit floating-point, SIMD or SVE vector register.
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index a7d1592..b59f8d7 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -10618,7 +10618,7 @@ AArch64TargetLowering::getConstraintType(StringRef Constraint) const {
     case 'Z':
       return C_Immediate;
     case 'z':
-    case 'S': // A symbolic address
+    case 'S': // A symbol or label reference with a constant offset
       return C_Other;
     }
   } else if (parsePredicateConstraint(Constraint))
@@ -10809,19 +10809,12 @@ void AArch64TargetLowering::LowerAsmOperandForConstraint(
       Result = DAG.getRegister(AArch64::WZR, MVT::i32);
     break;
   }
-  case 'S': {
-    // An absolute symbolic address or label reference.
-    if (const GlobalAddressSDNode *GA = dyn_cast<GlobalAddressSDNode>(Op)) {
-      Result = DAG.getTargetGlobalAddress(GA->getGlobal(), SDLoc(Op),
-                                          GA->getValueType(0));
-    } else if (const BlockAddressSDNode *BA =
-                   dyn_cast<BlockAddressSDNode>(Op)) {
-      Result =
-          DAG.getTargetBlockAddress(BA->getBlockAddress(), BA->getValueType(0));
-    } else
-      return;
+  case 'S':
+    // Use the generic code path for "s". In GCC's aarch64 port, "S" is
+    // supported for PIC while "s" isn't, making "s" less useful. We implement
+    // "S" but not "s".
+    TargetLowering::LowerAsmOperandForConstraint(Op, "s", Ops, DAG);
     break;
-  }
 
   case 'I':
   case 'J':
diff --git a/llvm/test/CodeGen/AArch64/inlineasm-S-constraint.ll b/llvm/test/CodeGen/AArch64/inlineasm-S-constraint.ll
index 18b4ef9..7664ad1 100644
--- a/llvm/test/CodeGen/AArch64/inlineasm-S-constraint.ll
+++ b/llvm/test/CodeGen/AArch64/inlineasm-S-constraint.ll
@@ -1,11 +1,15 @@
 ;RUN: llc -mtriple=aarch64-none-linux-gnu -mattr=+neon < %s | FileCheck %s
 @var = global i32 0
+@a = external global [2 x [2 x i32]], align 4
+
 define void @test_inline_constraint_S() {
 ; CHECK-LABEL: test_inline_constraint_S:
   call void asm sideeffect "adrp x0, $0", "S"(ptr @var)
   call void asm sideeffect "add x0, x0, :lo12:$0", "S"(ptr @var)
+  call void asm sideeffect "// $0", "S"(ptr getelementptr inbounds ([2 x [2 x i32]], ptr @a, i64 0, i64 1, i64 1))
 ; CHECK: adrp x0, var
 ; CHECK: add x0, x0, :lo12:var
+; CHECK: // a+12
   ret void
 }
 define i32 @test_inline_constraint_S_label(i1 %in) {