Import UnitLoopInfo and UnitLICM

[?]

8Cqo1QjSCQ3F66Vh3nan9N8dv1MjQSAEz8RNwt2cTacK

May 2, 2024, 11:15 PM

56LYABQH3HNWTV4MWQ2IWW7LEIYKMOW4ZVTRXKOSQID4JTQ5A35AC

Dependencies

[2] 7OVOGA6Q Import LLVM @ 15.0.7

Change contents

file addition: UnitLICM.cpp (----------)

[2.509877339]

#include <cassert>
#include <utility>
#include <functional>
#include <iterator>
#include <vector>
#include <ranges>
#include "llvm/Support/raw_ostream.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/Casting.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/IR/Value.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/Instruction.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/PatternMatch.h"
#include "llvm/IR/Dominators.h"
#include "llvm/Analysis/AliasAnalysis.h"
#include "llvm/Analysis/AliasSetTracker.h"
#include "llvm/Analysis/ValueTracking.h"

#include "llvm/Transforms/UnitLICM.h"
#include "llvm/Analysis/UnitLoopInfo.h"

#define DEBUG_TYPE "unit-licm"
STATISTIC(NumInstsHoisted, "Number of instructions hoisted");
STATISTIC(NumStoresHoisted, "Number of stores hoisted");
STATISTIC(NumLoadsHoisted, "Number of loads hoisted");
STATISTIC(NumComputationalHoisted, "Number of computational instructions (excluding casts and GEPs) hoisted");
STATISTIC(NumVisited, "Number of instructions considered for hoisting");

using namespace llvm;
using namespace ece479k;

/// Since all symbols are exported from the SO, we try to mark all symbols local
/// to this translation unit to improve diagnostics and improve performance.
/// Likewise, we try to keep functions pure.
///
/// To reduce confusion when describing LICM, we classify two types of "invariants":
/// - Invariant: a Value that is already invariant to the loop - it is outside the loop or not an Instruction
/// - Hoistable: a Value that can be hoisted outside the loop to become an invariant
///
/// (Un)Hoistable detection is implemented as DFS on every use-def chain in the loop,
/// stopping if we end up at an invariant or condition that prevents hoisting,
/// propagating the result back up the chain to every Value visited.
///
/// This is implemented as a series of tail-recursive calls to check each
/// instruction for any disqualifying cases or an invariant and then recursively
/// on each of its operands (the part that makes it a DFS).
///
/// Optimization misses are documented in debug output. For meaningful names,
/// make sure to run the `instnamer` pass before running `unit-licm`.
/// Example output:
/// UnitLICM: Can't hoist instruction: Depends on unsupported phi `p.0`: %idxprom = sext i32 %p.0 to i64
/// UnitLICM: Can't hoist instruction: Depends on not-hoistable sext `idxprom`: %arrayidx18 = getelementptr inbounds [8 x [3 x double]], ptr %position, i64 0, i64 %idxprom
/// UnitLICM: Can't hoist instruction: Depends on not-hoistable getelementptr `arrayidx18`: %arraydecay19 = getelementptr inbounds [3 x double], ptr %arrayidx18, i64 0, i64 0

static bool is_supported(Instruction const&) __attribute((pure));
static bool is_computational(Instruction const&) __attribute((pure));

namespace {
struct HoistArgs {
UnitLoopInfo const& loop;
DenseMap<Instruction const*, bool> cache;
DominatorTree const& DT;
};
}

static bool can_hoist(Instruction const&, HoistArgs &) __attribute((pure));
static void hoist(Instruction&, HoistArgs &);

