diff attribute

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 <typename T, int Size, bool WithHessian, typename HandleT>
+class DiffAttribute : public Attribute<Scalar<T, Size, WithHessian>, HandleT>
+{
+   public:
+    using PassiveType = T;
+    using ScalarType  = Scalar<T, Size, WithHessian>;
+
+
+    DiffAttribute() : Attribute<ScalarType, HandleT>()
+    {
+    }
+
+    explicit DiffAttribute(const char*         name,
+                           uint32_t            num_attributes,  // not used
+                           locationT           location,
+                           layoutT             layout,  // not used
+                           const RXMeshStatic* rxmesh)
+        : Attribute<ScalarType, HandleT>(name,
+                                         num_attributes,
+                                         location,
+                                         layout,
+                                         rxmesh)
+    {
+    }
+
+   private:
+};
+
+template <typename T, int Size, bool WithHessian>
+using DiffVertexAttribute = DiffAttribute<T, Size, WithHessian, VertexHandle>;
+
+template <typename T, int Size, bool WithHessian>
+using DiffEdgeAttribute = DiffAttribute<T, Size, WithHessian, EdgeHandle>;
+
+template <typename T, int Size, bool WithHessian>
+using DiffFaceAttribute = DiffAttribute<T, Size, WithHessian, FaceHandle>;
+
+}  // namespace rxmesh

include/rxmesh/diff/scalar.h

@@ -16,18 +16,14 @@
 
 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
  * vector. k: Size of variable vector at compile time. PassiveT: Internal
  * floating point type, e.g. double. WithHessian: Set to false for
  * gradient-only mode.
  */
-template <int k, typename PassiveT, bool WithHessian = true>
+template <typename PassiveT, int k, bool WithHessian = true>
 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<PassiveT, k, 1>;
-    using HessType = typename std::conditional_t<WithHessian,
+    using PassiveType = PassiveT;
+    using GradType    = Eigen::Matrix<PassiveT, k, 1>;
+    using HessType    = typename std::conditional_t<WithHessian,
                                                  Eigen::Matrix<PassiveT, k, k>,
                                                  Eigen::Matrix<PassiveT, 0, 0>>;
 
@@ -1142,16 +1139,16 @@ __host__ __device__ PassiveT atan2(const PassiveT& _y, const PassiveT& _x)
 
 template <int k, typename PassiveT, bool WithHessian>
 __host__ __device__ PassiveT
-to_passive(const Scalar<k, PassiveT, WithHessian>& a)
+to_passive(const Scalar<PassiveT, k, WithHessian>& a)
 {
     return a.val;
 }
 
-template <int k, int rows, int cols, typename ScalarT, bool WithHessian>
-__host__ __device__ Eigen::Matrix<ScalarT, rows, cols> to_passive(
-    const Eigen::Matrix<Scalar<k, ScalarT, WithHessian>, rows, cols>& A)
+template <int k, int rows, int cols, typename PassiveT, bool WithHessian>
+__host__ __device__ Eigen::Matrix<PassiveT, rows, cols> to_passive(
+    const Eigen::Matrix<Scalar<PassiveT, k, WithHessian>, rows, cols>& A)
 {
-    Eigen::Matrix<ScalarT, rows, cols> A_passive(A.rows(), A.cols());
+    Eigen::Matrix<PassiveT, rows, cols> 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<ScalarT, rows, cols> to_passive(
 }
 
 // ///////////////////////////////////////////////////////////////////////////
-// TinyAD::Scalar typedefs
+// Scalar typedefs
 // ///////////////////////////////////////////////////////////////////////////
 
 template <int k, bool WithHessian = true>
-using Float = Scalar<k, float, WithHessian>;
+using Float = Scalar<float, k, WithHessian>;
 template <int k, bool WithHessian = true>
-using Double = Scalar<k, double, WithHessian>;
+using Double = Scalar<double, k, WithHessian>;
 template <int k, bool WithHessian = true>
-using LongDouble = Scalar<k, long double, WithHessian>;
+using LongDouble = Scalar<long double, k, WithHessian>;
 
 }  // namespace rxmesh
 
