diff --git a/include/drjit-core/jit.h b/include/drjit-core/jit.h index f504aee2..e7514444 100644 --- a/include/drjit-core/jit.h +++ b/include/drjit-core/jit.h @@ -494,9 +494,15 @@ extern JIT_EXPORT void jit_registry_remove(const void *ptr); extern JIT_EXPORT uint32_t jit_registry_id(const void *ptr); /// Return the largest instance ID for the given domain +/// If the \c domain is a nullptr, it returns the number of active entries in +/// all domains for the given backend extern JIT_EXPORT uint32_t jit_registry_id_bound(JitBackend backend, const char *domain); +/// Fills the \c dest pointer array with all pointers registered in the registry +/// \c dest has to point to an array with \c jit_registry_id_bound(backend, nullptr) entries +extern JIT_EXPORT void jit_registry_get_pointers(JitBackend backend, void **dest); + /// Return the pointer value associated with a given instance ID extern JIT_EXPORT void *jit_registry_ptr(JitBackend backend, const char *domain, uint32_t id); diff --git a/src/api.cpp b/src/api.cpp index 553e84cc..d38c71b7 100644 --- a/src/api.cpp +++ b/src/api.cpp @@ -983,6 +983,11 @@ uint32_t jit_registry_id_bound(JitBackend backend, const char *domain) { return jitc_registry_id_bound(backend, domain); } +void jit_registry_get_pointers(JitBackend backend, void **dest) { + lock_guard guard(state.lock); + return jitc_registry_get_pointers(backend, dest); +} + void *jit_registry_ptr(JitBackend backend, const char *domain, uint32_t id) { lock_guard guard(state.lock); return jitc_registry_ptr(backend, domain, id); diff --git a/src/registry.cpp b/src/registry.cpp index 7479a09f..540d6f83 100644 --- a/src/registry.cpp +++ b/src/registry.cpp @@ -39,6 +39,7 @@ struct Ptr { // Per-domain information: a forward map (index-> pointer) and a list of unused entries struct Domain { const char *name; + JitBackend backend; uint32_t id_bound; std::vector fwd_map; std::priority_queue, std::greater> @@ -76,6 +77,7 @@ uint32_t jitc_registry_put(JitBackend backend, const char *domain_name, void *pt r.domains.emplace_back(); Domain &domain = r.domains.back(); domain.name = domain_name; + domain.backend = backend; domain.id_bound = 0; } @@ -149,6 +151,16 @@ uint32_t jitc_registry_id(const void *ptr) { uint32_t jitc_registry_id_bound(JitBackend backend, const char *domain) { Registry &r = registry; + if (!domain) { + uint32_t n = 0; + for (Domain &domain : r.domains) { + if (domain.backend == backend) + for (auto ptr : domain.fwd_map) + if (ptr.active) + n++; + } + return n; + } auto it = r.domain_ids.find(DomainKey{ backend, domain }); if (it == r.domain_ids.end()) return 0; @@ -156,6 +168,21 @@ uint32_t jitc_registry_id_bound(JitBackend backend, const char *domain) { return r.domains[it->second].id_bound; } +void jitc_registry_get_pointers(JitBackend backend, void **dest) { + const Registry &r = registry; + + uint32_t n = 0; + for (const Domain &domain : r.domains) { + if (domain.backend == backend) + for (auto ptr : domain.fwd_map) { + if (ptr.active) { + dest[n] = ptr.ptr; + n++; + } + } + } +} + void *jitc_registry_ptr(JitBackend backend, const char *domain_name, uint32_t id) { if (id == 0) return nullptr; diff --git a/src/registry.h b/src/registry.h index f2be89c5..befb2025 100644 --- a/src/registry.h +++ b/src/registry.h @@ -22,8 +22,14 @@ extern void jitc_registry_remove(const void *ptr); extern uint32_t jitc_registry_id(const void *ptr); /// Return the largest instance ID for the given domain +/// If the \c domain is a nullptr, it returns the number of active entries in +/// all domains for the given backend extern uint32_t jitc_registry_id_bound(JitBackend backend, const char *domain); +/// Fills the \c dest pointer array with all pointers registered in the registry +/// \c dest has to point to an array with \c jit_registry_id_bound(backend, nullptr) entries +void extern jitc_registry_get_pointers(JitBackend backend, void **dest); + /// Return the pointer value associated with a given instance ID extern void *jitc_registry_ptr(JitBackend backend, const char *domain, uint32_t id);