From 49693e47ad7ab6bfce940bebbce2c7ee5ed1ede6 Mon Sep 17 00:00:00 2001
From: Edd Barrett <vext01@gmail.com>
Date: Wed, 13 Nov 2024 13:35:34 +0000
Subject: [PATCH] Prevent JumpThreadingPass from messing with control points.

Requires moving the -yk-patch-ctrl-point flag so it's accessible in two
places.
---
 llvm/include/llvm/Support/Yk.h               |  1 +
 llvm/lib/CodeGen/TargetPassConfig.cpp        |  5 +----
 llvm/lib/Support/Yk.cpp                      | 15 +++++++++++++++
 llvm/lib/Transforms/Scalar/JumpThreading.cpp | 15 +++++++++++++++
 4 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/llvm/include/llvm/Support/Yk.h b/llvm/include/llvm/Support/Yk.h
index 2d8e95afb19bf7..dcd89a72aafca6 100644
--- 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
index c6b3e1058421dc..1f90426f97aa4f 100644
--- 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
index e74d8882169810..b07ffc17ca52b4 100644
--- 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
index f2b9d784ead8af..62d20f3f1f9601 100644
--- 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) ||