// ReSharper disable once CppMemberFunctionMayBeStatic
PreservedAnalyses UnitLICM::run(Function &F, FunctionAnalysisManager &FAM) // NOLINT(*-convert-member-functions-to-static)
{
for (auto const& loop: FAM.getResult<UnitLoopAnalysis>(F)) {

auto hoist_args = HoistArgs{
loop,
DenseMap<Instruction const*, bool>(32),
FAM.getResult<DominatorTreeAnalysis>(F)
};
hoist_args.cache.clear();

auto const is_hoistable = [&loop, &hoist_args](Instruction const& inst) {
return is_supported(inst) && can_hoist(inst, hoist_args);
};

// Hoist everything we can
for (auto i: loop.instructions(is_hoistable))
{
hoist(*i, hoist_args);

if (isa<LoadInst>(i))
++NumLoadsHoisted;
else if (isa<StoreInst>(i))
++NumStoresHoisted;
else if (is_computational(*i))
++NumComputationalHoisted;

++NumInstsHoisted;
}
}

// Invalidate analyses (we fucked the code up)
return PreservedAnalyses::none();
}

// This switch is inlined instead of chaining isXXX calls because it's easier to see what's going on
// and we get slightly better codegen (practically it doesn't matter)
static bool is_supported(Instruction const& inst)
{
switch (inst.getOpcode()) {
default:
return false;
// unary operator
case Instruction::FNeg:
// binary operator
case Instruction::Add:
case Instruction::FAdd:
case Instruction::Sub:
case Instruction::FSub:
case Instruction::Mul:
case Instruction::FMul:
case Instruction::UDiv:
case Instruction::SDiv:
case Instruction::FDiv:
case Instruction::URem:
case Instruction::SRem:
case Instruction::FRem:
// logical operator
case Instruction::Shl:
case Instruction::LShr:
case Instruction::AShr:
case Instruction::And:
case Instruction::Or:
case Instruction::Xor:
// memory operator
case Instruction::Load:
case Instruction::Store:
case Instruction::GetElementPtr:
// cast operator
case Instruction::Trunc:
case Instruction::ZExt:
case Instruction::SExt:
case Instruction::FPToUI:
case Instruction::FPToSI:
case Instruction::UIToFP:
case Instruction::SIToFP:
case Instruction::FPTrunc:
case Instruction::FPExt:
case Instruction::PtrToInt:
case Instruction::IntToPtr:
case Instruction::BitCast:
case Instruction::AddrSpaceCast:
// misc operator
case Instruction::Select:
case Instruction::ICmp:
case Instruction::FCmp:
return true;
}
}

static bool is_computational(Instruction const& inst)
{
switch (inst.getOpcode()) {
default:
return false;
// unary operator
case Instruction::FNeg:
// binary operator
case Instruction::Add:
case Instruction::FAdd:
case Instruction::Sub:
case Instruction::FSub:
case Instruction::Mul:
case Instruction::FMul:
case Instruction::UDiv:
case Instruction::SDiv:
case Instruction::FDiv:
case Instruction::URem:
case Instruction::SRem:
case Instruction::FRem:
// logical operator
case Instruction::Shl:
case Instruction::LShr:
case Instruction::AShr:
case Instruction::And:
case Instruction::Or:
case Instruction::Xor:
// misc operator
case Instruction::Select:
case Instruction::ICmp:
case Instruction::FCmp:
return true;
}
}

static void hoist(Instruction& inst, HoistArgs &args)
{
assert(!args.loop.is_invariant(inst) && "Cannot hoist variant");

auto* insertion_point = const_cast<Instruction *>(args.loop.preheader->getTerminator());
inst.moveBefore(insertion_point);

assert(args.loop.is_invariant(inst) && "Hoist did not make invariant");
}

// These are called in a tail-recursive (ish) manner
static bool can_hoist_impl(Instruction const&, HoistArgs &) __attribute((pure));
static bool can_hoist(Value const&, HoistArgs &) __attribute((pure));
static bool can_hoist_udiv(Instruction const&, HoistArgs &) __attribute((pure));
static bool can_hoist_sdiv(Instruction const&, HoistArgs &) __attribute((pure));
static bool can_hoist_store(Instruction const&, HoistArgs &) __attribute((pure));
static bool can_hoist_load(Instruction const&, HoistArgs &) __attribute((pure));
static bool can_hoist_load_impl(Instruction const&, HoistArgs &) __attribute((pure));
static bool can_hoist_sidefx(Instruction const&, HoistArgs&) __attribute((pure));
static bool can_hoist_operands(Instruction const&, HoistArgs &) __attribute((pure));

