BotUpParallelPass.h

//
// Created by yongchao on 12/20/22.
//
 
#ifndef CLEARBLUE_BOTUPPARALLELPASS_H
#define CLEARBLUE_BOTUPPARALLELPASS_H
 
#include <condition_variable>
#include <llvm/IR/Function.h>
#include <llvm/IR/Module.h>
#include <llvm/Pass.h>
#include <mutex>
#include <set>
 
#include "Analysis/CallGraph/CBCallGraph.h"
#include "IR/SEG/SymbolicExprGraphBuilder.h"
#include "Utils/ProgressBar.h"
 
using namespace llvm;
 
class TaskDepVertex {
private:
  Function *F;
  size_t Index;
 
  std::set<TaskDepVertex *> Parents;
 
  std::set<TaskDepVertex *> InternalChildren;
  std::set<TaskDepVertex *> ExternalChildren;
 
public:
  TaskDepVertex(Function *F, size_t I) : F(F), Index(I) {}
 
  void internallyDependsOn(TaskDepVertex *T) {
    assert(T);
    InternalChildren.insert(T);
    T->Parents.insert(this);
  }
 
  void externallyDependsOn(TaskDepVertex *T) {
    assert(T);
    ExternalChildren.insert(T);
    T->Parents.insert(this);
  }
 
  std::set<TaskDepVertex *>::iterator
  removeParent(std::set<TaskDepVertex *>::iterator It) {
    if ((*It)->F != F) {
      (*It)->ExternalChildren.erase(this);
    } else {
      (*It)->InternalChildren.erase(this);
    }
    return Parents.erase(It);
  }
 
  void removeParent(TaskDepVertex *T) {
    if (T->F != F) {
      T->ExternalChildren.erase(this);
    } else {
      T->InternalChildren.erase(this);
    }
    Parents.erase(T);
  }
 
  bool hasChild(TaskDepVertex *T) {
    if (T->F != F) {
      return ExternalChildren.count(T);
    }
    return InternalChildren.count(T);
  }
 
  size_t getIndex() { return Index; }
 
  Function *getFunction() { return F; }
 
  class ChildrenIterator {
  private:
    std::set<TaskDepVertex *> &InternalChildren;
    std::set<TaskDepVertex *> &ExternalChildren;
 
    std::set<TaskDepVertex *>::iterator It;
 
  public:
    ChildrenIterator(std::set<TaskDepVertex *> &InternalChildren,
                     std::set<TaskDepVertex *> &ExternalChildren, bool Begin)
        : InternalChildren(InternalChildren),
          ExternalChildren(ExternalChildren) {
      if (Begin) {
        It = InternalChildren.begin();
        if (It == InternalChildren.end())
          It = ExternalChildren.begin();
      } else {
        It = ExternalChildren.end();
      }
    }
 
    ChildrenIterator(const ChildrenIterator &ChildrenIt)
        : InternalChildren(ChildrenIt.InternalChildren),
          ExternalChildren(ChildrenIt.ExternalChildren), It(ChildrenIt.It) {}
 
    ChildrenIterator &operator++() {
      It++;
      if (It == InternalChildren.end()) {
        It = ExternalChildren.begin();
      }
      return *this;
    }
 
    ChildrenIterator operator++(int) {
      ChildrenIterator Old(*this);
      ++(*this);
      return Old;
    }
 
    TaskDepVertex *operator*() { return *It; }
 
    bool operator==(const ChildrenIterator &CIt) { return It == CIt.It; }
 
    bool operator!=(const ChildrenIterator &CIt) { return It != CIt.It; }
  };
 
  std::set<TaskDepVertex *>::iterator internal_children_begin() {
    return InternalChildren.begin();
  }
 
  std::set<TaskDepVertex *>::iterator internal_children_end() {
    return InternalChildren.end();
  }
 
  bool internal_children_empty() { return InternalChildren.empty(); }
 
  std::set<TaskDepVertex *>::iterator external_children_begin() {
    return ExternalChildren.begin();
  }
 
