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>
 
//  Clang headers will define IBAction and IBOutlet these so undefined 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"
#include "piecewise.hpp"
#include "timing.hpp"
 
#ifndef NDEBUG
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
LLVM_ATTRIBUTE_USED void linkComponents() {
    llvm::errs() << (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::unordered_map<graph::leaf_node<T, SAFE_MATH> *, std::vector<T>> kernel_arguments;
        std::unordered_map<graph::leaf_node<T, SAFE_MATH> *, std::vector<T>> host_buffers;
        std::unordered_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::input_nodes<T, SAFE_MATH> atomics,
                                                     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::unordered_map<size_t, T *> buffers;
 
            for (auto &input : inputs) {
                if (!kernel_arguments.contains(input.get())) {
                    std::vector<T> arg(input->size());
                    memcpy(arg.data(), input->data(), input->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 (!graph::atomic_accumulate_1D_cast(output).get()) {
                    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();
                }
            }
            for (auto &atomic : atomics) {
                if (!kernel_arguments.contains(atomic.get())) {
                    std::vector<T> arg(atomic->size());
                    memcpy(arg.data(), atomic->data(), atomic->size()*sizeof(T));
                    kernel_arguments[atomic.get()] = arg;
                }
                buffers[reinterpret_cast<size_t> (atomic.get())] = kernel_arguments[atomic.get()].data();
            }
 
            if (state.get()) {
                auto kernel = entry.toPtr<void(*)(std::unordered_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
#ifdef PROFILE_KERNELS
                        , kernel_name
#endif
                ] () mutable {
#ifdef PROFILE_KERNELS
                    timing::measure_diagnostic timer(kernel_name);
#endif
                    kernel(buffers, state->data());
#ifdef PROFILE_KERNELS
                    timer.print();
#endif
                };
            } else {
                auto kernel = entry.toPtr<void(*)(std::unordered_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
#ifdef PROFILE_KERNELS
                        , kernel_name
#endif
                ] () mutable {
#ifdef PROFILE_KERNELS
                    timing::measure_diagnostic timer(kernel_name);
#endif
                    kernel(buffers);
#ifdef PROFILE_KERNELS
                    timer.print();
#endif
                };
            }
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        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);
                }
            };
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        std::function<void(void)> create_zero_call(graph::input_nodes<T, SAFE_MATH> &inputs) {
            std::vector<T *> buffers;
            std::vector<size_t> sizes;
 
            for (auto &input : inputs) {
                if (!kernel_arguments.contains(input.get())) {
                    std::vector<T> arg(input->size());
                    memcpy(arg.data(), input->data(), input->size()*sizeof(T));
                    kernel_arguments[input.get()] = arg;
                }
                buffers.push_back(kernel_arguments[input.get()].data());
                sizes.push_back(input->size()*sizeof(T));
            }
 
            return [buffers, sizes] () mutable {
#ifdef PROFILE_KERNELS
                    timing::measure_diagnostic timer("zero buffer");
#endif
                for (size_t i = 0, ie = buffers.size(); i < ie; i++) {
                    std::memset(buffers[i], 0, sizes[i]);
                }
#ifdef PROFILE_KERNELS
                    timer.print();
#endif
            };
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        std::function<void(void)> create_copy_call(graph::copy_nodes<T, SAFE_MATH> &setters) {
            std::vector<T *> sources;
            std::vector<T *> destinations;
            std::vector<size_t> sizes;
 
            for (auto &[out, in] : setters) {
                if (!kernel_arguments.contains(in.get())) {
                    std::vector<T> arg(in->size());
                    memcpy(arg.data(), in->data(), in->size()*sizeof(T));
                    kernel_arguments[in.get()] = arg;
                }
                destinations.push_back(kernel_arguments[in.get()].data());
                sizes.push_back(in->size()*sizeof(T));
 
                if (!kernel_arguments.contains(out.get())) {
                    std::vector<T> arg(out->size());
                    memcpy(arg.data(), out->data(), out->size()*sizeof(T));
                    kernel_arguments[out.get()] = arg;
                }
                sources.push_back(kernel_arguments[out.get()].data());
            }
 
            return [sources, destinations, sizes] () mutable {
#ifdef PROFILE_KERNELS
                    timing::measure_diagnostic timer("copy buffer");
#endif
                for (size_t i = 0, ie = sources.size(); i < ie; i++) {
                    std::memcpy(destinations[i], sources[i], sizes[i]);
                }
#ifdef PROFILE_KERNELS
                    timer.print();
#endif
            };
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void wait() {
            for (auto &item : host_buffers) {
                memcpy(item.second.data(),
                       kernel_arguments[item.first].data(),
                       sizeof(T)*kernel_arguments[item.first].size());
            }
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        std::function<void(void)> run_function(std::function<void(void)> callback) {
#ifdef PROFILE_KERNELS
            return [callback]() {
                timing::measure_diagnostic timer("callback");
                callback();
                timer.print();
            };
#else
            return callback;
#endif
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        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 <unordered_map>" << std::endl
                          << "#include <array>"         << std::endl
                          << "#include <cstdint>"       << std::endl
                          << "#include <bit>"           << 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 << "#include <atomic>" << std::endl
                          << "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::input_nodes<T, SAFE_MATH> atomics,
                                  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,
                                  jit::argument_set &thread_shared,
                                  jit::register_map &thread_mem,
                                  const size_t iterations=1) {
            source_buffer << std::endl;
            source_buffer << "extern \"C\" void " << name << "(" << std::endl;
 
            source_buffer << "    unordered_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;
 
            jit::argument_set 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()) &&
                    !graph::atomic_accumulate_1D_cast(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());
                }
            }
            for (size_t i = 0, ie = atomics.size(); i < ie; i++) {
                if (!used_args.contains(atomics[i].get())) {
                    source_buffer << "    ";
                    jit::add_type<T> (source_buffer);
                    source_buffer << " *"
                                  << jit::to_string('v', atomics[i].get())
                                  << " = args["
                                  << reinterpret_cast<size_t> (atomics[i].get())
                                  << "];"
                                  << std::endl;
                    used_args.insert(atomics[i].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]";
                state->endline(source_buffer, usage);
            }
            source_buffer << "    for (size_t i = 0; i < " << size << "; i++) {" << std::endl;
            if (iterations > 1) {
                source_buffer << "    for (size_t j = 0; j < " << iterations << "; j++) {" << 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->endline(source_buffer, usage);
            }
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        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,
                                   const jit::register_usage &usage,
                                   const jit::argument_set &thread_shared,
                                   jit::register_map &thread_mem,
                                   const size_t iterations=1) {
            jit::argument_set out_registers;
            for (auto &[out, in] : setters) {
                if (!out->is_match(in)) {
                    auto a = out->compile(source_buffer, registers,
                                          thread_mem, 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 << " (";
                            if (!graph::random_cast(a).get()) {
                                source_buffer << "isnan(real("
                                              << registers[a.get()]
                                              << ")) ? 0.0 : real("
                                              << registers[a.get()]
                                              << "), isnan(imag("
                                              << registers[a.get()]
                                              << ")) ? 0.0 : imag("
                                              << registers[a.get()]
                                              << ")";
                            } else {
                                source_buffer << registers[a.get()];
                            }
                            source_buffer << ");" << std::endl;
                        } else {
                            if (!graph::random_cast(a).get()) {
                                source_buffer << "isnan(" << registers[a.get()]
                                              << ") ? 0.0 : ";
                            }
                            source_buffer << 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()             &&
                    !graph::atomic_accumulate_1D_cast(out).get() &&
                    !out_registers.contains(out.get())) {
                    auto a = out->compile(source_buffer, registers,
                                          thread_mem, 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 << " (";
                            if (!graph::random_cast(a).get()) {
                                source_buffer << "isnan(real("
                                              << registers[a.get()]
                                              << ")) ? 0.0 : real("
                                              << registers[a.get()]
                                              << "), isnan(imag("
                                              << registers[a.get()]
                                              << ")) ? 0.0 : imag("
                                              << registers[a.get()]
                                              << ")";
                            } else {
                                source_buffer << registers[a.get()];
                            }
                            source_buffer << ");" << std::endl;
                        } else {
                            if (!graph::random_cast(a).get()) {
                                source_buffer << "isnan(" << registers[a.get()]
                                              << ") ? 0.0 : ";
                            }
                            source_buffer << registers[a.get()]
                                          << ";" << std::endl;
                        }
                    } else {
                        source_buffer << registers[a.get()] << ";" << std::endl;
                    }
                    out_registers.insert(out.get());
                }
            }
 
            if (iterations > 1) {
                source_buffer << "    }" << std::endl;
            }
            source_buffer << "    }" << std::endl  << "}" << 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 */