Skip to content

Commit 1a67aaf

Browse files
kbayescopybara-github
kbayes
authored and
copybara-github
committed
Remove pointers from BVH for caching in mjCMesh.
PiperOrigin-RevId: 736849392 Change-Id: I684cca918c65718b29a1bcb41ecd6bad8d88f8f9
1 parent f25fc63 commit 1a67aaf

File tree

4 files changed

+169
-138
lines changed

4 files changed

+169
-138
lines changed

‎src/user/user_mesh.cc

+32-35
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
#include <cstdio>
2222
#include <cstring>
2323
#include <functional>
24+
#include <limits>
2425
#include <memory>
2526
#include <set>
2627
#include <string>
@@ -648,9 +649,10 @@ void mjCMesh::TryCompile(const mjVFS* vfs) {
648649

649650
// make bounding volume hierarchy
650651
if (tree_.Bvh().empty()) {
651-
face_aabb_.assign(6*nface(), 0);
652+
face_aabb_.clear();
653+
face_aabb_.reserve(3*face_.size());
652654
tree_.AllocateBoundingVolumes(nface());
653-
for (int i=0; i < nface(); i++) {
655+
for (int i = 0; i < nface(); i++) {
654656
SetBoundingVolume(i);
655657
}
656658
tree_.CreateBVH();
@@ -667,29 +669,28 @@ void mjCMesh::TryCompile(const mjVFS* vfs) {
667669

668670
// get bounding volume
669671
void mjCMesh::SetBoundingVolume(int faceid) {
670-
mjCBoundingVolume* node = tree_.GetBoundingVolume(faceid);
671-
node->SetId(faceid);
672-
node->conaffinity = 1;
673-
node->contype = 1;
674-
node->pos = center_ + 3*faceid;
675-
node->quat = NULL;
676-
double face_aamm[6] = {1E+10, 1E+10, 1E+10, -1E+10, -1E+10, -1E+10};
677-
for (int j=0; j<3; j++) {
678-
int vertid = face_[3*faceid+j];
679-
face_aamm[0] = mjMIN(face_aamm[0], vert_[3*vertid+0]);
680-
face_aamm[1] = mjMIN(face_aamm[1], vert_[3*vertid+1]);
681-
face_aamm[2] = mjMIN(face_aamm[2], vert_[3*vertid+2]);
682-
face_aamm[3] = mjMAX(face_aamm[3], vert_[3*vertid+0]);
683-
face_aamm[4] = mjMAX(face_aamm[4], vert_[3*vertid+1]);
684-
face_aamm[5] = mjMAX(face_aamm[5], vert_[3*vertid+2]);
685-
}
686-
face_aabb_[6*faceid+0] = .5 * (face_aamm[0] + face_aamm[3]);
687-
face_aabb_[6*faceid+1] = .5 * (face_aamm[1] + face_aamm[4]);
688-
face_aabb_[6*faceid+2] = .5 * (face_aamm[2] + face_aamm[5]);
689-
face_aabb_[6*faceid+3] = .5 * (face_aamm[3] - face_aamm[0]);
690-
face_aabb_[6*faceid+4] = .5 * (face_aamm[4] - face_aamm[1]);
691-
face_aabb_[6*faceid+5] = .5 * (face_aamm[5] - face_aamm[2]);
692-
node->aabb = face_aabb_.data() + 6*faceid;
672+
constexpr double kMaxVal = std::numeric_limits<double>::max();
673+
double face_aamm[6] = {kMaxVal, kMaxVal, kMaxVal, -kMaxVal, -kMaxVal, -kMaxVal};
674+
675+
for (int j = 0; j < 3; j++) {
676+
int vertid = face_[3*faceid + j];
677+
face_aamm[0] = std::min(face_aamm[0], vert_[3*vertid + 0]);
678+
face_aamm[1] = std::min(face_aamm[1], vert_[3*vertid + 1]);
679+
face_aamm[2] = std::min(face_aamm[2], vert_[3*vertid + 2]);
680+
face_aamm[3] = std::max(face_aamm[3], vert_[3*vertid + 0]);
681+
face_aamm[4] = std::max(face_aamm[4], vert_[3*vertid + 1]);
682+
face_aamm[5] = std::max(face_aamm[5], vert_[3*vertid + 2]);
683+
}
684+
685+
face_aabb_.push_back(.5 * (face_aamm[0] + face_aamm[3]));
686+
face_aabb_.push_back(.5 * (face_aamm[1] + face_aamm[4]));
687+
face_aabb_.push_back(.5 * (face_aamm[2] + face_aamm[5]));
688+
face_aabb_.push_back(.5 * (face_aamm[3] - face_aamm[0]));
689+
face_aabb_.push_back(.5 * (face_aamm[4] - face_aamm[1]));
690+
face_aabb_.push_back(.5 * (face_aamm[5] - face_aamm[2]));
691+
692+
tree_.AddBoundingVolume(faceid, 1, 1, center_ + 3*faceid, nullptr,
693+
&face_aabb_[6*faceid]);
693694
}
694695

695696

@@ -3565,7 +3566,7 @@ void mjCFlex::Compile(const mjVFS* vfs) {
35653566

35663567

35673568
// create flex BVH
3568-
void mjCFlex::CreateBVH(void) {
3569+
void mjCFlex::CreateBVH() {
35693570
int nbvh = 0;
35703571

35713572
// allocate element bounding boxes
@@ -3587,8 +3588,8 @@ void mjCFlex::CreateBVH(void) {
35873588
mjuu_copyvec(xmax, vertxpos.data() + 3*edata[0], 3);
35883589
for (int i=1; i <= dim; i++) {
35893590
for (int j=0; j<3; j++) {
3590-
xmin[j] = mjMIN(xmin[j], vertxpos[3*edata[i]+j]);
3591-
xmax[j] = mjMAX(xmax[j], vertxpos[3*edata[i]+j]);
3591+
xmin[j] = std::min(xmin[j], vertxpos[3*edata[i]+j]);
3592+
xmax[j] = std::max(xmax[j], vertxpos[3*edata[i]+j]);
35923593
}
35933594
}
35943595

@@ -3601,13 +3602,9 @@ void mjCFlex::CreateBVH(void) {
36013602
elemaabb_[6*e+5] = 0.5*(xmax[2]-xmin[2]) + radius;
36023603

36033604
// add bounding volume for this element
3604-
mjCBoundingVolume* bv = tree.GetBoundingVolume(nbvh++);
3605-
bv->contype = contype;
3606-
bv->conaffinity = conaffinity;
3607-
bv->quat = NULL;
3608-
bv->SetId(e);
3609-
bv->aabb = elemaabb_.data() + 6*e;
3610-
bv->pos = bv->aabb;
3605+
const double* aabb = elemaabb_.data() + 6*e;
3606+
tree.AddBoundingVolume(e, contype, conaffinity, aabb, nullptr, aabb);
3607+
nbvh++;
36113608
}
36123609

36133610
// create hierarchy

‎src/user/user_model.cc

+3-3
Original file line numberDiff line numberDiff line change
@@ -2218,7 +2218,7 @@ void mjCModel::CopyTree(mjModel* m) {
22182218
memcpy(m->bvh_child + 2*bvh_adr, pb->tree.Child().data(), 2*pb->tree.Nbvh()*sizeof(int));
22192219
memcpy(m->bvh_depth + bvh_adr, pb->tree.Level().data(), pb->tree.Nbvh()*sizeof(int));
22202220
for (int i=0; i<pb->tree.Nbvh(); i++) {
2221-
m->bvh_nodeid[i + bvh_adr] = pb->tree.Nodeid(i) ? *(pb->tree.Nodeid(i)) : -1;
2221+
m->bvh_nodeid[i + bvh_adr] = pb->tree.Nodeidptr(i) ? *(pb->tree.Nodeidptr(i)) : -1;
22222222
}
22232223
}
22242224
bvh_adr += pb->tree.Nbvh();
@@ -2838,7 +2838,7 @@ void mjCModel::CopyObjects(mjModel* m) {
28382838
memcpy(m->bvh_child + 2*bvh_adr, pme->tree().Child().data(), 2*pme->tree().Nbvh()*sizeof(int));
28392839
memcpy(m->bvh_depth + bvh_adr, pme->tree().Level().data(), pme->tree().Nbvh()*sizeof(int));
28402840
for (int j=0; j<pme->tree().Nbvh(); j++) {
2841-
m->bvh_nodeid[j + bvh_adr] = pme->tree().Nodeid(j) ? *(pme->tree().Nodeid(j)) : -1;
2841+
m->bvh_nodeid[j + bvh_adr] = pme->tree().Nodeid(j) > -1 ? pme->tree().Nodeid(j) : -1;
28422842
}
28432843
}
28442844

@@ -2955,7 +2955,7 @@ void mjCModel::CopyObjects(mjModel* m) {
29552955
memcpy(m->bvh_child + 2*bvh_adr, pfl->tree.Child().data(), 2*pfl->tree.Nbvh()*sizeof(int));
29562956
memcpy(m->bvh_depth + bvh_adr, pfl->tree.Level().data(), pfl->tree.Nbvh()*sizeof(int));
29572957
for (int i=0; i<pfl->tree.Nbvh(); i++) {
2958-
m->bvh_nodeid[i+ bvh_adr] = pfl->tree.Nodeid(i) ? *(pfl->tree.Nodeid(i)) : -1;
2958+
m->bvh_nodeid[i+ bvh_adr] = pfl->tree.Nodeidptr(i) ? *(pfl->tree.Nodeidptr(i)) : -1;
29592959
}
29602960
}
29612961

‎src/user/user_objects.cc

+67-62
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
#include "cc/array_safety.h"
3939
#include "engine/engine_passive.h"
4040
#include <mujoco/mjspec.h>
41+
#include <mujoco/mujoco.h>
4142
#include "user/user_api.h"
4243
#include "user/user_cache.h"
4344
#include "user/user_model.h"
@@ -347,12 +348,6 @@ const char* ResolveOrientation(double* quat, bool degree, const char* sequence,
347348

348349
//------------------------- class mjCBoundingVolumeHierarchy implementation ------------------------
349350

350-
// constructor
351-
mjCBoundingVolumeHierarchy::mjCBoundingVolumeHierarchy() {
352-
mjuu_setvec(ipos_, 0, 0, 0);
353-
mjuu_setvec(iquat_, 1, 0, 0, 0);
354-
}
355-
356351

357352
// assign position and orientation
358353
void mjCBoundingVolumeHierarchy::Set(double ipos_element[3], double iquat_element[4]) {
@@ -369,17 +364,29 @@ void mjCBoundingVolumeHierarchy::AllocateBoundingVolumes(int nleaf) {
369364
nodeid_.clear();
370365
level_.clear();
371366
bvleaf_.clear();
372-
bvleaf_.resize(nleaf);
367+
bvleaf_.reserve(nleaf);
373368
}
374369

375370

376371
void mjCBoundingVolumeHierarchy::RemoveInactiveVolumes(int nmax) {
377372
bvleaf_.erase(bvleaf_.begin() + nmax, bvleaf_.end());
378373
}
379374

375+
const mjCBoundingVolume*
376+
mjCBoundingVolumeHierarchy::AddBoundingVolume(int id, int contype, int conaffinity,
377+
const double* pos, const double* quat,
378+
const double* aabb) {
379+
bvleaf_.emplace_back(id, contype, conaffinity, pos, quat, aabb);
380+
return &bvleaf_.back();
381+
}
382+
380383

381-
mjCBoundingVolume* mjCBoundingVolumeHierarchy::GetBoundingVolume(int id) {
382-
return bvleaf_.data() + id;
384+
const mjCBoundingVolume*
385+
mjCBoundingVolumeHierarchy::AddBoundingVolume(const int* id, int contype, int conaffinity,
386+
const double* pos, const double* quat,
387+
const double* aabb) {
388+
bvleaf_.emplace_back(id, contype, conaffinity, pos, quat, aabb);
389+
return &bvleaf_.back();
383390
}
384391

385392

@@ -391,12 +398,12 @@ void mjCBoundingVolumeHierarchy::CreateBVH() {
391398
elements.reserve(bvleaf_.size());
392399
double qinv[4] = {iquat_[0], -iquat_[1], -iquat_[2], -iquat_[3]};
393400
for (int i = 0; i < bvleaf_.size(); i++) {
394-
if (bvleaf_[i].conaffinity || bvleaf_[i].contype) {
401+
if (bvleaf_[i].Conaffinity() || bvleaf_[i].Contype()) {
395402
BVElement element;
396403
element.e = &bvleaf_[i];
397-
double vert[3] = {element.e->pos[0] - ipos_[0],
398-
element.e->pos[1] - ipos_[1],
399-
element.e->pos[2] - ipos_[2]};
404+
double vert[3] = {element.e->Pos(0) - ipos_[0],
405+
element.e->Pos(1) - ipos_[1],
406+
element.e->Pos(2) - ipos_[2]};
400407
mjuu_rotVecQuat(element.lpos, vert, qinv);
401408
elements.push_back(std::move(element));
402409
}
@@ -412,34 +419,35 @@ int mjCBoundingVolumeHierarchy::MakeBVH(
412419
if (nelements == 0) {
413420
return -1;
414421
}
415-
double AAMM[6] = {mjMAXVAL, mjMAXVAL, mjMAXVAL, -mjMAXVAL, -mjMAXVAL, -mjMAXVAL};
422+
constexpr double kMaxVal = std::numeric_limits<double>::max();
423+
double AAMM[6] = {kMaxVal, kMaxVal, kMaxVal, -kMaxVal, -kMaxVal, -kMaxVal};
416424

417425
// inverse transformation
418426
double qinv[4] = {iquat_[0], -iquat_[1], -iquat_[2], -iquat_[3]};
419427

420428
// accumulate AAMM over elements
421429
for (auto element = elements_begin; element != elements_end; ++element) {
422430
// transform element aabb to aamm format
423-
double aamm[6] = {element->e->aabb[0] - element->e->aabb[3],
424-
element->e->aabb[1] - element->e->aabb[4],
425-
element->e->aabb[2] - element->e->aabb[5],
426-
element->e->aabb[0] + element->e->aabb[3],
427-
element->e->aabb[1] + element->e->aabb[4],
428-
element->e->aabb[2] + element->e->aabb[5]};
431+
double aamm[6] = {element->e->AABB(0) - element->e->AABB(3),
432+
element->e->AABB(1) - element->e->AABB(4),
433+
element->e->AABB(2) - element->e->AABB(5),
434+
element->e->AABB(0) + element->e->AABB(3),
435+
element->e->AABB(1) + element->e->AABB(4),
436+
element->e->AABB(2) + element->e->AABB(5)};
429437

430438
// update node AAMM
431-
for (int v=0; v<8; v++) {
439+
for (int v=0; v < 8; v++) {
432440
double vert[3], box[3];
433441
vert[0] = (v&1 ? aamm[3] : aamm[0]);
434442
vert[1] = (v&2 ? aamm[4] : aamm[1]);
435443
vert[2] = (v&4 ? aamm[5] : aamm[2]);
436444

437445
// rotate to the body inertial frame if specified
438-
if (element->e->quat) {
439-
mjuu_rotVecQuat(box, vert, element->e->quat);
440-
box[0] += element->e->pos[0] - ipos_[0];
441-
box[1] += element->e->pos[1] - ipos_[1];
442-
box[2] += element->e->pos[2] - ipos_[2];
446+
if (element->e->Quat()) {
447+
mjuu_rotVecQuat(box, vert, element->e->Quat());
448+
box[0] += element->e->Pos(0) - ipos_[0];
449+
box[1] += element->e->Pos(1) - ipos_[1];
450+
box[2] += element->e->Pos(2) - ipos_[2];
443451
mjuu_rotVecQuat(vert, box, qinv);
444452
}
445453

@@ -453,66 +461,73 @@ int mjCBoundingVolumeHierarchy::MakeBVH(
453461
}
454462

455463
// inflate flat AABBs
456-
for (int i=0; i<3; i++) {
457-
if (std::abs(AAMM[i]-AAMM[i+3])<mjEPS) {
458-
AAMM[i+0] -= mjEPS;
459-
AAMM[i+3] += mjEPS;
464+
for (int i = 0; i < 3; i++) {
465+
if (std::abs(AAMM[i] - AAMM[i+3]) < mjEPS) {
466+
AAMM[i + 0] -= mjEPS;
467+
AAMM[i + 3] += mjEPS;
460468
}
461469
}
462470

463471
// store current index
464472
int index = nbvh_++;
465473
child_.push_back(-1);
466474
child_.push_back(-1);
467-
nodeid_.push_back(nullptr);
475+
nodeid_.push_back(-1);
476+
nodeidptr_.push_back(nullptr);
468477
level_.push_back(lev);
469478

470479
// store bounding box of the current node
471-
for (int i=0; i<3; i++) {
472-
bvh_.push_back((AAMM[3+i] + AAMM[i]) / 2);
473-
}
474-
for (int i=0; i<3; i++) {
475-
bvh_.push_back((AAMM[3+i] - AAMM[i]) / 2);
476-
}
480+
bvh_.push_back((AAMM[3] + AAMM[0]) / 2);
481+
bvh_.push_back((AAMM[4] + AAMM[1]) / 2);
482+
bvh_.push_back((AAMM[5] + AAMM[2]) / 2);
483+
bvh_.push_back((AAMM[3] - AAMM[0]) / 2);
484+
bvh_.push_back((AAMM[4] - AAMM[1]) / 2);
485+
bvh_.push_back((AAMM[5] - AAMM[2]) / 2);
477486

478487
// leaf node, return
479-
if (nelements==1) {
480-
for (int i=0; i<2; i++) {
481-
child_[2*index+i] = -1;
482-
}
483-
nodeid_[index] = (int*)elements_begin->e->GetId();
488+
if (nelements == 1) {
489+
child_[2*index + 0] = -1;
490+
child_[2*index + 1] = -1;
491+
nodeid_[index] = *elements_begin->e->Id();
492+
nodeidptr_[index] = (int*)elements_begin->e->Id();
484493
return index;
485494
}
486495

487496
// find longest axis, by a margin of at least mjEPS, default to 0
488497
int axis = 0;
489-
double edges[3] = { AAMM[3]-AAMM[0], AAMM[4]-AAMM[1], AAMM[5]-AAMM[2] };
498+
double edges[3] = {AAMM[3] - AAMM[0], AAMM[4] - AAMM[1], AAMM[5] - AAMM[2]};
490499
if (edges[1] >= edges[0] + mjEPS) axis = 1;
491500
if (edges[2] >= edges[axis] + mjEPS) axis = 2;
492501

493502
// find median along the axis
503+
auto compare = [&](const BVElement& e1, const BVElement& e2) {
504+
if (std::abs(e1.lpos[axis] - e2.lpos[axis]) > mjEPS) {
505+
return e1.lpos[axis] < e2.lpos[axis];
506+
}
507+
// comparing pointers gives a stable sort, because they both come from the same array
508+
return e1.e < e2.e;
509+
};
510+
494511
// note: nth_element performs a partial sort of elements
495-
BVElementCompare compare;
496-
compare.axis = axis;
497512
int m = nelements / 2;
498513
std::nth_element(elements_begin, elements_begin + m, elements_end, compare);
499514

500515
// recursive calls
501516
if (m > 0) {
502-
child_[2*index+0] = MakeBVH(elements_begin, elements_begin + m, lev+1);
517+
child_[2*index + 0] = MakeBVH(elements_begin, elements_begin + m, lev + 1);
503518
}
504519

505520
if (m != nelements) {
506-
child_[2*index+1] = MakeBVH(elements_begin + m, elements_end, lev+1);
521+
child_[2*index + 1] = MakeBVH(elements_begin + m, elements_end, lev + 1);
507522
}
508523

509524
// SHOULD NOT OCCUR
510-
if (child_[2*index+0]==-1 && child_[2*index+1]==-1) {
525+
if (child_[2*index + 0] == -1 && child_[2*index + 1] == -1) {
511526
mju_error("this should have been a leaf, body=%s nelements=%d",
512527
name_.c_str(), nelements);
513528
}
514529

515-
if (lev>mjMAXTREEDEPTH) {
530+
if (lev > mjMAXTREEDEPTH) {
516531
mju_warning("max tree depth exceeded in body=%s", name_.c_str());
517532
}
518533

@@ -1595,8 +1610,9 @@ void mjCBody::ComputeBVH() {
15951610

15961611
tree.Set(ipos, iquat);
15971612
tree.AllocateBoundingVolumes(geoms.size());
1598-
for (int i=0; i<geoms.size(); i++) {
1599-
geoms[i]->SetBoundingVolume(tree.GetBoundingVolume(i));
1613+
for (const mjCGeom* geom : geoms) {
1614+
tree.AddBoundingVolume(&geom->id, geom->contype, geom->conaffinity,
1615+
geom->pos, geom->quat, geom->aabb);
16001616
}
16011617
tree.CreateBVH();
16021618
}
@@ -2463,17 +2479,6 @@ double mjCGeom::GetVolume() const {
24632479

24642480

24652481

2466-
void mjCGeom::SetBoundingVolume(mjCBoundingVolume* bv) const {
2467-
bv->SetId(&id);
2468-
bv->contype = contype;
2469-
bv->conaffinity = conaffinity;
2470-
bv->aabb = aabb;
2471-
bv->pos = pos;
2472-
bv->quat = quat;
2473-
}
2474-
2475-
2476-
24772482
// set geom diagonal inertia given density
24782483
void mjCGeom::SetInertia(void) {
24792484
// get from mesh

0 commit comments

Comments
 (0)