Adds ReexpressSpatialVector

Author: sherm1

math/fast_pose_composition_functions.h

@@ -1,5 +1,7 @@
 #pragma once
 
+#include "drake/common/eigen_types.h"
+
 namespace drake {
 namespace math {
 
@@ -66,6 +68,18 @@ void ComposeXinvX(const RigidTransform<double>& X_BA,
                   const RigidTransform<double>& X_BC,
                   RigidTransform<double>* X_AC);
 
+// TODO(sherm1) Consider making the signature take SpatialVector instead of
+//  Vector6.
+/* Re-express a spatial vector via `S_A = R_AB * S_B`. A spatial vector is
+a Vector6 composed of two 3-element Vector3's that are re-expressed
+independently. It is OK if S_A is exactly the same memory as S_B, but not if
+they partially overlap or overlap with R_AB.
+This requires 30 floating point operations but can be done very efficiently
+exploing SIMD instructions when available. */
+void ReexpressSpatialVector(const RotationMatrix<double>& R_AB,
+                            const Vector6<double>& V_B,
+                            Vector6<double>* V_A);
+
 }  // namespace internal
 }  // namespace math
 }  // namespace drake

math/fast_pose_composition_functions_avx2_fma.cc

@@ -591,6 +591,62 @@ void ComposeXinvXImpl(const double* X_BA, const double* X_BC, double* X_AC) {
   hn::StoreN(stu_, tag, X_AC + 9, 3);  // 3-wide write to stay in bounds
 }
 
