[InstCombine] Fix address space handling when removing allocas

[jprice]

Summary

This fixes the bug(s) originally described on llvm-dev in November (with some more discussion in the last couple of days):
http://lists.llvm.org/pipermail/llvm-dev/2016-November/106955.html

As described by @mehdi_amini in his reply, the approach here is to recursively propagate the address space of the global through the allocas uses, instead of introducing a potentially illegal addrspacecast.

Apologies for any obvious issues with either the approach or the style, I'm a first time contributor to LLVM. One thing in particular I'm not sure about is where I had to modify test2 to check for the nonnull attribute as well - I couldn't work out how to correctly retain the original function attributes.

[jprice]

Fix memtransfer function type deduction

[arsenm]

lib/Transforms/InstCombine/InstCombineInternal.h

447

SmallVector

471

SmallVector

472

nullptr

477–480

You should be able to do all of this in the initial constructor

487–488

This can be a C array

504–506

I think you'll also need to handle AtomicRMW and cmpxchg

lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

62

35 is a strange choice for vector size

[arsenm]

lib/Transforms/InstCombine/InstCombineInternal.h

444

This could also be a utility somewhere more general

[jprice]

• Address review comments from arsenm

lib/Transforms/InstCombine/InstCombineInternal.h

444

I did wonder this, but currently this function makes several assumptions that are specific to the particular optimisation that uses it (e.g. lack of handling for certain instructions as per comment below). I think it may need a fair bit more work to make it more general purpose.

504–506

At the moment, the isOnlyCopiedFromConstantGlobal method doesn't allow this optimisation to occur if these instructions are present.

[jprice]

Ping.

I don't know if it's too late, but I'd love for this to make it into the 4.0 release. Although this bug has been present for a while so perhaps isn't strictly a regression, there is a regression in downstream projects that I'm working on due to (valid) changes in the IR generated by Clang since 3.9 which expose this bug.

[hfinkel]

Is this optimization safe in general if the address spaces don't match? Could, for example, some of the users being affected be some target-specific intrinsic that is only valid for certain address spaces?

lib/Transforms/InstCombine/InstCombineInternal.h

502

llvm_unreachable

[jprice]

• Handle different global and function arg address spaces properly
• Use llvm_unreachable instead of assert(false)
• Rebase to current trunk

[jprice]

In D28300#656855, @hfinkel wrote:

Is this optimization safe in general if the address spaces don't match? Could, for example, some of the users being affected be some target-specific intrinsic that is only valid for certain address spaces?

This optimisation will always preserve the original address space of parameters passed to any call instructions, so I don't think this is an issue. I did just notice that this wasn't handled correctly by this patch in certain cases, so I've made sure that this is now the case (and added an appropriate test).

[mehdi_amini]

In D28300#651473, @jprice wrote:

Ping.

I don't know if it's too late, but I'd love for this to make it into the 4.0 release.

I don't believe it is too late at this point.

@arsenm, since you started reviewing, do you mind getting another look at it? (I also know that AMDGPU is a heavy user of address space)

[mehdi_amini]

A few inline comments

lib/Transforms/InstCombine/InstCombineInternal.h

442

Nit: Don't repeat the name of the function in the comment, we use to do this but we changed the doxygen settings some time ago and this is deprecated.

444

Can you turn this into a work-queue iterative loop? I'm worried about unbounded recursion.

445

Why not shortcut here:

If (OldPtr->getType()->getPointerAddressSpace() == NewAddrSpace) {
  replaceInstUsesWith(OldPtr, NewPtr);
  return;
}

448

Add a comment why you need to copy?

449

I haven't figured why it isn't a for-range loop?

457

