From 453f9617e1b161d8d03908b2971781de3a2d9f74 Mon Sep 17 00:00:00 2001 From: Pavel Durov Date: Fri, 6 Dec 2024 17:52:49 +0000 Subject: [PATCH] Ensure that unoptimised functions call only unoptimised functions in SWT. Switch from module linking to function cloning for simplicity and changed the function cloning prefix from `__yk_clone_` to `__yk_opt_` to better reflect the optimization intent in ModuleClone pass. --- llvm/include/llvm/Transforms/Yk/ModuleClone.h | 2 +- llvm/lib/Transforms/Yk/ControlPoint.cpp | 8 +- llvm/lib/Transforms/Yk/ModuleClone.cpp | 110 +++++++++++++----- llvm/test/Transforms/Yk/ModuleClone.ll | 49 +++++--- 4 files changed, 123 insertions(+), 46 deletions(-) diff --git a/llvm/include/llvm/Transforms/Yk/ModuleClone.h b/llvm/include/llvm/Transforms/Yk/ModuleClone.h index c2da8961400f4e..ab96d339fc0abf 100644 --- a/llvm/include/llvm/Transforms/Yk/ModuleClone.h +++ b/llvm/include/llvm/Transforms/Yk/ModuleClone.h @@ -3,7 +3,7 @@ #include "llvm/Pass.h" -#define YK_CLONE_PREFIX "__yk_clone_" +#define YK_CLONE_PREFIX "__yk_opt_" #define YK_CLONE_MODULE_CP_COUNT 2 namespace llvm { diff --git a/llvm/lib/Transforms/Yk/ControlPoint.cpp b/llvm/lib/Transforms/Yk/ControlPoint.cpp index c71fa9b29ae7f4..7d2a12ebb0ff70 100644 --- a/llvm/lib/Transforms/Yk/ControlPoint.cpp +++ b/llvm/lib/Transforms/Yk/ControlPoint.cpp @@ -134,9 +134,13 @@ class YkControlPoint : public ModulePass { DS_Warning)); return false; } - assert(ControlPointCalls.size() == controlPointCount && - "Unexpected number of control point calls"); + if (ControlPointCalls.size() != controlPointCount) { + Context.emitError("Unexpected number of control point calls: " + + Twine(ControlPointCalls.size()) + + " expected: " + Twine(controlPointCount)); + return false; + } unsigned CPStackMapID = 0; Function *NF = nullptr; diff --git a/llvm/lib/Transforms/Yk/ModuleClone.cpp b/llvm/lib/Transforms/Yk/ModuleClone.cpp index de32f8653eea17..0184d8ecda24f6 100644 --- a/llvm/lib/Transforms/Yk/ModuleClone.cpp +++ b/llvm/lib/Transforms/Yk/ModuleClone.cpp @@ -1,7 +1,8 @@ //===- ModuleClone.cpp - Yk Module Cloning Pass ---------------------------===// // -// This pass duplicates functions within the module, producing both the original -// and new versions of functions. the process is as follows: +// This pass duplicates functions within the module, producing both the +// original (unoptimised) and cloned (optimised) versions of these +// functions. The process is as follows: // // - **Cloning Criteria:** // - Functions **without** their address taken are cloned. This results in two @@ -11,18 +12,12 @@ // and are not cloned. // // - **Cloned Function Naming:** -// - The cloned functions are renamed by adding the prefix `__yk_clone_` to +// - The cloned functions are renamed by adding the prefix `__yk_opt_` to // their original names. This distinguishes them from the original // functions. // -// - **Clone Process:** -// - 1. The original module is cloned, creating two copies of the module. -// - 2. The functions in the cloned module that satisfy the cloning criteria -// are renamed. -// - 3. The cloned module is linked back into the original module. -// // - **Optimisation Intent:** -// - The **cloned functions** (those with the `__yk_clone_` prefix) are +// - The **cloned functions** (those with the `__yk_opt_` prefix) are // intended to be the **optimised versions** of the functions. // - The **original functions** remain **unoptimised**. // @@ -58,39 +53,96 @@ struct YkModuleClone : public ModulePass { YkModuleClone() : ModulePass(ID) { initializeYkModuleClonePass(*PassRegistry::getPassRegistry()); } - void updateClonedFunctions(Module &M) { + + /** + * @brief Clones eligible functions within the given module. + * + * This function iterates over all functions in the provided LLVM module `M` + * and clones those that meet the following criteria: + * + * - The function does **not** have external linkage and is **not** a + * declaration. + * - The function's address is **not** taken. + * + * @param M The LLVM module containing the functions to be cloned. + * @return A map where the keys are the original function names and the + * values are pointers to the cloned `Function` objects. Returns + * a map if cloning succeeds, or `nullopt` if cloning fails. + */ + std::optional> + cloneFunctionsInModule(Module &M) { + LLVMContext &Context = M.getContext(); + std::map ClonedFuncs; for (llvm::Function &F : M) { + // Skip external declarations. if (F.hasExternalLinkage() && F.isDeclaration()) { continue; } - // Skip functions that are address taken - if (!F.hasAddressTaken()) { - F.setName(Twine(YK_CLONE_PREFIX) + F.getName()); + // Skip already cloned functions or functions with address taken. + if (F.hasAddressTaken() || F.getName().startswith(YK_CLONE_PREFIX)) { + continue; + } + ValueToValueMapTy VMap; + Function *ClonedFunc = CloneFunction(&F, VMap); + if (ClonedFunc == nullptr) { + Context.emitError("Failed to clone function: " + F.getName()); + return std::nullopt; + } + // Copy arguments + auto DestArgIt = ClonedFunc->arg_begin(); + for (const Argument &OrigArg : F.args()) { + DestArgIt->setName(OrigArg.getName()); + VMap[&OrigArg] = &*DestArgIt++; } + // Rename function + auto originalName = F.getName().str(); + auto cloneName = Twine(YK_CLONE_PREFIX) + originalName; + ClonedFunc->setName(cloneName); + ClonedFuncs[originalName] = ClonedFunc; } + return ClonedFuncs; } - bool runOnModule(Module &M) override { - std::unique_ptr Cloned = CloneModule(M); - if (!Cloned) { - llvm::report_fatal_error("Error cloning the module"); - return false; + /** + * @brief Updates call instructions in cloned functions to reference + * other cloned functions. + * + * @param M The LLVM module containing the functions. + * @param ClonedFuncs A map of cloned function names to functions. + */ + void + updateClonedFunctionCalls(Module &M, + std::map &ClonedFuncs) { + for (auto &Entry : ClonedFuncs) { + Function *ClonedFunc = Entry.second; + for (BasicBlock &BB : *ClonedFunc) { + for (Instruction &I : BB) { + if (auto *Call = dyn_cast(&I)) { + Function *CalledFunc = Call->getCalledFunction(); + if (CalledFunc && !CalledFunc->isIntrinsic()) { + std::string CalledName = CalledFunc->getName().str(); + auto It = ClonedFuncs.find(CalledName); + if (It != ClonedFuncs.end()) { + Call->setCalledFunction(It->second); + } + } + } + } + } } - updateClonedFunctions(*Cloned); + } - // The `OverrideFromSrc` flag instructs the linker to prioritise - // definitions from the source module (the second argument) when - // conflicts arise. This means that if two global variables, functions, - // or constants have the same name in both the original and cloned modules, - // the version from the cloned module will overwrite the original. - if (Linker::linkModules(M, std::move(Cloned), - Linker::Flags::OverrideFromSrc)) { - llvm::report_fatal_error("Error linking the modules"); + bool runOnModule(Module &M) override { + LLVMContext &Context = M.getContext(); + auto clonedFunctions = cloneFunctionsInModule(M); + if (!clonedFunctions) { + Context.emitError("Failed to clone functions in module"); return false; } + updateClonedFunctionCalls(M, *clonedFunctions); if (verifyModule(M, &errs())) { - errs() << "Module verification failed!"; + Context.emitError("Module verification failed!"); return false; } diff --git a/llvm/test/Transforms/Yk/ModuleClone.ll b/llvm/test/Transforms/Yk/ModuleClone.ll index 5ea725ff9414f1..ff8bf898a8e074 100644 --- a/llvm/test/Transforms/Yk/ModuleClone.ll +++ b/llvm/test/Transforms/Yk/ModuleClone.ll @@ -16,6 +16,12 @@ declare dso_local void @yk_location_drop(i64) declare dso_local void @yk_mt_shutdown(ptr noundef) define dso_local i32 @func_inc_with_address_taken(i32 %x) { +entry: + %call = call i32 @inc(i32 %x) + ret i32 %call +} + +define dso_local i32 @inc(i32 %x) { entry: %0 = add i32 %x, 1 ret i32 %0 @@ -60,6 +66,20 @@ entry: ; CHECK: declare dso_local void @yk_location_drop(i64) ; CHECK: declare dso_local void @yk_mt_shutdown(ptr noundef) +; Check that func_inc_with_address_taken is present in its original form +; CHECK-LABEL: define dso_local i32 @func_inc_with_address_taken(i32 %x) +; CHECK-NEXT: entry: +; CHECK-NEXT: call void @__yk_trace_basicblock({{.*}}) +; CHECK-NEXT: %call = call i32 @inc(i32 %x) +; CHECK-NEXT: ret i32 %call + +; Check original function: inc +; CHECK-LABEL: define dso_local i32 @inc(i32 %x) +; CHECK-NEXT: entry: +; CHECK-NEXT: call void @__yk_trace_basicblock({{.*}}) +; CHECK-NEXT: %0 = add i32 %x, 1 +; CHECK-NEXT: ret i32 %0 + ; Check original function: my_func ; CHECK-LABEL: define dso_local i32 @my_func(i32 %x) ; CHECK-NEXT: entry: @@ -78,26 +98,27 @@ entry: ; CHECK-NEXT: %0 = call i32 @my_func(i32 10) ; CHECK-NEXT: %1 = load i32, ptr @my_global ; CHECK-NEXT: %2 = call i32 (ptr, ...) @printf - -; Check that func_inc_with_address_taken is present in its original form -; CHECK-LABEL: define dso_local i32 @func_inc_with_address_taken(i32 %x) -; CHECK-NEXT: entry: -; CHECK-NEXT: call void @__yk_trace_basicblock({{.*}}) -; CHECK-NEXT: %0 = add i32 %x, 1 -; CHECK-NEXT: ret i32 %0 +; CHECK-NEXT: ret i32 0 ; ====================================================================== ; Functions whose addresses are taken should not be cloned ; ====================================================================== -; Functions with their addresses taken should not be cloned. +; Functions with their addresses taken should not be cloned. ; `func_inc_with_address_taken` is used by pointer and thus remains unaltered. -; CHECK-NOT: define dso_local i32 @__yk_clone_func_inc_with_address_taken +; CHECK-NOT: define dso_local i32 @__yk_opt_func_inc_with_address_taken ; ====================================================================== ; Cloned functions - should have no trace calls ; ====================================================================== -; Check cloned function: __yk_clone_my_func -; CHECK-LABEL: define dso_local i32 @__yk_clone_my_func(i32 %x) +; Check cloned function: __yk_opt_inc +; CHECK-LABEL: define dso_local i32 @__yk_opt_inc(i32 %x) +; CHECK-NEXT: entry: +; CHECK-NOT: call void @yk_trace_basicblock({{.*}}) +; CHECK-NEXT: %0 = add i32 %x, 1 +; CHECK-NEXT: ret i32 %0 + +; Check cloned function: __yk_opt_my_func +; CHECK-LABEL: define dso_local i32 @__yk_opt_my_func(i32 %x) ; CHECK-NEXT: entry: ; CHECK-NOT: call void @__yk_trace_basicblock({{.*}}) ; CHECK-NEXT: %0 = add i32 %x, 1 @@ -107,10 +128,10 @@ entry: ; CHECK-NEXT: %2 = call i32 %1(i32 42) ; CHECK-NEXT: ret i32 %2 -; Check cloned function: __yk_clone_main -; CHECK-LABEL: define dso_local i32 @__yk_clone_main() +; Check cloned function: __yk_opt_main +; CHECK-LABEL: define dso_local i32 @__yk_opt_main() ; CHECK-NEXT: entry: ; CHECK-NOT: call void @__yk_trace_basicblock({{.*}}) -; CHECK-NEXT: %0 = call i32 @__yk_clone_my_func(i32 10) +; CHECK-NEXT: %0 = call i32 @__yk_opt_my_func(i32 10) ; CHECK-NEXT: %1 = load i32, ptr @my_global ; CHECK-NEXT: %2 = call i32 (ptr, ...) @printf