From 54dda7b7f3ac412d43e667ce07668494de8292c3 Mon Sep 17 00:00:00 2001
From: Wenzel Jakob <wenzel.jakob@epfl.ch>
Date: Sat, 21 Dec 2024 13:39:07 +0900
Subject: [PATCH] Fix builtin kd-tree min/max-binning precision problems

Mitsuba usually defers ray tracing to OptiX/Embree. It can, however,
also be built in a special mode that uses an internal kd-tree. This
is useful for things like double precision ray tracing, which is
unfortunately not possible in OptiX/Embree.

There was a long-standing bug in the builtin-kdtree implementation that
somehow got added in the port of this code from Mitsuba 0.5. When the
scene contains perfectly axis-aligned primitives, the min-max binning
builder could sometimes produce highly unbalanced trees. This turned out
to be due to an incorrect stopping criterion, and because of
inconsistencies caused by flush-to-zero (FTZ). Both issues are fixed as
of this commit, and nice trees are produced for two problematic test
workloads that were separately provided by different users.
---
 include/mitsuba/render/kdtree.h | 25 ++++++++++++++++++++++---
 1 file changed, 22 insertions(+), 3 deletions(-)

diff --git a/include/mitsuba/render/kdtree.h b/include/mitsuba/render/kdtree.h
index d5eaa84e1..1c6e2a9a3 100644
--- a/include/mitsuba/render/kdtree.h
+++ b/include/mitsuba/render/kdtree.h
@@ -190,6 +190,8 @@ class OrderedChunkAllocator {
 template <typename Value>
 struct ConcurrentVector {
     ConcurrentVector() : m_size_and_capacity(0) { }
+    ConcurrentVector(const ConcurrentVector &) = delete;
+    ConcurrentVector(ConcurrentVector &&) = delete;
     ~ConcurrentVector() { release(); }
 
     void reserve(uint32_t size) {
@@ -286,8 +288,8 @@ struct ConcurrentVector {
         #endif
     }
 
-    uint64_t size() const {
-        return m_size_and_capacity.load(std::memory_order_acquire);
+    uint32_t size() const {
+        return (uint32_t) m_size_and_capacity.load(std::memory_order_acquire);
     }
 
     void release() {
@@ -1117,11 +1119,28 @@ template <typename BoundingBox_, typename Index_, typename CostModel_,
             Assert(m_bbox.contains(tight_bbox));
         }
 
+        struct FTZGuard {
+            FTZGuard() {
+                #if defined(DRJIT_X86_64)
+                csr = _mm_getcsr();
+                _mm_setcsr(csr & ~(_MM_FLUSH_ZERO_ON | _MM_DENORMALS_ZERO_ON));
+                #endif
+            }
+
+            ~FTZGuard() {
+                #if defined(DRJIT_X86_64)
+                _mm_setcsr(csr);
+                #endif
+            }
+            int csr;
+        };
+
         /// Run one iteration of min-max binning and spawn recursive tasks
         void execute() {
             ScopedSetThreadEnvironment env(m_ctx.env);
             Size prim_count = Size(m_indices.size());
             const Derived &derived = m_ctx.derived;
+            FTZGuard g;
 
             m_ctx.work_units++;
 
@@ -1130,7 +1149,7 @@ template <typename BoundingBox_, typename Index_, typename CostModel_,
             /* ==================================================================== */
 
             if (prim_count <= derived.stop_primitives() ||
-                m_depth >= derived.max_depth() || m_tight_bbox.collapsed()) {
+                m_depth >= derived.max_depth() || dr::count(m_tight_bbox.max == m_tight_bbox.min) > 1) {
                 make_leaf(std::move(m_indices));
                 return;
             }