SDE

CMakeLists.txt

@@ -46,6 +46,8 @@ option(ARB_BACKTRACE "Enable stacktraces on assertion and exceptions (requires B
 
 set(ARB_GPU "none" CACHE STRING "GPU backend and compiler configuration")
 set_property(CACHE PROPERTY STRINGS "none" "cuda" "cuda-clang" "hip")
+option(ARB_USE_GPU_RNG
+    "Use GPU generated random numbers (only cuda, not bitwise equal to CPU version)" OFF)
 
 # Use bundled 3rd party libraries
 
@@ -405,6 +407,15 @@ if(ARB_VECTORIZE)
     list(APPEND ARB_MODCC_FLAGS "--simd")
 endif()
 
+# Random number creation
+# -----------------------------------------------
+
+if(ARB_USE_GPU_RNG AND (ARB_WITH_NVCC OR ARB_WITH_CUDA_CLANG))
+    set(ARB_USE_GPU_RNG_IMPL TRUE)
+else()
+    set(ARB_USE_GPU_RNG_IMPL FALSE)
+endif()
+
 #----------------------------------------------------------
 # Set up install paths, permissions.
 #----------------------------------------------------------

arbor/CMakeLists.txt

@@ -4,6 +4,7 @@ set(arbor_sources
     arbexcept.cpp
     assert.cpp
     backends/multicore/shared_state.cpp
+    backends/multicore/rand.cpp
     communication/communicator.cpp
     communication/dry_run_context.cpp
     benchmark_cell_group.cpp
@@ -81,8 +82,15 @@ if(ARB_WITH_GPU)
         backends/gpu/forest.cpp
         backends/gpu/stimulus.cu
         backends/gpu/threshold_watcher.cu
+        backends/gpu/rand.cpp
         memory/fill.cu
     )
+    if (ARB_USE_GPU_RNG_IMPL)
+        list(APPEND arbor_sources backends/gpu/rand.cu)
+    else()
+        set_source_files_properties(backends/gpu/rand.cpp PROPERTIES
+            COMPILE_DEFINITIONS ARB_ARBOR_NO_GPU_RAND)
+    endif()
 endif()
 
 if(ARB_WITH_MPI)

arbor/backends/gpu/chunk_writer.hpp

@@ -0,0 +1,54 @@
+#include <cstddef>
+
+#include "memory/memory.hpp"
+#include "memory/copy.hpp"
+#include "util/meta.hpp"
+
+#pragma once
+
+namespace arb {
+namespace gpu {
+namespace {
+template <typename T>
+struct chunk_writer {
+    T* end; // device ptr
+    const std::size_t stride;
+
+    chunk_writer(T* data, std::size_t stride): end(data), stride(stride) {}
+
+    template <typename Seq, typename = std::enable_if_t<util::is_contiguous_v<Seq>>>
+    T* append(Seq&& seq) {
+        arb_assert(std::size(seq)==stride);
+        return append_freely(std::forward<Seq>(seq));
+    }
+
+    template <typename Seq, typename = std::enable_if_t<util::is_contiguous_v<Seq>>>
+    T* append_freely(Seq&& seq) {
+        std::size_t n = std::size(seq);
+        memory::copy(memory::host_view<T>(const_cast<T*>(std::data(seq)), n), memory::device_view<T>(end, n));
+        auto p = end;
+        end += n;
+        return p;
+    }
+
+    T* fill(T value) {
+        memory::fill(memory::device_view<T>(end, stride), value);
+        auto p = end;
+        end += stride;
+        return p;
+    }
+
+    template <typename Seq, typename = std::enable_if_t<util::is_contiguous_v<Seq>>>
+    T* append_with_padding(Seq&& seq, typename util::sequence_traits<Seq>::value_type value) {
+        std::size_t n = std::size(seq);
+        arb_assert(n <= stride);
+        std::size_t r = stride - n;
+        auto p = append_freely(std::forward<Seq>(seq));
+        memory::fill(memory::device_view<typename util::sequence_traits<Seq>::value_type>(end, r), value);
+        end += r;
+        return p;
+    }
+};
+} // anonymous namespace
+} // namespace gpu
+} // namespace arb

arbor/backends/gpu/gpu_store_types.hpp

@@ -17,6 +17,7 @@ namespace gpu {
 
 using array  = memory::device_vector<arb_value_type>;
 using iarray = memory::device_vector<arb_index_type>;
+using sarray = memory::device_vector<arb_size_type>;
 
 using deliverable_event_stream = arb::gpu::multi_event_stream<deliverable_event>;
 using sample_event_stream = arb::gpu::multi_event_stream<sample_event>;

arbor/backends/gpu/rand.cpp

@@ -0,0 +1,62 @@
+#include <util/span.hpp>
+#include <backends/gpu/chunk_writer.hpp>
+#include <backends/gpu/rand.hpp>
+
+#ifdef ARB_ARBOR_NO_GPU_RAND
+#include <backends/gpu/rand_on_cpu.hpp>
+#else
+#include <backends/gpu/rand_on_gpu.hpp>
+#endif
+
+ARB_INSTANTIATE_PIMPL(arb::gpu::random_numbers::impl)
+
+namespace arb {
+namespace gpu {
+
+void random_numbers::instantiate(mechanism& m, std::size_t value_width_padded,
+    const mechanism_layout& pos_data, arb_seed_type seed) {
+
+    using util::make_span;
+
+    // Allocate view pointers for random nubers
+    std::size_t num_random_numbers_per_cv = m.mech_.n_random_variables;
+    std::size_t random_number_storage = num_random_numbers_per_cv*cbprng::cache_size();
+    for (auto& v : random_numbers_) v.resize(num_random_numbers_per_cv);
+
+    // Allocate bulk storage
+    std::size_t count = random_number_storage*value_width_padded;
+    data_ = array(count, NAN);
+
+    // Set random numbers
+    chunk_writer writer(data_.data(), value_width_padded);
+    for (auto idx_v: make_span(num_random_numbers_per_cv)) {
+        for (auto idx_c: make_span(cbprng::cache_size())) {
+            random_numbers_[idx_c][idx_v] = writer.fill(0);
+        }
+    }
+
+    // Shift data to GPU
+    for (auto idx_c: make_span(cbprng::cache_size()))
+        random_numbers_d_[idx_c] = memory::on_gpu(random_numbers_[idx_c]);
+
+    // Instantiate implementation details
+    impl_ = util::make_pimpl<impl>(random_numbers_[0][0], m, value_width_padded, pos_data, seed);
+}
+
+void random_numbers::update(mechanism& m) {
+    // Assign new random numbers by selecting the next cache
+    const auto counter = random_number_update_counter_++;
+    const auto cache_idx = cbprng::cache_index(counter);
+    m.ppack_.random_numbers = random_numbers_d_[cache_idx].data();
+
+    // Generate random numbers every cbprng::cache_size() iterations:
+    // For each random variable we will generate cbprng::cache_size() values per site
+    // and there are width sites.
+    // The RNG will be seeded by a global seed, the mechanism id, the variable index, the
+    // current site's global cell, the site index within its cell and a counter representing
+    // time.
+    if (cache_idx == 0) impl_->update(m, counter);
+}
+
+} // namespace gpu
+} // namespace arb

arbor/backends/gpu/rand.cu

@@ -0,0 +1,64 @@
+#include <arbor/gpu/gpu_api.hpp>
+#include <arbor/gpu/gpu_common.hpp>
+
+#include "backends/rand_impl.hpp"
+
+namespace arb {
+namespace gpu {
+
+namespace kernel {
+__global__
+void generate_random_numbers(
+    arb_value_type* __restrict__ dst,
+    std::size_t width,
+    std::size_t width_padded,
+    std::size_t num_rv,
+    arb::cbprng::value_type seed,
+    arb::cbprng::value_type mech_id,
+    arb::cbprng::value_type counter,
+    arb_size_type const * __restrict__ gids,
+    arb_size_type const * __restrict__ idxs,
+    unsigned cache_size) {
+    // location and variable number extracted from thread block
+    const int i = threadIdx.x + blockDim.x*blockIdx.x;
+    const arb::cbprng::value_type n = blockIdx.y;
+
+    if (i < width) {
+        const arb::cbprng::value_type gid = gids[i];
+        const arb::cbprng::value_type idx = idxs[i];
+        const auto r = arb::cbprng::generator{}(arb::cbprng::array_type{seed, mech_id, n, counter},
+            arb::cbprng::array_type{gid, idx, 0xdeadf00dull, 0xdeadbeefull});
+        const auto a = r123::boxmuller(r[0], r[1]);
+        const auto b = r123::boxmuller(r[2], r[3]);
+        dst[i + width_padded*(0 + cache_size*n)] = a.x;
+        dst[i + width_padded*(1 + cache_size*n)] = a.y;
+        dst[i + width_padded*(2 + cache_size*n)] = b.x;
+        dst[i + width_padded*(3 + cache_size*n)] = b.y;
+    }
+}
+} // namespace kernel
+
+void generate_random_numbers(
+    arb_value_type* dst,        // points to random number storage
+    std::size_t width,          // number of sites
+    std::size_t width_padded,   // padded number of sites
+    std::size_t num_rv,         // number of random variables
+    cbprng::value_type seed,    // simulation seed value
+    cbprng::value_type mech_id, // mechanism id
+    cbprng::value_type counter, // step counter
+    arb_size_type const * gid,  // global cell ids (size = width)
+    arb_size_type const * idx   // per-cell location index (size = width)
+    ) {
+    using impl::block_count;
+
+    unsigned const block_dim = 128;
+    unsigned const grid_dim_x = block_count(width, block_dim);
+    unsigned const grid_dim_y = num_rv;
+
+    kernel::generate_random_numbers<<<dim3{grid_dim_x, grid_dim_y, 1}, block_dim>>>(
+        dst, width, width_padded, num_rv, seed, mech_id, counter, gid, idx, cbprng::cache_size());
+}
+
+} // namespace gpu
+} // namespace arb
+

arbor/backends/gpu/rand.hpp

@@ -0,0 +1,37 @@
+#pragma once
+
+#include <array>
+#include <vector>
+
+#include <arbor/mechanism.hpp>
+
+#include <util/pimpl.hpp>
+#include <backends/rand_fwd.hpp>
+#include <backends/gpu/gpu_store_types.hpp>
+
+namespace arb {
+namespace gpu {
+
+class random_numbers {
+public:
+    void instantiate(mechanism& m, std::size_t value_width_padded, const mechanism_layout& pos_data,
+        arb_seed_type seed);
+
+    void update(mechanism& m);
+
+  private:
+    // random number device storage
+    array data_;
+    std::array<std::vector<arb_value_type*>, cbprng::cache_size()>           random_numbers_;
+    std::array<memory::device_vector<arb_value_type*>, cbprng::cache_size()> random_numbers_d_;
+
+    // time step counter
+    cbprng::counter_type random_number_update_counter_ = 0u;
+
+    // pimpl
+    struct impl;
+    util::pimpl<impl> impl_;
+};
+
+} // namespace gpu
+} // namespace arb

arbor/backends/gpu/rand_on_cpu.hpp

@@ -0,0 +1,47 @@
+#pragma once
+
+#include <util/pimpl_src.hpp>
+#include <memory/gpu_wrappers.hpp>
+#include <backends/gpu/rand.hpp>
+
+namespace arb {
+namespace gpu {
+
+struct random_numbers::impl {
+    // pointer to random number device storage
+    arb_value_type* random_numbers_;
+
+    // general parameters
+    std::size_t value_width_padded_;
+    arb_seed_type seed_;
+
+    // auxillary random number device storage
+    std::vector<arb_value_type> random_numbers_h_;
+    std::vector<arb_size_type> gid_;
+    std::vector<arb_size_type> idx_;
+
+    impl(arb_value_type* data, mechanism& m, std::size_t value_width_padded,
+        const mechanism_layout& pos_data, arb_seed_type seed) :
+        random_numbers_{data},
+        value_width_padded_{value_width_padded},
+        seed_{seed},
+        random_numbers_h_(m.mech_.n_random_variables * cbprng::cache_size() * value_width_padded),
+        gid_{pos_data.gid},
+        idx_{pos_data.idx} {}
+
+    void update(mechanism& m, cbprng::value_type counter) {
+        const auto num_rv = m.mech_.n_random_variables;
+        const auto width = m.ppack_.width;
+        const cbprng::value_type mech_id = m.mechanism_id();
+        // generate random numbers on the host
+        arb_value_type* dst = random_numbers_h_.data();
+        arb::multicore::generate_random_numbers(dst, width, value_width_padded_, num_rv, seed_,
+            mech_id, counter, gid_.data(), idx_.data());
+        // transfer values to device
+        memory::gpu_memcpy_h2d(random_numbers_, dst,
+            (num_rv*cbprng::cache_size()*value_width_padded_)*sizeof(arb_value_type));
+    }
+};
+
+} // namespace gpu
+} // namespace arb

arbor/backends/gpu/rand_on_gpu.hpp

@@ -0,0 +1,47 @@
+#pragma once
+
+#include <util/pimpl_src.hpp>
+#include <backends/gpu/rand.hpp>
+
+namespace arb {
+namespace gpu {
+
+struct random_numbers::impl {
+    // pointer to random number device storage
+    arb_value_type* random_numbers_;
+
+    // general parameters
+    std::size_t value_width_padded_;
+    arb_seed_type seed_;
+
+    // auxillary random number device storage
+    sarray sindices_;
+    std::vector<arb_size_type*> prng_indices_;
+    memory::device_vector<arb_size_type*> prng_indices_d_;
+
+    impl(arb_value_type* data, mechanism& m, std::size_t value_width_padded,
+        const mechanism_layout& pos_data, arb_seed_type seed) :
+        random_numbers_{data},
+        value_width_padded_{value_width_padded},
+        seed_{seed} {
+
+        sindices_ = sarray(2*value_width_padded);
+        chunk_writer writer(sindices_.data(), value_width_padded);
+        prng_indices_.resize(2);
+        prng_indices_[0] = writer.append_with_padding(pos_data.gid, 0);
+        prng_indices_[1] = writer.append_with_padding(pos_data.idx, 0);
+        prng_indices_d_ = memory::on_gpu(prng_indices_);
+    }
+
+    void update(mechanism& m, cbprng::value_type counter) {
+        const auto num_rv = m.mech_.n_random_variables;
+        const auto width = m.ppack_.width;
+        const cbprng::value_type mech_id = m.mechanism_id();
+        // generate random numbers on the device
+        generate_random_numbers(random_numbers_, width, value_width_padded_, num_rv, seed_,
+            mech_id, counter, prng_indices_[0], prng_indices_[1]);
+    }
+};
+
+} // namespace gpu
+} // namespace arb