线程安全 Cython RNG nogil 并行化

Question

我试图通过扩展numpy.random RNG 在 Cython 中的 nogil prange 循环内生成随机数。我正在尝试使用this 示例，并将其与here 的直觉结合起来，为每个线程创建单独的位生成器，以避免锁定要求。然而，我显然误解了一些东西，因为bitgen_t 实体的state 和内存位置在我认为是具有独立状态的不同位生成器中是相同的。结果是，例如，当以这种方式生成 1000 个随机数时，会有 40% 的重复数字（这是有道理的，因为实际上只使用了一个位生成器）。当生成大量随机数时，它也会出现段错误，我想这是相关的。

如果有更简单的方法，请告诉我！

这是当前代码：

#!/usr/bin/env python3
#cython: language_level=3
import numpy as np
cimport numpy as np
cimport cython
from cpython.pycapsule cimport PyCapsule_IsValid, PyCapsule_GetPointer
from cython.parallel import prange, threadid
from numpy.random import PCG64, SeedSequence, Generator
from numpy.random cimport bitgen_t
from libc.stdlib cimport malloc, free
from libc.stdint cimport uintptr_t


@cython.boundscheck(False)
@cython.wraparound(False)
def uniforms(Py_ssize_t n):
    """
    Create an array of `n` uniformly distributed doubles.
    A 'real' distribution would want to process the values into
    some non-uniform distribution
    """
    cdef Py_ssize_t i, thid
    cdef const char *capsule_name = "BitGenerator"
    cdef double[:] random_values = np.ones(n, dtype='float64') * -1.

    # Create an RNG for each thread so we don't have to lock them.
    cdef Py_ssize_t num_threads = 16
    capsules = [Generator(PCG64(s)).bit_generator.capsule for s in SeedSequence(1234).spawn(num_threads)]

    cdef bitgen_t **rngs_c

    try:
        # Cast capsules into ptr array for use outside of gil.
        rngs_c = <bitgen_t**>malloc(num_threads * sizeof(bitgen_t*))
        for i, capsule in enumerate(capsules):
            if not PyCapsule_IsValid(capsule, capsule_name):
                raise ValueError("Invalid pointer to anon_func_state")
            rngs_c[i] = <bitgen_t *> PyCapsule_GetPointer(capsule, capsule_name)
            print(<uintptr_t>rngs_c[i])

        with nogil:
            for i in prange(n, schedule="static", num_threads=num_threads):
                thid = threadid()
                rng = rngs_c[thid]
                random_values[i] = rng.next_double(rng.state)
    finally:
        free(rngs_c)

    return np.asarray(random_values)

如果我按原样运行，我会得到这个输出（注意 print(<uintptr_t>rngs_c[i]) 打印从不同胶囊中检索到的 bitgen_t 的地址。

>>> x = uniforms(1000)
4931169440
4931169440
4931169440
4931169440
4931169440
4931169440
4931169440
4931169440
4931169440
4931169440
4931169440
4931169440
4931169440
4931169440
4931169440
4931169440
>>> np.unique(x).size
594

请注意，编译这个（并观察与并行化相关的问题）需要使用 openmp 编译。详情可见here。

Answer 1

我现在无法检查我的假设，但我猜你有悬空指针，因为：

[Generator(PCG64(s)).bit_generator.capsule for s in SeedSequence(1234).spawn(num_threads)]

获取胶囊不会使生成器保持活动状态，生成器在创建后会被直接删除，释放的内存会被下一个临时生成器重用，因此所有capsules 都具有相同的地址。

所以你必须将生成器保存在一个列表中，这样只要你使用capsules，它们就可以存活：

 generators = [Generator(PCG64(s)) for s in SeedSequence(1234).spawn(num_threads)]
 capsules = [g.bit_generator.capsule for g in generators]

如果我的假设是错误的，我会删除这个答案。