[llvm-branch-commits] [llvm] [AArch64] Use __chkstk for dynamic stack allocation on UEFI (PR #232702)

Prabhu Rajasekaran via llvm-branch-commits llvm-branch-commits at lists.llvm.org
Tue Oct 20 11:55:56 PDT 2026


https://github.com/Prabhuk updated https://github.com/llvm/llvm-project/pull/232702

>From 243e81ef99bb9645810e495a34e6ebf4a7d715c4 Mon Sep 17 00:00:00 2001
From: prabhukr <prabhukr at google.com>
Date: Mon, 19 Oct 2026 17:27:04 -0700
Subject: [PATCH 1/2] [AArch64] Use __chkstk for dynamic stack allocation on
 UEFI

Register the Windows __chkstk stack probe helper for aarch64-unknown-uefi
and lower dynamic stack allocations through it, including the SME ZA
buffer allocation in LowerFormalArguments. GlobalISel falls back to
SelectionDAG for dynamic allocas on UEFI as it does on Windows.

Tests: win-alloca.ll, the vla in win64_vararg.ll (which now runs in full
on UEFI) and sme-lazy-save-windows.ll.

Assisted-by: Claude
---
 llvm/include/llvm/IR/RuntimeLibcalls.td            | 5 +++--
 llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp       | 4 ++--
 llvm/lib/Target/AArch64/AArch64ISelLowering.cpp    | 4 ++--
 llvm/lib/Target/AArch64/MachineSMEABIPass.cpp      | 2 +-
 llvm/test/CodeGen/AArch64/sme-lazy-save-windows.ll | 1 +
 llvm/test/CodeGen/AArch64/win-alloca.ll            | 1 +
 llvm/test/CodeGen/AArch64/win64_vararg.ll          | 1 +
 7 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/llvm/include/llvm/IR/RuntimeLibcalls.td b/llvm/include/llvm/IR/RuntimeLibcalls.td
index d05cd842bdb047..26fae82e95354c 100644
--- a/llvm/include/llvm/IR/RuntimeLibcalls.td
+++ b/llvm/include/llvm/IR/RuntimeLibcalls.td
@@ -2064,9 +2064,10 @@ defvar SecurityCheckCookieIfWinMSVC =
     LibcallImpls<(add __security_check_cookie, __security_cookie),
                  isWindowsMSVCOrItaniumEnvironment>;
 
-// The Windows CRT security cookie and stack probes are part of libc.
+// The Windows CRT security cookie and stack probes are part of libc. UEFI
+// follows the Windows ARM64 ABI and uses the same stack probe helper.
 def isWindowsAArch64NotArm64EC : RuntimeLibcallAvailability<
-  (all_of IsOSWindows, IsAArch64, (not IsWindowsArm64EC))>;
+  (all_of IsOSWindowsOrUEFI, IsAArch64, (not IsWindowsArm64EC))>;
 def isWindowsARM : RuntimeLibcallAvailability<
   (all_of IsOSWindows, (any_of IsARM, IsThumb))>;
 def WindowsCRTLibc : LibcallLibrary<"libc", (add SecurityCheckCookieIfWinMSVC)>;
diff --git a/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp b/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp
index c354c309921b87..18f5b1c8149711 100644
--- a/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp
@@ -4040,8 +4040,8 @@ bool IRTranslatorImpl::translateAlloca(const User &U,
     return true;
   }
 
-  // FIXME: support stack probing for Windows.
-  if (MF->getTarget().getTargetTriple().isOSWindows())
+  // FIXME: support stack probing for Windows and UEFI.
+  if (MF->getTarget().getTargetTriple().isOSWindowsOrUEFI())
     return false;
 
   // Now we're in the harder dynamic case.
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 9d47ef0a28c08e..008e11cfef697d 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -9729,7 +9729,7 @@ SDValue AArch64TargetLowering::LowerFormalArguments(
   if (Subtarget->hasCustomCallingConv())
     Subtarget->getRegisterInfo()->UpdateCustomCalleeSavedRegs(MF);
 
-  if (Subtarget->isTargetWindows() || hasInlineStackProbe(MF)) {
+  if (Subtarget->isTargetWindowsOrUEFI() || hasInlineStackProbe(MF)) {
     SDValue Size;
     if (Attrs.hasZAState()) {
       SDValue SVL = DAG.getNode(AArch64ISD::RDSVL, DL, MVT::i64,
@@ -18532,7 +18532,7 @@ AArch64TargetLowering::LowerDYNAMIC_STACKALLOC(SDValue Op,
                                                SelectionDAG &DAG) const {
   MachineFunction &MF = DAG.getMachineFunction();
 
-  if (Subtarget->isTargetWindows())
+  if (Subtarget->isTargetWindowsOrUEFI())
     return LowerWindowsDYNAMIC_STACKALLOC(Op, DAG);
   else if (hasInlineStackProbe(MF))
     return LowerInlineDYNAMIC_STACKALLOC(Op, DAG);
diff --git a/llvm/lib/Target/AArch64/MachineSMEABIPass.cpp b/llvm/lib/Target/AArch64/MachineSMEABIPass.cpp
index a09a9db6a3ad48..3f62a2ec97cfe9 100644
--- a/llvm/lib/Target/AArch64/MachineSMEABIPass.cpp
+++ b/llvm/lib/Target/AArch64/MachineSMEABIPass.cpp
@@ -887,7 +887,7 @@ void MachineSMEABI::emitAllocateLazySaveBuffer(
     // will always allocate the lazy save buffer (even if the function contains
     // no lazy saves). If we want to handle Windows here, we'll need to
     // implement something similar to LowerWindowsDYNAMIC_STACKALLOC.
-    assert(!Subtarget->isTargetWindows() &&
+    assert(!Subtarget->isTargetWindowsOrUEFI() &&
            "Lazy ZA save is not yet supported on Windows");
     Buffer = MRI->createVirtualRegister(&AArch64::GPR64RegClass);
     // Get original stack pointer.
diff --git a/llvm/test/CodeGen/AArch64/sme-lazy-save-windows.ll b/llvm/test/CodeGen/AArch64/sme-lazy-save-windows.ll
index 6e2d2c0269757a..42340a28801a2d 100644
--- a/llvm/test/CodeGen/AArch64/sme-lazy-save-windows.ll
+++ b/llvm/test/CodeGen/AArch64/sme-lazy-save-windows.ll
@@ -1,5 +1,6 @@
 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
 ; RUN: llc -mtriple=aarch64-windows-msvc -aarch64-streaming-hazard-size=0 -mattr=+sve,+sme < %s | FileCheck %s
+; RUN: llc -mtriple=aarch64-unknown-uefi -aarch64-streaming-hazard-size=0 -mattr=+sve,+sme < %s | FileCheck %s
 
 declare void @private_za_callee()
 declare void @shared_za_callee() "aarch64_inout_za"
diff --git a/llvm/test/CodeGen/AArch64/win-alloca.ll b/llvm/test/CodeGen/AArch64/win-alloca.ll
index 94c3d884581695..a3592308807757 100644
--- a/llvm/test/CodeGen/AArch64/win-alloca.ll
+++ b/llvm/test/CodeGen/AArch64/win-alloca.ll
@@ -1,5 +1,6 @@
 ; RUN: llc -mtriple aarch64-windows -verify-machineinstrs -filetype asm -o - %s | FileCheck %s -check-prefixes=CHECK,CHECK-OPT
 ; RUN: llc -mtriple aarch64-windows -verify-machineinstrs -filetype asm -o - %s -O0 | FileCheck %s
+; RUN: llc -mtriple aarch64-unknown-uefi -verify-machineinstrs -filetype asm -o - %s | FileCheck %s -check-prefixes=CHECK,CHECK-OPT
 ; RUN: llc -mtriple arm64ec-windows-msvc -verify-machineinstrs -filetype asm -o - %s | FileCheck -check-prefix=CHECK-ARM64EC %s
 
 define void @func(i64 %a) {
diff --git a/llvm/test/CodeGen/AArch64/win64_vararg.ll b/llvm/test/CodeGen/AArch64/win64_vararg.ll
index 0f4cb44a6f73b2..75f27ae97acdc8 100644
--- a/llvm/test/CodeGen/AArch64/win64_vararg.ll
+++ b/llvm/test/CodeGen/AArch64/win64_vararg.ll
@@ -1,5 +1,6 @@
 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc < %s -mtriple=aarch64-pc-win32 | FileCheck %s
+; RUN: llc < %s -mtriple=aarch64-unknown-uefi | FileCheck %s
 
 define void @pass_va(i32 %count, ...) nounwind {
 ; CHECK-LABEL: pass_va:

>From 7af2b7ecccb2242c05d2dec9190c2ba223b33cfc Mon Sep 17 00:00:00 2001
From: prabhukr <prabhukr at google.com>
Date: Tue, 20 Oct 2026 11:55:42 -0700
Subject: [PATCH 2/2] [AArch64] Address review feedback on UEFI dynamic stack
 allocation

Include UEFI in the MachineSMEABI lazy ZA save assertion message and add
an -O0 run line for aarch64-unknown-uefi in win-alloca.ll to exercise the
GlobalISel dynamic alloca fallback.

Assisted-by: Claude
---
 llvm/lib/Target/AArch64/MachineSMEABIPass.cpp | 2 +-
 llvm/test/CodeGen/AArch64/win-alloca.ll       | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Target/AArch64/MachineSMEABIPass.cpp b/llvm/lib/Target/AArch64/MachineSMEABIPass.cpp
index 3f62a2ec97cfe9..81fb7816e9c7a6 100644
--- a/llvm/lib/Target/AArch64/MachineSMEABIPass.cpp
+++ b/llvm/lib/Target/AArch64/MachineSMEABIPass.cpp
@@ -888,7 +888,7 @@ void MachineSMEABI::emitAllocateLazySaveBuffer(
     // no lazy saves). If we want to handle Windows here, we'll need to
     // implement something similar to LowerWindowsDYNAMIC_STACKALLOC.
     assert(!Subtarget->isTargetWindowsOrUEFI() &&
-           "Lazy ZA save is not yet supported on Windows");
+           "Lazy ZA save is not yet supported on Windows or UEFI");
     Buffer = MRI->createVirtualRegister(&AArch64::GPR64RegClass);
     // Get original stack pointer.
     BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::COPY), SP)
diff --git a/llvm/test/CodeGen/AArch64/win-alloca.ll b/llvm/test/CodeGen/AArch64/win-alloca.ll
index a3592308807757..2f162dd5301ae0 100644
--- a/llvm/test/CodeGen/AArch64/win-alloca.ll
+++ b/llvm/test/CodeGen/AArch64/win-alloca.ll
@@ -1,6 +1,7 @@
 ; RUN: llc -mtriple aarch64-windows -verify-machineinstrs -filetype asm -o - %s | FileCheck %s -check-prefixes=CHECK,CHECK-OPT
 ; RUN: llc -mtriple aarch64-windows -verify-machineinstrs -filetype asm -o - %s -O0 | FileCheck %s
 ; RUN: llc -mtriple aarch64-unknown-uefi -verify-machineinstrs -filetype asm -o - %s | FileCheck %s -check-prefixes=CHECK,CHECK-OPT
+; RUN: llc -mtriple aarch64-unknown-uefi -verify-machineinstrs -filetype asm -o - %s -O0 | FileCheck %s
 ; RUN: llc -mtriple arm64ec-windows-msvc -verify-machineinstrs -filetype asm -o - %s | FileCheck -check-prefix=CHECK-ARM64EC %s
 
 define void @func(i64 %a) {



More information about the llvm-branch-commits mailing list