Speed up MakeBVH.

1. Avoid creating unnecessary temporary vectors for the left and right sides of the recursive call, since std::nth_element already splits the array in the right way.
2. Filter contype and conaffinity before calling MakeBVH. This only affects the top level of the MakeBVH recursion.
3. Avoid repeatedly computing mju_rotVecQuat for each element at each nesting level.

This is roughly a 2x speedup.

PiperOrigin-RevId: 616152673
Change-Id: I6ce7d99cd4a923c43156beb3f2493c649602260e

src/user/user_objects.cc

@@ -292,43 +292,48 @@ mjCBoundingVolume* mjCBoundingVolumeHierarchy::GetBoundingVolume(int id) {
 
 // create bounding volume hierarchy
 void mjCBoundingVolumeHierarchy::CreateBVH() {
-  std::vector<const mjCBoundingVolume*> elements(bvleaf_.size());
+  // precompute the positions of each element in the hierarchy's axes, and drop
-  for (int i=0; i<bvleaf_.size(); i++) {
+  // visual-only elements.
-    elements[i] = bvleaf_.data() + i;
+  std::vector<BVElement> elements;
+  elements.reserve(bvleaf_.size());
+  mjtNum qinv[4] = {iquat_[0], -iquat_[1], -iquat_[2], -iquat_[3]};
+  for (int i = 0; i < bvleaf_.size(); i++) {
+    if (bvleaf_[i].conaffinity || bvleaf_[i].contype) {
+      BVElement element;
+      element.e = &bvleaf_[i];
+      element.index = i;
+      mjtNum vert[3] = {element.e->pos[0] - ipos_[0],
+                        element.e->pos[1] - ipos_[1],
+                        element.e->pos[2] - ipos_[2]};
+      mju_rotVecQuat(element.lpos, vert, qinv);
+      elements.push_back(std::move(element));
+    }
   }
-  MakeBVH(elements);
+  MakeBVH(elements.begin(), elements.end());
 }
 
-
 // compute bounding volume hierarchy
-int mjCBoundingVolumeHierarchy::MakeBVH(std::vector<const mjCBoundingVolume*>& elements, int lev) {
+int mjCBoundingVolumeHierarchy::MakeBVH(
-  if (elements.empty()) {
+    std::vector<BVElement>::iterator elements_begin,
+    std::vector<BVElement>::iterator elements_end, int lev) {
+  int nelements = elements_end - elements_begin;
+  if (nelements == 0) {
     return -1;
   }
-
-  bool is_visual = true;
-  int nelements = elements.size();
   mjtNum AAMM[6] = {mjMAXVAL, mjMAXVAL, mjMAXVAL, -mjMAXVAL, -mjMAXVAL, -mjMAXVAL};
 
   // inverse transformation
   mjtNum qinv[4] = {iquat_[0], -iquat_[1], -iquat_[2], -iquat_[3]};
 
   // accumulate AAMM over elements
-  for (int i=0; i<nelements; i++) {
+  for (auto element = elements_begin; element != elements_end; ++element) {
-    // skip visual objects
-    if (elements[i]->conaffinity==0 && elements[i]->contype==0) {
-      continue;
-    } else {
-      is_visual = false;
-    }
-
     // transform element aabb to aamm format
-    mjtNum aamm[6] = {elements[i]->aabb[0] - elements[i]->aabb[3],
+    mjtNum aamm[6] = {element->e->aabb[0] - element->e->aabb[3],
-                      elements[i]->aabb[1] - elements[i]->aabb[4],
+                      element->e->aabb[1] - element->e->aabb[4],
-                      elements[i]->aabb[2] - elements[i]->aabb[5],
+                      element->e->aabb[2] - element->e->aabb[5],
-                      elements[i]->aabb[0] + elements[i]->aabb[3],
+                      element->e->aabb[0] + element->e->aabb[3],
-                      elements[i]->aabb[1] + elements[i]->aabb[4],
+                      element->e->aabb[1] + element->e->aabb[4],
-                      elements[i]->aabb[2] + elements[i]->aabb[5]};
+                      element->e->aabb[2] + element->e->aabb[5]};
 
     // update node AAMM
     for (int v=0; v<8; v++) {
@@ -338,11 +343,11 @@ int mjCBoundingVolumeHierarchy::MakeBVH(std::vector<const mjCBoundingVolume*>& e
       vert[2] = (v&4 ? aamm[5] : aamm[2]);
 
       // rotate to the body inertial frame if specified
-      if (elements[i]->quat) {
+      if (element->e->quat) {
-        mju_rotVecQuat(box, vert, elements[i]->quat);
+        mju_rotVecQuat(box, vert, element->e->quat);
-        box[0] += elements[i]->pos[0] - ipos_[0];
+        box[0] += element->e->pos[0] - ipos_[0];
-        box[1] += elements[i]->pos[1] - ipos_[1];
+        box[1] += element->e->pos[1] - ipos_[1];
-        box[2] += elements[i]->pos[2] - ipos_[2];
+        box[2] += element->e->pos[2] - ipos_[2];
         mju_rotVecQuat(vert, box, qinv);
       }
 
@@ -355,11 +360,6 @@ int mjCBoundingVolumeHierarchy::MakeBVH(std::vector<const mjCBoundingVolume*>& e
     }
   }
 
-  // a body with only visual geoms does not have a bvh
-  if (is_visual) {
-    return nbvh;
-  }
-
   // inflate flat AABBs
   for (int i=0; i<3; i++) {
     if (mju_abs(AAMM[i]-AAMM[i+3])<mjEPS) {
@@ -388,7 +388,7 @@ int mjCBoundingVolumeHierarchy::MakeBVH(std::vector<const mjCBoundingVolume*>& e
     for (int i=0; i<2; i++) {
       child[2*index+i] = -1;
     }
-    nodeid[index] = (int*)elements[0]->GetId();
+    nodeid[index] = (int*)elements_begin->e->GetId();
     return index;
   }
 
@@ -398,66 +398,23 @@ int mjCBoundingVolumeHierarchy::MakeBVH(std::vector<const mjCBoundingVolume*>& e
   axis = edges[axis] > edges[2] ? axis : 2;
 
   // find median along the axis
-  std::vector<mjtNum> pos(nelements);
+  auto m = nelements/2;
-
+  // Note: nth element performs a partial sort of elements
-  for (int i=0; i<nelements; i++) {
+  BVElementCompare compare;
-    // get position in the body inertial frame
+  compare.axis = axis;
-    mjtNum vert[3] = {elements[i]->pos[0] - ipos_[0],
+  std::nth_element(elements_begin, elements_begin + m, elements_end, compare);
-                      elements[i]->pos[1] - ipos_[1],
-                      elements[i]->pos[2] - ipos_[2]};
-    mjtNum lpos[3];
-    mju_rotVecQuat(lpos, vert, qinv);
-    pos[i] = lpos[axis];
-  }
-
-  auto m = pos.size()/2;
-  std::nth_element(pos.begin(), pos.begin() + m, pos.end());
-  mjtNum threshold = pos[m];
-
-  // split using median
-  std::vector<const mjCBoundingVolume*> left;
-  std::vector<const mjCBoundingVolume*> right;
-  int skipped = 0;
-
-  for (int i=0; i<nelements; i++) {
-    // get position in the body inertial frame
-    mjtNum vert[3] = {elements[i]->pos[0] - ipos_[0],
-                      elements[i]->pos[1] - ipos_[1],
-                      elements[i]->pos[2] - ipos_[2]};
-    mjtNum lpos[3];
-    mju_rotVecQuat(lpos, vert, qinv);
-
-    // skip visual objects
-    if (elements[i]->conaffinity==0 && elements[i]->contype==0) {
-      skipped++;
-      continue;
-    }
-    if (lpos[axis] < threshold) {
-      left.push_back(elements[i]);
-    } else if (lpos[axis] > threshold) {
-      right.push_back(elements[i]);
-    } else {
-      if (left.size() < right.size()) left.push_back(elements[i]);
-      else right.push_back(elements[i]);
-    }
-  }
 
   // recursive calls
-  if (!left.empty()) {
+  if (m > 0) {
-    child[2*index+0] = MakeBVH(left, lev+1);
+    child[2*index+0] = MakeBVH(elements_begin, elements_begin + m, lev+1);
   }
 
-  if (!right.empty()) {
+  if (m != nelements) {
-    child[2*index+1] = MakeBVH(right, lev+1);
+    child[2*index+1] = MakeBVH(elements_begin + m, elements_end, lev+1);
   }
 
   // SHOULD NOT OCCUR
-  if (left.size()+right.size()+skipped != nelements) {
+  if (child[2*index+0]==-1 && child[2*index+1]==-1) {
-    mju_error("some elements were lost, body=%s parent=%d children=%lu",
-              name_.c_str(), nelements, left.size()+right.size()+skipped);
-  }
-
-  if (child[2*index+0]==-1 && child[2*index+1]==-1 && !skipped) {
     mju_error("this should have been a leaf, body=%s nelements=%d",
               name_.c_str(), nelements);
   }
@@ -469,8 +426,6 @@ int mjCBoundingVolumeHierarchy::MakeBVH(std::vector<const mjCBoundingVolume*>& e
   return index;
 }
 
-
-
 //------------------------- class mjCDef implementation --------------------------------------------
 
 // constructor

src/user/user_objects.h

@@ -15,10 +15,8 @@
 #ifndef MUJOCO_SRC_USER_USER_OBJECTS_H_
 #define MUJOCO_SRC_USER_USER_OBJECTS_H_
 
-#include <array>
 #include <functional>
 #include <map>
-#include <optional>
 #include <string>
 #include <string_view>
 #include <utility>
@@ -144,7 +142,28 @@ class mjCBoundingVolumeHierarchy : public mjCBoundingVolumeHierarchy_ {
   mjCBoundingVolume* GetBoundingVolume(int id);
 
  private:
-  int MakeBVH(std::vector<const mjCBoundingVolume*>& elements, int lev = 0);
+  // internal class used during BVH construction, for partial sorting of bounding volumes
+  struct BVElement {
+    const mjCBoundingVolume* e;
+    // index of the element in the original input to BVH, used to ensure a stable sort
+    int index;
+    // position of the element in the BVH axes
+    mjtNum lpos[3];
+  };
+
+  struct BVElementCompare {
+    int axis = 0;
+
+    bool operator()(const BVElement& e1, const BVElement& e2) const {
+      if (e1.lpos[axis] != e2.lpos[axis]) {
+        return e1.lpos[axis] < e2.lpos[axis];
+      }
+      return e1.index < e2.index;
+    }
+  };
+
+  int MakeBVH(std::vector<BVElement>::iterator elements_begin,
+              std::vector<BVElement>::iterator elements_end, int lev = 0);
 };