From 6ab9d78a52600b966dc67a3215f304c01a1b0f86 Mon Sep 17 00:00:00 2001 From: Lang Hames Date: Sat, 19 Sep 2020 18:56:49 -0700 Subject: [PATCH] [ORC][examples] Add an OrcV2 example for IR optimization via IRTransformLayer. Shows how to write a custom IR transform to apply a legacy::PassManager pipeline. --- examples/OrcV2Examples/CMakeLists.txt | 1 + .../CMakeLists.txt | 12 ++ .../LLJITWithOptimizingIRTransform.cpp | 122 ++++++++++++++++++ .../ExecutionEngine/Orc/IRTransformLayer.h | 3 +- 4 files changed, 137 insertions(+), 1 deletion(-) create mode 100644 examples/OrcV2Examples/LLJITWithOptimizingIRTransform/CMakeLists.txt create mode 100644 examples/OrcV2Examples/LLJITWithOptimizingIRTransform/LLJITWithOptimizingIRTransform.cpp diff --git a/examples/OrcV2Examples/CMakeLists.txt b/examples/OrcV2Examples/CMakeLists.txt index fb7fd41b970..1d87a84fee9 100644 --- a/examples/OrcV2Examples/CMakeLists.txt +++ b/examples/OrcV2Examples/CMakeLists.txt @@ -5,6 +5,7 @@ add_subdirectory(LLJITWithInitializers) add_subdirectory(LLJITWithLazyReexports) add_subdirectory(LLJITWithObjectCache) add_subdirectory(LLJITWithObjectLinkingLayerPlugin) +add_subdirectory(LLJITWithOptimizingIRTransform) add_subdirectory(LLJITWithTargetProcessControl) add_subdirectory(LLJITWithThinLTOSummaries) add_subdirectory(OrcV2CBindingsAddObjectFile) diff --git a/examples/OrcV2Examples/LLJITWithOptimizingIRTransform/CMakeLists.txt b/examples/OrcV2Examples/LLJITWithOptimizingIRTransform/CMakeLists.txt new file mode 100644 index 00000000000..63be55cfa37 --- /dev/null +++ b/examples/OrcV2Examples/LLJITWithOptimizingIRTransform/CMakeLists.txt @@ -0,0 +1,12 @@ +set(LLVM_LINK_COMPONENTS + Core + ExecutionEngine + IRReader + OrcJIT + Support + nativecodegen + ) + +add_llvm_example(LLJITWithOptimizingIRTransform + LLJITWithOptimizingIRTransform.cpp + ) diff --git a/examples/OrcV2Examples/LLJITWithOptimizingIRTransform/LLJITWithOptimizingIRTransform.cpp b/examples/OrcV2Examples/LLJITWithOptimizingIRTransform/LLJITWithOptimizingIRTransform.cpp new file mode 100644 index 00000000000..7ebf3e8e9a0 --- /dev/null +++ b/examples/OrcV2Examples/LLJITWithOptimizingIRTransform/LLJITWithOptimizingIRTransform.cpp @@ -0,0 +1,122 @@ +//===-- LLJITWithOptimizingIRTransform.cpp -- LLJIT with IR optimization --===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// In this example we will use an IR transform to optimize a module as it +// passes through LLJIT's IRTransformLayer. +// +//===----------------------------------------------------------------------===// + +#include "llvm/ExecutionEngine/Orc/LLJIT.h" +#include "llvm/IR/LegacyPassManager.h" +#include "llvm/Support/InitLLVM.h" +#include "llvm/Support/TargetSelect.h" +#include "llvm/Support/raw_ostream.h" +#include "llvm/Transforms/IPO.h" +#include "llvm/Transforms/Scalar.h" + +#include "../ExampleModules.h" + +using namespace llvm; +using namespace llvm::orc; + +ExitOnError ExitOnErr; + +// Example IR module. +// +// This IR contains a recursive definition of the factorial function: +// +// fac(n) | n == 0 = 1 +// | otherwise = n * fac(n - 1) +// +// It also contains an entry function which calls the factorial function with +// an input value of 5. +// +// We expect the IR optimization transform that we build below to transform +// this into a non-recursive factorial function and an entry function that +// returns a constant value of 5!, or 120. + +const llvm::StringRef MainMod = + R"( + + define i32 @fac(i32 %n) { + entry: + %tobool = icmp eq i32 %n, 0 + br i1 %tobool, label %return, label %if.then + + if.then: ; preds = %entry + %arg = add nsw i32 %n, -1 + %call_result = call i32 @fac(i32 %arg) + %result = mul nsw i32 %n, %call_result + br label %return + + return: ; preds = %entry, %if.then + %final_result = phi i32 [ %result, %if.then ], [ 1, %entry ] + ret i32 %final_result + } + + define i32 @entry() { + entry: + %result = call i32 @fac(i32 5) + ret i32 %result + } + +)"; + +// A function object that creates a simple pass pipeline to apply to each +// module as it passes through the IRTransformLayer. +class MyOptimizationTransform { +public: + MyOptimizationTransform() : PM(std::make_unique()) { + PM->add(createTailCallEliminationPass()); + PM->add(createFunctionInliningPass()); + PM->add(createIndVarSimplifyPass()); + PM->add(createCFGSimplificationPass()); + } + + Expected operator()(ThreadSafeModule TSM, + MaterializationResponsibility &R) { + TSM.withModuleDo([this](Module &M) { + dbgs() << "--- BEFORE OPTIMIZATION ---\n" << M << "\n"; + PM->run(M); + dbgs() << "--- AFTER OPTIMIZATION ---\n" << M << "\n"; + }); + return std::move(TSM); + } + +private: + std::unique_ptr PM; +}; + +int main(int argc, char *argv[]) { + // Initialize LLVM. + InitLLVM X(argc, argv); + + InitializeNativeTarget(); + InitializeNativeTargetAsmPrinter(); + + ExitOnErr.setBanner(std::string(argv[0]) + ": "); + + // (1) Create LLJIT instance. + auto J = ExitOnErr(LLJITBuilder().create()); + + // (2) Install transform to optimize modules when they're materialized. + J->getIRTransformLayer().setTransform(MyOptimizationTransform()); + + // (3) Add modules. + ExitOnErr(J->addIRModule(ExitOnErr(parseExampleModule(MainMod, "MainMod")))); + + // (4) Look up the JIT'd function and call it. + auto EntrySym = ExitOnErr(J->lookup("entry")); + auto *Entry = (int (*)())EntrySym.getAddress(); + + int Result = Entry(); + outs() << "--- Result ---\n" + << "entry() = " << Result << "\n"; + + return 0; +} diff --git a/include/llvm/ExecutionEngine/Orc/IRTransformLayer.h b/include/llvm/ExecutionEngine/Orc/IRTransformLayer.h index ee4ee3437fa..475d3f259ae 100644 --- a/include/llvm/ExecutionEngine/Orc/IRTransformLayer.h +++ b/include/llvm/ExecutionEngine/Orc/IRTransformLayer.h @@ -13,6 +13,7 @@ #ifndef LLVM_EXECUTIONENGINE_ORC_IRTRANSFORMLAYER_H #define LLVM_EXECUTIONENGINE_ORC_IRTRANSFORMLAYER_H +#include "llvm/ADT/FunctionExtras.h" #include "llvm/ExecutionEngine/JITSymbol.h" #include "llvm/ExecutionEngine/Orc/Layer.h" #include @@ -27,7 +28,7 @@ namespace orc { /// before operating on the module. class IRTransformLayer : public IRLayer { public: - using TransformFunction = std::function( + using TransformFunction = unique_function( ThreadSafeModule, MaterializationResponsibility &R)>; IRTransformLayer(ExecutionSession &ES, IRLayer &BaseLayer,