From 3a0e43e4e442717adbd9edacb2b0d9be067c82a2 Mon Sep 17 00:00:00 2001 From: Adrian-Diaz Date: Fri, 18 Oct 2024 23:17:19 -0600 Subject: [PATCH] WIP: tpetra mpi wrapper partitioning --- src/include/tpetra_wrapper_types.h | 147 ++++++++++++++++------------- 1 file changed, 84 insertions(+), 63 deletions(-) diff --git a/src/include/tpetra_wrapper_types.h b/src/include/tpetra_wrapper_types.h index 8e6e851..283d56b 100644 --- a/src/include/tpetra_wrapper_types.h +++ b/src/include/tpetra_wrapper_types.h @@ -120,7 +120,9 @@ class TpetraPartitionMap { TpetraPartitionMap(size_t global_length, MPI_Comm mpi_comm = MPI_COMM_WORLD, const std::string& tag_string = DEFAULTSTRINGARRAY); - TpetraPartitionMap(DCArrayKokkos &indices, MPI_Comm mpi_comm = MPI_COMM_WORLD); + TpetraPartitionMap(DCArrayKokkos &indices, MPI_Comm mpi_comm = MPI_COMM_WORLD, const std::string& tag_string = DEFAULTSTRINGARRAY); + + TpetraPartitionMap(Teuchos::RCP> input_tpetra_map, const std::string& tag_string = DEFAULTSTRINGARRAY); KOKKOS_INLINE_FUNCTION T& operator()(size_t i) const; @@ -155,9 +157,10 @@ class TpetraPartitionMap { T* device_pointer() const; // Method returns the raw host pointer of the Kokkos DualView - KOKKOS_INLINE_FUNCTION T* host_pointer() const; + void print() const; + // // Method returns kokkos dual view // KOKKOS_INLINE_FUNCTION // TArray1D get_kokkos_dual_view() const; @@ -195,7 +198,7 @@ TpetraPartitionMap::TpetraPartitionMap(size_t g // Constructor to pass matar dual view of indices template -TpetraPartitionMap::TpetraPartitionMap(DCArrayKokkos &indices, MPI_Comm mpi_comm) { +TpetraPartitionMap::TpetraPartitionMap(DCArrayKokkos &indices, MPI_Comm mpi_comm, const std::string& tag_string) { mpi_comm_ = mpi_comm; Teuchos::RCP> teuchos_comm = Teuchos::rcp(new Teuchos::MpiComm(mpi_comm)); tpetra_map = Teuchos::rcp(new Tpetra::Map(Teuchos::OrdinalTraits::invalid(), indices.get_kokkos_dual_view().d_view, 0, teuchos_comm)); @@ -206,6 +209,19 @@ TpetraPartitionMap::TpetraPartitionMap(DCArrayK set_mpi_type(); } +// Constructor to pass an existing Tpetra map +template +TpetraPartitionMap::TpetraPartitionMap(Teuchos::RCP> input_tpetra_map, const std::string& tag_string) { + tpetra_map = input_tpetra_map; + Teuchos::RCP> teuchos_comm = tpetra_map->getComm(); + mpi_comm_ = getRawMpiComm(*teuchos_comm); + TArray1D_host host = input_tpetra_map->getMyGlobalIndices(); + TArray1D_dev device = input_tpetra_map->getMyGlobalIndicesDevice(); + length_ = host.size(); + num_global_ = tpetra_map->getGlobalNumElements(); + set_mpi_type(); +} + template void TpetraPartitionMap::set_mpi_type() { if (typeid(T).name() == typeid(bool).name()) { @@ -278,30 +294,17 @@ T* TpetraPartitionMap::device_pointer() const { } template -KOKKOS_INLINE_FUNCTION T* TpetraPartitionMap::host_pointer() const { return host.data(); } -// template -// KOKKOS_INLINE_FUNCTION -// Kokkos::DualView TpetraPartitionMap::get_kokkos_dual_view() const { -// return this_array_; -// } - -// template -// void TpetraPartitionMap::update_host() { - -// this_array_.template modify(); -// this_array_.template sync(); -// } - -// template -// void TpetraPartitionMap::update_device() { - -// this_array_.template modify(); -// this_array_.template sync(); -// } +template +void TpetraPartitionMap::print() const { + std::ostream &out = std::cout; + Teuchos::RCP fos; + fos = Teuchos::fancyOStream(Teuchos::rcpFromRef(out)); + tpetra_map->describe(*fos,Teuchos::VERB_EXTREME); +} // Return local index (on this process/rank) corresponding to the input global index template @@ -385,8 +388,10 @@ class TpetraMVArray { bool own_comms; //This Mapped MPI Array contains its own communication plan; just call array_comms() void set_mpi_type(); - Teuchos::RCP> pmap; - Teuchos::RCP> comm_pmap; + TpetraPartitionMap pmap; + TpetraPartitionMap comm_pmap; + Teuchos::RCP> tpetra_pmap; + Teuchos::RCP> tpetra_comm_pmap; Teuchos::RCP tpetra_vector; Teuchos::RCP tpetra_sub_vector; //for owned comms situations @@ -496,7 +501,7 @@ class TpetraMVArray { // Default constructor template -TpetraMVArray::TpetraMVArray(): pmap(NULL){ +TpetraMVArray::TpetraMVArray(): tpetra_pmap(NULL){ length_ = order_ = 0; for (int i = 0; i < 7; i++) { dims_[i] = 0; @@ -509,15 +514,16 @@ TpetraMVArray::TpetraMVArray(size_t dim0, const mpi_comm_ = mpi_comm; global_dim1_ = dim0; Teuchos::RCP> teuchos_comm = Teuchos::rcp(new Teuchos::MpiComm(mpi_comm_)); - pmap = Teuchos::rcp(new Tpetra::Map((long long int) dim0, 0, teuchos_comm)); - dims_[0] = pmap->getLocalNumElements(); + tpetra_pmap = Teuchos::rcp(new Tpetra::Map((long long int) dim0, 0, teuchos_comm)); + pmap = TpetraPartitionMap(tpetra_pmap); + dims_[0] = tpetra_pmap->getLocalNumElements(); dims_[1] = 1; order_ = 1; length_ = dims_[0]; // Create host ViewCArray set_mpi_type(); this_array_ = TArray1D(tag_string, dims_[0], 1); - tpetra_vector = Teuchos::rcp(new MV(pmap, this_array_)); + tpetra_vector = Teuchos::rcp(new MV(tpetra_pmap, this_array_)); } // Overloaded 2D constructor where you provide dimensions, partitioning is done along first dimension @@ -526,15 +532,16 @@ TpetraMVArray::TpetraMVArray(size_t dim0, size_ mpi_comm_ = mpi_comm; global_dim1_ = dim0; Teuchos::RCP> teuchos_comm = Teuchos::rcp(new Teuchos::MpiComm(mpi_comm_)); - pmap = Teuchos::rcp(new Tpetra::Map((long long int) dim0, 0, teuchos_comm)); - dims_[0] = pmap->getLocalNumElements(); + tpetra_pmap = Teuchos::rcp(new Tpetra::Map((long long int) dim0, 0, teuchos_comm)); + pmap = TpetraPartitionMap(tpetra_pmap); + dims_[0] = tpetra_pmap->getLocalNumElements(); dims_[1] = dim1; order_ = 2; length_ = (dims_[0] * dims_[1]); // Create host ViewCArray set_mpi_type(); this_array_ = TArray1D(tag_string, dims_[0], dim1); - tpetra_vector = Teuchos::rcp(new MV(pmap, this_array_)); + tpetra_vector = Teuchos::rcp(new MV(tpetra_pmap, this_array_)); } // Overloaded 1D constructor where you provide a partition map @@ -543,15 +550,16 @@ TpetraMVArray::TpetraMVArray(TpetraPartitionMap const std::string& tag_string) { mpi_comm_ = input_pmap.mpi_comm_; global_dim1_ = input_pmap.num_global_; - pmap = input_pmap.tpetra_map; - dims_[0] = pmap->getLocalNumElements(); + tpetra_pmap = input_pmap.tpetra_map; + pmap = input_pmap; + dims_[0] = tpetra_pmap->getLocalNumElements(); dims_[1] = 1; order_ = 1; length_ = dims_[0]; // Create host ViewCArray set_mpi_type(); this_array_ = TArray1D(tag_string, dims_[0], 1); - tpetra_vector = Teuchos::rcp(new MV(pmap, this_array_)); + tpetra_vector = Teuchos::rcp(new MV(tpetra_pmap, this_array_)); } // Overloaded 2D constructor where you provide a partition map @@ -560,15 +568,16 @@ TpetraMVArray::TpetraMVArray(TpetraPartitionMap size_t dim1, const std::string& tag_string) { mpi_comm_ = input_pmap.mpi_comm_; global_dim1_ = input_pmap.num_global_; - pmap = input_pmap.tpetra_map; - dims_[0] = pmap->getLocalNumElements(); + tpetra_pmap = input_pmap.tpetra_map; + pmap = input_pmap; + dims_[0] = tpetra_pmap->getLocalNumElements(); dims_[1] = dim1; order_ = 2; length_ = (dims_[0] * dims_[1]); // Create host ViewCArray set_mpi_type(); this_array_ = TArray1D(tag_string, dims_[0], dim1); - tpetra_vector = Teuchos::rcp(new MV(pmap, this_array_)); + tpetra_vector = Teuchos::rcp(new MV(tpetra_pmap, this_array_)); } // Overloaded 1D constructor taking an RPC pointer to a Tpetra Map @@ -581,11 +590,12 @@ TpetraMVArray::TpetraMVArray(Teuchos::RCP(tpetra_pmap); // Create host ViewCArray set_mpi_type(); this_array_ = TArray1D(tag_string, dims_[0], 1); - tpetra_vector = Teuchos::rcp(new MV(pmap, this_array_)); + tpetra_vector = Teuchos::rcp(new MV(tpetra_pmap, this_array_)); } // Overloaded 2D constructor taking an RPC pointer to a Tpetra Map @@ -596,13 +606,14 @@ TpetraMVArray::TpetraMVArray(Teuchos::RCPgetGlobalNumElements(); dims_[0] = input_pmap->getLocalNumElements(); dims_[1] = dim1; - pmap = input_pmap; + tpetra_pmap = input_pmap; + pmap = TpetraPartitionMap(tpetra_pmap); order_ = 2; length_ = (dims_[0] * dims_[1]); // Create host ViewCArray set_mpi_type(); this_array_ = TArray1D(tag_string, dims_[0], dim1); - tpetra_vector = Teuchos::rcp(new MV(pmap, this_array_)); + tpetra_vector = Teuchos::rcp(new MV(tpetra_pmap, this_array_)); } // Tpetra vector argument constructor: CURRENTLY DOESN'T WORK SINCE WE CANT GET DUAL VIEW FROM THE MULTIVECTOR @@ -610,7 +621,7 @@ TpetraMVArray::TpetraMVArray(Teuchos::RCP::TpetraMVArray(Teuchos::RCP input_tpetra_vector, const std::string& tag_string){ // tpetra_vector = input_tpetra_vector; -// pmap = input_tpetra_vector->getMap(); +// tpetra_pmap = input_tpetra_vector->getMap(); // //this_array_ = tpetra_vector->getWrappedDualView(); // dims_[0] = tpetra_vector->getMap()->getLocalNumElements()(); // dims_[1] = tpetra_vector->getNumVectors(); @@ -672,7 +683,7 @@ T& TpetraMVArray::operator()(size_t i, size_t j template KOKKOS_INLINE_FUNCTION long long int TpetraMVArray::getSubMapGlobalIndex(int local_index) const { - long long int global_index = comm_pmap->getGlobalElement(local_index); + long long int global_index = tpetra_comm_pmap->getGlobalElement(local_index); return global_index; } @@ -680,7 +691,7 @@ long long int TpetraMVArray::getSubMapGlobalInd template KOKKOS_INLINE_FUNCTION long long int TpetraMVArray::getMapGlobalIndex(int local_index) const { - long long int global_index = pmap->getGlobalElement(local_index); + long long int global_index = tpetra_pmap->getGlobalElement(local_index); return global_index; } @@ -688,7 +699,7 @@ long long int TpetraMVArray::getMapGlobalIndex( template KOKKOS_INLINE_FUNCTION int TpetraMVArray::getSubMapLocalIndex(long long int global_index) const { - int local_index = comm_pmap->getLocalElement(global_index); + int local_index = tpetra_comm_pmap->getLocalElement(global_index); return local_index; } @@ -696,7 +707,7 @@ int TpetraMVArray::getSubMapLocalIndex(long lon template KOKKOS_INLINE_FUNCTION int TpetraMVArray::getMapLocalIndex(long long int global_index) const { - int local_index = pmap->getLocalElement(global_index); + int local_index = tpetra_pmap->getLocalElement(global_index); return local_index; } @@ -734,6 +745,8 @@ TpetraMVArray& TpetraMVArray::update_device() { template void TpetraMVArray::own_comm_setup(TpetraPartitionMap other_pmap) { own_comms = true; - comm_pmap = other_pmap.tpetra_map; - int local_offset = pmap->getLocalElement((comm_pmap->getMinGlobalIndex())); - tpetra_sub_vector = Teuchos::rcp(new MV(*tpetra_vector, comm_pmap, local_offset)); - submap_size_ = comm_pmap->getLocalNumElements(); - importer = Teuchos::rcp(new Tpetra::Import(comm_pmap, pmap)); + tpetra_comm_pmap = other_pmap.tpetra_map; + comm_pmap = TpetraPartitionMap(tpetra_comm_pmap); + int local_offset = tpetra_pmap->getLocalElement((tpetra_comm_pmap->getMinGlobalIndex())); + tpetra_sub_vector = Teuchos::rcp(new MV(*tpetra_vector, tpetra_comm_pmap, local_offset)); + submap_size_ = tpetra_comm_pmap->getLocalNumElements(); + importer = Teuchos::rcp(new Tpetra::Import(tpetra_comm_pmap, tpetra_pmap)); } -//requires both pmap and other_pmap to be contiguous and for other_pmap to be a subset of pmap on every process +//requires both tpetra_pmap and other_pmap to be contiguous and for other_pmap to be a subset of tpetra_pmap on every process template void TpetraMVArray::own_comm_setup(Teuchos::RCP> other_pmap) { own_comms = true; - comm_pmap = other_pmap; - int local_offset = pmap->getLocalElement((comm_pmap->getMinGlobalIndex())); - tpetra_sub_vector = Teuchos::rcp(new MV(*tpetra_vector, comm_pmap, local_offset)); - submap_size_ = comm_pmap->getLocalNumElements(); - importer = Teuchos::rcp(new Tpetra::Import(comm_pmap, pmap)); + tpetra_comm_pmap = other_pmap; + comm_pmap = TpetraPartitionMap(tpetra_comm_pmap); + int local_offset = tpetra_pmap->getLocalElement((tpetra_comm_pmap->getMinGlobalIndex())); + tpetra_sub_vector = Teuchos::rcp(new MV(*tpetra_vector, tpetra_comm_pmap, local_offset)); + submap_size_ = tpetra_comm_pmap->getLocalNumElements(); + importer = Teuchos::rcp(new Tpetra::Import(tpetra_comm_pmap, tpetra_pmap)); } template void TpetraMVArray::perform_comms() { - tpetra_vector->doImport(*tpetra_sub_vector, *importer, Tpetra::INSERT); + if(own_comms){ + tpetra_vector->doImport(*tpetra_sub_vector, *importer, Tpetra::INSERT); + } + else{} + } template @@ -907,20 +926,22 @@ void TpetraMVArray::repartition_vector() { problem_adapter->applyPartitioningSolution(*tpetra_vector, tpetra_vector, problem->getSolution()); - pmap = Teuchos::rcp(new Tpetra::Map(*(tpetra_vector->getMap()))); + tpetra_pmap = Teuchos::rcp(new Tpetra::Map(*(tpetra_vector->getMap()))); // *partitioned_node_coords_distributed = Xpetra::toTpetra(*xpartitioned_node_coords_distributed); // Teuchos::RCP> partitioned_map = Teuchos::rcp(new Tpetra::Map(*(partitioned_node_coords_distributed->getMap()))); Teuchos::RCP> partitioned_map_one_to_one; - partitioned_map_one_to_one = Tpetra::createOneToOne(pmap); + partitioned_map_one_to_one = Tpetra::createOneToOne(tpetra_pmap); //Teuchos::RCP partitioned_node_coords_one_to_one_distributed = Teuchos::rcp(new MV(partitioned_map_one_to_one, num_dim)); // Tpetra::Import importer_one_to_one(partitioned_map, partitioned_map_one_to_one); // partitioned_node_coords_one_to_one_distributed->doImport(*partitioned_node_coords_distributed, importer_one_to_one, Tpetra::INSERT); // node_coords_distributed = partitioned_node_coords_one_to_one_distributed; - pmap = Teuchos::rcp(new Tpetra::Map(*partitioned_map_one_to_one)); - dims_[0] = pmap->getLocalNumElements(); + tpetra_pmap = Teuchos::rcp(new Tpetra::Map(*partitioned_map_one_to_one)); + pmap = TpetraPartitionMap(tpetra_pmap); + own_comms = false; //reset submap setup now that full map is different + dims_[0] = tpetra_pmap->getLocalNumElements(); length_ = (dims_[0] * dims_[1]); // // migrate density vector if this is a restart file read // if (simparam.restart_file&&repartition_node_densities)