Allow tying an svm_allocation to a queue

Author: inducer

pyopencl/__init__.py

@@ -1146,12 +1146,13 @@ def memory_map_exit(self, exc_type, exc_val, exc_tb):
     if get_cl_header_version() >= (2, 0):
         svmallocation_old_init = SVMAllocation.__init__
 
-    def svmallocation_init(self, ctx, size, alignment, flags, _interface=None):
+    def svmallocation_init(self, ctx, size, alignment, flags, _interface=None,
+            queue=None):
         """
         :arg ctx: a :class:`Context`
         :arg flags: some of :class:`svm_mem_flags`.
         """
-        svmallocation_old_init(self, ctx, size, alignment, flags)
+        svmallocation_old_init(self, ctx, size, alignment, flags, queue)
 
         # mem_flags.READ_ONLY applies to kernels, not the host
         read_write = True
@@ -1996,7 +1997,7 @@ def enqueue_svm_migratemem(queue, svms, flags, wait_for=None):
             wait_for)
 
 
-def svm_empty(ctx, flags, shape, dtype, order="C", alignment=None):
+def svm_empty(ctx, flags, shape, dtype, order="C", alignment=None, queue=None):
     """Allocate an empty :class:`numpy.ndarray` of the given *shape*, *dtype*
     and *order*. (See :func:`numpy.empty` for the meaning of these arguments.)
     The array will be allocated in shared virtual memory belonging
@@ -2014,6 +2015,10 @@ def svm_empty(ctx, flags, shape, dtype, order="C", alignment=None):
     will likely want to wrap the returned array in an :class:`SVM` tag.
 
     .. versionadded:: 2016.2
+
+    .. versionchanged:: 2022.2
+
+        *queue* argument added.
     """
 
     dtype = np.dtype(dtype)
@@ -2060,7 +2065,8 @@ def svm_empty(ctx, flags, shape, dtype, order="C", alignment=None):
     if alignment is None:
         alignment = itemsize
 
-    svm_alloc = SVMAllocation(ctx, nbytes, alignment, flags, _interface=interface)
+    svm_alloc = SVMAllocation(ctx, nbytes, alignment, flags, _interface=interface,
+            queue=queue)
     return np.asarray(svm_alloc)
 
 

src/wrap_cl.hpp

@@ -3562,16 +3562,28 @@ namespace pyopencl
   };
 
 
-  class svm_allocation : noncopyable
+  class svm_allocation
   {
     private:
       std::shared_ptr<context> m_context;
       void *m_allocation;
+      command_queue_ref m_queue;
+      // FIXME Should keep a list of events so that we can wait for users
+      // to finish in the case of out-of-order queues.
 
     public:
-      svm_allocation(std::shared_ptr<context> const &ctx, size_t size, cl_uint alignment, cl_svm_mem_flags flags)
+      svm_allocation(std::shared_ptr<context> const &ctx, size_t size, cl_uint alignment,
+          cl_svm_mem_flags flags, const command_queue *queue = nullptr)
         : m_context(ctx)
       {
+        if (queue)
+        {
+          m_queue.set(queue->data());
+          if (!is_queue_in_order(m_queue.data()))
+            throw error("SVMAllocation.__init__", CL_INVALID_VALUE,
+                "supplying an out-of-order queue to SVMAllocation is invalid");
+        }
+
         PYOPENCL_PRINT_CALL_TRACE("clSVMalloc");
         m_allocation = clSVMAlloc(
             ctx->data(),
@@ -3581,6 +3593,9 @@ namespace pyopencl
           throw pyopencl::error("clSVMAlloc", CL_OUT_OF_RESOURCES);
       }
 
+      svm_allocation(const svm_allocation &) = delete;
+      svm_allocation &operator=(const svm_allocation &) = delete;
+
       ~svm_allocation()
       {
         if (m_allocation)
@@ -3593,8 +3608,20 @@ namespace pyopencl
           throw error("SVMAllocation.release", CL_INVALID_VALUE,
               "trying to double-unref svm allocation");
 
-        clSVMFree(m_context->data(), m_allocation);
-        m_allocation = nullptr;
+        if (m_queue.is_valid())
+        {
+          PYOPENCL_CALL_GUARDED_CLEANUP(clEnqueueSVMFree, (
+                m_queue.data(), 1, &m_allocation,
+                nullptr, nullptr,
+                0, nullptr, nullptr));
+          m_queue.reset();
+        }
+        else
+        {
+          PYOPENCL_PRINT_CALL_TRACE("clSVMFree");
+          clSVMFree(m_context->data(), m_allocation);
+          m_allocation = nullptr;
+        }
       }
 
       void enqueue_release(command_queue &queue, py::object py_wait_for)
@@ -3634,6 +3661,16 @@ namespace pyopencl
       {
         return m_allocation != other.m_allocation;
       }
+
+      void bind_to_queue(command_queue const &queue)
+      {
+        m_queue.set(queue.data());
+      }
+
+      void unbind_from_queue()
+      {
+        m_queue.reset();
+      }
   };
 
 

src/wrap_cl_part_2.cpp

@@ -305,7 +305,13 @@ void pyopencl_expose_part_2(py::module &m)
   {
     typedef svm_allocation cls;
     py::class_<cls>(m, "SVMAllocation", py::dynamic_attr())
-      .def(py::init<std::shared_ptr<context>, size_t, cl_uint, cl_svm_mem_flags>())
+      .def(py::init<std::shared_ptr<context>, size_t, cl_uint, cl_svm_mem_flags, const command_queue *>(),
+          py::arg("context"),
+          py::arg("size"),
+          py::arg("alignment"),
+          py::arg("flags"),
+          py::arg("queue").none(true)=py::none()
+          )
       .DEF_SIMPLE_METHOD(release)
       .def("enqueue_release", &cls::enqueue_release,
           ":returns: a :class:`pyopencl.Event`\n\n"
@@ -314,6 +320,8 @@ void pyopencl_expose_part_2(py::module &m)
       .def(py::self == py::self)
       .def(py::self != py::self)
       .def("__hash__", &cls::ptr_as_int)
+      .DEF_SIMPLE_METHOD(bind_to_queue)
+      .DEF_SIMPLE_METHOD(unbind_from_queue)
       ;
   }