[AArch64] Implement "rZ" inline asm constraint

miloserdow · miloserdow · commit 1653040f9f25 · 2025-11-01T23:52:29.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 clang/lib/Basic/Targets/AArch64.cpp: - validateAsmConstraint: recognize "rZ"/"rz" constraint - convertConstraint: convert to "^rZ" In llvm/lib/Target/AArch64/AArch64ISelLowering.cpp: - getConstraintType: return C_Other - LowerAsmOperandForConstraint: substitute XZR/WZR for literal zero values Fixes #162567.
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_Other;
   } 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':
@@ -13196,6 +13207,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;
