[llvm] [SafeStack] Don't move a scalable byval argument to the unsafe stack (PR #225067)

[Timur Baidusenov via llvm-commits]

https://github.com/bai-tim created https://github.com/llvm/llvm-project/pull/225067

The unsafe stack frame is laid out at compile time, but the size of a byval argument of a scalable type is not known then, so getTypeStoreSize() aborts with "Cannot implicitly convert a scalable size to a fixed-width size" before the argument is even classified. Skip it, the way MemCpyOpt bails out of the call-slot optimization for a byval argument of the same shape; the earlier 537f3d3a7588 fixed the other half of #182759, where the size of an access rather than of an argument was scalable.

Fixes #182759.

>From dbc21d3c752d1cae3c4defb31af2556328569801 Mon Sep 17 00:00:00 2001
From: Timur Baidusenov <timurbaidusenov at gmail.com>
Date: Fri, 18 Sep 2026 18:33:09 +0300
Subject: [PATCH] [SafeStack] Don't move a scalable byval argument to the
 unsafe stack

The unsafe stack frame is laid out at compile time, but the size of a byval
argument of a scalable type is not known then, so getTypeStoreSize() aborts
with "Cannot implicitly convert a scalable size to a fixed-width size" before
the argument is even classified. Skip it, the way MemCpyOpt bails out of the
call-slot optimization for a byval argument of the same shape; the earlier
537f3d3a7588 fixed the other half of #182759, where the size of an access
rather than of an argument was scalable.

Fixes #182759.
---
 llvm/lib/CodeGen/SafeStack.cpp                |  9 +++++--
 llvm/test/CodeGen/AArch64/safestack_scalar.ll | 26 +++++++++++++++++++
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/CodeGen/SafeStack.cpp b/llvm/lib/CodeGen/SafeStack.cpp
index 75a8762cc7090..07d8854a1629e 100644
--- a/llvm/lib/CodeGen/SafeStack.cpp
+++ b/llvm/lib/CodeGen/SafeStack.cpp
@@ -418,8 +418,13 @@ void SafeStack::findInsts(Function &F,
   for (Argument &Arg : F.args()) {
     if (!Arg.hasByValAttr())
       continue;
-    uint64_t Size = DL.getTypeStoreSize(Arg.getParamByValType());
-    if (IsSafeStackAlloca(&Arg, Size))
+    TypeSize Size = DL.getTypeStoreSize(Arg.getParamByValType());
+    // The unsafe stack frame is laid out at compile time, so an argument of a
+    // scalable type has no place on it. Leave it on the regular stack rather
+    // than reject a function that is otherwise valid.
+    if (Size.isScalable())
+      continue;
+    if (IsSafeStackAlloca(&Arg, Size.getFixedValue()))
       continue;
 
     ++NumUnsafeByValArguments;
diff --git a/llvm/test/CodeGen/AArch64/safestack_scalar.ll b/llvm/test/CodeGen/AArch64/safestack_scalar.ll
index f8675e7a709d3..0488218b6c848 100644
--- a/llvm/test/CodeGen/AArch64/safestack_scalar.ll
+++ b/llvm/test/CodeGen/AArch64/safestack_scalar.ll
@@ -15,3 +15,29 @@ entry:
 ; CHECK: load <vscale x 16 x i8>, ptr [[PTR]]
 ; CHECK: store ptr [[USP]], ptr @__safestack_unsafe_stack_ptr
 ; CHECK: ret void
+
+declare void @escape(ptr)
+
+; A byval argument of a scalable type has no size known at compile time, so it
+; cannot be placed on the unsafe stack frame and is left alone.
+
+define void @test_sve_byval(ptr byval(<vscale x 4 x i32>) %p) safestack {
+  call void @escape(ptr %p)
+  ret void
+}
+
+; CHECK-LABEL: define void @test_sve_byval(
+; CHECK-NEXT: call void @escape(ptr %p)
+; CHECK-NEXT: ret void
+
+; A byval argument of a fixed size is still copied onto the unsafe stack.
+
+define void @test_byval(ptr byval([16 x i8]) %p) safestack {
+  call void @escape(ptr %p)
+  ret void
+}
+
+; CHECK-LABEL: define void @test_byval(
+; CHECK: load ptr, ptr @__safestack_unsafe_stack_ptr
+; CHECK: call void @llvm.memcpy.p0.p0.i64(ptr align 4 [[SLOT:%.*]], ptr %p, i64 16, i1 false)
+; CHECK: call void @escape(ptr [[SLOT]])