static bool can_hoist(Instruction const& inst, HoistArgs &args)
{
if (args.loop.is_invariant(inst))
return true;

if (!is_supported(inst)) {
LLVM_DEBUG(dbgs() << "UnitLICM: Can't hoist instruction: Depends on unsupported " << inst.getOpcodeName() << " `" << inst.getName() << "`: ");
return false;
}

auto const it = args.cache.find(&inst);
if (it != args.cache.end()) {
LLVM_DEBUG(if (!it->second)
dbgs() << "UnitLICM: Can't hoist instruction: Depends on not-hoistable " << inst.getOpcodeName() << " `" << inst.getName() << "`: ");

return it->second;
}

bool const dfs = can_hoist_impl(inst, args);
args.cache.insert({&inst, dfs});
return dfs;
}

static bool can_hoist(Value const& val, HoistArgs &args)
{
if (auto const* inst = dyn_cast<Instruction>(&val))
return can_hoist(*inst, args);
return true;
}

/// This is effectively a specialization of isSafeToSpeculativelyExecute for our case:
/// - We support stores
/// - We don't bail directly if an instruction has sidefx - we check if it dominates all exits to try and hoist it anyway
/// - We don't check for instructions we don't support
///
/// This will dispatch to special cases for instructions that can have side effects
/// or cause the use-def chain to be marked as not-hoistable. Otherwise, this
/// will directly dispatch to the next step of the DFS
static bool can_hoist_impl(Instruction const& inst, HoistArgs &args)
{
assert(is_supported(inst) && "Instruction is not supported");

++NumVisited;

switch (inst.getOpcode()) {
default: // Instruction doesn't have edge cases
return can_hoist_operands(inst, args);

case Instruction::Load: // Need to do alias analysis
return can_hoist_load(inst, args);

case Instruction::Store: // This one always has sidefx, and we need to do alias analysis
return can_hoist_store(inst, args);

case Instruction::UDiv:
case Instruction::URem: // Division by zero can throw
return can_hoist_udiv(inst, args);

case Instruction::SDiv:
case Instruction::SRem: // Division by zero or of INT_MIN can throw
return can_hoist_sdiv(inst, args);

}
}

static bool can_hoist_sdiv(Instruction const& inst, HoistArgs &args)
{
// code stolen from isSafeToSpeculativelyExecute
using namespace llvm::PatternMatch;
// x / y is undefined if y == 0 or x == INT_MIN and y == -1
const APInt *Numerator, *Denominator;
if (!match(inst.getOperand(1), m_APInt(Denominator))) {
// LLVM_DEBUG(dbgs() << "UnitLICM: Can't hoist " << inst.getOpcodeName() << ": Undefined denominator " << inst.getOperand(1) << ":" << inst << "\n");
return can_hoist_sidefx(inst, args);
}
// We cannot hoist this division if the denominator is 0.
if (*Denominator == 0) {
// LLVM_DEBUG(dbgs() << "UnitLICM: Can't hoist " << inst.getOpcodeName() << ": Division by zero:" << inst << "\n");
return can_hoist_sidefx(inst, args);
}
// It's safe to hoist if the denominator is not 0 or -1.
if (!Denominator->isAllOnes()) {
return can_hoist_operands(inst, args);
}
// At this point we know that the denominator is -1. It is safe to hoist as
// long we know that the numerator is not INT_MIN.
if (match(inst.getOperand(0), m_APInt(Numerator))) {
if (Numerator->isMinSignedValue()) {
// LLVM_DEBUG(dbgs() << "UnitLICM: Can't hoist " << inst.getOpcodeName() << ": Division of INT_MIN by -1:" << inst << "\n");
return can_hoist_sidefx(inst, args);
}

return can_hoist_operands(inst, args);
}
// The numerator *might* be MinSignedValue.
// LLVM_DEBUG(dbgs() << "UnitLICM: Can't hoist " << inst.getOpcodeName() << ": Division of INT_MIN( " << Numerator << " ) by -1:" << inst << "\n");
return can_hoist_sidefx(inst, args);
}

