[MemCpyOpt] Check for throwing calls during call slot optimization

[nikic]

Summary

When performing call slot optimization for a non-local target, we need to check whether there may be throwing calls between the call and the copy. Otherwise, the copy may never be reached.

This was already done for call slot optimization of load/store, but not for memcpys. For the sake of clarity, I'm moving this check into the common optimization function, even if that does need an additional instruction scan for the load/store case.

[nikic]

llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp

835

Replying to @efriedma's comment on https://reviews.llvm.org/D88805#inline-824472: I believe checking for throwing instructions rather than guaranteed-to-transfer is correct here, because cpyDest cannot be read between the call and the copy, as a precondition for the call-slot optimization. So even if one of the instructions does not return, an early write to cpyDest will not be observable, with two caveats:

• The call itself might read cpyDest. We explicitly protect against this below (see the getModRefInfo and callCapturesBefore checks).
• We might throw and the catching code might read cpyDest. That's what we protect against here.

[efriedma]

llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp

835

The case I'm concerned about is a multithreaded context. Say there's an infinite loop between the call and the write to cpyDest. Then the current thread never touches cpyDest at all, so some other thread could legally use the memory.

If we're sure no other thread can access the memory, ensuring we don't throw/longjmp would be sufficient, for the reasons you describe.

[nikic]

Add more detailed comment and TODO.

[nikic]

llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp

835

I think you are right, and there is nothing that prevents that from happening right now. Unfortunately this also seems like a more significant limitation (not just because of the willreturn requirement, but also because this affects captured allocas, unlike the unwinding case). For now I've added a more extensive comment listing the different cases and marked this part as a TODO.

[efriedma]

LGTM; this seems like a step in the right direction in any case.

[jdoerfert]

llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp

847

You could use the fn attributes to avoid this loop all together. nounwind should imply this just fine.

llvm/include/llvm/Transforms/Scalar/MemCpyOptimizer.h

bool runImpl(Function &F, MemoryDependenceResults *MD_,
                AssumptionCache *AC_, DominatorTree *DT_, MemorySSA *MSSA_);
 
 private:
   // Helper functions
   bool processStore(StoreInst *SI, BasicBlock::iterator &BBI);
   bool processMemSet(MemSetInst *SI, BasicBlock::iterator &BBI);
   bool processMemCpy(MemCpyInst *M, BasicBlock::iterator &BBI);
   bool processMemMove(MemMoveInst *M);
-  bool performCallSlotOptzn(Instruction *cpy, Value *cpyDst, Value *cpySrc,
-                            uint64_t cpyLen, Align cpyAlign, CallInst *C);
+  bool performCallSlotOptzn(Instruction *cpyLoad, Instruction *cpyStore,
+                            Value *cpyDst, Value *cpySrc, uint64_t cpyLen,
+                            Align cpyAlign, CallInst *C);
   bool processMemCpyMemCpyDependence(MemCpyInst *M, MemCpyInst *MDep);
   bool processMemSetMemCpyDependence(MemCpyInst *MemCpy, MemSetInst *MemSet);
   bool performMemCpyToMemSetOptzn(MemCpyInst *MemCpy, MemSetInst *MemSet);
   bool processByValArgument(CallBase &CB, unsigned ArgNo);
   Instruction *tryMergingIntoMemset(Instruction *I, Value *StartPtr,
                                     Value *ByteVal);
   bool moveUp(StoreInst *SI, Instruction *P, const LoadInst *LI);
 
   void eraseInstruction(Instruction *I);
   bool iterateOnFunction(Function &F);
 };
 
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_SCALAR_MEMCPYOPTIMIZER_H

llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp

if (LI->isSimple() && LI->hasOneUse() &&
       MemDepResult ldep = MD->getDependency(LI);
       CallInst *C = nullptr;
       if (ldep.isClobber() && !isa<MemCpyInst>(ldep.getInst()))
         C = dyn_cast<CallInst>(ldep.getInst());
 
       if (C) {
         // Check that nothing touches the dest of the "copy" between
         // the call and the store.
-        Value *CpyDest = SI->getPointerOperand()->stripPointerCasts();
-        bool CpyDestIsLocal = isa<AllocaInst>(CpyDest);
         MemoryLocation StoreLoc = MemoryLocation::get(SI);
         for (BasicBlock::iterator I = --SI->getIterator(), E = C->getIterator();
              I != E; --I) {
           if (isModOrRefSet(AA->getModRefInfo(&*I, StoreLoc))) {
             C = nullptr;
             break;
           }
-          // The store to dest may never happen if an exception can be thrown
-          // between the load and the store.
-          if (I->mayThrow() && !CpyDestIsLocal) {
-            C = nullptr;
-            break;
-          }
         }
       }
 
       if (C) {
         bool changed = performCallSlotOptzn(
-            LI, SI->getPointerOperand()->stripPointerCasts(),
+            LI, SI, SI->getPointerOperand()->stripPointerCasts(),
             LI->getPointerOperand()->stripPointerCasts(),
             DL.getTypeStoreSize(SI->getOperand(0)->getType()),
             commonAlignment(SI->getAlign(), LI->getAlign()), C);
         if (changed) {
           eraseInstruction(SI);
           eraseInstruction(LI);
           ++NumMemCpyInstr;
           return true;
if (Instruction *I = tryMergingIntoMemset(MSI, MSI->getDest(),
       return true;
     }
   return false;
 }
 
 /// Takes a memcpy and a call that it depends on,
 /// and checks for the possibility of a call slot optimization by having
 /// the call write its result directly into the destination of the memcpy.
-bool MemCpyOptPass::performCallSlotOptzn(Instruction *cpy, Value *cpyDest,
+bool MemCpyOptPass::performCallSlotOptzn(Instruction *cpyLoad,
+                                         Instruction *cpyStore, Value *cpyDest,
                                          Value *cpySrc, uint64_t cpyLen,
                                          Align cpyAlign, CallInst *C) {
   // The general transformation to keep in mind is
   //
   //   call @func(..., src, ...)
   //   memcpy(dest, src, ...)
   //
   // ->
bool MemCpyOptPass::performCallSlotOptzn(Instruction *cpyLoad,
   AllocaInst *srcAlloca = dyn_cast<AllocaInst>(cpySrc);
   if (!srcAlloca)
     return false;
 
   ConstantInt *srcArraySize = dyn_cast<ConstantInt>(srcAlloca->getArraySize());
   if (!srcArraySize)
     return false;
 
-  const DataLayout &DL = cpy->getModule()->getDataLayout();
+  const DataLayout &DL = cpyLoad->getModule()->getDataLayout();
   uint64_t srcSize = DL.getTypeAllocSize(srcAlloca->getAllocatedType()) *
                      srcArraySize->getZExtValue();
 
   if (cpyLen < srcSize)
     return false;
 
   // Check that accessing the first srcSize bytes of dest will not cause a
   // trap.  Otherwise the transform is invalid since it might cause a trap
   // to occur earlier than it otherwise would.
+  // TODO: Use isDereferenceablePointer() API instead.
   if (AllocaInst *A = dyn_cast<AllocaInst>(cpyDest)) {
     // The destination is an alloca.  Check it is larger than srcSize.
     ConstantInt *destArraySize = dyn_cast<ConstantInt>(A->getArraySize());
     if (!destArraySize)
       return false;
 
     uint64_t destSize = DL.getTypeAllocSize(A->getAllocatedType()) *
                         destArraySize->getZExtValue();
 
     if (destSize < srcSize)
       return false;
   } else if (Argument *A = dyn_cast<Argument>(cpyDest)) {
-    // The store to dest may never happen if the call can throw.
-    if (C->mayThrow())
-      return false;
-
     if (A->getDereferenceableBytes() < srcSize) {
       // If the destination is an sret parameter then only accesses that are
       // outside of the returned struct type can trap.
       if (!A->hasStructRetAttr())
         return false;
 
       Type *StructTy = A->getParamStructRetType();
       if (!StructTy->isSized()) {
         // The call may never return and hence the copy-instruction may never
         // be executed, and therefore it's not safe to say "the destination
         // has at least <cpyLen> bytes, as implied by the copy-instruction",
         return false;
       }
 
       uint64_t destSize = DL.getTypeAllocSize(StructTy);
       if (destSize < srcSize)
         return false;
     }
   } else {
     return false;
   }
 
+  // Make sure that nothing can observe cpyDest being written early. There are
+  // a number of cases to consider:
+  //  1. cpyDest cannot be accessed between C and cpyStore as a precondition of
+  //     the transform.
+  //  2. C itself may not access cpyDest (prior to the transform). This is
+  //     checked further below.
+  //  3. If cpyDest is accessible to the caller of this function (potentially
+  //     captured and not based on an alloca), we need to ensure that we cannot
+  //     unwind between C and cpyStore. This is checked here.
+  //  4. If cpyDest is potentially captured, there may be accesses to it from
+  //     another thread. In this case, we need to check that cpyStore is
+  //     guaranteed to be executed if C is. As it is a non-atomic access, it
+  //     renders accesses from other threads undefined.
+  //     TODO: This is currently not checked.
+  if (!isa<AllocaInst>(cpyDest)) {
+    assert(C->getParent() == cpyStore->getParent() &&
+           "call and copy must be in the same block");
+    for (const Instruction &I : make_range(C->getIterator(),
+                                           cpyStore->getIterator())) {
+      if (I.mayThrow())
+        return false;
+    }
+  }
+
   // Check that dest points to memory that is at least as aligned as src.
   Align srcAlign = srcAlloca->getAlign();
   bool isDestSufficientlyAligned = srcAlign <= cpyAlign;
   // If dest is not aligned enough and we can't increase its alignment then
   // bail out.
   if (!isDestSufficientlyAligned && !isa<AllocaInst>(cpyDest))
     return false;
 
if (GetElementPtrInst *G = dyn_cast<GetElementPtrInst>(U)) {
       for (User *UU : U->users())
         srcUseList.push_back(UU);
       continue;
     }
     if (const IntrinsicInst *IT = dyn_cast<IntrinsicInst>(U))
       if (IT->isLifetimeStartOrEnd())
         continue;
 
-    if (U != C && U != cpy)
+    if (U != C && U != cpyLoad)
       return false;
   }
 
   // Check that src isn't captured by the called function since the
   // transformation can cause aliasing issues in that case.
   for (unsigned ArgI = 0, E = C->arg_size(); ArgI != E; ++ArgI)
     if (C->getArgOperand(ArgI) == cpySrc && !C->doesNotCapture(ArgI))
       return false;
bool MemCpyOptPass::performCallSlotOptzn(Instruction *cpyLoad,
 
   // Update AA metadata
   // FIXME: MD_tbaa_struct and MD_mem_parallel_loop_access should also be
   // handled here, but combineMetadata doesn't support them yet
   unsigned KnownIDs[] = {LLVMContext::MD_tbaa, LLVMContext::MD_alias_scope,
                          LLVMContext::MD_noalias,
                          LLVMContext::MD_invariant_group,
                          LLVMContext::MD_access_group};
-  combineMetadata(C, cpy, KnownIDs, true);
+  combineMetadata(C, cpyLoad, KnownIDs, true);
 
   return true;
 }
 
 /// We've found that the (upward scanning) memory dependence of memcpy 'M' is
 /// the memcpy 'MDep'. Try to simplify M to copy from MDep's input if we can.
 bool MemCpyOptPass::processMemCpyMemCpyDependence(MemCpyInst *M,
                                                   MemCpyInst *MDep) {
bool MemCpyOptPass::processMemCpy(MemCpyInst *M, BasicBlock::iterator &BBI) {
   //      memcpy in favor of the data that was already at the destination.
   //   d) memcpy from a just-memset'd source can be turned into memset.
   if (DepInfo.isClobber()) {
     if (CallInst *C = dyn_cast<CallInst>(DepInfo.getInst())) {
       // FIXME: Can we pass in either of dest/src alignment here instead
       // of conservatively taking the minimum?
       Align Alignment = std::min(M->getDestAlign().valueOrOne(),
                                  M->getSourceAlign().valueOrOne());
-      if (performCallSlotOptzn(M, M->getDest(), M->getSource(),
+      if (performCallSlotOptzn(M, M, M->getDest(), M->getSource(),
                                CopySize->getZExtValue(), Alignment, C)) {
         eraseInstruction(M);
         ++NumMemCpyInstr;
         return true;
       }
     }
   }
 

llvm/test/Transforms/MemCpyOpt/callslot.ll

;
   call void @llvm.memcpy.p0i8.p0i8.i64(i8* %dest.i8, i8* %src.i8, i64 16, i1 false)
   ret void
 }
 
 define void @throw_between_call_and_mempy(i8* dereferenceable(16) %dest.i8) {
 ; CHECK-LABEL: @throw_between_call_and_mempy(
 ; CHECK-NEXT:    [[SRC:%.*]] = alloca [16 x i8], align 1
 ; CHECK-NEXT:    [[SRC_I8:%.*]] = bitcast [16 x i8]* [[SRC]] to i8*
-; CHECK-NEXT:    [[DEST_I81:%.*]] = bitcast i8* [[DEST_I8:%.*]] to [16 x i8]*
-; CHECK-NEXT:    [[DEST_I812:%.*]] = bitcast [16 x i8]* [[DEST_I81]] to i8*
-; CHECK-NEXT:    call void @llvm.memset.p0i8.i64(i8* [[DEST_I812]], i8 0, i64 16, i1 false)
+; CHECK-NEXT:    call void @llvm.memset.p0i8.i64(i8* [[SRC_I8]], i8 0, i64 16, i1 false)
 ; CHECK-NEXT:    call void @may_throw() [[ATTR2:#.*]]
+; CHECK-NEXT:    call void @llvm.memset.p0i8.i64(i8* [[DEST_I8:%.*]], i8 0, i64 16, i1 false)
 ; CHECK-NEXT:    ret void
 ;
   %src = alloca [16 x i8]
   %src.i8 = bitcast [16 x i8]* %src to i8*
   call void @llvm.memset.p0i8.i64(i8* %src.i8, i8 0, i64 16, i1 false)
   call void @may_throw() readnone
   call void @llvm.memcpy.p0i8.p0i8.i64(i8* %dest.i8, i8* %src.i8, i64 16, i1 false)
   ret void