140 changes: 140 additions & 0 deletions llvm/lib/Transforms/IPO/AttributorAttributes.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@

#include "llvm/ADT/APInt.h"
#include "llvm/ADT/SCCIterator.h"
#include "llvm/ADT/SetOperations.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/Analysis/AliasAnalysis.h"
Expand All @@ -28,6 +29,7 @@
#include "llvm/Analysis/ScalarEvolution.h"
#include "llvm/Analysis/TargetTransformInfo.h"
#include "llvm/Analysis/ValueTracking.h"
#include "llvm/IR/Assumptions.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/Instruction.h"
Expand Down Expand Up @@ -146,6 +148,7 @@ PIPE_OPERATOR(AANoUndef)
PIPE_OPERATOR(AACallEdges)
PIPE_OPERATOR(AAFunctionReachability)
PIPE_OPERATOR(AAPointerInfo)
PIPE_OPERATOR(AAAssumptionInfo)

#undef PIPE_OPERATOR

Expand Down Expand Up @@ -9626,6 +9629,7 @@ struct AAFunctionReachabilityFunction : public AAFunctionReachability {
}

void trackStatistics() const override {}

private:
bool canReachUnknownCallee() const override {
return WholeFunction.CanReachUnknownCallee;
Expand All @@ -9639,6 +9643,140 @@ struct AAFunctionReachabilityFunction : public AAFunctionReachability {
DenseMap<CallBase *, QuerySet> CBQueries;
};

/// ---------------------- Assumption Propagation ------------------------------
struct AAAssumptionInfoImpl : public AAAssumptionInfo {
AAAssumptionInfoImpl(const IRPosition &IRP, Attributor &A,
const DenseSet<StringRef> &Known)
: AAAssumptionInfo(IRP, A, Known) {}

bool hasAssumption(const StringRef Assumption) const override {
return isValidState() && setContains(Assumption);
}

/// See AbstractAttribute::getAsStr()
const std::string getAsStr() const override {
const SetContents &Known = getKnown();
const SetContents &Assumed = getAssumed();

const std::string KnownStr =
llvm::join(Known.getSet().begin(), Known.getSet().end(), ",");
const std::string AssumedStr =
(Assumed.isUniversal())
? "Universal"
: llvm::join(Assumed.getSet().begin(), Assumed.getSet().end(), ",");

return "Known [" + KnownStr + "]," + " Assumed [" + AssumedStr + "]";
}
};

/// Propagates assumption information from parent functions to all of their
/// successors. An assumption can be propagated if the containing function
/// dominates the called function.
///
/// We start with a "known" set of assumptions already valid for the associated
/// function and an "assumed" set that initially contains all possible
/// assumptions. The assumed set is inter-procedurally updated by narrowing its
/// contents as concrete values are known. The concrete values are seeded by the
/// first nodes that are either entries into the call graph, or contains no
/// assumptions. Each node is updated as the intersection of the assumed state
/// with all of its predecessors.
struct AAAssumptionInfoFunction final : AAAssumptionInfoImpl {
AAAssumptionInfoFunction(const IRPosition &IRP, Attributor &A)
: AAAssumptionInfoImpl(IRP, A,
getAssumptions(*IRP.getAssociatedFunction())) {}

/// See AbstractAttribute::manifest(...).
ChangeStatus manifest(Attributor &A) override {
const auto &Assumptions = getKnown();

// Don't manifest a universal set if it somehow made it here.
if (Assumptions.isUniversal())
return ChangeStatus::UNCHANGED;

Function *AssociatedFunction = getAssociatedFunction();

bool Changed = addAssumptions(*AssociatedFunction, Assumptions.getSet());

return Changed ? ChangeStatus::CHANGED : ChangeStatus::UNCHANGED;
}

/// See AbstractAttribute::updateImpl(...).
ChangeStatus updateImpl(Attributor &A) override {
bool Changed = false;

auto CallSitePred = [&](AbstractCallSite ACS) {
const auto &AssumptionAA = A.getAAFor<AAAssumptionInfo>(
*this, IRPosition::callsite_function(*ACS.getInstruction()),
DepClassTy::REQUIRED);
// Get the set of assumptions shared by all of this function's callers.
Changed |= getIntersection(AssumptionAA.getAssumed());
return !getAssumed().empty() || !getKnown().empty();
};

bool AllCallSitesKnown;
// Get the intersection of all assumptions held by this node's predecessors.
// If we don't know all the call sites then this is either an entry into the
// call graph or an empty node. This node is known to only contain its own
// assumptions and can be propagated to its successors.
if (!A.checkForAllCallSites(CallSitePred, *this, true, AllCallSitesKnown))
return indicatePessimisticFixpoint();

return Changed ? ChangeStatus::CHANGED : ChangeStatus::UNCHANGED;
}

void trackStatistics() const override {}
};

/// Assumption Info defined for call sites.
struct AAAssumptionInfoCallSite final : AAAssumptionInfoImpl {

AAAssumptionInfoCallSite(const IRPosition &IRP, Attributor &A)
: AAAssumptionInfoImpl(IRP, A, getInitialAssumptions(IRP)) {}

/// See AbstractAttribute::initialize(...).
void initialize(Attributor &A) override {
const IRPosition &FnPos = IRPosition::function(*getAnchorScope());
A.getAAFor<AAAssumptionInfo>(*this, FnPos, DepClassTy::REQUIRED);
}

/// See AbstractAttribute::manifest(...).
ChangeStatus manifest(Attributor &A) override {
// Don't manifest a universal set if it somehow made it here.
if (getKnown().isUniversal())
return ChangeStatus::UNCHANGED;

CallBase &AssociatedCall = cast<CallBase>(getAssociatedValue());
bool Changed = addAssumptions(AssociatedCall, getAssumed().getSet());

return Changed ? ChangeStatus::CHANGED : ChangeStatus::UNCHANGED;
}

/// See AbstractAttribute::updateImpl(...).
ChangeStatus updateImpl(Attributor &A) override {
const IRPosition &FnPos = IRPosition::function(*getAnchorScope());
auto &AssumptionAA =
A.getAAFor<AAAssumptionInfo>(*this, FnPos, DepClassTy::REQUIRED);
bool Changed = getIntersection(AssumptionAA.getAssumed());
return Changed ? ChangeStatus::CHANGED : ChangeStatus::UNCHANGED;
}

/// See AbstractAttribute::trackStatistics()
void trackStatistics() const override {}

private:
/// Helper to initialized the known set as all the assumptions this call and
/// the callee contain.
DenseSet<StringRef> getInitialAssumptions(const IRPosition &IRP) {
const CallBase &CB = cast<CallBase>(IRP.getAssociatedValue());
auto Assumptions = getAssumptions(CB);
if (Function *F = IRP.getAssociatedFunction())
set_union(Assumptions, getAssumptions(*F));
if (Function *F = IRP.getAssociatedFunction())
set_union(Assumptions, getAssumptions(*F));
return Assumptions;
}
};

} // namespace

