Merge pull request #215 from vext01/jump-thread-ctrlp

Prevent JumpThreadingPass from messing with control points.
ykjit · Nov 18, 2024 · c73f96f · c73f96f
2 parents ce98289 + 49693e4
commit c73f96f
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 4 deletions.
diff --git a/llvm/include/llvm/Support/Yk.h b/llvm/include/llvm/Support/Yk.h
@@ -9,5 +9,6 @@ void initYkOptions(void);
 // of us randomly introducing `extern bool`s all over the place.
 extern bool YkOptNoneAfterIRPasses;
 extern bool YkDontOptFuncABI;
+extern bool YkPatchCtrlPoint;
 
 #endif
diff --git a/llvm/lib/CodeGen/TargetPassConfig.cpp b/llvm/lib/CodeGen/TargetPassConfig.cpp
@@ -44,6 +44,7 @@
 #include "llvm/Support/Threading.h"
 #include "llvm/Support/VirtualFileSystem.h"
 #include "llvm/Support/WithColor.h"
+#include "llvm/Support/Yk.h"
 #include "llvm/Target/CGPassBuilderOption.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Transforms/Scalar.h"
@@ -270,10 +271,6 @@ static cl::opt<bool>
     YkBlockDisambiguate("yk-block-disambiguate", cl::init(false), cl::NotHidden,
                         cl::desc("Disambiguate blocks for yk"));
 
-static cl::opt<bool> YkPatchCtrlPoint("yk-patch-control-point", cl::init(false),
-                                      cl::NotHidden,
-                                      cl::desc("Patch yk_mt_control_point()"));
-
 static cl::opt<bool> YkLinkage("yk-linkage", cl::init(false),
                                       cl::NotHidden,
                                       cl::desc("Change functions with internal linkage to have external linkage"));

diff --git a/llvm/lib/Support/Yk.cpp b/llvm/lib/Support/Yk.cpp
@@ -110,6 +110,20 @@ struct CreateYkDontOptFuncABIParser {
 } // namespace
 static ManagedStatic<cl::opt<bool, true>, CreateYkDontOptFuncABIParser> YkDontOptFuncABIParser;
 
+bool YkPatchCtrlPoint;
+namespace {
+struct CreateYkPatchCtrlPointParser {
+  static void *call() {
+    return new cl::opt<bool, true>(
+        "yk-patch-control-point",
+        cl::desc("Patch yk control points"),
+        cl::NotHidden, cl::location(YkPatchCtrlPoint));
+  }
+};
+} // namespace
+static ManagedStatic<cl::opt<bool, true>, CreateYkPatchCtrlPointParser> YkPatchCtrlPointParser;
+
+
 void llvm::initYkOptions() {
   *YkExtendedLLVMBBAddrMapSectionParser;
   *YkStackMapOffsetFixParser;
@@ -118,4 +132,5 @@ void llvm::initYkOptions() {
   *YkOptNoneAfterIRPassesParser;
   *YkEmbedIRParser;
   *YkDontOptFuncABIParser;
+  *YkPatchCtrlPointParser;
 }
diff --git a/llvm/lib/Transforms/Scalar/JumpThreading.cpp b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
@@ -64,11 +64,13 @@
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Support/Yk.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/SSAUpdater.h"
 #include "llvm/Transforms/Utils/ValueMapper.h"
+#include "llvm/Transforms/Yk/ControlPoint.h"
 #include <algorithm>
 #include <cassert>
 #include <cstdint>
@@ -963,6 +965,19 @@ static bool hasAddressTakenAndUsed(BasicBlock *BB) {
 /// processBlock - If there are any predecessors whose control can be threaded
 /// through to a successor, transform them now.
 bool JumpThreadingPass::processBlock(BasicBlock *BB) {
+  // We mustn't duplicate yk JIT control points.
+  if (YkPatchCtrlPoint) {
+    for (Instruction &I: *BB) {
+      if (CallInst *CI = dyn_cast<CallInst>(&I)) {
+        Function *CF = CI->getCalledFunction();
+        // The control point hasn't been patched yet, so we look for the
+        // pre-patched name.
+        if ((CF != nullptr) && (CF->getName() == YK_DUMMY_CONTROL_POINT)) {
+          return false;
+        }
+      }
+    }
+  }
   // If the block is trivially dead, just return and let the caller nuke it.
   // This simplifies other transformations.
   if (DTU->isBBPendingDeletion(BB) ||