[LLVM][IR] Add Scarr and block checkpoints to Transforms

[gAlfonso-bit]

Still a draft, I should probably add some tests.

llvm/include/llvm/Analysis/CFGPrinter.h

     raw_string_ostream OS(Str);
 
     Node->printAsOperand(OS, false);
+    OS << " " << cpToString(Node->getCheckpoint());
     return OS.str();
   }
 

     --I;
   }
 
+  static std::string cpToString(Checkpoint cp) {
+    switch(cp) {
+    case Checkpoint::NA:
+      return "";
+    case Checkpoint::ThreadStart:
+      return "ThreadStart";
+    case Checkpoint::ThreadEnd:
+      return "ThreadEnd";
+    case Checkpoint::ExitPoint:
+      return "ExitPoint";
+    case Checkpoint::Virtual:
+      return "Virtual";
+    default:
+      return "Unknown";
+    }
+  }
+
   static std::string getCompleteNodeLabel(
       const BasicBlock *Node, DOTFuncInfo *,
       llvm::function_ref<void(raw_string_ostream &, const BasicBlock &)>

 
     if (Node->getName().empty()) {
       Node->printAsOperand(OS, false);
+      if (Node->getCheckpoint() != Checkpoint::NA) {
+        OS << " " << cpToString(Node->getCheckpoint());
+      }
       OS << ":";
     }
 

llvm/include/llvm/IR/BasicBlock.h

 class PHINode;
 class ValueSymbolTable;
 
+// https://www.usenix.org/system/files/raid2019-toffalini.pdf
+enum class Checkpoint {
+  // Not a checkpoint
+  NA,
+  ThreadStart,
+  ThreadEnd,
+  ExitPoint,
+  Virtual
+};
+
 /// LLVM Basic Block Representation
 ///
 /// This represents a single basic block in LLVM. A basic block is simply a

 
   InstListType InstList;
   Function *Parent;
+  Checkpoint cp;
 
   void setParent(Function *parent);
 

 
   /// Get the context in which this basic block lives.
   LLVMContext &getContext() const;
+  void setCheckpoint(Checkpoint);
+  Checkpoint getCheckpoint() const;
 
   /// Instruction iterators...
   using iterator = InstListType::iterator;

llvm/include/llvm/Transforms/Scarr/ScarrCpMarker.h

+//===-- ScarrCpMarker.h - Transform IR with Checkpoint Info -----*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+///===----------------------------------------------------------------------===//
+//
+// Mark Basic Blocks into different ScaRR checkpoint types.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TRANSFORMS_SCARR_SCARRCPMARKER_H
+#define LLVM_TRANSFORMS_SCARR_SCARRCPMARKER_H
+
+#include "llvm/IR/PassManager.h"
+
+namespace llvm {
+
+class ScarrCpMarkerPass : public PassInfoMixin<ScarrCpMarkerPass> {
+public:
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+};
+
+} // namespace llvm
+
+#endif // LLVM_TRANSFORMS_SCARR_SCARRCPMARKER_H

llvm/include/llvm/Transforms/Scarr/ScarrLoaCollector.h

+//===- ScarrLoaCollector.h - Collect LoA from BasicBlock --------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+///===---------------------------------------------------------------------===//
+//
+// Collect LoA from BasicBlock. Must be run after scarr-cp-marker pass.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TRANSFORMS_SCARR_SCARRLOACOLLECTOR_H
+#define LLVM_TRANSFORMS_SCARR_SCARRLOACOLLECTOR_H
+
+#include "llvm/IR/PassManager.h"
+
+namespace llvm {
+
+class ScarrLoaCollectorPass : public PassInfoMixin<ScarrLoaCollectorPass> {
+public:
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+};
+
+} // namespace llvm
+
+#endif // LLVM_TRANSFORMS_SCARR_SCARRLOACOLLECTOR_H

llvm/lib/IR/BasicBlock.cpp

   setBasicBlockBits(Bits);
 }
 
+void BasicBlock::setCheckpoint(Checkpoint _cp) {
+  cp = _cp;
+}
+
+Checkpoint BasicBlock::getCheckpoint() const {
+  return cp;
+}
+
 #ifndef NDEBUG
 /// In asserts builds, this checks the numbering. In non-asserts builds, it
 /// is defined as a no-op inline function in BasicBlock.h.

llvm/lib/Passes/CMakeLists.txt

   Analysis
   Core
   Coroutines
+  Scarr
   IPO
   InstCombine
   ObjCARC

llvm/lib/Passes/PassBuilder.cpp

 #include "llvm/Transforms/Instrumentation/PoisonChecking.h"
 #include "llvm/Transforms/Instrumentation/SanitizerCoverage.h"
 #include "llvm/Transforms/Instrumentation/ThreadSanitizer.h"
