From 8d90314220987859bb9e21f062a983003bf50416 Mon Sep 17 00:00:00 2001 From: routhleck <1310722434@qq.com> Date: Fri, 22 Dec 2023 11:34:37 +0800 Subject: [PATCH] Add new function for taichi random seeds initialization --- .../_src/math/jitconn/_event_matvec_taichi.py | 179 +++++++++--------- brainpy/_src/math/jitconn/_matvec_taichi.py | 103 +++++----- brainpy/_src/math/taichi_random.py | 36 +++- .../math/tests/taichi_random_time_test.py | 96 ++++++++++ brainpy/_src/math/tests/test_taichi_random.py | 56 +++--- brainpy/math/taichi_random.py | 3 + 6 files changed, 309 insertions(+), 164 deletions(-) create mode 100644 brainpy/_src/math/tests/taichi_random_time_test.py diff --git a/brainpy/_src/math/jitconn/_event_matvec_taichi.py b/brainpy/_src/math/jitconn/_event_matvec_taichi.py index 9bf018720..645e98979 100644 --- a/brainpy/_src/math/jitconn/_event_matvec_taichi.py +++ b/brainpy/_src/math/jitconn/_event_matvec_taichi.py @@ -15,7 +15,8 @@ from brainpy._src.math.taichi_random import (taichi_uniform_int_distribution as uniform_int_distribution, taichi_uniform_real_distribution as uniform_real_distribution, taichi_normal_distribution as normal_distribution, - taichi_lfsr88 as random_generator) + taichi_lfsr88 as random_generator_lfsr88, + init_lfsr88_seeds) ti = import_taichi() @@ -51,15 +52,15 @@ def _event_mv_prob_homo_bool_cpu( # ti.loop_config(serialize=True) for i_col in range(num_col): - seeds = ti.math.uvec4(seed_value + 1 + i_col, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_col) - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_row = uniform_int_distribution(result, 1, clen_value) event = events[i_col] while i_row < num_row: if event: out[i_row] += weight_value - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_row += uniform_int_distribution(result, 1, clen_value) @@ -83,15 +84,15 @@ def _event_mv_prob_homo_bool_gpu( i_col = i >> 5 index = i & 31 - seeds = ti.math.uvec4(seed_value + 1 + i_col, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_col) - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_row = uniform_int_distribution(result, 1, clen_value) + avg_num_uniform * index event = events[i_col] while i_row < num_row: if event: out[i_row] += weight_value - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_row += uniform_int_distribution(result, 1, clen_value) * 32 @@ -112,16 +113,16 @@ def _event_mv_prob_homo_outdim_parallel_bool_cpu( # ti.loop_config(serialize=True) for i_row in range(num_row): - seeds = ti.math.uvec4(seed_value + 1 + i_row, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_row) r = 0. - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_col = uniform_int_distribution(result, 1, clen_value) event = events[i_col] while i_col < num_col: if event: r += weight_value - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_col += uniform_int_distribution(result, 1, clen_value) out[i_row] = r @@ -146,16 +147,16 @@ def _event_mv_prob_homo_outdim_parallel_bool_gpu( i_row = i >> 5 index = i & 31 - seeds = ti.math.uvec4(seed_value + 1 + i_row, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_row) r = 0. - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_col = uniform_int_distribution(result, 1, clen_value) + avg_num_uniform * index event = events[i_col] while i_col < num_col: if event: r += weight_value - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_col += uniform_int_distribution(result, 1, clen_value) * 32 out[i_row] += r @@ -177,15 +178,15 @@ def _event_mv_prob_homo_cpu( # ti.loop_config(serialize=True) for i_col in range(num_col): - seeds = ti.math.uvec4(seed_value + 1 + i_col, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_col) - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_row = uniform_int_distribution(result, 1, clen_value) event = events[i_col] while i_row < num_row: if event > 0.: out[i_row] += weight_value - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_row += uniform_int_distribution(result, 1, clen_value) @@ -209,15 +210,15 @@ def _event_mv_prob_homo_gpu( i_col = i >> 5 index = i & 31 - seeds = ti.math.uvec4(seed_value + 1 + i_col, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_col) - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_row = uniform_int_distribution(result, 1, clen_value) + avg_num_uniform * index event = events[i_col] while i_row < num_row: if event > 0.: out[i_row] += weight_value - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_row += uniform_int_distribution(result, 1, clen_value) * 32 @@ -238,16 +239,16 @@ def _event_mv_prob_homo_outdim_parallel_cpu( # ti.loop_config(serialize=True) for i_row in range(num_row): - seeds = ti.math.uvec4(seed_value + 1 + i_row, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_row) r = 0. - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_col = uniform_int_distribution(result, 1, clen_value) event = events[i_col] while i_col < num_col: if event > 0.: r += weight_value - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_col += uniform_int_distribution(result, 1, clen_value) out[i_row] = r @@ -272,16 +273,16 @@ def _event_mv_prob_homo_outdim_parallel_gpu( i_row = i >> 5 index = i & 31 - seeds = ti.math.uvec4(seed_value + 1 + i_row, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value+ i_row) r = 0. - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_col = uniform_int_distribution(result, 1, clen_value) + avg_num_uniform * index event = events[i_col] while i_col < num_col: if event > 0.: r += weight_value - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_col += uniform_int_distribution(result, 1, clen_value) * 32 out[i_row] += r @@ -531,16 +532,16 @@ def _event_mv_prob_uniform_bool_cpu( # ti.loop_config(serialize=True) for i_col in range(num_col): - seeds = ti.math.uvec4(seed_value + 1 + i_col, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_col) - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_row = uniform_int_distribution(result, 1, clen_value) event = events[i_col] while i_row < num_row: if event: - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) out[i_row] += uniform_real_distribution(result, w_min_value, w_max_value) - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_row += uniform_int_distribution(result, 1, clen_value) @@ -567,16 +568,16 @@ def _event_mv_prob_uniform_bool_gpu( i_col = i >> 5 index = i & 31 - seeds = ti.math.uvec4(seed_value + 1 + i_col, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_col) - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_row = uniform_int_distribution(result, 1, clen_value) + avg_num_uniform * index event = events[i_col] while i_row < num_row: if event: - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) out[i_row] += uniform_real_distribution(result, w_min_value, w_max_value) - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_row += uniform_int_distribution(result, 1, clen_value) * 32 @@ -599,17 +600,17 @@ def _event_mv_prob_uniform_outdim_parallel_bool_cpu( # ti.loop_config(serialize=True) for i_row in range(num_row): - seeds = ti.math.uvec4(seed_value + 1 + i_row, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_row) r = 0. - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_col = uniform_int_distribution(result, 1, clen_value) event = events[i_col] while i_col < num_col: if event: - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) r += uniform_real_distribution(result, w_min_value, w_max_value) - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_col += uniform_int_distribution(result, 1, clen_value) out[i_row] = r @@ -636,17 +637,17 @@ def _event_mv_prob_uniform_outdim_parallel_bool_gpu( i_row = i >> 5 index = i & 31 - seeds = ti.math.uvec4(seed_value + 1 + i_row, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_row) r = 0. - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_col = uniform_int_distribution(result, 1, clen_value) + avg_num_uniform * index event = events[i_col] while i_col < num_col: if event: - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) r += uniform_real_distribution(result, w_min_value, w_max_value) - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_col += uniform_int_distribution(result, 1, clen_value) * 32 out[i_row] += r @@ -671,16 +672,16 @@ def _event_mv_prob_uniform_cpu( # ti.loop_config(serialize=True) for i_col in range(num_col): - seeds = ti.math.uvec4(seed_value + 1 + i_col, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_col) - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_row = uniform_int_distribution(result, 1, clen_value) event = events[i_col] while i_row < num_row: if event > 0.: - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) out[i_row] += uniform_real_distribution(result, w_min_value, w_max_value) - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_row += uniform_int_distribution(result, 1, clen_value) @@ -707,16 +708,16 @@ def _event_mv_prob_uniform_gpu( i_col = i >> 5 index = i & 31 - seeds = ti.math.uvec4(seed_value + 1 + i_col, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_col) - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_row = uniform_int_distribution(result, 1, clen_value) + avg_num_uniform * index event = events[i_col] while i_row < num_row: if event > 0.: - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) out[i_row] += uniform_real_distribution(result, w_min_value, w_max_value) - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_row += uniform_int_distribution(result, 1, clen_value) * 32 @@ -739,17 +740,17 @@ def _event_mv_prob_uniform_outdim_parallel_cpu( # ti.loop_config(serialize=True) for i_row in range(num_row): - seeds = ti.math.uvec4(seed_value + 1 + i_row, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_row) r = 0. - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_col = uniform_int_distribution(result, 1, clen_value) event = events[i_col] while i_col < num_col: if event > 0.: - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) r += uniform_real_distribution(result, w_min_value, w_max_value) - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_col += uniform_int_distribution(result, 1, clen_value) out[i_row] = r @@ -776,17 +777,17 @@ def _event_mv_prob_uniform_outdim_parallel_gpu( i_row = i >> 5 index = i & 31 - seeds = ti.math.uvec4(seed_value + 1 + i_row, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_row) r = 0. - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_col = uniform_int_distribution(result, 1, clen_value) + avg_num_uniform * index event = events[i_col] while i_col < num_col: if event > 0.: - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) r += uniform_real_distribution(result, w_min_value, w_max_value) - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_col += uniform_int_distribution(result, 1, clen_value) * 32 out[i_row] += r @@ -1030,15 +1031,15 @@ def _event_mv_prob_normal_bool_cpu( # ti.loop_config(serialize=True) for i_col in range(num_col): - seeds = ti.math.uvec4(seed_value + 1 + i_col, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_col) - seeds, result1 = random_generator(seeds) + seeds, result1 = random_generator_lfsr88(seeds) i_row = uniform_int_distribution(result1, 1, clen_value) event = events[i_col] while i_row < num_row: if event: - seeds, result1 = random_generator(seeds) - seeds, result2 = random_generator(seeds) + seeds, result1 = random_generator_lfsr88(seeds) + seeds, result2 = random_generator_lfsr88(seeds) out[i_row] += normal_distribution(result1, result2, w_mu_value, w_sigma_value) i_row += uniform_int_distribution(result1, 1, clen_value) @@ -1065,15 +1066,15 @@ def _event_mv_prob_normal_bool_gpu( i_col = i >> 5 index = i & 31 - seeds = ti.math.uvec4(seed_value + 1 + i_col, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_col) - seeds, result1 = random_generator(seeds) + seeds, result1 = random_generator_lfsr88(seeds) i_row = uniform_int_distribution(result1, 1, clen_value) + avg_num_uniform * index event = events[i_col] while i_row < num_row: if event: - seeds, result1 = random_generator(seeds) - seeds, result2 = random_generator(seeds) + seeds, result1 = random_generator_lfsr88(seeds) + seeds, result2 = random_generator_lfsr88(seeds) out[i_row] += normal_distribution(result1, result2, w_mu_value, w_sigma_value) i_row += uniform_int_distribution(result1, 1, clen_value) * 32 @@ -1097,16 +1098,16 @@ def _event_mv_prob_normal_outdim_parallel_bool_cpu( # ti.loop_config(serialize=True) for i_row in range(num_row): - seeds = ti.math.uvec4(seed_value + 1 + i_row, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_row) r = 0. - seeds, result1 = random_generator(seeds) + seeds, result1 = random_generator_lfsr88(seeds) i_col = uniform_int_distribution(result1, 1, clen_value) event = events[i_col] while i_col < num_col: if event: - seeds, result1 = random_generator(seeds) - seeds, result2 = random_generator(seeds) + seeds, result1 = random_generator_lfsr88(seeds) + seeds, result2 = random_generator_lfsr88(seeds) r += normal_distribution(result1, result2, w_mu_value, w_sigma_value) i_col += uniform_int_distribution(result1, 1, clen_value) out[i_row] = r @@ -1134,16 +1135,16 @@ def _event_mv_prob_normal_outdim_parallel_bool_gpu( i_row = i >> 5 index = i & 31 - seeds = ti.math.uvec4(seed_value + 1 + i_row, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_row) r = 0. - seeds, result1 = random_generator(seeds) + seeds, result1 = random_generator_lfsr88(seeds) i_col = uniform_int_distribution(result1, 1, clen_value) + avg_num_uniform * index event = events[i_col] while i_col < num_col: if event: - seeds, result1 = random_generator(seeds) - seeds, result2 = random_generator(seeds) + seeds, result1 = random_generator_lfsr88(seeds) + seeds, result2 = random_generator_lfsr88(seeds) r += normal_distribution(result1, result2, w_mu_value, w_sigma_value) i_col += uniform_int_distribution(result1, 1, clen_value) * 32 out[i_row] += r @@ -1168,15 +1169,15 @@ def _event_mv_prob_normal_cpu( # ti.loop_config(serialize=True) for i_col in range(num_col): - seeds = ti.math.uvec4(seed_value + 1 + i_col, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_col) - seeds, result1 = random_generator(seeds) + seeds, result1 = random_generator_lfsr88(seeds) i_row = uniform_int_distribution(result1, 1, clen_value) event = events[i_col] while i_row < num_row: if event > 0.: - seeds, result1 = random_generator(seeds) - seeds, result2 = random_generator(seeds) + seeds, result1 = random_generator_lfsr88(seeds) + seeds, result2 = random_generator_lfsr88(seeds) out[i_row] += normal_distribution(result1, result2, w_mu_value, w_sigma_value) i_row += uniform_int_distribution(result1, 1, clen_value) @@ -1203,15 +1204,15 @@ def _event_mv_prob_normal_gpu( i_col = i >> 5 index = i & 31 - seeds = ti.math.uvec4(seed_value + 1 + i_col, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_col) - seeds, result1 = random_generator(seeds) + seeds, result1 = random_generator_lfsr88(seeds) i_row = uniform_int_distribution(result1, 1, clen_value) + avg_num_uniform * index event = events[i_col] while i_row < num_row: if event > 0.: - seeds, result1 = random_generator(seeds) - seeds, result2 = random_generator(seeds) + seeds, result1 = random_generator_lfsr88(seeds) + seeds, result2 = random_generator_lfsr88(seeds) out[i_row] += normal_distribution(result1, result2, w_mu_value, w_sigma_value) i_row += uniform_int_distribution(result1, 1, clen_value) * 32 @@ -1235,16 +1236,16 @@ def _event_mv_prob_normal_outdim_parallel_cpu( # ti.loop_config(serialize=True) for i_row in range(num_row): - seeds = ti.math.uvec4(seed_value + 1 + i_row, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_row) r = 0. - seeds, result1 = random_generator(seeds) + seeds, result1 = random_generator_lfsr88(seeds) i_col = uniform_int_distribution(result1, 1, clen_value) event = events[i_col] while i_col < num_col: if event > 0.: - seeds, result1 = random_generator(seeds) - seeds, result2 = random_generator(seeds) + seeds, result1 = random_generator_lfsr88(seeds) + seeds, result2 = random_generator_lfsr88(seeds) r += normal_distribution(result1, result2, w_mu_value, w_sigma_value) i_col += uniform_int_distribution(result1, 1, clen_value) out[i_row] = r @@ -1272,16 +1273,16 @@ def _event_mv_prob_normal_outdim_parallel_gpu( i_row = i >> 5 index = i & 31 - seeds = ti.math.uvec4(seed_value + 1 + i_row, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_row) r = 0. - seeds, result1 = random_generator(seeds) + seeds, result1 = random_generator_lfsr88(seeds) i_col = uniform_int_distribution(result1, 1, clen_value) + avg_num_uniform * index event = events[i_col] while i_col < num_col: if event > 0.: - seeds, result1 = random_generator(seeds) - seeds, result2 = random_generator(seeds) + seeds, result1 = random_generator_lfsr88(seeds) + seeds, result2 = random_generator_lfsr88(seeds) r += normal_distribution(result1, result2, w_mu_value, w_sigma_value) i_col += uniform_int_distribution(result1, 1, clen_value) * 32 out[i_row] += r diff --git a/brainpy/_src/math/jitconn/_matvec_taichi.py b/brainpy/_src/math/jitconn/_matvec_taichi.py index 7e3f8d29e..47b323ca9 100644 --- a/brainpy/_src/math/jitconn/_matvec_taichi.py +++ b/brainpy/_src/math/jitconn/_matvec_taichi.py @@ -15,7 +15,8 @@ from brainpy._src.math.taichi_random import (taichi_uniform_int_distribution as uniform_int_distribution, taichi_uniform_real_distribution as uniform_real_distribution, taichi_normal_distribution as normal_distribution, - taichi_lfsr88 as random_generator) + taichi_lfsr88 as random_generator_lfsr88, + init_lfsr88_seeds as init_lfsr88_seeds) ti = import_taichi() @@ -50,15 +51,15 @@ def _mv_prob_homo_cpu( # ti.loop_config(serialize=True) for i_col in range(num_col): - seeds = ti.math.uvec4(seed_value + 1 + i_col, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_col) - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_row = uniform_int_distribution(result, 1, clen_value) v = vector[i_col] * weight_value while i_row < num_row: - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) out[i_row] += uniform_int_distribution(result, 1, clen_value) * v - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_row += uniform_int_distribution(result, 1, clen_value) @ti.kernel @@ -81,15 +82,15 @@ def _mv_prob_homo_gpu( i_col = i >> 5 index = i & 31 - seeds = ti.math.uvec4(seed_value + 1 + i_col, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_col) - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_row = uniform_int_distribution(result, 1, clen_value) + avg_num_uniform * index v = vector[i_col] * weight_value while i_row < num_row: - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) out[i_row] += uniform_int_distribution(result, 1, clen_value) * v - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_row += uniform_int_distribution(result, 1, clen_value) * 32 @@ -110,14 +111,14 @@ def _mv_prob_homo_outdim_parallel_cpu( # ti.loop_config(serialize=True) for i_row in range(num_row): - seeds = ti.math.uvec4(seed_value + 1 + i_row, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_row) r = 0. - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_col = uniform_int_distribution(result, 1, clen_value) while i_col < num_col: r += vector[i_col] - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_col += uniform_int_distribution(result, 1, clen_value) out[i_row] = r * weight_value @@ -141,14 +142,14 @@ def _mv_prob_homo_outdim_parallel_gpu( i_row = i >> 5 index = i & 31 - seeds = ti.math.uvec4(seed_value + 1 + i_row, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_row) r = 0. - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_col = uniform_int_distribution(result, 1, clen_value) + avg_num_uniform * index while i_col < num_col: r += vector[i_col] - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_col += uniform_int_distribution(result, 1, clen_value) * 32 out[i_row] += r * weight_value @@ -375,14 +376,14 @@ def _mv_prob_uniform_cpu( # ti.loop_config(serialize=True) for i_col in range(num_col): - seeds = ti.math.uvec4(seed_value + 1 + i_col, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_col) - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_row = uniform_int_distribution(result, 1, clen_value) while i_row < num_row: - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) out[i_row] += uniform_real_distribution(result, w_min_value, w_max_value) * vector[i_col] - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_row += uniform_int_distribution(result, 1, clen_value) @ti.kernel @@ -407,14 +408,14 @@ def _mv_prob_uniform_gpu( i_col = i >> 5 index = i & 31 - seeds = ti.math.uvec4(seed_value + 1 + i_col, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_col) - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_row = uniform_int_distribution(result, 1, clen_value) + avg_num_uniform * index while i_row < num_row: - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) out[i_row] += uniform_real_distribution(result, w_min_value, w_max_value) * vector[i_col] - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_row += uniform_int_distribution(result, 1, clen_value) * 32 @ti.kernel @@ -436,15 +437,15 @@ def _mv_prob_uniform_outdim_parallel_cpu( # ti.loop_config(serialize=True) for i_row in range(num_row): - seeds = ti.math.uvec4(seed_value + 1 + i_row, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_row) r = 0. - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_col = uniform_int_distribution(result, 1, clen_value) while i_col < num_col: - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) r += uniform_real_distribution(result, w_min_value, w_max_value) * vector[i_col] - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_col += uniform_int_distribution(result, 1, clen_value) out[i_row] = r @@ -470,15 +471,15 @@ def _mv_prob_uniform_outdim_parallel_gpu( i_row = i >> 5 index = i & 31 - seeds = ti.math.uvec4(seed_value + 1 + i_row, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_row) r = 0. - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_col = uniform_int_distribution(result, 1, clen_value) + avg_num_uniform * index while i_col < num_col: - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) r += uniform_real_distribution(result, w_min_value, w_max_value) * vector[i_col] - seeds, result = random_generator(seeds) + seeds, result = random_generator_lfsr88(seeds) i_col += uniform_int_distribution(result, 1, clen_value) * 32 out[i_row] += r @@ -701,16 +702,16 @@ def _mv_prob_normal_cpu( # ti.loop_config(serialize=True) for i_col in range(num_col): - seeds = ti.math.uvec4(seed_value + 1 + i_col, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_col) r = 0. - seeds, result1 = random_generator(seeds) + seeds, result1 = random_generator_lfsr88(seeds) i_row = uniform_int_distribution(result1, 1, clen_value) while i_row < num_row: - seeds, result1 = random_generator(seeds) - seeds, result2 = random_generator(seeds) + seeds, result1 = random_generator_lfsr88(seeds) + seeds, result2 = random_generator_lfsr88(seeds) out[i_row] += normal_distribution(result1, result2, w_mu_value, w_sigma_value) * vector[i_col] - seeds, result1 = random_generator(seeds) + seeds, result1 = random_generator_lfsr88(seeds) i_row += uniform_int_distribution(result1, 1, clen_value) @ti.kernel @@ -735,16 +736,16 @@ def _mv_prob_normal_gpu( i_col = i >> 5 index = i & 31 - seeds = ti.math.uvec4(seed_value + 1 + i_col, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_col) r = 0. - seeds, result1 = random_generator(seeds) + seeds, result1 = random_generator_lfsr88(seeds) i_row = uniform_int_distribution(result1, 1, clen_value) + avg_num_uniform * index while i_row < num_row: - seeds, result1 = random_generator(seeds) - seeds, result2 = random_generator(seeds) + seeds, result1 = random_generator_lfsr88(seeds) + seeds, result2 = random_generator_lfsr88(seeds) out[i_row] += normal_distribution(result1, result2, w_mu_value, w_sigma_value) * vector[i_col] - seeds, result1 = random_generator(seeds) + seeds, result1 = random_generator_lfsr88(seeds) i_row += uniform_int_distribution(result1, 1, clen_value) * 32 @ti.kernel @@ -766,16 +767,16 @@ def _mv_prob_normal_outdim_parallel_cpu( # ti.loop_config(serialize=True) for i_row in range(num_row): - seeds = ti.math.uvec4(seed_value + 1 + i_row, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_row) r = 0. - seeds, result1 = random_generator(seeds) + seeds, result1 = random_generator_lfsr88(seeds) i_col = uniform_int_distribution(result1, 1, clen_value) while i_col < num_col: - seeds, result1 = random_generator(seeds) - seeds, result2 = random_generator(seeds) + seeds, result1 = random_generator_lfsr88(seeds) + seeds, result2 = random_generator_lfsr88(seeds) r += normal_distribution(result1, result2, w_mu_value, w_sigma_value) * vector[i_col] - seeds, result1 = random_generator(seeds) + seeds, result1 = random_generator_lfsr88(seeds) i_col += uniform_int_distribution(result1, 1, clen_value) out[i_row] = r @@ -801,16 +802,16 @@ def _mv_prob_normal_outdim_parallel_gpu( i_row = i >> 5 index = i & 31 - seeds = ti.math.uvec4(seed_value + 1 + i_row, seed_value + 7, seed_value + 15, ti.u32(0)) + seeds = init_lfsr88_seeds(seed_value + i_row) r = 0. - seeds, result1 = random_generator(seeds) + seeds, result1 = random_generator_lfsr88(seeds) i_col = uniform_int_distribution(result1, 1, clen_value) + avg_num_uniform * index while i_col < num_col: - seeds, result1 = random_generator(seeds) - seeds, result2 = random_generator(seeds) + seeds, result1 = random_generator_lfsr88(seeds) + seeds, result2 = random_generator_lfsr88(seeds) r += normal_distribution(result1, result2, w_mu_value, w_sigma_value) * vector[i_col] - seeds, result1 = random_generator(seeds) + seeds, result1 = random_generator_lfsr88(seeds) i_col += uniform_int_distribution(result1, 1, clen_value) * 32 out[i_row] += r diff --git a/brainpy/_src/math/taichi_random.py b/brainpy/_src/math/taichi_random.py index 3208255e1..87f49cacf 100644 --- a/brainpy/_src/math/taichi_random.py +++ b/brainpy/_src/math/taichi_random.py @@ -6,7 +6,7 @@ # taichi.func random generator implementation 'taichi_lcg_rand', 'taichi_uniform_int_distribution', 'taichi_uniform_real_distribution', 'taichi_normal_distribution', - 'taichi_lfsr88', + 'taichi_lfsr88', 'init_lfsr88_seeds', 'taichi_xorwow', 'init_xorwow_seeds' ] @@ -65,6 +65,40 @@ def taichi_lfsr88(seeds): s3 = (((seeds[2] & ti.u32(4294967280)) << 17) ^ b) return ti.math.uvec4(s1, s2, s3, b), ((s1 ^ s2 ^ s3) * ti.f32(2.3283064365386963e-10)) +@ti.func +def init_lfsr88_seeds(seed: int): + """ + Get the seeds for the lfsr88 random number generator. + + Parameters: + seed (int): The seed value for the random number generator. + + Returns: + ti.math.uvec4: The seeds for the lfsr88 random number generator. + """ + return ti.math.uvec4(seed + 1, seed + 7, seed + 15, ti.u32(0)) + +@ti.func +def taichi_xorwow(seeds1, seeds2): + t = (seeds1[0]) ^ (seeds1[0] >> ti.u32(2)) + v = (seeds1[1] ^ (seeds1[1] << ti.u32(2))) ^ (t ^ (t << ti.u32(1))) + + d = seeds2[2] + ti.u32(362437) + return ti.math.uvec3(seeds1[1], seeds1[2], seeds2[0]), ti.math.uvec3(seeds2[0], v, d), ((v + d) * ti.f32(2.3283064365386963e-10)) + +@ti.func +def init_xorwow_seeds(seed: int): + """ + Get the seeds for the xorwow random number generator. + + Parameters: + seed (int): The seed value for the random number generator. + + Returns: + ti.math.uvec4: The seeds for the xorwow random number generator. + """ + return ti.math.uvec3(seed, 362436069, 521288629), ti.math.uvec3(88675123, 5783321, 6615241) + @ti.func def taichi_uniform_int_distribution(state: ti.f32, low: ti.i32, high: ti.i32): diff --git a/brainpy/_src/math/tests/taichi_random_time_test.py b/brainpy/_src/math/tests/taichi_random_time_test.py new file mode 100644 index 000000000..e45c908c7 --- /dev/null +++ b/brainpy/_src/math/tests/taichi_random_time_test.py @@ -0,0 +1,96 @@ +import jax +import jax.numpy as jnp +import time + +import brainpy.math as bm +import taichi as ti +import matplotlib.pyplot as plt + +from brainpy._src.math.taichi_random import (taichi_lcg_rand as rand, + taichi_uniform_int_distribution as randint, + taichi_uniform_real_distribution as uniform, + taichi_normal_distribution as normal, + taichi_lfsr88, + init_lfsr88_seeds, + taichi_xorwow, + init_xorwow_seeds, + taichi_lfsr88_0, + taichi_xorwow_0, + init_lfsr88_seeds_0, + init_xorwow_seeds_0,) + +bm.set_platform('gpu') + +@ti.kernel +def test_taichi_lfsr88(seed: ti.types.ndarray(ndim=1, dtype=ti.u32), + out: ti.types.ndarray(ndim=1, dtype=ti.f32)): + seeds = init_lfsr88_seeds(seed[0]) + for i in range(out.shape[0]): + seeds, result = taichi_lfsr88(seeds) + out[i] = result + +@ti.kernel +def test_taichi_xorwow(seed: ti.types.ndarray(ndim=1, dtype=ti.u32), + out: ti.types.ndarray(ndim=1, dtype=ti.f32)): + seeds1, seeds2 = init_xorwow_seeds(seed[0]) + # print(seeds1, seeds2) + for i in range(out.shape[0]): + seeds1, seeds2, result = taichi_xorwow(seeds1, seeds2) + out[i] = result + +@ti.kernel +def test_taichi_lfsr88_0(seed: ti.types.ndarray(ndim=1, dtype=ti.u32), + out: ti.types.ndarray(ndim=1, dtype=ti.f32)): + s1, s2, s3, b = init_lfsr88_seeds_0(seed[0]) + for i in range(out.shape[0]): + s1, s2, s3, b, result = taichi_lfsr88_0(s1, s2, s3, b) + out[i] = result + +@ti.kernel +def test_taichi_xorwow_0(seed: ti.types.ndarray(ndim=1, dtype=ti.u32), + out: ti.types.ndarray(ndim=1, dtype=ti.f32)): + x, y, z, w, v, d = init_xorwow_seeds_0(seed[0]) + # print(seeds1, seeds2) + for i in range(out.shape[0]): + x, y, z, w, v, d, result = taichi_xorwow_0(x, y, z, w, v, d) + out[i] = result + +n = 100000000 +seed = jnp.array([1234, ], dtype=jnp.uint32) + +prim_lfsr88 = bm.XLACustomOp(cpu_kernel=test_taichi_lfsr88, + gpu_kernel=test_taichi_lfsr88) + +prim_xorwow = bm.XLACustomOp(cpu_kernel=test_taichi_xorwow, + gpu_kernel=test_taichi_xorwow) + +prim_lfsr88_0 = bm.XLACustomOp(cpu_kernel=test_taichi_lfsr88_0, + gpu_kernel=test_taichi_lfsr88_0) + +prim_xorwow_0 = bm.XLACustomOp(cpu_kernel=test_taichi_xorwow_0, + gpu_kernel=test_taichi_xorwow_0) + +out = jax.block_until_ready(prim_lfsr88(seed, outs=[jax.ShapeDtypeStruct((n,), jnp.float32)])) +time0 = time.time() +out = jax.block_until_ready(prim_lfsr88(seed, outs=[jax.ShapeDtypeStruct((n,), jnp.float32)])) +time1 = time.time() + +out = jax.block_until_ready(prim_xorwow(seed, outs=[jax.ShapeDtypeStruct((n,), jnp.float32)])) +time2 = time.time() +out = jax.block_until_ready(prim_xorwow(seed, outs=[jax.ShapeDtypeStruct((n,), jnp.float32)])) +time3 = time.time() + +out = jax.block_until_ready(prim_lfsr88_0(seed, outs=[jax.ShapeDtypeStruct((n,), jnp.float32)])) +time4 = time.time() +out = jax.block_until_ready(prim_lfsr88_0(seed, outs=[jax.ShapeDtypeStruct((n,), jnp.float32)])) +time5 = time.time() + +out = jax.block_until_ready(prim_xorwow_0(seed, outs=[jax.ShapeDtypeStruct((n,), jnp.float32)])) +time6 = time.time() +out = jax.block_until_ready(prim_xorwow_0(seed, outs=[jax.ShapeDtypeStruct((n,), jnp.float32)])) +time7 = time.time() + +print('lfsr88: ', time1 - time0) +print('xorwow: ', time3 - time2) +print('lfsr88_0: ', time5 - time4) +print('xorwow_0: ', time7 - time6) diff --git a/brainpy/_src/math/tests/test_taichi_random.py b/brainpy/_src/math/tests/test_taichi_random.py index 7c42c95cc..fdfd3b44e 100644 --- a/brainpy/_src/math/tests/test_taichi_random.py +++ b/brainpy/_src/math/tests/test_taichi_random.py @@ -13,7 +13,10 @@ taichi_uniform_int_distribution as randint, taichi_uniform_real_distribution as uniform, taichi_normal_distribution as normal, - taichi_lfsr88) + taichi_lfsr88, + init_lfsr88_seeds, + taichi_xorwow, + init_xorwow_seeds) bm.set_platform('cpu') @@ -22,12 +25,18 @@ def test_taichi_random(): @ti.kernel def test_taichi_lfsr88(seed: ti.types.ndarray(ndim=1, dtype=ti.u32), out: ti.types.ndarray(ndim=1, dtype=ti.f32)): - s1 = seed[0] + 1 - s2 = seed[0] + 7 - s3 = seed[0] + 15 - b = ti.u32(0) + seeds = init_lfsr88_seeds(seed[0]) for i in range(out.shape[0]): - s1, s2, s3, b, result = taichi_lfsr88(s1, s2, s3, b) + seeds, result = taichi_lfsr88(seeds) + out[i] = result + + @ti.kernel + def test_taichi_xorwow(seed: ti.types.ndarray(ndim=1, dtype=ti.u32), + out: ti.types.ndarray(ndim=1, dtype=ti.f32)): + seeds1, seeds2 = init_xorwow_seeds(seed[0]) + # print(seeds1, seeds2) + for i in range(out.shape[0]): + seeds1, seeds2, result = taichi_xorwow(seeds1, seeds2) out[i] = result @ti.kernel @@ -40,45 +49,36 @@ def test_taichi_lcg_rand(seed: ti.types.ndarray(ndim=1), def test_taichi_uniform_int_distribution(seed: ti.types.ndarray(ndim=1), low_high: ti.types.ndarray(ndim=1), out: ti.types.ndarray(ndim=1)): - s1 = seed[0] + 1 - s2 = seed[0] + 7 - s3 = seed[0] + 15 - b = ti.u32(0) + seeds = init_lfsr88_seeds(seed[0]) low = low_high[0] high = low_high[1] for i in range(out.shape[0]): - s1, s2, s3, b, result = taichi_lfsr88(s1, s2, s3, b) + seeds, result = taichi_lfsr88(seeds) out[i] = randint(result, low, high) @ti.kernel def test_taichi_uniform_real_distribution(seed: ti.types.ndarray(ndim=1), low_high: ti.types.ndarray(ndim=1), out: ti.types.ndarray(ndim=1)): - s1 = seed[0] + 1 - s2 = seed[0] + 7 - s3 = seed[0] + 15 - b = ti.u32(0) + seeds = init_lfsr88_seeds(seed[0]) low = low_high[0] high = low_high[1] for i in range(out.shape[0]): - s1, s2, s3, b, result = taichi_lfsr88(s1, s2, s3, b) + seeds, result = taichi_lfsr88(seeds) out[i] = uniform(result, low, high) @ti.kernel def test_taichi_normal_distribution(seed: ti.types.ndarray(ndim=1), mu_sigma: ti.types.ndarray(ndim=1), out: ti.types.ndarray(ndim=1)): - s1 = seed[0] + 1 - s2 = seed[0] + 7 - s3 = seed[0] + 15 - b = ti.u32(0) + seeds = init_lfsr88_seeds(seed[0]) mu = mu_sigma[0] sigma = mu_sigma[1] for i in range(out.shape[0]): - s1, s2, s3, b, r1 = taichi_lfsr88(s1, s2, s3, b) - s1, s2, s3, b, r2 = taichi_lfsr88(s1, s2, s3, b) - out[i] = normal(r1, r2, mu, sigma) + seeds, result1 = taichi_lfsr88(seeds) + seeds, result2 = taichi_lfsr88(seeds) + out[i] = normal(result1, result2, mu, sigma) n = 100000 seed = jnp.array([1234, ], dtype=jnp.uint32) @@ -87,6 +87,9 @@ def test_taichi_normal_distribution(seed: ti.types.ndarray(ndim=1), prim_lfsr88 = bm.XLACustomOp(cpu_kernel=test_taichi_lfsr88, gpu_kernel=test_taichi_lfsr88) + + prim_xorwow = bm.XLACustomOp(cpu_kernel=test_taichi_xorwow, + gpu_kernel=test_taichi_xorwow) prim_lcg_rand = bm.XLACustomOp(cpu_kernel=test_taichi_lcg_rand, gpu_kernel=test_taichi_lcg_rand) @@ -105,6 +108,13 @@ def test_taichi_normal_distribution(seed: ti.types.ndarray(ndim=1), plt.title("LFSR88 random number generator") plt.savefig(file_path + "/lfsr88.png") plt.close() + + out = prim_xorwow(seed, outs=[jax.ShapeDtypeStruct((n,), jnp.float32)]) + # show the distribution of out + plt.hist(out, bins=100) + plt.title("XORWOW random number generator") + plt.savefig(file_path + "/xorwow.png") + plt.close() out = prim_lcg_rand(seed, outs=[jax.ShapeDtypeStruct((n,), jnp.float32)]) diff --git a/brainpy/math/taichi_random.py b/brainpy/math/taichi_random.py index 6bd1cb343..29d8e76bd 100644 --- a/brainpy/math/taichi_random.py +++ b/brainpy/math/taichi_random.py @@ -3,7 +3,10 @@ from brainpy._src.math.taichi_random import ( taichi_lcg_rand as taichi_lcg_rand, taichi_lfsr88 as taichi_lfsr88, + init_lfsr88_seeds as init_lfsr88_seeds, taichi_uniform_int_distribution as taichi_uniform_int_distribution, taichi_uniform_real_distribution as taichi_uniform_real_distribution, taichi_normal_distribution as taichi_normal_distribution, + taichi_xorwow as taichi_xorwow, + init_xorwow_seeds as init_xorwow_seeds, ) \ No newline at end of file