normal.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 <chrono> //for std::time

namespace etl {

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

    const T mean;                                      
    const T stddev;                                    
    random_engine rand_engine;                         
    std::normal_distribution<value_type> distribution; 

    static constexpr bool gpu_computable = (is_single_precision_t<T> && curand_enabled) || (is_double_precision_t<T> && curand_enabled);

    normal_generator_op(T mean, T stddev) : mean(mean), stddev(stddev), rand_engine(std::time(nullptr)), distribution(mean, stddev) {}

    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
        impl::curand::generate_normal(gen, t1.gpu_memory(), etl::size(y), mean, stddev);

        // Destroy the generator
        curand_call(curandDestroyGenerator(gen));

        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
        impl::curand::generate_normal(gen, y.gpu_memory(), etl::size(y), mean, stddev);

        // Destroy the generator
        curand_call(curandDestroyGenerator(gen));

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

        return y;
    }

#endif

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

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

    const T mean;                                      
    const T stddev;                                    
    G& rand_engine;                                    
    std::normal_distribution<value_type> distribution; 

    static constexpr bool gpu_computable = (is_single_precision_t<T> && curand_enabled) || (is_double_precision_t<T> && curand_enabled);

    normal_generator_g_op(G& g, T mean, T stddev) : mean(mean), stddev(stddev), rand_engine(g), distribution(mean, stddev) {}

    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
        impl::curand::generate_normal(gen, t1.gpu_memory(), etl::size(y), mean, stddev);

        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
        impl::curand::generate_normal(gen, y.gpu_memory(), etl::size(y), mean, stddev);

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

        return y;
    }

#endif

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

} //end of namespace etl