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

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-aarch64

Author: Timur Baidusenov (bai-tim)

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/225067.diff


2 Files Affected:

- (modified) llvm/lib/CodeGen/SafeStack.cpp (+7-2) 
- (modified) llvm/test/CodeGen/AArch64/safestack_scalar.ll (+26) 


``````````diff
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]])

``````````

</details>


https://github.com/llvm/llvm-project/pull/225067