[llvm] [AArch64] Preserve undef on registers when expanding CALL_RVMARKER. (PR #71177)

[Florian Hahn via llvm-commits]

https://github.com/fhahn created https://github.com/llvm/llvm-project/pull/71177

Preserve the undef flag on register operands when expanding CALL_RVMARKER. This fixes a crash when the x0 operand is undef, e.g. because it was poison in the input IR.

Fixes https://github.com/llvm/llvm-project/issues/68005.

>From 2ccec19763026b07167d62cf6f4b8b3d793c7d12 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Fri, 3 Nov 2023 11:36:10 +0000
Subject: [PATCH] [AArch64] Preserve undef on registers when expanding
 CALL_RVMARKER.

Preserve the undef flag on register operands when expanding
CALL_RVMARKER. This fixes a crash when the x0 operand is undef, e.g.
because it was poison in the input IR.

Fixes https://github.com/llvm/llvm-project/issues/68005.
---
 .../AArch64/AArch64ExpandPseudoInsts.cpp      |  3 +-
 .../AArch64/expand-blr-rvmarker-pseudo.mir    | 31 ++++++++++++++++++-
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp b/llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp
index 42bc390a67b2b18..7fa12c6a2142537 100644
--- a/llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp
@@ -794,7 +794,8 @@ bool AArch64ExpandPseudo::expandCALL_RVMARKER(
     auto MOP = MI.getOperand(RegMaskStartIdx);
     assert(MOP.isReg() && "can only add register operands");
     OriginalCall->addOperand(MachineOperand::CreateReg(
-        MOP.getReg(), /*Def=*/false, /*Implicit=*/true));
+        MOP.getReg(), /*Def=*/false, /*Implicit=*/true, /*isKill=*/false,
+        /*isDead=*/false, /*isUndef=*/MOP.isUndef()));
     RegMaskStartIdx++;
   }
   for (const MachineOperand &MO :
diff --git a/llvm/test/CodeGen/AArch64/expand-blr-rvmarker-pseudo.mir b/llvm/test/CodeGen/AArch64/expand-blr-rvmarker-pseudo.mir
index 83dedcb111a9771..37d1623b2cd8db5 100644
--- a/llvm/test/CodeGen/AArch64/expand-blr-rvmarker-pseudo.mir
+++ b/llvm/test/CodeGen/AArch64/expand-blr-rvmarker-pseudo.mir
@@ -1,4 +1,4 @@
-# RUN: llc -run-pass=aarch64-expand-pseudo -mtriple=arm64-apple-ios -o - -emit-call-site-info %s | FileCheck %s
+# RUN: llc -run-pass=aarch64-expand-pseudo -mtriple=arm64-apple-ios -o - -emit-call-site-info %s -verify-machineinstrs | FileCheck %s
 
 --- |
   define void @test_1_callsite_info() {
@@ -25,7 +25,13 @@
     ret void
   }
 
+  define void @test_blr_undef_arg_in_x0() {
+    ret void
+  }
+
   declare i8* @attachedcall()
+
+  declare ptr @objc_retainAutoreleasedReturnValue()
 ...
 ---
 
@@ -132,3 +138,26 @@ body:             |
         BLR_RVMARKER @attachedcall, $x8, $w0, $w1, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp, implicit-def dead $x0
         RET_ReallyLR implicit killed $w0
 ...
+
+# Test for https://github.com/llvm/llvm-project/issues/68005.
+#
+# CHECK-LABEL: : test_blr_undef_arg_in_x0
+# CHECK:       bb.0:
+# CHECK-NEXT: liveins:
+# CHECK-NEXT: {{  $}}
+# CHECK-NEXT:    BUNDLE implicit-def $lr, implicit-def $w30, implicit-def $sp, implicit-def $wsp, implicit-def dead $x0, implicit-def $fp, implicit-def $w29, implicit $sp, implicit undef $x0, implicit $xzr, implicit $fp {
+# CHECK-NEXT:      BL @foo, csr_darwin_aarch64_aapcs, implicit-def $lr, implicit $sp, implicit undef $x0, implicit-def dead $lr, implicit $sp, implicit-def $sp, implicit-def dead $x0
+# CHECK-NEXT:      $fp = ORRXrs $xzr, $fp, 0
+# CHECK-NEXT:      BL @objc_retainAutoreleasedReturnValue, implicit-def $lr, implicit internal $sp
+# CHECK-NEXT:    }
+# CHECK-NEXT:    RET undef $lr
+#
+name: test_blr_undef_arg_in_x0
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $lr
+
+    BLR_RVMARKER @objc_retainAutoreleasedReturnValue, @foo, undef $x0, csr_darwin_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp, implicit-def dead $x0
+    RET_ReallyLR
+...