[InferAddrSpace] Teach to handle assumed address space.

[hliao]

Summary

• In certain cases, a generic pointer could be assumed as a pointer to the global memory space or other spaces. With a dedicated target hook to query that address space from a given value, infer-address-space pass could infer and propagate that to all its users.

[hliao]

Revise the fix.

[hliao]

Fix clang-tidy warnings.

[hliao]

Rebase

[hliao]

PING for review

[arsenm]

llvm/include/llvm/CodeGen/BasicTTIImpl.h

227

We already have a -1 as an invalid addrspace, so optional isn't necessary

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

533

I would assume !isPointerTy would be unreachable

540

I think this may be too aggressive and should check if it's specifically a kernel argument load. Contant address space should also not be necessary

llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp

562–570

Won't this cast be naturally inserted with the inferred addres space?

[arsenm]

Should also make a note about the restrictions in AMDGPUUsage

[hliao]

llvm/include/llvm/CodeGen/BasicTTIImpl.h

227

OK. Do we need to document that in the IR langref? That invalid address space ID is not documented anywhere.

[hliao]

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

540

If the variable locates in __constant__ memory, it could only be modified on the host side. It's safe to assume any generic pointer loaded from __constant__ memory is a global one.

[hliao]

llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp

562–570

We need to insert it explicitly as it's previously used as a generic pointer. Without this, we cannot change users of that generic pointer to a global one.

[hliao]

• Add a note in the AMDGPU usage document on the assumption made here.
• Revise the test in clang.

[hliao]

llvm/include/llvm/CodeGen/BasicTTIImpl.h

227

Shall we assume this before that reserved invalid address space ID is well-received and documented in LLVM IR ref? I will make either change in this change accordingly.

[hliao]

Revise the condition check.

[hliao]

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

533

With a condition evaluation order revised, we could assert that now.

[arsenm]

llvm/include/llvm/CodeGen/BasicTTIImpl.h

227

It's not really part of the IR. The address space is a 24-bit unsigned limit and this is already the pass convention

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

540

I thought we discussed the one possible case where this wasn't necessarily true

llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp

289–291

This can be return TII->...

llvm/test/Transforms/InferAddressSpaces/AMDGPU/assumed-addrspace.ll

13

Should have test with load from kernel argument

[hliao]

Revise the interface of that target hook.
Add a dedicated test case for value reading from parameter even though most cases are already covered in the clang test.

[hliao]

llvm/include/llvm/CodeGen/BasicTTIImpl.h

227

You are right. Forget that totally.

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

540

For __constant__ memory, the language itself guarantees it won't be changed on the device side. It's very safe to assume that. The frontend should issue an error on any attempt to write it. Also, the runtime and HW is supposed to provide the facility to protect any write into that space from the device side.

I thought the issue we may have the issue is the case for const variable in the global memory space and a const global pointer in the kernel argument. We need to investigate them further to identify potential risks.

llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp

289–291

good catch. The's the result of the previous Optional<unsigned>. Now, it's revised.

llvm/test/Transforms/InferAddressSpaces/AMDGPU/assumed-addrspace.ll

13

most cases are already captured in the clang test. A new function is added.

[hliao]

Kindly ping for review.

[tra]

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

533

Is there a suitable constant for don't know result?

[hliao]

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

533

I would say all generic pointers loaded from the constant memory are safe to be assumed as global ones (for AMDGPU). I explained the reason in the usage document. The constant memory could be modified by the host, where only global memory objects are visible.

[arsenm]

LGTM with nits

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

533

AMDGPUAS::UNKNOWN_ADDRESS_SPACE

540

This could be a bit broader but is fine for now

541

Ditto

[arsenm]

llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp

483

Checking this here looks strange, since it is also catching values that weren't processed as assumed address space

[hliao]

llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp

483

It stops tracking through instructions generating pointers with assumed AS. It's not necessary to check their operands as the result is assumed. For instructions without assumed AS results, we need to track them through their operands.

clang/test/CodeGenCUDA/amdgpu-kernel-arg-pointer-type.cu

 __device__ void func(int *x) {
   x[0]++;
 }
 
 struct S {
   int *x;
   float *y;
 };