Shouldn't this be if (NewPtr->getType() == ...?
(If so, you're missing a test to cover this)

479

Is it needed for loads? (Maybe because of typeless pointers?)

499

I'm worried about the validity of this. I would expect passing a pointer to a function to block this transformation when the address space mismatch.

505

Nit: Do not use braces for trivial blocks.

lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

62

Document it.

103

It seems that you will return false if there is any mismatch (line 168).

Do we need a vector? Can't we just keep one AS and return false as soon as we find a mismatch here?

165

Nit: period to end sentence.

314

Spurious change

[jprice]

• Address review comments from mehdi_amini

[jprice]

Thanks for the feedback, responses inline.

lib/Transforms/InstCombine/InstCombineInternal.h

445

This breaks certain cases where OldPtr is the source to an addrspacecast which may result in casting NewPtr to an address space that it shouldn't.

We can however shortcut inside the loop over users, as long as we handle addrspacecast instructions first. I've done this in the latest patch.

457

You're right.

There is a test that covers this case (test11), but since this check just results in a shortcut, it was passing anyway (by introducing an unnecessary NOP cast which was nuked anyway).

499

It does indeed block the transformation, unless we know we can use an addrspacecast because there is already one being used on the global when it is passed as a source to the memcpy. I've added a comment here to make this assumption clear.

lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

103

I'm not sure I see how this would work. There are either one or two valid address spaces that can be used for function arguments (the source and raw source of the memory transfer that we are trying to remove). We don't know whether we will see the memory transfer first, or the function calls that use the alloca.

[jprice]

• Remove accidental debug statement

[jprice]

Ping!

[jprice]

This patch is now obsolete - the core issue has now been fixed by D27283. While that fix is a little more conservative (i.e. it doesn't replace the alloca when it sometimes could), it's good enough to fix all of the issues that I was experiencing in my downstream projects, so I'm happy to leave it at that.

Closing.

lib/Transforms/InstCombine/InstCombineInternal.h

Instruction *replaceInstUsesWith(Instruction &I, Value *V) {
 
     DEBUG(dbgs() << "IC: Replacing " << I << "\n"
                  << "    with " << *V << '\n');
 
     I.replaceAllUsesWith(V);
     return &I;
   }
 
+  /// Recursively update the users of a pointer with a new value, propagating
+  /// the new address space to all users.
+  void replacePtrUsesAndPropagateAddrSpace(Value *OldPtr, Value *NewPtr) {
+    SmallVector<std::pair<Value *, Value *>, 4> WorkList;
+    WorkList.push_back({OldPtr, NewPtr});
+    do {
+      OldPtr = WorkList.back().first;
+      NewPtr = WorkList.back().second;
+      WorkList.pop_back();
+
+      unsigned NewAddrSpace = NewPtr->getType()->getPointerAddressSpace();
+
+      // Capture OldPtr users as we will be modifying the user list.
+      SmallVector<User *, 8> Users(OldPtr->users().begin(),
+                                   OldPtr->users().end());
+      for (auto U : Users) {
+        if (auto ASCI = dyn_cast<AddrSpaceCastInst>(U)) {
+          if (NewPtr->getType() == ASCI->getType()) {
+            // If the type already matches, remove the cast completely.
+            ASCI->replaceAllUsesWith(NewPtr);
+          } else {
+            // Create and insert new bitcast instruction with new address space.
+            auto NewBI = new BitCastInst(
+                NewPtr, ASCI->getType()->getPointerElementType()->getPointerTo(
+                            NewAddrSpace),
+                "", ASCI);
+            WorkList.push_back({ASCI, NewBI});
+          }
+          continue;
+        }
+
+        // If the address space already matches or the user is a load
+        // instruction, then just replace the use.
+        if (OldPtr->getType()->getPointerAddressSpace() == NewAddrSpace ||
+            isa<LoadInst>(U)) {
+          U->replaceUsesOfWith(OldPtr, NewPtr);
+          continue;
+        }
+
+        if (auto BI = dyn_cast<BitCastInst>(U)) {
+          // Create and insert new bitcast instruction
+          Type *DestType =
+              BI->getDestTy()->getPointerElementType()->getPointerTo(
+                  NewAddrSpace);
+          auto NewBI = new BitCastInst(NewPtr, DestType, "", BI);
+          WorkList.push_back({BI, NewBI});
+        } else if (auto GEP = dyn_cast<GetElementPtrInst>(U)) {
+          // Create and insert new getelementptr instruction
+          SmallVector<Value *, 4> Indices(GEP->idx_begin(), GEP->idx_end());
+          auto NewGEP =
+              GetElementPtrInst::Create(nullptr, NewPtr, Indices, "", GEP);
+          WorkList.push_back({GEP, NewGEP});
+        } else if (auto MI = dyn_cast<MemTransferInst>(U)) {
+          assert(MI->getArgOperand(1) == OldPtr);
+
+          // Get intrinsic with new address space
+          auto MIType = MI->getFunctionType();
+          Type *Types[3] = {MIType->getParamType(0), NewPtr->getType(),
+                            MIType->getParamType(2)};
+          auto NewMIFunc = Intrinsic::getDeclaration(
+              MI->getModule(), MI->getIntrinsicID(), Types);
+          NewMIFunc->setAttributes(MI->getCalledFunction()->getAttributes());
+
+          // Create and insert new call instruction
+          std::vector<Value *> Args(MI->arg_begin(), MI->arg_end());
+          Args[1] = NewPtr;
+          CallInst::Create(NewMIFunc, Args, "", MI);
+          eraseInstFromFunction(*MI);
+        } else if (auto CI = dyn_cast<CallInst>(U)) {
+
+          // Cast the source to the correct address space.
+          // We assume that this transformation would be blocked if this cast
+          // wasn't valid.
+          Value *Src;
+          if (auto ConstSrc = dyn_cast<Constant>(NewPtr))
+            Src = ConstantExpr::getAddrSpaceCast(ConstSrc, OldPtr->getType());
+          else
+            Src = new AddrSpaceCastInst(NewPtr, OldPtr->getType(), "", CI);
+          U->replaceUsesOfWith(OldPtr, Src);
+        } else {
+          llvm_unreachable("Can't update address space for this instruction");
+        }
+      }
+    } while (!WorkList.empty());
+  }
+
   /// Creates a result tuple for an overflow intrinsic \p II with a given
   /// \p Result and a constant \p Overflow value.
   Instruction *CreateOverflowTuple(IntrinsicInst *II, Value *Result,
                                    Constant *Overflow) {
     Constant *V[] = {UndefValue::get(Result->getType()), Overflow};
     StructType *ST = cast<StructType>(II->getType());
     Constant *Struct = ConstantStruct::get(ST, V);
     return InsertValueInst::Create(Struct, Result, 0);

lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

 /// can optimize this.
 static bool
 isOnlyCopiedFromConstantGlobal(Value *V, MemTransferInst *&TheCopy,
                                SmallVectorImpl<Instruction *> &ToDelete) {
   // We track lifetime intrinsics as we encounter them.  If we decide to go
   // ahead and replace the value with the global, this lets the caller quickly
   // eliminate the markers.
 
+  // Track the address spaces of any function arguments using any derived
+  // pointer to the alloca.
+  SmallDenseSet<unsigned, 4> ArgAddrSpaces;
+
   SmallVector<std::pair<Value *, bool>, 35> ValuesToInspect;
   ValuesToInspect.emplace_back(V, false);
   while (!ValuesToInspect.empty()) {
     auto ValuePair = ValuesToInspect.pop_back_val();
     const bool IsOffset = ValuePair.second;
     for (auto &U : ValuePair.first->uses()) {
       auto *I = cast<Instruction>(U.getUser());
 
for (auto &U : ValuePair.first->uses()) {
         // If this is the function being called then we treat it like a load and
         // ignore it.
         if (CS.isCallee(&U))
           continue;
 
         unsigned DataOpNo = CS.getDataOperandNo(&U);
         bool IsArgOperand = CS.isArgOperand(&U);
 
+        // Track the address spaces of function call arguments
+        if (IsArgOperand && !isa<MemTransferInst>(I))
+          ArgAddrSpaces.insert(U->getType()->getPointerAddressSpace());
+
         // Inalloca arguments are clobbered by the call.
         if (IsArgOperand && CS.isInAllocaArgument(DataOpNo))
           return false;
 
         // If this is a readonly/readnone call site, then we know it is just a
         // load (but one that potentially returns the value itself), so we can
         // ignore it if we know that the value isn't captured.
         if (CS.onlyReadsMemory() &&
for (auto &U : ValuePair.first->uses()) {
       // If the source of the memcpy/move is not a constant global, reject it.
       if (!pointsToConstantGlobal(MI->getSource()))
         return false;
 
       // Otherwise, the transform is safe.  Remember the copy instruction.
       TheCopy = MI;
     }
   }
+
+  // Check that the address spaces of all function call arguments match the
+  // source of the memory transfer.
+  if (TheCopy) {
+    for (unsigned AS : ArgAddrSpaces) {
+      if (AS != TheCopy->getSource()->getType()->getPointerAddressSpace() &&
+          AS != TheCopy->getRawSource()->getType()->getPointerAddressSpace())
+        return false;
+    }
+  }
+
   return true;
 }
 
 /// isOnlyCopiedFromConstantGlobal - Return true if the specified alloca is only
 /// modified by a copy from a constant global.  If we can prove this, we can
 /// replace any uses of the alloca with uses of the global directly.
 static MemTransferInst *
 isOnlyCopiedFromConstantGlobal(AllocaInst *AI,
if (AI.getAlignment()) {
     if (MemTransferInst *Copy = isOnlyCopiedFromConstantGlobal(&AI, ToDelete)) {
       unsigned SourceAlign = getOrEnforceKnownAlignment(
           Copy->getSource(), AI.getAlignment(), DL, &AI, &AC, &DT);
       if (AI.getAlignment() <= SourceAlign) {
         DEBUG(dbgs() << "Found alloca equal to global: " << AI << '\n');
         DEBUG(dbgs() << "  memcpy = " << *Copy << '\n');
         for (unsigned i = 0, e = ToDelete.size(); i != e; ++i)
           eraseInstFromFunction(*ToDelete[i]);
         Constant *TheSrc = cast<Constant>(Copy->getSource());
-        Constant *Cast
-          = ConstantExpr::getPointerBitCastOrAddrSpaceCast(TheSrc, AI.getType());
-        Instruction *NewI = replaceInstUsesWith(AI, Cast);
+
         eraseInstFromFunction(*Copy);
         ++NumGlobalCopies;
-        return NewI;
+
+        // Replace the alloca uses with the global (casting as necessary), and
+        // propagate the new address space through all of the alloca uses.
+        Constant *Cast = ConstantExpr::getPointerCast(
+            TheSrc,
+            AI.getType()->getPointerElementType()->getPointerTo(
+                TheSrc->getType()->getPointerAddressSpace()));
+        replacePtrUsesAndPropagateAddrSpace(&AI, Cast);
+
+        return &AI;
       }
     }
   }
 
   // At last, use the generic allocation site handler to aggressively remove
   // unused allocas.
   return visitAllocSite(AI);
 }

test/Transforms/InstCombine/memcpy-from-global.ll

 declare void @llvm.memcpy.p0i8.p1i8.i64(i8* nocapture, i8 addrspace(1)* nocapture, i64, i32, i1) nounwind
 declare void @llvm.memcpy.p1i8.p1i8.i64(i8 addrspace(1)* nocapture, i8 addrspace(1)* nocapture, i64, i32, i1) nounwind
 
 %T = type { i8, [123 x i8] }
 %U = type { i32, i32, i32, i32, i32 }
 
 @G = constant %T {i8 1, [123 x i8] zeroinitializer }
 @H = constant [2 x %U] zeroinitializer, align 16
+@I = addrspace(1) constant [2 x %U] zeroinitializer, align 16
 
 define void @test2() {
   %A = alloca %T
   %B = alloca %T
   %a = bitcast %T* %A to i8*
   %b = bitcast %T* %B to i8*
 
 ; CHECK-LABEL: @test2(

 ; %A alloca is deleted
 ; This doesn't exactly match what test2 does, because folding the type
 ; cast into the alloca doesn't work for the addrspacecast yet.
 ; CHECK-NEXT: alloca [124 x i8]
 ; CHECK-NEXT: getelementptr
 ; CHECK-NEXT: addrspacecast
 
 ; use @G instead of %A
-; CHECK-NEXT: call void @llvm.memcpy.p1i8.p1i8.i64(i8 addrspace(1)* %{{.*}},
+; CHECK-NEXT: call void @llvm.memcpy.p1i8.p0i8.i64(i8 addrspace(1)* %{{.*}},
   call void @llvm.memcpy.p1i8.p0i8.i64(i8 addrspace(1)* %a, i8* bitcast (%T* @G to i8*), i64 124, i32 4, i1 false)
   call void @llvm.memcpy.p1i8.p1i8.i64(i8 addrspace(1)* %b, i8 addrspace(1)* %a, i64 124, i32 4, i1 false)
   call void @bar_as1(i8 addrspace(1)* %b)
   ret void
 }
 
 declare void @bar(i8*)
 declare void @bar_as1(i8 addrspace(1)*)
define void @test9_addrspacecast() {
   %A = alloca %U, align 4
   %a = bitcast %U* %A to i8*
   call void @llvm.memcpy.p0i8.p1i8.i64(i8* %a, i8 addrspace(1)* addrspacecast (%U* getelementptr ([2 x %U], [2 x %U]* @H, i64 0, i32 1) to i8 addrspace(1)*), i64 20, i32 4, i1 false)
   call void @bar(i8* %a) readonly
 ; CHECK-LABEL: @test9_addrspacecast(
 ; CHECK-NEXT: call void @bar(i8* bitcast (%U* getelementptr inbounds ([2 x %U], [2 x %U]* @H, i64 0, i64 1) to i8*))
   ret void
 }
+
+define void @test10(i64 %index, i32* %out) {
+  %A = alloca %U, align 16
+  %a = bitcast %U* %A to i8*
+  call void @llvm.memcpy.p0i8.p1i8.i64(i8* %a, i8 addrspace(1)* bitcast ([2 x %U] addrspace(1)* @I to i8 addrspace(1)*), i64 20, i32 16, i1 false)
+  %tmp1 = bitcast i8* %a to i32*
+  %tmp2 = getelementptr i32, i32* %tmp1, i64 %index
+  %tmp3 = load i32, i32* %tmp2
+  store i32 %tmp3, i32* %out
+; CHECK-LABEL: @test10(
+; CHECK-NOT: alloca
+; CHECK-NOT: memcpy
+; CHECK-NOT: addrspacecast
+; Ensure that load still happens through addrspace(1)
+; CHECK: load i32, i32 addrspace(1)*
+  ret void
+}
+
+define void @test10_addrspacecast(i64 %index, i32* %out) {
+  %A = alloca %U, align 16
+  %a = bitcast %U* %A to i8*
+  call void @llvm.memcpy.p0i8.p1i8.i64(i8* %a, i8 addrspace(1)* bitcast ([2 x %U] addrspace(1)* @I to i8 addrspace(1)*), i64 20, i32 16, i1 false)
+  %tmp1 = addrspacecast i8* %a to i32 addrspace(2)*
+  %tmp2 = getelementptr i32, i32 addrspace(2)* %tmp1, i64 %index
+  %tmp3 = load i32, i32 addrspace(2)* %tmp2
+  store i32 %tmp3, i32* %out
+; CHECK-LABEL: @test10_addrspacecast(
+; CHECK-NOT: alloca
+; CHECK-NOT: memcpy
+; CHECK-NOT: addrspacecast
+; Similar to above but the load from the alloca is done through a different
+; address space.
+; Ensure that load still happens through addrspace(1)
+; CHECK: load i32, i32 addrspace(1)*
+  ret void
+}
+
+define void @test11() {
+  %A = alloca %U, align 16
+  %a = bitcast %U* %A to i8*
+  call void @llvm.memcpy.p0i8.p1i8.i64(i8* %a, i8 addrspace(1)* bitcast ([2 x %U] addrspace(1)* @I to i8 addrspace(1)*), i64 20, i32 16, i1 false)
+  %tmp = addrspacecast i8* %a to i8 addrspace(1)*
+  call void @bar_as1(i8 addrspace(1)* %tmp) readonly
+; CHECK-LABEL: @test11(
+; CHECK-NEXT: call void @bar_as1(i8 addrspace(1)* bitcast ([2 x %U] addrspace(1)* @I to i8 addrspace(1)*))
+  ret void
+}
+
+define void @test12() {
+  %A = alloca %U, align 16
+  %a = bitcast %U* %A to i8*
+  call void @llvm.memcpy.p0i8.p1i8.i64(i8* %a, i8 addrspace(1)* bitcast ([2 x %U] addrspace(1)* @I to i8 addrspace(1)*), i64 20, i32 16, i1 false)
+  call void @bar(i8* %a) readonly
+; Must retain memcpy as source address space doesn't match function argument.
+; CHECK-LABEL: @test12(
+; CHECK: llvm.memcpy
+; CHECK: bar
+  ret void
+}
+
+define void @test13() {
+  %A = alloca %U, align 16
+  %a = bitcast %U* %A to i8*
+  call void @llvm.memcpy.p0i8.p1i8.i64(i8* %a, i8 addrspace(1)* bitcast ([2 x %U] addrspace(1)* @I to i8 addrspace(1)*), i64 20, i32 16, i1 false)
+  %tmp = addrspacecast i8* %a to i8 addrspace(1)*
+  call void @bar_as1(i8 addrspace(1)* %tmp) readonly
+  call void @bar(i8* %a) readonly
+; Must retain memcpy as source address space doesn't match function argument.
+; Similar to test12, but ensures that we can handle multiple function calls
+; with different address spaces.
+; CHECK-LABEL: @test13(
+; CHECK: llvm.memcpy
+; CHECK: bar_as1
+; CHECK: bar
+  ret void
+}
+
+define void @test14() {
+  %A = alloca %U, align 16
+  %a = bitcast %U* %A to i8*
+  call void @llvm.memcpy.p0i8.p0i8.i64(i8* %a, i8* addrspacecast (i8 addrspace(1)* bitcast ([2 x %U] addrspace(1)* @I to i8 addrspace(1)*) to i8*), i64 20, i32 16, i1 false)
+  call void @bar(i8* %a) readonly
+; Similar to test12, but there is a pre-existing addresspacecast which we can
+; use to pass the global to the function.
+; CHECK-LABEL: @test14(
+; CHECK-NEXT: call void @bar(i8* addrspacecast (i8 addrspace(1)* bitcast ([2 x %U] addrspace(1)* @I to i8 addrspace(1)*) to i8*))
+  ret void
+}