[Attributor] Deduce "align" attribute

[uenoku]

Summary

Deduce "align" attribute in attributor.

[uenoku]

Fix patch

[uenoku]

@jdoerfert
I know align attribute is integer attribute and we need somewhat different deduction with other attributes. However, I implement AAAlignImpl as a subclass of BooleanState for now. Is it better to add some base class for the integer attribute?

llvm/test/Transforms/FunctionAttrs/align.ll

36

I want to test this for

if (ICS && ICS.hasRetAttr(Attribute::Alignment))
  return CompareOrSet(ICS.getRetAlignment());

but it does not work well now.
I find that ICS.hasRetAttr(Attribute::Alignment) returns true but ICS.getRetAlignment() returns 0 and get error in the current implementation. I think maybe that this is parse level problem... I'll investigate more.

[jdoerfert]

In D64152#1569030, @uenoku wrote:

@jdoerfert
I know align attribute is integer attribute and we need somewhat different deduction with other attributes. However, I implement AAAlignImpl as a subclass of BooleanState for now. Is it better to add some base class for the integer attribute?

Why not make it derive from IntegerState and use the integer as the alignment, zero would be the worst possible = unknown.

[jdoerfert]

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

696

Please add getters for the assumed and known alignment here.

llvm/lib/Transforms/IPO/Attributor.cpp

803

Can you unify the two initialize methods in the base class (AAAlignImpl) with the index trick used in addIfNotExistent?

827

What you want to do use to take the minimum of all align values and the current state (IntegerState). Remember the state at the beginning of this function to determine if it changed. We also need a test for this.

839

For the return value you should do two things (for now):

1. Use Value::getPointerAlignment to get the known alignment in the IR.
2. Use getAAFor<AAAlign>() to the the known and assumed alignment from the Attributor.

Later, 1) would be replaced by a new ValueTracking.h getPointerAlignment method that does traversal, e.g., look through casts, compute new alignment based on constant geps, etc.
Later later, the new 1) and 2) would be merged to use traversal and assumed alignment together.

You can work on the later points as well if you want.

859

There is (probably) no need for the ExistsAssume trick here, you can leave it but add a comment to explain that it is a shortcut to identify an optimistic fixpoint.

llvm/test/Transforms/FunctionAttrs/align.ll

23

We want align 4 here.

[uenoku]

Use IntegerState and take a minimum of alignment.

[uenoku]

Use getPointerAlignment.

[jdoerfert]

llvm/lib/Transforms/IPO/Attributor.cpp

257

I think you want to return -1 so it is clearly not a valid argument position.

772

return value or argument, right?

776

One could use the existing attribute as "known" and not indicate a fixpoint so that we can compute a better value.
We should have a test for it even if you keep it this way to describe the possibility.

786

Why the if? Can't you take the minimum of the knownAlign and the maximum allowed? Or initialize the assumed to the maximum allowed value already in the initialization?

787

style: remove braces around single statements.

790

I think you want to first look at the in-flight AAAlign attribute before you look at the IR.

838

If there is no top state, we can reach optimistic fixpoint earlier.

"If no assumed state was used for reasoning, an optimistic fixpoint is reached earlier."

llvm/test/Transforms/FunctionAttrs/align.ll

43

Can we have SCC tests as well

[uenoku]

Address comments.

[uenoku]

Rebase and add callsite argument and argument deduction. It is ready for review.

[uenoku]

Minor update

[jdoerfert]

llvm/lib/Transforms/IPO/Attributor.cpp

866

Are the statistics missing?

870

Make the 1U << 29 a static const member with a proper name and description.

874

The state is not initialized here as above. Call the general constructor from the specialized one (above) to avoid such inconsistencies.

912

I think we should, for consistency, call getAssumedAlign() here.

1019

remove isPointerTy, see above.

llvm/test/Transforms/FunctionAttrs/align.ll

44

What happened to test5 again? Why is it commented out?

58

Please add a more meaningful SCC test, 3 functions, conditionals, malloc, etc. This one is a no-return recursion for which we can deduce everything.

76

check the call sites.

[uenoku]

Address comments

[uenoku]

llvm/test/Transforms/FunctionAttrs/align.ll

44

The problem was that getPointerAlignment didn't work well in test5_2. This will be fixed in D65281.

[jdoerfert]

LGTM.

[uenoku]

Minor fix.

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