static bool can_hoist_udiv(Instruction const& inst, HoistArgs &args)
{
// code stolen from isSafeToSpeculativelyExecute
using namespace llvm::PatternMatch;
// x / y is undefined if y == 0.
const APInt *V;
if (match(inst.getOperand(1), m_APInt(V))) {
if (!*V) {
// LLVM_DEBUG(dbgs() << "UnitLICM: Can't hoist " << inst.getOpcodeName() << ": Division by zero:" << inst << "\n");
return can_hoist_sidefx(inst, args);
}

return can_hoist_operands(inst, args);
}
return false;
}

static bool can_hoist_load(Instruction const& inst, HoistArgs &args)
{
// code stolen from isSafeToSpeculativelyExecute
const auto *LI = dyn_cast<LoadInst>(&inst);
if (!LI) {
LLVM_DEBUG(dbgs() << "UnitLICM: Invalid instruction (expected load): " << inst << "\n");
return false;
}
if (mustSuppressSpeculation(*LI)) {
LLVM_DEBUG(dbgs() << "UnitLICM: Can't hoist load: Suppressed by sanitizer:" << inst << "\n");
return false;
}

return can_hoist_load_impl(inst, args);
}

static bool can_hoist_load_impl(Instruction const& inst, HoistArgs &args)
{
auto const *LI = dyn_cast<LoadInst>(&inst);
if (!LI) {
LLVM_DEBUG(dbgs() << "UnitLICM: Invalid instruction (expected load): " << inst << "\n");
return false;
}

auto const& aa = args.loop.alias_for(LI);

// If the load always reads the same mem loc, we can hoist it

if (aa.isMod()) {
LLVM_DEBUG(dbgs() << "UnitLICM: Can't hoist load: Load aliases modified memory:" << inst << "\n");
return false;
}

return can_hoist_operands(inst, args);
}

static bool can_hoist_store(Instruction const& inst, HoistArgs &args)
{
auto const *SI = dyn_cast<StoreInst>(&inst);
if (!SI) {
LLVM_DEBUG(dbgs() << "UnitLICM: Invalid instruction (expected store):\t" << inst << "\n");
return false;
}

// Volatile stores may not return, so we can't hoist them
if (SI->isVolatile())
return false;

auto const& aa = args.loop.alias_for(SI);

// If the store always modifies the same mem loc and it is never read, we can hoist it

if (!aa.isMustAlias()) {
LLVM_DEBUG(dbgs() << "UnitLICM: Can't hoist store: Store does not always write the same memory:" << inst << "\n");
return false;
}

if (aa.isRef()) {
LLVM_DEBUG(dbgs() << "UnitLICM: Can't hoist store: Store aliases read memory:" << inst << "\n");
return false;
}

// Writing to memory is side-effectful
return can_hoist_sidefx(inst, args);
}

/// We can only host expressions with side effects if they dominate all exits.
static bool can_hoist_sidefx(Instruction const& inst, HoistArgs& args)
{
if (!all_of(args.loop.exits, [&args, &inst](BasicBlock const* exit) {
return args.DT.dominates(&inst, exit);
})) {
LLVM_DEBUG(dbgs() << "UnitLICM: Can't hoist inst w/ sidefx: Does not dominate all exits:" << inst << "\n");
return false;
}

return can_hoist_operands(inst, args);
}

static bool can_hoist_operands(Instruction const& inst, HoistArgs &args)
{
return std::ranges::all_of(inst.operands(), [&args, &inst](Use const& V) {
bool const r = can_hoist(*V.get(), args);
LLVM_DEBUG(if (!r) dbgs() << inst << "\n");
return r;
});
}
edit in llvm/lib/Transforms/Scalar/CMakeLists.txt at line 80

[2.513308338]

[2.513308338]

