[InstCombineCalls] Unfold element atomic memcpy instruction

[igor-laevsky]

Summary

This change teaches instruction combiner on how to unfold element atomic memcpy intrinsic into sequence of explicit atomic memory operations.

Also I added new IntrinsicInst subclass to provide user friendly interface for the element atomic memcpy intrinsic. In the future we want to integrate this new class into MemIntrinsic hierarchy which will allow us to easily write optimizations covering both atomic and non-atomic versions. However as a first step I decided to frame it as a separate entity and refactor it later after we determine which use cases generalised interface should cover.

[igor-laevsky]

Add tests

[reames]

include/llvm/IR/IntrinsicInst.h

155

minor: instead of "class handles" use "represents"

lib/Transforms/InstCombine/InstCombineCalls.cpp

63

Units? The name should ideally reflect.

119

Please use the same convention used for the normal memcpy. Lower the atomic memcopy, set the length of the original to zero, then remove the zero length one. This also requires you to implement a currently missed optimization for zero length copies.

131

Why 16 elements as opposed to 8 bytes? (i.e. what the non-atomic form uses)

167

The source alignment may not hold for all elements. i.e. your alignment might be 1024 with your element size being 8.

[igor-laevsky]

lib/Transforms/InstCombine/InstCombineCalls.cpp

131

This transformation is not exactly similar to the non-atomic unfolding. Non atomic version only unfolds into single load/store pair. Here we can generate more operations. Number 16 means that no more than 16 load/store pairs will be generated. I picked it as an arbitrary small integer but we can easily tune it afterwards.

167

Yes, thanks for the catch!

[igor-laevsky]

ping?

[sanjoy]

lgtm with nits inline

include/llvm/IR/IntrinsicInst.h

171

No need to assert on isa -- that assert should implicitly happen during cast<>

lib/Transforms/InstCombine/InstCombineCalls.cpp

66

I'd rephrase this slightly as "Maximum number of elements in atomic memcpies the optimizer is allowed to unfold".

(i.e. no need to get into details -- this is a hidden option that only people who know what they're doing should use)

117

clang-format?

122

Minor: use auto * here.

132

s/ElementSize/ElementSizeInBits/

133

There is a DataLayout in InstCombiner.

165

I'd also prefer adding an assert here that checks what you said in the comment.

169

Note: you can do better here -- e.g. if the source is aligned to 16 bytes, and the element size is 4 bytes, then instead the alignment being 16, 4, 4, 4, ... it can be 16, 4, 8, 4, 16, ... etc.

For now I'll just leave a TODO.

181

Is it difficult to delete AMI right here? If not, I'd just do that -- no point in waiting for another iteration if it is obvious what will happen.

1538

This is really a separate transform -- I'd prefer splitting it out and landing it independently.

[igor-laevsky]

Second part: https://reviews.llvm.org/rL294453

lib/Transforms/InstCombine/InstCombineCalls.cpp

165

Proper checks of this invariants are quite verbose and will clobber the actual logic. I would prefer to rely on the verifier in this case.

181

Philip suggested that I should stick to the same process as other memory intrinsic optimizations do. I don't have a strong opinion here, but I would go with Philip's suggestion since I already implemented it.

include/llvm/IR/IntrinsicInst.h

public:
     static inline bool classof(const IntrinsicInst *I) {
       return I->getIntrinsicID() == Intrinsic::dbg_value;
     }
     static inline bool classof(const Value *V) {
       return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
     }
   };
 