AACallGraphNode *AACallEdgeIterator::operator*() const {
Expand Down Expand Up @@ -9674,6 +9812,7 @@ const char AANoUndef::ID = 0;
const char AACallEdges::ID = 0;
const char AAFunctionReachability::ID = 0;
const char AAPointerInfo::ID = 0;
const char AAAssumptionInfo::ID = 0;

// Macro magic to create the static generator function for attributes that
// follow the naming scheme.
Expand Down Expand Up @@ -9776,6 +9915,7 @@ CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANoReturn)
CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAReturnedValues)
CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAMemoryLocation)
CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AACallEdges)
CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAAssumptionInfo)

CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANonNull)
CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANoAlias)
Expand Down
95 changes: 95 additions & 0 deletions llvm/test/Transforms/Attributor/assumes_info.ll
Original file line number Diff line number Diff line change
@@ -0,0 +1,95 @@
; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --function-signature --check-attributes --check-globals
; RUN: opt -attributor -enable-new-pm=0 -attributor-manifest-internal -attributor-max-iterations-verify -attributor-annotate-decl-cs -attributor-max-iterations=3 -S < %s | FileCheck %s --check-prefixes=CHECK,NOT_CGSCC_NPM,NOT_CGSCC_OPM,NOT_TUNIT_NPM,IS__TUNIT____,IS________OPM,IS__TUNIT_OPM
; RUN: opt -aa-pipeline=basic-aa -passes=attributor -attributor-manifest-internal -attributor-max-iterations-verify -attributor-annotate-decl-cs -attributor-max-iterations=3 -S < %s | FileCheck %s --check-prefixes=CHECK,NOT_CGSCC_OPM,NOT_CGSCC_NPM,NOT_TUNIT_OPM,IS__TUNIT____,IS________NPM,IS__TUNIT_NPM
; RUN: opt -attributor-cgscc -enable-new-pm=0 -attributor-manifest-internal -attributor-annotate-decl-cs -S < %s | FileCheck %s --check-prefixes=CHECK,NOT_TUNIT_NPM,NOT_TUNIT_OPM,NOT_CGSCC_NPM,IS__CGSCC____,IS________OPM,IS__CGSCC_OPM
; RUN: opt -aa-pipeline=basic-aa -passes=attributor-cgscc -attributor-manifest-internal -attributor-annotate-decl-cs -S < %s | FileCheck %s --check-prefixes=CHECK,NOT_TUNIT_NPM,NOT_TUNIT_OPM,NOT_CGSCC_OPM,IS__CGSCC____,IS________NPM,IS__CGSCC_NPM

