[llvm] [IPO] Added attributor for identifying invariant loads (PR #141800)

[via llvm-commits]

================
@@ -12534,6 +12535,248 @@ struct AAIndirectCallInfoCallSite : public AAIndirectCallInfo {
 };
 } // namespace
 
+/// --------------------- Invariant Load Pointer -------------------------------
+namespace {
+
+struct AAInvariantLoadPointerImpl
+    : public StateWrapper<BitIntegerState<uint8_t, 7>, AAInvariantLoadPointer,
+                          uint8_t> {
+  // load invariance is implied by, but not equivalent to IS_NOALIAS |
+  // IS_READONLY, as load invariance is also implied by all underlying objects
+  // being load invariant.
+  //
+  // IS_INVARIANT is set to indicate that the contents of the pointer are
+  // *known* to be invariant.
+  enum {
+    IS_INVARIANT = 1 << 0,
+    IS_NOALIAS = 1 << 1,
+    IS_READONLY = 1 << 2,
+  };
+  static_assert(getBestState() == (IS_INVARIANT | IS_NOALIAS | IS_READONLY),
+                "Unexpected best state!");
+
+  using Base = StateWrapper<BitIntegerState<uint8_t, 7>, AAInvariantLoadPointer,
+                            uint8_t>;
+
+  // the BitIntegerState is optimistic about noalias and readonly, but
+  // pessimistic about invariance
+  AAInvariantLoadPointerImpl(const IRPosition &IRP, Attributor &A)
+      : Base(IRP, IS_NOALIAS | IS_READONLY) {}
+
+  void initialize(Attributor &A) final {
+    // conservatively assume that the pointer's contents are not invariant,
+    // until proven otherwise.
+    removeAssumedBits(IS_INVARIANT);
+  }
+
+  bool isKnownInvariant() const final {
+    return isKnown(IS_INVARIANT) || isKnown(IS_NOALIAS | IS_READONLY);
+  }
+
+  bool isAssumedInvariant() const final {
+    return isAssumed(IS_INVARIANT) || isAssumed(IS_NOALIAS | IS_READONLY);
+  }
+
+  ChangeStatus updateImpl(Attributor &A) override {
+    if (isKnownInvariant())
+      return ChangeStatus::UNCHANGED;
+
+    ChangeStatus Changed = ChangeStatus::UNCHANGED;
+
+    Changed |= updateNoAlias(A);
+    Changed |= updateReadOnly(A);
+
+    bool UsedAssumedInformation = false;
+    const auto IsInvariantLoadIfPointer = [&](const Value &V) {
+      if (!V.getType()->isPointerTy())
+        return true;
+      const auto *IsInvariantLoadPointer =
+          A.getOrCreateAAFor<AAInvariantLoadPointer>(IRPosition::value(V), this,
+                                                     DepClassTy::REQUIRED);
+      if (IsInvariantLoadPointer->isKnownInvariant())
+        return true;
+      if (!IsInvariantLoadPointer->isAssumedInvariant())
+        return false;
+
+      UsedAssumedInformation = true;
+      return true;
+    };
+
+    const auto *AUO = A.getOrCreateAAFor<AAUnderlyingObjects>(
+        getIRPosition(), this, DepClassTy::REQUIRED);
+
+    if (!AUO->forallUnderlyingObjects(IsInvariantLoadIfPointer)) {
+      removeAssumedBits(IS_INVARIANT);
+      return ChangeStatus::CHANGED;
+    }
+
+    if (!UsedAssumedInformation) {
+      // pointer is known (not assumed) to be invariant
+      addKnownBits(IS_INVARIANT);
+      return ChangeStatus::CHANGED;
+    }
+
+    return Changed;
+  }
+
+  ChangeStatus manifest(Attributor &A) override {
+    if (!isKnownInvariant())
+      return ChangeStatus::UNCHANGED;
+
+    ChangeStatus Changed = ChangeStatus::UNCHANGED;
+    Value *Ptr = &getAssociatedValue();
+    const auto TagInvariantLoads = [&](const Use &U, bool &) {
+      if (U.get() != Ptr)
+        return true;
+      auto *I = dyn_cast<Instruction>(U.getUser());
+      if (!I)
+        return true;
+
+      // Ensure that we are only changing uses from the corresponding callgraph
+      // SSC in the case that the AA isn't run on the entire module
+      if (!A.isRunOn(I->getFunction()))
+        return true;
+
+      if (I->hasMetadata(LLVMContext::MD_invariant_load))
+        return true;
+
+      if (auto *LI = dyn_cast<LoadInst>(I)) {
+        if (LI->isVolatile() || LI->isAtomic())
+          return true;
+
+        LI->setMetadata(LLVMContext::MD_invariant_load,
+                        MDNode::get(LI->getContext(), {}));
+        Changed = ChangeStatus::CHANGED;
+      }
+      return true;
+    };
+
+    (void)A.checkForAllUses(TagInvariantLoads, *this, *Ptr);
+    return Changed;
+  }
+
+  /// See AbstractAttribute::getAsStr().
+  const std::string getAsStr(Attributor *) const override {
+    std::string Str;
+    raw_string_ostream OS(Str);
+    OS << "load invariant pointer: " << isKnown() << '\n';
+    return Str;
+  }
+
+  /// See AbstractAttribute::trackStatistics().
+  void trackStatistics() const override {}
+
+protected:
+  ChangeStatus updateNoAlias(Attributor &A) {
+    if (isKnown(IS_NOALIAS) || !isAssumed(IS_NOALIAS))
+      return ChangeStatus::UNCHANGED;
+
+    const auto *ANoAlias = A.getOrCreateAAFor<AANoAlias>(getIRPosition(), this,
+                                                         DepClassTy::REQUIRED);
+    if (!ANoAlias)
+      return tryInferNoAlias(A);
+
+    if (!ANoAlias->isAssumedNoAlias()) {
+      removeAssumedBits(IS_NOALIAS);
+      return ChangeStatus::CHANGED;
+    }
+    if (ANoAlias->isKnownNoAlias())
+      addKnownBits(IS_NOALIAS);
----------------
zGoldthorpe wrote:

Yes, the intention was to return `UNCHANGED` after this. Based on my understanding of `indicateOptimisticFixpoint()` also returning `UNCHANGED`, I figured the same should happen here since the attributor already assumes `IS_NOALIAS` at this point.

(Or maybe I misunderstood the feedback?)

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