[Attributor] Force simplication callback updates

[ggeorgakoudis]

Summary

A simplification callback communicates to the Attributor a user-provided value simplification of an IR position based on an AbstractAttribute. This patch constraints callback simplification based on the assumption that there should be a single simplification callback per IRP and forces the Attributor to use the callback-returned simplification value, when there is one.

[jdoerfert]

The force update stuff does only work for one level but not more. So an AA depending on one that was forcefully updated is still going to see inconsistent values.
To really solve the issue we need to clear the state in aavaluesimplify (and friends) at the beginning of the update.

llvm/lib/Transforms/IPO/Attributor.cpp

662

Same as below you have to give up and not fallback to aavaluesimplify.

703

I think we just have to return SimplifiedV and be done with it. Fallback to AAValueSimplify is what can cause problems and inconsistencies.

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

struct Attributor {
   Optional<Constant *> getAssumedConstant(const Value &V,
                                           const AbstractAttribute &AA,
                                           bool &UsedAssumedInformation);
 
   /// If \p V is assumed simplified, return it, if it is unclear yet,
   /// return None, otherwise return `nullptr`.
   Optional<Value *> getAssumedSimplified(const IRPosition &IRP,
                                          const AbstractAttribute &AA,
+                                         bool &UsedAssumedInformation,
+                                         bool &ForceUpdate) {
+    return getAssumedSimplified(IRP, &AA, UsedAssumedInformation, ForceUpdate);
+  }
+  Optional<Value *> getAssumedSimplified(const Value &V,
+                                         const AbstractAttribute &AA,
+                                         bool &UsedAssumedInformation,
+                                         bool &ForceUpdate) {
+    return getAssumedSimplified(IRPosition::value(V), AA,
+                                UsedAssumedInformation, ForceUpdate);
+  }
+  Optional<Value *> getAssumedSimplified(const IRPosition &IRP,
+                                         const AbstractAttribute &AA,
                                          bool &UsedAssumedInformation) {
-    return getAssumedSimplified(IRP, &AA, UsedAssumedInformation);
+    bool ForceUpdate = false;
+    return getAssumedSimplified(IRP, &AA, UsedAssumedInformation, ForceUpdate);
   }
   Optional<Value *> getAssumedSimplified(const Value &V,
                                          const AbstractAttribute &AA,
                                          bool &UsedAssumedInformation) {
+    bool ForceUpdate = false;
     return getAssumedSimplified(IRPosition::value(V), AA,
-                                UsedAssumedInformation);
+                                UsedAssumedInformation, ForceUpdate);
   }
 
   /// Register \p CB as a simplification callback.
   /// `Attributor::getAssumedSimplified` will use these callbacks before
   /// we it will ask `AAValueSimplify`. It is important to ensure this
   /// is called before `identifyDefaultAbstractAttributes`, assuming the
   /// latter is called at all.
   using SimplifictionCallbackTy = std::function<Optional<Value *>(
       const IRPosition &, const AbstractAttribute *, bool &)>;
   void registerSimplificationCallback(const IRPosition &IRP,
                                       const SimplifictionCallbackTy &CB) {
+    assert(SimplificationCallbacks[IRP].size() == 0 &&
+           "Expected a single simplification callback for a given IRP");
     SimplificationCallbacks[IRP].emplace_back(CB);
   }
 
 private:
   /// The vector with all simplification callbacks registered by outside AAs.
   DenseMap<IRPosition, SmallVector<SimplifictionCallbackTy, 1>>
       SimplificationCallbacks;
 
   /// If \p V is assumed simplified, return it, if it is unclear yet,
   /// return None, otherwise return `nullptr`. Same as the public version
   /// except that it can be used without recording dependences on any \p AA.
   Optional<Value *> getAssumedSimplified(const IRPosition &V,
                                          const AbstractAttribute *AA,
-                                         bool &UsedAssumedInformation);
+                                         bool &UsedAssumedInformation,
+                                         bool &ForceUpdate);
+  Optional<Value *> getAssumedSimplified(const IRPosition &V,
+                                         const AbstractAttribute *AA,
+                                         bool &UsedAssumedInformation) {
+    bool ForceUpdate = false;
+    return getAssumedSimplified(V, AA, UsedAssumedInformation, ForceUpdate);
+  }
 
 public:
   /// Translate \p V from the callee context into the call site context.
   Optional<Value *>
   translateArgumentToCallSiteContent(Optional<Value *> V, CallBase &CB,
                                      const AbstractAttribute &AA,
                                      bool &UsedAssumedInformation);
 

