removes segfault caused by wrong object destruction order, adds global optimizations, toy JIT

Author: Akshay Deodhar

KaleidoscopeJIT.h

@@ -0,0 +1,100 @@
+//===- 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/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<JITEvaluatedSymbol> lookup(StringRef Name) {
+    return ES->lookup({&MainJD}, Mangle(Name.str()));
+  }
+};
+
+} // end namespace orc
+} // end namespace llvm
+
+#endif // LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H

Makefile

@@ -1,8 +1,8 @@
 kaleidoscope: kaleidoscope.cpp
-	clang++ -g3 -Wall kaleidoscope.cpp `llvm-config --cxxflags --ldflags --system-libs --libs core` -o kaleidoscope
+	clang++ -g3 -Wall kaleidoscope.cpp `llvm-config --cxxflags --ldflags --system-libs --libs core orcjit native` -o kaleidoscope
 
 theirkaleidoscope: theirkaleidoscope.cpp
-	clang++ -g3 -Wall theirkaleidoscope.cpp `llvm-config --cxxflags --ldflags --system-libs --libs core` -o theirkaleidoscope
+	clang++ -g3 -Wall theirkaleidoscope.cpp `llvm-config --cxxflags --ldflags --system-libs --libs core orcjit native` -o theirkaleidoscope
 
 test: kaleidoscope fib.k
 	cat fib.k | ./kaleidoscope

kaleidoscope.cpp

@@ -6,9 +6,15 @@
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/LegacyPassManager.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/Verifier.h"
+#include "llvm/Support/TargetSelect.h"
+#include "llvm/Target/TargetMachine.h"
+#include "llvm/Transforms/InstCombine/InstCombine.h"
+#include "llvm/Transforms/Scalar.h"
+#include "llvm/Transforms/Scalar/GVN.h"
 #include <iostream>
 #include <string>
 #include <vector>
@@ -18,6 +24,7 @@
 #include <cstdio>
 #include <cstdlib>
 #include <unordered_map>
+#include "./KaleidoscopeJIT.h"
 
 using namespace llvm;
 
@@ -107,9 +114,10 @@ static int gettok(void){
 
 
 // State variables
-static std::unique_ptr<IRBuilder<>> Builder; // for creating instructions, constants, etc
-static std::unique_ptr<Module> TheModule; // to hold blocks, definitions? (TODO), TODO: why does this have to be a pointer?
 static std::unique_ptr<LLVMContext> TheContext; 
+static std::unique_ptr<Module> TheModule; // to hold blocks, definitions? (TODO), TODO: why does this have to be a pointer?
+static std::unique_ptr<IRBuilder<>> Builder; // for creating instructions, constants, etc
+static std::unique_ptr<legacy::FunctionPassManager> TheFPM; // Function pass manager
 static std::unordered_map<std::string, Value *> Symbols; // Maps names inside function context to LLVM "values"
 
 
@@ -679,12 +687,6 @@ Value* VariableExprAST::codegen() {
 // Generates code for function call, returns `Value` of function call
 Value* CallExprAST::codegen() {
 
-	// Generate code for arguments, and get their values
-	std::vector<Value *> Argvec;
-	for (const auto& arg: Args) {
-		Argvec.push_back(arg->codegen());
-	}
-
 	// Obtain function with name `Callee` from Module
 	Function *func = TheModule->getFunction(Callee);
 	if (!func) {
@@ -699,16 +701,24 @@ Value* CallExprAST::codegen() {
 		);
 	}
 
+	// Generate code for arguments, and get their values
+	std::vector<Value *> Argvec;
+	for (const auto& arg: Args) {
+		Argvec.push_back(arg->codegen());
+	}
+
+
+
 	// TODO: why do I need to provide TheContext to getDoubleTy?
-	std::vector<Type *> ArgTypes = std::vector<Type *>(Args.size(), Builder->getDoubleTy());
+	//std::vector<Type *> ArgTypes = std::vector<Type *>(Args.size(), Builder->getDoubleTy());
 
 	// Create type for call
 	// TODO: this is not needed: CreateCall can be called without a type
 	// Why is this so? Is it because the function does not take variable arguments?
 	//FunctionType *func_type = FunctionType::get(Builder->getDoubleTy(), ArgTypes, false);
 
 	// Create call
-	return Builder->CreateCall(func, Argvec, Callee);
+	return Builder->CreateCall(func, Argvec, "call");
 }
 
 Value* BinaryExprAST::codegen() {
@@ -742,6 +752,7 @@ Value* BinaryExprAST::codegen() {
 			return Builder->CreateUIToFP(L, Builder->getDoubleTy(), "booltofp");
 		default:
 			return LogErrorV("Invalid Operator");
+			break;
 	}
 }
 
@@ -797,6 +808,9 @@ Function* FunctionAST::codegen() {
 		// TODO: What if this check fails? Do I still continue?
 		verifyFunction(*func);
 
+		// Run passes on function
+		TheFPM->run(*func);
+
 		return func;
 	}
 
@@ -873,15 +887,32 @@ static void HandleTopLevelExpression() {
 		}
 }
 
-static void InitializeModule() {
+static void InitializeModuleAndPassManager() {
 	TheContext = std::make_unique<LLVMContext>();
 	TheModule = std::make_unique<Module>("kaleidoscope", *TheContext);
+	// Why .get? Ahh- I want to pass a pointer. What about uniqueness?
+	TheFPM = std::make_unique<legacy::FunctionPassManager>(TheModule.get());
 	Builder = std::make_unique<IRBuilder<>>(*TheContext);
+
+	// Peephole optimizations
+	TheFPM->add(createInstructionCombiningPass());
+	
+	// ?
+	TheFPM->add(createReassociatePass());
+
+	// Global value numbering-> common subexpression elimination. Global is actually per-function
+	TheFPM->add(createGVNPass());
+
+	// Dead code elimination pass;
+	TheFPM->add(createCFGSimplificationPass());
+
+	// Run initalizers for all passes added to pass manager
+	TheFPM->doInitialization();
 }
 
 // top = definition | expression | external | ;
 static void MainLoop() {
-	while(1) {
+	while(true) {
 		fprintf(stderr, "ready>");
 		switch (CurTok) {
 				case tok_eof:
@@ -940,13 +971,13 @@ int main(void) {
 		BinopPrecedence['*'] = 40;
 		BinopPrecedence['/'] = 40;
 
-#if IRGEN
-		InitializeModule();
-#endif
-
 		fprintf(stderr, "ready>");
 		getNextToken();
 
+#if IRGEN
+		InitializeModuleAndPassManager();
+#endif
+
 		MainLoop();
 		//oldmain();