+#include "llvm/Transforms/Scarr/ScarrCpMarker.h"
+#include "llvm/Transforms/Scarr/ScarrLoaCollector.h"
 #include "llvm/Transforms/ObjCARC.h"
 #include "llvm/Transforms/Scalar/ADCE.h"
 #include "llvm/Transforms/Scalar/AlignmentFromAssumptions.h"

llvm/lib/Passes/PassRegistry.def

 FUNCTION_PASS("gvn-hoist", GVNHoistPass())
 FUNCTION_PASS("gvn-sink", GVNSinkPass())
 FUNCTION_PASS("helloworld", HelloWorldPass())
+FUNCTION_PASS("scarr-cp-marker", ScarrCpMarkerPass())
+FUNCTION_PASS("scarr-loa-collector", ScarrLoaCollectorPass())
 FUNCTION_PASS("infer-address-spaces", InferAddressSpacesPass())
 FUNCTION_PASS("instcombine", InstCombinePass())
 FUNCTION_PASS("instcount", InstCountPass())

llvm/lib/Transforms/CMakeLists.txt

 add_subdirectory(IPO)
 add_subdirectory(Vectorize)
 add_subdirectory(Hello)
+add_subdirectory(Scarr)
 add_subdirectory(ObjCARC)
 add_subdirectory(Coroutines)
 add_subdirectory(CFGuard)

llvm/lib/Transforms/Scarr/CMakeLists.txt

+add_llvm_component_library(LLVMScarr
+        ScarrCpMarker.cpp
+        ScarrLoaCollector.cpp
+
+        DEPENDS
+        intrinsics_gen
+
+        LINK_COMPONENTS
+        Core
+        Support
+        )

llvm/lib/Transforms/Scarr/ScarrCpMarker.cpp

+//===-- ScarrCpMarker.cpp - Transform IR with Checkpoint Info ---*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+///===----------------------------------------------------------------------===//
+//
+// Mark Basic Blocks into different ScaRR checkpoint types.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Scarr/ScarrCpMarker.h"
+#include "llvm/Analysis/LoopInfo.h"
+#include "llvm/IR/Dominators.h"
+
+using namespace llvm;
+
+void findVirtualCheckpoint(DominatorTree &DT, Function &F) {
+  DT.recalculate(F);
+  // generate the LoopInfoBase for the current function
+  LoopInfoBase<BasicBlock, Loop>* KLoop = new LoopInfoBase<BasicBlock, Loop>();
+  KLoop->releaseMemory();
+  KLoop->analyze(DT);
+  for (auto &bb : F) {
+    // Since the BasicBlock would have been inlined, just traverse from main function
+    if (F.getName() == "main") {
+      auto loop = KLoop->getLoopFor(&bb);
+      if (loop != nullptr) {
+        loop->getHeader()->setCheckpoint(Checkpoint::Virtual);
+      }
+    }
+  }
+}
+
+void findCheckpoints(DominatorTree &DT, Function &F, int nestedLevel) {
+  bool isThreadStartCheckpoint = F.getName() == "main";
+  for (auto &bb : F) {
+    bool isThreadEndCheckpoint = false;
+    bool isExitPointCheckpoint = false;
+    for (auto &i : bb) {
+      // if instruction is last in the block and has no more successor,
+      // then this will be thread end checkpoint
+      if (i.isTerminator()) {
+        // Thread end only in the original function (main)
+        if (i.getNumSuccessors() == 0 && nestedLevel == 0 && F.getName() == "main") {
+          isThreadEndCheckpoint |= true;
+        }
+      }
+      // Check if instruction is calling a function
+      if (isa<CallInst>(i)) {
+        auto *call = &cast<CallBase>(i);
+        // use this hack to check if function is external
+        if (call != nullptr && call->getCalledFunction() != nullptr && !call->getCalledFunction()->empty()) {
+          auto calledFunction = call->getCalledFunction()->getName();
+          if (calledFunction == F.getName()) {
+            // Recursion is detected
+            continue;
+          }
+          findCheckpoints(DT, *(call->getCalledFunction()), nestedLevel + 1);
+        } else {
+          // The function is outside of the translation unit, hence it is an exit point
+          if (!isThreadStartCheckpoint && !isThreadEndCheckpoint) {
+            isExitPointCheckpoint |= true;
+          }
+        }
+      }
+    }
+
+    if (isThreadStartCheckpoint) {
+      isThreadStartCheckpoint = false;
+      bb.setCheckpoint(Checkpoint::ThreadStart);
+    } else if (isThreadEndCheckpoint) {
+      bb.setCheckpoint(Checkpoint::ThreadEnd);
+    } else if (isExitPointCheckpoint) {
+      bb.setCheckpoint(Checkpoint::ExitPoint);
+    }
+  }
+
+  findVirtualCheckpoint(DT,F);
+}
+
+
+PreservedAnalyses ScarrCpMarkerPass::run(Function &F, FunctionAnalysisManager &AM) {
+  DominatorTree DT;
+  findCheckpoints(DT, F, 0);
+  return PreservedAnalyses::all();
+}

