[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.

4 files changed, 26 insertions, 16 deletions

diff --git a/clang/test/Sema/inline-asm-validate-aarch64.c b/clang/test/Sema/inline-asm-validate-aarch64.c
index 014767d..1e753d4 100644
--- a/clang/test/Sema/inline-asm-validate-aarch64.c
+++ b/clang/test/Sema/inline-asm-validate-aarch64.c
@@ -1,14 +1,24 @@
+// RUN: %clang_cc1 -triple aarch64 -fsyntax-only -verify -DVERIFY %s
 // RUN: %clang_cc1 -triple arm64-apple-darwin -fsyntax-only -fdiagnostics-parseable-fixits %s 2>&1 | FileCheck %s
 
 typedef unsigned char uint8_t;
 
+#ifdef VERIFY
+void test_s(int i) {
+  asm("" :: "s"(i)); // expected-error{{invalid input constraint 's' in asm}}
+
+  /// Codegen error
+  asm("" :: "S"(i));
+  asm("" :: "S"(test_s(i))); // expected-error{{invalid type 'void' in asm input for constraint 'S'}}
+}
+#else
 uint8_t constraint_r(uint8_t *addr) {
   uint8_t byte;
 
   __asm__ volatile("ldrb %0, [%1]" : "=r" (byte) : "r" (addr) : "memory");
 // CHECK: warning: value size does not match register size specified by the constraint and modifier
 // CHECK: note: use constraint modifier "w"
-// CHECK: fix-it:{{.*}}:{8:26-8:28}:"%w0"
+// CHECK: fix-it:{{.*}}:{[[#@LINE-3]]:26-[[#@LINE-3]]:28}:"%w0"
 
   return byte;
 }
@@ -19,7 +29,7 @@ uint8_t constraint_r_symbolic(uint8_t *addr) {
   __asm__ volatile("ldrb %[s0], [%[s1]]" : [s0] "=r" (byte) : [s1] "r" (addr) : "memory");
 // CHECK: warning: value size does not match register size specified by the constraint and modifier
 // CHECK: note: use constraint modifier "w"
-// CHECK: fix-it:{{.*}}:{19:26-19:31}:"%w[s0]"
+// CHECK: fix-it:{{.*}}:{[[#@LINE-3]]:26-[[#@LINE-3]]:31}:"%w[s0]"
 
   return byte;
 }
@@ -40,11 +50,11 @@ uint8_t constraint_r_symbolic_macro(uint8_t *addr) {
 // CHECK: warning: value size does not match register size specified by the constraint and modifier
 // CHECK: asm ("%w0 %w1 %2" : "+r" (one) : "r" (wide_two));
 // CHECK: note: use constraint modifier "w"
-// CHECK: fix-it:{{.*}}:{47:17-47:19}:"%w2"
 
 void read_write_modifier0(int one, int two) {
   long wide_two = two;
   asm ("%w0 %w1 %2" : "+r" (one) : "r" (wide_two));
+// CHECK: fix-it:{{.*}}:{[[#@LINE-1]]:17-[[#@LINE-1]]:19}:"%w2"
 }
 
 // CHECK-NOT: warning: 
@@ -52,3 +62,4 @@ void read_write_modifier1(int one, int two) {
   long wide_two = two;
   asm ("%w0 %1" : "+r" (one), "+r" (wide_two));
 }
+#endif
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) {