Reapply "[AMDGPU] Handle memcpy()-like ops in LowerBufferFatPointers (llvm#126621)" (llvm#129078)

This reverts commit 1559a65.

Fixed test (I suspect broken by unrelated change in the merge)

Author: krzysz00

llvm/lib/Target/AMDGPU/AMDGPULowerBufferFatPointers.cpp

@@ -45,16 +45,16 @@
 //
 // This pass proceeds in three main phases:
 //
-// ## Rewriting loads and stores of p7
+// ## Rewriting loads and stores of p7 and memcpy()-like handling
 //
 // The first phase is to rewrite away all loads and stors of `ptr addrspace(7)`,
 // including aggregates containing such pointers, to ones that use `i160`. This
-// is handled by `StoreFatPtrsAsIntsVisitor` , which visits loads, stores, and
-// allocas and, if the loaded or stored type contains `ptr addrspace(7)`,
-// rewrites that type to one where the p7s are replaced by i160s, copying other
-// parts of aggregates as needed. In the case of a store, each pointer is
-// `ptrtoint`d to i160 before storing, and load integers are `inttoptr`d back.
-// This same transformation is applied to vectors of pointers.
+// is handled by `StoreFatPtrsAsIntsAndExpandMemcpyVisitor` , which visits
+// loads, stores, and allocas and, if the loaded or stored type contains `ptr
+// addrspace(7)`, rewrites that type to one where the p7s are replaced by i160s,
+// copying other parts of aggregates as needed. In the case of a store, each
+// pointer is `ptrtoint`d to i160 before storing, and load integers are
+// `inttoptr`d back. This same transformation is applied to vectors of pointers.
 //
 // Such a transformation allows the later phases of the pass to not need
 // to handle buffer fat pointers moving to and from memory, where we load
@@ -66,6 +66,10 @@
 // Atomics operations on `ptr addrspace(7)` values are not suppported, as the
 // hardware does not include a 160-bit atomic.
 //
+// In order to save on O(N) work and to ensure that the contents type
+// legalizer correctly splits up wide loads, also unconditionally lower
+// memcpy-like intrinsics into loops here.
+//
 // ## Buffer contents type legalization
 //
 // The underlying buffer intrinsics only support types up to 128 bits long,
@@ -231,20 +235,24 @@
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/InstVisitor.h"
 #include "llvm/IR/Instructions.h"
+#include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/IntrinsicsAMDGPU.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Operator.h"
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/IR/ReplaceConstant.h"
+#include "llvm/IR/ValueHandle.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
+#include "llvm/Support/AMDGPUAddrSpace.h"
 #include "llvm/Support/Alignment.h"
 #include "llvm/Support/AtomicOrdering.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/Local.h"
+#include "llvm/Transforms/Utils/LowerMemIntrinsics.h"
 #include "llvm/Transforms/Utils/ValueMapper.h"
 
 #define DEBUG_TYPE "amdgpu-lower-buffer-fat-pointers"
@@ -431,14 +439,16 @@ namespace {
 /// marshalling costs when reading or storing these values, but since placing
 /// such pointers into memory is an uncommon operation at best, we feel that
 /// this cost is acceptable for better performance in the common case.
-class StoreFatPtrsAsIntsVisitor
-    : public InstVisitor<StoreFatPtrsAsIntsVisitor, bool> {
+class StoreFatPtrsAsIntsAndExpandMemcpyVisitor
+    : public InstVisitor<StoreFatPtrsAsIntsAndExpandMemcpyVisitor, bool> {
   BufferFatPtrToIntTypeMap *TypeMap;
 
   ValueToValueMapTy ConvertedForStore;
 
   IRBuilder<> IRB;
 
+  const TargetMachine *TM;
+
   // Convert all the buffer fat pointers within the input value to inttegers
   // so that it can be stored in memory.
   Value *fatPtrsToInts(Value *V, Type *From, Type *To, const Twine &Name);
@@ -448,20 +458,27 @@ class StoreFatPtrsAsIntsVisitor
   Value *intsToFatPtrs(Value *V, Type *From, Type *To, const Twine &Name);
 
 public:
-  StoreFatPtrsAsIntsVisitor(BufferFatPtrToIntTypeMap *TypeMap, LLVMContext &Ctx)
-      : TypeMap(TypeMap), IRB(Ctx) {}
+  StoreFatPtrsAsIntsAndExpandMemcpyVisitor(BufferFatPtrToIntTypeMap *TypeMap,
+                                           LLVMContext &Ctx,
+                                           const TargetMachine *TM)
+      : TypeMap(TypeMap), IRB(Ctx), TM(TM) {}
   bool processFunction(Function &F);
 
   bool visitInstruction(Instruction &I) { return false; }
   bool visitAllocaInst(AllocaInst &I);
   bool visitLoadInst(LoadInst &LI);
   bool visitStoreInst(StoreInst &SI);
   bool visitGetElementPtrInst(GetElementPtrInst &I);
+
+  bool visitMemCpyInst(MemCpyInst &MCI);
+  bool visitMemMoveInst(MemMoveInst &MMI);
+  bool visitMemSetInst(MemSetInst &MSI);
+  bool visitMemSetPatternInst(MemSetPatternInst &MSPI);
 };
 } // namespace
 
-Value *StoreFatPtrsAsIntsVisitor::fatPtrsToInts(Value *V, Type *From, Type *To,
-                                                const Twine &Name) {
+Value *StoreFatPtrsAsIntsAndExpandMemcpyVisitor::fatPtrsToInts(
+    Value *V, Type *From, Type *To, const Twine &Name) {
   if (From == To)
     return V;
   ValueToValueMapTy::iterator Find = ConvertedForStore.find(V);
@@ -498,8 +515,8 @@ Value *StoreFatPtrsAsIntsVisitor::fatPtrsToInts(Value *V, Type *From, Type *To,
   return Ret;
 }
 
-Value *StoreFatPtrsAsIntsVisitor::intsToFatPtrs(Value *V, Type *From, Type *To,
-                                                const Twine &Name) {
+Value *StoreFatPtrsAsIntsAndExpandMemcpyVisitor::intsToFatPtrs(
+    Value *V, Type *From, Type *To, const Twine &Name) {
   if (From == To)
     return V;
   if (isBufferFatPtrOrVector(To)) {
@@ -531,18 +548,25 @@ Value *StoreFatPtrsAsIntsVisitor::intsToFatPtrs(Value *V, Type *From, Type *To,
   return Ret;
 }
 
-bool StoreFatPtrsAsIntsVisitor::processFunction(Function &F) {
+bool StoreFatPtrsAsIntsAndExpandMemcpyVisitor::processFunction(Function &F) {
   bool Changed = false;
-  // The visitors will mutate GEPs and allocas, but will push loads and stores
-  // to the worklist to avoid invalidation.
+  // Process memcpy-like instructions after the main iteration because they can
+  // invalidate iterators.
+  SmallVector<WeakTrackingVH> CanBecomeLoops;
   for (Instruction &I : make_early_inc_range(instructions(F))) {
-    Changed |= visit(I);
+    if (isa<MemTransferInst, MemSetInst, MemSetPatternInst>(I))
+      CanBecomeLoops.push_back(&I);
+    else
+      Changed |= visit(I);
+  }
+  for (WeakTrackingVH VH : make_early_inc_range(CanBecomeLoops)) {
+    Changed |= visit(cast<Instruction>(VH));
   }
   ConvertedForStore.clear();
   return Changed;
 }
 
-bool StoreFatPtrsAsIntsVisitor::visitAllocaInst(AllocaInst &I) {
+bool StoreFatPtrsAsIntsAndExpandMemcpyVisitor::visitAllocaInst(AllocaInst &I) {
   Type *Ty = I.getAllocatedType();
   Type *NewTy = TypeMap->remapType(Ty);
   if (Ty == NewTy)
@@ -551,7 +575,8 @@ bool StoreFatPtrsAsIntsVisitor::visitAllocaInst(AllocaInst &I) {
   return true;
 }
 
-bool StoreFatPtrsAsIntsVisitor::visitGetElementPtrInst(GetElementPtrInst &I) {
+bool StoreFatPtrsAsIntsAndExpandMemcpyVisitor::visitGetElementPtrInst(
+    GetElementPtrInst &I) {
   Type *Ty = I.getSourceElementType();
   Type *NewTy = TypeMap->remapType(Ty);
   if (Ty == NewTy)
@@ -563,7 +588,7 @@ bool StoreFatPtrsAsIntsVisitor::visitGetElementPtrInst(GetElementPtrInst &I) {
   return true;
 }
 
-bool StoreFatPtrsAsIntsVisitor::visitLoadInst(LoadInst &LI) {
+bool StoreFatPtrsAsIntsAndExpandMemcpyVisitor::visitLoadInst(LoadInst &LI) {
   Type *Ty = LI.getType();
   Type *IntTy = TypeMap->remapType(Ty);
   if (Ty == IntTy)
@@ -581,7 +606,7 @@ bool StoreFatPtrsAsIntsVisitor::visitLoadInst(LoadInst &LI) {
   return true;
 }
 
-bool StoreFatPtrsAsIntsVisitor::visitStoreInst(StoreInst &SI) {
+bool StoreFatPtrsAsIntsAndExpandMemcpyVisitor::visitStoreInst(StoreInst &SI) {
   Value *V = SI.getValueOperand();
   Type *Ty = V->getType();
   Type *IntTy = TypeMap->remapType(Ty);
@@ -597,6 +622,47 @@ bool StoreFatPtrsAsIntsVisitor::visitStoreInst(StoreInst &SI) {
   return true;
 }
 
+bool StoreFatPtrsAsIntsAndExpandMemcpyVisitor::visitMemCpyInst(
+    MemCpyInst &MCI) {
+  // TODO: Allow memcpy.p7.p3 as a synonym for the direct-to-LDS copy, which'll
+  // need loop expansion here.
+  if (MCI.getSourceAddressSpace() != AMDGPUAS::BUFFER_FAT_POINTER &&
+      MCI.getDestAddressSpace() != AMDGPUAS::BUFFER_FAT_POINTER)
+    return false;
+  llvm::expandMemCpyAsLoop(&MCI,
+                           TM->getTargetTransformInfo(*MCI.getFunction()));
+  MCI.eraseFromParent();
+  return true;
+}
+
+bool StoreFatPtrsAsIntsAndExpandMemcpyVisitor::visitMemMoveInst(
+    MemMoveInst &MMI) {
+  if (MMI.getSourceAddressSpace() != AMDGPUAS::BUFFER_FAT_POINTER &&
+      MMI.getDestAddressSpace() != AMDGPUAS::BUFFER_FAT_POINTER)
+    return false;
+  report_fatal_error(
+      "memmove() on buffer descriptors is not implemented because pointer "
+      "comparison on buffer descriptors isn't implemented\n");
+}
+
+bool StoreFatPtrsAsIntsAndExpandMemcpyVisitor::visitMemSetInst(
+    MemSetInst &MSI) {
+  if (MSI.getDestAddressSpace() != AMDGPUAS::BUFFER_FAT_POINTER)
+    return false;
+  llvm::expandMemSetAsLoop(&MSI);
+  MSI.eraseFromParent();
+  return true;
+}
+
+bool StoreFatPtrsAsIntsAndExpandMemcpyVisitor::visitMemSetPatternInst(
+    MemSetPatternInst &MSPI) {
+  if (MSPI.getDestAddressSpace() != AMDGPUAS::BUFFER_FAT_POINTER)
+    return false;
+  llvm::expandMemSetPatternAsLoop(&MSPI);
+  MSPI.eraseFromParent();
+  return true;
+}
+
 namespace {
 /// Convert loads/stores of types that the buffer intrinsics can't handle into
 /// one ore more such loads/stores that consist of legal types.
@@ -1127,6 +1193,7 @@ bool LegalizeBufferContentTypesVisitor::visitStoreInst(StoreInst &SI) {
 
 bool LegalizeBufferContentTypesVisitor::processFunction(Function &F) {
   bool Changed = false;
+  // Note, memory transfer intrinsics won't
   for (Instruction &I : make_early_inc_range(instructions(F))) {
     Changed |= visit(I);
   }
@@ -2084,6 +2151,12 @@ static bool isRemovablePointerIntrinsic(Intrinsic::ID IID) {
   case Intrinsic::invariant_end:
   case Intrinsic::launder_invariant_group:
   case Intrinsic::strip_invariant_group:
+  case Intrinsic::memcpy:
+  case Intrinsic::memcpy_inline:
+  case Intrinsic::memmove:
+  case Intrinsic::memset:
+  case Intrinsic::memset_inline:
+  case Intrinsic::experimental_memset_pattern:
     return true;
   }
 }
@@ -2353,7 +2426,8 @@ bool AMDGPULowerBufferFatPointers::run(Module &M, const TargetMachine &TM) {
         /*RemoveDeadConstants=*/false, /*IncludeSelf=*/true);
   }
 
-  StoreFatPtrsAsIntsVisitor MemOpsRewrite(&IntTM, M.getContext());
+  StoreFatPtrsAsIntsAndExpandMemcpyVisitor MemOpsRewrite(&IntTM, M.getContext(),
+                                                         &TM);
   LegalizeBufferContentTypesVisitor BufferContentsTypeRewrite(DL,
                                                               M.getContext());
   for (Function &F : M.functions()) {