cpu_context.hpp

Go to the documentation of this file.

//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
 
#ifndef cpu_context_h
#define cpu_context_h
 
#include <fstream>
#include <cstdlib>
#include <cstring>
#include <thread>
#include <unordered_set>
 
//  Clang headers will define IBAction and IBOutlet these so undefine them here.
#undef IBAction
#undef IBOutlet
#include "llvm/Support/VirtualFileSystem.h"
#include "clang/Frontend/TextDiagnosticPrinter.h"
#include "clang/Frontend/CompilerInvocation.h"
#include "clang/Frontend/CompilerInstance.h"
#include "clang/Basic/TargetInfo.h"
#include "clang/CodeGen/CodeGenAction.h"
#include "clang/Lex/PreprocessorOptions.h"
#include "llvm/Support/TargetSelect.h"
#include "llvm/ExecutionEngine/Orc/LLJIT.h"
#ifndef NDEBUG
#include "llvm/ExecutionEngine/Orc/Debugging/DebuggerSupport.h"
#include "llvm/ExecutionEngine/Orc/TargetProcess/JITLoaderGDB.h"
#endif
#include "llvm/Support/raw_ostream.h"
#include "llvm/ADT/IntrusiveRefCntPtr.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/TargetParser/Host.h"
#include "llvm/ExecutionEngine/Orc/ThreadSafeModule.h"
 
#include "random.hpp"
 
#ifndef NDEBUG
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
LLVM_ATTRIBUTE_USED void linkComponents() {
    llvm::errs() << (void *)&llvm_orc_registerJITLoaderGDBWrapper
                 << (void *)&llvm_orc_registerJITLoaderGDBAllocAction;
}
#endif
 
