[InstCombine] Teach alloca replacement to handle `addrspacecast`

[hliao]

Summary

• As the address space cast may not be valid on a specific target, addrspacecast is not handled when an alloca is able to be replaced with the source of memcpy/memmove. This patch addresses that by querying a target hook on whether that address space cast is valid. For example, on most GPU targets, the cast from a global pointer to a generic pointer is valid.
• If that cast is allowedd (by querying isValidAddrSpaceCast), the replacement is enhanced to handle that addrspacecast as well.

[yaxunl]

llvm/test/Transforms/InstCombine/ptr-replace-alloca.ll

438

if half of %alloca is filled with @g1, and half of %alloca is filled with @g2, will it be replaced in the call of @readonly_caller? do we need a test for that?

[hliao]

Add one more keep test

[hliao]

a new keep test is added.

llvm/test/Transforms/InstCombine/ptr-replace-alloca.ll

438

it won't replace that alloca as isOnlyCopiedFromConstantMemory bails out if more than one copy to alloca (including pointers calculated from GEPs) is found. Even if we could find that's the safe case as the copies from g1 and g2 won't overlap, that initialization of the alloca is way complicated to be handled in this scenario (passing through a function call). For other scenarios, I hope SROA could slice that [32 x i8] alloca into 2 smaller ones. But, it still won't be handled here directly.

[hliao]

Fix typos

[hliao]

Ping for review

[yaxunl]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

265

can we make SrcAS a member of PointerReplacer since it seems not to change?

[hliao]

Revise following reviewer's comment

[nikic]

As far as I can tell, your motivating test cases do not require the TTI hook. They only need trivial address space casts (to the same address space).

Using TTI inside InstCombine is forbidden as a matter of policy. You would need to make an argument for why this exception should be allowed (it would probably go along the lines of "this should really be part of data layout"), so I would recommend to only handle this trivial case first and bypass the TTI issue.

[hliao]

In D147025#4233724, @nikic wrote:

As far as I can tell, your motivating test cases do not require the TTI hook. They only need trivial address space casts (to the same address space).

Using TTI inside InstCombine is forbidden as a matter of policy. You would need to make an argument for why this exception should be allowed (it would probably go along the lines of "this should really be part of data layout"), so I would recommend to only handle this trivial case first and bypass the TTI issue.

Yes, that specific test doesn't need a TTI hook to enhance the instcombine. But, that TTI hook is provided to handle other valid address spaces per target. For example, on most GPU devices, the cast from a global pointer to a generic pointer is valid but not in the reverse direction. However, that needs to change the target code and target-specific test. I do not own any target code to test that hook further. If required, I could modify AMDGPU as an example.

[hliao]

Add the target-specific usage

[hliao]

Kindly PING for review

[hliao]

Fix clang-format

[hliao]

Rebase and kindling PING for review

[hliao]

Rebase and kindly PING for review

[yaxunl]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

278

Do we really need to check whether the addrspacecast is valid here?

If it is invalid, then it is UB. Therefore we could just assume it is valid.

According to the LLVM IR manual, we can assume the memory pointed to by addrspacecasted pointers are the same, therefore we should be able to do this optimization always.

Another reason is that there should already be validity check for addrspacecast in some target specific passes, therefore there is no need to check it here.

[arsenm]

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

99

The langref/verifier do not allow addrspacecast between the same address space

llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.h

173–179 ↗