  std::set<TaskDepVertex *>::iterator external_children_end() {
    return ExternalChildren.end();
  }
 
  bool external_children_empty() { return ExternalChildren.empty(); }
 
  ChildrenIterator children_begin() {
    return ChildrenIterator(InternalChildren, ExternalChildren, true);
  }
 
  ChildrenIterator children_end() {
    return ChildrenIterator(InternalChildren, ExternalChildren, false);
  }
 
  bool children_empty() {
    return internal_children_empty() && external_children_empty();
  }
 
  std::set<TaskDepVertex *>::iterator parents_begin() {
    return Parents.begin();
  }
 
  std::set<TaskDepVertex *>::iterator parents_end() { return Parents.end(); }
};
 
class TaskDepGraph {
private:
  std::vector<TaskDepVertex *> Vertices;
 
public:
  TaskDepGraph() {}
 
  ~TaskDepGraph() {
    for (auto *V : Vertices) {
      delete V;
    }
  }
 
  void add(TaskDepVertex *V) { Vertices.push_back(V); }
 
  std::vector<TaskDepVertex *>::iterator begin() { return Vertices.begin(); }
 
  std::vector<TaskDepVertex *>::iterator end() { return Vertices.end(); }
 
  size_t size() { return Vertices.size(); }
 
  const std::vector<TaskDepVertex *> &getVertices() { return Vertices; }
};
 
class BotUpParallelPass : public ModulePass {
private:
  unsigned NumTasksEachFunc;
 
  ProgressBar Prog;
  std::mutex Mutex;
  std::condition_variable Cond;
  std::vector<TaskDepVertex *> FinishedTasks;
  std::map<Function *, unsigned> FunctionIndegreeMap;
 
protected:
  CBCallGraph *CG = nullptr;
  SymbolicExprGraphBuilder *SEGBuilder = nullptr;
 
public:
  BotUpParallelPass(char &ID, unsigned NumTasks)
      : ModulePass(ID), NumTasksEachFunc(NumTasks),
        Prog("[Sailfish Checking]", ProgressBar::PBS_CharacterStyle) {}
 
  virtual ~BotUpParallelPass() {}
 
  virtual void getAnalysisUsage(AnalysisUsage &AU) const override;
 
  virtual bool runOnModule(Module &M) override;
 
  virtual bool runOnFunction(Function &F) = 0;
 
  virtual bool runOnFunction(Function &F, unsigned TaskID) = 0;
 
  virtual bool releaseMemory(Function &F) = 0;
 
  virtual bool internallyDepends(unsigned TaskAID, unsigned TaskBID) = 0;
 
  virtual bool externallyDepends(unsigned CallerTaskID,
                                 unsigned CalleeTaskID) = 0;
  inline bool toCheck(const Function *Func) {
    if (Func && Func->hasName() &&
        Func->getName().startswith("Pinpoint.CondExit.")) {
      return false;
    }
    return Func && !Func->isIntrinsic() && !Func->isDeclaration();
  }
 
private:
  void buildTaskDependenceGraph(TaskDepGraph &TDG, CBCallGraph &CG, Module &M);
 
  void buildDepsTemplate(std::vector<TaskDepVertex *> &Tasks);
 
  void buildTasks(TaskDepGraph &TDG, Function *F,
                  std::vector<TaskDepVertex *> &Tasks,
                  std::vector<TaskDepVertex *> &Templates);
 
  void buildExternalDeps(std::vector<TaskDepVertex *> &CallerTasks,
                         std::vector<TaskDepVertex *> &CalleeTasks,
                         std::vector<TaskDepVertex *> &Templates);
 
  void execute(TaskDepVertex *);
 
  void finish(TaskDepVertex *);
 
  void postProcess(TaskDepVertex *, CBCallGraph &CG);
 
  void wait(TaskDepGraph &, CBCallGraph &CG);
};
 
#endif // CLEARBLUE_BOTUPPARALLELPASS_H