From 4da9193a6561b1a0f539ad8ef8c496e7b6bc0b14 Mon Sep 17 00:00:00 2001
From: Teresa Johnson <tejohnson@google.com>
Date: Sat, 1 Jul 2017 03:24:10 +0000
Subject: [PATCH] Recommit "r306541 - Add zero-length check to memcpy/memset
 load store loop expansion""

With fix for use-after-free errors. We can't add the new branch and
remove the old one until we are done with the Builder constructed for
the block.

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@306937 91177308-0d34-0410-b5e6-96231b3b80d8
---
 lib/Transforms/Utils/LowerMemIntrinsics.cpp | 19 ++++++++++++++-----
 test/CodeGen/NVPTX/lower-aggr-copies.ll     |  4 ++++
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/lib/Transforms/Utils/LowerMemIntrinsics.cpp b/lib/Transforms/Utils/LowerMemIntrinsics.cpp
index 0a51f9a0e4a..1c2a60a6b8b 100644
--- a/lib/Transforms/Utils/LowerMemIntrinsics.cpp
+++ b/lib/Transforms/Utils/LowerMemIntrinsics.cpp
@@ -27,7 +27,6 @@ void llvm::createMemCpyLoop(Instruction *InsertBefore,
   BasicBlock *LoopBB = BasicBlock::Create(F->getContext(), "loadstoreloop",
                                           F, NewBB);
 
-  OrigBB->getTerminator()->setSuccessor(0, LoopBB);
   IRBuilder<> Builder(OrigBB->getTerminator());
 
   // SrcAddr and DstAddr are expected to be pointer types,
@@ -39,6 +38,11 @@ void llvm::createMemCpyLoop(Instruction *InsertBefore,
   SrcAddr = Builder.CreateBitCast(SrcAddr, Builder.getInt8PtrTy(SrcAS));
   DstAddr = Builder.CreateBitCast(DstAddr, Builder.getInt8PtrTy(DstAS));
 
+  Builder.CreateCondBr(
+      Builder.CreateICmpEQ(ConstantInt::get(TypeOfCopyLen, 0), CopyLen), NewBB,
+      LoopBB);
+  OrigBB->getTerminator()->eraseFromParent();
+
   IRBuilder<> LoopBuilder(LoopBB);
   PHINode *LoopIndex = LoopBuilder.CreatePHI(TypeOfCopyLen, 0);
   LoopIndex->addIncoming(ConstantInt::get(TypeOfCopyLen, 0), OrigBB);
@@ -167,6 +171,7 @@ static void createMemMoveLoop(Instruction *InsertBefore,
 static void createMemSetLoop(Instruction *InsertBefore,
                              Value *DstAddr, Value *CopyLen, Value *SetValue,
                              unsigned Align, bool IsVolatile) {
+  Type *TypeOfCopyLen = CopyLen->getType();
   BasicBlock *OrigBB = InsertBefore->getParent();
   Function *F = OrigBB->getParent();
   BasicBlock *NewBB =
@@ -174,7 +179,6 @@ static void createMemSetLoop(Instruction *InsertBefore,
   BasicBlock *LoopBB
     = BasicBlock::Create(F->getContext(), "loadstoreloop", F, NewBB);
 
-  OrigBB->getTerminator()->setSuccessor(0, LoopBB);
   IRBuilder<> Builder(OrigBB->getTerminator());
 
   // Cast pointer to the type of value getting stored
@@ -182,9 +186,14 @@ static void createMemSetLoop(Instruction *InsertBefore,
   DstAddr = Builder.CreateBitCast(DstAddr,
                                   PointerType::get(SetValue->getType(), dstAS));
 
+  Builder.CreateCondBr(
+      Builder.CreateICmpEQ(ConstantInt::get(TypeOfCopyLen, 0), CopyLen), NewBB,
+      LoopBB);
+  OrigBB->getTerminator()->eraseFromParent();
+
   IRBuilder<> LoopBuilder(LoopBB);
-  PHINode *LoopIndex = LoopBuilder.CreatePHI(CopyLen->getType(), 0);
-  LoopIndex->addIncoming(ConstantInt::get(CopyLen->getType(), 0), OrigBB);
+  PHINode *LoopIndex = LoopBuilder.CreatePHI(TypeOfCopyLen, 0);
+  LoopIndex->addIncoming(ConstantInt::get(TypeOfCopyLen, 0), OrigBB);
 
   LoopBuilder.CreateStore(
       SetValue,
@@ -192,7 +201,7 @@ static void createMemSetLoop(Instruction *InsertBefore,
       IsVolatile);
 
   Value *NewIndex =
-      LoopBuilder.CreateAdd(LoopIndex, ConstantInt::get(CopyLen->getType(), 1));
+      LoopBuilder.CreateAdd(LoopIndex, ConstantInt::get(TypeOfCopyLen, 1));
   LoopIndex->addIncoming(NewIndex, LoopBB);
 
   LoopBuilder.CreateCondBr(LoopBuilder.CreateICmpULT(NewIndex, CopyLen), LoopBB,
diff --git a/test/CodeGen/NVPTX/lower-aggr-copies.ll b/test/CodeGen/NVPTX/lower-aggr-copies.ll
index 192d4becb05..f522c6722ee 100644
--- a/test/CodeGen/NVPTX/lower-aggr-copies.ll
+++ b/test/CodeGen/NVPTX/lower-aggr-copies.ll
@@ -17,6 +17,8 @@ entry:
   ret i8* %dst
 
 ; IR-LABEL:   @memcpy_caller
+; IR:         [[CMPREG:%[0-9]+]] = icmp eq i64 0, %n
+; IR:         br i1 [[CMPREG]], label %split, label %loadstoreloop
 ; IR:         loadstoreloop:
 ; IR:         [[LOADPTR:%[0-9]+]] = getelementptr inbounds i8, i8* %src, i64
 ; IR-NEXT:    [[VAL:%[0-9]+]] = load i8, i8* [[LOADPTR]]
@@ -73,6 +75,8 @@ entry:
 
 ; IR-LABEL:   @memset_caller
 ; IR:         [[VAL:%[0-9]+]] = trunc i32 %c to i8
+; IR:         [[CMPREG:%[0-9]+]] = icmp eq i64 0, %n
+; IR:         br i1 [[CMPREG]], label %split, label %loadstoreloop
 ; IR:         loadstoreloop:
 ; IR:         [[STOREPTR:%[0-9]+]] = getelementptr inbounds i8, i8* %dst, i64
 ; IR-NEXT:    store i8 [[VAL]], i8* [[STOREPTR]]
-- 
2.40.0