@@ -183,23 +183,43 @@ class MemoryToRegisters {
183183// / Returns true if \p I is an address of a LoadInst, skipping struct and
184184// / tuple address projections. Sets \p singleBlock to null if the load (or
185185// / it's address is not in \p singleBlock.
186- static bool isAddressForLoad (SILInstruction *I, SILBasicBlock *&singleBlock) {
187-
188- if (isa<LoadInst>(I))
186+ // / This function looks for these patterns:
187+ // / 1. (load %ASI)
188+ // / 2. (load (struct_element_addr/tuple_element_addr/unchecked_addr_cast %ASI))
189+ static bool isAddressForLoad (SILInstruction *I, SILBasicBlock *&singleBlock,
190+ bool &hasGuaranteedOwnership) {
191+
192+ if (isa<LoadInst>(I)) {
193+ // SILMem2Reg is disabled when we find:
194+ // (load [take] (struct_element_addr/tuple_element_addr %ASI))
195+ // struct_element_addr and tuple_element_addr are lowered into
196+ // struct_extract and tuple_extract and these SIL instructions have a
197+ // guaranteed ownership. For replacing load's users, we need an owned value.
198+ // We will need a new copy and destroy of the running val placed after the
199+ // last use. This is not implemented currently.
200+ if (hasGuaranteedOwnership && cast<LoadInst>(I)->getOwnershipQualifier () ==
201+ LoadOwnershipQualifier::Take) {
202+ return false ;
203+ }
189204 return true ;
205+ }
190206
191207 if (!isa<UncheckedAddrCastInst>(I) && !isa<StructElementAddrInst>(I) &&
192208 !isa<TupleElementAddrInst>(I))
193209 return false ;
194-
210+
211+ if (isa<StructElementAddrInst>(I) || isa<TupleElementAddrInst>(I)) {
212+ hasGuaranteedOwnership = true ;
213+ }
214+
195215 // Recursively search for other (non-)loads in the instruction's uses.
196216 for (auto UI : cast<SingleValueInstruction>(I)->getUses ()) {
197217 SILInstruction *II = UI->getUser ();
198218 if (II->getParent () != singleBlock)
199219 singleBlock = nullptr ;
200-
201- if (!isAddressForLoad (II, singleBlock))
202- return false ;
220+
221+ if (!isAddressForLoad (II, singleBlock, hasGuaranteedOwnership ))
222+ return false ;
203223 }
204224 return true ;
205225}
@@ -233,7 +253,8 @@ static bool isCaptured(AllocStackInst *ASI, bool &inSingleBlock) {
233253 singleBlock = nullptr ;
234254
235255 // Loads are okay.
236- if (isAddressForLoad (II, singleBlock))
256+ bool hasGuaranteedOwnership = false ;
257+ if (isAddressForLoad (II, singleBlock, hasGuaranteedOwnership))
237258 continue ;
238259
239260 // We can store into an AllocStack (but not the pointer).
@@ -302,8 +323,10 @@ promoteDebugValueAddr(DebugValueAddrInst *DVAI, SILValue Value, SILBuilder &B) {
302323 // Avoid inserting the same debug_value twice.
303324 for (Operand *Use : Value->getUses ())
304325 if (auto *DVI = dyn_cast<DebugValueInst>(Use->getUser ()))
305- if (*DVI->getVarInfo () == *DVAI->getVarInfo ())
326+ if (*DVI->getVarInfo () == *DVAI->getVarInfo ()) {
327+ DVAI->eraseFromParent ();
306328 return ;
329+ }
307330 B.setInsertionPoint (DVAI);
308331 B.setCurrentDebugScope (DVAI->getDebugScope ());
309332 B.createDebugValue (DVAI->getLoc (), Value, *DVAI->getVarInfo ());
@@ -348,21 +371,54 @@ static void collectLoads(SILInstruction *I, SmallVectorImpl<LoadInst *> &Loads)
348371static void replaceLoad (LoadInst *LI, SILValue val, AllocStackInst *ASI) {
349372 ProjectionPath projections (val->getType ());
350373 SILValue op = LI->getOperand ();
374+ SILBuilderWithScope builder (LI);
375+
351376 while (op != ASI) {
352377 assert (isa<UncheckedAddrCastInst>(op) || isa<StructElementAddrInst>(op) ||
353378 isa<TupleElementAddrInst>(op));
354379 auto *Inst = cast<SingleValueInstruction>(op);
355380 projections.push_back (Projection (Inst));
356381 op = Inst->getOperand (0 );
357382 }
358- SILBuilder builder (LI);
383+
384+ SmallVector<SILValue, 4 > borrowedVals;
359385 for (auto iter = projections.rbegin (); iter != projections.rend (); ++iter) {
360386 const Projection &projection = *iter;
387+ assert (projection.getKind () == ProjectionKind::BitwiseCast ||
388+ projection.getKind () == ProjectionKind::Struct ||
389+ projection.getKind () == ProjectionKind::Tuple);
390+
391+ // struct_extract and tuple_extract expect guaranteed operand ownership
392+ // non-trivial RunningVal is owned. Insert borrow operation to convert
393+ if (projection.getKind () == ProjectionKind::Struct ||
394+ projection.getKind () == ProjectionKind::Tuple) {
395+ SILValue opVal = builder.emitBeginBorrowOperation (LI->getLoc (), val);
396+ if (opVal != val) {
397+ borrowedVals.push_back (opVal);
398+ val = opVal;
399+ }
400+ }
361401 val = projection.createObjectProjection (builder, LI->getLoc (), val).get ();
362402 }
403+
363404 op = LI->getOperand ();
364- LI->replaceAllUsesWith (val);
405+ // Replace users of the loaded value with `val`
406+ // If we have a load [copy], replace the users with copy_value of `val`
407+ if (LI->getOwnershipQualifier () == LoadOwnershipQualifier::Copy) {
408+ LI->replaceAllUsesWith (builder.createCopyValue (LI->getLoc (), val));
409+ } else {
410+ assert (!ASI->getFunction ()->hasOwnership () ||
411+ val.getOwnershipKind () != ValueOwnershipKind::Guaranteed);
412+ LI->replaceAllUsesWith (val);
413+ }
414+
415+ for (auto borrowedVal : borrowedVals) {
416+ builder.emitEndBorrowOperation (LI->getLoc (), borrowedVal);
417+ }
418+
419+ // Delete the load
365420 LI->eraseFromParent ();
421+
366422 while (op != ASI && op->use_empty ()) {
367423 assert (isa<UncheckedAddrCastInst>(op) || isa<StructElementAddrInst>(op) ||
368424 isa<TupleElementAddrInst>(op));
@@ -399,6 +455,7 @@ StoreInst *
399455StackAllocationPromoter::promoteAllocationInBlock (SILBasicBlock *BB) {
400456 LLVM_DEBUG (llvm::dbgs () << " *** Promoting ASI in block: " << *ASI);
401457
458+ // RunningVal is the current value in the stack location.
402459 // We don't know the value of the alloca until we find the first store.
403460 SILValue RunningVal = SILValue ();
404461 // Keep track of the last StoreInst that we found.
@@ -415,12 +472,16 @@ StackAllocationPromoter::promoteAllocationInBlock(SILBasicBlock *BB) {
415472 // If we are loading from the AllocStackInst and we already know the
416473 // content of the Alloca then use it.
417474 LLVM_DEBUG (llvm::dbgs () << " *** Promoting load: " << *Load);
418-
419475 replaceLoad (Load, RunningVal, ASI);
420476 ++NumInstRemoved;
421- } else if (Load->getOperand () == ASI) {
477+ } else if (Load->getOperand () == ASI &&
478+ Load->getOwnershipQualifier () !=
479+ LoadOwnershipQualifier::Copy) {
422480 // If we don't know the content of the AllocStack then the loaded
423481 // value *is* the new value;
482+ // Don't use result of load [copy] as a RunningVal, it necessitates
483+ // additional logic for cleanup of consuming instructions of the result.
484+ // StackAllocationPromoter::fixBranchesAndUses will later handle it.
424485 LLVM_DEBUG (llvm::dbgs () << " *** First load: " << *Load);
425486 RunningVal = Load;
426487 }
@@ -433,16 +494,51 @@ StackAllocationPromoter::promoteAllocationInBlock(SILBasicBlock *BB) {
433494 if (SI->getDest () != ASI)
434495 continue ;
435496
436- // The stored value is the new running value.
437- RunningVal = SI->getSrc ();
497+ // Special handling of entry block
498+ // If we have a store [assign] in the first block, OSSA guarantees we can
499+ // find the previous value stored in the stack location in RunningVal.
500+ // Create destroy_value of the RunningVal.
501+ // For all other blocks we may not know the previous value stored in the
502+ // stack location. So we will create destroy_value in
503+ // StackAllocationPromoter::fixBranchesAndUses, by getting the live-in
504+ // value to the block.
505+ if (BB->isEntry ()) {
506+ if (SI->getOwnershipQualifier () == StoreOwnershipQualifier::Assign) {
507+ assert (RunningVal);
508+ SILBuilderWithScope (SI).createDestroyValue (SI->getLoc (), RunningVal);
509+ }
510+ }
438511
439512 // If we met a store before this one, delete it.
513+ // If the LastStore was a store with [assign], delete it only if we know
514+ // the RunningValue to destroy. If not, it will be deleted in
515+ // StackAllocationPromoter::fixBranchesAndUses.
440516 if (LastStore) {
441- ++NumInstRemoved;
442- LLVM_DEBUG (llvm::dbgs () << " *** Removing redundant store: "
443- << *LastStore);
444- LastStore->eraseFromParent ();
517+ if (LastStore->getOwnershipQualifier () ==
518+ StoreOwnershipQualifier::Assign) {
519+ if (RunningVal) {
520+ // For entry block, we would have already created the destroy_value,
521+ // skip it.
522+ if (!BB->isEntry ()) {
523+ SILBuilderWithScope (LastStore).createDestroyValue (
524+ LastStore->getLoc (), RunningVal);
525+ }
526+ LLVM_DEBUG (llvm::dbgs ()
527+ << " *** Removing redundant store: " << *LastStore);
528+ ++NumInstRemoved;
529+ LastStore->eraseFromParent ();
530+ }
531+ } else {
532+ LLVM_DEBUG (llvm::dbgs ()
533+ << " *** Removing redundant store: " << *LastStore);
534+ ++NumInstRemoved;
535+ LastStore->eraseFromParent ();
536+ }
445537 }
538+
539+ // The stored value is the new running value.
540+ RunningVal = SI->getSrc ();
541+ // The current store is now the LastStore
446542 LastStore = SI;
447543 continue ;
448544 }
@@ -466,10 +562,23 @@ StackAllocationPromoter::promoteAllocationInBlock(SILBasicBlock *BB) {
466562 continue ;
467563 }
468564
565+ if (auto *DVI = dyn_cast<DestroyValueInst>(Inst)) {
566+ if (DVI->getOperand () == RunningVal) {
567+ // Reset LastStore.
568+ // So that we don't end up passing dead values as phi args in
569+ // StackAllocationPromoter::fixBranchesAndUses
570+ LastStore = nullptr ;
571+ }
572+ }
573+
469574 // Stop on deallocation.
470575 if (auto *DSI = dyn_cast<DeallocStackInst>(Inst)) {
471- if (DSI->getOperand () == ASI)
576+ if (DSI->getOperand () == ASI) {
577+ // Reset LastStore.
578+ // So that we don't pass RunningVal as a phi arg beyond dealloc_stack
579+ LastStore = nullptr ;
472580 break ;
581+ }
473582 }
474583 }
475584 if (LastStore) {
@@ -512,6 +621,10 @@ void MemoryToRegisters::removeSingleBlockAllocation(AllocStackInst *ASI) {
512621 // value.
513622 if (auto *SI = dyn_cast<StoreInst>(Inst)) {
514623 if (SI->getDest () == ASI) {
624+ if (SI->getOwnershipQualifier () == StoreOwnershipQualifier::Assign) {
625+ assert (RunningVal);
626+ SILBuilderWithScope (SI).createDestroyValue (SI->getLoc (), RunningVal);
627+ }
515628 RunningVal = SI->getSrc ();
516629 Inst->eraseFromParent ();
517630 ++NumInstRemoved;
@@ -643,6 +756,21 @@ void StackAllocationPromoter::fixPhiPredBlock(BlockSet &PhiBlocks,
643756 TI->eraseFromParent ();
644757}
645758
759+ static bool hasOnlyUndefIncomingValues (SILPhiArgument *phiArg) {
760+ SmallVector<SILValue, 8 > incomingValues;
761+ phiArg->getIncomingPhiValues (incomingValues);
762+ for (auto predArg : incomingValues) {
763+ if (isa<SILUndef>(predArg))
764+ continue ;
765+ if (isa<SILPhiArgument>(predArg) &&
766+ hasOnlyUndefIncomingValues (cast<SILPhiArgument>(predArg))) {
767+ continue ;
768+ }
769+ return false ;
770+ }
771+ return true ;
772+ }
773+
646774void StackAllocationPromoter::fixBranchesAndUses (BlockSet &PhiBlocks) {
647775 // First update uses of the value.
648776 SmallVector<LoadInst *, 4 > collectedLoads;
@@ -679,6 +807,16 @@ void StackAllocationPromoter::fixBranchesAndUses(BlockSet &PhiBlocks) {
679807 // on.
680808 SILBasicBlock *BB = Inst->getParent ();
681809
810+ if (!BB->isEntry ()) {
811+ if (auto *SI = dyn_cast<StoreInst>(Inst)) {
812+ if (SI->getOwnershipQualifier () == StoreOwnershipQualifier::Assign) {
813+ SILValue Def = getLiveInValue (PhiBlocks, BB);
814+ SILBuilderWithScope (SI).createDestroyValue (SI->getLoc (), Def);
815+ continue ;
816+ }
817+ }
818+ }
819+
682820 if (auto *DVAI = dyn_cast<DebugValueAddrInst>(Inst)) {
683821 // Replace DebugValueAddr with DebugValue.
684822 SILValue Def = getLiveInValue (PhiBlocks, BB);
@@ -710,6 +848,22 @@ void StackAllocationPromoter::fixBranchesAndUses(BlockSet &PhiBlocks) {
710848 fixPhiPredBlock (PhiBlocks, Block, PBB);
711849 }
712850 }
851+
852+ // If the owned phi arg we added did not have any uses, create end_lifetime to
853+ // end its lifetime. In asserts mode, make sure we have only undef incoming
854+ // values for such phi args.
855+ if (ASI->getFunction ()->hasOwnership ()) {
856+ for (auto Block : PhiBlocks) {
857+ auto *phiArg = cast<SILPhiArgument>(
858+ Block->getArgument (Block->getNumArguments () - 1 ));
859+ if (phiArg->getOwnershipKind () == ValueOwnershipKind::Owned &&
860+ phiArg->use_empty ()) {
861+ assert (hasOnlyUndefIncomingValues (phiArg));
862+ SILBuilderWithScope (&Block->front ())
863+ .createEndLifetime (Block->front ().getLoc (), phiArg);
864+ }
865+ }
866+ }
713867}
714868
715869void StackAllocationPromoter::pruneAllocStackUsage () {
@@ -956,10 +1110,6 @@ class SILMem2Reg : public SILFunctionTransform {
9561110 void run () override {
9571111 SILFunction *F = getFunction ();
9581112
959- // FIXME: We should be able to support ownership.
960- if (F->hasOwnership ())
961- return ;
962-
9631113 LLVM_DEBUG (llvm::dbgs () << " ** Mem2Reg on function: " << F->getName ()
9641114 << " **\n " );
9651115
0 commit comments