Refactoring of runners' infrastructure and dumping object files (#620)

The following is added:

1. Dumping object files in JIT.

A functionality to dump (enabled by default) the generated from LLVM IR file binary to `.o` has been added to benchmarking. Now, in addition to logs, a `v<vector_width>_<mod_filename>.o` is generated. The reasons it is an object file and not an assembly (hence not included in logs) are the following:

- LLVM does not have library functions that take the object and turn back into assembly, but rather `object -> file -> assembly` path. It also has a `llvm-objdump` tool, but it is intended as a command-line utility and does not have a well-defined API.

- Writing custom functions to produce a readable assembly is not a priority. Also, mimicking `objdump` functionality would be difficult.

- Both `objdump` and `llvm-objdump` can be used to isnpect the `.o` file manually.

2. Refactoring of `Runner` class.

In addition to the support of dumping the binary, `Runner`and `JITDriver` classes were refactored to have a nicer OOP-style.

fixes #611

Co-authored-by: Pramod S Kumbhar <[email protected]>

src/codegen/llvm/jit_driver.cpp

@@ -15,6 +15,7 @@
 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
 #include "llvm/ExecutionEngine/Orc/LLJIT.h"
+#include "llvm/ExecutionEngine/Orc/ObjectTransformLayer.h"
 #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
 #include "llvm/Support/Host.h"
