support stream on curand
cimport cython
from cupy.cuda cimport driver
-
+from cupy.cuda import stream as stream_module
###############################################################################
# Extern
@@ -132,24 +132,28 @@ cpdef setGeneratorOrdering(size_t generator, int order):
###############################################################################
cpdef generate(size_t generator, size_t outputPtr, size_t num):
+ setStream(generator, stream_module.get_current_stream().ptr)
status = curandGenerate(
<Generator>generator, <unsigned int*>outputPtr, num)
check_status(status)
cpdef generateLongLong(size_t generator, size_t outputPtr, size_t num):
+ setStream(generator, stream_module.get_current_stream().ptr)
status = curandGenerateLongLong(
<Generator>generator, <unsigned long long*>outputPtr, num)
check_status(status)
cpdef generateUniform(size_t generator, size_t outputPtr, size_t num):
+ setStream(generator, stream_module.get_current_stream().ptr)
status = curandGenerateUniform(
<Generator>generator, <float*>outputPtr, num)
check_status(status)
cpdef generateUniformDouble(size_t generator, size_t outputPtr, size_t num):
+ setStream(generator, stream_module.get_current_stream().ptr)
status = curandGenerateUniformDouble(
<Generator>generator, <double*>outputPtr, num)
check_status(status)
@@ -161,6 +165,7 @@ cpdef generateNormal(size_t generator, size_t outputPtr, size_t n,
msg = ('curandGenerateNormal can only generate even number of '
'random variables simultaneously. See issue #390 for detail.')
raise ValueError(msg)
+ setStream(generator, stream_module.get_current_stream().ptr)
status = curandGenerateNormal(
<Generator>generator, <float*>outputPtr, n, mean, stddev)
check_status(status)
@@ -172,6 +177,7 @@ cpdef generateNormalDouble(size_t generator, size_t outputPtr, size_t n,
msg = ('curandGenerateNormalDouble can only generate even number of '
'random variables simultaneously. See issue #390 for detail.')
raise ValueError(msg)
+ setStream(generator, stream_module.get_current_stream().ptr)
status = curandGenerateNormalDouble(
<Generator>generator, <double*>outputPtr, n, mean, stddev)
check_status(status)
@@ -183,6 +189,7 @@ def generateLogNormal(size_t generator, size_t outputPtr, size_t n,
msg = ('curandGenerateLogNormal can only generate even number of '
'random variables simultaneously. See issue #390 for detail.')
raise ValueError(msg)
+ setStream(generator, stream_module.get_current_stream().ptr)
status = curandGenerateLogNormal(
<Generator>generator, <float*>outputPtr, n, mean, stddev)
check_status(status)
@@ -195,6 +202,7 @@ cpdef generateLogNormalDouble(size_t generator, size_t outputPtr, size_t n,
'of random variables simultaneously. See issue #390 for '
'detail.')
raise ValueError(msg)
+ setStream(generator, stream_module.get_current_stream().ptr)
status = curandGenerateLogNormalDouble(
<Generator>generator, <double*>outputPtr, n, mean, stddev)
check_status(status)
@@ -202,6 +210,7 @@ cpdef generateLogNormalDouble(size_t generator, size_t outputPtr, size_t n,
cpdef generatePoisson(size_t generator, size_t outputPtr, size_t n,
double lam):
+ setStream(generator, stream_module.get_current_stream().ptr)
status = curandGeneratePoisson(
<Generator>generator, <unsigned int*>outputPtr, n, lam)
check_status(status)
