SafeStack: Respect alloca addrspace

[arsenm]

Just insert addrspacecast in cases where the alloca uses a
different address space, since I don't know what else you
could possibly do.

[arsenm]

• SafeStack: Respect alloca addrspace #112536 👈
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

Join @arsenm and the rest of your teammates on Graphite

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Matt Arsenault (arsenm)

Changes

Just insert addrspacecast in cases where the alloca uses a
different address space, since I don't know what else you
could possibly do.

---

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

3 Files Affected:

• (modified) llvm/lib/CodeGen/SafeStack.cpp (+3-2)
• (modified) llvm/lib/CodeGen/TargetLoweringBase.cpp (+4-1)
• (added) llvm/test/Transforms/SafeStack/X86/alloca-addrspace.ll (+27)

diff --git a/llvm/lib/CodeGen/SafeStack.cpp b/llvm/lib/CodeGen/SafeStack.cpp
index a50909af8bfcfb..a20e72a8aff2a3 100644
--- a/llvm/lib/CodeGen/SafeStack.cpp
+++ b/llvm/lib/CodeGen/SafeStack.cpp
@@ -192,7 +192,7 @@ class SafeStack {
   SafeStack(Function &F, const TargetLoweringBase &TL, const DataLayout &DL,
             DomTreeUpdater *DTU, ScalarEvolution &SE)
       : F(F), TL(TL), DL(DL), DTU(DTU), SE(SE),
-        StackPtrTy(PointerType::getUnqual(F.getContext())),
+        StackPtrTy(DL.getAllocaPtrType(F.getContext())),
         IntPtrTy(DL.getIntPtrType(F.getContext())),
         Int32Ty(Type::getInt32Ty(F.getContext())) {}
 
@@ -617,7 +617,8 @@ Value *SafeStack::moveStaticAllocasToUnsafeStack(
       IRBuilder<> IRBUser(InsertBefore);
       Value *Off =
           IRBUser.CreatePtrAdd(BasePointer, ConstantInt::get(Int32Ty, -Offset));
-      Value *Replacement = IRBUser.CreateBitCast(Off, AI->getType(), Name);
+      Value *Replacement =
+          IRBUser.CreateAddrSpaceCast(Off, AI->getType(), Name);
 
       if (auto *PHI = dyn_cast<PHINode>(User))
         // PHI nodes may have multiple incoming edges from the same BB (why??),
diff --git a/llvm/lib/CodeGen/TargetLoweringBase.cpp b/llvm/lib/CodeGen/TargetLoweringBase.cpp
index 1f49d60c970593..04aa0d9a770b27 100644
--- a/llvm/lib/CodeGen/TargetLoweringBase.cpp
+++ b/llvm/lib/CodeGen/TargetLoweringBase.cpp
@@ -1848,7 +1848,8 @@ TargetLoweringBase::getDefaultSafeStackPointerLocation(IRBuilderBase &IRB,
   auto UnsafeStackPtr =
       dyn_cast_or_null<GlobalVariable>(M->getNamedValue(UnsafeStackPtrVar));
 
-  Type *StackPtrTy = PointerType::getUnqual(M->getContext());
+  const DataLayout &DL = M->getDataLayout();
+  PointerType *StackPtrTy = DL.getAllocaPtrType(M->getContext());
 
   if (!UnsafeStackPtr) {
     auto TLSModel = UseTLS ?
@@ -1862,6 +1863,8 @@ TargetLoweringBase::getDefaultSafeStackPointerLocation(IRBuilderBase &IRB,
         UnsafeStackPtrVar, nullptr, TLSModel);
   } else {
     // The variable exists, check its type and attributes.
+    //
+    // FIXME: Move to IR verifier.
     if (UnsafeStackPtr->getValueType() != StackPtrTy)
       report_fatal_error(Twine(UnsafeStackPtrVar) + " must have void* type");
     if (UseTLS != UnsafeStackPtr->isThreadLocal())
diff --git a/llvm/test/Transforms/SafeStack/X86/alloca-addrspace.ll b/llvm/test/Transforms/SafeStack/X86/alloca-addrspace.ll
new file mode 100644
index 00000000000000..efd6c64246102c
--- /dev/null
+++ b/llvm/test/Transforms/SafeStack/X86/alloca-addrspace.ll
@@ -0,0 +1,27 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -passes=safe-stack -mtriple=x86_64-pc-linux-gnu < %s -o - | FileCheck %s --check-prefix=TLS
+
+target datalayout = "A5"
+
+define void @correct_alloca_addrspace() nounwind uwtable safestack {
+; TLS-LABEL: define void @correct_alloca_addrspace(
+; TLS-SAME: ) #[[ATTR0:[0-9]+]] !annotation [[META0:![0-9]+]] {
+; TLS-NEXT:  [[ENTRY:.*:]]
+; TLS-NEXT:    [[UNSAFE_STACK_PTR:%.*]] = load ptr addrspace(5), ptr @__safestack_unsafe_stack_ptr, align 8
+; TLS-NEXT:    [[UNSAFE_STACK_STATIC_TOP:%.*]] = getelementptr i8, ptr addrspace(5) [[UNSAFE_STACK_PTR]], i32 -16
+; TLS-NEXT:    store ptr addrspace(5) [[UNSAFE_STACK_STATIC_TOP]], ptr @__safestack_unsafe_stack_ptr, align 8
+; TLS-NEXT:    [[TMP0:%.*]] = getelementptr i8, ptr addrspace(5) [[UNSAFE_STACK_PTR]], i32 -8
+; TLS-NEXT:    call void @Capture_as5(ptr addrspace(5) [[TMP0]])
+; TLS-NEXT:    store ptr addrspace(5) [[UNSAFE_STACK_PTR]], ptr @__safestack_unsafe_stack_ptr, align 8
+; TLS-NEXT:    ret void
+;
+entry:
+  %a = alloca i8, align 8, addrspace(5)
+  call void @Capture_as5(ptr addrspace(5) %a)
+  ret void
+}
+
+declare void @Capture_as5(ptr addrspace(5))
+;.
+; TLS: [[META0]] = !{!"unsafe-stack-size", i32 16}
+;.

[vitalybuka]

I suspect SafeStack require address space of UnsafeStackPtr and StackPtr is the same.

Maybe just fatal error if address space of PointerType::getUnqual != DL.getAllocaPtrType?

[arsenm]

Maybe just fatal error if address space of PointerType::getUnqual != DL.getAllocaPtrType?

That would defeat the point. It does produce an error if the realized __safestack_unsafe_stack_ptr type is wrong, test added

[vitalybuka]

I guess LGTM, not sure what we can do here better.

@eugenis do you see anything wrong with this approach?