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

[via llvm-commits]

================
@@ -12534,6 +12535,269 @@ 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_NOEFFECT, as load invariance is also implied by all underlying objects
+  // being load invariant.
+  //
+  // IS_KNOWN_INVARIANT is set to indicate that the contents of the pointer are
+  // *known* to be invariant, and is therefore a pessimistic bit.
+  enum {
+    IS_KNOWN_INVARIANT = 1 << 0,
+    IS_NOALIAS = 1 << 1,
+    IS_NOEFFECT = 1 << 2,
+
+    IS_IMPLIED_INVARIANT = IS_NOALIAS | IS_NOEFFECT,
+  };
+  static_assert(getBestState() == (IS_KNOWN_INVARIANT | IS_IMPLIED_INVARIANT),
+                "Unexpected best state!");
+
+  using Base = StateWrapper<BitIntegerState<uint8_t, 7>, AAInvariantLoadPointer,
+                            uint8_t>;
+
+  // the BitIntegerState is optimistic about IS_NOALIAS and IS_NOEFFECT, but
+  // pessimistic about IS_KNOWN_INVARIANT
+  AAInvariantLoadPointerImpl(const IRPosition &IRP, Attributor &A)
+      : Base(IRP, IS_IMPLIED_INVARIANT) {}
+
+  void initialize(Attributor &A) final {
+    removeAssumedBits(IS_KNOWN_INVARIANT);
+  }
+
+  bool isKnownInvariant() const final {
+    return isKnown(IS_KNOWN_INVARIANT) || isKnown(IS_IMPLIED_INVARIANT);
+  }
+
+  bool isAssumedInvariant() const final {
+    if (isAssumed(IS_KNOWN_INVARIANT) || isAssumed(IS_IMPLIED_INVARIANT))
+      return true;
+    // if the function is callable, optimistically assume that invariance can be
+    // inferred from the caller
+    const auto *F = getAssociatedFunction();
+    return F && isCallableCC(F->getCallingConv());
+  }
+
+  ChangeStatus updateImpl(Attributor &A) override {
+    if (isKnownInvariant())
+      return indicateOptimisticFixpoint();
+
+    ChangeStatus Changed = ChangeStatus::UNCHANGED;
+
+    Changed |= checkNoAlias(A);
+    Changed |= checkNoEffect(A);
+
+    // try to infer invariance from underlying objects
+    const auto *AUO = A.getOrCreateAAFor<AAUnderlyingObjects>(
+        getIRPosition(), this, DepClassTy::REQUIRED);
+
+    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;
+    };
+    if (!AUO->forallUnderlyingObjects(IsInvariantLoadIfPointer))
+      return indicatePessimisticFixpoint();
+
+    if (!UsedAssumedInformation) {
+      // pointer is known (not assumed) to be invariant
+      addKnownBits(IS_KNOWN_INVARIANT);
+      return indicateOptimisticFixpoint() | 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)) {
+
+        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 {}
+
+private:
+  ChangeStatus checkNoAlias(Attributor &A) {
+    if (isKnown(IS_NOALIAS) || !isAssumed(IS_NOALIAS))
+      return ChangeStatus::UNCHANGED;
+
+    const auto *F = getAssociatedFunction();
+
+    if (F && isCallableCC(F->getCallingConv())) {
+      // program-wide alias information cannot be inferred
+      removeAssumedBits(IS_NOALIAS);
+      return ChangeStatus::CHANGED;
+    }
+
+    // try to use AANoAlias
+    if (const auto *ANoAlias = A.getOrCreateAAFor<AANoAlias>(
+            getIRPosition(), this, DepClassTy::REQUIRED)) {
+      if (ANoAlias->isKnownNoAlias()) {
+        addKnownBits(IS_NOALIAS);
+        return ChangeStatus::UNCHANGED;
+      }
+
+      if (!ANoAlias->isAssumedNoAlias()) {
+        removeAssumedBits(IS_NOALIAS);
+        return ChangeStatus::CHANGED;
+      }
+
+      return ChangeStatus::UNCHANGED;
+    }
+
+    // if the function is not callable, try to infer noalias from argument
+    // attribute, since it is applicable for the duration of the function
+    if (const auto *Arg = getAssociatedArgument()) {
+      if (Arg->hasNoAliasAttr()) {
+        addKnownBits(IS_NOALIAS);
+        return ChangeStatus::UNCHANGED;
+      }
+
+      // noalias information is not provided, and cannot be inferred,
+      // so we conservatively assume the pointer aliases.
+      removeAssumedBits(IS_NOALIAS);
+      return ChangeStatus::CHANGED;
+    }
+
+    return ChangeStatus::UNCHANGED;
+  }
+
+  ChangeStatus checkNoEffect(Attributor &A) {
+    if (isKnown(IS_NOEFFECT) || !isAssumed(IS_NOEFFECT))
+      return ChangeStatus::UNCHANGED;
+
+    const auto *F = getAssociatedFunction();
+
+    if (!F)
+      return ChangeStatus::UNCHANGED;
+
+    if (isCallableCC(F->getCallingConv())) {
+      // effects cannot be tracked outside of function call;
+      // conservatively assume pointer has effectful uses
+      removeAssumedBits(IS_NOEFFECT);
----------------
zGoldthorpe wrote:

The `IS_NOEFFECT` bit checks for more than just communication, and also checks for other side effects involving the pointer (such as being written through).

As for `@test_noalias_buffer_desc`, I'm not sure if there is a nice generalisable way to mark the load as invariant in this case, besides explicitly checking for such intrinsics as `@llvm.amdgcn.make.buffer.rsrc.p7.p1` and forwarding information.

> Also, do we want to add some mechanism for declaring that ptr(4) and ptr(6) are all definitionally constant and use it here?

Is there existing infrastructure that indicates whether address spaces are definitionally constant based on the target architecture?

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