(On Diff #510350)

This misses a variety of valid casts. More or less anything except region should be a valid to flat case. Also the constant_32bit can cast to constant and to flat.

[arsenm]

This is just reinterpreting bits of a pointer in memory. This is not an addrspacecast. You can only use an addrspacecast in the trivial cases, otherwise you need to use ptrtoint+inttoptr. This is also why we should allow bitcasts between any pointers with the same size https://discourse.llvm.org/t/rfc-llvm-ir-should-allow-bitcast-between-address-spaces-with-the-same-size/5759

[hliao]

• Update TTI hook isValidAddrSpaceCast to be consistent with LLVM IR verifier
• Check the equal case in instcombine
• Add other valid cases in AMDGPU

[hliao]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

278

We should not transform valid code with well-defined behavior into one with UB. It's valid for the developer to copy data from an address space (not eligible to be cast into generic address space) into private memory and pass that data into a routine accepting generic pointers only. For this kind of case, we should not transform that code or generate a UB address space cast.

[hliao]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

278

For the last comment, once the alloca is replaced or removed, a new addrspacecast would be created from the source pointer (copied into alloca) to the pointer (previously cast from that alloca pointer). We need to ensure the resulting addrspacecast is valid.

[yaxunl]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

278

I c. I thought only operands of load is replaced. I did not notice that other uses of the pointer are replaced.

Probably need a test for that.

[hliao]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

278

This patch only handles the extra addrspacecast user. It doesn't introduce new users.

[hliao]

Rebase and kindly PING review. Please let me know any pending issue to be addressed.

[hliao]

Rebase and kindly PING for review

[nikic]

I'm approving this usage of TTI in InstCombine on the following grounds:

• This is a legality query, not a profitability query. It does not impact IR canonicality in any meaningful way.
• While legality queries should really not be part of TTI, I also think that including this in DataLayout would be somewhat awkward, and making this a TTI query is the more pragmatic choice for the time being. I'd still be interested in @arsenm's opinion on whether it would make sense to make this a data layout property in the future.

I'll leave final patch review to AMDGPU maintainers, just explicitly withdrawing my objection here.

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

352

Please perform the cast to AddrSpaceCastInst here, not inside the function.

547

Use SrcAddrSpace defined above.

[hliao]

Revise following comments from reviewers

[hliao]

Rebase and kindly PING for review

[hliao]

Rebase and kindly PING for review

[yaxunl]

LGTM. Thanks.

llvm/include/llvm/Analysis/TargetTransformInfo.h

public:
   /// algorithm starting with the sources of divergence.
   bool isSourceOfDivergence(const Value *V) const;
 
   // Returns true for the target specific
   // set of operations which produce uniform result
   // even taking non-uniform arguments
   bool isAlwaysUniform(const Value *V) const;
 
+  /// Query the target whether the specified address space cast from FromAS to
+  /// ToAS is valid.
+  bool isValidAddrSpaceCast(unsigned FromAS, unsigned ToAS) const;
+
   /// Returns the address space ID for a target's 'flat' address space. Note
   /// this is not necessarily the same as addrspace(0), which LLVM sometimes
   /// refers to as the generic address space. The flat address space is a
   /// generic address space that can be used access multiple segments of memory
   /// with different address spaces. Access of a memory location through a
   /// pointer with this address space is expected to be legal but slower
   /// compared to the same memory location accessed through a pointer with a
   /// different address space.
public:
   virtual InstructionCost getInstructionCost(const User *U,
                                              ArrayRef<const Value *> Operands,
                                              TargetCostKind CostKind) = 0;
   virtual BranchProbability getPredictableBranchThreshold() = 0;
   virtual bool hasBranchDivergence() = 0;
   virtual bool useGPUDivergenceAnalysis() = 0;
   virtual bool isSourceOfDivergence(const Value *V) = 0;
   virtual bool isAlwaysUniform(const Value *V) = 0;
+  virtual bool isValidAddrSpaceCast(unsigned FromAS, unsigned ToAS) const = 0;
   virtual unsigned getFlatAddressSpace() = 0;
   virtual bool collectFlatAddressOperands(SmallVectorImpl<int> &OpIndexes,
                                           Intrinsic::ID IID) const = 0;
   virtual bool isNoopAddrSpaceCast(unsigned FromAS, unsigned ToAS) const = 0;
   virtual bool
   canHaveNonUndefGlobalInitializerInAddressSpace(unsigned AS) const = 0;
   virtual unsigned getAssumedAddrSpace(const Value *V) const = 0;
   virtual bool isSingleThreaded() const = 0;
public:
   bool isSourceOfDivergence(const Value *V) override {
     return Impl.isSourceOfDivergence(V);
   }
 
   bool isAlwaysUniform(const Value *V) override {
     return Impl.isAlwaysUniform(V);
   }
 
+  bool isValidAddrSpaceCast(unsigned FromAS, unsigned ToAS) const override {
+    return Impl.isValidAddrSpaceCast(FromAS, ToAS);
+  }
+
   unsigned getFlatAddressSpace() override { return Impl.getFlatAddressSpace(); }
 
   bool collectFlatAddressOperands(SmallVectorImpl<int> &OpIndexes,
                                   Intrinsic::ID IID) const override {
     return Impl.collectFlatAddressOperands(OpIndexes, IID);
   }
 
   bool isNoopAddrSpaceCast(unsigned FromAS, unsigned ToAS) const override {

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

public:
   bool hasBranchDivergence() const { return false; }
 
   bool useGPUDivergenceAnalysis() const { return false; }
 
   bool isSourceOfDivergence(const Value *V) const { return false; }
 
   bool isAlwaysUniform(const Value *V) const { return false; }
 
+  bool isValidAddrSpaceCast(unsigned FromAS, unsigned ToAS) const {
+    return FromAS == ToAS;
+  }
+
   unsigned getFlatAddressSpace() const { return -1; }
 
   bool collectFlatAddressOperands(SmallVectorImpl<int> &OpIndexes,
                                   Intrinsic::ID IID) const {
     return false;
   }
 
   bool isNoopAddrSpaceCast(unsigned, unsigned) const { return false; }

llvm/include/llvm/CodeGen/BasicTTIImpl.h

public:
   bool hasBranchDivergence() { return false; }
 
   bool useGPUDivergenceAnalysis() { return false; }
 
   bool isSourceOfDivergence(const Value *V) { return false; }
 
   bool isAlwaysUniform(const Value *V) { return false; }
 
+  bool isValidAddrSpaceCast(unsigned FromAS, unsigned ToAS) const {
+    return FromAS == ToAS;
+  }
+
   unsigned getFlatAddressSpace() {
     // Return an invalid address space.
     return -1;
   }
 
   bool collectFlatAddressOperands(SmallVectorImpl<int> &OpIndexes,
                                   Intrinsic::ID IID) const {
     return false;

llvm/include/llvm/Transforms/InstCombine/InstCombiner.h

public:
 
   virtual bool SimplifyDemandedBits(Instruction *I, unsigned OpNo,
                                     const APInt &DemandedMask, KnownBits &Known,
                                     unsigned Depth = 0) = 0;
   virtual Value *
   SimplifyDemandedVectorElts(Value *V, APInt DemandedElts, APInt &UndefElts,
                              unsigned Depth = 0,
                              bool AllowMultipleUsers = false) = 0;
+
+  bool isValidAddrSpaceCast(unsigned FromAS, unsigned ToAS) const;
 };
 
 } // namespace llvm
 
 #undef DEBUG_TYPE
 
 #endif

llvm/lib/Analysis/TargetTransformInfo.cpp

 bool TargetTransformInfo::isSourceOfDivergence(const Value *V) const {
   return TTIImpl->isSourceOfDivergence(V);
 }
 
 bool llvm::TargetTransformInfo::isAlwaysUniform(const Value *V) const {
   return TTIImpl->isAlwaysUniform(V);
 }
 
+bool llvm::TargetTransformInfo::isValidAddrSpaceCast(unsigned FromAS,
+                                                     unsigned ToAS) const {
+  return TTIImpl->isValidAddrSpaceCast(FromAS, ToAS);
+}
+
 unsigned TargetTransformInfo::getFlatAddressSpace() const {
   return TTIImpl->getFlatAddressSpace();
 }
 
 bool TargetTransformInfo::collectFlatAddressOperands(
     SmallVectorImpl<int> &OpIndexes, Intrinsic::ID IID) const {
   return TTIImpl->collectFlatAddressOperands(OpIndexes, IID);
 }

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

 // the new pointer. If during the chasing it sees bitcast or GEP, it will
 // create new bitcast or GEP with the new pointer and use them in the load
 // instruction.
 class PointerReplacer {
 public:
   PointerReplacer(InstCombinerImpl &IC, Instruction &Root)
     : IC(IC), Root(Root) {}
 
-  bool collectUsers();
+  bool collectUsers(unsigned SrcAS);
   void replacePointer(Value *V);
 
 private:
-  bool collectUsersRecursive(Instruction &I);
+  bool collectUsersRecursive(unsigned SrcAS, Instruction &I);
   void replace(Instruction *I);
   Value *getReplacement(Value *I);
   bool isAvailable(Instruction *I) const {
     return I == &Root || Worklist.contains(I);
   }
 
+  bool isValidAddrSpaceCast(const Instruction *I, unsigned SrcAS) const {
+    const auto *ASC = dyn_cast<AddrSpaceCastInst>(I);
+    if (!ASC)
+      return false;
+    unsigned DstAS = ASC->getDestAddressSpace();
+    return IC.isValidAddrSpaceCast(DstAS, SrcAS);
+  }
+
   SmallPtrSet<Instruction *, 32> ValuesToRevisit;
   SmallSetVector<Instruction *, 4> Worklist;
   MapVector<Value *, Value *> WorkMap;
   InstCombinerImpl &IC;
   Instruction &Root;
 };
 } // end anonymous namespace
 