struct IntegerState : public AbstractState {
 
   /// Keep only "assumed bits" also set in \p BitsEncoding but all known ones.
   IntegerState &intersectAssumedBits(base_t BitsEncoding) {
     // Make sure we never loose any "known bits".
     Assumed = (Assumed & BitsEncoding) | Known;
     return *this;
   }
 
+  /// Take minimum of assumed and \p Value.
+  IntegerState &takeAssumedMinimum(base_t Value) {
+    // Make sure we never loose "known value".
+    Assumed = std::max(std::min(Assumed, Value), Known);
+    return *this;
+  }
+
 private:
   /// The known state encoding in an integer of type base_t.
   base_t Known = getWorstState();
 
   /// The assumed state encoding in an integer of type base_t.
   base_t Assumed;
 };
 
struct AAReturnedValues : public AbstractAttribute {
 
   /// See AbstractAttribute::getAttrKind()
   virtual Attribute::AttrKind getAttrKind() const override { return ID; }
 
   /// The identifier used by the Attributor for this class of attributes.
   static constexpr Attribute::AttrKind ID = Attribute::Returned;
 };
 
+/// An abstract interface for all align<n> attributes.
+struct AAAlign : public AbstractAttribute {
+
+  /// See AbstractAttribute::AbstractAttribute(...).
+  AAAlign(Value &V, InformationCache &InfoCache)
+      : AbstractAttribute(V, InfoCache) {}
+
+  /// See AbastractState::getAttrKind().
+  Attribute::AttrKind getAttrKind() const override { return ID; }
+
+  /// Return assumed alignment.
+  virtual unsigned getAssumedAlign() const = 0;
+
+  /// Return known alignemnt.
+  virtual unsigned getKnownAlign() const = 0;
+
+  /// The identifier used by the Attributor for this class of attributes.
+  static constexpr Attribute::AttrKind ID = Attribute::Alignment;
+};
+
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_IPO_FUNCTIONATTRS_H

llvm/lib/Transforms/IPO/Attributor.cpp

if (Attr.isStringAttribute()) {
     StringRef Kind = Attr.getKindAsString();
     if (Attrs.hasAttribute(AttrIdx, Kind))
       if (isEqualOrWorse(Attr, Attrs.getAttribute(AttrIdx, Kind)))
         return false;
     Attrs = Attrs.addAttribute(Ctx, AttrIdx, Attr);
     return true;
   }
 