+/* Re-express SpatialVector via V_A = R_AB * V_B.
+
+R_AB is 9 consecutive doubles in column-major order, the vectors are 6
+consecutive elements comprising two independent 3-vectors. It is allowable for
+V_A and V_B to be the same memory.
+
+  R_AB = abcdefghi
+  V_B = xyzrst
+  V_A = XYZRST
+
+We want to perform two matrix-vector products:
+
+   V_A     R_AB    V_B
+
+    X     a d g     x
+    Y  =  b e h  @  y      15 flops
+    Z     c f i     z
+
+    R     a d g     r
+    S  =  b e h  @  s      15 flops
+    T     c f i     t
+
+We can do this in 6 SIMD instructions. We end up doing 40 flops and throwing
+10 of them away.
+*/
+void ReexpressSpatialVectorImpl(const double* R_AB,
+                                const double* V_B,
+                                double* V_A) {
+  const hn::FixedTag<double, 4> tag;
+
+  const auto abc_ = hn::LoadU(tag, R_AB);      // (d is loaded but unused)
+  const auto def_ = hn::LoadU(tag, R_AB + 3);  // (g is loaded but unused)
+  const auto ghi_ = hn::LoadN(tag, R_AB + 6, 3);
+
+  const auto xxx_ = hn::Set(tag, V_B[0]);
+  const auto yyy_ = hn::Set(tag, V_B[1]);
+  const auto zzz_ = hn::Set(tag, V_B[2]);
+
+  // Vector XYZ:                            X   Y   Z   _
+  auto XYZ_ = hn::Mul(abc_, xxx_);      //  ax  bx  cx  _
+  XYZ_ = hn::MulAdd(def_, yyy_, XYZ_);  // +dy +ey +fy  _
+  XYZ_ = hn::MulAdd(ghi_, zzz_, XYZ_);  // +gz +hz +iz  _
+
+  const auto rrr_ = hn::Set(tag, V_B[3]);
+  const auto sss_ = hn::Set(tag, V_B[4]);
+  const auto ttt_ = hn::Set(tag, V_B[5]);
+
+  // Vector RST:                            R   S   T   _
+  auto RST_ = hn::Mul(abc_, rrr_);      //  ar  br  cr  _
+  RST_ = hn::MulAdd(def_, sss_, RST_);  // +ds +es +fs  _
+  RST_ = hn::MulAdd(ghi_, ttt_, RST_);  // +gt +ht +it  _
+
+  hn::StoreU(XYZ_, tag, V_A);         // 4-wide write temporarily overwrites R
+  hn::StoreN(RST_, tag, V_A + 3, 3);  // 3-wide write to stay in bounds
+}
+
 #else  // HWY_MAX_BYTES
 
 /* The portable versions are always defined. They should be written to maximize
@@ -697,6 +753,20 @@ void ComposeXinvXImpl(const double* X_BA, const double* X_BC, double* X_AC) {
   ComposeXinvXNoAlias(X_BA, X_BC, X_AC_temp);
   std::copy(X_AC_temp, X_AC_temp + 12, X_AC);
 }
+void ReexpressSpatialVectorImpl(const double* R_AB,
+                                const double* V_B,
+                                double* V_A) {
+  DRAKE_ASSERT(V_A != nullptr);
+  double x, y, z;  // Protect from overlap with V_B.
+  x = row_x_col(&R_AB[0], &V_B[0]);
+  y = row_x_col(&R_AB[1], &V_B[0]);
+  z = row_x_col(&R_AB[2], &V_B[0]);
+  V_A[0] = x; V_A[1] = y; V_A[2] = z;
+  x = row_x_col(&R_AB[0], &V_B[3]);
+  y = row_x_col(&R_AB[1], &V_B[3]);
+  z = row_x_col(&R_AB[2], &V_B[3]);
+  V_A[3] = x; V_A[4] = y; V_A[5] = z;
+}
 
 #endif  // HWY_MAX_BYTES
 
@@ -732,6 +802,10 @@ HWY_EXPORT(ComposeXinvXImpl);
 struct ChooseBestComposeXinvX {
   auto operator()() { return HWY_DYNAMIC_POINTER(ComposeXinvXImpl); }
 };
+HWY_EXPORT(ReexpressSpatialVectorImpl);
+struct ChooseBestReexpressSpatialVector {
+  auto operator()() { return HWY_DYNAMIC_POINTER(ReexpressSpatialVectorImpl); }
+};
 
 // These sugar functions convert C++ types into bare arrays.
 const double* GetRawData(const RotationMatrix<double>& R) {
@@ -750,6 +824,12 @@ double* GetRawData(RigidTransform<double>* X) {
   // the rotation matrix first followed by the translation.
   return const_cast<double*>(X->rotation().matrix().data());
 }
+const double* GetRawData(const Vector6<double>& V) {
+  return V.data();
+}
+double* GetRawData(Vector6<double>* V) {
+  return V->data();
+}
 
 }  // namespace
 
@@ -781,6 +861,13 @@ void ComposeXinvX(const RigidTransform<double>& X_BA,
       GetRawData(X_BA), GetRawData(X_BC), GetRawData(X_AC));
 }
 
+void ReexpressSpatialVector(const RotationMatrix<double>& R_AB,
+                            const Vector6<double>& V_B,
+                            Vector6<double>* V_A) {
+  LateBoundFunction<ChooseBestReexpressSpatialVector>::Call(
+      GetRawData(R_AB), GetRawData(V_B), GetRawData(V_A));
+}
+
 }  // namespace internal
 }  // namespace math
 }  // namespace drake

math/test/fast_pose_composition_functions_test.cc

@@ -171,6 +171,23 @@ TEST_P(FastPoseCompositionFunctions, ComposeXX) {
   EXPECT_TRUE(CompareMatrices(arg3->GetAsMatrix4(), expected));
 }
 
+GTEST_TEST(FastReexpressSpatialVector, ReexpressSpatialVector) {
+  constexpr double kTolerance = 8 * std::numeric_limits<double>::epsilon();
+  const RotationMatrixd R_AB(RollPitchYawd(1.0, 1.25, 1.5));
+  Vector6d V_B;
+  V_B << 1.0, 2.0, 3.0, 4.0, 5.0, 6.0;
+  Vector6d V_A_expected;
+  V_A_expected.head<3>() = R_AB * V_B.head<3>();
+  V_A_expected.tail<3>() = R_AB * V_B.tail<3>();
+  Vector6d V_A;
+  ReexpressSpatialVector(R_AB, V_B, &V_A);
+  EXPECT_TRUE(CompareMatrices(V_A, V_A_expected, kTolerance));
+
+  // Ensure that it works when V_A and V_B are the same vector.
+  ReexpressSpatialVector(R_AB, V_B, &V_B);
+  EXPECT_TRUE(CompareMatrices(V_B, V_A_expected, kTolerance));
+}
+
 }  // namespace
 }  // namespace internal
 }  // namespace math

multibody/math/spatial_vector.h

@@ -278,7 +278,15 @@ class SpatialVector {
   /// </pre>
   friend SpatialQuantity operator*(
       const math::RotationMatrix<T>& R_FE, const SpatialQuantity& V_E) {
-    return SpatialQuantity(R_FE * V_E.rotational(), R_FE * V_E.translational());
+    if constexpr (std::is_same_v<T, double>) {
+      SpatialQuantity V_F;
+      math::internal::ReexpressSpatialVector(R_FE, V_E.get_coeffs(),
+                                             &V_F.get_coeffs());
+      return V_F;
+    } else {
+      return SpatialQuantity(R_FE * V_E.rotational(),
+                             R_FE * V_E.translational());
+    }
   }
 
   /// Factory to create a _zero_ spatial vector, i.e., a %SpatialVector whose