[Attributor][FIX] Honor alloca address space in AAPrivatizablePtr

When we privatize a pointer (~argument promotion) we introduce new
private allocas as replacement. These need to be placed in the alloca
address space as later passes cannot properly deal with them otherwise.

Fixes #53725

Author: jdoerfert

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

@@ -236,7 +236,8 @@ static Value *constructPointer(Type *ResTy, Type *PtrElemTy, Value *Ptr,
   }
 
   // Ensure the result has the requested type.
-  Ptr = IRB.CreateBitOrPointerCast(Ptr, ResTy, Ptr->getName() + ".cast");
+  Ptr = IRB.CreatePointerBitCastOrAddrSpaceCast(Ptr, ResTy,
+                                                Ptr->getName() + ".cast");
 
   LLVM_DEBUG(dbgs() << "Constructed pointer: " << *Ptr << "\n");
   return Ptr;
@@ -6748,8 +6749,8 @@ struct AAPrivatizablePtrArgument final : public AAPrivatizablePtrImpl {
 
     Type *PrivPtrType = PrivType->getPointerTo();
     if (Base->getType() != PrivPtrType)
-      Base = BitCastInst::CreateBitOrPointerCast(Base, PrivPtrType, "",
-                                                 ACS.getInstruction());
+      Base = BitCastInst::CreatePointerBitCastOrAddrSpaceCast(
+          Base, PrivPtrType, "", ACS.getInstruction());
 
     // Traverse the type, build GEPs and loads.
     if (auto *PrivStructType = dyn_cast<StructType>(PrivType)) {
@@ -6816,14 +6817,16 @@ struct AAPrivatizablePtrArgument final : public AAPrivatizablePtrImpl {
             Function &ReplacementFn, Function::arg_iterator ArgIt) {
           BasicBlock &EntryBB = ReplacementFn.getEntryBlock();
           Instruction *IP = &*EntryBB.getFirstInsertionPt();
-          Instruction *AI = new AllocaInst(PrivatizableType.getValue(), 0,
+          const DataLayout &DL = IP->getModule()->getDataLayout();
+          unsigned AS = DL.getAllocaAddrSpace();
+          Instruction *AI = new AllocaInst(PrivatizableType.getValue(), AS,
                                            Arg->getName() + ".priv", IP);
           createInitialization(PrivatizableType.getValue(), *AI, ReplacementFn,
                                ArgIt->getArgNo(), *IP);
 
           if (AI->getType() != Arg->getType())
-            AI =
-                BitCastInst::CreateBitOrPointerCast(AI, Arg->getType(), "", IP);
+            AI = BitCastInst::CreatePointerBitCastOrAddrSpaceCast(
+                AI, Arg->getType(), "", IP);
           Arg->replaceAllUsesWith(AI);
 
           for (CallInst *CI : TailCalls)

llvm/test/Transforms/Attributor/ArgumentPromotion/alloca-as.ll

@@ -0,0 +1,78 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --function-signature --check-attributes --check-globals
+; RUN: opt -attributor -enable-new-pm=0 -attributor-manifest-internal  -attributor-max-iterations-verify -attributor-annotate-decl-cs -attributor-max-iterations=3 -S < %s | FileCheck %s --check-prefixes=CHECK,NOT_CGSCC_NPM,NOT_CGSCC_OPM,NOT_TUNIT_NPM,IS__TUNIT____,IS________OPM,IS__TUNIT_OPM
+; RUN: opt -aa-pipeline=basic-aa -passes=attributor -attributor-manifest-internal  -attributor-max-iterations-verify -attributor-annotate-decl-cs -attributor-max-iterations=3 -S < %s | FileCheck %s --check-prefixes=CHECK,NOT_CGSCC_OPM,NOT_CGSCC_NPM,NOT_TUNIT_OPM,IS__TUNIT____,IS________NPM,IS__TUNIT_NPM
+; RUN: opt -attributor-cgscc -enable-new-pm=0 -attributor-manifest-internal  -attributor-annotate-decl-cs -S < %s | FileCheck %s --check-prefixes=CHECK,NOT_TUNIT_NPM,NOT_TUNIT_OPM,NOT_CGSCC_NPM,IS__CGSCC____,IS________OPM,IS__CGSCC_OPM
+; RUN: opt -aa-pipeline=basic-aa -passes=attributor-cgscc -attributor-manifest-internal  -attributor-annotate-decl-cs -S < %s | FileCheck %s --check-prefixes=CHECK,NOT_TUNIT_NPM,NOT_TUNIT_OPM,NOT_CGSCC_OPM,IS__CGSCC____,IS________NPM,IS__CGSCC_NPM
+
+target datalayout = "A7"
+
+; Make sure we create allocas in AS 7 and cast them properly.
+
+define i32 @bar(i32 %arg) {
+; IS________OPM-LABEL: define {{[^@]+}}@bar
+; IS________OPM-SAME: (i32 [[ARG:%.*]]) {
+; IS________OPM-NEXT:  entry:
+; IS________OPM-NEXT:    [[STACK:%.*]] = alloca i32, align 4
+; IS________OPM-NEXT:    store i32 [[ARG]], i32* [[STACK]], align 4
+; IS________OPM-NEXT:    [[CALL:%.*]] = call i32 @foo(i32* noalias nocapture nofree noundef nonnull readonly align 4 dereferenceable(4) [[STACK]])
+; IS________OPM-NEXT:    ret i32 [[CALL]]
+;
+; IS__TUNIT_NPM-LABEL: define {{[^@]+}}@bar
+; IS__TUNIT_NPM-SAME: (i32 [[ARG:%.*]]) {
+; IS__TUNIT_NPM-NEXT:  entry:
+; IS__TUNIT_NPM-NEXT:    [[STACK:%.*]] = alloca i32, align 4
+; IS__TUNIT_NPM-NEXT:    store i32 [[ARG]], i32* [[STACK]], align 4
+; IS__TUNIT_NPM-NEXT:    [[TMP0:%.*]] = load i32, i32* [[STACK]], align 4
+; IS__TUNIT_NPM-NEXT:    [[CALL:%.*]] = call i32 @foo(i32 [[TMP0]])
+; IS__TUNIT_NPM-NEXT:    ret i32 [[CALL]]
+;
+; IS__CGSCC_NPM-LABEL: define {{[^@]+}}@bar
+; IS__CGSCC_NPM-SAME: (i32 returned [[ARG:%.*]]) {
+; IS__CGSCC_NPM-NEXT:  entry:
+; IS__CGSCC_NPM-NEXT:    [[STACK:%.*]] = alloca i32, align 4
+; IS__CGSCC_NPM-NEXT:    store i32 [[ARG]], i32* [[STACK]], align 4
+; IS__CGSCC_NPM-NEXT:    [[CALL:%.*]] = call i32 @foo(i32 [[ARG]])
+; IS__CGSCC_NPM-NEXT:    ret i32 [[ARG]]
+;
+entry:
+  %stack = alloca i32
+  store i32 %arg, i32* %stack
+  %call = call i32 @foo(i32* %stack)
+  ret i32 %call
+}
+
+define internal i32 @foo(i32* %arg) {
+; IS________OPM-LABEL: define {{[^@]+}}@foo
+; IS________OPM-SAME: (i32* noalias nocapture nofree noundef nonnull readonly align 4 dereferenceable(4) [[ARG:%.*]]) {
+; IS________OPM-NEXT:  entry:
+; IS________OPM-NEXT:    [[L:%.*]] = load i32, i32* [[ARG]], align 4
+; IS________OPM-NEXT:    call void @use(i32 [[L]])
+; IS________OPM-NEXT:    ret i32 [[L]]
+;
+; IS__TUNIT_NPM-LABEL: define {{[^@]+}}@foo
+; IS__TUNIT_NPM-SAME: (i32 [[TMP0:%.*]]) {
+; IS__TUNIT_NPM-NEXT:  entry:
+; IS__TUNIT_NPM-NEXT:    [[ARG_PRIV:%.*]] = alloca i32, align 4, addrspace(7)
+; IS__TUNIT_NPM-NEXT:    store i32 [[TMP0]], i32 addrspace(7)* [[ARG_PRIV]], align 4
+; IS__TUNIT_NPM-NEXT:    [[TMP1:%.*]] = addrspacecast i32 addrspace(7)* [[ARG_PRIV]] to i32*
+; IS__TUNIT_NPM-NEXT:    [[L:%.*]] = load i32, i32* [[TMP1]], align 4
+; IS__TUNIT_NPM-NEXT:    call void @use(i32 [[L]])
+; IS__TUNIT_NPM-NEXT:    ret i32 [[L]]
+;
+; IS__CGSCC_NPM-LABEL: define {{[^@]+}}@foo
+; IS__CGSCC_NPM-SAME: (i32 returned [[TMP0:%.*]]) {
+; IS__CGSCC_NPM-NEXT:  entry:
+; IS__CGSCC_NPM-NEXT:    [[ARG_PRIV:%.*]] = alloca i32, align 4, addrspace(7)
+; IS__CGSCC_NPM-NEXT:    store i32 [[TMP0]], i32 addrspace(7)* [[ARG_PRIV]], align 4
+; IS__CGSCC_NPM-NEXT:    [[TMP1:%.*]] = addrspacecast i32 addrspace(7)* [[ARG_PRIV]] to i32*
+; IS__CGSCC_NPM-NEXT:    [[L:%.*]] = load i32, i32* [[TMP1]], align 4
+; IS__CGSCC_NPM-NEXT:    call void @use(i32 [[TMP0]])
+; IS__CGSCC_NPM-NEXT:    ret i32 [[TMP0]]
+;
+entry:
+  %l = load i32, i32* %arg
+  call void @use(i32 %l)
+  ret i32 %l
+}
+
+declare void @use(i32)