[llvm] Attributor: Infer noalias.addrspace metadata for memory instructions (PR #136553)

[Matt Arsenault via llvm-commits]

================
@@ -12778,6 +12779,217 @@ struct AAAddressSpaceCallSiteArgument final : AAAddressSpaceImpl {
 };
 } // namespace
 
+/// ------------------------ No Alias Address Space  ---------------------------
+namespace {
+struct AANoAliasAddrSpaceImpl : public AANoAliasAddrSpace {
+  AANoAliasAddrSpaceImpl(const IRPosition &IRP, Attributor &A)
+      : AANoAliasAddrSpace(IRP, A) {}
+
+  void initialize(Attributor &A) override {
+    assert(getAssociatedType()->isPtrOrPtrVectorTy() &&
+           "Associated value is not a pointer");
+
+    resetASRanges(A);
+
+    auto FlatAS = A.getInfoCache().getFlatAddressSpace();
+    if (!FlatAS.has_value()) {
+      indicatePessimisticFixpoint();
+      return;
+    }
+
+    removeAS(FlatAS.value());
+
+    unsigned AS = getAssociatedType()->getPointerAddressSpace();
+    if (AS != FlatAS.value()) {
+      removeAS(AS);
+      indicateOptimisticFixpoint();
+    }
+  }
+
+  ChangeStatus updateImpl(Attributor &A) override {
+    unsigned FlatAS = A.getInfoCache().getFlatAddressSpace().value();
+    uint32_t OrigAssumed = getAssumed();
+
+    auto CheckAddressSpace = [&](Value &Obj) {
+      if (isa<PoisonValue>(&Obj))
+        return true;
+      // Handle argument in flat address space only has addrspace cast uses
+      if (auto *Arg = dyn_cast<Argument>(&Obj)) {
+        if (Arg->getType()->getPointerAddressSpace() == FlatAS &&
+            getAssociatedFunction()->hasKernelCallingConv()) {
+          for (auto *U : Arg->users()) {
+            auto *ASCI = dyn_cast<AddrSpaceCastInst>(U);
+            if (!ASCI)
+              return false;
+            if (ASCI->getDestAddressSpace() == FlatAS)
+              return false;
+            removeAS(ASCI->getDestAddressSpace());
+          }
+          return true;
+        }
+      }
+
+      unsigned AS = Obj.getType()->getPointerAddressSpace();
+      if (AS == FlatAS)
+        return false;
+
+      removeAS(Obj.getType()->getPointerAddressSpace());
+      return true;
+    };
+
+    auto *AUO = A.getOrCreateAAFor<AAUnderlyingObjects>(getIRPosition(), this,
+                                                        DepClassTy::REQUIRED);
+    if (!AUO->forallUnderlyingObjects(CheckAddressSpace)) {
+      resetASRanges(A);
+      return indicatePessimisticFixpoint();
+    }
+
+    return OrigAssumed == getAssumed() ? ChangeStatus::UNCHANGED
+                                       : ChangeStatus::CHANGED;
+  }
+
+  /// See AbstractAttribute::manifest(...).
+  ChangeStatus manifest(Attributor &A) override {
+    auto FlatAS = A.getInfoCache().getFlatAddressSpace();
+    if (!FlatAS.has_value())
+      llvm_unreachable("Must have flat address space!");
+
+    unsigned AS = getAssociatedType()->getPointerAddressSpace();
+    if (AS != FlatAS.value())
+      return ChangeStatus::UNCHANGED;
+
+    LLVMContext &Ctx = getAssociatedValue().getContext();
+    MDNode *NoAliasASNode = nullptr;
+    MDBuilder MDB(Ctx);
+    for (std::pair<unsigned, unsigned> Range : ASRanges) {
+      if (NoAliasASNode == nullptr) {
+        NoAliasASNode =
+            MDB.createRange(APInt(32, Range.first), APInt(32, Range.second));
+      } else {
+        MDNode *ASRange =
+            MDB.createRange(APInt(32, Range.first), APInt(32, Range.second));
+        NoAliasASNode = MDNode::getMostGenericRange(NoAliasASNode, ASRange);
+      }
+    }
+
+    if (!NoAliasASNode)
+      return ChangeStatus::UNCHANGED;
+
+    Value *AssociatedValue = &getAssociatedValue();
+    bool Changed = false;
+
+    auto Pred = [&](const Use &U, bool &) {
+      if (U.get() != AssociatedValue)
+        return true;
+      auto *Inst = dyn_cast<Instruction>(U.getUser());
+      if (!Inst)
+        return true;
+      if (!isa<LoadInst>(Inst) && !isa<StoreInst>(Inst) &&
+          !isa<AtomicCmpXchgInst>(Inst) && !isa<AtomicRMWInst>(Inst))
+        return true;
+      if (!A.isRunOn(Inst->getFunction()))
+        return true;
+      Inst->setMetadata(LLVMContext::MD_noalias_addrspace, NoAliasASNode);
+      Changed = true;
+      return true;
+    };
+    (void)A.checkForAllUses(Pred, *this, getAssociatedValue(),
+                            /* CheckBBLivenessOnly=*/true);
+    return Changed ? ChangeStatus::CHANGED : ChangeStatus::UNCHANGED;
+  }
+
+  /// See AbstractAttribute::getAsStr().
+  const std::string getAsStr(Attributor *A) const override {
+    if (!isValidState())
+      return "noaliasaddrspace(<invalid>)";
+    std::string Str;
+    raw_string_ostream OS(Str);
+    OS << "noaliasaddrspace(";
+    for (auto range : ASRanges)
+      OS << ' ' << '[' << range.first << ',' << range.second << ')';
+    OS << " )";
+    return OS.str();
+  }
+
+private:
+  void removeAS(unsigned AS) {
+    for (auto it = ASRanges.begin(); it != ASRanges.end();) {
+      if (it->first == AS) {
+        uint32_t Upper = it->second;
+        ASRanges.erase(it);
+        if (AS + 1 < Upper)
+          ASRanges.push_back(std::pair(AS + 1, Upper));
+        return;
+      } else if (it->second - 1 == AS) {
----------------
arsenm wrote:

No else after return

https://github.com/llvm/llvm-project/pull/136553