Skip to content

Commit a1754dd

Browse files
committed
Adds ReexpressSpatialVector
1 parent 70c40ab commit a1754dd

File tree

4 files changed

+127
-1
lines changed

4 files changed

+127
-1
lines changed

math/fast_pose_composition_functions.h

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
#pragma once
22

3+
#include "drake/common/eigen_types.h"
4+
35
namespace drake {
46
namespace math {
57

@@ -66,6 +68,18 @@ void ComposeXinvX(const RigidTransform<double>& X_BA,
6668
const RigidTransform<double>& X_BC,
6769
RigidTransform<double>* X_AC);
6870

71+
// TODO(sherm1) Consider making the signature take SpatialVector instead of
72+
// Vector6.
73+
/* Re-express a spatial vector via `S_A = R_AB * S_B`. A spatial vector is
74+
a Vector6 composed of two 3-element Vector3's that are re-expressed
75+
independently. It is OK if S_A is exactly the same memory as S_B, but not if
76+
they partially overlap or overlap with R_AB.
77+
This requires 30 floating point operations but can be done very efficiently
78+
exploing SIMD instructions when available. */
79+
void ReexpressSpatialVector(const RotationMatrix<double>& R_AB,
80+
const Vector6<double>& V_B,
81+
Vector6<double>* V_A);
82+
6983
} // namespace internal
7084
} // namespace math
7185
} // namespace drake

math/fast_pose_composition_functions_avx2_fma.cc

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -591,6 +591,62 @@ void ComposeXinvXImpl(const double* X_BA, const double* X_BC, double* X_AC) {
591591
hn::StoreN(stu_, tag, X_AC + 9, 3); // 3-wide write to stay in bounds
592592
}
593593

594+
/* Re-express SpatialVector via V_A = R_AB * V_B.
595+
596+
R_AB is 9 consecutive doubles in column-major order, the vectors are 6
597+
consecutive elements comprising two independent 3-vectors. It is allowable for
598+
V_A and V_B to be the same memory.
599+
600+
R_AB = abcdefghi
601+
V_B = xyzrst
602+
V_A = XYZRST
603+
604+
We want to perform two matrix-vector products:
605+
606+
V_A R_AB V_B
607+
608+
X a d g x
609+
Y = b e h @ y 15 flops
610+
Z c f i z
611+
612+
R a d g r
613+
S = b e h @ s 15 flops
614+
T c f i t
615+
616+
We can do this in 6 SIMD instructions. We end up doing 40 flops and throwing
617+
10 of them away.
618+
*/
619+
void ReexpressSpatialVectorImpl(const double* R_AB,
620+
const double* V_B,
621+
double* V_A) {
622+
const hn::FixedTag<double, 4> tag;
623+
624+
const auto abc_ = hn::LoadU(tag, R_AB); // (d is loaded but unused)
625+
const auto def_ = hn::LoadU(tag, R_AB + 3); // (g is loaded but unused)
626+
const auto ghi_ = hn::LoadN(tag, R_AB + 6, 3);
627+
628+
const auto xxx_ = hn::Set(tag, V_B[0]);
629+
const auto yyy_ = hn::Set(tag, V_B[1]);
630+
const auto zzz_ = hn::Set(tag, V_B[2]);
631+
632+
// Vector XYZ: X Y Z _
633+
auto XYZ_ = hn::Mul(abc_, xxx_); // ax bx cx _
634+
XYZ_ = hn::MulAdd(def_, yyy_, XYZ_); // +dy +ey +fy _
635+
XYZ_ = hn::MulAdd(ghi_, zzz_, XYZ_); // +gz +hz +iz _
636+
637+
const auto rrr_ = hn::Set(tag, V_B[3]);
638+
const auto sss_ = hn::Set(tag, V_B[4]);
639+
const auto ttt_ = hn::Set(tag, V_B[5]);
640+
641+
// Vector RST: R S T _
642+
auto RST_ = hn::Mul(abc_, rrr_); // ar br cr _
643+
RST_ = hn::MulAdd(def_, sss_, RST_); // +ds +es +fs _
644+
RST_ = hn::MulAdd(ghi_, ttt_, RST_); // +gt +ht +it _
645+
646+
hn::StoreU(XYZ_, tag, V_A); // 4-wide write temporarily overwrites R
647+
hn::StoreN(RST_, tag, V_A + 3, 3); // 3-wide write to stay in bounds
648+
}
649+
594650
#else // HWY_MAX_BYTES
595651

