[AArch64] Implement "rZ" inline asm constraint

miloserdow · miloserdow · commit 9482083bb6cb · 2025-11-01T21:40:02.000Z
Add support for the "rZ" inline assembly constraint.
The constraint accepts literal zero values and emits the arch zero
register (xzr/wzr) instead of materializing zero in a gpr.

In AArch64.cpp:
  - validateAsmConstraint: recognize "rZ"/"rz" constraint
  - convertConstraint: convert to "^rZ"

In AArch64ISelLowering.cpp:
  - getConstraintType: return C_RegisterClass
  - getRegForInlineAsmConstraint: return appropriate register class based
    on value type (GPR32/GPR64/GPR64x8 for LS64, reject scalable vectors)
  - LowerAsmOperandForConstraint: substitute XZR/WZR for literal zero
    values
diff --git a/clang/lib/Basic/Targets/AArch64.cpp b/clang/lib/Basic/Targets/AArch64.cpp
@@ -1495,6 +1495,17 @@ std::string
 AArch64TargetInfo::convertConstraint(const char *&Constraint) const {
   std::string R;
   switch (*Constraint) {
+  case 'r':
+    // Check for "rZ" or "rz" constraint (register or zero)
+    if (Constraint[1] == 'Z' || Constraint[1] == 'z') {
+      // Return with "^" prefix to indicate 2-character constraint
+      R = "^r";
+      R += Constraint[1];
+      Constraint += 1;
+      return R;
+    }
+    R = TargetInfo::convertConstraint(Constraint);
+    break;
   case 'U': // Three-character constraint; add "@3" hint for later parsing.
     R = std::string("@3") + std::string(Constraint, 3);
     Constraint += 2;
@@ -1518,6 +1529,14 @@ bool AArch64TargetInfo::validateAsmConstraint(
   switch (*Name) {
   default:
     return false;
+  case 'r':
+    // Check if this is "rZ" constraint (register or zero)
+    if (Name[1] == 'Z' || Name[1] == 'z') {
+      Info.setAllowsRegister();
+      Name++;
+      return true;
+    }
+    return false;
   case 'w': // Floating point and SIMD registers (V0-V31)
     Info.setAllowsRegister();
     return true;
diff --git a/clang/test/CodeGen/aarch64-inline-asm-rZ-constraint-error.c b/clang/test/CodeGen/aarch64-inline-asm-rZ-constraint-error.c
@@ -0,0 +1,19 @@
+// RUN: not %clang_cc1 -triple aarch64-linux-gnu -O2 -S -o /dev/null %s 2>&1 | FileCheck %s
+
+// Test that the "rZ" inline assembly constraint properly rejects non-constant values.
+// The "rZ" constraint is only valid for literal zero values.
+
+// CHECK: error: invalid operand for inline asm constraint 'rZ'
+void test_rZ_runtime_value(long *addr, long val) {
+    __asm__ volatile("str %1, [%0]" : : "r"(addr), "rZ"(val));
+}
+
+// CHECK: error: invalid operand for inline asm constraint 'rZ'
+void test_rZ_runtime_i32(int *addr, int val) {
+    __asm__ volatile("str %w1, [%0]" : : "r"(addr), "rZ"(val));
+}
+
+// CHECK: error: invalid operand for inline asm constraint 'rZ'
+void test_rZ_non_constant(long val) {
+    __asm__ volatile("mov x2, %0" : : "rZ"(val));
+}
diff --git a/clang/test/CodeGen/aarch64-inline-asm-rZ-constraint.c b/clang/test/CodeGen/aarch64-inline-asm-rZ-constraint.c
@@ -0,0 +1,86 @@
+// RUN: %clang_cc1 -triple aarch64-linux-gnu -O2 -emit-llvm -o - %s | FileCheck %s --check-prefix=CHECK-IR
+// RUN: %clang_cc1 -triple aarch64-linux-gnu -O2 -S -o - %s | FileCheck %s --check-prefix=CHECK-ASM
+
+// Test the "rZ" inline assembly constraint for AArch64.
+
+// CHECK-IR-LABEL: define dso_local void @test_rZ_zero_i64(
+// CHECK-IR: tail call void asm sideeffect "str $1, [$0]", "r,^rZ"(ptr %addr, i64 0)
+//
+// CHECK-ASM-LABEL: test_rZ_zero_i64:
+// CHECK-ASM: str xzr, [x0]
+void test_rZ_zero_i64(long *addr) {
+    __asm__ volatile("str %1, [%0]" : : "r"(addr), "rZ"(0L));
+}
+
+// CHECK-IR-LABEL: define dso_local void @test_rZ_zero_i32(
+// CHECK-IR: tail call void asm sideeffect "str ${1:w}, [$0]", "r,^rZ"(ptr %addr, i32 0)
+//
+// CHECK-ASM-LABEL: test_rZ_zero_i32:
+// CHECK-ASM: str wzr, [x0]
+void test_rZ_zero_i32(int *addr) {
+    __asm__ volatile("str %w1, [%0]" : : "r"(addr), "rZ"(0));
+}
+
+// CHECK-IR-LABEL: define dso_local void @test_rZ_zero_i16(
+// CHECK-IR: tail call void asm sideeffect "strh ${1:w}, [$0]", "r,^rZ"(ptr %addr, i16 0)
+//
+// CHECK-ASM-LABEL: test_rZ_zero_i16:
+// CHECK-ASM: strh wzr, [x0]
+void test_rZ_zero_i16(short *addr) {
+    __asm__ volatile("strh %w1, [%0]" : : "r"(addr), "rZ"((short)0));
+}
+
+// CHECK-IR-LABEL: define dso_local void @test_rZ_zero_i8(
+// CHECK-IR: tail call void asm sideeffect "strb ${1:w}, [$0]", "r,^rZ"(ptr %addr, i8 0)
+//
+// CHECK-ASM-LABEL: test_rZ_zero_i8:
+// CHECK-ASM: strb wzr, [x0]
+void test_rZ_zero_i8(char *addr) {
+    __asm__ volatile("strb %w1, [%0]" : : "r"(addr), "rZ"((char)0));
+}
+
+// CHECK-IR-LABEL: define dso_local void @test_rz_lowercase(
+// CHECK-IR: tail call void asm sideeffect "str $1, [$0]", "r,^rz"(ptr %addr, i64 0)
+//
+// CHECK-ASM-LABEL: test_rz_lowercase:
+// CHECK-ASM: str xzr, [x0]
+void test_rz_lowercase(long *addr) {
+    __asm__ volatile("str %1, [%0]" : : "r"(addr), "rz"(0L));
+}
+
+// CHECK-IR-LABEL: define dso_local void @test_rZ_explicit_x(
+// CHECK-IR: tail call void asm sideeffect "mov ${0:x}, xzr", "^rZ"(i64 0)
+//
+// CHECK-ASM-LABEL: test_rZ_explicit_x:
+// CHECK-ASM: mov xzr, xzr
+void test_rZ_explicit_x(void) {
+    __asm__ volatile("mov %x0, xzr" : : "rZ"(0L));
+}
+
+// CHECK-IR-LABEL: define dso_local void @test_rZ_explicit_w(
+// CHECK-IR: tail call void asm sideeffect "mov ${0:w}, wzr", "^rZ"(i32 0)
+//
+// CHECK-ASM-LABEL: test_rZ_explicit_w:
+// CHECK-ASM: mov wzr, wzr
+void test_rZ_explicit_w(void) {
+    __asm__ volatile("mov %w0, wzr" : : "rZ"(0));
+}
+
+// CHECK-IR-LABEL: define dso_local void @test_rZ_x_modifier(
+// CHECK-IR: tail call void asm sideeffect "add x2, x1, ${0:x}", "^rZ"(i64 0)
+//
+// CHECK-ASM-LABEL: test_rZ_x_modifier:
+// CHECK-ASM: add x2, x1, xzr
+void test_rZ_x_modifier(void) {
+    __asm__ volatile("add x2, x1, %x0" : : "rZ"(0L));
+}
+
+// CHECK-IR-LABEL: define dso_local void @test_rZ_w_modifier(
+// CHECK-IR: tail call void asm sideeffect "add w2, w1, ${0:w}", "^rZ"(i32 0)
+//
+// CHECK-ASM-LABEL: test_rZ_w_modifier:
+// CHECK-ASM: add w2, w1, wzr
+void test_rZ_w_modifier(void) {
+    __asm__ volatile("add w2, w1, %w0" : : "rZ"(0));
+}
+
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -13018,6 +13018,9 @@ AArch64TargetLowering::getConstraintType(StringRef Constraint) const {
     case 'S': // A symbol or label reference with a constant offset
       return C_Other;
     }
+  } else if (Constraint.size() == 2 &&
+             (Constraint == "rZ" || Constraint == "rz")) {
+    return C_RegisterClass;
   } else if (parsePredicateConstraint(Constraint))
     return C_RegisterClass;
   else if (parseReducedGprConstraint(Constraint))
@@ -13045,6 +13048,14 @@ AArch64TargetLowering::getSingleConstraintMatchWeight(
   default:
     weight = TargetLowering::getSingleConstraintMatchWeight(info, constraint);
     break;
+  case 'r':
+    // Check for "rZ" or "rz" constraint (register or zero)
+    if (constraint[1] == 'Z' || constraint[1] == 'z') {
+      weight = CW_Register;
+      break;
+    }
+    weight = TargetLowering::getSingleConstraintMatchWeight(info, constraint);
+    break;
   case 'x':
   case 'w':
   case 'y':
@@ -13066,6 +13077,18 @@ AArch64TargetLowering::getSingleConstraintMatchWeight(
 std::pair<unsigned, const TargetRegisterClass *>
 AArch64TargetLowering::getRegForInlineAsmConstraint(
     const TargetRegisterInfo *TRI, StringRef Constraint, MVT VT) const {
+  // Handle "rZ" and "rz" constraints
+  if (Constraint.size() == 2 && Constraint[0] == 'r' &&
+      (Constraint[1] == 'Z' || Constraint[1] == 'z')) {
+    if (VT.isScalableVector())
+      return std::make_pair(0U, nullptr);
+    if (Subtarget->hasLS64() && VT.getSizeInBits() == 512)
+      return std::make_pair(0U, &AArch64::GPR64x8ClassRegClass);
+    if (VT.getFixedSizeInBits() == 64)
+      return std::make_pair(0U, &AArch64::GPR64commonRegClass);
+    return std::make_pair(0U, &AArch64::GPR32commonRegClass);
+  }
+
   if (Constraint.size() == 1) {
     switch (Constraint[0]) {
     case 'r':
@@ -13196,6 +13219,23 @@ void AArch64TargetLowering::LowerAsmOperandForConstraint(
     SelectionDAG &DAG) const {
   SDValue Result;
 
+  // Handle "rZ" and "rz" constraints (register or zero)
+  if (Constraint.size() == 2 && Constraint[0] == 'r' &&
+      (Constraint[1] == 'Z' || Constraint[1] == 'z')) {
+    if (isNullConstant(Op)) {
+      if (Op.getValueType() == MVT::i64)
+        Result = DAG.getRegister(AArch64::XZR, MVT::i64);
+      else
+        Result = DAG.getRegister(AArch64::WZR, MVT::i32);
+
+      if (Result.getNode()) {
+        Ops.push_back(Result);
+        return;
+      }
+    }
+    return;
+  }
+
   // Currently only support length 1 constraints.
   if (Constraint.size() != 1)
     return;
