SparseGrid and Transfers for MPM

multibody/mpm/BUILD.bazel

@@ -0,0 +1,165 @@
+#- * - python - * -
+
+load("@drake//tools/skylark:cc.bzl", "cc_library", "cc_binary")
+load("//tools/lint:lint.bzl", "add_lint_tests")
+load(
+ "//tools/skylark:drake_cc.bzl",
+ "drake_cc_googletest",
+ "drake_cc_library",
+ "drake_cc_package_library",
+)
+
+load(
+ "//tools/performance:defs.bzl",
+ "drake_cc_googlebench_binary",
+ "drake_py_experiment_binary",
+)
+
+package(
+ default_visibility = ["//visibility:public"],
+)
+
+cc_library(
+ name = "sparse_grid",
+ srcs = [
+ "sparse_grid.cc",
+ ],
+ # TODO(xuchenhan-tri): Remove the hard-coded flag when we build this as a drake_cc_library.
+ copts = ["-fopenmp", "-march=native"],
+ hdrs = [
+ "sparse_grid.h",
+ ],
+ linkopts = ["-fopenmp"],
+ deps = [
+ ":math",
+ "//common:essential",
+ "//common:parallelism",
+ "//tools/workspace/spgrid",
+ "//tools/workspace/ips2ra_internal",
+ ],
+)
+
+drake_cc_library(
+ name = "particles",
+ hdrs = [
+ "particles.h",
+ ],
+ deps = [
+ "//common:essential",
+ ],
+)
+
+drake_cc_library(
+ name = "math",
+ hdrs = [
+ "math.h",
+ ],
+ deps = [
+ ":simd_scalar",
+ "//common:essential",
+ ],
+)
+
+cc_library(
+ name = "transfer",
+ srcs = [
+ "transfer.cc",
+ ],
+ hdrs = [
+ "transfer.h",
+ ],
+ # TODO(xuchenhan-tri): Remove the hard-coded flag when we build this as a drake_cc_library.
+ copts = ["-fopenmp", "-march=native"],
+ linkopts = ["-fopenmp"],
+ deps = [
+ ":particles",
+ ":sparse_grid",
+ ":simd_scalar",
+ "//common:essential",
+ "//common:parallelism",
+ ],
+)
+
+# This setting governs when we'll compile with Intel AVX2 and FMA enabled.
+# Compiling for Broadwell (or later) gets those instructions.
+#
+# Note that we have runtime detection of CPU support; this flag only affects
+# what happens at build time, i.e., will the compiler support it.
+config_setting(
+ name = "build_avx2_fma",
+ constraint_values = [
+ "@platforms//cpu:x86_64",
+ # On macOS, we opt-out of this feature (even for Apple hardware that
+ # supports it) to reduce our test matrix burden for the deprecated
+ # architecture.
+ "@platforms//os:linux",
+ ],
+)
+
+drake_cc_library(
+ name = "simd_scalar",
+ srcs = ["simd_scalar.cc", "simd_operations.cc"],
+ hdrs = ["simd_scalar.h", "simd_operations.h", "eigen_specializations.h"],
+ copts = select({
+ ":build_avx2_fma": ["-march=broadwell"],
+ "//conditions:default": [],
+ }),
+ deps = [
+ "@highway_internal//:hwy",
+ "//common:essential",
+ ],
+)
+
+drake_cc_googletest(
+ name = "sparse_grid_test",
+ deps = [
+ ":sparse_grid",
+ "//common/test_utilities:eigen_matrix_compare",
+ ],
+)
+
+drake_cc_googletest(
+ name = "math_test",
+ deps = [
+ ":math",
+ ],
+)
+
+drake_cc_googletest(
+ name = "simd_scalar_test",
+ copts = ["-march=native"],
+ deps = [
+ ":simd_scalar",
+ ],
+)
+
+drake_cc_googletest(
+ name = "transfer_test",
+ srcs = [
+ "test/transfer_test.cc",
+ ],
+ copts = ["-march=native"],
+ deps = [
+ ":transfer",
+ "//common/test_utilities:eigen_matrix_compare",
+ ],
+)
+
+drake_cc_googlebench_binary(
+ name = "transfer_benchmark",
+ srcs = ["transfer_benchmark.cc"],
+ add_test_rule = True,
+ deps = [
+ ":transfer",
+ "//common:add_text_logging_gflags",
+ "//tools/performance:fixture_common",
+ "//tools/performance:gflags_main",
+ ],
+)
+
+drake_py_experiment_binary(
+ name = "transfer_experiment",
+ googlebench_binary = ":transfer_benchmark",
+)
+
+# add_lint_tests()

