LLVM 8 and later ORC JIT problems with math library - llvm

Question

Since LLVM 8 (including current LLVM trunk, aka LLVM 9) I run into a problem when using the ORC JIT ExecutionEngine with functions which contain a call to the standard math library.
The JIT compiler is able to find the symbol for the function but fails to get the address for it in case the function makes a call to the math library.
I attach a simple exploit that shows the problem. The program test.cc reads in an IR file which contains a single function in LLVM's Intermediate Representation:
The function takes a single argument, a floating point number, and returns in case of
"func_works.ll" the argument itself, and in case of
"func_cos_fails.ll" the cosine of the argument.
I didn't implement the choice between the two files at runtime, so the program needs to be rebuilt when switching to the other case.
The program uses the standard KaleidoscopeJIT.h that comes with LLVM (except that I had to expose the Datalayout).
If you build the program with the "func_works.ll" and run it, the program succeeds with:
symbol found!
address found!
If you build the program with the "func_cos_fails.ll" and run it, the program fails with:
symbol found!
Failure value returned from cantFail wrapped call
UNREACHABLE executed at install/llvm-8.0-x86-debug/include/llvm/Support/Error.h:732!
This happens with LLVM 8 release and the current LLVM trunk.
Does anyone see what's going on?
This test was run on an x86 Linux Ubuntu system with LLVM configured with
cmake -G "Unix Makefiles" \
-DBUILD_SHARED_LIBS="ON" \
-DLLVM_ENABLE_RTTI="ON" \
-DLLVM_ENABLE_ZLIB="OFF" \
-DLLVM_ENABLE_TERMINFO="OFF" \
-DCMAKE_BUILD_TYPE="Debug" \
-DCMAKE_INSTALL_PREFIX=$CMAKE_INSTALL_PREFIX \
-DLLVM_TARGETS_TO_BUILD="X86" \
$SRC
test.cc:
#include "KaleidoscopeJIT.h"
#include "llvm/Analysis/Passes.h"
#include "llvm/ExecutionEngine/ExecutionEngine.h"
#include "llvm/ExecutionEngine/MCJIT.h"
#include "llvm/ExecutionEngine/ObjectCache.h"
#include "llvm/ExecutionEngine/SectionMemoryManager.h"
#include "llvm/IR/DataLayout.h"
#include "llvm/IR/DerivedTypes.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/LegacyPassManager.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/Verifier.h"
#include "llvm/IRReader/IRReader.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/Path.h"
#include "llvm/Support/SourceMgr.h"
#include "llvm/Support/TargetSelect.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Transforms/Scalar.h"
#include "llvm/Support/InitLLVM.h"
#include <iostream>
using namespace llvm;
using namespace llvm::orc;
int main(int argc, char **argv) {
InitLLVM X(argc, argv);
EnableDebugBuffering = true;
LLVMContext Context;
InitializeNativeTarget();
InitializeNativeTargetAsmPrinter();
InitializeNativeTargetAsmParser();
cl::ParseCommandLineOptions(argc, argv, "Kaleidoscope example program\n");
SMDiagnostic Err;
std::unique_ptr<llvm::Module> M = parseIRFile( "func_cos_fails.ll" , Err, Context, false);
//std::unique_ptr<llvm::Module> M = parseIRFile( "func_works.ll" , Err, Context, false);
if (!M) {
Err.print("IR parsing failed: ", errs());
return 0;
}
std::unique_ptr<KaleidoscopeJIT> TheJIT;
TheJIT = llvm::make_unique<KaleidoscopeJIT>();
auto H = TheJIT->addModule(std::move(M));
std::string MangledName;
llvm::raw_string_ostream MangledNameStream(MangledName);
llvm::Mangler::getNameWithPrefix(MangledNameStream, "func_ir" , TheJIT->getDL() );
if (auto Sym = TheJIT->findSymbol(MangledNameStream.str()))
{
std::cout << "symbol found!\n";
void* fptr = (void *)cantFail(Sym.getAddress());
std::cout << "address found!\n";
}
else
{
std::cout << "symbol not found!\n";
}
return 0;
}
func_cos_fails.ll:
source_filename = "module"
target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
declare float #cosf(float)
define float #func_ir(float %arg0) {
entrypoint:
%0 = call float #cosf(float %arg0)
ret float %0
}
func_works.ll:
source_filename = "module"
target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
define float #func_ir(float %arg0) {
entrypoint:
ret float %arg0
}
KaleidoscopeJIT.h:
#ifndef LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H
#define LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/iterator_range.h"
#include "llvm/ExecutionEngine/ExecutionEngine.h"
#include "llvm/ExecutionEngine/JITSymbol.h"
#include "llvm/ExecutionEngine/Orc/CompileUtils.h"
#include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
#include "llvm/ExecutionEngine/Orc/LambdaResolver.h"
#include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
#include "llvm/ExecutionEngine/RTDyldMemoryManager.h"
#include "llvm/ExecutionEngine/SectionMemoryManager.h"
#include "llvm/IR/DataLayout.h"
#include "llvm/IR/Mangler.h"
#include "llvm/Support/DynamicLibrary.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Target/TargetMachine.h"
#include <algorithm>
#include <map>
#include <memory>
#include <string>
#include <vector>
namespace llvm {
namespace orc {
class KaleidoscopeJIT {
public:
using ObjLayerT = LegacyRTDyldObjectLinkingLayer;
using CompileLayerT = LegacyIRCompileLayer<ObjLayerT, SimpleCompiler>;
KaleidoscopeJIT()
: Resolver(createLegacyLookupResolver(
ES,
[this](const std::string &Name) {
return ObjectLayer.findSymbol(Name, true);
},
[](Error Err) { cantFail(std::move(Err), "lookupFlags failed"); })),
TM(EngineBuilder().selectTarget()), DL(TM->createDataLayout()),
ObjectLayer(ES,
[this](VModuleKey) {
return ObjLayerT::Resources{
std::make_shared<SectionMemoryManager>(), Resolver};
}),
CompileLayer(ObjectLayer, SimpleCompiler(*TM)) {
llvm::sys::DynamicLibrary::LoadLibraryPermanently(nullptr);
}
TargetMachine &getTargetMachine() { return *TM; }
VModuleKey addModule(std::unique_ptr<Module> M) {
auto K = ES.allocateVModule();
cantFail(CompileLayer.addModule(K, std::move(M)));
ModuleKeys.push_back(K);
return K;
}
void removeModule(VModuleKey K) {
ModuleKeys.erase(find(ModuleKeys, K));
cantFail(CompileLayer.removeModule(K));
}
JITSymbol findSymbol(const std::string Name) {
return findMangledSymbol(mangle(Name));
}
const DataLayout& getDL() const {
return DL;
}
private:
std::string mangle(const std::string &Name) {
std::string MangledName;
{
raw_string_ostream MangledNameStream(MangledName);
Mangler::getNameWithPrefix(MangledNameStream, Name, DL);
}
return MangledName;
}
JITSymbol findMangledSymbol(const std::string &Name) {
#ifdef _WIN32
// The symbol lookup of ObjectLinkingLayer uses the SymbolRef::SF_Exported
// flag to decide whether a symbol will be visible or not, when we call
// IRCompileLayer::findSymbolIn with ExportedSymbolsOnly set to true.
//
// But for Windows COFF objects, this flag is currently never set.
// For a potential solution see: https://reviews.llvm.org/rL258665
// For now, we allow non-exported symbols on Windows as a workaround.
const bool ExportedSymbolsOnly = false;
#else
const bool ExportedSymbolsOnly = true;
#endif
// Search modules in reverse order: from last added to first added.
// This is the opposite of the usual search order for dlsym, but makes more
// sense in a REPL where we want to bind to the newest available definition.
for (auto H : make_range(ModuleKeys.rbegin(), ModuleKeys.rend()))
if (auto Sym = CompileLayer.findSymbolIn(H, Name, ExportedSymbolsOnly))
return Sym;
// If we can't find the symbol in the JIT, try looking in the host process.
if (auto SymAddr = RTDyldMemoryManager::getSymbolAddressInProcess(Name))
return JITSymbol(SymAddr, JITSymbolFlags::Exported);
#ifdef _WIN32
// For Windows retry without "_" at beginning, as RTDyldMemoryManager uses
// GetProcAddress and standard libraries like msvcrt.dll use names
// with and without "_" (for example "_itoa" but "sin").
if (Name.length() > 2 && Name[0] == '_')
if (auto SymAddr =
RTDyldMemoryManager::getSymbolAddressInProcess(Name.substr(1)))
return JITSymbol(SymAddr, JITSymbolFlags::Exported);
#endif
return nullptr;
}
ExecutionSession ES;
std::shared_ptr<SymbolResolver> Resolver;
std::unique_ptr<TargetMachine> TM;
const DataLayout DL;
ObjLayerT ObjectLayer;
CompileLayerT CompileLayer;
std::vector<VModuleKey> ModuleKeys;
};
} // end namespace orc
} // end namespace llvm
#endif // LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H
for convenience I provide a Makefile:
LLVM_CONFIG = ${LLVM_INSTALL_PATH}
LLVM_CXXFLAGS = $(shell $(LLVM_CONFIG) --cxxflags)
LLVM_LDFLAGS = $(shell $(LLVM_CONFIG) --ldflags)
LLVM_LIBS = $(shell $(LLVM_CONFIG) --libs)
all: test
test.o: test.cc KaleidoscopeJIT.h
g++ -c -o $# $< $(LLVM_CXXFLAGS)
test: test.o
g++ -o $# $< $(LLVM_LDFLAGS) $(LLVM_LIBS)
clean:
rm -f *.o
rm -f test

Answer 1

I believe the solution to this (for llvm 7 and 8 anyway) was found here:
https://stackoverflow.com/a/56862433/2310373
Namely, replace:
[this](const std::string &Name) {
return ObjectLayer.findSymbol(Name, true);
},
With something like
[this](const std::string &Name) {
auto FoundSymbol = ObjectLayer.findSymbol(Name, true);
if (!FoundSymbol) {
if (auto SymAddr = RTDyldMemoryManager::getSymbolAddressInProcess(Name))
return JITSymbol(SymAddr, JITSymbolFlags::Exported);
}
return FoundSymbol;
},