[MemCpyOptimizer] Optimize passing byref function arguments down the stack

[atrosinenko]

Summary

When a byref function argument is passed through like this:

void leaf(struct X);
void middle(struct X a) {
  leaf(a);
}

... an unnecessary temporary copy may be introduced inside middle().

This patch makes that copy being eliminated by MemCpyOptimizer unless size of
the byref argument is 8 bytes or less (in that case memcpy() is eliminated by
InstCombine first preventing later optimization).

[rjmccall]

If this optimization is valid, it's valid regardless of byref.

[jdoerfert]

This needs an LLVM IR test.

llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp

1348

I thought we have a shortcut to check for both attributes. If not, we might want one.

[atrosinenko]

Restrict transformation to passing through byref arguments, not arbitrary pointers.

[jdoerfert]

This looks reasonable on first glance. @arsenm ?

[arsenm]

llvm/test/Transforms/MemCpyOpt/byref-memcpy.ll

44

I think this is missing an alignment check. Can you add a test where the callee requires a higher parameter alignment than the caller?

[rjmccall]

Sorry, I tried to say thiis more succinctly before, but what exactly is the semantic property of byref that allows this optimization?

[atrosinenko]

Add a test for insufficiently aligned source.

[atrosinenko]

@rjmccall

Maybe I overestimated similarity of byval and recently introduced byref... Looks like some aliasing restrictions are not mentioned in LLVM Language Reference. For example, the only way for Clang to emit the byref attribute to LLVM IR I know about is via ABIArgInfo with Kind == IndirectAliased introduced in D79744: clang: Use byref for aggregate kernel arguments - and it has a note that "The object is known to not be through any other references for the duration of the call, and the callee must not itself modify the object". This seems to be necessary for correctness of this transformation. If LLVM Language Reference does not mention such restrictions intentionally, then it may be clang that should be responsible for such optimizations.

llvm/test/Transforms/MemCpyOpt/byref-memcpy.ll

56

Interestingly, if I change the alignment of the second argument from 2 back to 4, then memcpy() is successfully dropped (as if everything has big-enough alignment) while it seemingly shouldn't be. Is it just due to the code being ill-formed?

[arsenm]

In D86020#2222103, @atrosinenko wrote:

@rjmccall

Maybe I overestimated similarity of byval and recently introduced byref... Looks like some aliasing restrictions are not mentioned in LLVM Language Reference. For example, the only way for Clang to emit the byref attribute to LLVM IR I know about is via ABIArgInfo with Kind == IndirectAliased introduced in D79744: clang: Use byref for aggregate kernel arguments - and it has a note that "The object is known to not be through any other references for the duration of the call, and the callee must not itself modify the object". This seems to be necessary for correctness of this transformation. If LLVM Language Reference does not mention such restrictions intentionally, then it may be clang that should be responsible for such optimizations.

I think this lost the word written or stored

llvm/test/Transforms/MemCpyOpt/byref-memcpy.ll

9

Typo attrbiute

56

I'm not sure I follow, if you change it to 4 it is big enough?

llvm/include/llvm/IR/InstrTypes.h

bool doesNotCapture(unsigned OpNo) const {
     return dataOperandHasImpliedAttr(OpNo + 1, Attribute::NoCapture);
   }
 
   /// Determine whether this argument is passed by value.
   bool isByValArgument(unsigned ArgNo) const {
     return paramHasAttr(ArgNo, Attribute::ByVal);
   }
 
+  /// Determine whether this argument is passed via pointer annotated with
+  /// `byref` attribute.
+  bool isByRefArgument(unsigned ArgNo) const {
+    return paramHasAttr(ArgNo, Attribute::ByRef);
+  }
+
   /// Determine whether this argument is passed in an alloca.
   bool isInAllocaArgument(unsigned ArgNo) const {
     return paramHasAttr(ArgNo, Attribute::InAlloca);
   }
 
   /// Determine whether this argument is passed by value, in an alloca, or is
   /// preallocated.
   bool isPassPointeeByValueArgument(unsigned ArgNo) const {

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

private:
   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 processMemCpyMemCpyDependence(MemCpyInst *M, MemCpyInst *MDep);
   bool processMemSetMemCpyDependence(MemCpyInst *MemCpy, MemSetInst *MemSet);
   bool performMemCpyToMemSetOptzn(MemCpyInst *MemCpy, MemSetInst *MemSet);
-  bool processByValArgument(CallBase &CB, unsigned ArgNo);
+  bool processByValOrByRefArgument(CallBase &CB, unsigned ArgNo);
   Instruction *tryMergingIntoMemset(Instruction *I, Value *StartPtr,
                                     Value *ByteVal);
 
   bool iterateOnFunction(Function &F);
 };
 
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_SCALAR_MEMCPYOPTIMIZER_H

llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp

bool MemCpyOptPass::processMemMove(MemMoveInst *M) {
   // MemDep may have over conservative information about this instruction, just
   // conservatively flush it from the cache.
   MD->removeInstruction(M);
 
   ++NumMoveToCpy;
   return true;
 }
 
-/// This is called on every byval argument in call sites.
-bool MemCpyOptPass::processByValArgument(CallBase &CB, unsigned ArgNo) {
+/// This is called on every byval/byref argument in call sites.
+bool MemCpyOptPass::processByValOrByRefArgument(CallBase &CB, unsigned ArgNo) {
   const DataLayout &DL = CB.getCaller()->getParent()->getDataLayout();
-  // Find out what feeds this byval argument.
+  // Find out what feeds this byval/byref argument.
   Value *ByValArg = CB.getArgOperand(ArgNo);
   Type *ByValTy = cast<PointerType>(ByValArg->getType())->getElementType();
   uint64_t ByValSize = DL.getTypeAllocSize(ByValTy);
   MemDepResult DepInfo = MD->getPointerDependencyFrom(
       MemoryLocation(ByValArg, LocationSize::precise(ByValSize)), true,
       CB.getIterator(), CB.getParent());
   if (!DepInfo.isClobber())
     return false;
 
   // If the byval argument isn't fed by a memcpy, ignore it.  If it is fed by
   // a memcpy, see if we can byval from the source of the memcpy instead of the
   // result.
   MemCpyInst *MDep = dyn_cast<MemCpyInst>(DepInfo.getInst());
   if (!MDep || MDep->isVolatile() ||
       ByValArg->stripPointerCasts() != MDep->getDest())
     return false;
 
+  // ByRef arguments should be handled with care: they are not implicitly copied
+  // to the calling function's stack frame, so they may be accessible from other
+  // threads, coroutines, etc. On the other hand, there is a special case: just
+  // passing through an argument down the stack like this:
+  //     void leaf(struct X /* byref(%struct.X) */);
+  //     void intermediate(struct X a /* byref(%struct.X) */) {
+  //       /* prologue: `a` is implicitly copied by clang in the callee frame */
+  //       leaf(a);
+  //     }
+  // So, checking whether the memcpy() source is byref argument on itself.
+  if (CB.isByRefArgument(ArgNo)) {
+    Argument *SourceAsArgument =
+        dyn_cast<Argument>(MDep->getSource()->stripPointerCasts());
+    if (!SourceAsArgument || !SourceAsArgument->hasByRefAttr())
+      return false;
+  }
+
   // The length of the memcpy must be larger or equal to the size of the byval.
   ConstantInt *C1 = dyn_cast<ConstantInt>(MDep->getLength());
   if (!C1 || C1->getValue().getZExtValue() < ByValSize)
     return false;
 
   // Get the alignment of the byval.  If the call doesn't specify the alignment,
   // then it is some target specific value that we can't know.
   MaybeAlign ByValAlign = CB.getParamAlign(ArgNo);
bool MemCpyOptPass::processByValOrByRefArgument(CallBase &CB, unsigned ArgNo) {
   if (MDep->getSource()->getType() != ByValArg->getType()) {
     BitCastInst *TmpBitCast = new BitCastInst(MDep->getSource(), ByValArg->getType(),
                                               "tmpcast", &CB);
     // Set the tmpcast's DebugLoc to MDep's
     TmpBitCast->setDebugLoc(MDep->getDebugLoc());
     TmpCast = TmpBitCast;
   }
 
-  LLVM_DEBUG(dbgs() << "MemCpyOptPass: Forwarding memcpy to byval:\n"
+  LLVM_DEBUG(dbgs() << "MemCpyOptPass: Forwarding memcpy to byval or byref:\n"
                     << "  " << *MDep << "\n"
                     << "  " << CB << "\n");
 
   // Otherwise we're good!  Update the byval argument.
   CB.setArgOperand(ArgNo, TmpCast);
   ++NumMemCpyInstr;
   return true;
 }
for (BasicBlock::iterator BI = BB.begin(), BE = BB.end(); BI != BE;) {
         MadeChange |= processStore(SI, BI);
       else if (MemSetInst *M = dyn_cast<MemSetInst>(I))
         RepeatInstruction = processMemSet(M, BI);
       else if (MemCpyInst *M = dyn_cast<MemCpyInst>(I))
         RepeatInstruction = processMemCpy(M, BI);
       else if (MemMoveInst *M = dyn_cast<MemMoveInst>(I))
         RepeatInstruction = processMemMove(M);
       else if (auto *CB = dyn_cast<CallBase>(I)) {
-        for (unsigned i = 0, e = CB->arg_size(); i != e; ++i)
-          if (CB->isByValArgument(i))
-            MadeChange |= processByValArgument(*CB, i);
+        for (unsigned i = 0, e = CB->arg_size(); i != e; ++i) {
+          if (CB->isByValArgument(i) || CB->isByRefArgument(i))
+            MadeChange |= processByValOrByRefArgument(*CB, i);
+        }
       }
 
       // Reprocess the instruction if desired.
       if (RepeatInstruction) {
         if (BI != BB.begin())
           --BI;
         MadeChange = true;
       }

llvm/test/Transforms/MemCpyOpt/byref-memcpy.ll

+; RUN: opt -memcpyopt -dse -S < %s | FileCheck %s
+
+%struct.S = type { i32 }
+
+declare void @llvm.memcpy.p0i8.p0i8.i16(i8* noalias nocapture writeonly, i8* noalias nocapture readonly, i16, i1 immarg)
+
+; A caller of this function has to arrange everything for the pointer it passes
+; to be usable with byref attrbiute anyway, so just pass the pointer through.
+define void @function_with_byref_arg(%struct.S* byref(%struct.S) align 4 %0) {
+; CHECK:       define void @function_with_byref_arg{{.*}}byref(%struct.S){{.*}} {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    call void @leaf
+; CHECK-NEXT:    ret void
+; CHECK-NEXT:  }
+entry:
+  %x = alloca %struct.S, align 4
+  %1 = bitcast %struct.S* %x to i8*
+  %2 = bitcast %struct.S* %0 to i8*
+  call void @llvm.memcpy.p0i8.p0i8.i16(i8* align 4 %1, i8* align 4 %2, i16 16, i1 false)
+  call void @leaf(%struct.S* byref(%struct.S) align 4 %x)
+  ret void
+}
+
+; A generic pointer may point to memory being modified by other threads,
+; coroutines, etc. during execution of @leaf().
+define void @function_with_pointer_arg(%struct.S* align 4 %0) {
+; CHECK:       define void @function_with_pointer_arg(%struct.S* align 4 [[V0:%.+]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[ALLOCA:%.+]] = alloca %struct.S, align 4
+; CHECK-NEXT:    [[V1:%.+]] = bitcast %struct.S* [[ALLOCA]] to i8*
+; CHECK-NEXT:    [[V2:%.+]] = bitcast %struct.S* [[V0]] to i8*
+; CHECK-NEXT:    call void @llvm.memcpy.p0i8.p0i8.i16(i8* align 4 %1, i8* align 4 %2, i16 16, i1 false)
+; CHECK-NEXT:    call void @leaf(%struct.S* byref(%struct.S) align 4 [[ALLOCA]])
+; CHECK-NEXT:    ret void
+; CHECK-NEXT:  }
+entry:
+  %x = alloca %struct.S, align 4
+  %1 = bitcast %struct.S* %x to i8*
+  %2 = bitcast %struct.S* %0 to i8*
+  call void @llvm.memcpy.p0i8.p0i8.i16(i8* align 4 %1, i8* align 4 %2, i16 16, i1 false)
+  call void @leaf(%struct.S* byref(%struct.S) align 4 %x)
+  ret void
+}
+
+declare void @leaf(%struct.S* byref(%struct.S) align 4)