jit.hpp

Go to the documentation of this file.

//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
 
#ifndef jit_h
#define jit_h
 
#include <algorithm>
#include <iterator>
#include <thread>
 
#ifdef USE_METAL
#include "metal_context.hpp"
#elif defined(USE_CUDA)
#include "cuda_context.hpp"
#endif
#include "cpu_context.hpp"
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
#ifdef USE_METAL
#define START_GPU @autoreleasepool {
#define END_GPU }
#else
#define START_GPU
#define END_GPU
#endif
 
namespace jit {
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<float_scalar T, bool SAFE_MATH=false>
    class context {
    private:
        std::ostringstream source_buffer;
        register_map registers;
        preamble_map pre_funcs;
        std::vector<std::string> kernel_names;
        std::unordered_map<std::string, texture1d_list> kernel_1dtextures;
        std::unordered_map<std::string, texture2d_list> kernel_2dtextures;
 
        using gpu_context_type = typename std::conditional<use_gpu<T> (),
#ifdef USE_CUDA
                                                           gpu::cuda_context<T, SAFE_MATH>,
#elif defined(USE_METAL)
                                                           gpu::metal_context<SAFE_MATH>,
#else
                                                           gpu::cpu_context<T, SAFE_MATH>,
#endif
                                                           gpu::cpu_context<T, SAFE_MATH>>::type;
 
        gpu_context_type gpu_context;
 
    public:
        constexpr static size_t random_state_size = gpu_context_type::random_state_size;
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        static size_t max_random_state_size(const size_t size) {
            return std::min(size, random_state_size);
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        static size_t max_concurrency() {
            const size_t num = gpu_context_type::max_concurrency();
            std::cout << "Located " << num << " "
                      << gpu_context_type::device_type() << " device"
                      << (num == 1 ? "." : "s.")
                      << std::endl;
            return num;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        context(const size_t index) : gpu_context(index) {
            source_buffer << std::setprecision(max_digits10<T> ());
            gpu_context.create_header(source_buffer);
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void add_kernel(const std::string name,
                        graph::input_nodes<T, SAFE_MATH> &inputs,
                        graph::output_nodes<T, SAFE_MATH> &outputs,
                        graph::map_nodes<T, SAFE_MATH> &setters,
                        graph::input_nodes<T, SAFE_MATH> &atomics,
                        graph::shared_random_state<T, SAFE_MATH> &state,
                        const size_t size,
                        const size_t iterations=1) {
            kernel_names.push_back(name);
 
            std::vector<bool> is_constant(inputs.size(), true);
            visiter_map visited;
            register_usage usage;
            kernel_1dtextures[name] = texture1d_list();
            kernel_2dtextures[name] = texture2d_list();
            for (auto &[out, in] : setters) {
                auto found = std::distance(inputs.begin(),
                                           std::find(inputs.begin(),
                                                     inputs.end(), in));
                if (found < is_constant.size()) {
                    is_constant[found] = false;
                }
                out->compile_preamble(source_buffer, registers,
                                      visited, usage,
                                      kernel_1dtextures[name],
                                      kernel_2dtextures[name],
                                      pre_funcs,
                                      gpu_context.remaining_const_memory);
            }
            for (auto &out : outputs) {
                out->compile_preamble(source_buffer, registers,
                                      visited, usage,
                                      kernel_1dtextures[name],
                                      kernel_2dtextures[name],
                                      pre_funcs,
                                      gpu_context.remaining_const_memory);
            }
 
            for (auto &in : inputs) {
                if (usage.find(in.get()) == usage.end()) {
                    usage[in.get()] = 0;
                }
            }
 
            argument_set thread_shared;
            jit::register_map thread_mem;
 
            gpu_context.create_kernel_prefix(source_buffer,
                                             name, inputs, outputs,
                                             atomics, state,
                                             size, is_constant,
                                             registers, usage,
                                             kernel_1dtextures[name],
                                             kernel_2dtextures[name],
                                             thread_shared,
                                             thread_mem,
                                             iterations);
 
            for (auto &[out, in] : setters) {
                out->compile(source_buffer, registers, thread_mem, usage);
            }
            for (auto &out : outputs) {
                out->compile(source_buffer, registers, thread_mem, usage);
            }
 
            gpu_context.create_kernel_postfix(source_buffer, outputs,
                                              setters, state,
                                              registers, usage,
                                              thread_shared,
                                              thread_mem,
                                              iterations);
 
//  Delete the registers so that they can be used again in other kernels.
            std::vector<void *> removed_elements;
            for (auto &[key, value] : registers) {
                if (value[0] == 'r' ||
                    value[0] == 'l' ||
                    value[0] == 'i') {
                    removed_elements.push_back(key);
                }
            }
            for (auto &out : outputs) {
                if (graph::atomic_accumulate_1D_cast(out).get()) {
                    removed_elements.push_back(out.get());
                }
            }
 
            for (auto &key : removed_elements) {
                registers.erase(key);
            }
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void add_max_reduction(const size_t size) {
            gpu_context.create_reduction(source_buffer, size);
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        std::function<void(void)> create_zero_call(graph::input_nodes<T, SAFE_MATH> inputs) {
            return gpu_context.create_zero_call(inputs);
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        std::function<void(void)> create_copy_call(graph::copy_nodes<T, SAFE_MATH> setters) {
            return gpu_context.create_copy_call(setters);
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void print_source() {
            std::cout << std::endl << source_buffer.str() << std::endl;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void save_source() {
            std::string source = source_buffer.str();
            std::ostringstream filename;
            filename << std::hash<std::string> {} (source)
                     << std::hash<std::thread::id>{}(std::this_thread::get_id());
            if constexpr (use_cuda()) {
                filename << ".cu";
            } else if constexpr (use_metal<T> ()) {
                filename << ".metal";
            } else {
                filename << ".cpp";
            }
 
            std::ofstream outFile(filename.str());
            outFile << source;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void compile(const bool add_reduction=false) {
#ifdef SAVE_KERNEL_SOURCE
            save_source();
#endif
            gpu_context.compile(source_buffer.str(),
                                kernel_names,
                                add_reduction);
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        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) {
            return gpu_context.create_kernel_call(kernel_name, inputs, outputs,
                                                  atomics, state, num_rays,
                                                  kernel_1dtextures[kernel_name],
                                                  kernel_2dtextures[kernel_name]);
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        std::function<T(void)> create_max_call(graph::shared_leaf<T, SAFE_MATH> &argument,
                                               std::function<void(void)> run) {
            return gpu_context.create_max_call(argument, run);
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void print(const size_t index,
                   const graph::output_nodes<T, SAFE_MATH> &nodes) {
            gpu_context.print_results(index, nodes);
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        T check_value(const size_t index,
                      const graph::shared_leaf<T, SAFE_MATH> &node) {
            return gpu_context.check_value(index, node);
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void wait() {
            gpu_context.wait();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        std::function<void(void)> run_function(std::function<void(void)> callback) {
            return gpu_context.run_function(callback);
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void copy_to_device(graph::shared_leaf<T, SAFE_MATH> &node,
                            T *source) {
            gpu_context.copy_to_device(node, source);
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void copy_to_host(graph::shared_leaf<T, SAFE_MATH> &node,
                          T *destination) {
            gpu_context.copy_to_host(node, destination);
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        T *get_buffer(graph::shared_leaf<T, SAFE_MATH> &node) {
            return gpu_context.get_buffer(node);
        }
    };
}
 
#endif /* jit_h */