From 2b0629d75f2a752c2b8db86d4b5c2df657680ca7 Mon Sep 17 00:00:00 2001 From: Kyle Bayes Date: Tue, 22 Oct 2024 06:04:10 -0700 Subject: [PATCH] Replace `mjQUICKSORT` with faster, native sorting function. fixes #1638 PiperOrigin-RevId: 688519096 Change-Id: I05ab1576703c2458968ef81915ac682a8beac391 --- src/engine/engine_collision_driver.c | 43 +++++++------ src/engine/engine_collision_sdf.c | 17 ++++-- src/engine/engine_sort.h | 90 ++++++++++++++++------------ src/render/render_gl3.c | 15 +++-- test/engine/engine_sort_test.cc | 89 ++++++++++++++++++++++----- 5 files changed, 172 insertions(+), 82 deletions(-) diff --git a/src/engine/engine_collision_driver.c b/src/engine/engine_collision_driver.c index b49d54ec..616dd294 100644 --- a/src/engine/engine_collision_driver.c +++ b/src/engine/engine_collision_driver.c @@ -224,10 +224,8 @@ int mj_isElemActive(const mjModel* m, int f, int e) { //----------------------------- collision detection entry point ------------------------------------ // compare contact pairs by their geom/elem/vert IDs -quicksortfunc(contactcompare, context, el1, el2) { +static inline int contactcompare(const mjContact* c1, const mjContact* c2, void* context) { const mjModel* m = (const mjModel*) context; - mjContact* c1 = (mjContact*)el1; - mjContact* c2 = (mjContact*)el2; // get colliding object ids int con1_obj1 = c1->geom[0] >= 0 ? c1->geom[0] : (c1->elem[0] >= 0 ? c1->elem[0] : c1->vert[0]); @@ -258,6 +256,9 @@ quicksortfunc(contactcompare, context, el1, el2) { return 0; } +// define contactSort function for sorting contacts +mjSORT(contactSort, mjContact, contactcompare) + // main collision function @@ -364,8 +365,14 @@ void mj_collision(const mjModel* m, mjData* d) { int ncon_after = d->ncon; // sort contacts - mjQUICKSORT(d->contact + ncon_before, ncon_after - ncon_before, - sizeof(mjContact), contactcompare, (void*) m); + int n = ncon_after - ncon_before; + if (n > 1) { + mj_markStack(d); + mjContact* buf = (mjContact*)mj_stackAllocByte(d, n * sizeof(mjContact), + _Alignof(mjContact)); + contactSort(d->contact + ncon_before, buf, n, (void*)m); + mj_freeStack(d); + } } // process bodyflex pair: all-to-all @@ -1006,10 +1013,7 @@ typedef struct _mjtSAP mjtSAP; // comparison function for SAP -quicksortfunc(SAPcompare, context, el1, el2) { - mjtSAP* obj1 = (mjtSAP*)el1; - mjtSAP* obj2 = (mjtSAP*)el2; - +static inline int SAPcmp(mjtSAP* obj1, mjtSAP* obj2, void* context) { if (obj1->value < obj2->value) { return -1; } else if (obj1->value == obj2->value) { @@ -1019,6 +1023,8 @@ quicksortfunc(SAPcompare, context, el1, el2) { } } +// define SAPsort function for sorting SAP sorting +mjSORT(SAPsort, mjtSAP, SAPcmp) // given list of axis-aligned bounding boxes in AAMM (xmin[3], xmax[3]) format, @@ -1043,7 +1049,8 @@ static int mj_SAP(mjData* d, const mjtNum* aamm, int n, int axis, int* pair, int } // sort along specified axis - mjQUICKSORT(sortbuf, 2*n, sizeof(mjtSAP), SAPcompare, 0); + mjtSAP* buf = (mjtSAP*) mj_stackAllocByte(d, 2*n*sizeof(mjtSAP), _Alignof(mjtSAP)); + SAPsort(sortbuf, buf, 2*n, NULL); // define the other two axes int axisA, axisB; @@ -1133,19 +1140,18 @@ static void updateCov(mjtNum cov[9], const mjtNum vec[3], const mjtNum cen[3]) { // comparison function for unsigned ints -quicksortfunc(uintcompare, context, el1, el2) { - unsigned int n1 = *(unsigned int*)el1; - unsigned int n2 = *(unsigned int*)el2; - - if (n1 < n2) { +static inline int uintcmp(int* i, int* j, void* context) { + if ((unsigned) *i < (unsigned) *j) { return -1; - } else if (n1 == n2) { + } else if (*i == *j) { return 0; } else { return 1; } } +// define bfsort function for sorting bodyflex pairs +mjSORT(bfsort, int, uintcmp) // broadphase collision detector @@ -1281,8 +1287,9 @@ int mj_broadphase(const mjModel* m, mjData* d, int* bfpair, int maxpair) { endbroad: // sort bodyflex pairs by signature - if (npair) { - mjQUICKSORT(bfpair, npair, sizeof(int), uintcompare, 0); + if (npair > 1) { + int* buf = mj_stackAllocInt(d, npair); + bfsort(bfpair, buf, npair, NULL); } mj_freeStack(d); diff --git a/src/engine/engine_collision_sdf.c b/src/engine/engine_collision_sdf.c index 4174d34b..e01de86f 100644 --- a/src/engine/engine_collision_sdf.c +++ b/src/engine/engine_collision_sdf.c @@ -306,11 +306,10 @@ static void undoTransformation(const mjModel* m, const mjData* d, int g, //---------------------------- narrow phase ----------------------------------------------- // comparison function for contact sorting -quicksortfunc(distcompare, dist, i1, i2) { - mjtNum d1 = ((mjtNum*)dist)[*(int*)i1]; - mjtNum d2 = ((mjtNum*)dist)[*(int*)i2]; - - if (d1 < d2) { +static inline int distcmp(int* i, int* j, void* context) { + mjtNum d1 = ((mjtNum*)context)[*i]; + mjtNum d2 = ((mjtNum*)context)[*j]; + if (d1 < d2) { return -1; } else if (d1 == d2) { return 0; @@ -319,6 +318,9 @@ quicksortfunc(distcompare, dist, i1, i2) { } } +// define distSort function for contact sorting +mjSORT(distSort, int, distcmp) + // check if the collision point already exists static int isknown(const mjtNum* points, const mjtNum x[3], int cnt) { for (int i = 0; i < cnt; i++) { @@ -642,7 +644,10 @@ int mjc_MeshSDF(const mjModel* m, const mjData* d, mjContact* con, int g1, int g } // sort contacts using depth - mjQUICKSORT(index, ncandidate, sizeof(int), distcompare, dist); + if (ncandidate > 1) { + int buf[MAXMESHPNT]; + distSort(index, buf, ncandidate, dist); + } // add only the first mjMAXCONPAIR pairs for (int i=0; i < mju_min(ncandidate, mjMAXCONPAIR); i++) { diff --git a/src/engine/engine_sort.h b/src/engine/engine_sort.h index 049a8d9d..e4a7bf06 100644 --- a/src/engine/engine_sort.h +++ b/src/engine/engine_sort.h @@ -15,46 +15,60 @@ #ifndef MUJOCO_SRC_ENGINE_ENGINE_SORT_H_ #define MUJOCO_SRC_ENGINE_ENGINE_SORT_H_ -#if !defined(__cplusplus) -#include -#include +// threshold size of a run to do insertion sort on +#define _mjRUNSIZE 32 -// sorting functions using q_sort_s/r -#ifdef _WIN32 -#define mjQUICKSORT(buf, elnum, elsz, func, context) \ - qsort_s(buf, elnum, elsz, func, context) -#define quicksortfunc(name, context, el1, el2) \ - static int name(void* context, const void* el1, const void* el2) -#else // assumes POSIX -#ifdef __APPLE__ -#define mjQUICKSORT(buf, elnum, elsz, func, context) \ - qsort_r(buf, elnum, elsz, context, func) -#define quicksortfunc(name, context, el1, el2) \ - static int name(void* context, const void* el1, const void* el2) -#else // non-Apple -#define mjQUICKSORT(buf, elnum, elsz, func, context) \ - qsort_r(buf, elnum, elsz, func, context) -#define quicksortfunc(name, context, el1, el2) \ - static int name(const void* el1, const void* el2, void* context) -#endif -#endif -#else -#include -#include -#include - -// sorting function using std::sort -template -void mjQUICKSORT(T* buf, size_t elnum, size_t elsz, - int (*compare)(const void* a, const void* b, void* c), - void* context) { - std::sort(buf, buf + elnum, [compare, context](const T& a, const T& b) { - return compare(&a, &b, context) < 0; - }); +// insertion sort sub-macro that runs on a sub-array [start, ..., end) +#define _mjINSERTION_SORT(type, arr, start, end, cmp, context) \ +{ \ + for (int j = start + 1; j < end; j++) { \ + type tmp = arr[j]; \ + int k = j - 1; \ + for (; k >= start && cmp(arr + k, &tmp, context) > 0; k--) { \ + arr[k + 1] = arr[k]; \ + } \ + arr[k + 1] = tmp; \ + } \ } -#define quicksortfunc(name, context, el1, el2) \ - static int name(const void* el1, const void* el2, void* context) -#endif +// sub-macro that merges two sub-sorted arrays [start, ..., mid), [mid, ..., end) together +#define _mjMERGE(type, arr, buf, start, mid, end, cmp, context) \ +{ \ + int len1 = mid - start, len2 = end - mid; \ + type* left = buf, *right = buf + len1; \ + for (int i = 0; i < len1; i++) left[i] = arr[start + i]; \ + for (int i = 0; i < len2; i++) right[i] = arr[mid + i]; \ + int i = 0, j = 0, k = start; \ + while (i < len1 && j < len2) { \ + if (cmp(left + i, right + j, context) <= 0) { \ + arr[k++] = left[i++]; \ + } else { \ + arr[k++] = right[j++]; \ + } \ + } \ + while (i < len1) arr[k++] = left[i++]; \ + while (j < len2) arr[k++] = right[j++]; \ +} + +// defines an inline stable sorting function via tiled merge sorting (timsort) +// function is of form: +// void name(type* arr, type* buf, int n, void* context) +// where arr is the array of size n to be sorted inplace and buf is a buffer of size n. +#define mjSORT(name, type, cmp) \ + static inline void name(type* arr, type* buf, int n, void* context) { \ + for (int start = 0; start < n; start += _mjRUNSIZE) { \ + int end = (start + _mjRUNSIZE < n) ? start + _mjRUNSIZE : n; \ + _mjINSERTION_SORT(type, arr, start, end, cmp, context); \ + } \ + for (int len = _mjRUNSIZE; len < n; len *= 2) { \ + for (int start = 0; start < n; start += 2*len) { \ + int mid = start + len; \ + int end = (start + 2*len < n) ? start + 2*len : n; \ + if (mid < end) { \ + _mjMERGE(type, arr, buf, start, mid, end, cmp, context); \ + } \ + } \ + } \ + } #endif // MUJOCO_SRC_ENGINE_ENGINE_SORT_H_ diff --git a/src/render/render_gl3.c b/src/render/render_gl3.c index 44ac1596..5c3b9571 100644 --- a/src/render/render_gl3.c +++ b/src/render/render_gl3.c @@ -766,10 +766,10 @@ static void setView(int view, mjrRect viewport, const mjvScene* scn, const mjrCo // comparison function for geom sorting -quicksortfunc(geomcompare, context, el1, el2) { +static inline int geomcmp(int* i, int* j, void* context) { mjvGeom* geom = (mjvGeom*) context; - float d1 = geom[*(int*)el1].camdist; - float d2 = geom[*(int*)el2].camdist; + float d1 = geom[*i].camdist; + float d2 = geom[*j].camdist; if (d1 < d2) { return -1; @@ -780,6 +780,9 @@ quicksortfunc(geomcompare, context, el1, el2) { } } +// define geomSort function for sorting geoms +mjSORT(geomSort, int, geomcmp) + // adjust light n position and direction @@ -907,7 +910,11 @@ void mjr_render(mjrRect viewport, mjvScene* scn, const mjrContext* con) { } // sort transparent geoms according to distance to camera - mjQUICKSORT(scn->geomorder, nt, sizeof(int), geomcompare, scn->geoms); + if (nt > 1) { + int *buf = (int*) mju_malloc(nt * sizeof(int)); + geomSort(scn->geomorder, buf, nt, scn->geoms); + mju_free(buf); + } // allow only one reflective geom int j = 0; diff --git a/test/engine/engine_sort_test.cc b/test/engine/engine_sort_test.cc index bfb69ea3..a68d77c8 100644 --- a/test/engine/engine_sort_test.cc +++ b/test/engine/engine_sort_test.cc @@ -14,11 +14,12 @@ #include "src/engine/engine_sort.h" -#include +#include +#include +#include #include #include -#include "src/engine/engine_sort.h" #include "test/fixture.h" namespace mujoco { @@ -26,23 +27,57 @@ namespace { using EngineSortTest = MujocoTest; +constexpr int kPermutationSize = 5; // 5! = 120 checks + +constexpr int factorial() { + int n = 1; + for (int i = 1; i <= kPermutationSize; ++i) { + n *= i; + } + return n; +} + struct IntStruct { int value; }; -quicksortfunc(int_compare, context, x, y) { - int a = *(int*)x; - int b = *(int*)y; - if (a < b) { +// computes next permutation in lexicographical order in place +// returns false if there is no next permutation +bool NextPermutation(std::array& arr) { + int j; + for (j = arr.size() - 2; j >= 0; --j) { + if (arr[j] < arr[j + 1]) { + break; + } + } + + // last permutation i.e. [n, n - 1, ..., 3, 2, 1] + if (j < 0) { + return false; + } + + for (int l = arr.size() - 1; l > -1; --l) { + if (arr[j] < arr[l]) { + std::swap(arr[j], arr[l]); + std::reverse(arr.begin() + j + 1, arr.end()); + return true; + } + } + return true; +} + +int IntCompare(int* i, int* j, void* context) { + if (*i < *j) { return -1; - } else if (a == b) { + } else if (*i == *j) { return 0; } else { return 1; } } +mjSORT(IntSort, int, IntCompare) -quicksortfunc(intstruct_compare, context, x, y) { +int IntStructCompare(const IntStruct* x, const IntStruct* y, void* context) { IntStruct* a = (IntStruct*)x; IntStruct* b = (IntStruct*)y; if (a->value < b->value) { @@ -53,18 +88,40 @@ quicksortfunc(intstruct_compare, context, x, y) { return 1; } } +mjSORT(IntStructSort, IntStruct, IntStructCompare) TEST_F(EngineSortTest, Sort) { - // test int - std::vector x = {1, 3, 2}; - mjQUICKSORT(x.data(), x.size(), sizeof(int), int_compare, x.data()); - EXPECT_EQ(x[0], 1); - EXPECT_EQ(x[1], 2); - EXPECT_EQ(x[2], 3); + int total = 0; + std::array initial, buf; + for (int i = 0; i < kPermutationSize; ++i) { + initial[i] = i + 1; + } + std::array arr = initial; + do { + ++total; + std::array sorted_arr = arr; + IntSort(sorted_arr.data(), buf.data(), sorted_arr.size(), nullptr); + EXPECT_EQ(sorted_arr, initial); + } while (NextPermutation(arr)); + ASSERT_EQ(total, factorial()); +} - // test custom struct with mjQUICKSORT +TEST_F(EngineSortTest, LargeSort) { + std::array arr, buf; + int n = 1; + for (int i = arr.size() - 1; i >= 0; --i) { + arr[i] = n++; + } + IntSort(arr.data(), buf.data(), arr.size(), nullptr); + for (int i = 0; i < arr.size(); ++i) { + EXPECT_EQ(arr[i], i + 1); + } +} + +TEST_F(EngineSortTest, SortStruct) { std::vector y = {{1}, {3}, {2}}; - mjQUICKSORT(y.data(), y.size(), sizeof(IntStruct), intstruct_compare, NULL); + std::vector buf(3); + IntStructSort(y.data(), buf.data(), y.size(), nullptr); EXPECT_EQ(y[0].value, 1); EXPECT_EQ(y[1].value, 2); EXPECT_EQ(y[2].value, 3);