nanny_event

pyopencl/_cffi.py

@@ -182,6 +182,6 @@ def _get_insert_func(obj):
     return _insert
 
 def _find_obj(_id):
-    return _pyref[_id]
+    return _pyref.get(_id, None)
 
 _lib.set_deref(_py_deref)

src/c_wrapper/pyhelper.cpp

@@ -4,11 +4,15 @@
 
 namespace pyopencl {
 
-static std::atomic<unsigned long> pyobj_id = ATOMIC_VAR_INIT(0ul);
+static std::atomic<unsigned long> pyobj_id = ATOMIC_VAR_INIT(1ul);
 unsigned long
 next_obj_id()
 {
-    return std::atomic_fetch_add(&pyobj_id, 1ul);
+    unsigned long id;
+    do {
+        id = std::atomic_fetch_add(&pyobj_id, 1ul);
+    } while (id == 0);
+    return id;
 }
 
 static int

src/c_wrapper/wrap_cl.cpp

@@ -743,15 +743,57 @@ public:
         }
     }
     virtual void
-    wait() const
+    finished()
+    {}
+    void
+    wait()
     {
         pyopencl_call_guarded(clWaitForEvents, 1, &data());
+        finished();
+    }
+};
+
+class nanny_event : public event {
+private:
+    unsigned int m_ward;
+public:
+    nanny_event(cl_event evt, bool retain, void (*reffunc)(unsigned long)=0)
+        : event(evt, retain), m_ward(0)
+    {
+        if (reffunc) {
+            m_ward = next_obj_id();
+            reffunc(m_ward);
+        }
+    }
+    ~nanny_event()
+    {
+        if (m_ward) {
+            wait();
+        }
+    }
+    unsigned int
+    get_ward() const
+    {
+        return m_ward;
+    }
+    void
+    finished()
+    {
+        // No lock needed because multiple release is safe here.
+        unsigned long ward = m_ward;
+        m_ward = 0;
+        python_deref(ward);
     }
 };
+
 static inline event*
-new_event(cl_event evt)
+new_event(cl_event evt, void (*reffunc)(unsigned long)=0)
 {
-    return pyopencl_convert_obj(event, clReleaseEvent, evt);
+    if (reffunc) {
+        return pyopencl_convert_obj(nanny_event, clReleaseEvent, evt, reffunc);
+    } else {
+        return pyopencl_convert_obj(event, clReleaseEvent, evt);
+    }
 }
 
 // }}}