manually inline mju_combineSparse

PiperOrigin-RevId: 720326669
Change-Id: I03727291f133e3d481e73ad88735e54af8f62f1f

src/engine/engine_util_sparse.c

@@ -193,19 +193,6 @@ void mju_mulMatTVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int
 
 
 
-// res = res*scl1 + vec*scl2
-static void mju_addToSclScl(mjtNum* res, const mjtNum* vec, mjtNum scl1, mjtNum scl2, int n) {
-#ifdef mjUSEAVX
-  mju_addToSclScl_avx(res, vec, scl1, scl2, n);
-#else
-  for (int i=0; i < n; i++) {
-    res[i] = res[i]*scl1 + vec[i]*scl2;
-  }
-#endif  // mjUSEAVX
-}
-
-
-
 // count the number of non-zeros in the sum of two sparse vectors
 int mju_combineSparseCount(int a_nnz, int b_nnz, const int* a_ind, const int* b_ind) {
   int a = 0, b = 0, c_nnz = 0;
@@ -233,69 +220,6 @@ int mju_combineSparseCount(int a_nnz, int b_nnz, const int* a_ind, const int* b_
 
 
 
-// combine two sparse vectors: dst = a*dst + b*src, return nnz of result
-int mju_combineSparse(mjtNum* dst, const mjtNum* src, mjtNum a, mjtNum b,
-                      int dst_nnz, int src_nnz, int* dst_ind, const int* src_ind,
-                      mjtNum* buf, int* buf_ind) {
-  // check for identical pattern
-  if (dst_nnz == src_nnz) {
-    if (mju_compare(dst_ind, src_ind, dst_nnz)) {
-      // combine mjtNum data directly
-      mju_addToSclScl(dst, src, a, b, dst_nnz);
-      return dst_nnz;
-    }
-  }
-
-  // copy dst into buf
-  if (dst_nnz) {
-    mju_copy(buf, dst, dst_nnz);
-    mju_copyInt(buf_ind, dst_ind, dst_nnz);
-  }
-
-  // prepare to merge buf and src into dst
-  int bi = 0, si = 0, nnz = 0;
-  int buf_nnz = dst_nnz;
-
-  // merge vectors
-  while (bi < buf_nnz && si < src_nnz) {
-    int badr = buf_ind[bi];
-    int sadr = src_ind[si];
-
-    if (badr == sadr) {
-      dst[nnz] = a*buf[bi++] + b*src[si++];
-      dst_ind[nnz++] = badr;
-    }
-
-    // buf only
-    else if (badr < sadr) {
-      dst[nnz] = a*buf[bi++];
-      dst_ind[nnz++] = badr;
-    }
-
-    // src only
-    else {
-      dst[nnz] = b*src[si++];
-      dst_ind[nnz++] = sadr;
-    }
-  }
-
-  // the rest of src only
-  while (si < src_nnz) {
-    dst[nnz] = b*src[si];
-    dst_ind[nnz++] = src_ind[si++];
-  }
-
-  // the rest of buf only
-  while (bi < buf_nnz) {
-    dst[nnz] = a*buf[bi];
-    dst_ind[nnz++] = buf_ind[bi++];
-  }
-
-  return nnz;
-}
-
-
-
 // incomplete combine sparse: dst = a*dst + b*src at common indices
 void mju_combineSparseInc(mjtNum* dst, const mjtNum* src, int n, mjtNum a, mjtNum b,
                           int dst_nnz, int src_nnz, const int* dst_ind, const int* src_ind) {

src/engine/engine_util_sparse.h

@@ -20,7 +20,7 @@
 #include <mujoco/mjdata.h>
 #include <mujoco/mjexport.h>
 #include <mujoco/mjtnum.h>
-#include "engine/engine_util_sparse_avx.h"  // IWYU pragma: keep
+#include "engine/engine_util_sparse_avx.h"
 
 #ifdef __cplusplus
 extern "C" {
@@ -57,11 +57,6 @@ MJAPI void mju_compressSparse(mjtNum* mat, int nr, int nc,
 // count the number of non-zeros in the sum of two sparse vectors
 MJAPI int mju_combineSparseCount(int a_nnz, int b_nnz, const int* a_ind, const int* b_ind);
 
-// combine two sparse vectors: dst = a*dst + b*src, return nnz of result
-int mju_combineSparse(mjtNum* dst, const mjtNum* src, mjtNum a, mjtNum b,
-                      int dst_nnz, int src_nnz, int* dst_ind, const int* src_ind,
-                      mjtNum* buf, int* buf_ind);
-
 // incomplete combine sparse: dst = a*dst + b*src at common indices
 void mju_combineSparseInc(mjtNum* dst, const mjtNum* src, int n, mjtNum a, mjtNum b,
                           int dst_nnz, int src_nnz, const int* dst_ind, const int* src_ind);
@@ -163,6 +158,8 @@ mjtNum mju_dotSparse(const mjtNum* vec1, const mjtNum* vec2, int nnz1, const int
 #endif  // mjUSEAVX
 }
 
+
+
 // return 1 if vec1==vec2, 0 otherwise
 static inline
 int mju_compare(const int* vec1, const int* vec2, int n) {
@@ -174,6 +171,7 @@ int mju_compare(const int* vec1, const int* vec2, int n) {
 }
 
 
+
 // merge unique sorted integers, merge array must be large enough (not checked for)
 static inline
 int mj_mergeSorted(int* merge, const int* chain1, int n1, const int* chain2, int n2) {
@@ -227,6 +225,83 @@ int mj_mergeSorted(int* merge, const int* chain1, int n1, const int* chain2, int
 }
 
 
+
+// res = res*scl1 + vec*scl2
+static inline
+void mju_addToSclScl(mjtNum* res, const mjtNum* vec, mjtNum scl1, mjtNum scl2, int n) {
+#ifdef mjUSEAVX
+  mju_addToSclScl_avx(res, vec, scl1, scl2, n);
+#else
+  for (int i=0; i < n; i++) {
+    res[i] = res[i]*scl1 + vec[i]*scl2;
+  }
+#endif  // mjUSEAVX
+}
+
+
+
+// combine two sparse vectors: dst = a*dst + b*src, return nnz of result
+static inline
+int mju_combineSparse(mjtNum* dst, const mjtNum* src, mjtNum a, mjtNum b,
+                      int dst_nnz, int src_nnz, int* dst_ind, const int* src_ind,
+                      mjtNum* buf, int* buf_ind) {
+  // check for identical pattern
+  if (dst_nnz == src_nnz) {
+    if (mju_compare(dst_ind, src_ind, dst_nnz)) {
+      // combine mjtNum data directly
+      mju_addToSclScl(dst, src, a, b, dst_nnz);
+      return dst_nnz;
+    }
+  }
+
+  // copy dst into buf
+  if (dst_nnz) {
+    memcpy(buf, dst, dst_nnz * sizeof(mjtNum));
+    memcpy(buf_ind, dst_ind, dst_nnz * sizeof(int));
+  }
+
+  // prepare to merge buf and src into dst
+  int bi = 0, si = 0, nnz = 0;
+  int buf_nnz = dst_nnz;
+
+  // merge vectors
+  while (bi < buf_nnz && si < src_nnz) {
+    int badr = buf_ind[bi];
+    int sadr = src_ind[si];
+
+    if (badr == sadr) {
+      dst[nnz] = a*buf[bi++] + b*src[si++];
+      dst_ind[nnz++] = badr;
+    }
+
+    // buf only
+    else if (badr < sadr) {
+      dst[nnz] = a*buf[bi++];
+      dst_ind[nnz++] = badr;
+    }
+
+    // src only
+    else {
+      dst[nnz] = b*src[si++];
+      dst_ind[nnz++] = sadr;
+    }
+  }
+
+  // the rest of src only
+  while (si < src_nnz) {
+    dst[nnz] = b*src[si];
+    dst_ind[nnz++] = src_ind[si++];
+  }
+
+  // the rest of buf only
+  while (bi < buf_nnz) {
+    dst[nnz] = a*buf[bi];
+    dst_ind[nnz++] = buf_ind[bi++];
+  }
+
+  return nnz;
+}
+
 #ifdef __cplusplus
 }
 #endif