@@ -1183,11 +1180,11 @@ namespace Eigen {
  * and https://eigen.tuxfamily.org/dox/structEigen_1_1NumTraits.html
  */
 template <int k, typename PassiveT, bool WithHessian>
-struct NumTraits<rxmesh::Scalar<k, PassiveT, WithHessian>> : NumTraits<PassiveT>
+struct NumTraits<rxmesh::Scalar<PassiveT, k, WithHessian>> : NumTraits<PassiveT>
 {
-    typedef rxmesh::Scalar<k, PassiveT, WithHessian> Real;
-    typedef rxmesh::Scalar<k, PassiveT, WithHessian> NonInteger;
-    typedef rxmesh::Scalar<k, PassiveT, WithHessian> Nested;
+    typedef rxmesh::Scalar<PassiveT, k, WithHessian> Real;
+    typedef rxmesh::Scalar<PassiveT, k, WithHessian> NonInteger;
+    typedef rxmesh::Scalar<PassiveT, k, WithHessian> Nested;
 
     enum
     {
@@ -1206,19 +1203,19 @@ struct NumTraits<rxmesh::Scalar<k, PassiveT, WithHessian>> : NumTraits<PassiveT>
  * allowed, and that the return type is rxmesh::Scalar.
  */
 template <typename BinaryOp, int k, typename PassiveT, bool WithHessian>
-struct ScalarBinaryOpTraits<rxmesh::Scalar<k, PassiveT, WithHessian>,
+struct ScalarBinaryOpTraits<rxmesh::Scalar<PassiveT, k, WithHessian>,
                             PassiveT,
                             BinaryOp>
 {
-    typedef rxmesh::Scalar<k, PassiveT, WithHessian> ReturnType;
+    typedef rxmesh::Scalar<PassiveT, k, WithHessian> ReturnType;
 };
 
 template <typename BinaryOp, int k, typename PassiveT, bool WithHessian>
 struct ScalarBinaryOpTraits<PassiveT,
-                            rxmesh::Scalar<k, PassiveT, WithHessian>,
+                            rxmesh::Scalar<PassiveT, k, WithHessian>,
                             BinaryOp>
 {
-    typedef rxmesh::Scalar<k, PassiveT, WithHessian> ReturnType;
+    typedef rxmesh::Scalar<PassiveT, k, WithHessian> ReturnType;
 };
 
 }  // namespace Eigen

include/rxmesh/diff/util.h

@@ -11,7 +11,7 @@
 
 namespace rxmesh {
 
-template <int k, typename PassiveT, bool WithHessian>
+template <typename PassiveT, int k, bool WithHessian>
 struct Scalar;
 
 
@@ -97,7 +97,7 @@ __host__ __device__ __inline__ bool is_sym(const T& A, S eps = double(1e-6))
  */
 template <int k, typename PassiveT, bool WithHessian>
 __device__ __host__ __inline__ bool is_finite_scalar(
-    const Scalar<k, PassiveT, WithHessian>& s)
+    const Scalar<PassiveT, k, WithHessian>& s)
 {
     if (!is_finite(s.val) || !is_finite_mat(s.grad) || !is_finite_mat(s.Hess)) {
         return false;

include/rxmesh/rxmesh_static.h

@@ -7,6 +7,7 @@
 #include <cuda_profiler_api.h>
 
 #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 <class T, int Size, bool WithHessian>
+    std::shared_ptr<DiffFaceAttribute<T, Size, WithHessian>>
+    add_diff_face_attribute(const std::string& name,
+                            locationT          location = LOCATION_ALL)
+    {
+        return m_attr_container
+            ->template add<DiffFaceAttribute<T, Size, WithHessian>>(
+                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 <class T, int Size, bool WithHessian>
+    std::shared_ptr<DiffEdgeAttribute<T, Size, WithHessian>>
+    add_diff_edge_attribute(const std::string& name,
+                            locationT          location = LOCATION_ALL)
+    {
+        return m_attr_container
+            ->template add<DiffEdgeAttribute<T, Size, WithHessian>>(
+                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 <class T, int Size, bool WithHessian>
+    std::shared_ptr<DiffVertexAttribute<T, Size, WithHessian>>
+    add_diff_vertex_attribute(const std::string& name,
+                              locationT          location = LOCATION_ALL)
+    {
+        return m_attr_container
+            ->template add<DiffVertexAttribute<T, Size, WithHessian>>(
+                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

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 

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 <typename T, int Size, bool WithHessian>
+void populate(rxmesh::RXMeshStatic&                              rx,
+              rxmesh::DiffVertexAttribute<T, Size, WithHessian>& v,
+              rxmesh::Scalar<T, Size, WithHessian>               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<float, 1, true>("v");
+
+    rx.for_each_vertex(HOST, [&](const VertexHandle& vh) { v_attr(vh) = 1; });
+
+    auto val = Scalar<float, 1, true>::known_derivatives(2.0, 3.0, 4.0);
+
+    populate<float>(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
+}