UnitLICM.cpp
file addition: UnitLoopInfo.cpp (----------)

[2.602795540]

#include <cassert>
#include <utility>
#include <functional>
#include <iterator>
#include <ranges>
#include "llvm/Support/raw_ostream.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/Casting.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/IR/Value.h"
#include "llvm/IR/User.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/Instruction.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/Dominators.h"
#include "llvm/Analysis/AliasAnalysis.h"
#include "llvm/Analysis/AliasSetTracker.h"

#include "llvm/Analysis/UnitLoopInfo.h"

#define DEBUG_TYPE "unit-loop"
STATISTIC(NumLoopsDetected, "Number of loops detected");

using namespace llvm;
using namespace ece479k;

// Need to do this conversion because predecessors doesn't satisfy forward-iterator
// so we can use it in a filter view
static auto preds(BasicBlock const* block)
{
SmallVector<BasicBlock const*> r;

auto const p = predecessors(block);
r.append(p.begin(), p.end());

return r;
}

#pragma region "UnitLoopAnalysis"
/// Main function for running the Loop Identification analysis. This function
/// returns information about the loops in the function via the UnitLoopInfo
/// object
UnitLoopAnalysis::Result UnitLoopAnalysis::run(Function &F, FunctionAnalysisManager &FAM)
{
auto &DT = FAM.getResult<DominatorTreeAnalysis>(F);
auto &AA = FAM.getResult<AAManager>(F);

Result Loops;
// Need to have this ugliness to support both LLVM 15 and LLVM 16-19
#if LLVM_VERSION_MAJOR > 15
UnitLoopInfo::alias_results alias = BatchAAResults(AA);
#else
UnitLoopInfo::alias_results alias = AA;
#endif

auto const is_reachable = [&DT](BasicBlock const& bb) {
return DT.isReachableFromEntry(&bb);
};

// Find all back-edges in F
for (BasicBlock& BB: F | std::views::filter(is_reachable)) {

// ReSharper disable once CppDeclarationHidesLocal
auto const is_reachable = [&DT](BasicBlock const* bb) {
return DT.isReachableFromEntry(bb);
};

auto const is_backedge = [&DT, &BB](BasicBlock const* bb) {
return DT.dominates(&BB, bb);
};

// We need to filter out unreachble blocks twice because it's possible
// for an unreachble block to be the predecessor of a reachable block.
// See bad_loop1.ll for canonical case of this.

for (BasicBlock const* P: preds(&BB)
| std::views::filter(is_reachable)
| std::views::filter(is_backedge))
{
LLVM_DEBUG(dbgs()
<< "UnitLoopAnalysis: Found loop in " << F.getName() << " from `"
<< BB.getName() << "` to `" << P->getName() << "`");
Loops.emplace_back(&BB, const_cast<BasicBlock *>(P), alias, DT);
++NumLoopsDetected;
}
}
return Loops;
}

AnalysisKey UnitLoopAnalysis::Key;
#pragma endregion

#pragma region "UnitLoopInfo"
static BasicBlock* get_preheader(BasicBlock* header, BasicBlock* end)
{
BasicBlock* preheader;
// LoopSimplify guarantees 1 preheader and 1 backedge
// hence the preds of hedaer are [preheader, end]
assert(header->hasNPredecessors(2) && "Loop header preds were not [preheader, backedge]");

for (auto* pred : predecessors(header))
if (pred != end)
preheader = pred;

assert(preheader && "Preheader not found");
assert(preheader->isLegalToHoistInto() && "Preheader is not hoistable");
assert(preheader->getSingleSuccessor() == header && "Preheader is invalid");

return preheader;
}