multibody/mpm/eigen_specializations.h

@@ -0,0 +1,77 @@
+#pragma once
+
+#include <limits>
+
+#include <Eigen/Core>
+
+/* This file contains Eigen-related specializations for Drake's SimdScalar
+scalar type (plus the intimately related std::numeric_limits specialization).
+
+NOTE: This file should never be included directly, rather only from
+simd_scalar.h in a very specific order. */
+
+#ifndef DRAKE_DOXYGEN_CXX
+
+namespace std {
+template <typename T>
+class numeric_limits<drake::multibody::mpm::internal::SimdScalar<T>>
+ : public numeric_limits<double> {};
+} // namespace std
+
+namespace Eigen {
+
+// === See Eigen/src/Core/NumTraits.h ===
+
+// See https://eigen.tuxfamily.org/dox/structEigen_1_1NumTraits.html.
+// We'll Inherit the constants from `T`, but be sure to fatten a few types
+// up to full SimdScalar where necessary.
+// @tparam T double or float.
+template <typename T>
+struct NumTraits<drake::multibody::mpm::internal::SimdScalar<T>>
+ : public NumTraits<T> {
+ // This refers to the "real part" of a complex number (e.g., std::complex).
+ // Because we're not a complex number, it's just the same type as ourselves.
+ using Real = drake::multibody::mpm::internal::SimdScalar<T>;
+
+ // This promotes integer types during operations like quotients, square roots,
+ // etc. We're already floating-point, so it's just the same type as ourselves.
+ using NonInteger = drake::multibody::mpm::internal::SimdScalar<T>;
+
+ // Eigen says "If you don't know what this means, just use [your type] here."
+ using Nested = drake::multibody::mpm::internal::SimdScalar<T>;
+
+ // Our constructor is required during matrix storage initialization.
+ enum { RequireInitialization = 1 };
+};
+
+// Computing "SimdScalar [op] double" yields an SimdScalar.
+template <typename BinOp, typename T>
+struct ScalarBinaryOpTraits<drake::multibody::mpm::internal::SimdScalar<T>,
+ double, BinOp> {
+ using ReturnType = drake::multibody::mpm::internal::SimdScalar<T>;
+};
+
+// Computing "double [op] SimdScalar" yields an SimdScalar.
+template <typename BinOp, typename T>
+struct ScalarBinaryOpTraits<
+ double, drake::multibody::mpm::internal::SimdScalar<T>, BinOp> {
+ using ReturnType = drake::multibody::mpm::internal::SimdScalar<T>;
+};
+
+// Computing "SimdScalar [op] float" yields an SimdScalar.
+template <typename BinOp, typename T>
+struct ScalarBinaryOpTraits<drake::multibody::mpm::internal::SimdScalar<T>,
+ float, BinOp> {
+ using ReturnType = drake::multibody::mpm::internal::SimdScalar<T>;
+};
+
+// Computing "float [op] SimdScalar" yields an SimdScalar.
+template <typename BinOp, typename T>
+struct ScalarBinaryOpTraits<
+ float, drake::multibody::mpm::internal::SimdScalar<T>, BinOp> {
+ using ReturnType = drake::multibody::mpm::internal::SimdScalar<T>;
+};
+
+} // namespace Eigen
+
+#endif // DRAKE_DOXYGEN_CXX

multibody/mpm/math.h

