roland/xo-expression2
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

xo-jit: incorporate kaleidoscope jit for codegen

Browse source
This commit is contained in:
Roland Conybeare 2024-06-15 13:17:14 -04:00
commit fa0104422f
5 changed files with 319 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

12
example/ex_kaleidoscope4/CMakeLists.txt Normal file
View file

@ -0,0 +1,12 @@
# xo-jit/example/ex1/CMakeLists.txt
set(SELF_EXE xo_kaleidoscope4)
set(SELF_SRCS ex_kaleidoscope4.cpp)
if (XO_ENABLE_EXAMPLES)
xo_add_executable(${SELF_EXE} ${SELF_SRCS})
xo_self_dependency(${SELF_EXE} xo_jit)
#xo_dependency(${SELF_EXE} xo_expression)
endif()
# end CMakeLists.txt

16
example/ex_kaleidoscope4/ex_kaleidoscope4.cpp Normal file
View file

@ -0,0 +1,16 @@
/** ex_kaleidoscop4.cpp **/
#include "xo/jit/KaleidoscopeJit.hpp"
#include <iostream>
int
main() {
using std::cerr;
using std::endl;
auto jit = xo::jit::KaleidoscopeJIT::Create();
cerr << "created kaleidoscope jit successfully" << endl;
}
/** end ex_kaleidoscope4.cpp **/

50
include/xo/jit/Jit.hpp
View file

