uniform.hpp

Go to the documentation of this file.

//=======================================================================
// Copyright (c) 2014-2023 Baptiste Wicht
// Distributed under the terms of the MIT License.
// (See accompanying file LICENSE or copy at
//  http://opensource.org/licenses/MIT)
//=======================================================================

#pragma once

#ifdef ETL_CURAND_MODE
#include "etl/impl/curand/curand.hpp"
#endif

#include "etl/impl/egblas/scalar_add.hpp"
#include "etl/impl/egblas/scalar_mul.hpp"

#include <chrono> //for std::time

namespace etl {

template <typename T>
using uniform_distribution = std::conditional_t<std::is_floating_point_v<T>, std::uniform_real_distribution<T>, std::uniform_int_distribution<T>>;

template <typename T = double>
struct uniform_generator_op {
    using value_type = T; 

    const T start;                                 
    const T end;                                   
    random_engine rand_engine;                     
    uniform_distribution<value_type> distribution; 

    static constexpr bool gpu_computable = curand_enabled
                                           && ((is_single_precision_t<T> && impl::egblas::has_scalar_sadd && impl::egblas::has_scalar_smul)
                                               || (is_double_precision_t<T> && impl::egblas::has_scalar_dadd && impl::egblas::has_scalar_dmul));

    uniform_generator_op(T start, T end) : start(start), end(end), rand_engine(std::time(nullptr)), distribution(start, end) {}

    value_type operator()() {
        return distribution(rand_engine);
    }

#ifdef ETL_CURAND_MODE

    template <typename Y>
    auto gpu_compute_hint(Y& y) noexcept {
        auto t1 = etl::force_temporary_gpu_dim_only_t<T>(y);

        curandGenerator_t gen;

        // Create the generator
        curand_call(curandCreateGenerator(&gen, CURAND_RNG_PSEUDO_DEFAULT));

        // Seed it with the internal random engine
        std::uniform_int_distribution<long> seed_dist;
        curand_call(curandSetPseudoRandomGeneratorSeed(gen, seed_dist(rand_engine)));

        // Generate the random numbers in [0,1]
        impl::curand::generate_uniform(gen, t1.gpu_memory(), etl::size(y));

        // mul by b-a => [0,b-a]
        auto s1 = T(end) - T(start);
        impl::egblas::scalar_mul(t1.gpu_memory(), etl::size(y), 1, s1);

        // Add a => [a,b]
        auto s2 = T(start);
        impl::egblas::scalar_add(t1.gpu_memory(), etl::size(y), 1, s2);

        return t1;
    }

    template <typename Y>
    Y& gpu_compute(Y& y) noexcept {
        y.ensure_gpu_allocated();

        curandGenerator_t gen;

        // Create the generator
        curand_call(curandCreateGenerator(&gen, CURAND_RNG_PSEUDO_DEFAULT));

        // Seed it with the internal random engine
        std::uniform_int_distribution<long> seed_dist;
        curand_call(curandSetPseudoRandomGeneratorSeed(gen, seed_dist(rand_engine)));

        // Generate the random numbers in [0,1]
        impl::curand::generate_uniform(gen, y.gpu_memory(), etl::size(y));

        // mul by b-a => [0,b-a]
        auto s1 = T(end) - T(start);
        impl::egblas::scalar_mul(y.gpu_memory(), etl::size(y), 1, s1);

        // Add a => [a,b]
        auto s2 = T(start);
        impl::egblas::scalar_add(y.gpu_memory(), etl::size(y), 1, s2);

        y.validate_gpu();
        y.invalidate_cpu();

        return y;
    }

#endif

    friend std::ostream& operator<<(std::ostream& os, [[maybe_unused]] const uniform_generator_op& s) {
        return os << "U(0,1)";
    }
};

template <typename G, typename T = double>
struct uniform_generator_g_op {
    using value_type = T; 

    const T start;                                 
    const T end;                                   
    G& rand_engine;                                
    uniform_distribution<value_type> distribution; 

    static constexpr bool gpu_computable = curand_enabled
                                           && ((is_single_precision_t<T> && impl::egblas::has_scalar_sadd && impl::egblas::has_scalar_smul)
                                               || (is_double_precision_t<T> && impl::egblas::has_scalar_dadd && impl::egblas::has_scalar_dmul));

    uniform_generator_g_op(G& g, T start, T end) : start(start), end(end), rand_engine(g), distribution(start, end) {}

    value_type operator()() {
        return distribution(rand_engine);
    }

#ifdef ETL_CURAND_MODE

    template <typename Y>
    auto gpu_compute_hint(Y& y) noexcept {
        auto t1 = etl::force_temporary_gpu_dim_only_t<T>(y);

        curandGenerator_t gen;

        // Create the generator
        curand_call(curandCreateGenerator(&gen, CURAND_RNG_PSEUDO_DEFAULT));

        // Seed it with the internal random engine
        std::uniform_int_distribution<long> seed_dist;
        curand_call(curandSetPseudoRandomGeneratorSeed(gen, seed_dist(rand_engine)));

        // Generate the random numbers in [0,1]
        impl::curand::generate_uniform(gen, t1.gpu_memory(), etl::size(y));

        // mul by b-a => [0,b-a]
        auto s1 = T(end) - T(start);
        impl::egblas::scalar_mul(t1.gpu_memory(), etl::size(y), 1, s1);

        // Add a => [a,b]
        auto s2 = T(start);
        impl::egblas::scalar_add(t1.gpu_memory(), etl::size(y), 1, s2);

        return t1;
    }

    template <typename Y>
    Y& gpu_compute(Y& y) noexcept {
        y.ensure_gpu_allocated();

        curandGenerator_t gen;

        // Create the generator
        curand_call(curandCreateGenerator(&gen, CURAND_RNG_PSEUDO_DEFAULT));

        // Seed it with the internal random engine
        std::uniform_int_distribution<long> seed_dist;
        curand_call(curandSetPseudoRandomGeneratorSeed(gen, seed_dist(rand_engine)));

        // Generate the random numbers in [0,1]
        impl::curand::generate_uniform(gen, y.gpu_memory(), etl::size(y));

        // mul by b-a => [0,b-a]
        auto s1 = T(end) - T(start);
        impl::egblas::scalar_mul(y.gpu_memory(), etl::size(y), 1, s1);

        // Add a => [a,b]
        auto s2 = T(start);
        impl::egblas::scalar_add(y.gpu_memory(), etl::size(y), 1, s2);

        y.validate_gpu();
        y.invalidate_cpu();

        return y;
    }

#endif

    friend std::ostream& operator<<(std::ostream& os, [[maybe_unused]] const uniform_generator_g_op& s) {
        return os << "U(0,1)";
    }
};

} //end of namespace etl