From c5aa22b177676928e5f1d7ffdcc2195657fe472c Mon Sep 17 00:00:00 2001 From: Nimrod Gileadi Date: Fri, 15 Mar 2024 09:23:49 -0700 Subject: [PATCH] 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 | 135 +++++++++++++-------------------------- src/user/user_objects.h | 25 +++++++- 2 files changed, 67 insertions(+), 93 deletions(-) diff --git a/src/user/user_objects.cc b/src/user/user_objects.cc index db38cad3..dd21fab0 100644 --- a/src/user/user_objects.cc +++ b/src/user/user_objects.cc @@ -292,43 +292,48 @@ mjCBoundingVolume* mjCBoundingVolumeHierarchy::GetBoundingVolume(int id) { // create bounding volume hierarchy void mjCBoundingVolumeHierarchy::CreateBVH() { - std::vector elements(bvleaf_.size()); - for (int i=0; i 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& elements, int lev) { - if (elements.empty()) { +int mjCBoundingVolumeHierarchy::MakeBVH( + std::vector::iterator elements_begin, + std::vector::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; iconaffinity==0 && elements[i]->contype==0) { - continue; - } else { - is_visual = false; - } - + for (auto element = elements_begin; element != elements_end; ++element) { // transform element aabb to aamm format - mjtNum aamm[6] = {elements[i]->aabb[0] - elements[i]->aabb[3], - elements[i]->aabb[1] - elements[i]->aabb[4], - elements[i]->aabb[2] - elements[i]->aabb[5], - elements[i]->aabb[0] + elements[i]->aabb[3], - elements[i]->aabb[1] + elements[i]->aabb[4], - elements[i]->aabb[2] + elements[i]->aabb[5]}; + mjtNum aamm[6] = {element->e->aabb[0] - element->e->aabb[3], + element->e->aabb[1] - element->e->aabb[4], + element->e->aabb[2] - element->e->aabb[5], + element->e->aabb[0] + element->e->aabb[3], + element->e->aabb[1] + element->e->aabb[4], + 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& e vert[2] = (v&4 ? aamm[5] : aamm[2]); // rotate to the body inertial frame if specified - if (elements[i]->quat) { - mju_rotVecQuat(box, vert, elements[i]->quat); - box[0] += elements[i]->pos[0] - ipos_[0]; - box[1] += elements[i]->pos[1] - ipos_[1]; - box[2] += elements[i]->pos[2] - ipos_[2]; + if (element->e->quat) { + mju_rotVecQuat(box, vert, element->e->quat); + box[0] += element->e->pos[0] - ipos_[0]; + box[1] += element->e->pos[1] - ipos_[1]; + box[2] += element->e->pos[2] - ipos_[2]; mju_rotVecQuat(vert, box, qinv); } @@ -355,11 +360,6 @@ int mjCBoundingVolumeHierarchy::MakeBVH(std::vector& 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])& 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& e axis = edges[axis] > edges[2] ? axis : 2; // find median along the axis - std::vector pos(nelements); - - for (int i=0; ipos[0] - ipos_[0], - 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 left; - std::vector right; - int skipped = 0; - - for (int i=0; ipos[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]); - } - } + auto m = nelements/2; + // Note: nth element performs a partial sort of elements + BVElementCompare compare; + compare.axis = axis; + std::nth_element(elements_begin, elements_begin + m, elements_end, compare); // recursive calls - if (!left.empty()) { - child[2*index+0] = MakeBVH(left, lev+1); + if (m > 0) { + child[2*index+0] = MakeBVH(elements_begin, elements_begin + m, lev+1); } - if (!right.empty()) { - child[2*index+1] = MakeBVH(right, lev+1); + if (m != nelements) { + child[2*index+1] = MakeBVH(elements_begin + m, elements_end, lev+1); } // SHOULD NOT OCCUR - if (left.size()+right.size()+skipped != nelements) { - 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) { + if (child[2*index+0]==-1 && child[2*index+1]==-1) { 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& e return index; } - - //------------------------- class mjCDef implementation -------------------------------------------- // constructor diff --git a/src/user/user_objects.h b/src/user/user_objects.h index 988a084d..2d24ea85 100644 --- a/src/user/user_objects.h +++ b/src/user/user_objects.h @@ -15,10 +15,8 @@ #ifndef MUJOCO_SRC_USER_USER_OBJECTS_H_ #define MUJOCO_SRC_USER_USER_OBJECTS_H_ -#include #include #include -#include #include #include #include @@ -144,7 +142,28 @@ class mjCBoundingVolumeHierarchy : public mjCBoundingVolumeHierarchy_ { mjCBoundingVolume* GetBoundingVolume(int id); private: - int MakeBVH(std::vector& 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::iterator elements_begin, + std::vector::iterator elements_end, int lev = 0); };