6262#include " llvm/IR/Instruction.h"
6363#include " llvm/IR/Instructions.h"
6464#include " llvm/IR/IntrinsicInst.h"
65+ #include " llvm/IR/Intrinsics.h"
6566#include " llvm/IR/LLVMContext.h"
6667#include " llvm/IR/Metadata.h"
6768#include " llvm/IR/Module.h"
@@ -535,9 +536,11 @@ class Slice {
535536public:
536537 Slice () = default ;
537538
538- Slice (uint64_t BeginOffset, uint64_t EndOffset, Use *U, bool IsSplittable)
539+ Slice (uint64_t BeginOffset, uint64_t EndOffset, Use *U, bool IsSplittable,
540+ Value *ProtectedFieldDisc)
539541 : BeginOffset(BeginOffset), EndOffset(EndOffset),
540- UseAndIsSplittable (U, IsSplittable) {}
542+ UseAndIsSplittable (U, IsSplittable),
543+ ProtectedFieldDisc(ProtectedFieldDisc) {}
541544
542545 uint64_t beginOffset () const { return BeginOffset; }
543546 uint64_t endOffset () const { return EndOffset; }
@@ -550,6 +553,10 @@ class Slice {
550553 bool isDead () const { return getUse () == nullptr ; }
551554 void kill () { UseAndIsSplittable.setPointer (nullptr ); }
552555
556+ // When this access is via an llvm.protected.field.ptr intrinsic, contains
557+ // the second argument to the intrinsic, the discriminator.
558+ Value *ProtectedFieldDisc;
559+
553560 // / Support for ordering ranges.
554561 // /
555562 // / This provides an ordering over ranges such that start offsets are
@@ -641,6 +648,10 @@ class AllocaSlices {
641648 // / Access the dead users for this alloca.
642649 ArrayRef<Instruction *> getDeadUsers () const { return DeadUsers; }
643650
651+ // / Access the users for this alloca that are llvm.protected.field.ptr
652+ // / intrinsics.
653+ ArrayRef<IntrinsicInst *> getPFPUsers () const { return PFPUsers; }
654+
644655 // / Access Uses that should be dropped if the alloca is promotable.
645656 ArrayRef<Use *> getDeadUsesIfPromotable () const {
646657 return DeadUseIfPromotable;
@@ -701,6 +712,10 @@ class AllocaSlices {
701712 // / they come from outside of the allocated space.
702713 SmallVector<Instruction *, 8 > DeadUsers;
703714
715+ // / Users that are llvm.protected.field.ptr intrinsics. These will be RAUW'd
716+ // / to their first argument if we rewrite the alloca.
717+ SmallVector<IntrinsicInst *, 0 > PFPUsers;
718+
704719 // / Uses which will become dead if can promote the alloca.
705720 SmallVector<Use *, 8 > DeadUseIfPromotable;
706721
@@ -1029,6 +1044,10 @@ class AllocaSlices::SliceBuilder : public PtrUseVisitor<SliceBuilder> {
10291044 // / Set to de-duplicate dead instructions found in the use walk.
10301045 SmallPtrSet<Instruction *, 4 > VisitedDeadInsts;
10311046
1047+ // When this access is via an llvm.protected.field.ptr intrinsic, contains
1048+ // the second argument to the intrinsic, the discriminator.
1049+ Value *ProtectedFieldDisc = nullptr ;
1050+
10321051public:
10331052 SliceBuilder (const DataLayout &DL, AllocaInst &AI, AllocaSlices &AS)
10341053 : PtrUseVisitor<SliceBuilder>(DL),
@@ -1074,7 +1093,8 @@ class AllocaSlices::SliceBuilder : public PtrUseVisitor<SliceBuilder> {
10741093 EndOffset = AllocSize;
10751094 }
10761095
1077- AS.Slices .push_back (Slice (BeginOffset, EndOffset, U, IsSplittable));
1096+ AS.Slices .push_back (
1097+ Slice (BeginOffset, EndOffset, U, IsSplittable, ProtectedFieldDisc));
10781098 }
10791099
10801100 void visitBitCastInst (BitCastInst &BC) {
@@ -1274,6 +1294,27 @@ class AllocaSlices::SliceBuilder : public PtrUseVisitor<SliceBuilder> {
12741294 return ;
12751295 }
12761296
1297+ if (II.getIntrinsicID () == Intrinsic::protected_field_ptr) {
1298+ // We only handle loads and stores as users of llvm.protected.field.ptr.
1299+ // Other uses may add items to the worklist, which will cause
1300+ // ProtectedFieldDisc to be tracked incorrectly.
1301+ AS.PFPUsers .push_back (&II);
1302+ ProtectedFieldDisc = II.getArgOperand (1 );
1303+ for (Use &U : II.uses ()) {
1304+ this ->U = &U;
1305+ if (auto *LI = dyn_cast<LoadInst>(U.getUser ()))
1306+ visitLoadInst (*LI);
1307+ else if (auto *SI = dyn_cast<StoreInst>(U.getUser ()))
1308+ visitStoreInst (*SI);
1309+ else
1310+ PI.setAborted (&II);
1311+ if (PI.isAborted ())
1312+ break ;
1313+ }
1314+ ProtectedFieldDisc = nullptr ;
1315+ return ;
1316+ }
1317+
12771318 Base::visitIntrinsicInst (II);
12781319 }
12791320
@@ -4948,7 +4989,7 @@ bool SROA::presplitLoadsAndStores(AllocaInst &AI, AllocaSlices &AS) {
49484989 NewSlices.push_back (
49494990 Slice (BaseOffset + PartOffset, BaseOffset + PartOffset + PartSize,
49504991 &PLoad->getOperandUse (PLoad->getPointerOperandIndex ()),
4951- /* IsSplittable*/ false ));
4992+ /* IsSplittable*/ false , nullptr ));
49524993 LLVM_DEBUG (dbgs () << " new slice [" back ().beginOffset ()
49534994 << " , " back ().endOffset ()
49544995 << " ): " " \n "
@@ -5104,10 +5145,12 @@ bool SROA::presplitLoadsAndStores(AllocaInst &AI, AllocaSlices &AS) {
51045145 LLVMContext::MD_access_group});
51055146
51065147 // Now build a new slice for the alloca.
5148+ // ProtectedFieldDisc==nullptr is a lie, but it doesn't matter because we
5149+ // already determined that all accesses are consistent.
51075150 NewSlices.push_back (
51085151 Slice (BaseOffset + PartOffset, BaseOffset + PartOffset + PartSize,
51095152 &PStore->getOperandUse (PStore->getPointerOperandIndex ()),
5110- /* IsSplittable*/ false ));
5153+ /* IsSplittable*/ false , nullptr ));
51115154 LLVM_DEBUG (dbgs () << " new slice [" back ().beginOffset ()
51125155 << " , " back ().endOffset ()
51135156 << " ): " " \n "
@@ -5875,6 +5918,30 @@ SROA::runOnAlloca(AllocaInst &AI) {
58755918 return {Changed, CFGChanged};
58765919 }
58775920
5921+ for (auto &P : AS.partitions ()) {
5922+ // For now, we can't split if a field is accessed both via protected field
5923+ // and not, because that would mean that we would need to introduce sign and
5924+ // auth operations to convert between the protected and non-protected uses,
5925+ // and this pass doesn't know how to do that. Also, this case is unlikely to
5926+ // occur in normal code.
5927+ std::optional<Value *> ProtectedFieldDisc;
5928+ auto SliceHasMismatch = [&](Slice &S) {
5929+ if (auto *II = dyn_cast<IntrinsicInst>(S.getUse ()->getUser ()))
5930+ if (II->getIntrinsicID () == Intrinsic::lifetime_start ||
5931+ II->getIntrinsicID () == Intrinsic::lifetime_end)
5932+ return false ;
5933+ if (!ProtectedFieldDisc)
5934+ ProtectedFieldDisc = S.ProtectedFieldDisc ;
5935+ return *ProtectedFieldDisc != S.ProtectedFieldDisc ;
5936+ };
5937+ for (Slice &S : P)
5938+ if (SliceHasMismatch (S))
5939+ return {Changed, CFGChanged};
5940+ for (Slice *S : P.splitSliceTails ())
5941+ if (SliceHasMismatch (*S))
5942+ return {Changed, CFGChanged};
5943+ }
5944+
58785945 // Delete all the dead users of this alloca before splitting and rewriting it.
58795946 for (Instruction *DeadUser : AS.getDeadUsers ()) {
58805947 // Free up everything used by this instruction.
@@ -5892,6 +5959,12 @@ SROA::runOnAlloca(AllocaInst &AI) {
58925959 clobberUse (*DeadOp);
58935960 Changed = true ;
58945961 }
5962+ for (IntrinsicInst *PFPUser : AS.getPFPUsers ()) {
5963+ PFPUser->replaceAllUsesWith (PFPUser->getArgOperand (0 ));
5964+
5965+ DeadInsts.push_back (PFPUser);
5966+ Changed = true ;
5967+ }
58955968
58965969 // No slices to split. Leave the dead alloca for a later pass to clean up.
58975970 if (AS.begin () == AS.end ())
0 commit comments