-// `by-val` struct will be coerced into a similar struct with all generic
-// pointers lowerd into global ones.
+// `by-val` struct is passed by-indirect-alias (a mix of by-ref and indirect
+// by-val). However, the enhanced address inferring pass should be able to
+// assume they are global pointers.
+//
 // HOST: define void @_Z22__device_stub__kernel41S(i32* %s.coerce0, float* %s.coerce1)
 // COMMON-LABEL: define amdgpu_kernel void @_Z7kernel41S(%struct.S addrspace(4)*{{.*}} byref(%struct.S) align 8 %0)
 // OPT: [[R0:%.*]] = getelementptr inbounds %struct.S, %struct.S addrspace(4)* %0, i64 0, i32 0
 // OPT: [[P0:%.*]] = load i32*, i32* addrspace(4)* [[R0]], align 8
+// OPT: [[G0:%.*]] = addrspacecast i32* [[P0]] to i32 addrspace(1)*
 // OPT: [[R1:%.*]] = getelementptr inbounds %struct.S, %struct.S addrspace(4)* %0, i64 0, i32 1
 // OPT: [[P1:%.*]] = load float*, float* addrspace(4)* [[R1]], align 8
-// OPT: [[V0:%.*]] = load i32, i32* [[P0]], align 4
+// OPT: [[G1:%.*]] = addrspacecast float* [[P1]] to float addrspace(1)*
+// OPT: [[V0:%.*]] = load i32, i32 addrspace(1)* [[G0]], align 4
 // OPT: [[INC:%.*]] = add nsw i32 [[V0]], 1
-// OPT: store i32 [[INC]], i32* [[P0]], align 4
-// OPT: [[V1:%.*]] = load float, float* [[P1]], align 4
+// OPT: store i32 [[INC]], i32 addrspace(1)* [[G0]], align 4
+// OPT: [[V1:%.*]] = load float, float addrspace(1)* [[G1]], align 4
 // OPT: [[ADD:%.*]] = fadd contract float [[V1]], 1.000000e+00
-// OPT: store float [[ADD]], float* [[P1]], align 4
+// OPT: store float [[ADD]], float addrspace(1)* [[G1]], align 4
 // OPT: ret void
 __global__ void kernel4(struct S s) {
   s.x[0]++;
   s.y[0] += 1.f;
 }
 
 // If a pointer to struct is passed, only the pointer itself is coerced into the global one.
 // HOST: define void @_Z22__device_stub__kernel5P1S(%struct.S* %s)
 // COMMON-LABEL: define amdgpu_kernel void @_Z7kernel5P1S(%struct.S addrspace(1)*{{.*}} %s.coerce)
 __global__ void kernel5(struct S *s) {
   s->x[0]++;
   s->y[0] += 1.f;
 }
 
 struct T {
   float *x[2];
 };
-// `by-val` array is also coerced.
+// `by-val` array is passed by-indirect-alias (a mix of by-ref and indirect
+// by-val). However, the enhanced address inferring pass should be able to
+// assume they are global pointers.
+//
 // HOST: define void @_Z22__device_stub__kernel61T(float* %t.coerce0, float* %t.coerce1)
 // COMMON-LABEL: define amdgpu_kernel void @_Z7kernel61T(%struct.T addrspace(4)*{{.*}} byref(%struct.T) align 8 %0)
 // OPT: [[R0:%.*]] = getelementptr inbounds %struct.T, %struct.T addrspace(4)* %0, i64 0, i32 0, i64 0
 // OPT: [[P0:%.*]] = load float*, float* addrspace(4)* [[R0]], align 8
+// OPT: [[G0:%.*]] = addrspacecast float* [[P0]] to float addrspace(1)*
 // OPT: [[R1:%.*]] = getelementptr inbounds %struct.T, %struct.T addrspace(4)* %0, i64 0, i32 0, i64 1
 // OPT: [[P1:%.*]] = load float*, float* addrspace(4)* [[R1]], align 8
-// OPT: [[V0:%.*]] = load float, float* [[P0]], align 4
+// OPT: [[G1:%.*]] = addrspacecast float* [[P1]] to float addrspace(1)*
+// OPT: [[V0:%.*]] = load float, float addrspace(1)* [[G0]], align 4
 // OPT: [[ADD0:%.*]] = fadd contract float [[V0]], 1.000000e+00
