chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions 10 Packages Projects Releases Wiki Activity

Replace mjQUICKSORT with faster, native sorting function. fixes #1638

Browse Source 
PiperOrigin-RevId: 688519096
Change-Id: I05ab1576703c2458968ef81915ac682a8beac391
This commit is contained in:
Kyle Bayes
2024-10-22 06:04:10 -07:00
committed by Copybara-Service
parent 9e1aa37b55
commit 2b0629d75f
5 changed files with 172 additions and 82 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_collision_driver.c
+25 -18
View File
@@ -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);
src/engine/engine_collision_sdf.c
+11 -6
View File
@@ -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++) {
src/engine/engine_sort.h
+52 -38
View File
@@ -15,46 +15,60 @@
#ifndef MUJOCO_SRC_ENGINE_ENGINE_SORT_H_
#define MUJOCO_SRC_ENGINE_ENGINE_SORT_H_
#if !defined(__cplusplus)
#include <stddef.h>
#include <stdlib.h>
// 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 <algorithm>
#include <cstddef>
#include <cstdlib>
// sorting function using std::sort
template <typename T>
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_
src/render/render_gl3.c
+11 -4
View File
@@ -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;
test/engine/engine_sort_test.cc
+73 -16
View File
@@ -14,11 +14,12 @@
#include "src/engine/engine_sort.h"
#include <cstddef>
#include <algorithm>
#include <array>
#include <utility>
#include <vector>
#include <gtest/gtest.h>
#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<int, kPermutationSize>& 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<int> 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<int, kPermutationSize> initial, buf;
for (int i = 0; i < kPermutationSize; ++i) {
initial[i] = i + 1;
}
std::array<int, kPermutationSize> arr = initial;
do {
++total;
std::array<int, kPermutationSize> 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<int, 2500> 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<IntStruct> y = {{1}, {3}, {2}};
mjQUICKSORT(y.data(), y.size(), sizeof(IntStruct), intstruct_compare, NULL);
std::vector<IntStruct> 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);