Add a pass for basic block tracing.

This pass is a crucial part of the Software Tracer implementation.
Software Tracer will allow us to execute tracing on platforms that are
not supported by the Hardware Tracer and compare performance between
both.

The pass iterates over module basic blocks and injects externally
linked function yk_trace_basicblock call instruction into the beginning
of each block.

Author: Pavel-Durov

llvm/include/llvm/InitializePasses.h

@@ -342,7 +342,7 @@ void initializeWinEHPreparePass(PassRegistry&);
 void initializeWriteBitcodePassPass(PassRegistry&);
 void initializeXRayInstrumentationPass(PassRegistry&);
 void initializeYkSplitBlocksAfterCallsPass(PassRegistry&);
-
+void initializeYkBasicBlockTracerPass(PassRegistry&);
 } // end namespace llvm
 
 #endif // LLVM_INITIALIZEPASSES_H

llvm/include/llvm/Transforms/Yk/BasicBlockTracer.h

@@ -0,0 +1,13 @@
+#ifndef LLVM_TRANSFORMS_YK_HELLOWORLD_H
+#define LLVM_TRANSFORMS_YK_HELLOWORLD_H
+
+#include "llvm/Pass.h"
+
+// The name of the trace function
+#define YK_TRACE_FUNCTION "yk_trace_basicblock"
+
+namespace llvm {
+ModulePass *createYkBasicBlockTracerPass();
+} // namespace llvm
+
+#endif

llvm/lib/CodeGen/CodeGen.cpp

@@ -151,4 +151,5 @@ void llvm::initializeCodeGen(PassRegistry &Registry) {
   initializeWinEHPreparePass(Registry);
   initializeXRayInstrumentationPass(Registry);
   initializeYkSplitBlocksAfterCallsPass(Registry);
+  initializeYkBasicBlockTracerPass(Registry);
 }

llvm/lib/CodeGen/TargetPassConfig.cpp

@@ -55,6 +55,7 @@
 #include "llvm/Transforms/Yk/SplitBlocksAfterCalls.h"
 #include "llvm/Transforms/Yk/Stackmaps.h"
 #include "llvm/Transforms/Yk/NoCallsInEntryBlocks.h"
+#include "llvm/Transforms/Yk/BasicBlockTracer.h"
 #include <cassert>
 #include <optional>
 #include <string>
@@ -292,6 +293,11 @@ static cl::opt<bool>
     YkSplitBlocksAfterCalls("yk-split-blocks-after-calls", cl::init(false), cl::NotHidden,
                       cl::desc("Split blocks after function calls."));
 
+static cl::opt<bool>
+    YkBasicBlockTracer("yk-basicblock-tracer", cl::init(false), cl::NotHidden,
+                      cl::desc("Enables YK Software Tracer capability"));
+
+
 /// Allow standard passes to be disabled by command line options. This supports
 /// simple binary flags that either suppress the pass or do nothing.
 /// i.e. -disable-mypass=false has no effect.
@@ -1181,6 +1187,10 @@ bool TargetPassConfig::addISelPasses() {
     addPass(createYkStackmapsPass());
   }
 
+  if (YkBasicBlockTracer) {
+    addPass(createYkBasicBlockTracerPass());
+  }
+
   addISelPrepare();
   return addCoreISelPasses();
 }

llvm/lib/Transforms/Yk/BasicBlockTracer.cpp

@@ -0,0 +1,69 @@
+#include "llvm/Transforms/Yk/BasicBlockTracer.h"
+#include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Instruction.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Module.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/Pass.h"
+
+#define DEBUG_TYPE "yk-basicblock-tracer-pass"
+
+using namespace llvm;
+
+namespace llvm {
+void initializeYkBasicBlockTracerPass(PassRegistry &);
+} // namespace llvm
+
+namespace {
+struct YkBasicBlockTracer : public ModulePass {
+  static char ID;
+
+  YkBasicBlockTracer() : ModulePass(ID) {
+    initializeYkBasicBlockTracerPass(*PassRegistry::getPassRegistry());
+  }
+
+  bool runOnModule(Module &M) override {
+    LLVMContext &Context = M.getContext();
+    // Create externally linked function declaration:
+    //   void yk_trace_basicblock(int functionIndex, int blockIndex)
+    Type *ReturnType = Type::getVoidTy(Context);
+    Type *FunctionIndexArgType = Type::getInt32Ty(Context);
+    Type *BlockIndexArgType = Type::getInt32Ty(Context);
+
+    FunctionType *FType = FunctionType::get(
+        ReturnType, {FunctionIndexArgType, BlockIndexArgType}, false);
+    Function *TraceFunc = Function::Create(
+        FType, GlobalVariable::ExternalLinkage, YK_TRACE_FUNCTION, M);
+
+    IRBuilder<> builder(Context);
+    uint32_t FunctionIndex = 0;
+    for (auto &F : M) {
+      uint32_t BlockIndex = 0;
+      for (auto &BB : F) {
+        builder.SetInsertPoint(&*BB.getFirstInsertionPt());
+        builder.CreateCall(TraceFunc, {builder.getInt32(FunctionIndex),
+                                       builder.getInt32(BlockIndex)});
+        assert(BlockIndex != UINT32_MAX &&
+               "Expected BlockIndex to not overflow");
+        BlockIndex++;
+      }
+      assert(FunctionIndex != UINT32_MAX &&
+             "Expected FunctionIndex to not overflow");
+      FunctionIndex++;
+    }
+    return true;
+  }
+};
+} // namespace
+
+char YkBasicBlockTracer::ID = 0;
+
+INITIALIZE_PASS(YkBasicBlockTracer, DEBUG_TYPE, "yk basicblock tracer", false,
+                false)
+
+ModulePass *llvm::createYkBasicBlockTracerPass() {
+  return new YkBasicBlockTracer();
+}