@@ -0,0 +1,117 @@
+#pragma once
+
+#include <array>
+
+#include "drake/common/eigen_types.h"
+#include "drake/multibody/mpm/simd_scalar.h"
+
+namespace drake {
+namespace multibody {
+namespace mpm {
+namespace internal {
+
+/* Computes A: ε where ε is the Levi-Civita tensor. */
+template <typename T>
+Vector3<T> ContractWithLeviCivita(const Matrix3<T>& A) {
+ Vector3<T> A_dot_eps = {0.0, 0.0, 0.0};
+ A_dot_eps(0) = A(1, 2) - A(2, 1);
+ A_dot_eps(1) = A(2, 0) - A(0, 2);
+ A_dot_eps(2) = A(0, 1) - A(1, 0);
+ return A_dot_eps;
+}
+
+// TODO(xuchenhan-tri): Move this into a separate file.
+/* MassAndMomentum stores the mass, linear momentum, and angular momentum of an
+ object, represented by a collection of particles or a grid. */
+template <typename T>
+struct MassAndMomentum {
+ T mass{0.0};
+ Vector3<T> linear_momentum{Vector3<T>::Zero()};
+ Vector3<T> angular_momentum{Vector3<T>::Zero()};
+};
+
+/* Computes the the 1D base node of a point x in reference space.
+ @tparam double or float. */
+template <typename T>
+int ComputeBaseNode(const T& x) {
+ return static_cast<int>(std::floor(x + static_cast<T>(0.5)));
+}
+
+/* Computes the the 1D base node of a few points in reference space.
+ @param x A vector of points in reference space.
+ @pre all points in `x` are have the same base node.
+ @tparam double or float. */
+template <typename T>
+inline int ComputeBaseNode(const SimdScalar<T>& x) {
+ T x0 = x.get_lane();
+ return static_cast<int>(std::floor(x0 + static_cast<T>(0.5)));
+}
+
+/* Computes the the 3D base node of a point x reference space.
+ @tparam double, float, SimdScalar<double>, or SimdScalar<float> */
+template <typename T>
+Vector3<int> ComputeBaseNode(const Vector3<T>& x) {
+ Vector3<int> result;
+ result[0] = ComputeBaseNode(x[0]);
+ result[1] = ComputeBaseNode(x[1]);
+ result[2] = ComputeBaseNode(x[2]);
+ return result;
+}
+
+/* BsplineWeights computes and stores the weights between a particle and the
+grid nodes in its support in the cartesian grid with grid spacing dx (meter).
+The weights are computed using the quadratic B-spline kernel, and each particle
+has 3x3x3=27 grid nodes in its support and the weights can be queried with the
+`weight()` function.
+@tparam double, float, Simdscalar<double>, or Simdscalar<float>. */
+template <typename T>
+struct BsplineWeights {
+ DRAKE_DEFAULT_COPY_AND_MOVE_AND_ASSIGN(BsplineWeights);
+
+ /* Constructs the BsplineWeights between a particle with world frame positions
+ x and the grid nodes in its support in the cartesian grid with grid spacing
+ dx (meter).
+ @pre if T is a SimdScalar, all particles in x have the same base node. */
+ BsplineWeights(const Vector3<T>& x, const T& dx) {
+ const Vector3<T> x_reference = x / dx;
+ const Vector3<int> base_node = ComputeBaseNode(x_reference);
+ for (int d = 0; d < 3; ++d) {
+ data_[d] = Compute(x_reference[d], base_node[d]);
+ }
+ }
+
+ /* Returns between the particle and the ijk-th grid node in its support.
+ @pre 0 <= i, j, k < 3. */
+ T weight(int i, int j, int k) const {
+ return data_[0](i) * data_[1](j) * data_[2](k);
+ }
+
+ private:
+ /* Computes the weights in a single dimension.
+ @param[in] x_reference The particle position in the reference frame.
+ @param[in] base_node The base grid node position in reference frame.
+ @returns the weights between the particle and the three grid nodes in its
+ support in a single dimension. */
+ Vector3<T> Compute(const T& x_reference, int base_node) const {
+ Vector3<T> result;
+ const T d0 = x_reference - static_cast<T>(base_node);
+ const T z = static_cast<T>(0.5) - d0;
+ const T z2 = z * z;
+ result(0) = static_cast<T>(0.5) * z2;
+ result(1) = static_cast<T>(0.75) - d0 * d0;
+ const T d1 = static_cast<T>(1.0) - d0;
+ const T zz = static_cast<T>(1.5) - d1;
+ const T zz2 = zz * zz;
+ result(2) = static_cast<T>(0.5) * zz2;
+
+ return result;
+ }
+
+ static constexpr int kDim = 3;
+ std::array<Vector3<T>, kDim> data_;
+};
+
+} // namespace internal
+} // namespace mpm
+} // namespace multibody
+} // namespace drake