define dso_local void @entry(i1 %cond) #0 {
; CHECK-LABEL: define {{[^@]+}}@entry
; CHECK-SAME: (i1 [[COND:%.*]]) #[[ATTR0:[0-9]+]] {
; CHECK-NEXT: entry:
; CHECK-NEXT: call void @foo(i1 [[COND]]) #[[ATTR1:[0-9]+]]
; CHECK-NEXT: call void @bar() #[[ATTR2:[0-9]+]]
; CHECK-NEXT: call void @qux() #[[ATTR1]]
; CHECK-NEXT: ret void
;
entry:
call void @foo(i1 %cond)
call void @bar()
call void @qux() #1
ret void
}

define internal void @foo(i1 %cond) #1 {
; CHECK-LABEL: define {{[^@]+}}@foo
; CHECK-SAME: (i1 [[COND:%.*]]) #[[ATTR1]] {
; CHECK-NEXT: entry:
; CHECK-NEXT: call void @baz(i1 [[COND]]) #[[ATTR1]]
; CHECK-NEXT: ret void
;
entry:
call void @baz(i1 %cond)
ret void
}

define internal void @bar() #2 {
; CHECK-LABEL: define {{[^@]+}}@bar
; CHECK-SAME: () #[[ATTR2]] {
; CHECK-NEXT: entry:
; CHECK-NEXT: call void @baz(i1 noundef false) #[[ATTR2]]
; CHECK-NEXT: ret void
;
entry:
call void @baz(i1 0)
ret void
}

define internal void @baz(i1 %Cond) {
; CHECK-LABEL: define {{[^@]+}}@baz
; CHECK-SAME: (i1 [[COND:%.*]]) #[[ATTR1]] {
; CHECK-NEXT: entry:
; CHECK-NEXT: [[TOBOOL:%.*]] = icmp ne i1 [[COND]], false
; CHECK-NEXT: br i1 [[TOBOOL]], label [[IF_THEN:%.*]], label [[IF_END:%.*]]
; CHECK: if.then:
; CHECK-NEXT: call void @baz(i1 noundef false) #[[ATTR1]]
; CHECK-NEXT: br label [[IF_END]]
; CHECK: if.end:
; CHECK-NEXT: call void @qux() #[[ATTR1]]
; CHECK-NEXT: ret void
;
entry:
%tobool = icmp ne i1 %Cond, 0
br i1 %tobool, label %if.then, label %if.end

if.then:
call void @baz(i1 0)
br label %if.end

if.end:
call void @qux()
ret void
}