UnitLoopInfo::UnitLoopInfo(BasicBlock* header, BasicBlock* end, alias_results AA, DominatorTree& DT)
:header(header), preheader(get_preheader(header, end)), alias_sets(std::make_unique<AliasSetTracker>(AA))
{
assert(header && "Can't record a loop that starts from a non-existent BasicBlock");
assert(end && "Can't record a loop that ends with a non-existent BasicBlock");
assert(preheader && "Loop does not have a valid preheader");

SmallVector<BasicBlock*, 16> worklist;
worklist.emplace_back(end);

// Populate loop body by walking up from the end to the header
while (!worklist.empty()) {
BasicBlock* cur = worklist.back(); worklist.pop_back();

// Make sure that we don't add the preds of the header
if (cur == header) {
set.insert(header);
blocks.emplace_back(header);
break;
}

if (set.insert(cur).second) {
blocks.emplace_back(cur);
auto preds = predecessors(cur);
worklist.append(preds.begin(), preds.end());
}
}

assert(!set.empty() && "Loops must contain at least one BasicBlock");
assert(blocks.back() == header && "Loop must start at the header");
assert(blocks.front() == end && "Loop must end at the end");

for (BasicBlock* block: blocks) {
assert(DT.dominates(preheader, block) && "All blocks in the loop must be dominated by the preheader");
assert(DT.dominates(header, block) && "All blocks in the loop must be dominated by the header");
}

// Populate AliasSets for this loop
for (BasicBlock* block: blocks | std::views::reverse)
alias_sets->add(*block);

// Populate the loop exits
for (BasicBlock* block : blocks)
// if any of the successors of a block point outside the loop, it is an exit from the loop
if (any_of(successors(block), [end, this](BasicBlock const* bb) {
return !this->contains(bb);
}))
exits.insert(block);

LLVM_DEBUG(
dbgs() << " with preheader `" << preheader->getName()
<< "`, header `" << header->getName()
<< "`, exits [ ";
for (BasicBlock const* exit: exits)
dbgs() << "`" << exit->getName() << "` ";

dbgs() << "]\n";
);
}

bool UnitLoopInfo::contains(BasicBlock const *BB) const
{
assert(!set.empty() && "Can't check an empty loop");
// assert(BB && "Can't look-up a non-existent BasicBlock");
return set.contains(BB);
}

bool UnitLoopInfo::contains(Instruction const *I) const
{
assert(!set.empty() && "Can't check an empty loop");
// assert(I && "Can't look-up a non-existent Instruction");
return contains(I->getParent());
}

bool UnitLoopInfo::contains(Instruction const& I) const
{
assert(!set.empty() && "Can't check an empty loop");
// assert(I && "Can't look-up a non-existent Instruction");
return contains(I.getParent());
}

bool UnitLoopInfo::is_invariant(Value const *V) const
{
assert(!set.empty() && "Can't be invariant to an empty loop");
if (auto const* i = dyn_cast<Instruction>(V))
return !contains(i);
return true;
}

bool UnitLoopInfo::is_invariant(Instruction const *I) const
{
assert(!set.empty() && "Can't be invariant to an empty loop");
return !contains(I);
}

bool UnitLoopInfo::is_invariant(Instruction const& I) const
{
assert(!set.empty() && "Can't be invariant to an empty loop");
return !contains(I);
}

bool UnitLoopInfo::has_invariant_operands(Instruction const *I) const
{
assert(!set.empty() && "Can't be invariant to an empty loop");
return std::ranges::all_of(I->operands(), [this](Value *v){ return is_invariant(v); });
}

std::vector<Instruction*> UnitLoopInfo::instructions() const
{
std::vector<Instruction*> r;

for (BasicBlock* block : blocks | std::views::reverse)
for (Instruction& inst: *block)
r.emplace_back(&inst);

return r;
}

std::vector<Instruction*> UnitLoopInfo::instructions(std::function<bool(Instruction const&)> const& filter) const
{
std::vector<Instruction*> r;

for (BasicBlock* block : blocks | std::views::reverse)
for (Instruction& inst: *block | std::views::filter(filter))
r.emplace_back(&inst);

return r;
}

AliasSet const& UnitLoopInfo::alias_for(LoadInst const* inst) const
{
return alias_sets->getAliasSetFor(MemoryLocation::get(inst));
}