596652
/* 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) {
697753
ComposeXinvXNoAlias(X_BA, X_BC, X_AC_temp);
698754
std::copy(X_AC_temp, X_AC_temp + 12, X_AC);
699755
}
756+
void ReexpressSpatialVectorImpl(const double* R_AB,
757+
const double* V_B,
758+
double* V_A) {
759+
DRAKE_ASSERT(V_A != nullptr);
760+
double x, y, z; // Protect from overlap with V_B.
761+
x = row_x_col(&R_AB[0], &V_B[0]);
762+
y = row_x_col(&R_AB[1], &V_B[0]);
763+
z = row_x_col(&R_AB[2], &V_B[0]);
764+
V_A[0] = x; V_A[1] = y; V_A[2] = z;
765+
x = row_x_col(&R_AB[0], &V_B[3]);
766+
y = row_x_col(&R_AB[1], &V_B[3]);
767+
z = row_x_col(&R_AB[2], &V_B[3]);
768+
V_A[3] = x; V_A[4] = y; V_A[5] = z;
769+
}
700770

701771
#endif // HWY_MAX_BYTES
702772

@@ -732,6 +802,10 @@ HWY_EXPORT(ComposeXinvXImpl);
732802
struct ChooseBestComposeXinvX {
733803
auto operator()() { return HWY_DYNAMIC_POINTER(ComposeXinvXImpl); }
734804
};
805+
HWY_EXPORT(ReexpressSpatialVectorImpl);
806+
struct ChooseBestReexpressSpatialVector {
807+
auto operator()() { return HWY_DYNAMIC_POINTER(ReexpressSpatialVectorImpl); }
808+
};
735809

736810
// These sugar functions convert C++ types into bare arrays.
737811
const double* GetRawData(const RotationMatrix<double>& R) {
@@ -750,6 +824,12 @@ double* GetRawData(RigidTransform<double>* X) {
750824
// the rotation matrix first followed by the translation.
751825
return const_cast<double*>(X->rotation().matrix().data());
752826
}
827+
const double* GetRawData(const Vector6<double>& V) {
828+
return V.data();
829+
}
830+
double* GetRawData(Vector6<double>* V) {
831+
return V->data();
832+
}
753833

754834
} // namespace
755835

@@ -781,6 +861,13 @@ void ComposeXinvX(const RigidTransform<double>& X_BA,
781861
GetRawData(X_BA), GetRawData(X_BC), GetRawData(X_AC));
782862
}
783863

864+
void ReexpressSpatialVector(const RotationMatrix<double>& R_AB,
865+
const Vector6<double>& V_B,
866+
Vector6<double>* V_A) {
867+
LateBoundFunction<ChooseBestReexpressSpatialVector>::Call(
868+
GetRawData(R_AB), GetRawData(V_B), GetRawData(V_A));
869+
}
870+
784871
} // namespace internal
785872
} // namespace math
786873
} // namespace drake

math/test/fast_pose_composition_functions_test.cc

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -171,6 +171,23 @@ TEST_P(FastPoseCompositionFunctions, ComposeXX) {
171171
EXPECT_TRUE(CompareMatrices(arg3->GetAsMatrix4(), expected));
172172
}
173173

174+
GTEST_TEST(FastReexpressSpatialVector, ReexpressSpatialVector) {
175+
constexpr double kTolerance = 8 * std::numeric_limits<double>::epsilon();
176+
const RotationMatrixd R_AB(RollPitchYawd(1.0, 1.25, 1.5));
177+
Vector6d V_B;
178+
V_B << 1.0, 2.0, 3.0, 4.0, 5.0, 6.0;
179+
Vector6d V_A_expected;
180+
V_A_expected.head<3>() = R_AB * V_B.head<3>();
181+
V_A_expected.tail<3>() = R_AB * V_B.tail<3>();
182+
Vector6d V_A;
183+
ReexpressSpatialVector(R_AB, V_B, &V_A);
184+
EXPECT_TRUE(CompareMatrices(V_A, V_A_expected, kTolerance));
185+
186+
// Ensure that it works when V_A and V_B are the same vector.
187+
ReexpressSpatialVector(R_AB, V_B, &V_B);
188+
EXPECT_TRUE(CompareMatrices(V_B, V_A_expected, kTolerance));
189+
}
190+
174191
} // namespace
175192
} // namespace internal
176193
} // namespace math

multibody/math/spatial_vector.h

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -278,7 +278,15 @@ class SpatialVector {
278278
/// </pre>
279279
friend SpatialQuantity operator*(
280280
const math::RotationMatrix<T>& R_FE, const SpatialQuantity& V_E) {
281-
return SpatialQuantity(R_FE * V_E.rotational(), R_FE * V_E.translational());
281+
if constexpr (std::is_same_v<T, double>) {
282+
SpatialQuantity V_F;
283+
math::internal::ReexpressSpatialVector(R_FE, V_E.get_coeffs(),
284+
&V_F.get_coeffs());
285+
return V_F;
286+
} else {
287+
return SpatialQuantity(R_FE * V_E.rotational(),
288+
R_FE * V_E.translational());
289+
}
282290
}
283291

284292
/// Factory to create a _zero_ spatial vector, i.e., a %SpatialVector whose

0 commit comments

Comments
 (0)