Add internal functions for 3x3 matmul.

PiperOrigin-RevId: 616749614
Change-Id: I7f3ac2ffc826076b6c97424324f03ecd9ccd4655

src/engine/engine_collision_box.c

@@ -638,7 +638,7 @@ int mjc_BoxBox(const mjModel* M, const mjData* D, mjContact* con, int g1, int g2
   mju_sub3(tmp1, pos1, pos2);
   mju_rotVecMatT(pos12, tmp1, mat2);
 
-  mju_mulMatTMat(rot, mat1, mat2, 3, 3, 3);
+  mju_mulMatTMat3(rot, mat1, mat2);
   mju_transpose(rott, rot, 3, 3);
 
   for (i = 0; i < 9; i++)
@@ -804,7 +804,7 @@ int mjc_BoxBox(const mjModel* M, const mjData* D, mjContact* con, int g1, int g2
 }
 
   if (q2) {
-    mju_mulMatMatT(r, rotmore, rot, 3, 3, 3);
+    mju_mulMatMatT3(r, rotmore, rot);
 
     // mju_rotVecMat(p,pos12,rotmore);
     // mju_rotVecMat(tmp1,size2,rotmore);
@@ -956,7 +956,7 @@ int mjc_BoxBox(const mjModel* M, const mjData* D, mjContact* con, int g1, int g2
   }
 
 
-  mju_mulMatMatT(r, q2 ? mat2 : mat1, rotmore, 3, 3, 3);
+  mju_mulMatMatT3(r, q2 ? mat2 : mat1, rotmore);
   mju_copy3(p, q2 ? pos2 : pos1);
 
   tmp2[0] = (q2 ? -1 : 1) * r[2];
@@ -1319,7 +1319,7 @@ edgeedge:
     }
   }
 
-  mju_mulMatMatT(r, mat1, rotmore, 3, 3, 3);
+  mju_mulMatMatT3(r, mat1, rotmore);
 
   mju_rotVecMat(tmp1, rnorm, r);
 

src/engine/engine_collision_convex.c