AliasSet const& UnitLoopInfo::alias_for(StoreInst const* inst) const
{
return alias_sets->getAliasSetFor(MemoryLocation::get(inst));
}
#pragma endregion
edit in llvm/lib/Analysis/CMakeLists.txt at line 138

[2.606172191]

[2.606172191]

UnitLoopInfo.cpp
file addition: UnitLICM.h (----------)

[2.607628340]

#ifndef INCLUDE_UNIT_LICM_H
#define INCLUDE_UNIT_LICM_H

#include "llvm/IR/PassManager.h"

namespace llvm {
class Function;
}

using namespace llvm;

namespace ece479k
{
/// Loop Invariant Code Motion Optimization Pass
///
/// Implements the relaxed-case:
/// - We move pure expressions even if they don't dominate all exits
/// - We move expressions with side effects if they dominate all exits
struct UnitLICM : PassInfoMixin<UnitLICM> {
PreservedAnalyses run(Function &F, FunctionAnalysisManager &FAM);
};
} // namespace

#endif // INCLUDE_UNIT_LICM_H
file addition: UnitLoopInfo.h (----------)

[2.622017596]

#ifndef INCLUDE_UNIT_LOOP_INFO_H
#define INCLUDE_UNIT_LOOP_INFO_H

#include <vector>
#include <memory>
#include "llvm/IR/PassManager.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/Analysis/AliasAnalysis.h"
#include "llvm/Analysis/AliasSetTracker.h"

namespace llvm {
class Value;
class Instruction;
class BasicBlock;
class DominatorTree;
class LoadInst;
class StoreInst;
}

using namespace llvm;

namespace ece479k {

struct UnitLoopInfo {
private:
// The basic blocks that form this loop
SmallPtrSet<BasicBlock*, 16> set;
// The blocks in *reverse* order
SmallVector<BasicBlock*, 16> blocks;
public:
BasicBlock const* header;
BasicBlock const* preheader;
SmallPtrSet<BasicBlock const*, 4> exits;
private:
std::unique_ptr<AliasSetTracker> alias_sets;
public:
using alias_results = AAResults&;

// Record a natural loop given its backedge from end to header
UnitLoopInfo(BasicBlock* header, BasicBlock* end, alias_results AA, DominatorTree& DT);

/*loop properties*/
[[nodiscard]] bool contains(BasicBlock const*) const;
[[nodiscard]] bool contains(BasicBlock const&) const;
[[nodiscard]] bool contains(Instruction const*) const;
[[nodiscard]] bool contains(Instruction const&) const;
[[nodiscard]] bool contains(Value const*) const;
[[nodiscard]] bool contains(Value const&) const;

/*loop invariants*/
[[nodiscard]] bool is_invariant(Value const*) const;
[[nodiscard]] bool is_invariant(Value const&) const;
[[nodiscard]] bool is_invariant(Instruction const*) const;
[[nodiscard]] bool is_invariant(Instruction const&) const;
[[nodiscard]] bool has_invariant_operands(Instruction const*) const;
[[nodiscard]] bool has_invariant_operands(Instruction const&) const;

/*utility*/
[[nodiscard]] std::vector<Instruction*> instructions() const;
[[nodiscard]] std::vector<Instruction*> instructions(std::function<bool(Instruction const&)> const&) const;
[[nodiscard]] AliasSet const& alias_for(LoadInst const*) const;
[[nodiscard]] AliasSet const& alias_for(StoreInst const*) const;
};

/// Loop Identification Analysis Pass. Produces a UnitLoopInfo object which
/// should contain any information about the loops in the function which is
/// needed for your implementation of LICM
class UnitLoopAnalysis : public AnalysisInfoMixin<UnitLoopAnalysis> {
friend AnalysisInfoMixin<UnitLoopAnalysis>;
static AnalysisKey Key;

public:
using Result = std::vector<UnitLoopInfo>;

Result run(Function &F, FunctionAnalysisManager &AM);
};

} // namespace ece479k
#endif // INCLUDE_UNIT_LOOP_INFO_H