namespace gpu {
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    llvm::SmallVector<const char *, 8> split_string(char *string) {
        llvm::SmallVector<const char *, 8> args = {string};
 
        while (*(++string) != '\0') {
            if (*string == ' ') {
                *string = '\0';
                args.push_back(++string);
            }
        }
 
        return args;
    }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    class cpu_context {
    private:
        std::unique_ptr<llvm::orc::LLJIT> jit;
        std::map<graph::leaf_node<T, SAFE_MATH> *, std::vector<T>> kernel_arguments;
        std::map<graph::leaf_node<T, SAFE_MATH> *, std::vector<T>> host_buffers;
        std::map<graph::leaf_node<T, SAFE_MATH> *, size_t> arg_index;
 
    public:
        constexpr static size_t random_state_size = 1;
 
        int remaining_const_memory;
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        static size_t max_concurrency() {
            return std::thread::hardware_concurrency();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        static std::string device_type() {
            return "CPU";
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        cpu_context(const size_t index) {
            llvm::InitializeNativeTarget();
            llvm::InitializeNativeTargetAsmPrinter();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void compile(const std::string kernel_source,
                     std::vector<std::string> names,
                     const bool add_reduction=false) {
            std::ostringstream temp_stream;
            temp_stream << reinterpret_cast<size_t> (this);
            const std::string thread_id = temp_stream.str();
 
            temp_stream.str(std::string());
            temp_stream.clear();
 
            temp_stream << "temp_" << thread_id << ".cpp";
 
            const std::string filename = temp_stream.str();
 
            if (jit::verbose) {
                std::cout << "CPU info." << std::endl;
                std::cout << "  Command Line    : " << std::endl;
            }
 
            char arg_string[] = CXX_ARGS;
            llvm::SmallVector<const char *, 8> args = split_string(arg_string);
            args.push_back(filename.c_str());
#ifdef NDEBUG
            args.push_back("-ffast-math");
            args.push_back("-O3");
#else
            args.push_back("-debug-info-kind=standalone");
#endif
            if (jit::verbose) {
                for (auto &arg : args) {
                    std::cout << "    " << arg << std::endl;
                }
            }
 
            clang::DiagnosticOptions diagnostic_options;
            auto diagnostic_printer = std::make_unique<clang::TextDiagnosticPrinter> (llvm::errs(),
                                                                                      diagnostic_options);
 
            auto diagnostic_ids = llvm::makeIntrusiveRefCnt<clang::DiagnosticIDs> ();
            clang::DiagnosticsEngine diagnostic_engine(diagnostic_ids,
                                                       diagnostic_options,
                                                       diagnostic_printer.release());
 
            auto invocation = std::make_shared<clang::CompilerInvocation> ();
            clang::CompilerInvocation::CreateFromArgs(*(invocation.get()), args,
                                                      diagnostic_engine);
 
            llvm::StringRef source_code_data(kernel_source);
            auto buffer = llvm::MemoryBuffer::getMemBuffer(source_code_data);
            invocation->getPreprocessorOpts().addRemappedFile(filename.c_str(),
                                                              buffer.release());
 
            clang::CompilerInstance clang(invocation);
            clang.createDiagnostics();
 
            clang::TargetOptions target_options;
            target_options.Triple = llvm::sys::getProcessTriple();
            auto *target_info = clang::TargetInfo::CreateTargetInfo(diagnostic_engine,
                                                                    target_options);
            clang.setTarget(target_info);
 
            clang::EmitLLVMOnlyAction action;
            clang.ExecuteAction(action);
 
            auto ir_module = action.takeModule();
            auto context = std::unique_ptr<llvm::LLVMContext> (action.takeLLVMContext());
 
            auto jit_try = llvm::orc::LLJITBuilder()
#ifndef NDEBUG
                               .setPrePlatformSetup([](llvm::orc::LLJIT &J) {
                                   return llvm::orc::enableDebuggerSupport(J);
                               })
#endif
                               .create();
            if (auto jiterror = jit_try.takeError()) {
                std::cerr << "Failed to build JIT : " << toString(std::move(jiterror)) << std::endl;
                exit(-1);
            }
            jit = std::move(jit_try.get());
 
            auto error = jit->addIRModule(llvm::orc::ThreadSafeModule(std::move(ir_module),
                                                                      llvm::orc::ThreadSafeContext(std::move(context))));
 
#ifdef MACOS_LIB_RT
            error = jit->linkStaticLibraryInto(jit->getMainJITDylib(), MACOS_LIB_RT);
#endif
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        std::function<void(void)> create_kernel_call(const std::string kernel_name,
                                                     graph::input_nodes<T, SAFE_MATH> inputs,
                                                     graph::output_nodes<T, SAFE_MATH> outputs,
                                                     graph::shared_random_state<T, SAFE_MATH> state,
                                                     const size_t num_rays,
                                                     const jit::texture1d_list &tex1d_list,
                                                     const jit::texture2d_list &tex2d_list) {
            auto entry = std::move(jit->lookup(kernel_name)).get();
 
            std::map<size_t, T *> buffers;
 
            for (auto &input : inputs) {
                if (!kernel_arguments.contains(input.get())) {
                    backend::buffer<T> buffer = input->evaluate();
                    std::vector<T> arg(buffer.size());
                    memcpy(arg.data(), buffer.data(), buffer.size()*sizeof(T));
                    kernel_arguments[input.get()] = arg;
                }
                buffers[reinterpret_cast<size_t> (input.get())] = kernel_arguments[input.get()].data();
            }
            for (auto &output : outputs) {
                if (!kernel_arguments.contains(output.get())) {
                    std::vector<T> arg(num_rays);
                    kernel_arguments[output.get()] = arg;
                }
                buffers[reinterpret_cast<size_t> (output.get())] = kernel_arguments[output.get()].data();
            }
 
            if (state.get()) {
                auto kernel = entry.toPtr<void(*)(std::map<size_t, T *> &, typename graph::random_state_node<T, SAFE_MATH>::mt_state *)> ();
 
                if (!kernel) {
                    std::cerr << "Failed to load function. " << kernel_name
                              << std::endl;
                    exit(-1);
                }
 
                if (jit::verbose) {
                    std::cout << "  Function pointer: "
                              << reinterpret_cast<size_t> (kernel)
                              << std::endl;
                }
 
                return [kernel, buffers, state] () mutable {
                    kernel(buffers, state->data());
                };
            } else {
                auto kernel = entry.toPtr<void(*)(std::map<size_t, T *> &)> ();
 
                if (!kernel) {
                    std::cerr << "Failed to load function. " << kernel_name
                              << std::endl;
                    exit(-1);
                }
 
                if (jit::verbose) {
                    std::cout << "  Function pointer: "
                              << reinterpret_cast<size_t> (kernel)
                              << std::endl;
                }
 
                return [kernel, buffers] () mutable {
                    kernel(buffers);
                };
            }
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        std::function<T(void)> create_max_call(graph::shared_leaf<T, SAFE_MATH> &argument,
                                               std::function<void(void)> run) {
            auto begin = kernel_arguments[argument.get()].cbegin();
            auto end = kernel_arguments[argument.get()].cend();
 
            return [run, begin, end] () mutable {
                run();
                if constexpr (jit::complex_scalar<T>) {
                    return *std::max_element(begin, end,
                                             [] (const T a, const T b) {
                        return std::abs(a) < std::abs(b);
                    });
                } else {
                    return *std::max_element(begin, end);
                }
            };
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void wait() {
            for (auto &item : host_buffers) {
                memcpy(item.second.data(),
                       kernel_arguments[item.first].data(),
                       sizeof(T)*kernel_arguments[item.first].size());
            }
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void print_results(const size_t index,
                           const graph::output_nodes<T, SAFE_MATH> &nodes) {
            for (auto &out : nodes) {
                const T temp = kernel_arguments[out.get()][index];
                if constexpr (jit::complex_scalar<T>) {
                    std::cout << std::real(temp) << " " << std::imag(temp) << " ";
                } else {
                    std::cout << temp << " ";
                }
            }
            std::cout << std::endl;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        T check_value(const size_t index,
                      const graph::shared_leaf<T, SAFE_MATH> &node) {
            return kernel_arguments[node.get()][index];
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void copy_to_device(graph::shared_leaf<T, SAFE_MATH> node,
                            T *source) {
            memcpy(kernel_arguments[node.get()].data(),
                   source,
                   sizeof(T)*kernel_arguments[node.get()].size());
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void copy_to_host(const graph::shared_leaf<T, SAFE_MATH> node,
                          T *destination) {
            memcpy(destination,
                   kernel_arguments[node.get()].data(),
                   sizeof(T)*kernel_arguments[node.get()].size());
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void create_header(std::ostringstream &source_buffer) {
            source_buffer << "#include <map>"     << std::endl
                          << "#include <array>"   << std::endl
                          << "#include <cstdint>" << std::endl;
            if (jit::complex_scalar<T>) {
                source_buffer << "#include <complex>" << std::endl;
                source_buffer << "#include <special_functions.hpp>" << std::endl;
            } else {
                source_buffer << "#include <cmath>" << std::endl;
            }
            source_buffer << "using namespace std;" << std::endl;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void create_kernel_prefix(std::ostringstream &source_buffer,
                                  const std::string name,
                                  graph::input_nodes<T, SAFE_MATH> &inputs,
                                  graph::output_nodes<T, SAFE_MATH> &outputs,
                                  graph::shared_random_state<T, SAFE_MATH> state,
                                  const size_t size,
                                  const std::vector<bool> &is_constant,
                                  jit::register_map &registers,
                                  const jit::register_usage &usage,
                                  jit::texture1d_list &textures1d,
                                  jit::texture2d_list &textures2d) {
            source_buffer << std::endl;
            source_buffer << "extern \"C\" void " << name << "(" << std::endl;
 
            source_buffer << "    map<size_t, ";
            jit::add_type<T> (source_buffer);
            source_buffer << " *> &args";
            if (state.get()) {
                source_buffer << "," << std::endl
                              << "    mt_state *" << jit::to_string('s', state.get());
            }
            source_buffer << ") {" << std::endl;
 
            std::unordered_set<void *> used_args;
            for (size_t i = 0, ie = inputs.size(); i < ie; i++) {
                if (!used_args.contains(inputs[i].get())) {
                    source_buffer << "    ";
                    if (is_constant[i]) {
                        source_buffer << "const ";
                    }
                    jit::add_type<T> (source_buffer);
                    source_buffer << " *" << jit::to_string('v', inputs[i].get())
                                  << " = args["
                                  << reinterpret_cast<size_t> (inputs[i].get())
                                  << "];" << std::endl;
                    used_args.insert(inputs[i].get());
                }
            }
            for (auto &output : outputs) {
                if (!used_args.contains(output.get())) {
                    source_buffer << "    ";
                    jit::add_type<T> (source_buffer);
                    source_buffer << " *" << jit::to_string('o', output.get())
                                  << " = args["
                                  << reinterpret_cast<size_t> (output.get())
                                  << "];" << std::endl;
                    used_args.insert(output.get());
                }
            }
            if (state.get()) {
                registers[state.get()] = jit::to_string('r', state.get());
                source_buffer << "    mt_state &"
                              << registers[state.get()] << " = "
                              << jit::to_string('s', state.get()) << "[0];"
#ifdef SHOW_USE_COUNT
                              << " // used " << usage.at(state.get())
#endif
                              << std::endl;
            }
            source_buffer << "    for (size_t i = 0; i < " << size << "; i++) {" << std::endl;
 
            for (auto &input : inputs) {
                registers[input.get()] = jit::to_string('r', input.get());
                source_buffer << "        const ";
                jit::add_type<T> (source_buffer);
                source_buffer << " " << registers[input.get()]
                              << " = " << jit::to_string('v', input.get())
                              << "[i]; // " << input->get_symbol()
#ifdef SHOW_USE_COUNT
                              << " used " << usage.at(input.get())
#endif
                              << std::endl;
            }
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void create_kernel_postfix(std::ostringstream &source_buffer,
                                   graph::output_nodes<T, SAFE_MATH> &outputs,
                                   graph::map_nodes<T, SAFE_MATH> &setters,
                                   graph::shared_random_state<T, SAFE_MATH> state,
                                   jit::register_map &registers,
                                   jit::register_map &indices,
                                   const jit::register_usage &usage) {
            std::unordered_set<void *> out_registers;
            for (auto &[out, in] : setters) {
                if (!out->is_match(in) &&
                    !out_registers.contains(out.get())) {
                    graph::shared_leaf<T, SAFE_MATH> a = out->compile(source_buffer,
                                                                      registers,
                                                                      indices,
                                                                      usage);
                    source_buffer << "        " << jit::to_string('v', in.get());
                    source_buffer << "[i] = ";
                    if constexpr (SAFE_MATH) {
                        if constexpr (jit::complex_scalar<T>) {
                            jit::add_type<T> (source_buffer);
                            source_buffer << " (";
                            source_buffer << "isnan(real(" << registers[a.get()]
                            << ")) ? 0.0 : real(" << registers[a.get()]
                            << "), ";
                            source_buffer << "isnan(imag(" << registers[a.get()]
                            << ")) ? 0.0 : imag(" << registers[a.get()]
                            << "));" << std::endl;
                        } else {
                            source_buffer << "isnan(" << registers[a.get()]
                            << ") ? 0.0 : " << registers[a.get()]
                            << ";" << std::endl;
                        }
                    } else {
                        source_buffer << registers[a.get()] << ";" << std::endl;
                    }
                    out_registers.insert(out.get());
                }
            }
            for (auto &out : outputs) {
                if (!graph::variable_cast(out).get() &&
                    !out_registers.contains(out.get())) {
                    graph::shared_leaf<T, SAFE_MATH> a = out->compile(source_buffer,
                                                                      registers,
                                                                      indices,
                                                                      usage);
                    source_buffer << "        " << jit::to_string('o', out.get());
                    source_buffer << "[i] = ";
                    if constexpr (SAFE_MATH) {
                        if constexpr (jit::complex_scalar<T>) {
                            jit::add_type<T> (source_buffer);
                            source_buffer << " (";
                            source_buffer << "isnan(real(" << registers[a.get()]
                            << ")) ? 0.0 : real(" << registers[a.get()]
                            << "), ";
                            source_buffer << "isnan(imag(" << registers[a.get()]
                            << ")) ? 0.0 : imag(" << registers[a.get()]
                            << "));" << std::endl;
                        } else {
                            source_buffer << "isnan(" << registers[a.get()]
                            << ") ? 0.0 : " << registers[a.get()]
                            << ";" << std::endl;
                        }
                    } else {
                        source_buffer << registers[a.get()] << ";" << std::endl;
                    }
                    out_registers.insert(out.get());
                }
            }
 
            source_buffer << "    }" << std::endl;
            source_buffer << "}" << std::endl;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void create_reduction(std::ostringstream &source_buffer,
                              const size_t size) {}
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        T *get_buffer(graph::shared_leaf<T, SAFE_MATH> &node) {
            if (!host_buffers.contains(node.get())) {
                host_buffers[node.get()] = kernel_arguments[node.get()];
            }
            return host_buffers[node.get()].data();
        }
    };
}
 
#endif /* cpu_context_h */