[Instrumentor] Add branch instrumentation

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

@@ -913,6 +913,39 @@ struct UnreachableIO : public InstructionIO<Instruction::Unreachable> {
   }
 };
 
+struct BranchIO : public InstructionIO<Instruction::Br> {
+  BranchIO() : InstructionIO<Instruction::Br>(/*IsPRE*/ true) {}
+  virtual ~BranchIO(){};
+
+  void init(InstrumentationConfig &IConf, LLVMContext &Ctx) {
+    IRTArgs.push_back(IRTArg(IntegerType::getInt8Ty(Ctx), "is_conditional",
+                             "Flag indicating a conditional branch.",
+                             IRTArg::NONE, isConditional));
+    IRTArgs.push_back(IRTArg(IntegerType::getInt8Ty(Ctx), "value",
+                             "Value of condition.", IRTArg::REPLACABLE,
+                             getValue, setValue));
+    IRTArgs.push_back(IRTArg(PointerType::getInt64Ty(Ctx), "num_successors",
+                             "Number of branch successors.", IRTArg::NONE,
+                             getNumSuccessors));
+    IConf.addChoice(*this);
+  }
+
+  static void populate(InstrumentationConfig &IConf, LLVMContext &Ctx) {
+    auto *AIC = IConf.allocate<BranchIO>();
+    AIC->init(IConf, Ctx);
+  }
+
+  static Value *isConditional(Value &V, Type &Ty, InstrumentationConfig &IConf,
+                              InstrumentorIRBuilderTy &IIRB);
+  static Value *getValue(Value &V, Type &Ty, InstrumentationConfig &IConf,
+                         InstrumentorIRBuilderTy &IIRB);
+  static Value *setValue(Value &V, Value &NewV, InstrumentationConfig &IConf,
+                         InstrumentorIRBuilderTy &IIRB);
+  static Value *getNumSuccessors(Value &V, Type &Ty,
+                                 InstrumentationConfig &IConf,
+                                 InstrumentorIRBuilderTy &IIRB);
+};
+
 struct ICmpIO : public InstructionIO<Instruction::ICmp> {
   ICmpIO(bool IsPRE) : InstructionIO<Instruction::ICmp>(IsPRE) {}
   virtual ~ICmpIO() {};

llvm/lib/Transforms/Instrumentation/Instrumentor.cpp

@@ -777,6 +777,7 @@ void InstrumentationConfig::populate(InstrumentorIRBuilderTy &IIRB) {
   ModuleIO::populate(*this, IIRB.Ctx);
   GlobalIO::populate(*this, IIRB.Ctx);
   AllocaIO::populate(*this, IIRB.Ctx);
+  BranchIO::populate(*this, IIRB.Ctx);
   StoreIO::populate(*this, IIRB);
   LoadIO::populate(*this, IIRB);
   CallIO::populate(*this, IIRB.Ctx);
@@ -1544,6 +1545,38 @@ Value *CallIO::isDefinition(Value &V, Type &Ty, InstrumentationConfig &IConf,
   return getCI(&Ty, !CI.getCalledFunction()->isDeclaration());
 }
 
+Value *BranchIO::isConditional(Value &V, Type &Ty, InstrumentationConfig &IConf,
+                               InstrumentorIRBuilderTy &IIRB) {
+  auto &BI = cast<BranchInst>(V);
+  return getCI(&Ty, BI.isConditional());
+}
+
+Value *BranchIO::getValue(Value &V, Type &Ty, InstrumentationConfig &IConf,
+                          InstrumentorIRBuilderTy &IIRB) {
+  auto &BI = cast<BranchInst>(V);
+  if (BI.isUnconditional())
+    return getCI(&Ty, 1);
+  return BI.getCondition();
+}
+
+Value *BranchIO::setValue(Value &V, Value &NewV, InstrumentationConfig &IConf,
+                          InstrumentorIRBuilderTy &IIRB) {
+  auto *BI = cast<BranchInst>(&V);
+  if (BI->isConditional()) {
+    auto Cast = tryToCast(IIRB.IRB, &NewV, BI->getCondition()->getType(),
+                          IIRB.DL, true);
+    BI->setCondition(Cast);
+  }
+  return BI;
+}
+
+Value *BranchIO::getNumSuccessors(Value &V, Type &Ty,
+                                  InstrumentationConfig &IConf,
+                                  InstrumentorIRBuilderTy &IIRB) {
+  auto &BI = cast<BranchInst>(V);
+  return getCI(&Ty, BI.getNumSuccessors());
+}
+
 Value *ICmpIO::getCmpPredicate(Value &V, Type &Ty, InstrumentationConfig &IConf,
                                InstrumentorIRBuilderTy &IIRB) {
   auto &II = cast<ICmpInst>(V);

llvm/test/Instrumentation/Instrumentor/branch.ll

@@ -0,0 +1,34 @@
+; 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:%.*]] = alloca <{ i32, i32, [7 x i8], i1 }>, align 8
+; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i64(ptr align 8 [[TMP0]], ptr @__instrumentor_value_pack, i64 16, i1 false)
+; CHECK-NEXT:    [[TMP3:%.*]] = getelementptr inbounds nuw <{ i32, i32, [7 x i8], i1 }>, ptr [[TMP0]], i32 0, i32 3
+; CHECK-NEXT:    store i1 [[C]], ptr [[TMP3]], align 1
+; CHECK-NEXT:    call void @__instrumentor_pre_function(ptr @foo, ptr @__instrumentor_.str, i32 1, ptr [[TMP0]]) #[[ATTR1:[0-9]+]]
+; CHECK-NEXT:    [[TMP4:%.*]] = getelementptr inbounds i8, ptr [[TMP0]], i32 8
+; CHECK-NEXT:    [[TMP5:%.*]] = load i1, ptr [[TMP4]], align 1
+; CHECK-NEXT:    [[TMP6:%.*]] = zext i1 [[TMP5]] to i8
+; CHECK-NEXT:    [[TMP8:%.*]] = call i8 @__instrumentor_pre_br(i8 1, i8 [[TMP6]], i64 2) #[[ATTR1]]
+; CHECK-NEXT:    [[TMP9:%.*]] = trunc i8 [[TMP8]] to i1
+; CHECK-NEXT:    br i1 [[TMP9]], label %[[A:.*]], label %[[B:.*]]
+; CHECK:       [[A]]:
+; CHECK-NEXT:    ret i32 0
+; CHECK:       [[B]]:
+; CHECK-NEXT:    [[TMP7:%.*]] = call i8 @__instrumentor_pre_br(i8 0, i8 1, i64 1) #[[ATTR1]]
+; 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

@@ -104,6 +104,16 @@
       "pointer": true,
       "pointer.description": "Input pointer of the ptr to int."
     },
+    "br": {
+      "enabled": true,
+      "is_conditional": true,
+      "is_conditional.description": "Flag indicating a conditional branch.",
+      "value": true,
+      "value.replace": true,
+      "value.description": "Value of condition.",
+      "num_successors": true,
+      "num_successors.description": "Number of branch successors."
+    },
     "alloca": {
       "enabled": true,
       "size": true,