-// OPT: store float [[ADD0]], float* [[P0]], align 4
-// OPT: [[V1:%.*]] = load float, float* [[P1]], align 4
+// OPT: store float [[ADD0]], float addrspace(1)* [[G0]], align 4
+// OPT: [[V1:%.*]] = load float, float addrspace(1)* [[G1]], align 4
 // OPT: [[ADD1:%.*]] = fadd contract float [[V1]], 2.000000e+00
-// OPT: store float [[ADD1]], float* [[P1]], align 4
+// OPT: store float [[ADD1]], float addrspace(1)* [[G1]], align 4
 // OPT: ret void
 __global__ void kernel6(struct T t) {
   t.x[0][0] += 1.f;
   t.x[1][0] += 2.f;
 }
 
 // Check that coerced pointers retain the noalias attribute when qualified with __restrict.
 // HOST: define void @_Z22__device_stub__kernel7Pi(i32* noalias %x)

llvm/docs/AMDGPUUsage.rst

 **Global and Constant**
   The global and constant address spaces both use global virtual addresses,
   which are the same virtual address space used by the CPU. However, some
   virtual addresses may only be accessible to the CPU, some only accessible
   by the GPU, and some by both.
 
   Using the constant address space indicates that the data will not change
   during the execution of the kernel. This allows scalar read instructions to
-  be used. The vector and scalar L1 caches are invalidated of volatile data
-  before each kernel dispatch execution to allow constant memory to change
-  values between kernel dispatches.
+  be used. As the constant address space could only be modified on the host
+  side, a generic pointer loaded from the constant address space is safe to be
+  assumed as a global pointer since only the device global memory is visible
+  and managed on the host side. The vector and scalar L1 caches are invalidated
+  of volatile data before each kernel dispatch execution to allow constant
+  memory to change values between kernel dispatches.
 
 **Region**
   The region address space uses the hardware Global Data Store (GDS). All
   wavefronts executing on the same device will access the same memory for any
   given region address. However, the same region address accessed by wavefronts
   executing on different devices will access different memory. It is higher
   performance than global memory. It is allocated by the runtime. The data
   store (DS) instructions can be used to access it.

llvm/include/llvm/Analysis/TargetTransformInfo.h

public:
   /// they use a flat address space pointer.
   ///
   /// \returns true if the intrinsic was handled.
   bool collectFlatAddressOperands(SmallVectorImpl<int> &OpIndexes,
                                   Intrinsic::ID IID) const;
 
   bool isNoopAddrSpaceCast(unsigned FromAS, unsigned ToAS) const;
 
+  unsigned getAssumedAddrSpace(const Value *V) const;
+
   /// Rewrite intrinsic call \p II such that \p OldV will be replaced with \p
   /// NewV, which has a different address space. This should happen for every
   /// operand index that collectFlatAddressOperands returned for the intrinsic.
   /// \returns nullptr if the intrinsic was not handled. Otherwise, returns the
   /// new value (which may be the original \p II with modified operands).
   Value *rewriteIntrinsicWithAddressSpace(IntrinsicInst *II, Value *OldV,
                                           Value *NewV) const;
 
public:
   virtual bool hasBranchDivergence() = 0;
   virtual bool useGPUDivergenceAnalysis() = 0;
   virtual bool isSourceOfDivergence(const Value *V) = 0;
   virtual bool isAlwaysUniform(const Value *V) = 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 unsigned getAssumedAddrSpace(const Value *V) const = 0;
   virtual Value *rewriteIntrinsicWithAddressSpace(IntrinsicInst *II,
                                                   Value *OldV,
                                                   Value *NewV) const = 0;
   virtual bool isLoweredToCall(const Function *F) = 0;
   virtual void getUnrollingPreferences(Loop *L, ScalarEvolution &,
                                        UnrollingPreferences &UP) = 0;
   virtual void getPeelingPreferences(Loop *L, ScalarEvolution &SE,
                                      PeelingPreferences &PP) = 0;
bool collectFlatAddressOperands(SmallVectorImpl<int> &OpIndexes,
                                   Intrinsic::ID IID) const override {
     return Impl.collectFlatAddressOperands(OpIndexes, IID);
   }
 
   bool isNoopAddrSpaceCast(unsigned FromAS, unsigned ToAS) const override {
     return Impl.isNoopAddrSpaceCast(FromAS, ToAS);
   }
 
