Skip to content

Commit

Permalink
simple llvm kaleidoscope jit setup for testing
Browse files Browse the repository at this point in the history
  • Loading branch information
jeremy-rifkin committed Nov 21, 2024
1 parent 379a0fa commit 73690b7
Show file tree
Hide file tree
Showing 4 changed files with 1,574 additions and 1 deletion.
2 changes: 1 addition & 1 deletion Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ help: # with thanks to Ben Rady
build: debug ## build in debug mode

build/configured-debug:
cmake -S . -B build -GNinja -DCMAKE_BUILD_TYPE=Debug -DCMAKE_EXPORT_COMPILE_COMMANDS=On -DCPPTRACE_BUILD_TESTING=On
cmake -S . -B build -GNinja -DCMAKE_BUILD_TYPE=Debug -DCMAKE_EXPORT_COMPILE_COMMANDS=On -DCPPTRACE_BUILD_TESTING=On -DCPPTRACE_BUILD_SHARED=On -DCMAKE_PREFIX_PATH=~/thirdparty/llvm-project/build/foo
rm -f build/configured-release
touch build/configured-debug

Expand Down
8 changes: 8 additions & 0 deletions test/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -103,3 +103,11 @@ if(NOT CPPTRACE_SKIP_UNIT)
endif()
add_test(NAME unittest COMMAND unittest)
endif()

find_package(LLVM REQUIRED CONFIG)
add_executable(jit_test jit/main.cpp)
add_test_dependencies(jit_test)
target_include_directories(jit_test PUBLIC ${LLVM_INCLUDE_DIRS})
target_compile_definitions(jit_test PUBLIC ${LLVM_DEFINITIONS})
llvm_map_components_to_libnames(llvm_libs support core orcjit native)
target_link_libraries(jit_test PUBLIC ${llvm_libs})
101 changes: 101 additions & 0 deletions test/jit/KaleidoscopeJIT.hpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,101 @@
//===- 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/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 llvm {
namespace orc {

class KaleidoscopeJIT {
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())));
}

~KaleidoscopeJIT() {
if (auto Err = ES->endSession())
ES->reportError(std::move(Err));
}

static 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; }

Error addModule(ThreadSafeModule TSM, ResourceTrackerSP RT = nullptr) {
if (!RT)
RT = MainJD.getDefaultResourceTracker();
return CompileLayer.add(RT, std::move(TSM));
}

Expected<ExecutorSymbolDef> lookup(StringRef Name) {
return ES->lookup({&MainJD}, Mangle(Name.str()));
}
};

} // end namespace orc
} // end namespace llvm

#endif // LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H
Loading

0 comments on commit 73690b7

Please sign in to comment.