-bool PointerReplacer::collectUsers() {
-  if (!collectUsersRecursive(Root))
+bool PointerReplacer::collectUsers(unsigned SrcAS) {
+  if (!collectUsersRecursive(SrcAS, Root))
     return false;
 
   // Ensure that all outstanding (indirect) users of I
   // are inserted into the Worklist. Return false
   // otherwise.
   for (auto *Inst : ValuesToRevisit)
     if (!Worklist.contains(Inst))
       return false;
   return true;
 }
 
-bool PointerReplacer::collectUsersRecursive(Instruction &I) {
+bool PointerReplacer::collectUsersRecursive(unsigned SrcAS, Instruction &I) {
   for (auto *U : I.users()) {
     auto *Inst = cast<Instruction>(&*U);
     if (auto *Load = dyn_cast<LoadInst>(Inst)) {
       if (Load->isVolatile())
         return false;
       Worklist.insert(Load);
     } else if (auto *PHI = dyn_cast<PHINode>(Inst)) {
       // All incoming values must be instructions for replacability
       if (any_of(PHI->incoming_values(),
                  [](Value *V) { return !isa<Instruction>(V); }))
         return false;
 
       // If at least one incoming value of the PHI is not in Worklist,
       // store the PHI for revisiting and skip this iteration of the
       // loop.
       if (any_of(PHI->incoming_values(), [this](Value *V) {
             return !isAvailable(cast<Instruction>(V));
           })) {
         ValuesToRevisit.insert(Inst);
         continue;
       }
 
       Worklist.insert(PHI);
-      if (!collectUsersRecursive(*PHI))
+      if (!collectUsersRecursive(SrcAS, *PHI))
         return false;
     } else if (auto *SI = dyn_cast<SelectInst>(Inst)) {
       if (!isa<Instruction>(SI->getTrueValue()) ||
           !isa<Instruction>(SI->getFalseValue()))
         return false;
 
       if (!isAvailable(cast<Instruction>(SI->getTrueValue())) ||
           !isAvailable(cast<Instruction>(SI->getFalseValue()))) {
         ValuesToRevisit.insert(Inst);
         continue;
       }
       Worklist.insert(SI);
-      if (!collectUsersRecursive(*SI))
+      if (!collectUsersRecursive(SrcAS, *SI))
         return false;
     } else if (isa<GetElementPtrInst, BitCastInst>(Inst)) {
       Worklist.insert(Inst);
-      if (!collectUsersRecursive(*Inst))
+      if (!collectUsersRecursive(SrcAS, *Inst))
         return false;
     } else if (auto *MI = dyn_cast<MemTransferInst>(Inst)) {
       if (MI->isVolatile())
         return false;
       Worklist.insert(Inst);
+    } else if (isValidAddrSpaceCast(Inst, SrcAS)) {
+      Worklist.insert(Inst);
     } else if (Inst->isLifetimeStartOrEnd()) {
       continue;
     } else {
       LLVM_DEBUG(dbgs() << "Cannot handle pointer user: " << *U << '\n');
       return false;
     }
   }
 
auto *NewI = IC.Builder.CreateMemTransferInst(
         SrcV, MemCpy->getSourceAlign(), MemCpy->getLength(),
         MemCpy->isVolatile());
     AAMDNodes AAMD = MemCpy->getAAMetadata();
     if (AAMD)
       NewI->setAAMetadata(AAMD);
 
     IC.eraseInstFromFunction(*MemCpy);
     WorkMap[MemCpy] = NewI;
+  } else if (auto *ASC = dyn_cast<AddrSpaceCastInst>(I)) {
+    auto *V = getReplacement(ASC->getPointerOperand());
+    assert(V && "Operand not replaced");
+    assert(isValidAddrSpaceCast(ASC, V->getType()->getPointerAddressSpace()) &&
+           "Invalid address space cast!");
+    auto *NewV = V;
+    if (V->getType()->getPointerAddressSpace() !=
+        ASC->getType()->getPointerAddressSpace()) {
+      auto *NewI = new AddrSpaceCastInst(V, ASC->getType(), "");
+      NewI->takeName(ASC);
+      IC.InsertNewInstWith(NewI, *ASC);
+      NewV = NewI;
+    }
+    IC.replaceInstUsesWith(*ASC, NewV);
+    IC.eraseInstFromFunction(*ASC);
   } else {
     llvm_unreachable("should never reach here");
   }
 }
 
 void PointerReplacer::replacePointer(Value *V) {
 #ifndef NDEBUG
   auto *PT = cast<PointerType>(Root.getType());
if (AllocaAlign <= SourceAlign &&
         Value *Cast = Builder.CreateBitCast(TheSrc, DestTy);
         Instruction *NewI = replaceInstUsesWith(AI, Cast);
         eraseInstFromFunction(*Copy);
         ++NumGlobalCopies;
         return NewI;
       }
 
       PointerReplacer PtrReplacer(*this, AI);
-      if (PtrReplacer.collectUsers()) {
+      if (PtrReplacer.collectUsers(
+              TheSrc->getType()->getPointerAddressSpace())) {
         for (Instruction *Delete : ToDelete)
           eraseInstFromFunction(*Delete);
 
         Value *Cast = Builder.CreateBitCast(TheSrc, DestTy);
         PtrReplacer.replacePointer(Cast);
         ++NumGlobalCopies;
       }
     }

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

std::optional<Value *> InstCombiner::targetSimplifyDemandedVectorEltsIntrinsic(
   if (II.getCalledFunction()->isTargetIntrinsic()) {
     return TTI.simplifyDemandedVectorEltsIntrinsic(
         *this, II, DemandedElts, UndefElts, UndefElts2, UndefElts3,
         SimplifyAndSetOp);
   }
   return std::nullopt;
 }
 
+bool InstCombiner::isValidAddrSpaceCast(unsigned FromAS, unsigned ToAS) const {
+  return TTI.isValidAddrSpaceCast(FromAS, ToAS);
+}
+
 Value *InstCombinerImpl::EmitGEPOffset(User *GEP) {
   return llvm::emitGEPOffset(&Builder, DL, GEP);
 }
 
 /// Legal integers and common types are considered desirable. This is used to
 /// avoid creating instructions with types that may not be supported well by the
 /// the backend.
 /// NOTE: This treats i8, i16 and i32 specially because they are common

llvm/test/Transforms/InstCombine/ptr-replace-alloca.ll

entry:
   call void @llvm.memcpy.p0.p1.i64(ptr %alloca, ptr addrspace(1) @g2, i64 32, i1 false)
   %gep = getelementptr inbounds [32 x i8], ptr %alloca, i32 0, i32 2
   %sel = select i1 %cond, ptr %alloca, ptr %gep
   %gep2 = getelementptr inbounds i8, ptr %sel, i64 4
   %load = load i8, ptr %gep2
   ret i8 %load
 }
 
+declare i8 @readonly_callee(ptr readonly nocapture)
+
+define i8 @call_readonly_remove_alloca() {
+; CHECK-LABEL: @call_readonly_remove_alloca(
+; CHECK-NEXT:    [[V:%.*]] = call i8 @readonly_callee(ptr nonnull @g1)
+; CHECK-NEXT:    ret i8 [[V]]
+;
+  %alloca = alloca [32 x i8], addrspace(1)
+  call void @llvm.memcpy.p1.p0.i64(ptr addrspace(1) %alloca, ptr @g1, i64 32, i1 false)
+  %p = addrspacecast ptr addrspace(1) %alloca to ptr
+  %v = call i8 @readonly_callee(ptr %p)
+  ret i8 %v
+}
+
 declare void @llvm.memcpy.p1.p0.i64(ptr addrspace(1), ptr, i64, i1)
 declare void @llvm.memcpy.p0.p0.i64(ptr, ptr, i64, i1)
 declare void @llvm.memcpy.p0.p1.i64(ptr, ptr addrspace(1), i64, i1)