From e6b8650551d8ebe3a5e32858c4ab283882ea9a15 Mon Sep 17 00:00:00 2001 From: PikachuHy Date: Sat, 20 Jan 2024 14:28:09 +0800 Subject: [PATCH] [codegen] generate to LLVM IR directly --- CMakeLists.txt | 2 +- include/pscm/codegen/codegen.h | 3 +- include/pscm/scm_utils.h | 2 +- src/codegen/codegen.cpp | 11 +- src/codegen/llvm_ir/llvm_ir_codegen.cpp | 156 ++++++++++++++++++++++++ 5 files changed, 169 insertions(+), 5 deletions(-) create mode 100644 src/codegen/llvm_ir/llvm_ir_codegen.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 6124e399..ce1d022b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -104,7 +104,7 @@ target_sources(pscm PRIVATE ${PSCM_SRCS}) if (PSCM_ENABLE_MLIR_CODEGEN) message(STATUS "Enable codegen with MLIR") target_compile_definitions(pscm PRIVATE PSCM_ENABLE_MLIR_CODEGEN) - file(GLOB CODEGEN_SRCS "src/codegen/*.cpp" src/codegen/mlir/*.cpp) + file(GLOB CODEGEN_SRCS src/codegen/*.cpp src/codegen/mlir/*.cpp src/codegen/llvm_ir/*.cpp) target_sources(pscm PRIVATE ${CODEGEN_SRCS}) set(CMAKE_MODULE_PATH /usr/local/opt/llvm/lib/cmake/clang /usr/local/opt/llvm/lib/cmake/lld /usr/local/opt/llvm/lib/cmake/llvm /usr/local/opt/llvm/lib/cmake/mlir) set(CMAKE_PREFIX_PATH /usr/local/opt/llvm/lib/cmake/clang /usr/local/opt/llvm/lib/cmake/lld /usr/local/opt/llvm/lib/cmake/llvm /usr/local/opt/llvm/lib/cmake/mlir) diff --git a/include/pscm/codegen/codegen.h b/include/pscm/codegen/codegen.h index 441caeff..71e72347 100644 --- a/include/pscm/codegen/codegen.h +++ b/include/pscm/codegen/codegen.h @@ -4,4 +4,5 @@ namespace pscm { std::optional mlir_codegen_and_run_jit(Cell expr); -} \ No newline at end of file +std::optional llvm_ir_codegen_and_run_jit(Cell expr); +} // namespace pscm \ No newline at end of file diff --git a/include/pscm/scm_utils.h b/include/pscm/scm_utils.h index 51962021..edac76c3 100644 --- a/include/pscm/scm_utils.h +++ b/include/pscm/scm_utils.h @@ -123,6 +123,6 @@ class fmt::formatter { auto format(const pscm::Cell& cell, format_context& ctx) const { std::string str; cell.to_string().toUTF8String(str); - return format_to(ctx.out(), "{}", str); + return fmt::format_to(ctx.out(), "{}", str); } }; diff --git a/src/codegen/codegen.cpp b/src/codegen/codegen.cpp index 5fd0d757..75065ee7 100644 --- a/src/codegen/codegen.cpp +++ b/src/codegen/codegen.cpp @@ -24,7 +24,10 @@ using namespace mlir; #include "mlir/Target/LLVMIR/Export.h" #include "mlir/Transforms/Passes.h" +#include "llvm/ADT/ScopedHashTable.h" #include "llvm/ADT/StringRef.h" +#include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h" +#include "llvm/IR/IRBuilder.h" #include "llvm/IR/Module.h" #include "llvm/Support/CommandLine.h" #include "llvm/Support/ErrorOr.h" @@ -33,8 +36,6 @@ using namespace mlir; #include "llvm/Support/TargetSelect.h" #include "llvm/Support/raw_ostream.h" -#include "llvm/ADT/ScopedHashTable.h" - #include "pscm/Number.h" #include "pscm/Symbol.h" #include "pscm/common_def.h" @@ -229,6 +230,12 @@ int create_mlir_add(mlir::MLIRContext& ctx, mlir::OwningOpRef& m } std::optional mlir_codegen_and_run_jit(Cell expr) { + if (auto ret = llvm_ir_codegen_and_run_jit(expr); ret.has_value()) { + llvm::errs() << "get result from llvm IR" + << "\n"; + return ret; + } + mlir::registerAsmPrinterCLOptions(); mlir::registerMLIRContextCLOptions(); mlir::registerPassManagerCLOptions(); diff --git a/src/codegen/llvm_ir/llvm_ir_codegen.cpp b/src/codegen/llvm_ir/llvm_ir_codegen.cpp new file mode 100644 index 00000000..dd9409e6 --- /dev/null +++ b/src/codegen/llvm_ir/llvm_ir_codegen.cpp @@ -0,0 +1,156 @@ + +#include "llvm/ADT/ScopedHashTable.h" +#include "llvm/ADT/StringRef.h" +#include "llvm/ExecutionEngine/Orc/CompileUtils.h" +#include "llvm/ExecutionEngine/Orc/Core.h" +#include "llvm/ExecutionEngine/Orc/ExecutionUtils.h" +#include "llvm/ExecutionEngine/Orc/IRCompileLayer.h" +#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h" +#include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h" +#include "llvm/ExecutionEngine/Orc/ThreadSafeModule.h" +#include "llvm/ExecutionEngine/SectionMemoryManager.h" +#include "llvm/IR/IRBuilder.h" +#include "llvm/IR/Module.h" +#include "llvm/Support/CommandLine.h" +#include "llvm/Support/ErrorOr.h" +#include "llvm/Support/MemoryBuffer.h" +#include "llvm/Support/SourceMgr.h" +#include "llvm/Support/TargetSelect.h" +#include "llvm/Support/raw_ostream.h" + +#include "pscm/Number.h" +#include "pscm/Symbol.h" +#include "pscm/common_def.h" +#include "pscm/scm_utils.h" + +#include +#include +#include +#include +#include +#include + +#include +PSCM_INLINE_LOG_DECLARE("pscm.codegen"); + +namespace pscm { + +llvm::ExitOnError exit_on_err; + +int create_llvm_ir_add(llvm::LLVMContext& ctx, llvm::Module& m) { + std::vector func_args(2, llvm::Type::getInt64Ty(ctx)); + auto func_type = llvm::FunctionType::get(llvm::Type::getInt64Ty(ctx), func_args, false); + auto func = llvm::Function::Create(func_type, llvm::Function::ExternalLinkage, "pscm_jit_add2", &m); + if (!func) { + llvm::errs() << "create Function error" + << "\n"; + return -1; + } + std::vector func_args_names; + func_args_names.push_back("lhs"); + func_args_names.push_back("rhs"); + int idx = 0; + for (auto& arg : func->args()) { + arg.setName(func_args_names[idx++]); + } + + auto basic_block = llvm::BasicBlock::Create(ctx, "entry", func); + auto builder = std::make_unique>(ctx); + builder->SetInsertPoint(basic_block); + auto ret = builder->CreateAdd(func->getArg(0), func->getArg(1), "addtmp"); + builder->CreateRet(ret); + return 0; +} + +std::optional llvm_ir_codegen_and_run_jit(Cell expr) { + llvm::InitializeNativeTarget(); + llvm::InitializeNativeTargetAsmPrinter(); + llvm::InitializeNativeTargetAsmParser(); + + std::cout << "expr: " << expr.to_std_string() << std::endl; + if (!expr.is_pair()) { + return std::nullopt; + } + Cell fakeResult; + // hardcode only suport (+ num1 num2) + if (car(expr).is_sym() && *car(expr).to_sym() == "+"_sym) { + auto num1 = cadr(expr); + auto num2 = caddr(expr); + if (!cdddr(expr).is_nil()) { + return std::nullopt; + } + if (!num1.is_num() || !num2.is_num()) { + return std::nullopt; + } + auto n1 = num1.to_num(); + auto n2 = num2.to_num(); + if (!n1->is_int() || !n2->is_int()) { + return std::nullopt; + } + auto ctx = std::make_unique(); + auto module = std::make_unique("pscm jit", *ctx); + + auto jit_target_machine_builder = llvm::orc::JITTargetMachineBuilder::detectHost(); + if (!jit_target_machine_builder) { + llvm::errs() << "JITTargetMachineBuilder::detectHost() error" + << "\n"; + return std::nullopt; + } + + auto default_data_layout = jit_target_machine_builder->getDefaultDataLayoutForTarget(); + if (!default_data_layout) { + llvm::errs() << "getDefaultDataLayoutForTarget error" + << "\n"; + return std::nullopt; + } + + module->setDataLayout(*default_data_layout); + + if (auto err = create_llvm_ir_add(*ctx, *module)) { + llvm::errs() << "create mlir error" + << "\n"; + return std::nullopt; + } + + auto self_executor_process_control = llvm::orc::SelfExecutorProcessControl::Create(); + if (!self_executor_process_control) { + llvm::errs() << "Could not create SelfExecutorProcessControl" + << "\n"; + return std::nullopt; + } + + auto execution_session = std::make_unique(std::move(*self_executor_process_control)); + auto get_section_memory_manager = []() { + return std::make_unique(); + }; + auto rt_dyld_object_linking_layer = + std::make_unique(*execution_session, get_section_memory_manager); + llvm::orc::IRCompileLayer ir_compile_layer( + *execution_session, *rt_dyld_object_linking_layer, + std::make_unique(std::move(*jit_target_machine_builder))); + + auto& main_jit_dylib = execution_session->createBareJITDylib("main"); + main_jit_dylib.addGenerator(exit_on_err( + llvm::orc::DynamicLibrarySearchGenerator::GetForCurrentProcess(default_data_layout->getGlobalPrefix()))); + + llvm::SMDiagnostic sm_diagnostic; + auto llvm_ctx = std::make_unique(); + llvm::errs() << *module << "\n"; + auto thread_safe_module = llvm::orc::ThreadSafeModule(std::move(module), std::move(llvm_ctx)); + + exit_on_err(ir_compile_layer.add(main_jit_dylib, std::move(thread_safe_module))); + + auto func_sym = exit_on_err(execution_session->lookup({ &main_jit_dylib }, "_pscm_jit_add2")); + auto func_ptr = func_sym.getAddress().toPtr(); + auto ret = func_ptr(n1->to_int(), n2->to_int()); + execution_session->endSession(); + return new pscm::Number(ret); + } + else { + llvm::errs() << "not supported now" + << "\n"; + return std::nullopt; + } +} + +} // namespace pscm