From 09c0b714138f30f75a868e7dd79e6887104a5faa Mon Sep 17 00:00:00 2001
From: TB Schardl <neboat@mit.edu>
Date: Sat, 7 Dec 2024 15:45:12 -0500
Subject: [PATCH] [clang] Fix handling of Cilkscale tool variants, specifically
 cilkscale and cilkscale-instructions, to ensure instrumentation is inserted
 appropriately when using LLD with LTO.

---
 clang/lib/CodeGen/BackendUtil.cpp          | 2 +-
 clang/lib/Driver/ToolChains/CommonArgs.cpp | 4 ++++
 clang/lib/Driver/ToolChains/Darwin.cpp     | 4 ++++
 llvm/lib/Passes/PassBuilderPipelines.cpp   | 4 ++--
 4 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/clang/lib/CodeGen/BackendUtil.cpp b/clang/lib/CodeGen/BackendUtil.cpp
index 337034544a44..9d232d05dbd1 100644
--- a/clang/lib/CodeGen/BackendUtil.cpp
+++ b/clang/lib/CodeGen/BackendUtil.cpp
@@ -1078,7 +1078,7 @@ void EmitAssemblyHelper::RunOptimizationPipeline(
             });
         break;
       case LangOptions::CilktoolKind::Cilktool_Cilkscale_Benchmark:
-          PB.registerTapirLoopEndEPCallback(
+          PB.registerTapirLateEPCallback(
             [&PB](ModulePassManager &MPM, OptimizationLevel Level) {
               MPM.addPass(CSISetupPass(getCSIOptionsForCilkscaleBenchmark()));
               MPM.addPass(ComprehensiveStaticInstrumentationPass(
diff --git a/clang/lib/Driver/ToolChains/CommonArgs.cpp b/clang/lib/Driver/ToolChains/CommonArgs.cpp
index 9688b811c294..b18b7f20b2c7 100644
--- a/clang/lib/Driver/ToolChains/CommonArgs.cpp
+++ b/clang/lib/Driver/ToolChains/CommonArgs.cpp
@@ -301,6 +301,10 @@ static void renderTapirLoweringOptions(const ArgList &Args,
     } else if (Args.hasArg(options::OPT_fcilkplus))
       CmdArgs.push_back("--plugin-opt=tapir-target=cilkplus");
   }
+
+  if (const Arg *A = Args.getLastArg(options::OPT_fcilktool_EQ))
+    CmdArgs.push_back(Args.MakeArgString(
+        Twine("--plugin-opt=cilktool=") + A->getValue()));
 }
 
 void tools::addPathIfExists(const Driver &D, const Twine &Path,
diff --git a/clang/lib/Driver/ToolChains/Darwin.cpp b/clang/lib/Driver/ToolChains/Darwin.cpp
index 4439b2fa366e..71854e9eb893 100644
--- a/clang/lib/Driver/ToolChains/Darwin.cpp
+++ b/clang/lib/Driver/ToolChains/Darwin.cpp
@@ -576,6 +576,10 @@ static void renderTapirLoweringOptions(const ArgList &Args,
     } else if (Args.hasArg(options::OPT_fcilkplus))
       CmdArgs.push_back("--tapir-target=cilkplus");
   }
+
+  if (const Arg *A = Args.getLastArg(options::OPT_fcilktool_EQ))
+    CmdArgs.push_back(
+        Args.MakeArgString(Twine("--plugin-opt=cilktool=") + A->getValue()));
 }
 
 static void AppendPlatformPrefix(SmallString<128> &Path, const llvm::Triple &T);
diff --git a/llvm/lib/Passes/PassBuilderPipelines.cpp b/llvm/lib/Passes/PassBuilderPipelines.cpp
index f84de1e60c01..0a4055dce924 100644
--- a/llvm/lib/Passes/PassBuilderPipelines.cpp
+++ b/llvm/lib/Passes/PassBuilderPipelines.cpp
@@ -1777,7 +1777,7 @@ PassBuilder::buildPerModuleDefaultPipeline(OptimizationLevel Level,
     MPM.addPass(buildTapirLoweringPipeline(
         Level, LTOPreLink ? ThinOrFullLTOPhase::FullLTOPreLink
                           : ThinOrFullLTOPhase::None));
-  else
+  else if (!LTOPreLink)
     invokeTapirLoopEndEPCallbacks(MPM, Level);
 
   return MPM;
@@ -2359,7 +2359,7 @@ ModulePassManager PassBuilder::buildO0DefaultPipeline(OptimizationLevel Level,
     MPM.addPass(buildTapirLoweringPipeline(
         Level, LTOPreLink ? ThinOrFullLTOPhase::FullLTOPreLink
                           : ThinOrFullLTOPhase::None));
-  else
+  else if (!LTOPreLink)
     invokeTapirLoopEndEPCallbacks(MPM, Level);
 
   invokeOptimizerLastEPCallbacks(MPM, Level);