rayd 0.2.0__tar.gz → 0.2.2__tar.gz

{rayd-0.2.0 → rayd-0.2.2}/CMakeLists.txt

@@ -143,6 +143,8 @@ set(RAYD_CORE_SOURCE_FILES
     ${RAYD_SOURCE_DIR}/scene/scene.cpp
 
     ${RAYD_SOURCE_DIR}/optix.cpp
+    ${RAYD_SOURCE_DIR}/native_launch_audit.h
+    ${RAYD_SOURCE_DIR}/native_launch_audit.cpp
 
     ${RAYD_INCLUDE_DIR}/slang/interop.h
     ${RAYD_INCLUDE_DIR}/slang/interop_types.h

{rayd-0.2.0 → rayd-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: rayd
-Version: 0.2.0
+Version: 0.2.2
 Summary: RayD: minimalist differentiable ray tracing package wrapping Dr.Jit and OptiX.
 Author-Email: Xingyu Chen <xic063@ucsd.edu>
 License: BSD-3-Clause

{rayd-0.2.0 → rayd-0.2.2}/include/rayd/scene/scene_edge.h

@@ -64,6 +64,7 @@ public:
 
 private:
     void build_bvh(const SecondaryEdgeInfo &edge_info);
+    void set_all_active_state();
     void update_active_counts_from_mask(const MaskDetached &mask);
     IntDetached refit_leaf_nodes_from_primitive_indices(const SecondaryEdgeInfo &edge_info,
                                                         const IntDetached &primitive_indices);

{rayd-0.2.0 → rayd-0.2.2}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "scikit_build_core.build"
 
 [project]
 name = "rayd"
-version = "0.2.0"
+version = "0.2.2"
 description = "RayD: minimalist differentiable ray tracing package wrapping Dr.Jit and OptiX."
 readme = "README.md"
 authors = [{ name = "Xingyu Chen", email = "xic063@ucsd.edu" }]

{rayd-0.2.0 → rayd-0.2.2}/src/multipath/reflection_dedup.cu

@@ -8,6 +8,7 @@
 #include <string>
 
 #include <rayd/rayd.h>
+#include "../native_launch_audit.h"
 
 namespace rayd {
 
@@ -366,17 +367,24 @@ int reflection_dedup_gpu(
                                     sizeof(int) * static_cast<size_t>(n_rays),
                                     stream),
                     "reflection_dedup_gpu(): failed to clear discovery counts");
+    audit_cuda_memset_async();
     check_cuda_call(cudaMemsetAsync(out_representative_ray_index,
                                     0xFF,
                                     sizeof(int) * static_cast<size_t>(n_rays),
                                     stream),
                     "reflection_dedup_gpu(): failed to clear representative indices");
+    audit_cuda_memset_async();
     check_cuda_call(cudaMemsetAsync(unique_count_device.get(),
                                     0,
                                     sizeof(int),
                                     stream),
                     "reflection_dedup_gpu(): failed to clear unique counter");
+    audit_cuda_memset_async();
 
+    audit_cuda_kernel_launch("reflection_dedup_build_keys_kernel",
+                             static_cast<uint32_t>(block_count), 1, 1,
+                             block_size, 1, 1,
+                             static_cast<uint64_t>(n_rays));
     reflection_dedup_build_keys_kernel<<<block_count, block_size, 0, stream>>>(
         n_rays,
         max_bounces,
@@ -392,6 +400,7 @@ int reflection_dedup_gpu(
     check_cuda_last_error("reflection_dedup_gpu(): failed to launch build-keys kernel");
 
     size_t sort_temp_size = 0;
+    audit_cub_sort();
     check_cuda_call(cub::DeviceRadixSort::SortPairs(nullptr,
                                                     sort_temp_size,
                                                     keys_in.get(),
@@ -404,6 +413,7 @@ int reflection_dedup_gpu(
                                                     stream),
                     "reflection_dedup_gpu(): failed to size first radix sort");
     CudaBuffer<char> sort_temp(std::max<size_t>(sort_temp_size, 1));
+    audit_cub_sort();
     check_cuda_call(cub::DeviceRadixSort::SortPairs(sort_temp.get(),
                                                     sort_temp_size,
                                                     keys_in.get(),
@@ -416,11 +426,16 @@ int reflection_dedup_gpu(
                                                     stream),
                     "reflection_dedup_gpu(): failed to run first radix sort");
 
+    audit_cuda_kernel_launch("reflection_dedup_mark_boundaries_kernel",
+                             static_cast<uint32_t>(block_count), 1, 1,
+                             block_size, 1, 1,
+                             static_cast<uint64_t>(n_rays));
     reflection_dedup_mark_boundaries_kernel<<<block_count, block_size, 0, stream>>>(
         n_rays, keys_out.get(), boundary_flags.get());
     check_cuda_last_error("reflection_dedup_gpu(): failed to launch first boundary kernel");
 
     size_t scan_temp_size = 0;
+    audit_cub_scan();
     check_cuda_call(cub::DeviceScan::InclusiveSum(nullptr,
                                                   scan_temp_size,
                                                   boundary_flags.get(),
@@ -429,6 +444,7 @@ int reflection_dedup_gpu(
                                                   stream),
                     "reflection_dedup_gpu(): failed to size first scan");
     CudaBuffer<char> scan_temp(std::max<size_t>(scan_temp_size, 1));
+    audit_cub_scan();
     check_cuda_call(cub::DeviceScan::InclusiveSum(scan_temp.get(),
                                                   scan_temp_size,
                                                   boundary_flags.get(),
@@ -436,10 +452,18 @@ int reflection_dedup_gpu(
                                                   n_rays,
                                                   stream),
                     "reflection_dedup_gpu(): failed to run first scan");
+    audit_cuda_kernel_launch("reflection_dedup_zero_base_ids_kernel",
+                             static_cast<uint32_t>(block_count), 1, 1,
+                             block_size, 1, 1,
+                             static_cast<uint64_t>(n_rays));
     reflection_dedup_zero_base_ids_kernel<<<block_count, block_size, 0, stream>>>(
         n_rays, keys_out.get(), hash_group_ids.get());
     check_cuda_last_error("reflection_dedup_gpu(): failed to launch first id-fix kernel");
 
+    audit_cuda_kernel_launch("reflection_dedup_sub_cluster_kernel",
+                             static_cast<uint32_t>(block_count), 1, 1,
+                             block_size, 1, 1,
+                             static_cast<uint64_t>(n_rays));
     reflection_dedup_sub_cluster_kernel<<<block_count, block_size, 0, stream>>>(
         n_rays,
         max_bounces,
@@ -456,6 +480,7 @@ int reflection_dedup_gpu(
     check_cuda_last_error("reflection_dedup_gpu(): failed to launch sub-cluster kernel");
 
     size_t cluster_sort_temp_size = 0;
+    audit_cub_sort();
     check_cuda_call(cub::DeviceRadixSort::SortPairs(nullptr,
                                                     cluster_sort_temp_size,
                                                     cluster_keys_in.get(),
@@ -468,6 +493,7 @@ int reflection_dedup_gpu(
                                                     stream),
                     "reflection_dedup_gpu(): failed to size second radix sort");
     CudaBuffer<char> cluster_sort_temp(std::max<size_t>(cluster_sort_temp_size, 1));
+    audit_cub_sort();
     check_cuda_call(cub::DeviceRadixSort::SortPairs(cluster_sort_temp.get(),
                                                     cluster_sort_temp_size,
                                                     cluster_keys_in.get(),
@@ -480,10 +506,15 @@ int reflection_dedup_gpu(
                                                     stream),
                     "reflection_dedup_gpu(): failed to run second radix sort");
 
+    audit_cuda_kernel_launch("reflection_dedup_mark_boundaries_kernel",
+                             static_cast<uint32_t>(block_count), 1, 1,
+                             block_size, 1, 1,
+                             static_cast<uint64_t>(n_rays));
     reflection_dedup_mark_boundaries_kernel<<<block_count, block_size, 0, stream>>>(
         n_rays, cluster_keys_out.get(), boundary_flags.get());
     check_cuda_last_error("reflection_dedup_gpu(): failed to launch second boundary kernel");
 
+    audit_cub_scan();
     check_cuda_call(cub::DeviceScan::InclusiveSum(scan_temp.get(),
                                                   scan_temp_size,
                                                   boundary_flags.get(),
@@ -491,10 +522,18 @@ int reflection_dedup_gpu(
                                                   n_rays,
                                                   stream),
                     "reflection_dedup_gpu(): failed to run second scan");
+    audit_cuda_kernel_launch("reflection_dedup_zero_base_ids_kernel",
+                             static_cast<uint32_t>(block_count), 1, 1,
+                             block_size, 1, 1,
+                             static_cast<uint64_t>(n_rays));
     reflection_dedup_zero_base_ids_kernel<<<block_count, block_size, 0, stream>>>(
         n_rays, cluster_keys_out.get(), unique_path_ids.get());
     check_cuda_last_error("reflection_dedup_gpu(): failed to launch second id-fix kernel");
 
+    audit_cuda_kernel_launch("reflection_dedup_compact_kernel",
+                             static_cast<uint32_t>(block_count), 1, 1,
+                             block_size, 1, 1,
+                             static_cast<uint64_t>(n_rays));
     reflection_dedup_compact_kernel<<<block_count, block_size, 0, stream>>>(
         n_rays,
         max_bounces,
@@ -537,12 +576,14 @@ int reflection_dedup_gpu(
     check_cuda_last_error("reflection_dedup_gpu(): failed to launch compact kernel");
 
     int unique_count = 0;
+    audit_cuda_memcpy_async();
     check_cuda_call(cudaMemcpyAsync(&unique_count,
                                     unique_count_device.get(),
                                     sizeof(int),
                                     cudaMemcpyDeviceToHost,
                                     stream),
                     "reflection_dedup_gpu(): failed to copy unique count");
+    audit_cuda_stream_synchronize();
     check_cuda_call(cudaStreamSynchronize(stream),
                     "reflection_dedup_gpu(): failed to finish dedup stream");
     return unique_count;

{rayd-0.2.0 → rayd-0.2.2}/src/multipath/reflection_trace_host.cpp

@@ -7,6 +7,7 @@
 #include <vector>
 
 #include "reflection_trace_ptx.h"
+#include "../native_launch_audit.h"
 
 namespace rayd {
 
@@ -161,6 +162,7 @@ void ReflectionTracePipeline::build(OptixDeviceContext context,
     check_optix(optixSbtRecordPackHeader(pg_raygen_, &raygen_record),
                 "optixSbtRecordPackHeader(raygen)");
     sbt_raygen_record_ = jit_malloc(AllocType::Device, sizeof(EmptySbtRecord));
+    audit_jit_memcpy();
     jit_memcpy(JitBackend::CUDA,
                sbt_raygen_record_,
                &raygen_record,
@@ -170,6 +172,7 @@ void ReflectionTracePipeline::build(OptixDeviceContext context,
     check_optix(optixSbtRecordPackHeader(pg_miss_, &miss_record),
                 "optixSbtRecordPackHeader(miss)");
     sbt_miss_record_ = jit_malloc(AllocType::Device, sizeof(EmptySbtRecord));
+    audit_jit_memcpy();
     jit_memcpy(JitBackend::CUDA,
                sbt_miss_record_,
                &miss_record,
@@ -182,6 +185,7 @@ void ReflectionTracePipeline::build(OptixDeviceContext context,
     }
     sbt_hitgroup_records_ = jit_malloc(AllocType::Device,
                                        sizeof(EmptySbtRecord) * hitgroup_records.size());
+    audit_jit_memcpy();
     jit_memcpy(JitBackend::CUDA,
                sbt_hitgroup_records_,
                hitgroup_records.data(),
@@ -195,6 +199,7 @@ void ReflectionTracePipeline::build(OptixDeviceContext context,
 void ReflectionTracePipeline::launch(const ReflectionTraceParams &params) const {
     require(ready_, "ReflectionTracePipeline::launch(): pipeline is not ready.");
 
+    audit_jit_memcpy_async();
     jit_memcpy_async(JitBackend::CUDA,
                      params_buffer_,
                      &params,
@@ -209,6 +214,7 @@ void ReflectionTracePipeline::launch(const ReflectionTraceParams &params) const
     sbt.hitgroupRecordStrideInBytes = sizeof(EmptySbtRecord);
     sbt.hitgroupRecordCount = static_cast<unsigned int>(hitgroup_record_count_);
 
+    audit_optix_launch();
     check_optix(optixLaunch(pipeline_,
                             jit_cuda_stream(),
                             reinterpret_cast<CUdeviceptr>(params_buffer_),

rayd-0.2.2/src/native_launch_audit.cpp

@@ -0,0 +1,170 @@
+#include "native_launch_audit.h"
+
+#include <algorithm>
+#include <mutex>
+
+namespace rayd {
+
+namespace {
+
+std::mutex &audit_mutex() {
+    static std::mutex mutex;
+    return mutex;
+}
+
+NativeLaunchAuditSnapshot &audit_snapshot_storage() {
+    static NativeLaunchAuditSnapshot snapshot;
+    return snapshot;
+}
+
+thread_local NativeLaunchStage current_stage = NativeLaunchStage::Unknown;
+
+NativeLaunchStageStats &stage_stats(NativeLaunchAuditSnapshot &snapshot,
+                                    NativeLaunchStage stage) {
+    switch (stage) {
+    case NativeLaunchStage::Build:
+        return snapshot.build;
+    case NativeLaunchStage::Sync:
+        return snapshot.sync;
+    case NativeLaunchStage::TraceReflections:
+        return snapshot.trace_reflections;
+    case NativeLaunchStage::Unknown:
+    default:
+        return snapshot.unknown;
+    }
+}
+
+void clear_stage_stats(NativeLaunchStageStats &stats) {
+    stats = NativeLaunchStageStats();
+}
+
+void update_kernel_stat(NativeLaunchStageStats &stats,
+                        const char *label,
+                        uint64_t threads,
+                        uint64_t items) {
+    const std::string key = label != nullptr ? std::string(label) : std::string("unnamed");
+    auto it = std::find_if(stats.kernels.begin(),
+                           stats.kernels.end(),
+                           [&key](const NativeKernelLaunchStat &entry) {
+                               return entry.label == key;
+                           });
+    if (it == stats.kernels.end()) {
+        stats.kernels.push_back(NativeKernelLaunchStat { key, 1, threads, threads, items, items });
+        return;
+    }
+
+    it->launches += 1;
+    it->total_threads += threads;
+    it->max_threads = std::max(it->max_threads, threads);
+    it->total_items += items;
+    it->max_items = std::max(it->max_items, items);
+}
+
+template <typename Member>
+void increment_counter(Member NativeLaunchStageStats::*member) {
+    std::lock_guard<std::mutex> guard(audit_mutex());
+    NativeLaunchStageStats &stats = stage_stats(audit_snapshot_storage(), current_stage);
+    stats.*member += 1;
+}
+
+} // namespace
+
+ScopedNativeLaunchStage::ScopedNativeLaunchStage(NativeLaunchStage stage)
+    : previous_(current_stage) {
+    current_stage = stage;
+}
+
+ScopedNativeLaunchStage::~ScopedNativeLaunchStage() {
+    current_stage = previous_;
+}
+
+void native_launch_audit_clear() {
+    std::lock_guard<std::mutex> guard(audit_mutex());
+    NativeLaunchAuditSnapshot &snapshot = audit_snapshot_storage();
+    clear_stage_stats(snapshot.unknown);
+    clear_stage_stats(snapshot.build);
+    clear_stage_stats(snapshot.sync);
+    clear_stage_stats(snapshot.trace_reflections);
+}
+
+NativeLaunchAuditSnapshot native_launch_audit_snapshot() {
+    std::lock_guard<std::mutex> guard(audit_mutex());
+    return audit_snapshot_storage();
+}
+
+void audit_cuda_kernel_launch(const char *label,
+                              uint32_t grid_x,
+                              uint32_t grid_y,
+                              uint32_t grid_z,
+                              uint32_t block_x,
+                              uint32_t block_y,
+                              uint32_t block_z,
+                              uint64_t items) {
+    const uint64_t threads =
+        static_cast<uint64_t>(grid_x) * static_cast<uint64_t>(grid_y) * static_cast<uint64_t>(grid_z) *
+        static_cast<uint64_t>(block_x) * static_cast<uint64_t>(block_y) * static_cast<uint64_t>(block_z);
+
+    std::lock_guard<std::mutex> guard(audit_mutex());
+    NativeLaunchStageStats &stats = stage_stats(audit_snapshot_storage(), current_stage);
+    stats.cuda_kernel_launches += 1;
+    stats.cuda_kernel_total_threads += threads;
+    update_kernel_stat(stats, label, threads, items);
+}
+
+void audit_cuda_memcpy() {
+    increment_counter(&NativeLaunchStageStats::cuda_memcpy);
+}
+
+void audit_cuda_memcpy_async() {
+    increment_counter(&NativeLaunchStageStats::cuda_memcpy_async);
+}
+
+void audit_cuda_memset_async() {
+    increment_counter(&NativeLaunchStageStats::cuda_memset_async);
+}
+
+void audit_cuda_stream_synchronize() {
+    increment_counter(&NativeLaunchStageStats::cuda_stream_synchronize);
+}
+
+void audit_cuda_event_record() {
+    increment_counter(&NativeLaunchStageStats::cuda_event_record);
+}
+
+void audit_cuda_stream_wait_event() {
+    increment_counter(&NativeLaunchStageStats::cuda_stream_wait_event);
+}
+
+void audit_cub_reduce() {
+    increment_counter(&NativeLaunchStageStats::cub_reduce);
+}
+
+void audit_cub_sort() {
+    increment_counter(&NativeLaunchStageStats::cub_sort);
+}
+
+void audit_cub_scan() {
+    increment_counter(&NativeLaunchStageStats::cub_scan);
+}
+
+void audit_jit_memcpy() {
+    increment_counter(&NativeLaunchStageStats::jit_memcpy);
+}
+
+void audit_jit_memcpy_async() {
+    increment_counter(&NativeLaunchStageStats::jit_memcpy_async);
+}
+
+void audit_optix_accel_build() {
+    increment_counter(&NativeLaunchStageStats::optix_accel_build);
+}
+
+void audit_optix_accel_compact() {
+    increment_counter(&NativeLaunchStageStats::optix_accel_compact);
+}
+
+void audit_optix_launch() {
+    increment_counter(&NativeLaunchStageStats::optix_launch);
+}
+
+} // namespace rayd

rayd-0.2.2/src/native_launch_audit.h

@@ -0,0 +1,90 @@
+#pragma once
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+namespace rayd {
+
+enum class NativeLaunchStage {
+    Unknown = 0,
+    Build,
+    Sync,
+    TraceReflections
+};
+
+struct NativeKernelLaunchStat {
+    std::string label;
+    uint64_t launches = 0;
+    uint64_t total_threads = 0;
+    uint64_t max_threads = 0;
+    uint64_t total_items = 0;
+    uint64_t max_items = 0;
+};
+
+struct NativeLaunchStageStats {
+    uint64_t cuda_kernel_launches = 0;
+    uint64_t cuda_kernel_total_threads = 0;
+    uint64_t cuda_memcpy = 0;
+    uint64_t cuda_memcpy_async = 0;
+    uint64_t cuda_memset_async = 0;
+    uint64_t cuda_stream_synchronize = 0;
+    uint64_t cuda_event_record = 0;
+    uint64_t cuda_stream_wait_event = 0;
+    uint64_t cub_reduce = 0;
+    uint64_t cub_sort = 0;
+    uint64_t cub_scan = 0;
+    uint64_t jit_memcpy = 0;
+    uint64_t jit_memcpy_async = 0;
+    uint64_t optix_accel_build = 0;
+    uint64_t optix_accel_compact = 0;
+    uint64_t optix_launch = 0;
+    std::vector<NativeKernelLaunchStat> kernels;
+};
+
+struct NativeLaunchAuditSnapshot {
+    NativeLaunchStageStats unknown;
+    NativeLaunchStageStats build;
+    NativeLaunchStageStats sync;
+    NativeLaunchStageStats trace_reflections;
+};
+
+class ScopedNativeLaunchStage {
+public:
+    explicit ScopedNativeLaunchStage(NativeLaunchStage stage);
+    ~ScopedNativeLaunchStage();
+
+    ScopedNativeLaunchStage(const ScopedNativeLaunchStage &) = delete;
+    ScopedNativeLaunchStage &operator=(const ScopedNativeLaunchStage &) = delete;
+
+private:
+    NativeLaunchStage previous_;
+};
+
+void native_launch_audit_clear();
+NativeLaunchAuditSnapshot native_launch_audit_snapshot();
+
+void audit_cuda_kernel_launch(const char *label,
+                              uint32_t grid_x,
+                              uint32_t grid_y,
+                              uint32_t grid_z,
+                              uint32_t block_x,
+                              uint32_t block_y,
+                              uint32_t block_z,
+                              uint64_t items = 0);
+void audit_cuda_memcpy();
+void audit_cuda_memcpy_async();
+void audit_cuda_memset_async();
+void audit_cuda_stream_synchronize();
+void audit_cuda_event_record();
+void audit_cuda_stream_wait_event();
+void audit_cub_reduce();
+void audit_cub_sort();
+void audit_cub_scan();
+void audit_jit_memcpy();
+void audit_jit_memcpy_async();
+void audit_optix_accel_build();
+void audit_optix_accel_compact();
+void audit_optix_launch();
+
+} // namespace rayd

{rayd-0.2.0 → rayd-0.2.2}/src/rayd.cpp

@@ -15,6 +15,8 @@
 #include <rayd/camera.h>
 #include <rayd/scene/scene.h>
 
+#include "native_launch_audit.h"
+
 namespace nb = nanobind;
 using namespace nb::literals;
 using namespace rayd;
@@ -139,6 +141,44 @@ struct type_caster<T, std::enable_if_t<drjit::is_array_v<T>, int>> {
 } // namespace nanobind::detail
 
 NB_MODULE(rayd, m) {
+    auto native_kernel_dict = [](const NativeKernelLaunchStat &entry) {
+        nb::dict result;
+        result["label"] = entry.label;
+        result["launches"] = entry.launches;
+        result["total_threads"] = entry.total_threads;
+        result["max_threads"] = entry.max_threads;
+        result["total_items"] = entry.total_items;
+        result["max_items"] = entry.max_items;
+        return result;
+    };
+
+    auto native_stage_dict = [&](const NativeLaunchStageStats &stats) {
+        nb::list kernels;
+        for (const NativeKernelLaunchStat &entry : stats.kernels) {
+            kernels.append(native_kernel_dict(entry));
+        }
+
+        nb::dict result;
+        result["cuda_kernel_launches"] = stats.cuda_kernel_launches;
+        result["cuda_kernel_total_threads"] = stats.cuda_kernel_total_threads;
+        result["cuda_memcpy"] = stats.cuda_memcpy;
+        result["cuda_memcpy_async"] = stats.cuda_memcpy_async;
+        result["cuda_memset_async"] = stats.cuda_memset_async;
+        result["cuda_stream_synchronize"] = stats.cuda_stream_synchronize;
+        result["cuda_event_record"] = stats.cuda_event_record;
+        result["cuda_stream_wait_event"] = stats.cuda_stream_wait_event;
+        result["cub_reduce"] = stats.cub_reduce;
+        result["cub_sort"] = stats.cub_sort;
+        result["cub_scan"] = stats.cub_scan;
+        result["jit_memcpy"] = stats.jit_memcpy;
+        result["jit_memcpy_async"] = stats.jit_memcpy_async;
+        result["optix_accel_build"] = stats.optix_accel_build;
+        result["optix_accel_compact"] = stats.optix_accel_compact;
+        result["optix_launch"] = stats.optix_launch;
+        result["kernels"] = kernels;
+        return result;
+    };
+
     auto bind_section = [](const char *name, auto &&fn) {
         try {
             fn();
@@ -164,6 +204,20 @@ NB_MODULE(rayd, m) {
     m.def("current_device",
           []() { return checked_cuda_device_count() > 0 ? jit_cuda_device() : 0; },
           "Return the current thread's active Dr.Jit CUDA device index.");
+    m.def("native_launch_audit_clear",
+          &native_launch_audit_clear,
+          "Clear grouped native launch audit counters.");
+    m.def("native_launch_audit",
+          [native_stage_dict]() {
+              const NativeLaunchAuditSnapshot snapshot = native_launch_audit_snapshot();
+              nb::dict result;
+              result["unknown"] = native_stage_dict(snapshot.unknown);
+              result["build"] = native_stage_dict(snapshot.build);
+              result["sync"] = native_stage_dict(snapshot.sync);
+              result["trace_reflections"] = native_stage_dict(snapshot.trace_reflections);
+              return result;
+          },
+          "Return grouped native launch audit counters.");
     m.def("set_device",
           &set_rayd_device,
           "device"_a,
@@ -448,6 +502,61 @@ NB_MODULE(rayd, m) {
     bind_section("mesh", [&]() {
         nb::class_<Mesh>(m, "Mesh")
             .def(nb::init<>())
+            .def("__init__",
+                 [](Mesh *mesh,
+                    nb::handle v_obj,
+                    const Vector3iDetached &f,
+                    nb::handle uv_obj,
+                    const Vector3iDetached &f_uv,
+                    bool verbose) {
+                     const std::string v_module_name =
+                         nb::cast<std::string>(v_obj.type().attr("__module__"));
+                     const std::string v_type_name =
+                         nb::cast<std::string>(v_obj.type().attr("__name__"));
+                     const std::string uv_module_name =
+                         nb::cast<std::string>(uv_obj.type().attr("__module__"));
+                     const std::string uv_type_name =
+                         nb::cast<std::string>(uv_obj.type().attr("__name__"));
+
+                     const bool has_live_v =
+                         v_module_name == "drjit.cuda.ad" && v_type_name == "Array3f";
+                     const bool has_live_uv =
+                         uv_module_name == "drjit.cuda.ad" && uv_type_name == "Array2f";
+                     if (!has_live_v && !has_live_uv) {
+                         throw nb::next_overload();
+                     }
+
+                     Vector3fDetached v_detached;
+                     if (has_live_v) {
+                         const Vector3f v_live = nb::cast<Vector3f>(v_obj);
+                         v_detached = detach<false>(v_live);
+                     } else if (!drjit_try_load<Vector3fDetached>(v_obj, v_detached, true)) {
+                         throw nb::next_overload();
+                     }
+
+                     Vector2fDetached uv_detached;
+                     if (has_live_uv) {
+                         const Vector2f uv_live = nb::cast<Vector2f>(uv_obj);
+                         uv_detached = detach<false>(uv_live);
+                     } else if (!drjit_try_load<Vector2fDetached>(uv_obj, uv_detached, true)) {
+                         throw nb::next_overload();
+                     }
+
+                     new (mesh) Mesh(v_detached, f, uv_detached, f_uv, verbose);
+                     if (has_live_v) {
+                         const Vector3f v_live = nb::cast<Vector3f>(v_obj);
+                         mesh->set_vertex_positions(v_live);
+                     }
+                     if (has_live_uv) {
+                         const Vector2f uv_live = nb::cast<Vector2f>(uv_obj);
+                         mesh->set_vertex_uv(uv_live);
+                     }
+                 },
+                 "v"_a,
+                 "f"_a,
+                 "uv"_a = Vector2fDetached(),
+                 "f_uv"_a = Vector3iDetached(),
+                 "verbose"_a = false)
             .def("__init__",
                  [](Mesh *mesh,
                     const Vector3fDetached &v,