llvm/lib/Transforms/Yk/CMakeLists.txt

@@ -7,6 +7,7 @@ add_llvm_component_library(LLVMYkPasses
   ShadowStack.cpp
   NoCallsInEntryBlocks.cpp
   SplitBlocksAfterCalls.cpp
+  BasicBlockTracer.cpp
 
   DEPENDS
   intrinsics_gen

llvm/test/Transforms/Yk/BasicBlockTracer.ll

@@ -0,0 +1,86 @@
+; RUNNING TEST EXAMPLE: llvm-lit llvm/test/Transforms/Yk/BasicBlockTracer.ll
+; RUN: llc -stop-after yk-basicblock-tracer-pass --yk-basicblock-tracer < %s  | FileCheck %s
+
+; CHECK-LABEL: define dso_local noundef i32 @main()
+; CHECK-NEXT:  call void @yk_trace_basicblock(i32 0, i32 0)
+define dso_local noundef i32 @main() #0 {
+  %1 = alloca i32, align 4
+  %2 = alloca i32, align 4
+  %3 = alloca i32, align 4
+  store i32 0, i32* %1, align 4
+  store i32 0, i32* %2, align 4
+  store i32 0, i32* %3, align 4
+  br label %4
+
+; CHECK-LABEL: 4:{{.*}}
+; CHECK-NEXT:  call void @yk_trace_basicblock(i32 0, i32 1)
+4:                                                ; preds = %13, %0
+  %5 = load i32, i32* %3, align 4
+  %6 = icmp slt i32 %5, 43
+  br i1 %6, label %7, label %16
+
+; CHECK-LABEL: 7:{{.*}}
+; CHECK-NEXT:  call void @yk_trace_basicblock(i32 0, i32 2)
+7:                                                ; preds = %4
+  %8 = load i32, i32* %3, align 4
+  %9 = icmp eq i32 %8, 42
+  br i1 %9, label %10, label %12
+
+; CHECK-LABEL: 10:{{.*}}
+; CHECK-NEXT:  call void @yk_trace_basicblock(i32 0, i32 3)
+10:                                               ; preds = %7
+  %11 = load i32, i32* %3, align 4
+  store i32 %11, i32* %1, align 4
+  br label %17
+
+; CHECK-LABEL: 12:{{.*}}
+; CHECK-NEXT:  call void @yk_trace_basicblock(i32 0, i32 4)
+12:                                               ; preds = %7
+  br label %13
+
+; CHECK-LABEL: 13:{{.*}}
+; CHECK-NEXT:  call void @yk_trace_basicblock(i32 0, i32 5)
+13:                                               ; preds = %12
+  %14 = load i32, i32* %3, align 4
+  %15 = add nsw i32 %14, 1
+  store i32 %15, i32* %3, align 4
+  br label %4, !llvm.loop !6
+
+; CHECK-LABEL: 16:{{.*}}
+; CHECK-NEXT:  call void @yk_trace_basicblock(i32 0, i32 6)
+16:                                               ; preds = %4
+  store i32 0, i32* %1, align 4
+  br label %17
+
+; CHECK-LABEL: 17:{{.*}}
+; CHECK-NEXT:  call void @yk_trace_basicblock(i32 0, i32 7)
+17:                                               ; preds = %16, %10
+  %18 = load i32, i32* %1, align 4
+  ret i32 %18
+}
+
+; CHECK-LABEL: define dso_local noundef i32 @_Z5checki(i32 noundef %0)
+; CHECK-NEXT:  call void @yk_trace_basicblock(i32 1, i32 0)
+define dso_local noundef i32 @_Z5checki(i32 noundef %0) #1 {
+  %2 = alloca i32, align 4
+  store i32 %0, i32* %2, align 4
+  %3 = load i32, i32* %2, align 4
+  %4 = icmp eq i32 %3, 42
+  %5 = zext i1 %4 to i32
+  ret i32 %5
+}
+
+attributes #0 = { mustprogress noinline norecurse nounwind optnone uwtable "frame-pointer"="all" "min-legal-vector-width"="0" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" }
+attributes #1 = { mustprogress noinline nounwind optnone uwtable "frame-pointer"="all" "min-legal-vector-width"="0" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" }
+
+!llvm.module.flags = !{!0, !1, !2, !3, !4}
+!llvm.ident = !{!5}
+
+!0 = !{i32 1, !"wchar_size", i32 4}
+!1 = !{i32 7, !"PIC Level", i32 2}
+!2 = !{i32 7, !"PIE Level", i32 2}
+!3 = !{i32 7, !"uwtable", i32 1}
+!4 = !{i32 7, !"frame-pointer", i32 2}
+!5 = !{!"Debian clang version 14.0.6"}
+!6 = distinct !{!6, !7}
+!7 = !{!"llvm.loop.mustprogress"}