roland/xo-alloc
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

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