define internal void @qux() {
; CHECK-LABEL: define {{[^@]+}}@qux
; CHECK-SAME: () #[[ATTR1]] {
; CHECK-NEXT: entry:
; CHECK-NEXT: call void @call() #[[ATTR2]]
; CHECK-NEXT: ret void
;
entry:
call void @call()
ret void
}

declare void @call() #3

attributes #0 = { "llvm.assume"="A" }
attributes #1 = { "llvm.assume"="B" }
attributes #2 = { "llvm.assume"="B,C" }
attributes #3 = { "llvm.assume"="B,C,A" }
;.
; CHECK: attributes #[[ATTR0]] = { "llvm.assume"="A" }
; CHECK: attributes #[[ATTR1]] = { "llvm.assume"="B,A" }
; CHECK: attributes #[[ATTR2]] = { "llvm.assume"="B,C,A" }
;.
4 changes: 4 additions & 0 deletions llvm/test/Transforms/Attributor/depgraph.ll
Original file line number Diff line number Diff line change
Expand Up @@ -123,6 +123,8 @@ define i32* @checkAndAdvance(i32* align 16 %0) {
; GRAPH-NEXT: updates [AAMemoryLocation] for CtxI ' %6 = call i32* @checkAndAdvance(i32* %5)' at position {cs: [@-1]} with state memory:argument
; GRAPH-NEXT: updates [AAMemoryLocation] for CtxI ' %6 = call i32* @checkAndAdvance(i32* %5)' at position {cs: [@-1]} with state memory:argument
; GRAPH-EMPTY:
; GRAPH-NEXT: [AAAssumptionInfo] for CtxI ' %2 = load i32, i32* %0, align 4' at position {fn:checkAndAdvance [checkAndAdvance@-1]} with state Known [], Assumed []
; GRAPH-EMPTY:
; GRAPH-NEXT: [AAHeapToStack] for CtxI ' %2 = load i32, i32* %0, align 4' at position {fn:checkAndAdvance [checkAndAdvance@-1]} with state [H2S] Mallocs Good/Bad: 0/0
; GRAPH-EMPTY:
; GRAPH-NEXT: [AAValueSimplify] for CtxI ' %2 = load i32, i32* %0, align 4' at position {fn_ret:checkAndAdvance [checkAndAdvance@-1]} with state not-simple
Expand Down Expand Up @@ -181,6 +183,8 @@ define i32* @checkAndAdvance(i32* align 16 %0) {
; GRAPH-NEXT: updates [AAIsDead] for CtxI ' %6 = call i32* @checkAndAdvance(i32* %5)' at position {cs_ret: [@-1]} with state assumed-live
; GRAPH-NEXT: updates [AAMemoryBehavior] for CtxI ' %2 = load i32, i32* %0, align 4' at position {fn:checkAndAdvance [checkAndAdvance@-1]} with state readonly
; GRAPH-EMPTY:
; GRAPH-NEXT: [AAAssumptionInfo] for CtxI ' %6 = call i32* @checkAndAdvance(i32* %5)' at position {cs: [@-1]} with state Known [], Assumed []
; GRAPH-EMPTY:
; GRAPH-NEXT: [AAValueSimplify] for CtxI ' %6 = call i32* @checkAndAdvance(i32* %5)' at position {cs_ret: [@-1]} with state not-simple
; GRAPH-EMPTY:
; GRAPH-NEXT: [AAIsDead] for CtxI ' %6 = call i32* @checkAndAdvance(i32* %5)' at position {cs_arg: [@0]} with state assumed-live
Expand Down