[Instrumentor] Add branch instrumentation

llvm/include/llvm/Transforms/Instrumentation/Instrumentor.h

@@ -60,6 +60,7 @@ using GlobalCallbackTy = std::function<bool(GlobalObject &)>;
 struct InstrumentorConfig {
   static Type *getType(LLVMContext &Ctx, StringRef S) {
     auto *Ty = StringSwitch<Type *>(S)
+                   .Case("int1_t", Type::getInt1Ty(Ctx))
                    .Case("int8_t", Type::getInt8Ty(Ctx))
                    .Case("int16_t", Type::getInt16Ty(Ctx))
                    .Case("int32_t", Type::getInt32Ty(Ctx))

llvm/include/llvm/Transforms/Instrumentation/InstrumentorConfig.def

@@ -352,6 +352,30 @@ CVALUE(base_pointer, bool, SkipUnused, true)
 SECTION_END(base_pointer)
 ///}
 
+SECTION_START(branch, InstrumentorConfig::PRE)
+
+/// Selection of information passed to the runtime.
+///{
+/// Whether or not the branch instruction is conditional
+RTVALUE(branch, IsConditional, true, "int8_t", BOOLEAN)
+
+/// The boolean value of the conditional
+RTVALUE(branch, Value, true, "int8_t", BOOLEAN)
+
+/// Number of branch successors
+RTVALUE(branch, Successors, true, "int64_t", PLAIN)
+///}
+
+/// Should conditional branches be instrumented?
+CVALUE(branch, bool, SkipConditional, false)
+
+/// Should unconditional branches be instrumented?
+CVALUE(branch, bool, SkipUnconditional, false)
+
+CVALUE_INTERNAL(branch, CallbackTy, CB, nullptr)
+
+SECTION_END(branch)
+
 #undef SECTION_START
 #undef CVALUE
 #undef CVALUE_INTERNAL

llvm/lib/Transforms/Instrumentation/Instrumentor.cpp

@@ -486,6 +486,7 @@ class InstrumentorImpl final {
   bool instrumentStore(StoreInst &I, InstrumentorConfig::Position P);
   bool instrumentUnreachable(UnreachableInst &I,
                              InstrumentorConfig::Position P);
+  bool instrumentBranch(BranchInst &I, InstrumentorConfig::Position P);
   bool instrumentMainFunction(Function &MainFn);
   bool instrumentModule(InstrumentorConfig::Position P);
 
@@ -1475,6 +1476,36 @@ bool InstrumentorImpl::instrumentUnreachable(UnreachableInst &I,
   return true;
 }
 
+bool InstrumentorImpl::instrumentBranch(BranchInst &I,
+                                        InstrumentorConfig::Position P) {
+  if (!IC.branch.isEnabled(P))
+    return false;
+  if (IC.branch.CB && !IC.branch.CB(I))
+    return false;
+  if (I.isConditional() && IC.branch.SkipConditional)
+    return false;
+  if (I.isUnconditional() && IC.branch.SkipUnconditional)
+    return false;
+
+  assert(P != InstrumentorConfig::PRE_AND_POST &&
+         P != InstrumentorConfig::POST);
+
+  RTArgumentPack RTArgPack(*this);
+  setInsertPoint(I, P);
+  SmallVector<RTArgument> RTArgs;
+
+  RTArgPack.addCI(IC.branch.IsConditional, I.isConditional(), P);
+  if (I.isConditional())
+    RTArgPack.addVal(IC.branch.Value, I.getCondition(), P);
+  else
+    RTArgPack.addCI(IC.branch.Value, 1, P);
+  RTArgPack.addCI(IC.branch.Successors, I.getNumSuccessors(), P);
+
+  getCall(IC.branch.SectionName, RTArgPack, P);
+
+  return true;
+}
+
 bool InstrumentorImpl::instrumentFunction(Function &Fn) {
   bool Changed = false;
   if (!shouldInstrumentFunction(&Fn))
@@ -1555,6 +1586,10 @@ bool InstrumentorImpl::instrumentFunction(Function &Fn) {
         Changed |= instrumentUnreachable(cast<UnreachableInst>(I),
                                          InstrumentorConfig::POST);
         break;
+      case Instruction::Br:
+        Changed |=
+            instrumentBranch(cast<BranchInst>(I), InstrumentorConfig::PRE);
+        break;
       default:
         break;
       }

llvm/test/Instrumentation/Instrumentor/branch.ll

@@ -0,0 +1,26 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt < %s -passes=instrumentor -S | FileCheck %s
+define i32 @foo(i1 %c) {
+; CHECK-LABEL: define i32 @foo(
+; CHECK-SAME: i1 [[C:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[TMP0:%.*]] = zext i1 [[C]] to i8
+; CHECK-NEXT:    call void @__instrumentor_pre_branch(i8 1, i8 [[TMP0]], i64 2)
+; CHECK-NEXT:    br i1 [[C]], label %[[A:.*]], label %[[B:.*]]
+; CHECK:       [[A]]:
+; CHECK-NEXT:    ret i32 0
+; CHECK:       [[B]]:
+; CHECK-NEXT:    call void @__instrumentor_pre_branch(i8 0, i8 1, i64 1)
+; CHECK-NEXT:    br label %[[D:.*]]
+; CHECK:       [[D]]:
+; CHECK-NEXT:    ret i32 1
+;
+entry:
+  br i1 %c, label %a, label %b
+a:
+  ret i32 0
+b:
+  br label %d
+d:
+  ret i32 1
+}

llvm/test/Instrumentation/Instrumentor/default_config.json

@@ -114,5 +114,13 @@
     "EnabledPost": true,
     "Value": true,
     "SkipUnused": true
+  },
+  "branch": {
+    "EnabledPre": true,
+    "IsConditional": true,
+    "Value": true,
+    "Successors": true,
+    "SkipConditional": false,
+    "SkipUnconditional": false
   }
 }