Panzer tangent unit tests (Scatter Tpetra) (#13583)

* Panzer Tangents :: ScatterResidual_Tpetra updated for device along with unit test --------- Signed-off-by: Bryan Reuter <[email protected]>
trilinos · Nov 13, 2024 · b3df351 · b3df351
1 parent ccf9207
commit b3df351
Show file tree

Hide file tree

Showing 8 changed files with 1,028 additions and 134 deletions.
diff --git a/packages/panzer/adapters-stk/test/evaluator_tests/CMakeLists.txt b/packages/panzer/adapters-stk/test/evaluator_tests/CMakeLists.txt
@@ -57,6 +57,13 @@ TRIBITS_ADD_EXECUTABLE_AND_TEST(
   NUM_MPI_PROCS 2
   )
 
+TRIBITS_ADD_EXECUTABLE_AND_TEST(
+  tScatterResidual_BlockedTpetra
+  SOURCES tpetra_blocked_scatter_residual.cpp ${UNIT_TEST_DRIVER}
+  COMM serial mpi
+  NUM_MPI_PROCS 2
+  )
+
 TRIBITS_ADD_EXECUTABLE_AND_TEST(
   tScatterDirichletResidual_Tpetra
   SOURCES tpetra_scatter_dirichlet_residual.cpp ${UNIT_TEST_DRIVER}

diff --git a/packages/panzer/adapters-stk/test/evaluator_tests/tpetra_blocked_scatter_residual.cpp b/packages/panzer/adapters-stk/test/evaluator_tests/tpetra_blocked_scatter_residual.cpp
diff --git a/packages/panzer/adapters-stk/test/evaluator_tests/tpetra_scatter_residual.cpp b/packages/panzer/adapters-stk/test/evaluator_tests/tpetra_scatter_residual.cpp
diff --git a/packages/panzer/disc-fe/src/evaluators/Panzer_GatherSolution_BlockedTpetra_impl.hpp b/packages/panzer/disc-fe/src/evaluators/Panzer_GatherSolution_BlockedTpetra_impl.hpp
@@ -305,9 +305,6 @@ postRegistrationSetup(typename TRAITS::SetupData d,
                                                 maxElementBlockGIDCount);
 
   // Set up storage for tangentFields using view of views
-  // We also need storage for the number of tangent fields associated with
-  // each gatherField
-
   if (has_tangent_fields_) {
 
     size_t inner_vector_max_size = 0;

diff --git a/packages/panzer/disc-fe/src/evaluators/Panzer_GatherSolution_Tpetra_impl.hpp b/packages/panzer/disc-fe/src/evaluators/Panzer_GatherSolution_Tpetra_impl.hpp
@@ -263,8 +263,8 @@ postRegistrationSetup(typename TRAITS::SetupData /* d */,
   }
   Kokkos::deep_copy(tangentInnerVectorSizes_,tangentInnerVectorSizes_host);
 
-  gatherFieldsVoV_.initialize("GatherSolution_Teptra<Tangent>::gatherFieldsVoV_",gatherFields_.size());
-  tangentFieldsVoV_.initialize("GatherSolution_Teptra<Tangent>::tangentFieldsVoV_",gatherFields_.size(),inner_vector_max_size);
+  gatherFieldsVoV_.initialize("GatherSolution_Tpetra<Tangent>::gatherFieldsVoV_",gatherFields_.size());
+  tangentFieldsVoV_.initialize("GatherSolution_Tpetra<Tangent>::tangentFieldsVoV_",gatherFields_.size(),inner_vector_max_size);
 
   for (std::size_t fd = 0; fd < gatherFields_.size(); ++fd) {
     const std::string& fieldName = indexerNames_[fd];

diff --git a/packages/panzer/disc-fe/src/evaluators/Panzer_ScatterResidual_BlockedTpetra_impl.hpp b/packages/panzer/disc-fe/src/evaluators/Panzer_ScatterResidual_BlockedTpetra_impl.hpp
@@ -472,6 +472,4 @@ evaluateFields(typename TRAITS::EvalData workset)
 
 }
 
-// **********************************************************************
-
 #endif
diff --git a/packages/panzer/disc-fe/src/evaluators/Panzer_ScatterResidual_Tpetra_decl.hpp b/packages/panzer/disc-fe/src/evaluators/Panzer_ScatterResidual_Tpetra_decl.hpp
@@ -14,6 +14,7 @@
 #include "Phalanx_config.hpp"
 #include "Phalanx_Evaluator_Macros.hpp"
 #include "Phalanx_MDField.hpp"
+#include "Phalanx_KokkosViewOfViews.hpp"
 
 #include "Teuchos_ParameterList.hpp"
 
@@ -134,6 +135,7 @@ class ScatterResidual_Tpetra<panzer::Traits::Tangent,TRAITS,LO,GO,NodeT>
 
 private:
   typedef typename panzer::Traits::Tangent::ScalarT ScalarT;
+  typedef typename panzer::Traits::RealType RealT;
 
   // dummy field so that the evaluator will have something to do
   Teuchos::RCP<PHX::FieldTag> scatterHolder_;
@@ -153,8 +155,13 @@ class ScatterResidual_Tpetra<panzer::Traits::Tangent,TRAITS,LO,GO,NodeT>
   Teuchos::RCP<const std::map<std::string,std::string> > fieldMap_;
 
   std::string globalDataKey_; // what global data does this fill?
+  Teuchos::RCP<const TpetraLinearObjContainer<double,LO,GO,NodeT> > tpetraContainer_;
+
+  /// Storage for the tangent data
+  PHX::ViewOfViews<1,Kokkos::View<RealT**,Kokkos::LayoutLeft,PHX::Device>> dfdpFieldsVoV_;
 
-  std::vector< Teuchos::ArrayRCP<double> > dfdp_vectors_;
+  PHX::View<int**> scratch_lids_;
+  std::vector<PHX::View<int*> > scratch_offsets_;
 
   ScatterResidual_Tpetra();
 };

diff --git a/packages/panzer/disc-fe/src/evaluators/Panzer_ScatterResidual_Tpetra_impl.hpp b/packages/panzer/disc-fe/src/evaluators/Panzer_ScatterResidual_Tpetra_impl.hpp
@@ -130,6 +130,7 @@ ScatterResidual_Tpetra(const Teuchos::RCP<const panzer::GlobalIndexer> & indexer
   : globalIndexer_(indexer)
   , globalDataKey_("Residual Scatter Container")
 {
+
   std::string scatterName = p.get<std::string>("Scatter Name");
   scatterHolder_ =
     Teuchos::rcp(new PHX::Tag<ScalarT>(scatterName,Teuchos::rcp(new PHX::MDALayout<Dummy>(0))));
@@ -146,6 +147,7 @@ ScatterResidual_Tpetra(const Teuchos::RCP<const panzer::GlobalIndexer> & indexer
 
   // build the vector of fields that this is dependent on
   scatterFields_.resize(names.size());
+  scratch_offsets_.resize(names.size());
   for (std::size_t eq = 0; eq < names.size(); ++eq) {
     scatterFields_[eq] = PHX::MDField<const ScalarT,Cell,NODE>(names[eq],dl);
 
@@ -165,16 +167,25 @@ ScatterResidual_Tpetra(const Teuchos::RCP<const panzer::GlobalIndexer> & indexer
 // **********************************************************************
 template<typename TRAITS,typename LO,typename GO,typename NodeT>
 void panzer::ScatterResidual_Tpetra<panzer::Traits::Tangent, TRAITS,LO,GO,NodeT>::
-postRegistrationSetup(typename TRAITS::SetupData /* d */,
+postRegistrationSetup(typename TRAITS::SetupData d,
                       PHX::FieldManager<TRAITS>& /* fm */)
 {
   fieldIds_.resize(scatterFields_.size());
+  const Workset & workset_0 = (*d.worksets_)[0];
+  std::string blockId = this->wda(workset_0).block_id;
+
   // load required field numbers for fast use
   for(std::size_t fd=0;fd<scatterFields_.size();++fd) {
     // get field ID from DOF manager
     std::string fieldName = fieldMap_->find(scatterFields_[fd].fieldTag().name())->second;
     fieldIds_[fd] = globalIndexer_->getFieldNum(fieldName);
+
+    const std::vector<int> & offsets = globalIndexer_->getGIDFieldOffsets(blockId,fieldIds_[fd]);
+    scratch_offsets_[fd] = PHX::View<int*>("offsets",offsets.size());
+    Kokkos::deep_copy(scratch_offsets_[fd], Kokkos::View<const int*, Kokkos::HostSpace, Kokkos::MemoryUnmanaged>(offsets.data(), offsets.size()));
   }
+  scratch_lids_ = PHX::View<LO**>("lids",scatterFields_[0].extent(0),
+                                                 globalIndexer_->getElementBlockGIDCount(blockId));
 }
 
 // **********************************************************************
@@ -191,50 +202,26 @@ preEvaluate(typename TRAITS::PreEvalData d)
   std::vector<std::string> activeParameters =
     rcp_dynamic_cast<ParameterList_GlobalEvaluationData>(d.gedc->getDataObject("PARAMETER_NAMES"))->getActiveParameters();
 
-  dfdp_vectors_.clear();
+  dfdpFieldsVoV_.initialize("ScatterResidual_Tpetra<Tangent>::dfdpFieldsVoV_",activeParameters.size());
+
   for(std::size_t i=0;i<activeParameters.size();i++) {
     RCP<typename LOC::VectorType> vec =
       rcp_dynamic_cast<LOC>(d.gedc->getDataObject(activeParameters[i]),true)->get_f();
-    Teuchos::ArrayRCP<double> vec_array = vec->get1dViewNonConst();
-    dfdp_vectors_.push_back(vec_array);
+    auto dfdp_view = vec->getLocalViewDevice(Tpetra::Access::ReadWrite);
+
+    dfdpFieldsVoV_.addView(dfdp_view,i);
   }
-}
 
-// **********************************************************************
-template<typename TRAITS,typename LO,typename GO,typename NodeT>
-void panzer::ScatterResidual_Tpetra<panzer::Traits::Tangent, TRAITS,LO,GO,NodeT>::
-evaluateFields(typename TRAITS::EvalData workset)
-{
-   // for convenience pull out some objects from workset
-   std::string blockId = this->wda(workset).block_id;
-   const std::vector<std::size_t> & localCellIds = this->wda(workset).cell_local_ids;
-
-   // NOTE: A reordering of these loops will likely improve performance
-   //       The "getGIDFieldOffsets may be expensive.  However the
-   //       "getElementGIDs" can be cheaper. However the lookup for LIDs
-   //       may be more expensive!
-
-   // scatter operation for each cell in workset
-   for(std::size_t worksetCellIndex=0;worksetCellIndex<localCellIds.size();++worksetCellIndex) {
-      std::size_t cellLocalId = localCellIds[worksetCellIndex];
-
-      auto LIDs = globalIndexer_->getElementLIDs(cellLocalId);
-
-      // loop over each field to be scattered
-      for (std::size_t fieldIndex = 0; fieldIndex < scatterFields_.size(); fieldIndex++) {
-         int fieldNum = fieldIds_[fieldIndex];
-         const std::vector<int> & elmtOffset = globalIndexer_->getGIDFieldOffsets(blockId,fieldNum);
-
-         // loop over basis functions
-         for(std::size_t basis=0;basis<elmtOffset.size();basis++) {
-            int offset = elmtOffset[basis];
-            LO lid = LIDs[offset];
-            ScalarT value = (scatterFields_[fieldIndex])(worksetCellIndex,basis);
-            for(int d=0;d<value.size();d++)
-              dfdp_vectors_[d][lid] += value.fastAccessDx(d);
-         }
-      }
-   }
+  dfdpFieldsVoV_.syncHostToDevice();
+
+  // extract linear object container
+  tpetraContainer_ = Teuchos::rcp_dynamic_cast<LOC>(d.gedc->getDataObject(globalDataKey_));
+
+  if(tpetraContainer_==Teuchos::null) {
+    // extract linear object container
+    Teuchos::RCP<LinearObjContainer> loc = Teuchos::rcp_dynamic_cast<LOCPair_GlobalEvaluationData>(d.gedc->getDataObject(globalDataKey_),true)->getGhostedLOC();
+    tpetraContainer_ = Teuchos::rcp_dynamic_cast<LOC>(loc);
+  }  
 }
 
 // **********************************************************************
@@ -392,7 +379,6 @@ class ScatterResidual_Residual_Functor {
   PHX::View<const int*> offsets; // how to get a particular field
   FieldType field;
 
-
   KOKKOS_INLINE_FUNCTION
   void operator()(const unsigned int cell) const
   {
@@ -407,6 +393,44 @@ class ScatterResidual_Residual_Functor {
   }
 };
 
+template <typename ScalarT,typename LO,typename GO,typename NodeT>
+class ScatterResidual_Tangent_Functor {
+public:
+  typedef typename PHX::Device execution_space;
+  typedef PHX::MDField<const ScalarT,Cell,NODE> FieldType;
+
+  bool fillResidual;
+  Kokkos::View<double**, Kokkos::LayoutLeft,PHX::Device> r_data;
+
+  Kokkos::View<const LO**> lids; // local indices for unknowns.
+  PHX::View<const int*> offsets; // how to get a particular field
+  FieldType field;
+  double num_params;
+
+  Kokkos::View<Kokkos::View<double**,Kokkos::LayoutLeft,PHX::Device>*>  dfdp_fields; // tangent fields
+
+  KOKKOS_INLINE_FUNCTION
+  void operator()(const unsigned int cell) const
+  {
+
+    // loop over the basis functions (currently they are nodes)
+    for(std::size_t basis=0; basis < offsets.extent(0); basis++) {
+       typename FieldType::array_type::reference_type scatterField = field(cell,basis);
+       int offset = offsets(basis);
+       LO lid    = lids(cell,offset);
+
+       // Sum residual
+       if(fillResidual)
+         Kokkos::atomic_add(&r_data(lid,0), scatterField.val());
+
+       // loop over the tangents
+       for(int i_param=0; i_param<num_params; i_param++)
+          dfdp_fields(i_param)(lid,0) += scatterField.fastAccessDx(i_param);
+
+    } // end basis
+  }
+};
+
 }
 }
 
@@ -483,6 +507,35 @@ evaluateFields(typename TRAITS::EvalData workset)
 
 }
 
+// **********************************************************************
+template<typename TRAITS,typename LO,typename GO,typename NodeT>
+void panzer::ScatterResidual_Tpetra<panzer::Traits::Tangent, TRAITS,LO,GO,NodeT>::
+evaluateFields(typename TRAITS::EvalData workset)
+{
+  typedef TpetraLinearObjContainer<double,LO,GO,NodeT> LOC;
+
+  // for convenience pull out some objects from workset
+  std::string blockId = this->wda(workset).block_id;
+
+  Teuchos::RCP<typename LOC::VectorType> r = tpetraContainer_->get_f();
+
+  globalIndexer_->getElementLIDs(this->wda(workset).getLocalCellIDs(),scratch_lids_);
+
+  ScatterResidual_Tangent_Functor<ScalarT,LO,GO,NodeT> functor;
+  functor.r_data = r->getLocalViewDevice(Tpetra::Access::ReadWrite);
+  functor.lids = scratch_lids_;
+  functor.dfdp_fields = dfdpFieldsVoV_.getViewDevice();
+
+  // for each field, do a parallel for loop
+  for(std::size_t fieldIndex = 0; fieldIndex < scatterFields_.size(); fieldIndex++) {
+    functor.offsets = scratch_offsets_[fieldIndex];
+    functor.field = scatterFields_[fieldIndex];
+    functor.num_params = Kokkos::dimension_scalar(scatterFields_[fieldIndex].get_view())-1;
+
+    Kokkos::parallel_for(workset.num_cells,functor);
+  }
+}
+
 // **********************************************************************
 
 #endif
Original file line number	Diff line number	Diff line change
Expand Up		@@ -472,6 +472,4 @@ evaluateFields(typename TRAITS::EvalData workset)

		}

		// **********************************************************************

		#endif