From 728033a2fec329bb9a16e903caec7579c54f0669 Mon Sep 17 00:00:00 2001 From: ahmed Date: Mon, 23 Sep 2024 20:24:37 -0400 Subject: [PATCH] diff attribute --- include/rxmesh/diff/diff_attribute.h | 47 +++++++++++++++++++ include/rxmesh/diff/scalar.h | 45 +++++++++--------- include/rxmesh/diff/util.h | 4 +- include/rxmesh/rxmesh_static.h | 60 ++++++++++++++++++++++++ tests/RXMesh_test/CMakeLists.txt | 1 + tests/RXMesh_test/test_diff_attribute.cu | 59 +++++++++++++++++++++++ tests/RXMesh_test/test_scalar.cu | 42 ++++++++--------- 7 files changed, 211 insertions(+), 47 deletions(-) create mode 100644 include/rxmesh/diff/diff_attribute.h create mode 100644 tests/RXMesh_test/test_diff_attribute.cu diff --git a/include/rxmesh/diff/diff_attribute.h b/include/rxmesh/diff/diff_attribute.h new file mode 100644 index 00000000..dba8cdc8 --- /dev/null +++ b/include/rxmesh/diff/diff_attribute.h @@ -0,0 +1,47 @@ +#pragma once + +#include "rxmesh/attribute.h" +#include "rxmesh/diff/scalar.h" + +namespace rxmesh { + +class RXMeshStatic; + +template +class DiffAttribute : public Attribute, HandleT> +{ + public: + using PassiveType = T; + using ScalarType = Scalar; + + + DiffAttribute() : Attribute() + { + } + + explicit DiffAttribute(const char* name, + uint32_t num_attributes, // not used + locationT location, + layoutT layout, // not used + const RXMeshStatic* rxmesh) + : Attribute(name, + num_attributes, + location, + layout, + rxmesh) + { + } + + private: +}; + +template +using DiffVertexAttribute = DiffAttribute; + +template +using DiffEdgeAttribute = DiffAttribute; + +template +using DiffFaceAttribute = DiffAttribute; + +} // namespace rxmesh \ No newline at end of file diff --git a/include/rxmesh/diff/scalar.h b/include/rxmesh/diff/scalar.h index f3b46e40..80d42a11 100644 --- a/include/rxmesh/diff/scalar.h +++ b/include/rxmesh/diff/scalar.h @@ -16,10 +16,6 @@ namespace rxmesh { -#define TINYAD_CHECK_FINITE_IF_ENABLED_AD(exp) \ - { \ - } - /** * Forward-differentiable scalar type with constructors for passive and active * variables. Each scalar carries its gradient and Hessian w.r.t. a variable @@ -27,7 +23,7 @@ namespace rxmesh { * floating point type, e.g. double. WithHessian: Set to false for * gradient-only mode. */ -template +template struct Scalar { // Make template arguments available as members @@ -38,8 +34,9 @@ struct Scalar // Determine derivative data types at compile time. Use 0-by-0 if no Hessian // required. - using GradType = Eigen::Matrix; - using HessType = typename std::conditional_t; + using HessType = typename std::conditional_t, Eigen::Matrix>; @@ -1142,16 +1139,16 @@ __host__ __device__ PassiveT atan2(const PassiveT& _y, const PassiveT& _x) template __host__ __device__ PassiveT -to_passive(const Scalar& a) +to_passive(const Scalar& a) { return a.val; } -template -__host__ __device__ Eigen::Matrix to_passive( - const Eigen::Matrix, rows, cols>& A) +template +__host__ __device__ Eigen::Matrix to_passive( + const Eigen::Matrix, rows, cols>& A) { - Eigen::Matrix A_passive(A.rows(), A.cols()); + Eigen::Matrix A_passive(A.rows(), A.cols()); for (Eigen::Index i = 0; i < A.rows(); ++i) { for (Eigen::Index j = 0; j < A.cols(); ++j) A_passive(i, j) = A(i, j).val; @@ -1161,15 +1158,15 @@ __host__ __device__ Eigen::Matrix to_passive( } // /////////////////////////////////////////////////////////////////////////// -// TinyAD::Scalar typedefs +// Scalar typedefs // /////////////////////////////////////////////////////////////////////////// template -using Float = Scalar; +using Float = Scalar; template -using Double = Scalar; +using Double = Scalar; template -using LongDouble = Scalar; +using LongDouble = Scalar; } // namespace rxmesh @@ -1183,11 +1180,11 @@ namespace Eigen { * and https://eigen.tuxfamily.org/dox/structEigen_1_1NumTraits.html */ template -struct NumTraits> : NumTraits +struct NumTraits> : NumTraits { - typedef rxmesh::Scalar Real; - typedef rxmesh::Scalar NonInteger; - typedef rxmesh::Scalar Nested; + typedef rxmesh::Scalar Real; + typedef rxmesh::Scalar NonInteger; + typedef rxmesh::Scalar Nested; enum { @@ -1206,19 +1203,19 @@ struct NumTraits> : NumTraits * allowed, and that the return type is rxmesh::Scalar. */ template -struct ScalarBinaryOpTraits, +struct ScalarBinaryOpTraits, PassiveT, BinaryOp> { - typedef rxmesh::Scalar ReturnType; + typedef rxmesh::Scalar ReturnType; }; template struct ScalarBinaryOpTraits, + rxmesh::Scalar, BinaryOp> { - typedef rxmesh::Scalar ReturnType; + typedef rxmesh::Scalar ReturnType; }; } // namespace Eigen diff --git a/include/rxmesh/diff/util.h b/include/rxmesh/diff/util.h index 796f6d53..4462f3f9 100644 --- a/include/rxmesh/diff/util.h +++ b/include/rxmesh/diff/util.h @@ -11,7 +11,7 @@ namespace rxmesh { -template +template struct Scalar; @@ -97,7 +97,7 @@ __host__ __device__ __inline__ bool is_sym(const T& A, S eps = double(1e-6)) */ template __device__ __host__ __inline__ bool is_finite_scalar( - const Scalar& s) + const Scalar& s) { if (!is_finite(s.val) || !is_finite_mat(s.grad) || !is_finite_mat(s.Hess)) { return false; diff --git a/include/rxmesh/rxmesh_static.h b/include/rxmesh/rxmesh_static.h index f82379f2..f762e81e 100644 --- a/include/rxmesh/rxmesh_static.h +++ b/include/rxmesh/rxmesh_static.h @@ -7,6 +7,7 @@ #include #include "rxmesh/attribute.h" +#include "rxmesh/diff/diff_attribute.h" #include "rxmesh/handle.h" #include "rxmesh/kernels/for_each.cuh" #include "rxmesh/kernels/shmem_allocator.cuh" @@ -846,6 +847,25 @@ class RXMeshStatic : public RXMesh return ret; } + /** + * @brief Adding a new differentiable face attribute + * @tparam T the underlying type of the attribute + * @tparam Size the number of components per face + * @tparam WithHessian if hessian is required + * @param name of the attribute. Should not collide with other attributes + * names + * @param location where to allocate the attributes + */ + template + std::shared_ptr> + add_diff_face_attribute(const std::string& name, + locationT location = LOCATION_ALL) + { + return m_attr_container + ->template add>( + name.c_str(), 1, location, SoA, this); + } + /** * @brief Adding a new edge attribute * @tparam T type of the attribute @@ -868,6 +888,25 @@ class RXMeshStatic : public RXMesh name.c_str(), num_attributes, location, layout, this); } + /** + * @brief Adding a new differentiable edge attribute + * @tparam T the underlying type of the attribute + * @tparam Size the number of components per edge + * @tparam WithHessian if hessian is required + * @param name of the attribute. Should not collide with other attributes + * names + * @param location where to allocate the attributes + */ + template + std::shared_ptr> + add_diff_edge_attribute(const std::string& name, + locationT location = LOCATION_ALL) + { + return m_attr_container + ->template add>( + name.c_str(), 1, location, SoA, this); + } + /** * @brief Adding a new vertex attribute * @tparam T type of the attribute @@ -890,6 +929,27 @@ class RXMeshStatic : public RXMesh name.c_str(), num_attributes, location, layout, this); } + + /** + * @brief Adding a new differentiable vertex attribute + * @tparam T the underlying type of the attribute + * @tparam Size the number of components per vertex, e.g., 3 for vertex + * coordinates + * @tparam WithHessian if hessian is required + * @param name of the attribute. Should not collide with other attributes + * names + * @param location where to allocate the attributes + */ + template + std::shared_ptr> + add_diff_vertex_attribute(const std::string& name, + locationT location = LOCATION_ALL) + { + return m_attr_container + ->template add>( + name.c_str(), 1, location, SoA, this); + } + /** * @brief Adding a new vertex attribute by reading values from a host buffer * v_attributes where the order of vertices is the same as the order of diff --git a/tests/RXMesh_test/CMakeLists.txt b/tests/RXMesh_test/CMakeLists.txt index 4372711b..e90d1f05 100644 --- a/tests/RXMesh_test/CMakeLists.txt +++ b/tests/RXMesh_test/CMakeLists.txt @@ -28,6 +28,7 @@ set( SOURCE_LIST test_export.cu test_svd.cu test_scalar.cu + test_diff_attribute.cu ) target_sources( RXMesh_test diff --git a/tests/RXMesh_test/test_diff_attribute.cu b/tests/RXMesh_test/test_diff_attribute.cu new file mode 100644 index 00000000..4b365132 --- /dev/null +++ b/tests/RXMesh_test/test_diff_attribute.cu @@ -0,0 +1,59 @@ +#include "gtest/gtest.h" + +#include "rxmesh/rxmesh_static.h" + +#include "rxmesh/diff/scalar.h" + + +template +void populate(rxmesh::RXMeshStatic& rx, + rxmesh::DiffVertexAttribute& v, + rxmesh::Scalar val) +{ + rx.for_each_vertex( + rxmesh::DEVICE, + [v, val] __device__(const rxmesh::VertexHandle vh) { v(vh) = val; }); + + EXPECT_EQ(cudaDeviceSynchronize(), cudaSuccess); +} + + +TEST(DiffAttribute, Simple) +{ + // write diff vertex attribute on the device and verify it on the host + using namespace rxmesh; + + RXMeshStatic rx(STRINGIFY(INPUT_DIR) "sphere3.obj"); + + auto v_attr = *rx.add_diff_vertex_attribute("v"); + + rx.for_each_vertex(HOST, [&](const VertexHandle& vh) { v_attr(vh) = 1; }); + + auto val = Scalar::known_derivatives(2.0, 3.0, 4.0); + + populate(rx, v_attr, val); + + v_attr.move(DEVICE, HOST); + + bool is_okay = true; + + rx.for_each_vertex( + HOST, + [&](const VertexHandle& vh) { + if (v_attr(vh).val != val.val || v_attr(vh).grad != val.grad || + v_attr(vh).Hess != val.Hess) { + is_okay = false; + } + }, + NULL, + false); + + EXPECT_TRUE(is_okay); + + EXPECT_EQ(cudaDeviceSynchronize(), cudaSuccess); + + // #if USE_POLYSCOPE + // rx.get_polyscope_mesh()->addVertexScalarQuantity("vAttr", v_attr); + // polyscope::show(); + // #endif +} diff --git a/tests/RXMesh_test/test_scalar.cu b/tests/RXMesh_test/test_scalar.cu index 5dab3cf3..c95d438b 100644 --- a/tests/RXMesh_test/test_scalar.cu +++ b/tests/RXMesh_test/test_scalar.cu @@ -24,7 +24,7 @@ template __inline__ __device__ void test_unary_minus(int* d_err, T eps) { using namespace rxmesh; - using Real1 = Scalar<1, T, WithHessian>; + using Real1 = Scalar; // a(x) = x^2 + x + 2 at x = 1 const Real1 a = Real1::known_derivatives(4.0, 3.0, 2.0); @@ -42,7 +42,7 @@ template __inline__ __device__ void test_sqrt(int* d_err, T eps) { using namespace rxmesh; - using Real1 = Scalar<1, T, WithHessian>; + using Real1 = Scalar; // a(x) = x^2 + x + 2 at x=1 const Real1 a = Real1::known_derivatives(4.0, 3.0, 2.0); @@ -60,7 +60,7 @@ template __inline__ __device__ void test_sqr(int* d_err, T eps) { using namespace rxmesh; - using Real2 = Scalar<2, T, WithHessian>; + using Real2 = Scalar; Real2 x = Real2::Scalar(4.0, 0); Real2 y = Real2::Scalar(6.0, 1); @@ -96,7 +96,7 @@ template __inline__ __device__ void test_fabs(int* d_err, T eps) { using namespace rxmesh; - using Real1 = Scalar<1, T, WithHessian>; + using Real1 = Scalar; { // a(x) = x^3 at x = 1 const Real1 a = Real1::known_derivatives(1.0, 3.0, 6.0); @@ -125,7 +125,7 @@ template __inline__ __device__ void test_abs(int* d_err, T eps) { using namespace rxmesh; - using Real1 = Scalar<1, T, WithHessian>; + using Real1 = Scalar; { // a(x) = x^3 at x = 1 const Real1 a = Real1::known_derivatives(1.0, 3.0, 6.0); @@ -154,7 +154,7 @@ template __inline__ __device__ void test_exp(int* d_err, T eps) { using namespace rxmesh; - using Real1 = Scalar<1, T, WithHessian>; + using Real1 = Scalar; // a(x) = x^2 + x + 2 at x=1 const Real1 a = Real1::known_derivatives(4.0, 3.0, 2.0); @@ -172,7 +172,7 @@ template __inline__ __device__ void test_log(int* d_err, T eps) { using namespace rxmesh; - using Real1 = Scalar<1, T, WithHessian>; + using Real1 = Scalar; // a(x) = x^2 + x + 2 at x=1 const Real1 a = Real1::known_derivatives(4.0, 3.0, 2.0); @@ -190,7 +190,7 @@ template __inline__ __device__ void test_log2(int* d_err, T eps) { using namespace rxmesh; - using Real1 = Scalar<1, T, WithHessian>; + using Real1 = Scalar; // a(x) = x^2 + x + 2 at x=1 const Real1 a = Real1::known_derivatives(4.0, 3.0, 2.0); @@ -208,7 +208,7 @@ template __inline__ __device__ void test_log10(int* d_err, T eps) { using namespace rxmesh; - using Real1 = Scalar<1, T, WithHessian>; + using Real1 = Scalar; // a(x) = x^2 + x + 2 at x=1 const Real1 a = Real1::known_derivatives(4.0, 3.0, 2.0); @@ -226,7 +226,7 @@ template __inline__ __device__ void test_sin(int* d_err, T eps) { using namespace rxmesh; - using Real1 = Scalar<1, T, WithHessian>; + using Real1 = Scalar; // a(x) = x^2 + x + 2 at x=1 const Real1 a = Real1::known_derivatives(4.0, 3.0, 2.0); @@ -245,7 +245,7 @@ template __inline__ __device__ void test_cos(int* d_err, T eps) { using namespace rxmesh; - using Real1 = Scalar<1, T, WithHessian>; + using Real1 = Scalar; // a(x) = x^2 + x + 2 at x=1 const Real1 a = Real1::known_derivatives(4.0, 3.0, 2.0); @@ -264,7 +264,7 @@ template __inline__ __device__ void test_tan(int* d_err, T eps) { using namespace rxmesh; - using Real1 = Scalar<1, T, WithHessian>; + using Real1 = Scalar; // a(x) = x^2 + x + 2 at x=1 const Real1 a = Real1::known_derivatives(4.0, 3.0, 2.0); @@ -285,7 +285,7 @@ template __inline__ __device__ void test_asin(int* d_err, T eps) { using namespace rxmesh; - using Real1 = Scalar<1, T, WithHessian>; + using Real1 = Scalar; // a(x) = x^2 + x - 1.5 at x=1 const Real1 a = Real1::known_derivatives(0.5, 3.0, 2.0); @@ -305,7 +305,7 @@ template __inline__ __device__ void test_acos(int* d_err, T eps) { using namespace rxmesh; - using Real1 = Scalar<1, T, WithHessian>; + using Real1 = Scalar; // a(x) = x^2 + x - 1.5 at x=1 const Real1 a = Real1::known_derivatives(0.5, 3.0, 2.0); @@ -325,7 +325,7 @@ template __inline__ __device__ void test_atan(int* d_err, T eps) { using namespace rxmesh; - using Real1 = Scalar<1, T, WithHessian>; + using Real1 = Scalar; // a(x) = x^2 + x - 1.5 at x=1 const Real1 a = Real1::known_derivatives(0.5, 3.0, 2.0); @@ -344,7 +344,7 @@ template __inline__ __device__ void test_sinh(int* d_err, T eps) { using namespace rxmesh; - using Real1 = Scalar<1, T, WithHessian>; + using Real1 = Scalar; // a(x) = x^2 + x + 2 at x=1 const Real1 a = Real1::known_derivatives(4.0, 3.0, 2.0); @@ -363,7 +363,7 @@ template __inline__ __device__ void test_cosh(int* d_err, T eps) { using namespace rxmesh; - using Real1 = Scalar<1, T, WithHessian>; + using Real1 = Scalar; // a(x) = x^2 + x + 2 at x=1 const Real1 a = Real1::known_derivatives(4.0, 3.0, 2.0); @@ -383,7 +383,7 @@ template __inline__ __device__ void test_tanh(int* d_err, T eps) { using namespace rxmesh; - using Real1 = Scalar<1, T, WithHessian>; + using Real1 = Scalar; // a(x) = x^2 + x + 2 at x=1 const Real1 a = Real1::known_derivatives(4.0, 3.0, 2.0); @@ -403,7 +403,7 @@ template __inline__ __device__ void test_asinh(int* d_err, T eps) { using namespace rxmesh; - using Real1 = Scalar<1, T, WithHessian>; + using Real1 = Scalar; // a(x) = x^2 + x - 1.5 at x=1 const Real1 a = Real1::known_derivatives(0.5, 3.0, 2.0); @@ -423,7 +423,7 @@ template __inline__ __device__ void test_acosh(int* d_err, T eps) { using namespace rxmesh; - using Real1 = Scalar<1, T, WithHessian>; + using Real1 = Scalar; // a(x) = x^2 + x + 2 at x=1 const Real1 a = Real1::known_derivatives(4.0, 3.0, 2.0); @@ -442,7 +442,7 @@ template __inline__ __device__ void test_atanh(int* d_err, T eps) { using namespace rxmesh; - using Real1 = Scalar<1, T, WithHessian>; + using Real1 = Scalar; // a(x) = x^2 + x - 1.5 at x=1 const Real1 a = Real1::known_derivatives(0.5, 3.0, 2.0);