+  unsigned getAssumedAddrSpace(const Value *V) const override {
+    return Impl.getAssumedAddrSpace(V);
+  }
+
   Value *rewriteIntrinsicWithAddressSpace(IntrinsicInst *II, Value *OldV,
                                           Value *NewV) const override {
     return Impl.rewriteIntrinsicWithAddressSpace(II, OldV, NewV);
   }
 
   bool isLoweredToCall(const Function *F) override {
     return Impl.isLoweredToCall(F);
   }

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

public:
 
   bool collectFlatAddressOperands(SmallVectorImpl<int> &OpIndexes,
                                   Intrinsic::ID IID) const {
     return false;
   }
 
   bool isNoopAddrSpaceCast(unsigned, unsigned) const { return false; }
 
+  unsigned getAssumedAddrSpace(const Value *V) const { return -1; }
+
   Value *rewriteIntrinsicWithAddressSpace(IntrinsicInst *II, Value *OldV,
                                           Value *NewV) const {
     return nullptr;
   }
 
   bool isLoweredToCall(const Function *F) {
     assert(F && "A concrete function must be provided to this routine.");
 

llvm/include/llvm/CodeGen/BasicTTIImpl.h

bool collectFlatAddressOperands(SmallVectorImpl<int> &OpIndexes,
                                   Intrinsic::ID IID) const {
     return false;
   }
 
   bool isNoopAddrSpaceCast(unsigned FromAS, unsigned ToAS) const {
     return getTLI()->getTargetMachine().isNoopAddrSpaceCast(FromAS, ToAS);
   }
 
+  unsigned getAssumedAddrSpace(const Value *V) const {
+    return getTLI()->getTargetMachine().getAssumedAddrSpace(V);
+  }
+
   Value *rewriteIntrinsicWithAddressSpace(IntrinsicInst *II, Value *OldV,
                                           Value *NewV) const {
     return nullptr;
   }
 
   bool isLegalAddImmediate(int64_t imm) {
     return getTLI()->isLegalAddImmediate(imm);
   }

llvm/include/llvm/Target/TargetMachine.h

const MemoryBuffer *getBBSectionsFuncListBuf() const {
     return Options.BBSectionsFuncListBuf.get();
   }
 
   /// Returns true if a cast between SrcAS and DestAS is a noop.
   virtual bool isNoopAddrSpaceCast(unsigned SrcAS, unsigned DestAS) const {
     return false;
   }
 
+  /// If the specified generic pointer could be assumed as a pointer to a
+  /// specific address space, return that address space.
+  ///
+  /// Under offloading programming, the offloading target may be passed with
+  /// values only prepared on the host side and could assume certain
+  /// properties.
+  virtual unsigned getAssumedAddrSpace(const Value *V) const { return -1; }
+
   /// Get a \c TargetIRAnalysis appropriate for the target.
   ///
   /// This is used to construct the new pass manager's target IR analysis pass,
   /// set up appropriately for this target machine. Even the old pass manager
   /// uses this to answer queries about the IR.
   TargetIRAnalysis getTargetIRAnalysis();
 
   /// Return a TargetTransformInfo for a given function.

llvm/lib/Analysis/TargetTransformInfo.cpp

bool TargetTransformInfo::collectFlatAddressOperands(
   return TTIImpl->collectFlatAddressOperands(OpIndexes, IID);
 }
 
 bool TargetTransformInfo::isNoopAddrSpaceCast(unsigned FromAS,
                                               unsigned ToAS) const {
   return TTIImpl->isNoopAddrSpaceCast(FromAS, ToAS);
 }
 
+unsigned TargetTransformInfo::getAssumedAddrSpace(const Value *V) const {
+  return TTIImpl->getAssumedAddrSpace(V);
+}
+
 Value *TargetTransformInfo::rewriteIntrinsicWithAddressSpace(
     IntrinsicInst *II, Value *OldV, Value *NewV) const {
   return TTIImpl->rewriteIntrinsicWithAddressSpace(II, OldV, NewV);
 }
 
 bool TargetTransformInfo::isLoweredToCall(const Function *F) const {
   return TTIImpl->isLoweredToCall(F);
 }

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.h

public:
   /// Get the integer value of a null pointer in the given address space.
   static int64_t getNullPointerValue(unsigned AddrSpace) {
     return (AddrSpace == AMDGPUAS::LOCAL_ADDRESS ||
             AddrSpace == AMDGPUAS::PRIVATE_ADDRESS ||
             AddrSpace == AMDGPUAS::REGION_ADDRESS) ? -1 : 0;
   }
 
   bool isNoopAddrSpaceCast(unsigned SrcAS, unsigned DestAS) const override;
+
+  unsigned getAssumedAddrSpace(const Value *V) const override;
 };
 
 //===----------------------------------------------------------------------===//
 // R600 Target Machine (R600 -> Cayman)
 //===----------------------------------------------------------------------===//
 
 class R600TargetMachine final : public AMDGPUTargetMachine {
 private:

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

 }
 
 bool AMDGPUTargetMachine::isNoopAddrSpaceCast(unsigned SrcAS,
                                               unsigned DestAS) const {
   return AMDGPU::isFlatGlobalAddrSpace(SrcAS) &&
          AMDGPU::isFlatGlobalAddrSpace(DestAS);
 }
 
