Merge pull request rust-lang#111 from Pavel-Durov/basic-block-tracer-llvm-pass

Implement BasicBlockTracer pass

Author: ltratt

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"}