+  /// This class handles atomic memcpy intrinsic
+  /// TODO: Integrate this class into MemIntrinsic hierarchy.
+  class ElementAtomicMemCpyInst : public IntrinsicInst {
+  public:
+    Value *getRawDest() const { return getArgOperand(0); }
+    Value *getRawSource() const { return getArgOperand(1); }
+    Value *getNumElements() const { return getArgOperand(2); }
+
+    uint64_t getSrcAlignment() const { return getParamAlignment(1); }
+    uint64_t getDstAlignment() const { return getParamAlignment(2); }
+
+    uint64_t getElementSizeInBytes() const {
+      Value *Arg = getArgOperand(3);
+      assert(isa<ConstantInt>(Arg));
+      return cast<ConstantInt>(Arg)->getZExtValue();
+    }
+
+    static inline bool classof(const IntrinsicInst *I) {
+      return I->getIntrinsicID() == Intrinsic::memcpy_element_atomic;
+    }
+    static inline bool classof(const Value *V) {
+      return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
+    }
+  };
+
   /// This is the common base class for memset/memcpy/memmove.
   class MemIntrinsic : public IntrinsicInst {
   public:
     Value *getRawDest() const { return const_cast<Value*>(getArgOperand(0)); }
     const Use &getRawDestUse() const { return getArgOperandUse(0); }
     Use &getRawDestUse() { return getArgOperandUse(0); }
 
     Value *getLength() const { return const_cast<Value*>(getArgOperand(2)); }

lib/Transforms/InstCombine/InstCombineCalls.cpp

 
 using namespace llvm;
 using namespace PatternMatch;
 
 #define DEBUG_TYPE "instcombine"
 
 STATISTIC(NumSimplified, "Number of library calls simplified");
 
+static cl::opt<uint64_t> UnfoldElementAtomicMemcpyThreshold(
+    "element-atomic-memcpy-threshold",
+    cl::init(16),
+    cl::desc("Threshold to determine when optimizer is allowed to unfold "
+             "atomic memcpy into sequence of atomic loads and stores"));
+
 /// Return the specified type promoted as it would be to pass though a va_arg
 /// area.
 static Type *getPromotedType(Type *Ty) {
   if (IntegerType* ITy = dyn_cast<IntegerType>(Ty)) {
     if (ITy->getBitWidth() < 32)
       return Type::getInt32Ty(Ty->getContext());
   }
   return Ty;
for (unsigned I = 0, E = V->getNumElements(); I != E; ++I) {
     bool Sign = V->getElementType()->isIntegerTy()
                     ? cast<ConstantInt>(Elt)->isNegative()
                     : cast<ConstantFP>(Elt)->isNegative();
     BoolVec.push_back(ConstantInt::get(BoolTy, Sign));
   }
   return ConstantVector::get(BoolVec);
 }
 
+/// Try to simplify element atomic memcpy intrinsic.
+/// Return true if changes were made and false otherwise.
+/// All newly inserted instructions will be added to the worklist.
+bool
+InstCombiner::SimplifyElementAtomicMemCpy(ElementAtomicMemCpyInst *AMI) {
+  // Try to unfold this intrinsic into sequence of explicit atomic loads and
+  // stores.
+  // First check that number of elements is compile time constant.
+  ConstantInt *NumElementsCI = dyn_cast<ConstantInt>(AMI->getNumElements());
+  if (!NumElementsCI)
+    return false;
+
+  // Check that there are not too many elements.
+  uint64_t NumElements = NumElementsCI->getZExtValue();
+  if (NumElements >= UnfoldElementAtomicMemcpyThreshold)
+    return false;
+
+  // Don't unfold into illegal integers
+  uint64_t ElementSize = AMI->getElementSizeInBytes() * 8;
+  if (!AMI->getModule()->getDataLayout().isLegalInteger(ElementSize))
+    return false;
+
+  // Cast source and destination to the correct type. Intrinsic input arguments
+  // are usually represented as i8*.
+  // Often operands will be explicitly casted to i8* and we can just strip
+  // those casts instead of inserting new ones. However it's easier to rely on
+  // other InstCombine rules which will cover trivial cases anyway.
+  Value *Src = AMI->getRawSource();
+  Value *Dst = AMI->getRawDest();
+  Type *ElementPointerType = Type::getIntNPtrTy(
+      AMI->getContext(), ElementSize, Src->getType()->getPointerAddressSpace());
+
+  Value *SrcCasted = Builder->CreatePointerCast(Src, ElementPointerType,
+                                                "memcpy_unfold.src_casted");
+  Value *DstCasted = Builder->CreatePointerCast(Dst, ElementPointerType,
+                                                "memcpy_unfold.dst_casted");
+
+  for (uint64_t i = 0; i < NumElements; ++i) {
+    // Get current element addresses
+    ConstantInt *ElementIdxCI =
+        ConstantInt::get(AMI->getContext(), APInt(64, i));
+    Value *SrcElementAddr =
+        Builder->CreateGEP(SrcCasted, ElementIdxCI, "memcpy_unfold.src_addr");
+    Value *DstElementAddr =
+        Builder->CreateGEP(DstCasted, ElementIdxCI, "memcpy_unfold.dst_addr");
+
+    // Load from the source. Transfer alignment information and mark load as
+    // unordered atomic.
+    LoadInst *Load = Builder->CreateLoad(SrcElementAddr, "memcpy_unfold.val");
+    Load->setOrdering(AtomicOrdering::Unordered);
+    Load->setAlignment(AMI->getSrcAlignment());
+    Load->setDebugLoc(AMI->getDebugLoc());
+
+    // Store loaded value via unordered atomic store.
+    StoreInst *Store = Builder->CreateStore(Load, DstElementAddr);
+    Store->setOrdering(AtomicOrdering::Unordered);
+    Store->setAlignment(AMI->getDstAlignment());
+    Store->setDebugLoc(AMI->getDebugLoc());
+  }
+
+  // Explicitly remove stale intrinsic because we can't DCE it.
+  eraseInstFromFunction(*AMI);
+
+  return true;
+}
+
 Instruction *InstCombiner::SimplifyMemTransfer(MemIntrinsic *MI) {
   unsigned DstAlign = getKnownAlignment(MI->getArgOperand(0), DL, MI, &AC, &DT);
   unsigned SrcAlign = getKnownAlignment(MI->getArgOperand(1), DL, MI, &AC, &DT);
   unsigned MinAlign = std::min(DstAlign, SrcAlign);
   unsigned CopyAlign = MI->getAlignment();
 
   if (CopyAlign < MinAlign) {
     MI->setAlignment(ConstantInt::get(MI->getAlignmentType(), MinAlign, false));
if (MemIntrinsic *MI = dyn_cast<MemIntrinsic>(II)) {
     } else if (MemSetInst *MSI = dyn_cast<MemSetInst>(MI)) {
       if (Instruction *I = SimplifyMemSet(MSI))
         return I;
     }
 
     if (Changed) return II;
   }
 
+  if (auto *AMI = dyn_cast<ElementAtomicMemCpyInst>(II)) {
+    if (SimplifyElementAtomicMemCpy(AMI))
+      // SimplifyElementAtomicMemCpy already added all relevant instructions
+      // to the worklist.
+      return nullptr;
+  }
+
   auto SimplifyDemandedVectorEltsLow = [this](Value *Op, unsigned Width,
                                               unsigned DemandedWidth) {
     APInt UndefElts(Width, 0);
     APInt DemandedElts = APInt::getLowBitsSet(Width, DemandedWidth);
     return SimplifyDemandedVectorElts(Op, DemandedElts, UndefElts);
   };
 
   switch (II->getIntrinsicID()) {

lib/Transforms/InstCombine/InstCombineInternal.h

private:
 
   Value *FoldLogicalPlusAnd(Value *LHS, Value *RHS, ConstantInt *Mask,
                             bool isSub, Instruction &I);
   Value *insertRangeTest(Value *V, const APInt &Lo, const APInt &Hi,
                          bool isSigned, bool Inside);
   Instruction *PromoteCastOfAllocation(BitCastInst &CI, AllocaInst &AI);
   Instruction *MatchBSwap(BinaryOperator &I);
   bool SimplifyStoreAtEndOfBlock(StoreInst &SI);
+
+  bool SimplifyElementAtomicMemCpy(ElementAtomicMemCpyInst *AMI);
   Instruction *SimplifyMemTransfer(MemIntrinsic *MI);
   Instruction *SimplifyMemSet(MemSetInst *MI);
 
   Value *EvaluateInDifferentType(Value *V, Type *Ty, bool isSigned);
 
   /// \brief Returns a value X such that Val = X * Scale, or null if none.
   ///
   /// If the multiplication is known not to overflow then NoSignedWrap is set.
   Value *Descale(Value *Val, APInt Scale, bool &NoSignedWrap);
 };
 
 } // end namespace llvm.
 
 #undef DEBUG_TYPE
 
 #endif

test/Transforms/InstCombine/element-atomic-memcpy-to-loads.ll

+; RUN: opt -instcombine -element-atomic-memcpy-threshold=8 -S < %s | FileCheck %s
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+
+; Test basic unfolding
+define void @test1(i8* %Src, i8* %Dst) {
+; CHECK-LABEL: test1
+
+; CHECK-DAG: %memcpy_unfold.src_casted = bitcast i8* %Src to i32*
+; CHECK-DAG: %memcpy_unfold.dst_casted = bitcast i8* %Dst to i32*
+
+; CHECK-DAG: [[VAL1:%[^\s]+]] =  load atomic i32, i32* %memcpy_unfold.src_casted unordered, align 4
+; CHECK-DAG: store atomic i32 [[VAL1]], i32* %memcpy_unfold.dst_casted unordered, align 8
+
+; CHECK-DAG: [[VAL2:%[^\s]+]] =  load atomic i32, i32* %{{[^\s]+}} unordered, align 4
+; CHECK-DAG: store atomic i32 [[VAL2]], i32* %{{[^\s]+}} unordered, align 8
+
+; CHECK-DAG: [[VAL3:%[^\s]+]] =  load atomic i32, i32* %{{[^\s]+}} unordered, align 4
+; CHECK-DAG: store atomic i32 [[VAL3]], i32* %{{[^\s]+}} unordered, align 8
+
+; CHECK-DAG: [[VAL4:%[^\s]+]] =  load atomic i32, i32* %{{[^\s]+}} unordered, align 4
+; CHECK-DAG: store atomic i32 [[VAL4]], i32* %{{[^\s]+}} unordered, align 8
+entry:
+  call void @llvm.memcpy.element.atomic.p0i8.p0i8(i8* align 4 %Dst, i8* align 8 %Src, i64 4, i32 4)
+  ret void
+}
+
+; Test that we don't unfold too much
+define void @test2(i8* %Src, i8* %Dst) {
+; CHECK-LABEL: test2
+
+; CHECK-NOT: load
+; CHECK-NOT: store
+; CHECK: llvm.memcpy.element.atomic
+entry:
+  call void @llvm.memcpy.element.atomic.p0i8.p0i8(i8* align 4 %Dst, i8* align 4 %Src, i64 1000, i32 4)
+  ret void
+}
+
+; Test that we will not unfold into non native integers
+define void @test3(i8* %Src, i8* %Dst) {
+; CHECK-LABEL: test3
+
+; CHECK-NOT: load
+; CHECK-NOT: store
+; CHECK: llvm.memcpy.element.atomic
+entry:
+  call void @llvm.memcpy.element.atomic.p0i8.p0i8(i8* align 64 %Dst, i8* align 64 %Src, i64 4, i32 64)
+  ret void
+}
+
+; Test that we will eliminate redundant bitcasts
+define void @test4(i64* %Src, i64* %Dst) {
+; CHECK-LABEL: test4
+
+; CHECK-NOT: bitcast
+
+; CHECK-DAG: [[VAL1:%[^\s]+]] =  load atomic i64, i64* %Src unordered, align 8
+; CHECK-DAG: store atomic i64 [[VAL1]], i64* %Dst unordered, align 8
+
+; CHECK-DAG: [[SRC_ADDR2:%[^ ]+]] = getelementptr i64, i64* %Src, i64 1
+; CHECK-DAG: [[DST_ADDR2:%[^ ]+]] = getelementptr i64, i64* %Dst, i64 1
+; CHECK-DAG: [[VAL2:%[^\s]+]] =  load atomic i64, i64* [[SRC_ADDR2]] unordered, align 8
+; CHECK-DAG: store atomic i64 [[VAL2]], i64* [[DST_ADDR2]] unordered, align 8
+
+; CHECK-DAG: [[SRC_ADDR3:%[^ ]+]] = getelementptr i64, i64* %Src, i64 2
+; CHECK-DAG: [[DST_ADDR3:%[^ ]+]] = getelementptr i64, i64* %Dst, i64 2
+; CHECK-DAG: [[VAL3:%[^ ]+]] =  load atomic i64, i64* [[SRC_ADDR3]] unordered, align 8
+; CHECK-DAG: store atomic i64 [[VAL3]], i64* [[DST_ADDR3]] unordered, align 8
+
+; CHECK-DAG: [[SRC_ADDR4:%[^ ]+]] = getelementptr i64, i64* %Src, i64 3
+; CHECK-DAG: [[DST_ADDR4:%[^ ]+]] = getelementptr i64, i64* %Dst, i64 3
+; CHECK-DAG: [[VAL4:%[^ ]+]] =  load atomic i64, i64* [[SRC_ADDR4]] unordered, align 8
+; CHECK-DAG: store atomic i64 [[VAL4]], i64* [[DST_ADDR4]] unordered, align 8
+entry:
+  %Src.casted = bitcast i64* %Src to i8*
+  %Dst.casted = bitcast i64* %Dst to i8*
+  call void @llvm.memcpy.element.atomic.p0i8.p0i8(i8* align 8 %Dst.casted, i8* align 8 %Src.casted, i64 4, i32 8)
+  ret void
+}
+
+declare void @llvm.memcpy.element.atomic.p0i8.p0i8(i8* nocapture, i8* nocapture, i64, i32)