Bindings for sync_device() and sync_all_devices()

Author: merlinND

docs/reference.rst

@@ -620,6 +620,8 @@ Low-level bits
 .. autofunction:: thread_count
 .. autofunction:: set_thread_count
 .. autofunction:: sync_thread
+.. autofunction:: sync_device
+.. autofunction:: sync_all_devices
 .. autofunction:: flush_kernel_cache
 .. autofunction:: flush_malloc_cache
 .. autofunction:: expand_threshold

src/python/docstr.rst

@@ -7130,6 +7130,16 @@
     then you have found a bug. Please report it on the project's
     `GitHub issue tracker <https://github.com/mitsuba-renderer/drjit>`__.
 
+
+.. topic:: sync_device
+
+    Wait for all computation on the current device to finish.
+
+.. topic:: sync_all_devices
+
+    Wait for all computation on *all devices* to finish.
+
+
 .. topic:: flush_malloc_cache
 
     Free the memory allocation cache maintained by Dr.Jit.

src/python/main.cpp

@@ -171,6 +171,8 @@ NB_MODULE(_drjit_ext, m_) {
     m.def("has_backend", &jit_has_backend, doc_has_backend);
 
     m.def("sync_thread", &jit_sync_thread, doc_sync_thread)
+     .def("sync_device", &jit_sync_device, doc_sync_device)
+     .def("sync_all_devices", &jit_sync_all_devices, doc_sync_all_devices)
      .def("flush_kernel_cache", &jit_flush_kernel_cache, doc_flush_kernel_cache)
      .def("flush_malloc_cache", &jit_flush_malloc_cache, doc_flush_malloc_cache)
      .def("malloc_clear_statistics", &jit_malloc_clear_statistics)