Add LLVM support for JIT compiling (#174)

* add llvm to cmake scripts and add Dockerfile

* add JIT compiler

* add scalar JIT type tests

* add pointer type JIT tests

* add loops to jit function class

* add JIT function tests

* fix JIT loop bug

* Update include/micm/jit/jit_function.hpp

Co-authored-by: Kyle Shores <kyle.shores44@gmail.com>

* Update cmake/test_util.cmake

Co-authored-by: Kyle Shores <kyle.shores44@gmail.com>

* add assert on regenerating JIT function

---------

Co-authored-by: Kyle Shores <kyle.shores44@gmail.com>

.github/workflows/test.yml

@@ -14,10 +14,11 @@ jobs:
         dockerfile:
           - Dockerfile
           - Dockerfile.coverage
+          - Dockerfile.llvm
           - Dockerfile.memcheck
           - Dockerfile.no_json
           - Dockerfile.nvhpc
-          - Dockerfile.intel
+          # - Dockerfile.intel # intel image is too large for GH action
           # - Dockerfile.openmp
           # - Dockerfile.mpi
     steps:

CMakeLists.txt

@@ -40,6 +40,7 @@ option(ENABLE_JSON "Enable json configureation file reading" ON)
 option(BUILD_DOCS "Build the documentation" OFF)
 option(ENABLE_CUDA "Build with Cuda support" OFF)
 option(ENABLE_OPENACC "Build with OpenACC Support" OFF)
+option(ENABLE_LLVM "Build with LLVM support for JIT-compiling" OFF)
 
 include(CMakeDependentOption)
 # Option to collect custom OpenACC flags

Dockerfile.coverage

@@ -7,6 +7,8 @@ RUN dnf -y update \
         git \
         lcov \
         make \
+        zlib-devel \
+        llvm-devel \
         openmpi-devel \
     && dnf clean all
 
@@ -24,6 +26,7 @@ RUN mkdir /build \
         -D CMAKE_BUILD_TYPE=debug \
         -D ENABLE_CLANG_TIDY:BOOL=FALSE \
         -D ENABLE_COVERAGE:BOOL=TRUE \
+        -D ENABLE_LLVM:BOOL=TRUE \
         # -D ENABLE_MPI:BOOL=TRUE \
         # -D ENABLE_OPENMP:BOOL=TRUE \
         ../micm \

Dockerfile.llvm

@@ -0,0 +1,29 @@
+FROM fedora:37
+
+RUN dnf -y update \
+    && dnf -y install \
+        cmake \
+        gcc-c++ \
+        gdb \
+        git \
+        make \
+        zlib-devel \
+        llvm-devel \
+        valgrind \
+    && dnf clean all
+
+# copy the MICM code
+COPY . /micm/
+
+# build the library and run the tests
+RUN mkdir /build \
+      && cd /build \
+      && cmake \
+        -D CMAKE_BUILD_TYPE=debug \
+        -D ENABLE_CLANG_TIDY:BOOL=FALSE \
+        -D ENABLE_LLVM:BOOL=TRUE \
+        -D ENABLE_MEMCHECK:BOOL=TRUE \
+        ../micm \
+      && make install -j 8
+
+WORKDIR /build

cmake/dependencies.cmake

@@ -92,3 +92,27 @@ endif()
 if(ENABLE_OPENACC)
   find_package(OpenACC REQUIRED)
 endif()
+
+################################################################################
+# LLVM Support
+#
+# TODO: Try to use fetch content for LLVM libraries
+
+if(ENABLE_LLVM)
+  find_package(LLVM REQUIRED CONFIG)
+  if(LLVM_FOUND)
+    message(STATUS "Found LLVM ${LLVM_PACKAGE_VERSION}")
+    message(STATUS "Using LLVMConfig.cmake in: ${LLVM_DIR}")
+
+    include_directories(${LLVM_INCLUDE_DIRS})
+    separate_arguments(LLVM_DEFINITIONS_LIST NATIVE_COMMAND ${LLVM_DEFINITIONS})
+    add_definitions(${LLVM_DEFINITIONS_LIST})
+
+    llvm_map_components_to_libnames(llvm_libs support core orcjit native irreader)
+  else()
+    set(LLVM_CMD "llvm-config --cxxflags --ldflags --system-libs --libs support core orcjit native irreader | tr '\\n' ' '")
+    execute_process(COMMAND bash "-c" ${LLVM_CMD}
+                    OUTPUT_VARIABLE llvm_libs)
+    separate_arguments(llvm_libs)
+  endif()
+endif()

cmake/test_util.cmake

@@ -27,7 +27,10 @@ function(create_standard_test)
 
   if(ENABLE_JSON)
     target_link_libraries(test_${TEST_NAME} PRIVATE nlohmann_json::nlohmann_json)
-    target_compile_definitions(test_${TEST_NAME} PUBLIC USE_JSON)
+  endif()
+
+  if(ENABLE_LLVM)
+    target_link_libraries(test_${TEST_NAME} PRIVATE ${llvm_libs})
   endif()
 
   if(NOT DEFINED TEST_WORKING_DIRECTORY)

include/micm/jit/jit_compiler.hpp

@@ -0,0 +1,149 @@
+// Copyright (C) 2023 National Center for Atmospheric Research,
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+// Based on examples from 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
+#pragma once
+
+#include <memory>
+
+#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/IRTransformLayer.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 "llvm/IR/LegacyPassManager.h"
+#include "llvm/Support/TargetSelect.h"
+#include "llvm/Transforms/InstCombine/InstCombine.h"
+#include "llvm/Transforms/Scalar.h"
+#include "llvm/Transforms/Scalar/GVN.h"
+
+namespace micm
+{
+
+  class JitCompiler
+  {
+   private:
+    std::unique_ptr<llvm::orc::ExecutionSession> execution_session_;
+
+    llvm::DataLayout data_layout_;
+    llvm::orc::MangleAndInterner mangle_;
+
+    llvm::orc::RTDyldObjectLinkingLayer object_layer_;
+    llvm::orc::IRCompileLayer compile_layer_;
+    llvm::orc::IRTransformLayer optimize_layer_;
+
+    llvm::orc::JITDylib &main_lib_;
+
+   public:
+    JitCompiler(
+        std::unique_ptr<llvm::orc::ExecutionSession> execution_session,
+        llvm::orc::JITTargetMachineBuilder machine_builder,
+        llvm::DataLayout data_layout)
+        : execution_session_(std::move(execution_session)),
+          data_layout_(std::move(data_layout)),
+          mangle_(*this->execution_session_, this->data_layout_),
+          object_layer_(*this->execution_session_, []() { return std::make_unique<llvm::SectionMemoryManager>(); }),
+          compile_layer_(
+              *this->execution_session_,
+              object_layer_,
+              std::make_unique<llvm::orc::ConcurrentIRCompiler>(std::move(machine_builder))),
+          optimize_layer_(*this->execution_session_, compile_layer_, OptimizeModule),
+          main_lib_(this->execution_session_->createBareJITDylib("<main>"))
+    {
+      main_lib_.addGenerator(
+          llvm::cantFail(llvm::orc::DynamicLibrarySearchGenerator::GetForCurrentProcess(data_layout_.getGlobalPrefix())));
+    }
+
+    ~JitCompiler()
+    {
+      if (auto Err = execution_session_->endSession())
+        execution_session_->reportError(std::move(Err));
+    }
+
+    static llvm::Expected<std::shared_ptr<JitCompiler>> create()
+    {
+      llvm::InitializeNativeTarget();
+      llvm::InitializeNativeTargetAsmPrinter();
+      llvm::InitializeNativeTargetAsmParser();
+
+      auto EPC = llvm::orc::SelfExecutorProcessControl::Create();
+      if (!EPC)
+        return EPC.takeError();
+
+      auto execution_session = std::make_unique<llvm::orc::ExecutionSession>(std::move(*EPC));
+
+      llvm::orc::JITTargetMachineBuilder machine_builder(execution_session->getExecutorProcessControl().getTargetTriple());
+
+      auto data_layout = machine_builder.getDefaultDataLayoutForTarget();
+      if (!data_layout)
+        return data_layout.takeError();
+
+      return std::make_shared<JitCompiler>(
+          std::move(execution_session), std::move(machine_builder), std::move(*data_layout));
+    }
+
+    const llvm::DataLayout &GetDataLayout() const
+    {
+      return data_layout_;
+    }
+
+    llvm::orc::JITDylib &GetMainJITDylib()
+    {
+      return main_lib_;
+    }
+
+    llvm::Error AddModule(
+        llvm::orc::ThreadSafeModule threadsafe_module,
+        llvm::orc::ResourceTrackerSP resource_tracker = nullptr)
+    {
+      if (!resource_tracker)
+        resource_tracker = main_lib_.getDefaultResourceTracker();
+      return optimize_layer_.add(resource_tracker, std::move(threadsafe_module));
+    }
+
+    llvm::Expected<llvm::JITEvaluatedSymbol> Lookup(llvm::StringRef name)
+    {
+      return execution_session_->lookup({ &main_lib_ }, mangle_(name.str()));
+    }
+
+   private:
+    static llvm::Expected<llvm::orc::ThreadSafeModule> OptimizeModule(
+        llvm::orc::ThreadSafeModule threadsafe_module,
+        const llvm::orc::MaterializationResponsibility &responsibility)
+    {
+      threadsafe_module.withModuleDo(
+          [](llvm::Module &module)
+          {
+            // Create a function pass manager.
+            auto pass_manager = std::make_unique<llvm::legacy::FunctionPassManager>(&module);
+
+            // Add some optimizations.
+            pass_manager->add(llvm::createInstructionCombiningPass());
+            pass_manager->add(llvm::createReassociatePass());
+            pass_manager->add(llvm::createGVNPass());
+            pass_manager->add(llvm::createCFGSimplificationPass());
+            pass_manager->doInitialization();
+
+            // Run the optimizations over all functions in the module being added to
+            // the JIT.
+            for (auto &function : module)
+              pass_manager->run(function);
+          });
+
+      return std::move(threadsafe_module);
+    }
+  };
+
+}  // end namespace micm