Replace ADAMANTINE_HOST_DEV with KOKKOS_FUNCTION

adamantine-sim · Jan 24, 2024 · e7b23ec · e7b23ec
1 parent 0f68aba
commit e7b23ec
Show file tree

Hide file tree

Showing 4 changed files with 60 additions and 77 deletions.
diff --git a/source/MaterialProperty.hh b/source/MaterialProperty.hh
@@ -141,7 +141,7 @@ public:
   /**
    * Compute a material property at a quadrature point for a mix of states.
    */
-  ADAMANTINE_HOST_DEV
+  KOKKOS_FUNCTION
   double compute_material_property(StateProperty state_property,
                                    dealii::types::material_id const material_id,
                                    double const *state_ratios,
@@ -208,7 +208,7 @@ public:
   /**
    * Compute a property from a table given the temperature.
    */
-  static ADAMANTINE_HOST_DEV double compute_property_from_table(
+  static KOKKOS_FUNCTION double compute_property_from_table(
       MemoryBlockView<double, MemorySpaceType> const
           &state_property_tables_view,
       unsigned int const material_id, unsigned int const material_state,

diff --git a/source/MaterialProperty.templates.hh b/source/MaterialProperty.templates.hh
@@ -277,9 +277,7 @@ void MaterialProperty<dim, MemorySpaceType>::update(
 
   bool use_table = _use_table;
   for_each(
-      MemorySpaceType{}, material_ids_size,
-      [=] ADAMANTINE_HOST_DEV(int i)
-      {
+      MemorySpaceType{}, material_ids_size, KOKKOS_LAMBDA(int i) {
         unsigned int constexpr liquid =
             static_cast<unsigned int>(MaterialState::liquid);
         unsigned int constexpr powder =
@@ -546,7 +544,7 @@ MaterialProperty<dim, MemorySpaceType>::compute_material_property(
 }
 
 template <int dim, typename MemorySpaceType>
-ADAMANTINE_HOST_DEV double
+KOKKOS_FUNCTION double
 MaterialProperty<dim, MemorySpaceType>::compute_material_property(
     StateProperty state_property, dealii::types::material_id const material_id,
     double const *state_ratios, double temperature) const
@@ -685,25 +683,24 @@ void MaterialProperty<dim, MemorySpaceType>::set_state_device(
   auto const powder_state = static_cast<unsigned int>(MaterialState::powder);
   auto const liquid_state = static_cast<unsigned int>(MaterialState::liquid);
   auto const solid_state = static_cast<unsigned int>(MaterialState::solid);
-  for_each(MemorySpaceType{}, cell_i,
-           [=] ADAMANTINE_HOST_DEV(int i) mutable
-           {
-             double liquid_ratio_sum = 0.;
-             double powder_ratio_sum = 0.;
-             for (unsigned int q = 0; q < n_q_points; ++q)
-             {
-               liquid_ratio_sum += liquid_ratio_view(mapping_view(i, q));
-               powder_ratio_sum += powder_ratio_view(mapping_view(i, q));
-             }
-             state_view(liquid_state, mp_dof_view(i)) =
-                 liquid_ratio_sum / n_q_points;
-             state_view(powder_state, mp_dof_view(i)) =
-                 powder_ratio_sum / n_q_points;
-             state_view(solid_state, mp_dof_view(i)) =
-                 std::max(1. - state_view(liquid_state, mp_dof_view(i)) -
-                              state_view(powder_state, mp_dof_view(i)),
-                          0.);
-           });
+  for_each(
+      MemorySpaceType{}, cell_i, KOKKOS_LAMBDA(int i) mutable {
+        double liquid_ratio_sum = 0.;
+        double powder_ratio_sum = 0.;
+        for (unsigned int q = 0; q < n_q_points; ++q)
+        {
+          liquid_ratio_sum += liquid_ratio_view(mapping_view(i, q));
+          powder_ratio_sum += powder_ratio_view(mapping_view(i, q));
+        }
+        state_view(liquid_state, mp_dof_view(i)) =
+            liquid_ratio_sum / n_q_points;
+        state_view(powder_state, mp_dof_view(i)) =
+            powder_ratio_sum / n_q_points;
+        state_view(solid_state, mp_dof_view(i)) =
+            std::max(1. - state_view(liquid_state, mp_dof_view(i)) -
+                         state_view(powder_state, mp_dof_view(i)),
+                     0.);
+      });
 }
 #endif
 
@@ -733,9 +730,10 @@ void MaterialProperty<dim, MemorySpaceType>::set_initial_state()
 
   _state.set_zero();
   MemoryBlockView<double, MemorySpaceType> state_view(_state);
-  for_each(MemorySpaceType{}, user_indices.size(),
-           [=] ADAMANTINE_HOST_DEV(int i) mutable
-           { state_view(user_indices_view(i), mp_dofs_view(i)) = 1.; });
+  for_each(
+      MemorySpaceType{}, user_indices.size(), KOKKOS_LAMBDA(int i) mutable {
+        state_view(user_indices_view(i), mp_dofs_view(i)) = 1.;
+      });
 }
 
 template <int dim, typename MemorySpaceType>
@@ -1052,7 +1050,7 @@ MaterialProperty<dim, MemorySpaceType>::compute_average_temperature(
 }
 
 template <int dim, typename MemorySpaceType>
-ADAMANTINE_HOST_DEV double
+KOKKOS_FUNCTION double
 MaterialProperty<dim, MemorySpaceType>::compute_property_from_table(
     MemoryBlockView<double, MemorySpaceType> const &state_property_tables_view,
     unsigned int const material_id, unsigned int const material_state,

diff --git a/source/MemoryBlockView.hh b/source/MemoryBlockView.hh
@@ -25,77 +25,77 @@ public:
   /**
    * Default constructor.
    */
-  ADAMANTINE_HOST_DEV
+  KOKKOS_FUNCTION
   MemoryBlockView() = default;
 
   /**
    * Constructor.
    */
-  ADAMANTINE_HOST_DEV
+  KOKKOS_FUNCTION
   MemoryBlockView(MemoryBlock<Number, MemorySpaceType> const &memory_block);
 
   /**
    * Copy constructor.
    */
-  ADAMANTINE_HOST_DEV MemoryBlockView(
+  KOKKOS_FUNCTION MemoryBlockView(
       MemoryBlockView<Number, MemorySpaceType> const &memory_block_view);
 
   /**
    * Reinitialize the pointer to the underlying data and the size.
    */
-  ADAMANTINE_HOST_DEV
+  KOKKOS_FUNCTION
   void reinit(MemoryBlock<Number, MemorySpaceType> const &memory_block);
 
   /**
    * Assignment operator.
    */
-  ADAMANTINE_HOST_DEV
+  KOKKOS_FUNCTION
   MemoryBlockView<Number, MemorySpaceType> &
   operator=(MemoryBlockView<Number, MemorySpaceType> const &memory_block_view);
 
   /**
    * Access operator for a 1D MemoryBlock.
    */
-  ADAMANTINE_HOST_DEV Number &operator()(unsigned int i) const;
+  KOKKOS_FUNCTION Number &operator()(unsigned int i) const;
 
   /**
    * Access operator for a 2D MemoryBlock.
    */
-  ADAMANTINE_HOST_DEV Number &operator()(unsigned int i, unsigned int j) const;
+  KOKKOS_FUNCTION Number &operator()(unsigned int i, unsigned int j) const;
 
   /**
    * Access operator for a 3D MemoryBlock.
    */
-  ADAMANTINE_HOST_DEV Number &operator()(unsigned int i, unsigned int j,
-                                         unsigned int k) const;
+  KOKKOS_FUNCTION Number &operator()(unsigned int i, unsigned int j,
+                                     unsigned int k) const;
 
   /**
    * Access operator for a 4D MemoryBlock.
    */
-  ADAMANTINE_HOST_DEV Number &operator()(unsigned int i, unsigned int j,
-                                         unsigned int k, unsigned int l) const;
+  KOKKOS_FUNCTION Number &operator()(unsigned int i, unsigned int j,
+                                     unsigned int k, unsigned int l) const;
 
   /**
    * Access operator for a 5D MemoryBlock.
    */
-  ADAMANTINE_HOST_DEV Number &operator()(unsigned int i, unsigned int j,
-                                         unsigned int k, unsigned int l,
-                                         unsigned int m) const;
+  KOKKOS_FUNCTION Number &operator()(unsigned int i, unsigned int j,
+                                     unsigned int k, unsigned int l,
+                                     unsigned int m) const;
 
   /**
    * Return the number of accessible elements.
    */
-  ADAMANTINE_HOST_DEV unsigned int size() const;
+  KOKKOS_FUNCTION unsigned int size() const;
 
   /**
    * Return the @p i dimension.
    */
-  ADAMANTINE_HOST_DEV unsigned int extent(unsigned int i) const;
+  KOKKOS_FUNCTION unsigned int extent(unsigned int i) const;
 
   /**
    * Return the pointer to the underlying data.
    */
-  ADAMANTINE_HOST_DEV Number *data() const;
+  KOKKOS_FUNCTION Number *data() const;
 
 private:
   unsigned int _size = 0;
@@ -108,7 +108,7 @@ private:
 };
 
 template <typename Number, typename MemorySpaceType>
-ADAMANTINE_HOST_DEV MemoryBlockView<Number, MemorySpaceType>::MemoryBlockView(
+KOKKOS_FUNCTION MemoryBlockView<Number, MemorySpaceType>::MemoryBlockView(
     MemoryBlock<Number, MemorySpaceType> const &memory_block)
 {
   _size = memory_block._size;
@@ -121,7 +121,7 @@ ADAMANTINE_HOST_DEV MemoryBlockView<Number, MemorySpaceType>::MemoryBlockView(
 }
 
 template <typename Number, typename MemorySpaceType>
-ADAMANTINE_HOST_DEV MemoryBlockView<Number, MemorySpaceType>::MemoryBlockView(
+KOKKOS_FUNCTION MemoryBlockView<Number, MemorySpaceType>::MemoryBlockView(
     MemoryBlockView<Number, MemorySpaceType> const &memory_block_view)
 {
   _size = memory_block_view._size;
@@ -134,7 +134,7 @@ ADAMANTINE_HOST_DEV MemoryBlockView<Number, MemorySpaceType>::MemoryBlockView(
 }
 
 template <typename Number, typename MemorySpaceType>
-ADAMANTINE_HOST_DEV void MemoryBlockView<Number, MemorySpaceType>::reinit(
+KOKKOS_FUNCTION void MemoryBlockView<Number, MemorySpaceType>::reinit(
     MemoryBlock<Number, MemorySpaceType> const &memory_block)
 {
   _size = memory_block._size;
@@ -147,7 +147,7 @@ ADAMANTINE_HOST_DEV void MemoryBlockView<Number, MemorySpaceType>::reinit(
 }
 
 template <typename Number, typename MemorySpaceType>
-ADAMANTINE_HOST_DEV MemoryBlockView<Number, MemorySpaceType> &
+KOKKOS_FUNCTION MemoryBlockView<Number, MemorySpaceType> &
 MemoryBlockView<Number, MemorySpaceType>::operator=(
     MemoryBlockView<Number, MemorySpaceType> const &memory_block_view)
 {
@@ -163,7 +163,7 @@ MemoryBlockView<Number, MemorySpaceType>::operator=(
 }
 
 template <typename Number, typename MemorySpaceType>
-ADAMANTINE_HOST_DEV Number &
+KOKKOS_FUNCTION Number &
 MemoryBlockView<Number, MemorySpaceType>::operator()(unsigned int i) const
 {
   ASSERT(i < _dim_0, "Out-of-bound access.");
@@ -172,7 +172,7 @@ MemoryBlockView<Number, MemorySpaceType>::operator()(unsigned int i) const
 }
 
 template <typename Number, typename MemorySpaceType>
-ADAMANTINE_HOST_DEV Number &
+KOKKOS_FUNCTION Number &
 MemoryBlockView<Number, MemorySpaceType>::operator()(unsigned int i,
                                                      unsigned int j) const
 {
@@ -183,10 +183,8 @@ MemoryBlockView<Number, MemorySpaceType>::operator()(unsigned int i,
 }
 
 template <typename Number, typename MemorySpaceType>
-ADAMANTINE_HOST_DEV Number &
-MemoryBlockView<Number, MemorySpaceType>::operator()(unsigned int i,
-                                                     unsigned int j,
-                                                     unsigned int k) const
+KOKKOS_FUNCTION Number &MemoryBlockView<Number, MemorySpaceType>::operator()(
+    unsigned int i, unsigned int j, unsigned int k) const
 {
   ASSERT(i < _dim_0, "Out-of-bound access.");
   ASSERT(j < _dim_1, "Out-of-bound access.");
@@ -196,11 +194,8 @@ MemoryBlockView<Number, MemorySpaceType>::operator()(unsigned int i,
 }
 
 template <typename Number, typename MemorySpaceType>
-ADAMANTINE_HOST_DEV Number &
-MemoryBlockView<Number, MemorySpaceType>::operator()(unsigned int i,
-                                                     unsigned int j,
-                                                     unsigned int k,
-                                                     unsigned int l) const
+KOKKOS_FUNCTION Number &MemoryBlockView<Number, MemorySpaceType>::operator()(
+    unsigned int i, unsigned int j, unsigned int k, unsigned int l) const
 {
   ASSERT(i < _dim_0, "Out-of-bound access.");
   ASSERT(j < _dim_1, "Out-of-bound access.");
@@ -212,12 +207,9 @@ MemoryBlockView<Number, MemorySpaceType>::operator()(unsigned int i,
 }
 
 template <typename Number, typename MemorySpaceType>
-ADAMANTINE_HOST_DEV Number &
-MemoryBlockView<Number, MemorySpaceType>::operator()(unsigned int i,
-                                                     unsigned int j,
-                                                     unsigned int k,
-                                                     unsigned int l,
-                                                     unsigned int m) const
+KOKKOS_FUNCTION Number &MemoryBlockView<Number, MemorySpaceType>::operator()(
+    unsigned int i, unsigned int j, unsigned int k, unsigned int l,
+    unsigned int m) const
 {
   ASSERT(i < _dim_0, "Out-of-bound access.");
   ASSERT(j < _dim_1, "Out-of-bound access.");
@@ -231,14 +223,14 @@ MemoryBlockView<Number, MemorySpaceType>::operator()(unsigned int i,
 }
 
 template <typename Number, typename MemorySpaceType>
-ADAMANTINE_HOST_DEV unsigned int
+KOKKOS_FUNCTION unsigned int
 MemoryBlockView<Number, MemorySpaceType>::size() const
 {
   return _size;
 }
 
 template <typename Number, typename MemorySpaceType>
-ADAMANTINE_HOST_DEV unsigned int
+KOKKOS_FUNCTION unsigned int
 MemoryBlockView<Number, MemorySpaceType>::extent(unsigned int i) const
 {
   switch (i)
@@ -259,8 +251,7 @@ MemoryBlockView<Number, MemorySpaceType>::extent(unsigned int i) const
 }
 
 template <typename Number, typename MemorySpaceType>
-ADAMANTINE_HOST_DEV Number *
-MemoryBlockView<Number, MemorySpaceType>::data() const
+KOKKOS_FUNCTION Number *MemoryBlockView<Number, MemorySpaceType>::data() const
 {
   return _data;
 }

diff --git a/source/utils.hh b/source/utils.hh
@@ -24,12 +24,6 @@
 
 namespace adamantine
 {
-#ifdef __CUDACC__
-#define ADAMANTINE_HOST_DEV __host__ __device__
-#else
-#define ADAMANTINE_HOST_DEV
-#endif
-
 template <typename Number>
 inline void deep_copy(Number *output, dealii::MemorySpace::Host const &,
                       Number const *input, dealii::MemorySpace::Host const &,