@@ -24,7 +25,9 @@
 namespace nmodl {
 namespace runner {
 
-void JITDriver::init(std::string features, std::vector<std::string>& lib_paths) {
+void JITDriver::init(std::string features,
+                     std::vector<std::string> lib_paths,
+                     ObjDumpInfo* dump_info) {
     llvm::InitializeNativeTarget();
     llvm::InitializeNativeTargetAsmPrinter();
 
@@ -83,6 +86,12 @@ void JITDriver::init(std::string features, std::vector<std::string>& lib_paths)
     llvm::orc::JITDylib& sym_tab = jit->getMainJITDylib();
     sym_tab.addGenerator(cantFail(llvm::orc::DynamicLibrarySearchGenerator::GetForCurrentProcess(
         data_layout.getGlobalPrefix())));
+
+    // Optionally, dump the binary to the object file.
+    if (dump_info) {
+        jit->getObjTransformLayer().setTransform(
+            llvm::orc::DumpObjects(dump_info->output_dir, dump_info->filename));
+    }
 }
 
 std::unique_ptr<llvm::TargetMachine> JITDriver::create_target(

src/codegen/llvm/jit_driver.hpp

@@ -20,9 +20,18 @@
 namespace nmodl {
 namespace runner {
 
+/// A struct to hold the information for dumping object file.
+struct ObjDumpInfo {
+    /// Object file name.
+    std::string filename;
+
+    /// Object file output directory.
+    std::string output_dir;
+};
+
 /**
  * \class JITDriver
- * \brief Driver to execute MOD file function via LLVM IR backend
+ * \brief Driver to execute a MOD file function via LLVM IR backend.
  */
 class JITDriver {
   private:
@@ -33,13 +42,15 @@ class JITDriver {
     std::unique_ptr<llvm::Module> module;
 
   public:
-    JITDriver(std::unique_ptr<llvm::Module> m)
+    explicit JITDriver(std::unique_ptr<llvm::Module> m)
         : module(std::move(m)) {}
 
-    /// Initialize the JIT.
-    void init(std::string features, std::vector<std::string>& lib_paths);
+    /// Initializes the JIT.
+    void init(std::string features = "",
+              std::vector<std::string> lib_paths = {},
+              ObjDumpInfo* dump_info = nullptr);
 
-    /// Lookup the entry-point without arguments in the JIT and execute it, returning the result.
+    /// Lookups the entry-point without arguments in the JIT and executes it, returning the result.
     template <typename ReturnType>
     ReturnType execute_without_arguments(const std::string& entry_point) {
         auto expected_symbol = jit->lookup(entry_point);
@@ -51,7 +62,7 @@ class JITDriver {
         return result;
     }
 
-    /// Lookup the entry-point with an argument in the JIT and execute it, returning the result.
+    /// Lookups the entry-point with an argument in the JIT and executes it, returning the result.
     template <typename ReturnType, typename ArgType>
     ReturnType execute_with_arguments(const std::string& entry_point, ArgType arg) {
         auto expected_symbol = jit->lookup(entry_point);
@@ -63,7 +74,8 @@ class JITDriver {
         return result;
     }
 
-    /// A wrapper around llvm::createTargetMachine to turn on/off certain CPU features.
+  private:
+    /// Creates llvm::TargetMachine with certain CPU features turned on/off.
     std::unique_ptr<llvm::TargetMachine> create_target(llvm::orc::JITTargetMachineBuilder* builder,
                                                        const std::string& features);
 
@@ -72,35 +84,79 @@ class JITDriver {
 };
 
 /**
- * \class Runner
- * \brief A wrapper around JITDriver to execute an entry point in the LLVM IR module.
+ * \class BaseRunner
+ * \brief A base runner class that provides functionality to execute an
+ * entry point in the LLVM IR module.
  */
-class Runner {
-  private:
-    std::unique_ptr<llvm::Module> module;
+class BaseRunner {
+  protected:
+    std::unique_ptr<JITDriver> driver;
 
-    std::unique_ptr<JITDriver> driver = std::make_unique<JITDriver>(std::move(module));
+    explicit BaseRunner(std::unique_ptr<llvm::Module> m)
+        : driver(std::make_unique<JITDriver>(std::move(m))) {}
 
   public:
-    Runner(std::unique_ptr<llvm::Module> m,
-           std::string features = "",
-           std::vector<std::string> lib_paths = {})
-        : module(std::move(m)) {
-        driver->init(features, lib_paths);
-    }
+    /// Sets up the JIT driver.
+    virtual void initialize_driver() = 0;
 
-    /// Run the entry-point function without arguments.
+    /// Runs the entry-point function without arguments.
     template <typename ReturnType>
     ReturnType run_without_arguments(const std::string& entry_point) {
         return driver->template execute_without_arguments<ReturnType>(entry_point);
     }
 
-    /// Run the entry-point function with a pointer to the data as an argument.
+    /// Runs the entry-point function with a pointer to the data as an argument.
     template <typename ReturnType, typename ArgType>
     ReturnType run_with_argument(const std::string& entry_point, ArgType arg) {
         return driver->template execute_with_arguments<ReturnType, ArgType>(entry_point, arg);
     }
 };
 
+/**
+ * \class TestRunner
+ * \brief A simple runner for testing purposes.
+ */
+class TestRunner: public BaseRunner {
+  public:
+    explicit TestRunner(std::unique_ptr<llvm::Module> m)
+        : BaseRunner(std::move(m)) {}
+
+    virtual void initialize_driver() {
+        driver->init();
+    }
+};
+
+/**
+ * \class BenchmarkRunner
+ * \brief A runner with benchmarking functionality. It takes user-specified CPU
+ * features into account, as well as it can link against shared libraries.
+ */
+class BenchmarkRunner: public BaseRunner {
+  private:
+    /// Information on dumping object file generated from LLVM IR.
+    ObjDumpInfo dump_info;
+
+    /// CPU features specified by the user.
+    std::string features;
+
+    /// Shared libraries' paths to link against.
+    std::vector<std::string> shared_lib_paths;
+
+  public:
+    BenchmarkRunner(std::unique_ptr<llvm::Module> m,
+                    std::string filename,
+                    std::string output_dir,
+                    std::string features = "",
+                    std::vector<std::string> lib_paths = {})
+        : BaseRunner(std::move(m))
+        , dump_info{filename, output_dir}
+        , features(features)
+        , shared_lib_paths(lib_paths) {}
+
+    virtual void initialize_driver() {
+        driver->init(features, shared_lib_paths, &dump_info);
+    }
+};
+
 }  // namespace runner
 }  // namespace nmodl

src/codegen/llvm/llvm_benchmark.cpp

@@ -104,7 +104,11 @@ void LLVMBenchmark::run_benchmark(codegen::CodegenLLVMVisitor& visitor,
 
     std::string features_str = llvm::join(features.begin(), features.end(), ",");
     std::unique_ptr<llvm::Module> m = visitor.get_module();
-    runner::Runner runner(std::move(m), features_str, shared_libs);
+
+    // Create the benchmark runner and intialize it.
+    std::string filename = "v" + std::to_string(llvm_build_info.vector_width) + "_" + mod_filename;
+    runner::BenchmarkRunner runner(std::move(m), filename, output_dir, features_str, shared_libs);
+    runner.initialize_driver();
 
     // Benchmark every kernel.
     for (const auto& kernel_name: kernel_names) {

src/codegen/llvm/main.cpp

@@ -64,7 +64,8 @@ int main(int argc, const char* argv[]) {
         throw std::runtime_error(
             "Error: entry-point functions with non-double return type are not supported\n");
 
-    Runner runner(std::move(module));
+    TestRunner runner(std::move(module));
+    runner.initialize_driver();
 
     // Since only double type is supported, provide explicit double type to the running function.
     auto r = runner.run_without_arguments<double>(entry_point_name);

test/unit/codegen/codegen_llvm_execution.cpp

@@ -130,7 +130,8 @@ SCENARIO("Arithmetic expression", "[llvm][runner]") {
         llvm_visitor.visit_program(*ast);
 
         std::unique_ptr<llvm::Module> m = llvm_visitor.get_module();
-        Runner runner(std::move(m));
+        TestRunner runner(std::move(m));
+        runner.initialize_driver();
 
         THEN("functions are evaluated correctly") {
             auto exp_result = runner.run_without_arguments<double>("exponential");
@@ -231,7 +232,8 @@ SCENARIO("Optimised arithmetic expression", "[llvm][runner]") {
         llvm_visitor.visit_program(*ast);
 
         std::unique_ptr<llvm::Module> m = llvm_visitor.get_module();
-        Runner runner(std::move(m));
+        TestRunner runner(std::move(m));
+        runner.initialize_driver();
 
         THEN("optimizations preserve function results") {
             // Check exponential is turned into a constant.
@@ -325,7 +327,8 @@ SCENARIO("Simple scalar kernel", "[llvm][runner]") {
 
         // Set up the JIT runner.
         std::unique_ptr<llvm::Module> module = llvm_visitor.get_module();
-        Runner runner(std::move(module));
+        TestRunner runner(std::move(module));
+        runner.initialize_driver();
 
         THEN("Values in struct have changed according to the formula") {
             runner.run_with_argument<int, void*>("__nrn_state_test_wrapper",
@@ -412,7 +415,8 @@ SCENARIO("Simple vectorised kernel", "[llvm][runner]") {
 
         // Set up the JIT runner.
         std::unique_ptr<llvm::Module> module = llvm_visitor.get_module();
-        Runner runner(std::move(module));
+        TestRunner runner(std::move(module));
+        runner.initialize_driver();
 
         THEN("Values in struct have changed according to the formula") {
             runner.run_with_argument<int, void*>("__nrn_state_test_wrapper",