@ -14,6 +14,27 @@
#include "xo/expression/Apply.hpp"
#include "xo/expression/Lambda.hpp"
#include "xo/expression/Variable.hpp"
#include "KaleidoscopeJit.hpp"
#ifdef NOT_USING
/* stuff from KaleidoscopeJIT.hpp */
#include "llvm/ADT/StringRef.h"
#include "llvm/ExecutionEngine/JITSymbol.h"
#include "llvm/ExecutionEngine/Orc/CompileUtils.h"
#include "llvm/ExecutionEngine/Orc/Core.h"
#include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
#include "llvm/ExecutionEngine/Orc/ExecutorProcessControl.h"
#include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
#include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
#include "llvm/ExecutionEngine/Orc/Shared/ExecutorSymbolDef.h"
#include "llvm/ExecutionEngine/SectionMemoryManager.h"
#include "llvm/IR/DataLayout.h"
#include "llvm/IR/LLVMContext.h"
#include <memory>
#endif
/* stuff from kaleidoscope.cpp */
#include "llvm/ADT/APFloat.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/IR/BasicBlock.h"
@ -36,6 +57,7 @@
#include "llvm/Transforms/Scalar/Reassociate.h"
#include "llvm/Transforms/Scalar/SimplifyCFG.h"
namespace xo {
namespace jit {
/** @class Jit
@ -49,7 +71,8 @@ namespace xo {
//using ConstantInterface = xo::ast::ConstantInterface;
public:
static ref::rp<Jit> make() { return new Jit(); }
/* tracking KaleidoscopeJIT::Create() here.. */
static llvm::Expected<std::unique_ptr<Jit>> make_aux();
llvm::Value * codegen_constant(ref::brw<xo::ast::ConstantInterface> expr);
llvm::Function * codegen_primitive(ref::brw<xo::ast::PrimitiveInterface> expr);
@ -63,9 +86,32 @@ namespace xo {
virtual std::string display_string() const;
private:
Jit();
Jit(
std::unique_ptr<KaleidoscopeJIT> kal_jit
#ifdef NOT_USING
std::unique_ptr<llvm::orc::ExecutionSession> es,
llvm::orc::JITTargetMachineBuilder jtmb,
llvm::DataLayout dl
#endif
);
private:
// ----- this part adapted from LLVM 19.0 KaleidoscopeJIT.hpp [wip] -----
std::unique_ptr<KaleidoscopeJIT> kal_jit_;
#ifdef NOT_USING
std::unique_ptr<llvm::orc::ExecutionSession> jit_es_;
llvm::DataLayout jit_data_layout_;
llvm::orc::MangleAndInterner jit_mangle_;
llvm::orc::RTDyldObjectLinkingLayer jit_object_layer_;
llvm::orc::IRCompileLayer jit_compile_layer_;
/** reference here. looks like storage owned by .jit_es **/
llvm::orc::JITDylib & jit_our_dynamic_lib_;
#endif
// ----- this part adapted from kaleidoscope.cpp -----
/** owns + manages core "global" llvm data,
* including type- and constant- unique-ing tables.
*

126
include/xo/jit/KaleidoscopeJit.hpp Normal file
View file

@ -0,0 +1,126 @@
//===- KaleidoscopeJIT.h - A simple JIT for Kaleidoscope --------*- C++ -*-===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// Contains a simple JIT definition for use in the kaleidoscope tutorials.
//
//===----------------------------------------------------------------------===//
#ifndef LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H
#define LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H
#include "llvm/ADT/StringRef.h"
#include "llvm/ExecutionEngine/JITSymbol.h"
#include "llvm/ExecutionEngine/Orc/CompileUtils.h"
#include "llvm/ExecutionEngine/Orc/Core.h"
#include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
#include "llvm/ExecutionEngine/Orc/ExecutorProcessControl.h"
#include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
#include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
#include "llvm/ExecutionEngine/Orc/Shared/ExecutorSymbolDef.h"
#include "llvm/ExecutionEngine/SectionMemoryManager.h"
#include "llvm/IR/DataLayout.h"
#include "llvm/IR/LLVMContext.h"
#include <memory>
namespace xo {
namespace jit {
class KaleidoscopeJIT {
private:
using StringRef = llvm::StringRef;
using SectionMemoryManager = llvm::SectionMemoryManager;
using DynamicLibrarySearchGenerator = llvm::orc::DynamicLibrarySearchGenerator;
using ConcurrentIRCompiler = llvm::orc::ConcurrentIRCompiler;
using ExecutionSession = llvm::orc::ExecutionSession;
using DataLayout = llvm::DataLayout;
using MangleAndInterner = llvm::orc::MangleAndInterner;
using RTDyldObjectLinkingLayer = llvm::orc::RTDyldObjectLinkingLayer;
using IRCompileLayer = llvm::orc::IRCompileLayer;
using JITDylib = llvm::orc::JITDylib;
using JITTargetMachineBuilder = llvm::orc::JITTargetMachineBuilder;
using ThreadSafeModule = llvm::orc::ThreadSafeModule;
using ResourceTrackerSP = llvm::orc::ResourceTrackerSP;
using ExecutorSymbolDef = llvm::orc::ExecutorSymbolDef;
using SelfExecutorProcessControl = llvm::orc::SelfExecutorProcessControl;
private:
std::unique_ptr<ExecutionSession> ES;
DataLayout DL;
MangleAndInterner Mangle;
RTDyldObjectLinkingLayer ObjectLayer;
IRCompileLayer CompileLayer;
JITDylib &MainJD;
public:
KaleidoscopeJIT(std::unique_ptr<ExecutionSession> ES,
JITTargetMachineBuilder JTMB,
DataLayout DL)
: ES(std::move(ES)),
DL(std::move(DL)),
Mangle(*this->ES, this->DL),
ObjectLayer(*this->ES,
[]() { return std::make_unique<SectionMemoryManager>(); }),
CompileLayer(*this->ES, ObjectLayer,
std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
MainJD(this->ES->createBareJITDylib("<main>"))
{
MainJD.addGenerator(
cantFail(DynamicLibrarySearchGenerator::GetForCurrentProcess(
DL.getGlobalPrefix())));
if (JTMB.getTargetTriple().isOSBinFormatCOFF()) {
ObjectLayer.setOverrideObjectFlagsWithResponsibilityFlags(true);
ObjectLayer.setAutoClaimResponsibilityForObjectSymbols(true);
}
}
~KaleidoscopeJIT() {
if (auto Err = ES->endSession())
ES->reportError(std::move(Err));
}
static llvm::Expected<std::unique_ptr<KaleidoscopeJIT>> Create() {
auto EPC = SelfExecutorProcessControl::Create();
if (!EPC)
return EPC.takeError();
auto ES = std::make_unique<ExecutionSession>(std::move(*EPC));
JITTargetMachineBuilder JTMB(
ES->getExecutorProcessControl().getTargetTriple());
auto DL = JTMB.getDefaultDataLayoutForTarget();
if (!DL)
return DL.takeError();
return std::make_unique<KaleidoscopeJIT>(std::move(ES), std::move(JTMB),
std::move(*DL));
}
const DataLayout &getDataLayout() const { return DL; }
JITDylib &getMainJITDylib() { return MainJD; }
llvm::Error addModule(ThreadSafeModule TSM, ResourceTrackerSP RT = nullptr) {
if (!RT)
RT = MainJD.getDefaultResourceTracker();
return CompileLayer.add(RT, std::move(TSM));
}
llvm::Expected<ExecutorSymbolDef> lookup(StringRef Name) {
return ES->lookup({&MainJD}, Mangle(Name.str()));
}
};
} // end namespace jit
} // end namespace xo
#endif // LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H

126
src/jit/Jit.cpp
View file

@ -15,11 +15,96 @@ namespace xo {
using std::endl;
namespace jit {
Jit::Jit()
: llvm_cx_{std::make_unique<llvm::LLVMContext>()},
/* tracking KaleidoscopeJIT::Create() here..
*
* Verified:
* + 'execution session' as per Kaleidoscope JIT,
* can instantiate from python
* + 'jit object layer'
* (realtime dynamic library object linking layer)
* + 'jit_copmile_layer'
* + 'jit_our_dynamic_lib'
*/
llvm::Expected<std::unique_ptr<Jit>>
Jit::make_aux()
{
#ifdef NOT_USING
/* 'executor process control' */
auto epc = llvm::orc::SelfExecutorProcessControl::Create();
if (!epc) {
return epc.takeError();
//throw std::runtime_error("Jit::make: failed to establish executor process control");
}
/* 'execution session' */
auto es = std::make_unique<llvm::orc::ExecutionSession>(std::move(*epc));
/* 'jit target machine builder' */
llvm::orc::JITTargetMachineBuilder jtmb(es
->getExecutorProcessControl()
.getTargetTriple());
auto dl = jtmb.getDefaultDataLayoutForTarget();
if (!dl) {
return dl.takeError();
//throw std::runtime_error("Jit::make: failed to establish data layout object"
// " for target machine");
}
#endif
static llvm::ExitOnError llvm_exit_on_err;
std::unique_ptr<KaleidoscopeJIT> kal_jit = llvm_exit_on_err(KaleidoscopeJIT::Create());
return std::unique_ptr<Jit>(new Jit(std::move(kal_jit)
#ifdef NOT_USING
std::move(es),
std::move(jtmb),
std::move(*dl)
#endif
));
} /*make*/
Jit::Jit(
std::unique_ptr<KaleidoscopeJIT> kal_jit
#ifdef NOT_USING
std::unique_ptr<llvm::orc::ExecutionSession> jit_es,
llvm::orc::JITTargetMachineBuilder jtmb,
llvm::DataLayout dl
#endif
)
: kal_jit_{std::move(kal_jit)},
#ifdef NOT_USING
jit_es_(std::move(jit_es)),
jit_data_layout_(std::move(dl)),
jit_mangle_(*this->jit_es_, this->jit_data_layout_),
jit_object_layer_(*this->jit_es_,
[]() { return std::make_unique<llvm::SectionMemoryManager>(); }),
jit_compile_layer_(*this->jit_es_,
jit_object_layer_,
std::make_unique<llvm::orc::ConcurrentIRCompiler>(std::move(jtmb))),
/* note: string passed to createBareJITDyLib must be unique
* (within running process address space?)
*/
jit_our_dynamic_lib_(this->jit_es_->createBareJITDylib("<xojitlib>")), /*was MainJD*/
#endif
llvm_cx_{std::make_unique<llvm::LLVMContext>()},
llvm_ir_builder_{std::make_unique<llvm::IRBuilder<>>(*llvm_cx_)},
llvm_module_{std::make_unique<llvm::Module>("xojit", *llvm_cx_)}
{}
{
#ifdef NOT_USING
jit_our_dynamic_lib_.addGenerator
(cantFail(llvm::orc::DynamicLibrarySearchGenerator::GetForCurrentProcess
(jit_data_layout_.getGlobalPrefix())));
if(jtmb.getTargetTriple().isOSBinFormatCOFF()) {
jit_object_layer_.setOverrideObjectFlagsWithResponsibilityFlags(true);
jit_object_layer_.setAutoClaimResponsibilityForObjectSymbols(true);
}
#endif
}
llvm::Value *
Jit::codegen_constant(ref::brw<ConstantInterface> expr)
@ -40,9 +125,12 @@ namespace xo {
llvm::Function *
Jit::codegen_primitive(ref::brw<PrimitiveInterface> expr)
{
constexpr bool c_debug_flag = true;
using xo::scope;
/** note: documentation (such as it is) for llvm::Function here:
*
* https://llvm.org/doxygen/classllvm_1_1Function.html
* https://llvm.org/doxygenL/classllvm_1_1Function.html
**/
auto * fn = llvm_module_->getFunction(expr->name());
@ -64,12 +152,19 @@ namespace xo {
TypeDescr fn_td = expr->value_td();
int n_fn_arg = fn_td->n_fn_arg();
std::vector<llvm::Type *> llvm_argtype_v(n_fn_arg);
scope log(XO_DEBUG(c_debug_flag),
xtag("fn_td", fn_td->short_name()),
xtag("n_fn_arg", n_fn_arg));
std::vector<llvm::Type *> llvm_argtype_v;
llvm_argtype_v.reserve(n_fn_arg);
/** check function args are all doubles **/
for (int i = 0; i < n_fn_arg; ++i) {
TypeDescr arg_td = fn_td->fn_arg(i);
log && log(xtag("i_arg", i), xtag("arg_td", arg_td->short_name()));
if (arg_td->is_native<double>()) {
llvm_argtype_v.push_back(llvm::Type::getDoubleTy(*llvm_cx_));
@ -87,6 +182,9 @@ namespace xo {
//std::vector<llvm::Type *> double_v(n_fn_arg, llvm::Type::getDoubleTy(*llvm_cx_));
TypeDescr retval_td = fn_td->fn_retval();
log && log(xtag("retval_td", retval_td->short_name()));
llvm::Type * llvm_retval = nullptr;
if (retval_td->is_native<double>()) {
@ -108,12 +206,22 @@ namespace xo {
expr->name(),
llvm_module_.get());
// set names for arguments (for diagonostics?). Money-see-kaleidoscope-monkey-do here
#ifdef NOT_USING
for (auto & arg : fn->args())
arg.setName(formalnameofthisarg);
// set names for arguments (for diagnostics?). Monkey-see-kaleidoscope-monkey-do here
{
int i_arg = 0;
for (auto & arg : fn->args()) {
std::stringstream ss;
ss << "x_" << i_arg;
arg.setName(ss.str());
++i_arg;
}
}
#endif
log && log("returning llvm function");
return fn;
} /*codegen_primitive*/
@ -189,7 +297,7 @@ namespace xo {
/* establish prototype for this function */
// PLACEHOLDER
// just make prototype for function :: double -> double
// just handle double arguments + return type for now
std::vector<llvm::Type *> double_v(1, llvm::Type::getDoubleTy(*llvm_cx_));