2 files changed, 36 insertions, 12 deletions

diff --git a/llvm/lib/Target/AMDGPU/AMDGPULowerIntrinsics.cpp b/llvm/lib/Target/AMDGPU/AMDGPULowerIntrinsics.cpp
index 846e7dff5f8..7e0e9802c0e 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPULowerIntrinsics.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPULowerIntrinsics.cpp
@@ -10,6 +10,7 @@
 #include "AMDGPU.h"
 #include "AMDGPUSubtarget.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
@@ -34,9 +35,14 @@ public:
   AMDGPULowerIntrinsics() : ModulePass(ID) {}
 
   bool runOnModule(Module &M) override;
+  bool expandMemIntrinsicUses(Function &F);
   StringRef getPassName() const override {
     return "AMDGPU Lower Intrinsics";
   }
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<TargetTransformInfoWrapperPass>();
+  }
 };
 
 }
@@ -55,7 +61,7 @@ static bool shouldExpandOperationWithSize(Value *Size) {
   return !CI || (CI->getZExtValue() > MaxStaticSize);
 }
 
-static bool expandMemIntrinsicUses(Function &F) {
+bool AMDGPULowerIntrinsics::expandMemIntrinsicUses(Function &F) {
   Intrinsic::ID ID = F.getIntrinsicID();
   bool Changed = false;
 
@@ -67,7 +73,10 @@ static bool expandMemIntrinsicUses(Function &F) {
     case Intrinsic::memcpy: {
       auto *Memcpy = cast<MemCpyInst>(Inst);
       if (shouldExpandOperationWithSize(Memcpy->getLength())) {
-        expandMemCpyAsLoop(Memcpy);
+        Function *ParentFunc = Memcpy->getParent()->getParent();
+        const TargetTransformInfo &TTI =
+            getAnalysis<TargetTransformInfoWrapperPass>().getTTI(*ParentFunc);
+        expandMemCpyAsLoop(Memcpy, TTI);
         Changed = true;
         Memcpy->eraseFromParent();
       }
diff --git a/llvm/lib/Target/NVPTX/NVPTXLowerAggrCopies.cpp b/llvm/lib/Target/NVPTX/NVPTXLowerAggrCopies.cpp
index 3be291b48b8..989f0a3aba2 100644
--- a/llvm/lib/Target/NVPTX/NVPTXLowerAggrCopies.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXLowerAggrCopies.cpp
@@ -14,6 +14,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "NVPTXLowerAggrCopies.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/CodeGen/StackProtector.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
@@ -42,6 +43,7 @@ struct NVPTXLowerAggrCopies : public FunctionPass {
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addPreserved<StackProtector>();
+    AU.addRequired<TargetTransformInfoWrapperPass>();
   }
 
   bool runOnFunction(Function &F) override;
@@ -61,6 +63,8 @@ bool NVPTXLowerAggrCopies::runOnFunction(Function &F) {
 
   const DataLayout &DL = F.getParent()->getDataLayout();
   LLVMContext &Context = F.getParent()->getContext();
+  const TargetTransformInfo &TTI =
+      getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
 
   // Collect all aggregate loads and mem* calls.
   for (Function::iterator BI = F.begin(), BE = F.end(); BI != BE; ++BI) {
@@ -104,15 +108,26 @@ bool NVPTXLowerAggrCopies::runOnFunction(Function &F) {
     Value *SrcAddr = LI->getOperand(0);
     Value *DstAddr = SI->getOperand(1);
     unsigned NumLoads = DL.getTypeStoreSize(LI->getType());
-    Value *CopyLen = ConstantInt::get(Type::getInt32Ty(Context), NumLoads);
-
-    createMemCpyLoop(/* ConvertedInst */ SI,
-                     /* SrcAddr */ SrcAddr, /* DstAddr */ DstAddr,
-                     /* CopyLen */ CopyLen,
-                     /* SrcAlign */ LI->getAlignment(),
-                     /* DestAlign */ SI->getAlignment(),
-                     /* SrcIsVolatile */ LI->isVolatile(),
-                     /* DstIsVolatile */ SI->isVolatile());
+    ConstantInt *CopyLen =
+        ConstantInt::get(Type::getInt32Ty(Context), NumLoads);
+
+    if (!TTI.useWideIRMemcpyLoopLowering()) {
+      createMemCpyLoop(/* ConvertedInst */ SI,
+                       /* SrcAddr */ SrcAddr, /* DstAddr */ DstAddr,
+                       /* CopyLen */ CopyLen,
+                       /* SrcAlign */ LI->getAlignment(),
+                       /* DestAlign */ SI->getAlignment(),
+                       /* SrcIsVolatile */ LI->isVolatile(),
+                       /* DstIsVolatile */ SI->isVolatile());
+    } else {
+      createMemCpyLoopKnownSize(/* ConvertedInst */ SI,
+                                /* SrcAddr */ SrcAddr, /* DstAddr */ DstAddr,
+                                /* CopyLen */ CopyLen,
+                                /* SrcAlign */ LI->getAlignment(),
+                                /* DestAlign */ SI->getAlignment(),
+                                /* SrcIsVolatile */ LI->isVolatile(),
+                                /* DstIsVolatile */ SI->isVolatile(), TTI);
+    }
 
     SI->eraseFromParent();
     LI->eraseFromParent();
@@ -121,7 +136,7 @@ bool NVPTXLowerAggrCopies::runOnFunction(Function &F) {
   // Transform mem* intrinsic calls.
   for (MemIntrinsic *MemCall : MemCalls) {
     if (MemCpyInst *Memcpy = dyn_cast<MemCpyInst>(MemCall)) {
-      expandMemCpyAsLoop(Memcpy);
+      expandMemCpyAsLoop(Memcpy, TTI);
     } else if (MemMoveInst *Memmove = dyn_cast<MemMoveInst>(MemCall)) {
       expandMemMoveAsLoop(Memmove);
     } else if (MemSetInst *Memset = dyn_cast<MemSetInst>(MemCall)) {