Fix loadStateField evaluator to address #1092 (#1096)

Apparently linear access of Kokkos dynamic rank views is no longer working
sandialabs · Dec 4, 2024 · 238f75f · 238f75f
1 parent 26d3dda
commit 238f75f
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 26 deletions.
diff --git a/src/evaluators/state/PHAL_LoadStateField.hpp b/src/evaluators/state/PHAL_LoadStateField.hpp
@@ -37,13 +37,11 @@ class LoadStateFieldBase : public PHX::EvaluatorWithBaseImpl<Traits>,
 
   using ExecutionSpace = typename PHX::Device::execution_space;
 
-  PHX::MDField<ScalarType> data;
+  PHX::MDField<ScalarType> field;
   std::string fieldName;
   std::string stateName;
 
   MDFieldMemoizer<Traits> memoizer;
-
-  MDFieldVectorRight<ScalarType> dataVec;
 };
 
 template<typename EvalT, typename Traits>
@@ -65,13 +63,11 @@ class LoadStateField : public PHX::EvaluatorWithBaseImpl<Traits>,
 
   using ExecutionSpace = typename PHX::Device::execution_space;
 
-  PHX::MDField<ParamScalarT> data;
+  PHX::MDField<ParamScalarT> field;
   std::string fieldName;
   std::string stateName;
 
   MDFieldMemoizer<Traits> memoizer;
-
-  MDFieldVectorRight<ParamScalarT> dataVec;
 };
 
 // Shortcut names

diff --git a/src/evaluators/state/PHAL_LoadStateField_Def.hpp b/src/evaluators/state/PHAL_LoadStateField_Def.hpp
@@ -22,10 +22,9 @@ LoadStateFieldBase(const Teuchos::ParameterList& p)
   fieldName =  p.get<std::string>("Field Name");
   stateName =  p.get<std::string>("State Name");
 
-  PHX::MDField<ScalarType> f(fieldName, p.get<Teuchos::RCP<PHX::DataLayout> >("State Field Layout") );
-  data = f;
+  field = PHX::MDField<ScalarType>(fieldName, p.get<Teuchos::RCP<PHX::DataLayout> >("State Field Layout") );
 
-  this->addEvaluatedField(data);
+  this->addEvaluatedField(field);
   this->setName("LoadStateField("+stateName+")"+PHX::print<EvalT>());
 }
 
@@ -34,7 +33,7 @@ template<typename EvalT, typename Traits, typename ScalarType>
 void LoadStateFieldBase<EvalT, Traits, ScalarType>::postRegistrationSetup(typename Traits::SetupData d,
                       PHX::FieldManager<Traits>& fm)
 {
-  this->utils.setFieldData(data,fm);
+  this->utils.setFieldData(field,fm);
 
   d.fill_field_dependencies(this->dependentFields(),this->evaluatedFields());
   if (d.memoizer_active()) memoizer.enable_memoizer();
@@ -51,15 +50,13 @@ void LoadStateFieldBase<EvalT, Traits, ScalarType>::evaluateFields(typename Trai
   //       whomever changed the data.
   const auto& stateToLoad = (*workset.stateArrayPtr)[stateName];
   auto stateData = stateToLoad.dev();
-  const int stateToLoad_size = stateToLoad.size();
 
-  MDFieldVectorRight<ScalarType> g(data);
-  dataVec = g;
+  ALBANY_ASSERT (stateData.rank() <= 3, "Current implementation supports only views with rank up to 3. If larger rank is needed modify code below");
 
   Kokkos::parallel_for(this->getName(),
-                       Kokkos::RangePolicy<ExecutionSpace>(0,data.size()),
-                       KOKKOS_CLASS_LAMBDA(const int i) {
-    dataVec[i] = (i < stateToLoad_size) ? stateData(i) : 0.0;
+                       Kokkos::MDRangePolicy<ExecutionSpace, Kokkos::Rank<3>>({0,0,0},{stateData.extent(0),stateData.extent(1),stateData.extent(2)}),
+                       KOKKOS_CLASS_LAMBDA(const int i, const int j, const int k) {
+    field.access(i,j,k) = stateData.access(i,j,k);  //works also when rank is less than 3
   });
 }
 
@@ -70,10 +67,10 @@ LoadStateField(const Teuchos::ParameterList& p)
   fieldName =  p.get<std::string>("Field Name");
   stateName =  p.get<std::string>("State Name");
 
-  PHX::MDField<ParamScalarT> f(fieldName, p.get<Teuchos::RCP<PHX::DataLayout> >("State Field Layout") );
-  data = f;
 
-  this->addEvaluatedField(data);
+  field = PHX::MDField<ParamScalarT>(fieldName, p.get<Teuchos::RCP<PHX::DataLayout> >("State Field Layout") );
+
+  this->addEvaluatedField(field);
   this->setName("Load State Field"+PHX::print<EvalT>());
 }
 
@@ -82,7 +79,7 @@ template<typename EvalT, typename Traits>
 void LoadStateField<EvalT, Traits>::postRegistrationSetup(typename Traits::SetupData d,
                       PHX::FieldManager<Traits>& fm)
 {
-  this->utils.setFieldData(data,fm);
+  this->utils.setFieldData(field,fm);
 
   d.fill_field_dependencies(this->dependentFields(),this->evaluatedFields());
   if (d.memoizer_active()) memoizer.enable_memoizer();
@@ -99,15 +96,13 @@ void LoadStateField<EvalT, Traits>::evaluateFields(typename Traits::EvalData wor
   //       whomever changed the data.
   const auto& stateToLoad = (*workset.stateArrayPtr)[stateName];
   auto stateData = stateToLoad.dev();
-  const int stateToLoad_size = stateToLoad.size();
 
-  MDFieldVectorRight<ParamScalarT> g(data);
-  dataVec = g;
+  ALBANY_ASSERT (stateData.rank() <= 3, "Current implementation supports only views with rank up to 3. If larger rank is needed modify code below");
 
   Kokkos::parallel_for(this->getName(),
-                       Kokkos::RangePolicy<ExecutionSpace>(0,data.size()),
-                       KOKKOS_CLASS_LAMBDA(const int i) {
-    dataVec[i] = (i < stateToLoad_size) ? stateData(i) : 0.0;
+                       Kokkos::MDRangePolicy<ExecutionSpace, Kokkos::Rank<3>>({0,0,0},{stateData.extent(0),stateData.extent(1),stateData.extent(2)}),
+                       KOKKOS_CLASS_LAMBDA(const int i, const int j, const int k) {
+    field.access(i,j,k) = stateData.access(i,j,k); //works also when rank is less than 3
   });
 }