[AMDGPU] Fix promote alloca with double use in a same insn

[rampitec]

Summary

If we have an instruction where more than one pointer operands
are derived from the same promoted alloca, we are fixing it for
one argument and do not fix a second use considering this user
done.

Fix this by deferring processing of memory intrinsics until all
potential operands are replaced.

Fixes: SWDEV-271358

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

603–605

I don't see how this sorts it, the ordering is still determined by the arbitrary ordering in users. Can't the replacement just check all candidate operands

[rampitec]

llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

603–605

The function is recursive and does DFS. This code pushes any later use to the end.

Checking all uses at replacement is practically impossible because it may be a long use-def chain. You would need to do DFS again at every replacement.

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

603–605

This is a DFS on the user list, not the function. I do not think this will give you the guarantee that the first user dominates the second. But you also just need to ensure the transitive users are enqueued regardless. Would it be easier to just use a SetVector instead of is_contained and a regular worklist

llvm/test/CodeGen/AMDGPU/promote-alloca-mem-intrinsics.ll

69

I noticed this dropped dereferenceable, should maybe fix that

80

Can you also add tests with select and phi both derived from the same

[rampitec]

I actually think this creates topological problem. If I move a use forward I also need to move the whole chain after it forward. So it doesn't seem to work unfortunately. Uses create trees and somewhere in these trees we might need a sort of barriers. Essentially all use chains have to come to an instruction useing them.

[rampitec]

llvm/test/CodeGen/AMDGPU/promote-alloca-mem-intrinsics.ll

80

We have these tests. One in promote-alloca-to-lds-select.ll @lds_promote_alloca_select_two_derived_pointers, one in the promote-alloca-to-lds-phi.ll @branch_ptr_var_same_alloca. This works because we only update operands not replacing the instruction. Problem with the memcpy is that we actually create a new call which is not going to be hit when we touch second operand.

One possible solution I am exploring is to postpone patching memory intrinsics until the end. Another is to patch the call in place.

[rampitec]

Changed to just defer processing memcpy and memove.
Restored dereferenceable attributes.

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

1004

In principle this could happen for all multi-operand instructions, like select and icmp

1007

I'm not sure I understand why it really needs to be deferred. If we tracked the specific use, you could just replace the one operand and then encounter the instruction again for the second?

1014–1016

This is a separate patch (also more attributes still apply)

[rampitec]

llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

1004

We have tests for this. It works fine. Operands are just updated with replaceAllUsesWith(). With memcpy and memmove there are two issues:

1. We need to replace called function with different mangling. We are only replacing it for one of the arguments. Second argument is also updated but we do not visit this intrinsic again.
2. Even if we have these calls in the worklist twice that still would not work. Other instructions just get their operands updated, but here we are replacing call instruction itself with a new one. No pointer in the worklist would point to it.

1007

No, because instruction is dropped.

[rampitec]

Dropped memset attribute parts.

[rampitec]

Here is what happens w/o the patch:

call void @llvm.memcpy.p0i8.p3i8.i64(i8 addrspace(3)* align 8 %i, i8 addrspace(3)* align 8 %i1, i64 16, i1 false)

As you may see both operands are correctly updated. That is mangling of the memcpy is wrong (p0i8 for the first operand).

llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