-  llvm_unreachable("Expected enum or string attribute!");
+  if (Attr.isIntAttribute()) {
+    Attribute::AttrKind Kind = Attr.getKindAsEnum();
+    if (Attrs.hasAttribute(AttrIdx, Kind))
+      if (isEqualOrWorse(Attr, Attrs.getAttribute(AttrIdx, Kind)))
+        return false;
+    Attrs = Attrs.addAttribute(Ctx, AttrIdx, Attr);
+    return true;
+  }
+
+  llvm_unreachable("Expected enum, string or int attribute!");
+}
+
+// Helper function that returns argument index of value.
+// If the value is not an argument, this returns 0.
+static int getArgNo(Value &V) {
+  if (auto *Arg = dyn_cast<Argument>(&V))
+    return Arg->getArgNo();
+  return 0;
 }
 
 ChangeStatus AbstractAttribute::update(Attributor &A) {
   ChangeStatus HasChanged = ChangeStatus::UNCHANGED;
   if (getState().isAtFixpoint())
     return HasChanged;
 
   LLVM_DEBUG(dbgs() << "[Attributor] Update: " << *this << "\n");
ChangeStatus AAReturnedValuesImpl::updateImpl(Attributor &A) {
   if (!HasCallSite) {
     indicateOptimisticFixpoint();
     return ChangeStatus::CHANGED;
   }
 
   return Changed;
 }
 
+/// ------------------------ Align Argument Attribute ------------------------
+
+struct AAAlignImpl : AAAlign, IntegerState {
+  AAAlignImpl(Value &V, InformationCache &InfoCache) : AAAlign(V, InfoCache) {}
+
+  /// See AbstractAttribute::getState()
+  /// {
+  AbstractState &getState() override { return *this; }
+  const AbstractState &getState() const override { return *this; }
+  /// }
+
+  virtual const std::string getAsStr() const override {
+    return getAssumedAlign()
+               ? (getKnownAlign()
+                      ? "align<" + std::to_string(getKnownAlign()) + ">"
+                      : "maybe-align<" + std::to_string(getAssumedAlign()) +
+                            ">")
+               : "unknown-align";
+  }
+
+  /// See AAAlign::getAssumedAlign().
+  unsigned getAssumedAlign() const override { return getAssumed(); }
+
+  /// See AAAlign::getKnownAlign().
+  unsigned getKnownAlign() const override { return getKnown(); }
+
+  /// See AbstractAttriubute::initialize(...).
+  void initialize(Attributor &A) override {
+    Function &F = getAnchorScope();
+
+    unsigned AttrIdx =
+        getAttrIndex(getManifestPosition(), getArgNo(getAnchoredValue()));
+
+    // Already the function has align attribute on return value.
+    if (F.getAttributes().hasAttribute(AttrIdx, Attribute::Alignment)) {
+      takeAssumedMinimum(
+          F.getAttribute(AttrIdx, Attribute::Alignment).getAlignment());
+      indicateOptimisticFixpoint();
+    }
+  }
+
+  /// See AbstractAttribute::getDeducedAttributes
+  virtual void
+  getDeducedAttributes(SmallVectorImpl<Attribute> &Attrs) const override {
+    LLVMContext &Ctx = AnchoredVal.getContext();
+
+    if (getKnownAlign() <= (1U << 29)) {
+      Attrs.emplace_back(Attribute::getWithAlignment(Ctx, getKnownAlign()));
+    }
+  }
+};
+
+/// Align attribute for function return value.
+struct AAAlignReturned : AAAlignImpl {
+
+  AAAlignReturned(Function &F, InformationCache &InfoCache)
+      : AAAlignImpl(F, InfoCache) {}
+
+  /// See AbstractAttribute::getManifestPosition().
+  virtual ManifestPosition getManifestPosition() const override {
+    return MP_RETURNED;
+  }
+
+  /// See AbstractAttribute::updateImpl(...).
+  virtual ChangeStatus updateImpl(Attributor &A) override;
+};
+
+/// Align attribute for function argument.
+struct AAAlignArgument : AAAlignImpl {
+
+  AAAlignArgument(Argument &A, InformationCache &InfoCache)
+      : AAAlignImpl(A, InfoCache) {}
+
+  /// See AbstractAttribute::getManifestPosition().
+  virtual ManifestPosition getManifestPosition() const override {
+    return MP_ARGUMENT;
+  }
+
+  /// See AbstractAttribute::updateImpl(...).
+  virtual ChangeStatus updateImpl(Attributor &A) override;
+};
+
+ChangeStatus AAAlignArgument::updateImpl(Attributor &A) {
+  // TODO: Currently, we have no deduction for argument align attribute.
+  indicatePessimisticFixpoint();
+  return ChangeStatus::CHANGED;
+}
+
+ChangeStatus AAAlignReturned::updateImpl(Attributor &A) {
+  Function &F = getAnchorScope();
+  auto *AARetValImpl = A.getAAFor<AAReturnedValuesImpl>(*this, F);
+  if (!AARetValImpl) {
+    indicatePessimisticFixpoint();
+    return ChangeStatus::CHANGED;
+  }
+
+  // Currently, align<n> is deduced if alignments in return values are known as
+  // greater than n. We reach pessimistic fixpoint if any of the return value
+  // wouldn't have align. If there is no top state, we can reach optimistic
+  // fixpoint earlier.
 
+  base_t BeforeState = getAssumed();
+  std::function<bool(Value &)> Pred = [&](Value &RV) -> bool {
+    auto *AlignAA = A.getAAFor<AAAlignImpl>(*this, RV);
+
+    if (AlignAA && AlignAA->isValidState()) {
+      takeAssumedMinimum(AlignAA->getAssumedAlign());
+      return true;
+    }
+    ImmutableCallSite ICS(&RV);
+
+    if (ICS && ICS.hasRetAttr(Attribute::Alignment)) {
+      takeAssumedMinimum(ICS.getRetAlignment());
+      return true;
+    }
+
+    return false;
+  };
+  if (!AARetValImpl->checkForallReturnedValues(Pred)) {
+    indicatePessimisticFixpoint();
+    return ChangeStatus::CHANGED;
+  }
+
+  return (getAssumed() != BeforeState) ? ChangeStatus::CHANGED
+                                       : ChangeStatus::UNCHANGED;
+}
 
 /// ----------------------------------------------------------------------------
 ///                               Attributor
 /// ----------------------------------------------------------------------------
 
 ChangeStatus Attributor::run() {
   // Initialize all abstract attributes.
   for (AbstractAttribute *AA : AllAbstractAttributes)
     AA->initialize(*this);
 
   LLVM_DEBUG(dbgs() << "[Attributor] Identified and initialized "
                     << AllAbstractAttributes.size()
                     << " abstract attributes.\n");
 
   // Now that all abstract attributes are collected and initialized we start
   // the abstract analysis.
 
   unsigned IterationCounter = 1;
do {
     for (AbstractAttribute *AA : Worklist)
       if (AA->update(*this) == ChangeStatus::CHANGED)
         ChangedAAs.push_back(AA);
 
     // Reset the work list and repopulate with the changed abstract attributes.
     // Note that dependent ones are added above.
     Worklist.clear();
     Worklist.insert(ChangedAAs.begin(), ChangedAAs.end());
 
   } while (!Worklist.empty() && ++IterationCounter < MaxFixpointIterations);
 
   LLVM_DEBUG(dbgs() << "\n[Attributor] Fixpoint iteration done after: "
                     << IterationCounter << "/" << MaxFixpointIterations
                     << " iterations\n");
 
   bool FinishedAtFixpoint = Worklist.empty();
 
void Attributor::identifyDefaultAbstractAttributes(
 
   // Return attributes are only appropriate if the return type is non void.
   Type *ReturnType = F.getReturnType();
   if (!ReturnType->isVoidTy()) {
     // Argument attribute "returned" --- Create only one per function even
     // though it is an argument attribute.
     if (!Whitelist || Whitelist->count(AAReturnedValues::ID))
       registerAA(*new AAReturnedValuesImpl(F, InfoCache));
+
+    // Every function with pointer return type might be marked align.
+    if (ReturnType->isPointerTy() &&
+        (!Whitelist || Whitelist->count(AAAlignReturned::ID)))
+      registerAA(*new AAAlignReturned(F, InfoCache));
   }
 
+  // Every argument with pointer type might be marked align.
+  for (Argument &Arg : F.args()) {
+    if (Arg.getType()->isPointerTy())
+      registerAA(*new AAAlignArgument(Arg, InfoCache));
+  }
   // Walk all instructions to find more attribute opportunities and also
   // interesting instructions that might be queried by abstract attributes
   // during their initialization or update.
   auto &ReadOrWriteInsts = InfoCache.FuncRWInstsMap[&F];
   auto &InstOpcodeMap = InfoCache.FuncInstOpcodeMap[&F];
 
   for (Instruction &I : instructions(&F)) {
     bool IsInterestingOpcode = false;

llvm/test/Transforms/FunctionAttrs/align.ll

+; RUN: opt -attributor -attributor-disable=false -S < %s | FileCheck %s --check-prefix=ATTRIBUTOR
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+
+; Test cases specifically designed for "align" attribute.
+; We use FIXME's to indicate problems and missing attributes.
+
+
+; TEST 1 
+; ATTRIBUTOR: define align 8 i32* @test1(i32* returned align 8)
+define i32* @test1(i32* align 8) #0 {
+  ret i32* %0
+}
+
+; TEST 2
+; ATTRIBUTOR: define i32* @test2(i32* returned)
+define i32* @test2(i32*) #0 {
+  ret i32* %0
+}
+
+; TEST 3
+; ATTRIBUTOR: define align 4 i32* @test3(i32* align 8, i32* align 4, i1)
+define i32* @test3(i32* align 8, i32* align 4, i1) #0 {
+  %ret = select i1 %2, i32* %0, i32* %1
+  ret i32* %ret
+}
+
+; TEST 4
+; ATTRIBUTOR: define align 32 i32* @test4(i32* align 32, i32* align 32, i1)
+define i32* @test4(i32* align 32, i32* align 32, i1) #0 {
+  %ret = select i1 %2, i32* %0, i32* %1
+  ret i32* %ret
+}
+
+; TEST 5
+; declare align 8 i32* @align8() #0  
+
+; define i32* @test5() {
+;  %ret = tail call i32* @align8()
+;  ret i32* %ret
+; }
+
+
+attributes #0 = { nounwind uwtable noinline }
+attributes #1 = { uwtable noinline }
+