Implement stream priority feature

inducer · Nov 23, 2021 · c08a582 · c08a582
1 parent 351bc65
commit c08a582
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 5 deletions.
diff --git a/doc/driver.rst b/doc/driver.rst
@@ -331,7 +331,7 @@ Constants
         CUDA 6.0 and above.
 
         .. versionadded:: 2014.1
-        
+
     .. attribute :: HOST_NATIVE_ATOMIC_SUPPORTED
         SINGLE_TO_DOUBLE_PRECISION_PERF_RATIO
         PAGEABLE_MEMORY_ACCESS
@@ -813,7 +813,7 @@ Devices and Contexts
 Concurrency and Streams
 -----------------------
 
-.. class:: Stream(flags=0)
+.. class:: Stream(flags=0, priority=0)
 
     A handle for a queue of operations that will be carried out in order.
 

diff --git a/src/cpp/cuda.hpp b/src/cpp/cuda.hpp
@@ -997,8 +997,8 @@ namespace pycuda
       CUstream m_stream;
 
     public:
-      stream(unsigned int flags=0)
-      { CUDAPP_CALL_GUARDED(cuStreamCreate, (&m_stream, flags)); }
+      stream(unsigned int flags=0, int priority=0)
+      { CUDAPP_CALL_GUARDED(cuStreamCreateWithPriority, (&m_stream, flags, priority)); }
 
       ~stream()
       {

diff --git a/src/wrapper/wrap_cudadrv.cpp b/src/wrapper/wrap_cudadrv.cpp
@@ -1199,7 +1199,7 @@ BOOST_PYTHON_MODULE(_driver)
   {
     typedef stream cl;
     py::class_<cl, boost::noncopyable, shared_ptr<cl> >
-      ("Stream", py::init<unsigned int>(py::arg("flags")=0))
+      ("Stream", py::init<unsigned int, int>(py::arg("flags")=0, py::arg("priority")=0))
       .DEF_SIMPLE_METHOD(synchronize)
       .DEF_SIMPLE_METHOD(is_done)
 #if CUDAPP_CUDA_VERSION >= 3020