random.hpp

Go to the documentation of this file.

//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
 
#ifndef random_h
#define random_h
 
#include "node.hpp"
 
namespace graph {
//******************************************************************************
//******************************************************************************
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    class random_state_node final : public leaf_node<T, SAFE_MATH> {
    public:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        static void compile_random_state(std::ostringstream &stream) {
            stream << "struct mt_state {"               << std::endl
                   << "    array<uint32_t, 624> array;" << std::endl
                   << "    uint16_t index;"             << std::endl
#ifdef USE_CUDA
                   << "    uint16_t padding[3];"        << std::endl
#endif
                   << "};"                              << std::endl;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        struct mt_state {
            std::array<uint32_t, 624> array;
            uint16_t index;
#ifdef USE_CUDA
            uint16_t padding[3];
#endif
        };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        random_state_node(const size_t size,
                          const uint32_t seed=0) :
        leaf_node<T, SAFE_MATH> (random_state_node::to_string(), 1, false) {
            for (uint32_t i = 0; i < size; i++) {
                states.push_back(initialize_state(seed + i));
            }
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual backend::buffer<T> evaluate() {
            backend::buffer<T> result;
            return result;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T, SAFE_MATH> df(shared_leaf<T, SAFE_MATH> x) {
            return zero<T, SAFE_MATH> ();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual void compile_preamble(std::ostringstream &stream,
                                      jit::register_map &registers,
                                      jit::visiter_map &visited,
                                      jit::register_usage &usage,
                                      jit::texture1d_list &textures1d,
                                      jit::texture2d_list &textures2d,
                                      jit::preamble_map &pre_funcs,
                                      int &avail_const_mem) {
            if (!pre_funcs.contains("random_state")) {
                pre_funcs.insert("random_state");
 
                random_state_node::compile_random_state(stream);
            }
 
            if (!visited.contains(this)) {
                visited.insert(this);
#ifdef SHOW_USE_COUNT
                usage[this] = 1;
            } else {
                ++usage[this];
#endif
            }
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T, SAFE_MATH>
        compile(std::ostringstream &stream,
                jit::register_map &registers,
                const jit::register_map &thread_mem,
                const jit::register_usage &usage) {
            return this->shared_from_this();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual void to_latex() const {
            std::cout << "state";
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T, SAFE_MATH> to_vizgraph(std::stringstream &stream,
                                                      jit::register_map &registers) {
            if (registers.find(this) == registers.end()) {
                const std::string name = jit::to_string('r', this);
                registers[this] = name;
                stream << "    " << name
                       << " [label = \"state\", shape = box, style = \"rounded,filled\", fillcolor = black, fontcolor = white];" << std::endl;
            }
 
            return this->shared_from_this();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T, SAFE_MATH> get_power_exponent() const {
            return constant<T, SAFE_MATH> (static_cast<T> (1.0));
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        size_t size() {
            return states.size();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        size_t get_size_bytes() {
            return size()*sizeof(mt_state);
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        mt_state *data() {
            return states.data();
        }
 
    private:
        std::vector<mt_state> states;
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        static std::string to_string() {
            return "random_state";
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        mt_state initialize_state(const uint32_t seed) {
            mt_state state;
            state.array[0] = seed;
            for (uint16_t i = 1, ie = state.array.size(); i < ie; i++) {
                state.array[i] = 1812433253U*(state.array[i - 1]^(state.array[i - 1] >> 30)) + i;
            }
            state.index = 0;
 
            return state;
        }
    };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    shared_leaf<T, SAFE_MATH> random_state(const size_t size,
                                           const uint32_t seed=0) {
        auto temp = std::make_shared<random_state_node<T, SAFE_MATH>> (size, seed)->reduce();
//  Test for hash collisions.
        for (size_t i = temp->get_hash();
             i < std::numeric_limits<size_t>::max(); i++) {
            if (leaf_node<T, SAFE_MATH>::caches.nodes.find(i) ==
                leaf_node<T, SAFE_MATH>::caches.nodes.end()) {
                leaf_node<T, SAFE_MATH>::caches.nodes[i] = temp;
                return temp;
            } else if (temp->is_match(leaf_node<T, SAFE_MATH>::caches.nodes[i])) {
                return leaf_node<T, SAFE_MATH>::caches.nodes[i];
            }
        }
#if defined(__clang__) || defined(__GNUC__)
        __builtin_unreachable();
#else
        assert(false && "Should never reach.");
#endif
    }
 
    template<jit::float_scalar T, bool SAFE_MATH=false>
    using shared_random_state = std::shared_ptr<random_state_node<T, SAFE_MATH>>;
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    shared_random_state<T, SAFE_MATH> random_state_cast(shared_leaf<T, SAFE_MATH> x) {
        return std::dynamic_pointer_cast<random_state_node<T, SAFE_MATH>> (x);
    }
 
//******************************************************************************
//  Random constant.
//******************************************************************************
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    class random_node final : public straight_node<T, SAFE_MATH> {
    private:
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        static std::string to_string() {
            return "random";
        }
 
    public:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        static void compile_random(std::ostringstream &stream) {
            stream << "uint32_t random(";
            if constexpr (jit::use_metal<T> ()) {
                stream << "device ";
            }
            stream << "mt_state &state) {"                                         << std::endl
                   << "    const uint16_t k = state.index;"                        << std::endl
                   << "    state.index = (k + 1) % 624;"                           << std::endl
                   << "    uint32_t x = (state.array[k] & 0x80000000U) |"          << std::endl
                   << "                 (state.array[state.index] & 0x7fffffffU);" << std::endl
                   << "    uint32_t xA = x >> 1;"                                  << std::endl
                   << "    xA = x & 0x1U ? xA^0x9908b0dfU : xA;"                   << std::endl
                   << "    const uint16_t j = (k + 397) % 624;"                    << std::endl
                   << "    x = state.array[j]^xA;"                                 << std::endl
                   << "    state.array[k] = x;"                                    << std::endl
                   << "    uint32_t y = x^(x >> 11);"                              << std::endl
                   << "    y = y^((y << 7) & 0x9d2c5680U);"                        << std::endl
                   << "    y = y^((y << 15) & 0xefc60000U);"                       << std::endl
                   << "    return y^(y >> 18);"                                    << std::endl
                   << "}"                                                          << std::endl;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        random_node(shared_random_state<T, SAFE_MATH> x) :
        straight_node<T, SAFE_MATH> (x, random_node::to_string()) {}
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual backend::buffer<T> evaluate() {
            backend::buffer<T> result;
            return result;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T, SAFE_MATH> df(shared_leaf<T, SAFE_MATH> x) {
            return zero<T, SAFE_MATH> ();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual void compile_preamble(std::ostringstream &stream,
                                      jit::register_map &registers,
                                      jit::visiter_map &visited,
                                      jit::register_usage &usage,
                                      jit::texture1d_list &textures1d,
                                      jit::texture2d_list &textures2d,
                                      jit::preamble_map &pre_funcs,
                                      int &avail_const_mem) {
            if (!visited.contains(this)) {
                this->arg->compile_preamble(stream, registers,
                                            visited, usage,
                                            textures1d, textures2d,
                                            pre_funcs, avail_const_mem);
 
                visited.insert(this);
#ifdef SHOW_USE_COUNT
                usage[this] = 1;
            } else {
                ++usage[this];
#endif
            }
 
//  Need to do this after visited was checked so the random_state is created
//  first.
            if (!pre_funcs.contains("random")) {
                pre_funcs.insert("random");
 
                random_node::compile_random(stream);
            }
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T, SAFE_MATH>
        compile(std::ostringstream &stream,
                jit::register_map &registers,
                const jit::register_map &thread_mem,
                const jit::register_usage &usage) {
            if (registers.find(this) == registers.end()) {
                auto a = this->arg->compile(stream, registers,
                                            thread_mem, usage);
 
                if constexpr (jit::complex_scalar<T>) {
                    registers[this] = "static_cast<"
                                    + jit::get_type_string<T> ()
                                    + "> (random("
                                    + registers[a.get()] + "))";
                } else {
                    registers[this] = "random(" + registers[a.get()] + ")";
                }
            }
 
            return this->shared_from_this();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual bool is_match(shared_leaf<T, SAFE_MATH> x) {
            return false;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual void to_latex() const {
            std::cout << "random(";
            this->arg->to_latex();
            std::cout << ")";
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T, SAFE_MATH> to_vizgraph(std::stringstream &stream,
                                                      jit::register_map &registers) {
            if (registers.find(this) == registers.end()) {
                const std::string name = jit::to_string('r', this);
                registers[this] = name;
                stream << "    " << name
                       << " [label = \"state\", shape = box, style = \"rounded,filled\", fillcolor = black, fontcolor = white];" << std::endl;
 
                auto a = this->arg->to_vizgraph(stream, registers);
                stream << "    " << name << " -- " << registers[a.get()] << ";" << std::endl;
            }
 
            return this->shared_from_this();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T, SAFE_MATH> get_power_exponent() const {
            return constant<T, SAFE_MATH> (static_cast<T> (1.0));
        }
    };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    shared_leaf<T, SAFE_MATH> random(shared_random_state<T, SAFE_MATH> state) {
        auto temp = std::make_shared<random_node<T, SAFE_MATH>> (state)->reduce();
//  Test for hash collisions.
        for (size_t i = temp->get_hash();
             i < std::numeric_limits<size_t>::max(); i++) {
            if (leaf_node<T, SAFE_MATH>::caches.nodes.find(i) ==
                leaf_node<T, SAFE_MATH>::caches.nodes.end()) {
                leaf_node<T, SAFE_MATH>::caches.nodes[i] = temp;
                return temp;
            } else if (temp->is_match(leaf_node<T, SAFE_MATH>::caches.nodes[i])) {
                return leaf_node<T, SAFE_MATH>::caches.nodes[i];
            }
        }
#if defined(__clang__) || defined(__GNUC__)
        __builtin_unreachable();
#else
        assert(false && "Should never reach.");
#endif
    }
 
    template<jit::float_scalar T, bool SAFE_MATH=false>
    using shared_random = std::shared_ptr<random_node<T, SAFE_MATH>>;
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    shared_random<T, SAFE_MATH> random_cast(shared_leaf<T, SAFE_MATH> x) {
        return std::dynamic_pointer_cast<random_node<T, SAFE_MATH>> (x);
    }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    constexpr shared_leaf<T, SAFE_MATH> random_scale() {
        return constant<T, SAFE_MATH> (static_cast<T> (std::numeric_limits<uint32_t>::max()));
    }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, jit::float_scalar R, bool SAFE_MATH=false>
    constexpr shared_leaf<T, SAFE_MATH> uniform_random(const R min,
                                                       const R max,
                                                       shared_random_state<T, SAFE_MATH> state) {
        auto random = graph::random<T> (state);
        return (max - min)/graph::random_scale<T> ()*random + min;
    }
}
 
#endif /* random_h */