@@ -318,7 +318,7 @@ static void mju_rotateFrame(const mjtNum origin[3], const mjtNum rot[9],
   mjtNum mat[9], vec[3], rel[3];
 
   // rotate frame: xmat = rot*xmat
-  mju_mulMatMat(mat, rot, xmat, 3, 3, 3);
+  mju_mulMatMat3(mat, rot, xmat);
   mju_copy(xmat, mat, 9);
 
   // vector to rotation origin: rel = origin - xpos
@@ -682,7 +682,7 @@ int mjc_ConvexHField(const mjModel* m, const mjData* d,
   }
 
   // express geom2 mat in heightfield frame
-  mju_mulMatTMat(mat, mat1, mat2, 3, 3, 3);
+  mju_mulMatTMat3(mat, mat1, mat2);
 
   //------------------------------------- AABB computation, box-box test
 

src/engine/engine_derivative.c

@@ -251,7 +251,7 @@ void mjd_subQuat(const mjtNum qa[4], const mjtNum qb[4], mjtNum Da[9], mjtNum Db
 
   // add term linear in K * K
   mjtNum KK[9];
-  mju_mulMatMat(KK, K, K, 3, 3, 3);
+  mju_mulMatMat3(KK, K, K);
   mjtNum coef = 1.0 - (half_angle < 6e-8 ? 1.0 : half_angle / mju_tan(half_angle));
   mju_addToScl(Da_tmp, KK, coef, 9);
 

src/engine/engine_support.c

@@ -815,8 +815,8 @@ void mj_angmomMat(const mjModel* m, mjData* d, mjtNum* mat, int body) {
 
     // term1 = body angular momentum about self COM in world frame
     mjtNum tmp1[9], tmp2[9];
-    mju_mulMatMat(tmp1, ximat, inertia, 3, 3, 3);  // tmp1  = ximat * inertia
-    mju_mulMatMatT(tmp2, tmp1, ximat, 3, 3, 3);    // tmp2  = ximat * inertia * ximat^T
+    mju_mulMatMat3(tmp1, ximat, inertia);          // tmp1  = ximat * inertia
+    mju_mulMatMatT3(tmp2, tmp1, ximat);            // tmp2  = ximat * inertia * ximat^T
     mju_mulMatMat(term1, tmp2, jacr, 3, 3, nv);    // term1 = ximat * inertia * ximat^T * jacr
 
     // location of body COM w.r.t subtree COM

src/engine/engine_util_blas.c

@@ -180,6 +180,51 @@ void mju_rotVecMatT(mjtNum res[3], const mjtNum vec[3], const mjtNum mat[9]) {
 
 
 
+// multiply 3x3 matrices,
+void mju_mulMatMat3(mjtNum res[9], const mjtNum a[3], const mjtNum b[9]) {
+  res[0] = a[0]*b[0] + a[1]*b[3] + a[2]*b[6];
+  res[1] = a[0]*b[1] + a[1]*b[4] + a[2]*b[7];
+  res[2] = a[0]*b[2] + a[1]*b[5] + a[2]*b[8];
+  res[3] = a[3]*b[0] + a[4]*b[3] + a[5]*b[6];
+  res[4] = a[3]*b[1] + a[4]*b[4] + a[5]*b[7];
+  res[5] = a[3]*b[2] + a[4]*b[5] + a[5]*b[8];
+  res[6] = a[6]*b[0] + a[7]*b[3] + a[8]*b[6];
+  res[7] = a[6]*b[1] + a[7]*b[4] + a[8]*b[7];
+  res[8] = a[6]*b[2] + a[7]*b[5] + a[8]*b[8];
+}
+
+
+
+// multiply 3x3 matrices, first argument transposed
+void mju_mulMatTMat3(mjtNum res[9], const mjtNum a[3], const mjtNum b[9]) {
+  res[0] = a[0]*b[0] + a[3]*b[3] + a[6]*b[6];
+  res[1] = a[0]*b[1] + a[3]*b[4] + a[6]*b[7];
+  res[2] = a[0]*b[2] + a[3]*b[5] + a[6]*b[8];
+  res[3] = a[1]*b[0] + a[4]*b[3] + a[7]*b[6];
+  res[4] = a[1]*b[1] + a[4]*b[4] + a[7]*b[7];
+  res[5] = a[1]*b[2] + a[4]*b[5] + a[7]*b[8];
+  res[6] = a[2]*b[0] + a[5]*b[3] + a[8]*b[6];
+  res[7] = a[2]*b[1] + a[5]*b[4] + a[8]*b[7];
+  res[8] = a[2]*b[2] + a[5]*b[5] + a[8]*b[8];
+}
+
+
+
+// multiply 3x3 matrices, second argument transposed
+void mju_mulMatMatT3(mjtNum res[9], const mjtNum a[3], const mjtNum b[9]) {
+  res[0] = a[0]*b[0] + a[1]*b[1] + a[2]*b[2];
+  res[1] = a[0]*b[3] + a[1]*b[4] + a[2]*b[5];
+  res[2] = a[0]*b[6] + a[1]*b[7] + a[2]*b[8];
+  res[3] = a[3]*b[0] + a[4]*b[1] + a[5]*b[2];
+  res[4] = a[3]*b[3] + a[4]*b[4] + a[5]*b[5];
+  res[5] = a[3]*b[6] + a[4]*b[7] + a[5]*b[8];
+  res[6] = a[6]*b[0] + a[7]*b[1] + a[8]*b[2];
+  res[7] = a[6]*b[3] + a[7]*b[4] + a[8]*b[5];
+  res[8] = a[6]*b[6] + a[7]*b[7] + a[8]*b[8];
+}
+
+
+
 //------------------------------ 4D vector and matrix-vector operations ----------------------------
 
 // res = 0

src/engine/engine_util_blas.h

@@ -109,6 +109,14 @@ MJAPI void mju_rotVecMat(mjtNum res[3], const mjtNum vec[3], const mjtNum mat[9]
 // multiply vector by transposed 3D rotation matrix
 MJAPI void mju_rotVecMatT(mjtNum res[3], const mjtNum vec[3], const mjtNum mat[9]);
 
+// multiply 3x3 matrices
+MJAPI void mju_mulMatMat3(mjtNum res[9], const mjtNum mat1[3], const mjtNum mat2[9]);
+
+// multiply 3x3 matrices, first argument transposed
+MJAPI void mju_mulMatTMat3(mjtNum res[9], const mjtNum a[3], const mjtNum b[9]);
+
+// multiply 3x3 matrices, second argument transposed
+MJAPI void mju_mulMatMatT3(mjtNum res[9], const mjtNum mat1[3], const mjtNum mat2[9]);
 
 //------------------------------ 4D/quaternion operations ------------------------------------------
 

src/engine/engine_util_solve.c

@@ -721,8 +721,8 @@ int mju_eig3(mjtNum eigval[3], mjtNum eigvec[9], mjtNum quat[4], const mjtNum ma
   for (iter=0; iter < 500; iter++) {
     // make quaternion matrix eigvec, compute D = eigvec'*mat*eigvec
     mju_quat2Mat(eigvec, quat);
-    mju_mulMatTMat(tmp, eigvec, mat, 3, 3, 3);
-    mju_mulMatMat(D, tmp, eigvec, 3, 3, 3);
+    mju_mulMatTMat3(tmp, eigvec, mat);
+    mju_mulMatMat3(D, tmp, eigvec);
 
     // assign eigenvalues
     eigval[0] = D[0];

test/engine/engine_util_blas_test.cc

@@ -95,5 +95,38 @@ TEST_F(EngineUtilBlasTest, MjuSymmetrize) {
                                3, 3.5, 3));
 }
 
+TEST_F(EngineUtilBlasTest, MjuMulMat3) {
+  const mjtNum mat1[9] = {
+    1, 2, 3,
+    4, 5, 6,
+    7, 8, 9
+  };
+  const mjtNum mat2[9] = {
+    2, 3, 4,
+    5, 6, 7,
+    8, 9, 10
+  };
+  mjtNum res1[9] = {0};
+  mjtNum res2[9] = {0};
+
+  mju_mulMatMat3(res1, mat1, mat2);
+  mju_mulMatMat(res2, mat1, mat2, 3, 3, 3);
+  for (int i = 0; i < 9; ++i) {
+    EXPECT_EQ(res1[i], res2[i]);
+  }
+
+  mju_mulMatTMat3(res1, mat1, mat2);
+  mju_mulMatTMat(res2, mat1, mat2, 3, 3, 3);
+  for (int i = 0; i < 9; ++i) {
+    EXPECT_EQ(res1[i], res2[i]);
+  }
+
+  mju_mulMatMatT3(res1, mat1, mat2);
+  mju_mulMatMatT(res2, mat1, mat2, 3, 3, 3);
+  for (int i = 0; i < 9; ++i) {
+    EXPECT_EQ(res1[i], res2[i]);
+  }
+}
+
 }  // namespace
 }  // namespace mujoco