llvm/lib/Transforms/IPO/Attributor.cpp

Attributor::getAssumedConstant(const Value &V, const AbstractAttribute &AA,
   // First check all callbacks provided by outside AAs. If any of them returns
   // a non-null value that is different from the associated value, or None, we
   // assume it's simpliied.
   IRPosition IRP = IRPosition::value(V, AA.getCallBaseContext());
   for (auto &CB : SimplificationCallbacks[IRP]) {
     Optional<Value *> SimplifiedV = CB(IRP, &AA, UsedAssumedInformation);
     if (!SimplifiedV.hasValue())
       return llvm::None;
-    if (*SimplifiedV && *SimplifiedV != &IRP.getAssociatedValue() &&
-        isa<Constant>(*SimplifiedV))
+    if (*SimplifiedV && isa<Constant>(*SimplifiedV))
       return cast<Constant>(*SimplifiedV);
   }
   const auto &ValueSimplifyAA =
       getAAFor<AAValueSimplify>(AA, IRP, DepClassTy::NONE);
   Optional<Value *> SimplifiedV =
       ValueSimplifyAA.getAssumedSimplifiedValue(*this);
   bool IsKnown = ValueSimplifyAA.isAtFixpoint();
   UsedAssumedInformation |= !IsKnown;
   if (!SimplifiedV.hasValue()) {
if (CI && CI->getType() != V.getType()) {
     // TODO: Check for a save conversion.
     return nullptr;
   }
   if (CI)
     recordDependence(ValueSimplifyAA, AA, DepClassTy::OPTIONAL);
   return CI;
 }
 
-Optional<Value *>
-Attributor::getAssumedSimplified(const IRPosition &IRP,
+Optional<Value *> Attributor::getAssumedSimplified(const IRPosition &IRP,
                                                    const AbstractAttribute *AA,
-                                 bool &UsedAssumedInformation) {
+                                                   bool &UsedAssumedInformation,
+                                                   bool &ForceUpdate) {
   // First check all callbacks provided by outside AAs. If any of them returns
   // a non-null value that is different from the associated value, or None, we
   // assume it's simpliied.
   for (auto &CB : SimplificationCallbacks[IRP]) {
     Optional<Value *> SimplifiedV = CB(IRP, AA, UsedAssumedInformation);
-    if (!SimplifiedV.hasValue() ||
-        (*SimplifiedV && *SimplifiedV != &IRP.getAssociatedValue()))
+    if (!SimplifiedV.hasValue())
       return SimplifiedV;
+
+    if (*SimplifiedV) {
+      ForceUpdate = true;
+      return SimplifiedV;
+    }
   }
 
   // If no high-level/outside simplification occured, use AAValueSimplify.
   const auto &ValueSimplifyAA =
       getOrCreateAAFor<AAValueSimplify>(IRP, AA, DepClassTy::NONE);
   Optional<Value *> SimplifiedV =
       ValueSimplifyAA.getAssumedSimplifiedValue(*this);
   bool IsKnown = ValueSimplifyAA.isAtFixpoint();

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

Value *getReplacementValue(Attributor &A) const {
     return NewV;
   }
 
   /// Helper function for querying AAValueSimplify and updating candicate.
   /// \param IRP The value position we are trying to unify with SimplifiedValue
   bool checkAndUpdate(Attributor &A, const AbstractAttribute &QueryingAA,
                       const IRPosition &IRP, bool Simplify = true) {
     bool UsedAssumedInformation = false;
+    bool ForceUpdate = false;
     Optional<Value *> QueryingValueSimplified = &IRP.getAssociatedValue();
     if (Simplify)
-      QueryingValueSimplified =
-          A.getAssumedSimplified(IRP, QueryingAA, UsedAssumedInformation);
+      QueryingValueSimplified = A.getAssumedSimplified(
+          IRP, QueryingAA, UsedAssumedInformation, ForceUpdate);
+
+    if (ForceUpdate) {
+      SimplifiedAssociatedValue = QueryingValueSimplified;
+      return true;
+    }
+
     return unionAssumed(QueryingValueSimplified);
   }
 
   /// Returns a candidate is found or not
   template <typename AAType> bool askSimplifiedValueFor(Attributor &A) {
     if (!getAssociatedValue().getType()->isIntegerTy())
       return false;