[Attributor] Use more appropriate approach to check flat address space (

llvm#108713)

llvm/include/llvm/Transforms/IPO/Attributor.h

@@ -1341,6 +1341,9 @@ struct InformationCache {
  const ArrayRef<Function *>
  getIndirectlyCallableFunctions(Attributor &A) const;
 
+ /// Return the flat address space if the associated target has.
+ std::optional<unsigned> getFlatAddressSpace() const;
+
 private:
  struct FunctionInfo {
  ~FunctionInfo();
@@ -6267,11 +6270,12 @@ struct AAAddressSpace : public StateWrapper<BooleanState, AbstractAttribute> {
  return (AA->getIdAddr() == &ID);
  }
 
- // No address space which indicates the associated value is dead.
- static const uint32_t NoAddressSpace = ~0U;
-
  /// Unique ID (due to the unique address)
  static const char ID;
+
+protected:
+ // Invalid address space which indicates the associated value is dead.
+ static const uint32_t InvalidAddressSpace = ~0U;
 };
 
 struct AAAllocationInfo : public StateWrapper<BooleanState, AbstractAttribute> {

llvm/lib/Transforms/IPO/Attributor.cpp

@@ -3294,6 +3294,12 @@ InformationCache::getIndirectlyCallableFunctions(Attributor &A) const {
  return IndirectlyCallableFunctions;
 }
 
+std::optional<unsigned> InformationCache::getFlatAddressSpace() const {
+ if (TargetTriple.isAMDGPU() || TargetTriple.isNVPTX())
+ return 0;
+ return std::nullopt;
+}
+
 void Attributor::recordDependence(const AbstractAttribute &FromAA,
  const AbstractAttribute &ToAA,
  DepClassTy DepClass) {

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

@@ -12580,8 +12580,19 @@ struct AAAddressSpaceImpl : public AAAddressSpace {
  void initialize(Attributor &A) override {
  assert(getAssociatedType()->isPtrOrPtrVectorTy() &&
  "Associated value is not a pointer");
- if (getAssociatedType()->getPointerAddressSpace())
+
+ if (!A.getInfoCache().getFlatAddressSpace().has_value()) {
+ indicatePessimisticFixpoint();
+ return;
+ }
+
+ unsigned FlatAS = A.getInfoCache().getFlatAddressSpace().value();
+ unsigned AS = getAssociatedType()->getPointerAddressSpace();
+ if (AS != FlatAS) {
+ [[maybe_unused]] bool R = takeAddressSpace(AS);
+ assert(R && "The take should happen");
  indicateOptimisticFixpoint();
+ }
  }
 
  ChangeStatus updateImpl(Attributor &A) override {
@@ -12603,12 +12614,13 @@ struct AAAddressSpaceImpl : public AAAddressSpace {
 
  /// See AbstractAttribute::manifest(...).
  ChangeStatus manifest(Attributor &A) override {
- Value *AssociatedValue = &getAssociatedValue();
- Value *OriginalValue = peelAddrspacecast(AssociatedValue);
- if (getAddressSpace() == NoAddressSpace ||
+ if (getAddressSpace() == InvalidAddressSpace ||
  getAddressSpace() == getAssociatedType()->getPointerAddressSpace())
  return ChangeStatus::UNCHANGED;
 
+ Value *AssociatedValue = &getAssociatedValue();
+ Value *OriginalValue = peelAddrspacecast(AssociatedValue);
+
  PointerType *NewPtrTy =
  PointerType::get(getAssociatedType()->getContext(), getAddressSpace());
  bool UseOriginalValue =
@@ -12655,17 +12667,17 @@ struct AAAddressSpaceImpl : public AAAddressSpace {
  if (!isValidState())
  return "addrspace(<invalid>)";
  return "addrspace(" +
- (AssumedAddressSpace == NoAddressSpace
+ (AssumedAddressSpace == InvalidAddressSpace
  ? "none"
  : std::to_string(AssumedAddressSpace)) +
  ")";
  }
 
 private:
- uint32_t AssumedAddressSpace = NoAddressSpace;
+ uint32_t AssumedAddressSpace = InvalidAddressSpace;
 
  bool takeAddressSpace(uint32_t AS) {
- if (AssumedAddressSpace == NoAddressSpace) {
+ if (AssumedAddressSpace == InvalidAddressSpace) {
  AssumedAddressSpace = AS;
  return true;
  }

llvm/test/Transforms/Attributor/address_space_info.ll

@@ -1,5 +1,7 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --function-signature --check-attributes --check-globals --prefix-filecheck-ir-name true
-; RUN: opt -aa-pipeline=basic-aa -passes=attributor -attributor-manifest-internal -attributor-annotate-decl-cs -S < %s | FileCheck %s --check-prefixes=CHECK
+; RUN: opt -mtriple=amdgcn-amd-amdhsa -aa-pipeline=basic-aa -passes=attributor -attributor-manifest-internal -attributor-annotate-decl-cs -S < %s | FileCheck %s --check-prefix=CHECK
+
+; REQUIRES: amdgpu-registered-target
 
 @dst = dso_local addrspace(1) externally_initialized global i32 0, align 4
 @g1 = dso_local addrspace(1) externally_initialized global ptr null, align 4

llvm/test/Transforms/Attributor/nocapture-1.ll

@@ -257,7 +257,7 @@ define i32 @nc1_addrspace(ptr %q, ptr addrspace(1) %p, i1 %b) {
 ; TUNIT-NEXT: [[TMP:%.*]] = addrspacecast ptr addrspace(1) [[P]] to ptr
 ; TUNIT-NEXT: [[TMP2:%.*]] = select i1 [[B]], ptr [[TMP]], ptr [[Q]]
 ; TUNIT-NEXT: [[VAL:%.*]] = load i32, ptr [[TMP2]], align 4
-; TUNIT-NEXT: store i32 0, ptr addrspace(1) [[P]], align 4
+; TUNIT-NEXT: store i32 0, ptr [[TMP]], align 4
 ; TUNIT-NEXT: store ptr [[Q]], ptr @g, align 8
 ; TUNIT-NEXT: ret i32 [[VAL]]
 ;
@@ -272,7 +272,7 @@ define i32 @nc1_addrspace(ptr %q, ptr addrspace(1) %p, i1 %b) {
 ; CGSCC-NEXT: [[TMP:%.*]] = addrspacecast ptr addrspace(1) [[P]] to ptr
 ; CGSCC-NEXT: [[TMP2:%.*]] = select i1 [[B]], ptr [[TMP]], ptr [[Q]]
 ; CGSCC-NEXT: [[VAL:%.*]] = load i32, ptr [[TMP2]], align 4
-; CGSCC-NEXT: store i32 0, ptr addrspace(1) [[P]], align 4
+; CGSCC-NEXT: store i32 0, ptr [[TMP]], align 4
 ; CGSCC-NEXT: store ptr [[Q]], ptr @g, align 8
 ; CGSCC-NEXT: ret i32 [[VAL]]
 ;

llvm/test/Transforms/Attributor/value-simplify.ll

@@ -838,8 +838,7 @@ define void @user() {
 ; TUNIT: Function Attrs: mustprogress nofree norecurse nosync nounwind willreturn memory(write)
 ; TUNIT-LABEL: define {{[^@]+}}@user
 ; TUNIT-SAME: () #[[ATTR5]] {
-; TUNIT-NEXT: [[TMP1:%.*]] = addrspacecast ptr addrspacecast (ptr addrspace(3) @ConstAS3Ptr to ptr) to ptr addrspace(3)
-; TUNIT-NEXT: store i32 0, ptr addrspace(3) [[TMP1]], align 4
+; TUNIT-NEXT: store i32 0, ptr addrspacecast (ptr addrspace(3) @ConstAS3Ptr to ptr), align 4
 ; TUNIT-NEXT: ret void
 ;
 ; CGSCC: Function Attrs: mustprogress nofree nosync nounwind willreturn memory(write)