kildall.h

// This file if part of the llir-opt project.
// Licensing information can be found in the LICENSE file.
// (C) 2018 Nandor Licker. All rights reserved.
 
#pragma once
 
#include <llvm/ADT/PostOrderIterator.h>
#include <queue>
#include "core/cfg.h"
 
 
 
enum class Direction {
  FORWARD,
  BACKWARD,
};
 
template <typename FlowSet, typename GenSet, typename KillSet, Direction Dir>
class KillGenSolver {
protected:
  struct InstInfo {
    Inst *I;
 
    GenSet Gen;
    KillSet Kill;
 
    InstInfo(Inst *I) : I(I) {}
  };
 
public:
  KillGenSolver(Func &func);
 
  void Solve();
 
protected:
  virtual void Build(Inst &inst) = 0;
 
  virtual void Traverse(Inst *inst, const FlowSet &set) = 0;
 
  InstInfo &Info(Inst *I)
  {
    return blocks_[blockToIndex_[I->getParent()]].Insts.emplace_back(I);
  }
 
private:
  struct BlockInfo {
    Block *B;
    llvm::SmallVector<uint64_t, 5> Preds;
    llvm::SmallVector<uint64_t, 5> Succs;
    std::vector<InstInfo> Insts;
    GenSet Gen;
    KillSet Kill;
    FlowSet Flow;
 
    BlockInfo(Block *B) : B(B) {}
  };
 
protected:
  Func &func_;
 
private:
  std::vector<BlockInfo> blocks_;
  std::unordered_map<const Block *, uint64_t> blockToIndex_;
};
 
template <typename FlowSet, typename GenSet, typename KillSet, Direction Dir>
KillGenSolver<FlowSet, GenSet, KillSet, Dir>::KillGenSolver(Func &func)
  : func_(func)
{
  // Generate unique block IDs.
  for (Block &block : func) {
    blockToIndex_.insert({ &block, blocks_.size() });
    blocks_.push_back(&block);
  }
 
  // Build the graph.
  for (Block &block : func) {
    BlockInfo &blockInfo = blocks_[blockToIndex_[&block]];
 
    // Construct fast pred/succ information.
    for (const Block *pred : block.predecessors()) {
      blockInfo.Preds.push_back(blockToIndex_[pred]);
    }
    for (const Block *succ : block.successors()) {
      blockInfo.Succs.push_back(blockToIndex_[succ]);
    }
  }
}
 
template <typename FlowSet, typename GenSet, typename KillSet, Direction Dir>
void KillGenSolver<FlowSet, GenSet, KillSet, Dir>::Solve()
{
  // Build constraints.
  for (Block &block : func_) {
    for (Inst &inst : block) {
      Build(inst);
    }
  }
 
  // Construct per-block kill/gen.
  for (BlockInfo &block : blocks_) {
    auto Step = [&block](InstInfo &inst) {
      // kill' = kill U killNew
      block.Kill.Union(inst.Kill);
      // gen' = (gen - killNew) U genNew
      block.Gen.Minus(inst.Kill);
      block.Gen.Union(inst.Gen);
    };
 
    switch (Dir) {
      case Direction::FORWARD: {
        for (auto it = block.Insts.begin(); it != block.Insts.end(); ++it) {
          Step(*it);
        }
        break;
      }
      case Direction::BACKWARD: {
        for (auto it = block.Insts.rbegin(); it != block.Insts.rend(); ++it) {
          Step(*it);
        }
        break;
      }
    }
  }
 
  // Populate the worklist.
  std::queue<BlockInfo *> queue_;
  std::set<BlockInfo *> inQueue_;
  std::vector<llvm::GraphTraits<Func *>::NodeRef> blockOrder_;
  {
    auto Add = [&queue_, &inQueue_] (BlockInfo *info) {
      inQueue_.insert(info);
      queue_.push(info);
    };
 
    // Find the order of the nodes.
    auto Entry = llvm::GraphTraits<Func *>::getEntryNode(&func_);
    std::copy(po_begin(Entry), po_end(Entry), std::back_inserter(blockOrder_));
 
    // Add nodes to queue.
    switch (Dir) {
      case Direction::FORWARD: {
        for (auto it = blockOrder_.rbegin(); it != blockOrder_.rend(); ++it) {
          auto *block = &blocks_[blockToIndex_[*it]];
 
          std::optional<FlowSet> init;
          for (unsigned index : block->Preds) {
            auto *pred = &blocks_[index];
 
            FlowSet out = pred->Flow;
            out.Minus(pred->Kill);
            out.Union(pred->Gen);
            if (init) {
              init->Union(out);
            } else {
              init = out;
            }
          }
 
          if (init) {
            block->Flow = *init;
          }
 
          Add(block);
        }
        break;
      }
      case Direction::BACKWARD: {
        for (auto it = blockOrder_.begin(); it != blockOrder_.end(); ++it) {
          auto *block = &blocks_[blockToIndex_[*it]];
 
          std::optional<FlowSet> init;
          for (unsigned index : block->Preds) {
            auto *succ = &blocks_[index];
 
            FlowSet out = succ->Flow;
            out.Minus(succ->Kill);
            out.Union(succ->Gen);
            if (init) {
              init->Union(out);
            } else {
              init = out;
            }
          }
 
          if (init) {
            block->Flow = *init;
          }
 
          Add(block);
        }
        break;
      }
    }
  }
 
  // Iterate until worklist has elements.
  while (!queue_.empty()) {
    BlockInfo *block = queue_.front();
    queue_.pop();
    inQueue_.erase(block);
 
    FlowSet out = block->Flow;
    out.Minus(block->Kill);
    out.Union(block->Gen);
 
    // If inputs to a block change, continue.
    auto Transfer = [&out, &inQueue_, &queue_] (BlockInfo *next) {
      FlowSet in = next->Flow;
      in.Union(out);
      if (!(in == next->Flow)) {
        next->Flow = in;
        if (inQueue_.count(next) == 0) {
          queue_.push(next);
        }
      }
    };
 
    switch (Dir) {
      case Direction::FORWARD: {
        for (unsigned succ : block->Succs) {
          Transfer(&blocks_[succ]);
        }
        break;
      }
      case Direction::BACKWARD: {
        for (unsigned pred : block->Preds) {
          Transfer(&blocks_[pred]);
        }
        break;
      }
    }
  }
 
  // Traverse nodes, applying the transformation.
  for (Block *block : blockOrder_) {
    BlockInfo &info = blocks_[blockToIndex_[block]];
 
    FlowSet set = info.Flow;
 
    auto Step = [this, &set] (InstInfo &i) {
      set.Minus(i.Kill);
      set.Union(i.Gen);
      Traverse(i.I, set);
    };
 
    switch (Dir) {
      case Direction::FORWARD: {
        for (auto it = info.Insts.begin(); it != info.Insts.end(); ++it) {
          Step(*it);
        }
        break;
      }
      case Direction::BACKWARD: {
        for (auto it = info.Insts.rbegin(); it != info.Insts.rend(); ++it) {
          Step(*it);
        }
        break;
      }
    }
  }
}