llvm/lib/Transforms/Scarr/ScarrLoaCollector.cpp

+//===- ScarrLoaCollector.cpp - Collect LoA from BasicBlock ------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+///===---------------------------------------------------------------------===//
+//
+// Collect LoA from BasicBlock. Must be run after scarr-cp-marker pass.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Scarr/ScarrLoaCollector.h"
+#include "llvm/ADT/DepthFirstIterator.h"
+#include "llvm/ADT/SCCIterator.h"
+#include "llvm/IR/CFG.h"
+#include <iostream>
+#include <vector>
+
+using namespace llvm;
+
+// We will store ScaRR measurements in this type
+using MeasurementMap = std::map<std::pair<BasicBlock *, BasicBlock *>, std::vector<BasicBlock *>>;
+
+// For each BasicBlock find checkpoint and collect LoA that directs
+// control flow from previous Checkpoint to the next one.
+void handle(BasicBlock *firstCp,
+            BasicBlock *successor,
+            MeasurementMap &measurements,
+            std::vector<BasicBlock *> LoA) {
+  // This checkpoint is branch, hence we need to collect LoA
+  if (firstCp->getTerminator()->getNumSuccessors() > 1) {
+    // We only add firstCp to LoA if LoA is still empty
+    if (LoA.empty()) {
+      LoA.push_back(firstCp);
+    }
+  }
+  for (auto succ : successors(successor)) {
+    // We need a copy of LoA in every loop
+    auto succLoA = LoA;
+    // Successor is a checkpoint
+    if (succ->getCheckpoint() != Checkpoint::NA) {
+      if (succLoA.size() == 1) {
+        succLoA.push_back(succ);
+      }
+      auto cp = std::make_pair(firstCp, succ);
+      measurements[cp] = succLoA;
+    } else {
+      if (!succLoA.empty() &&
+          succLoA.back()->getCheckpoint() != Checkpoint::NA) {
+        succLoA.push_back(succ);
+      }
+      handle(firstCp, succ, measurements, succLoA);
+    }
+  }
+}
+
+void collectListOfActions(Function &function) {
+  // We will run only in main function
+  if (function.getName() == "main") {
+    // This is the list of checkpoints which we will later iterate
+    std::vector<BasicBlock *> basicBlockCheckpoints;
+    for (auto it : depth_first(&function.getEntryBlock())) {
+      if (it->getCheckpoint() != Checkpoint::NA) {
+        basicBlockCheckpoints.push_back(it);
+      }
+    }
+
+    // List of Action
+    std::vector<BasicBlock *> LoA;
+    // We will store the measurement here
+    MeasurementMap measurements;
+    for (auto cp : basicBlockCheckpoints) {
+      handle(cp, cp, measurements, LoA);
+    }
+
+    auto loaCount = 0;
+    for (auto measurement : measurements) {
+      auto listOfActions = measurement.second;
+      loaCount+= listOfActions.size();
+    }
+
+    // Printing the result
+    std::cout << "=============================================================" << std::endl;
+    std::cout << "ScaRR Offline Measurement Statistics" << std::endl;
+    std::cout << "=============================================================" << std::endl;
+    std::cout << "Offline Measurement Size: " << measurements.size() << std::endl;
+    std::cout << "Number of Checkpoints: " << basicBlockCheckpoints.size() << std::endl;
+    std::cout << "Number of List of Actions: " << loaCount << std::endl;
+    std::cout << "=============================================================" << std::endl;
+    std::cout << "Checkpoints and LoA Details: " << std::endl;
+    auto mIndex = 0;
+    for (auto measurement : measurements) {
+      auto cpPair = measurement.first;
+      auto listOfActions = measurement.second;
+      std::cout << "=============================================================" << std::endl;
+      std::cout << "Measurement " << mIndex << std::endl;
+      std::cout << "LoA Size: " << listOfActions.size() << std::endl << std::endl;
+      outs() << "Checkpoint_" << mIndex << "_A: " << *cpPair.first << "\n";
+      outs() << "Checkpoint_" << mIndex << "_B: " << *cpPair.second << "\n";
+
+      if (!listOfActions.empty()) {
+        std::cout << "LoA Details: " << std::endl;
+      }
+
+      int idx = 0;
+      for (auto loa : listOfActions) {
+        outs() << "LOA_" << idx << ": " << *loa << "\n";
+        idx++;
+      }
+      std::cout << "=============================================================" << std::endl;
+      mIndex++;
+    }
+  }
+}
+
+PreservedAnalyses ScarrLoaCollectorPass::run(Function &F, FunctionAnalysisManager &AM) {
+  collectListOfActions(F);
+  return PreservedAnalyses::all();
+}