+unsigned AMDGPUTargetMachine::getAssumedAddrSpace(const Value *V) const {
+  const auto *LD = dyn_cast<LoadInst>(V);
+  if (!LD)
+    return AMDGPUAS::UNKNOWN_ADDRESS_SPACE;
+
+  // It must be a generic pointer loaded.
+  assert(V->getType()->isPointerTy() &&
+         V->getType()->getPointerAddressSpace() == AMDGPUAS::FLAT_ADDRESS);
+
+  const auto *Ptr = LD->getPointerOperand();
+  if (Ptr->getType()->getPointerAddressSpace() != AMDGPUAS::CONSTANT_ADDRESS)
+    return AMDGPUAS::UNKNOWN_ADDRESS_SPACE;
+  // For a generic pointer loaded from the constant memory, it could be assumed
+  // as a global pointer since the constant memory is only populated on the
+  // host side. As implied by the offload programming model, only global
+  // pointers could be referenced on the host side.
+  return AMDGPUAS::GLOBAL_ADDRESS;
+}
+
 TargetTransformInfo
 R600TargetMachine::getTargetTransformInfo(const Function &F) {
   return TargetTransformInfo(R600TTIImpl(this, F));
 }
 
 //===----------------------------------------------------------------------===//
 // GCN Target Machine (SI+)
 //===----------------------------------------------------------------------===//

llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp

case Instruction::Select:
     return Op->getType()->isPointerTy();
   case Instruction::Call: {
     const IntrinsicInst *II = dyn_cast<IntrinsicInst>(&V);
     return II && II->getIntrinsicID() == Intrinsic::ptrmask;
   }
   case Instruction::IntToPtr:
     return isNoopPtrIntCastPair(Op, DL, TTI);
   default:
-    return false;
+    // That value is an address expression if it has an assumed address space.
+    return TTI->getAssumedAddrSpace(&V) != UninitializedAddressSpace;
   }
 }
 
 // Returns the pointer operands of V.
 //
 // Precondition: V is an address expression.
 static SmallVector<Value *, 2>
 getPointerOperands(const Value &V, const DataLayout &DL,
                    const TargetTransformInfo *TTI) {
void InferAddressSpaces::appendsFlatAddressExpressionToPostorderStack(
   if (ConstantExpr *CE = dyn_cast<ConstantExpr>(V)) {
     // TODO: Look in non-address parts, like icmp operands.
     if (isAddressExpression(*CE, *DL, TTI) && Visited.insert(CE).second)
       PostorderStack.emplace_back(CE, false);
 
     return;
   }
 
-  if (isAddressExpression(*V, *DL, TTI) &&
-      V->getType()->getPointerAddressSpace() == FlatAddrSpace) {
+  if (V->getType()->getPointerAddressSpace() == FlatAddrSpace &&
+      isAddressExpression(*V, *DL, TTI)) {
     if (Visited.insert(V).second) {
       PostorderStack.emplace_back(V, false);
 
       Operator *Op = cast<Operator>(V);
       for (unsigned I = 0, E = Op->getNumOperands(); I != E; ++I) {
         if (ConstantExpr *CE = dyn_cast<ConstantExpr>(Op->getOperand(I))) {
           if (isAddressExpression(*CE, *DL, TTI) && Visited.insert(CE).second)
             PostorderStack.emplace_back(CE, false);
while (!PostorderStack.empty()) {
     if (PostorderStack.back().getInt()) {
       if (TopVal->getType()->getPointerAddressSpace() == FlatAddrSpace)
         Postorder.push_back(TopVal);
       PostorderStack.pop_back();
       continue;
     }
     // Otherwise, adds its operands to the stack and explores them.
     PostorderStack.back().setInt(true);
+    // Skip values with an assumed address space.
+    if (TTI->getAssumedAddrSpace(TopVal) == UninitializedAddressSpace) {
       for (Value *PtrOperand : getPointerOperands(*TopVal, *DL, TTI)) {
         appendsFlatAddressExpressionToPostorderStack(PtrOperand, PostorderStack,
                                                      Visited);
       }
     }
+  }
   return Postorder;
 }
 
 // A helper function for cloneInstructionWithNewAddressSpace. Returns the clone
 // of OperandUse.get() in the new address space. If the clone is not ready yet,
 // returns an undef in the new address space as a placeholder.
 static Value *operandWithNewAddressSpaceOrCreateUndef(
     const Use &OperandUse, unsigned NewAddrSpace,
if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(I)) {
     if (Rewrite) {
       assert(Rewrite != II && "cannot modify this pointer operation in place");
       return Rewrite;
     }
 
     return nullptr;
   }
 
+  unsigned AS = TTI->getAssumedAddrSpace(I);
+  if (AS != UninitializedAddressSpace) {
+    // For the assumed address space, insert an `addrspacecast` to make that
+    // explicit.
+    auto *NewPtrTy = I->getType()->getPointerElementType()->getPointerTo(AS);
+    auto *NewI = new AddrSpaceCastInst(I, NewPtrTy);
+    NewI->insertAfter(I);
+    return NewI;
+  }
+
   // Computes the converted pointer operands.
   SmallVector<Value *, 4> NewPointerOperands;
   for (const Use &OperandUse : I->operands()) {
     if (!OperandUse.get()->getType()->isPointerTy())
       NewPointerOperands.push_back(nullptr);
     else
       NewPointerOperands.push_back(operandWithNewAddressSpaceOrCreateUndef(
                                      OperandUse, NewAddrSpace, ValueWithNewAddrSpace, UndefUsesToFix));

 // expression whose address space needs to be modified, in postorder.
 //
 // See cloneInstructionWithNewAddressSpace for the meaning of UndefUsesToFix.
 Value *InferAddressSpaces::cloneValueWithNewAddressSpace(
   Value *V, unsigned NewAddrSpace,
   const ValueToValueMapTy &ValueWithNewAddrSpace,
   SmallVectorImpl<const Use *> *UndefUsesToFix) const {
   // All values in Postorder are flat address expressions.
-  assert(isAddressExpression(*V, *DL, TTI) &&
-         V->getType()->getPointerAddressSpace() == FlatAddrSpace);
+  assert(V->getType()->getPointerAddressSpace() == FlatAddrSpace &&
+         isAddressExpression(*V, *DL, TTI));
 
   if (Instruction *I = dyn_cast<Instruction>(V)) {
     Value *NewV = cloneInstructionWithNewAddressSpace(
       I, NewAddrSpace, ValueWithNewAddrSpace, UndefUsesToFix);
     if (Instruction *NewI = dyn_cast_or_null<Instruction>(NewV)) {
       if (NewI->getParent() == nullptr) {
         NewI->insertBefore(I);
         NewI->takeName(I);
if (Op.getOpcode() == Instruction::Select) {
 
     if (C0 && isSafeToCastConstAddrSpace(C0, Src1AS))
       NewAS = Src1AS;
     else if (C1 && isSafeToCastConstAddrSpace(C1, Src0AS))
       NewAS = Src0AS;
     else
       NewAS = joinAddressSpaces(Src0AS, Src1AS);
   } else {
+    unsigned AS = TTI->getAssumedAddrSpace(&V);
+    if (AS != UninitializedAddressSpace) {
+      // Use the assumed address space directly.
+      NewAS = AS;
+    } else {
+      // Otherwise, infer the address space from its pointer operands.
       for (Value *PtrOperand : getPointerOperands(V, *DL, TTI)) {
         auto I = InferredAddrSpace.find(PtrOperand);
-      unsigned OperandAS = I != InferredAddrSpace.end() ?
-        I->second : PtrOperand->getType()->getPointerAddressSpace();
+        unsigned OperandAS =
+            I != InferredAddrSpace.end()
+                ? I->second
+                : PtrOperand->getType()->getPointerAddressSpace();
 
         // join(flat, *) = flat. So we can break if NewAS is already flat.
         NewAS = joinAddressSpaces(NewAS, OperandAS);
         if (NewAS == FlatAddrSpace)
           break;
       }
     }
+  }
 
   unsigned OldAS = InferredAddrSpace.lookup(&V);
   assert(OldAS != FlatAddrSpace);
   if (OldAS == NewAS)
     return None;
   return NewAS;
 }
 
for (I = V->use_begin(), E = V->use_end(); I != E; ) {
         // If V is used as the pointer operand of a compatible memory operation,
         // sets the pointer operand to NewV. This replacement does not change
         // the element type, so the resultant load/store is still valid.
         U.set(NewV);
         continue;
       }
 
       User *CurUser = U.getUser();
+      // Skip if the current user is the new value itself.
+      if (CurUser == NewV)
+        continue;
       // Handle more complex cases like intrinsic that need to be remangled.
       if (auto *MI = dyn_cast<MemIntrinsic>(CurUser)) {
         if (!MI->isVolatile() && handleMemIntrinsicPtrUse(MI, V, NewV))
           continue;
       }
 
       if (auto *II = dyn_cast<IntrinsicInst>(CurUser)) {
         if (rewriteIntrinsicOperands(II, V, NewV))

llvm/test/CodeGen/AMDGPU/GlobalISel/divergent-control-flow.ll

 ; CHECK-NEXT:    s_load_dwordx2 s[4:5], s[4:5], 0x0
 ; CHECK-NEXT:    s_waitcnt lgkmcnt(0)
 ; CHECK-NEXT:    s_load_dword s4, s[4:5], 0x0
 ; CHECK-NEXT:    s_waitcnt lgkmcnt(0)
 ; CHECK-NEXT:    s_cmp_lg_u32 s4, 0
 ; CHECK-NEXT:    s_cselect_b32 s4, 1, 0
 ; CHECK-NEXT:    s_and_b32 s4, s4, 1
 ; CHECK-NEXT:    s_cmp_lg_u32 s4, 0
-; CHECK-NEXT:    s_cbranch_scc1 BB4_6
+; CHECK-NEXT:    s_cbranch_scc1 BB4_4
 ; CHECK-NEXT:  ; %bb.1: ; %bb2
 ; CHECK-NEXT:    s_getpc_b64 s[6:7]
 ; CHECK-NEXT:    s_add_u32 s6, s6, const.ptr@gotpcrel32@lo+4
 ; CHECK-NEXT:    s_addc_u32 s7, s7, const.ptr@gotpcrel32@hi+12
 ; CHECK-NEXT:    s_load_dwordx2 s[6:7], s[6:7], 0x0
+; CHECK-NEXT:    v_mov_b32_e32 v0, 0
 ; CHECK-NEXT:    s_mov_b32 s4, -1
 ; CHECK-NEXT:    s_waitcnt lgkmcnt(0)
 ; CHECK-NEXT:    s_load_dwordx2 s[6:7], s[6:7], 0x0
 ; CHECK-NEXT:    s_waitcnt lgkmcnt(0)
-; CHECK-NEXT:    v_mov_b32_e32 v0, s6
-; CHECK-NEXT:    v_mov_b32_e32 v1, s7
-; CHECK-NEXT:    flat_load_dword v0, v[0:1]
-; CHECK-NEXT:    s_waitcnt vmcnt(0) lgkmcnt(0)
-; CHECK-NEXT:    v_cmp_ngt_f32_e32 vcc, 1.0, v0
-; CHECK-NEXT:    s_and_saveexec_b64 s[6:7], vcc
+; CHECK-NEXT:    global_load_dword v0, v0, s[6:7]
+; CHECK-NEXT:    s_waitcnt vmcnt(0)
+; CHECK-NEXT:    v_cmp_gt_f32_e32 vcc, 1.0, v0
+; CHECK-NEXT:    s_cbranch_vccnz BB4_3
 ; CHECK-NEXT:  ; %bb.2: ; %bb7
 ; CHECK-NEXT:    s_mov_b32 s4, 0
-; CHECK-NEXT:  ; %bb.3: ; %bb8
-; CHECK-NEXT:    s_or_b64 exec, exec, s[6:7]
-; CHECK-NEXT:    v_cmp_eq_u32_e64 s[6:7], s4, 0
-; CHECK-NEXT:    s_and_saveexec_b64 s[4:5], s[6:7]
-; CHECK-NEXT:    s_cbranch_execz BB4_5
-; CHECK-NEXT:  ; %bb.4: ; %bb11
+; CHECK-NEXT:  BB4_3: ; %bb8
+; CHECK-NEXT:    s_cmp_lg_u32 s4, 0
+; CHECK-NEXT:    s_cselect_b32 s4, 1, 0
+; CHECK-NEXT:    s_and_b32 s4, s4, 1
+; CHECK-NEXT:    s_cmp_lg_u32 s4, 0
+; CHECK-NEXT:    s_cbranch_scc0 BB4_5
+; CHECK-NEXT:  BB4_4: ; %bb12
+; CHECK-NEXT:    s_setpc_b64 s[30:31]
+; CHECK-NEXT:  BB4_5: ; %bb11
 ; CHECK-NEXT:    v_mov_b32_e32 v0, 4.0
 ; CHECK-NEXT:    buffer_store_dword v0, v0, s[0:3], 0 offen
-; CHECK-NEXT:  BB4_5: ; %Flow
-; CHECK-NEXT:    s_or_b64 exec, exec, s[4:5]
-; CHECK-NEXT:  BB4_6: ; %bb12
 ; CHECK-NEXT:    s_waitcnt vmcnt(0)
 ; CHECK-NEXT:    s_setpc_b64 s[30:31]
 bb:
   %tmp = load i32, i32 addrspace(4)* @external_constant
   %ptr = load float*, float* addrspace(4)* @const.ptr
   %tmp1 = icmp ne i32 %tmp, 0
   br i1 %tmp1, label %bb12, label %bb2
 

llvm/test/Transforms/InferAddressSpaces/AMDGPU/assumed-addrspace.ll

+; RUN: opt -S -mtriple=amdgcn-amd-amdhsa -infer-address-spaces -o - %s | FileCheck %s
+
+@c0 = addrspace(4) global float* undef
+
+; CHECK-LABEL: @generic_ptr_from_constant
+; CHECK: addrspacecast float* %p to float addrspace(1)*
+; CHECK-NEXT: load float, float addrspace(1)*
+define float @generic_ptr_from_constant() {
+  %p = load float*, float* addrspace(4)* @c0
+  %v = load float, float* %p
+  ret float %v
+}
+
+%struct.S = type { i32*, float* }
+
+; CHECK-LABEL: @generic_ptr_from_aggregate_argument
+; CHECK: addrspacecast i32* %p0 to i32 addrspace(1)*
+; CHECK: addrspacecast float* %p1 to float addrspace(1)*
+; CHECK: load i32, i32 addrspace(1)*
+; CHECK: store float %v1, float addrspace(1)*
+; CHECK: ret
+define amdgpu_kernel void @generic_ptr_from_aggregate_argument(%struct.S addrspace(4)* byref(%struct.S) align 8 %0) {
+  %f0 = getelementptr inbounds %struct.S, %struct.S addrspace(4)* %0, i64 0, i32 0
+  %p0 = load i32*, i32* addrspace(4)* %f0
+  %f1 = getelementptr inbounds %struct.S, %struct.S addrspace(4)* %0, i64 0, i32 1
+  %p1 = load float*, float* addrspace(4)* %f1
+  %v0 = load i32, i32* %p0
+  %v1 = sitofp i32 %v0 to float
+  store float %v1, float* %p1
+  ret void
+}