bool AMDGPUPromoteAllocaImpl::binaryOpIsDerivedFromSameAlloca(
 
   return true;
 }
 
 bool AMDGPUPromoteAllocaImpl::collectUsesWithPtrTypes(
     Value *BaseAlloca, Value *Val, std::vector<Value *> &WorkList) const {
 
   for (User *User : Val->users()) {
     if (is_contained(WorkList, User))
       continue;
 
     if (CallInst *CI = dyn_cast<CallInst>(User)) {
       if (!isCallPromotable(CI))
         return false;
 
       WorkList.push_back(User);
       continue;
     }
 
Value *Indices[] = {
     TID
   };
 
   Value *Offset = Builder.CreateInBoundsGEP(GVTy, GV, Indices);
   I.mutateType(Offset->getType());
   I.replaceAllUsesWith(Offset);
   I.eraseFromParent();
 
+  SmallVector<IntrinsicInst *> DeferredIntrs;
+
   for (Value *V : WorkList) {
     CallInst *Call = dyn_cast<CallInst>(V);
     if (!Call) {
       if (ICmpInst *CI = dyn_cast<ICmpInst>(V)) {
         Value *Src0 = CI->getOperand(0);
         Type *EltTy = Src0->getType()->getPointerElementType();
         PointerType *NewTy = PointerType::get(EltTy, AMDGPUAS::LOCAL_ADDRESS);
 
for (Value *V : WorkList) {
     IntrinsicInst *Intr = cast<IntrinsicInst>(Call);
     Builder.SetInsertPoint(Intr);
     switch (Intr->getIntrinsicID()) {
     case Intrinsic::lifetime_start:
     case Intrinsic::lifetime_end:
       // These intrinsics are for address space 0 only
       Intr->eraseFromParent();
       continue;
-    case Intrinsic::memcpy: {
-      MemCpyInst *MemCpy = cast<MemCpyInst>(Intr);
-      Builder.CreateMemCpy(MemCpy->getRawDest(), MemCpy->getDestAlign(),
-                           MemCpy->getRawSource(), MemCpy->getSourceAlign(),
-                           MemCpy->getLength(), MemCpy->isVolatile());
-      Intr->eraseFromParent();
-      continue;
-    }
-    case Intrinsic::memmove: {
-      MemMoveInst *MemMove = cast<MemMoveInst>(Intr);
-      Builder.CreateMemMove(MemMove->getRawDest(), MemMove->getDestAlign(),
-                            MemMove->getRawSource(), MemMove->getSourceAlign(),
-                            MemMove->getLength(), MemMove->isVolatile());
-      Intr->eraseFromParent();
+    case Intrinsic::memcpy:
+    case Intrinsic::memmove:
+      // These have 2 pointer operands. In case if second pointer also needs
+      // to be replaced we defer processing of these intrinsics until all
+      // other values are processed.
+      DeferredIntrs.push_back(Intr);
       continue;
-    }
     case Intrinsic::memset: {
       MemSetInst *MemSet = cast<MemSetInst>(Intr);
       Builder.CreateMemSet(
           MemSet->getRawDest(), MemSet->getValue(), MemSet->getLength(),
           MaybeAlign(MemSet->getDestAlignment()), MemSet->isVolatile());
       Intr->eraseFromParent();
       continue;
     }
     case Intrinsic::invariant_start:
     case Intrinsic::invariant_end:
     case Intrinsic::launder_invariant_group:
     case Intrinsic::strip_invariant_group:
       Intr->eraseFromParent();
       // FIXME: I think the invariant marker should still theoretically apply,
       // but the intrinsics need to be changed to accept pointers with any
       // address space.
case Intrinsic::objectsize: {
       Intr->eraseFromParent();
       continue;
     }
     default:
       Intr->print(errs());
       llvm_unreachable("Don't know how to promote alloca intrinsic use.");
     }
   }
+
+  for (IntrinsicInst *Intr : DeferredIntrs) {
+    Builder.SetInsertPoint(Intr);
+    Intrinsic::ID ID = Intr->getIntrinsicID();
+    assert(ID == Intrinsic::memcpy || ID == Intrinsic::memmove);
+
+    MemTransferInst *MI = cast<MemTransferInst>(Intr);
+    auto *B =
+      Builder.CreateMemTransferInst(ID, MI->getRawDest(), MI->getDestAlign(),
+                                    MI->getRawSource(), MI->getSourceAlign(),
+                                    MI->getLength(), MI->isVolatile());
+
+    for (unsigned I = 1; I != 3; ++I) {
+      if (uint64_t Bytes = Intr->getDereferenceableBytes(I)) {
+        B->addDereferenceableAttr(I, Bytes);
+      }
+    }
+
+    Intr->eraseFromParent();
+  }
+
   return true;
 }
 
 bool handlePromoteAllocaToVector(AllocaInst &I, unsigned MaxVGPRs) {
   // Array allocations are probably not worth handling, since an allocation of
   // the array type is the canonical form.
   if (!I.isStaticAlloca() || I.isArrayAllocation())
     return false;

llvm/test/CodeGen/AMDGPU/promote-alloca-mem-intrinsics.ll

 ; RUN: opt -S -mtriple=amdgcn-unknown-amdhsa -mcpu=kaveri -amdgpu-promote-alloca < %s | FileCheck %s
 
 declare void @llvm.memcpy.p0i8.p1i8.i32(i8* nocapture, i8 addrspace(1)* nocapture, i32, i1) #0
 declare void @llvm.memcpy.p1i8.p0i8.i32(i8 addrspace(1)* nocapture, i8* nocapture, i32, i1) #0
+declare void @llvm.memcpy.p0i8.p0i8.i64(i8* nocapture, i8* nocapture, i64, i1) #0
 
 declare void @llvm.memmove.p0i8.p1i8.i32(i8* nocapture, i8 addrspace(1)* nocapture, i32, i1) #0
 declare void @llvm.memmove.p1i8.p0i8.i32(i8 addrspace(1)* nocapture, i8* nocapture, i32, i1) #0
+declare void @llvm.memmove.p0i8.p0i8.i64(i8* nocapture, i8* nocapture, i64, i1) #0
 
 declare void @llvm.memset.p0i8.i32(i8* nocapture, i8, i32, i1) #0
 
 declare i32 @llvm.objectsize.i32.p0i8(i8*, i1, i1, i1) #1
 
 ; CHECK-LABEL: @promote_with_memcpy(
 ; CHECK: getelementptr inbounds [64 x [17 x i32]], [64 x [17 x i32]] addrspace(3)* @promote_with_memcpy.alloca, i32 0, i32 %{{[0-9]+}}
 ; CHECK: call void @llvm.memcpy.p3i8.p1i8.i32(i8 addrspace(3)* align 4 %alloca.bc, i8 addrspace(1)* align 4 %in.bc, i32 68, i1 false)

 define amdgpu_kernel void @promote_with_objectsize(i32 addrspace(1)* %out) #0 {
   %alloca = alloca [17 x i32], align 4
   %alloca.bc = bitcast [17 x i32]* %alloca to i8*
   %size = call i32 @llvm.objectsize.i32.p0i8(i8* %alloca.bc, i1 false, i1 false, i1 false)
   store i32 %size, i32 addrspace(1)* %out
   ret void
 }
 
+; CHECK-LABEL: @promote_alloca_used_twice_in_memcpy(
+; CHECK: %i = bitcast double addrspace(3)* %arrayidx1 to i8 addrspace(3)*
+; CHECK: %i1 = bitcast double addrspace(3)* %arrayidx2 to i8 addrspace(3)*
+; CHECK: call void @llvm.memcpy.p3i8.p3i8.i64(i8 addrspace(3)* align 8 dereferenceable(16) %i, i8 addrspace(3)* align 8 dereferenceable(16) %i1, i64 16, i1 false)
+define amdgpu_kernel void @promote_alloca_used_twice_in_memcpy(i32 %c) {
+entry:
+  %r = alloca double, align 8
+  %arrayidx1 = getelementptr inbounds double, double* %r, i32 1
+  %i = bitcast double* %arrayidx1 to i8*
+  %arrayidx2 = getelementptr inbounds double, double* %r, i32 %c
+  %i1 = bitcast double* %arrayidx2 to i8*
+  call void @llvm.memcpy.p0i8.p0i8.i64(i8* align 8 dereferenceable(16) %i, i8* align 8 dereferenceable(16) %i1, i64 16, i1 false)
+  ret void
+}
+
+; CHECK-LABEL: @promote_alloca_used_twice_in_memmove(
+; CHECK: %i = bitcast double addrspace(3)* %arrayidx1 to i8 addrspace(3)*
+; CHECK: %i1 = bitcast double addrspace(3)* %arrayidx2 to i8 addrspace(3)*
+; CHECK: call void @llvm.memmove.p3i8.p3i8.i64(i8 addrspace(3)* align 8 dereferenceable(16) %i, i8 addrspace(3)* align 8 dereferenceable(16) %i1, i64 16, i1 false)
+define amdgpu_kernel void @promote_alloca_used_twice_in_memmove(i32 %c) {
+entry:
+  %r = alloca double, align 8
+  %arrayidx1 = getelementptr inbounds double, double* %r, i32 1
+  %i = bitcast double* %arrayidx1 to i8*
+  %arrayidx2 = getelementptr inbounds double, double* %r, i32 %c
+  %i1 = bitcast double* %arrayidx2 to i8*
+  call void @llvm.memmove.p0i8.p0i8.i64(i8* align 8 dereferenceable(16) %i, i8* align 8 dereferenceable(16) %i1, i64 16, i1 false)
+  ret void
+}
+
 attributes #0 = { nounwind "amdgpu-flat-work-group-size"="64,64" "amdgpu-waves-per-eu"="1,3" }
 attributes #1 = { nounwind readnone }