From 4696b3b6e48757f6bfd23d86bd089670a75b66f2 Mon Sep 17 00:00:00 2001 From: Adrian-Diaz Date: Fri, 18 Oct 2024 10:12:03 -0600 Subject: [PATCH] WIP: partition support --- examples/ann_distributed.cpp | 14 +++- src/include/tpetra_wrapper_types.h | 116 +++++++++++++++++++++++++++++ 2 files changed, 128 insertions(+), 2 deletions(-) diff --git a/examples/ann_distributed.cpp b/examples/ann_distributed.cpp index 48f88e8b..5b9a899f 100644 --- a/examples/ann_distributed.cpp +++ b/examples/ann_distributed.cpp @@ -63,7 +63,7 @@ using namespace mtr; // matar namespace // Number of nodes in each layer including inputs and outputs // // ================================================================= -std::vector num_nodes_in_layer = {64000, 30000, 8000, 4000, 100, 25, 6} ; +std::vector num_nodes_in_layer = {64000, 30000, 8000, 4000, 2000, 1000, 100} ; //std::vector num_nodes_in_layer = {50, 25} ; // {9, 50, 100, 300, 200, 100, 20, 6} @@ -383,7 +383,17 @@ int main(int argc, char* argv[]) int local_index = ANNLayers(num_layers-1).distributed_outputs.getMapLocalIndex(global_index); std::cout << " " << ANNLayers(num_layers-1).distributed_outputs.host(local_index) << std::endl; } // end for - + + //test repartition; assume a 10 by 10 grid of outputs from ANN + //assign coords to each grid point, find a partition of the grid, then repartition output layer using new map + TpetraMVArray output_grid(100, 2); //array of 2D coordinates for 10 by 10 grid of points + //populate coords + FOR_ALL(i,0,output_grid.dims(0), { + output_grid(i, 0) = i/10; + output_grid(i, 1) = i%10; + }); // end parallel for + output_grid.repartition_vector(); + } // end of kokkos scope Kokkos::finalize(); diff --git a/src/include/tpetra_wrapper_types.h b/src/include/tpetra_wrapper_types.h index 918dc7e8..8e6e851a 100644 --- a/src/include/tpetra_wrapper_types.h +++ b/src/include/tpetra_wrapper_types.h @@ -58,6 +58,12 @@ #include "Tpetra_Details_DefaultTypes.hpp" #include "Tpetra_Import.hpp" +// Repartition Package +#include +#include +#include +#include + // Trilinos type definitions typedef Tpetra::Map<>::local_ordinal_type tpetra_LO; typedef Tpetra::Map<>::global_ordinal_type tpetra_GO; @@ -417,6 +423,8 @@ class TpetraMVArray { void perform_comms(); + void repartition_vector(); //repartitions this vector using the zoltan2 multijagged algorithm on its data + KOKKOS_INLINE_FUNCTION T& operator()(size_t i) const; @@ -824,6 +832,114 @@ void TpetraMVArray::perform_comms() { tpetra_vector->doImport(*tpetra_sub_vector, *importer, Tpetra::INSERT); } +template +void TpetraMVArray::repartition_vector() { + + int num_dim = dims_[1]; + int nranks, process_rank; + MPI_Comm_rank(mpi_comm_, &process_rank); + MPI_Comm_size(mpi_comm_, &nranks); + // construct input adapted needed by Zoltan2 problem + //typedef Xpetra::MultiVector xvector_t; + typedef Zoltan2::XpetraMultiVectorAdapter inputAdapter_t; + typedef Zoltan2::EvaluatePartition quality_t; + + // Teuchos::RCP xpetra_vector = + // Teuchos::rcp(new Xpetra::TpetraMultiVector(tpetra_vector)); + + //Teuchos::RCP problem_adapter = Teuchos::rcp(new inputAdapter_t(xpetra_vector)); + Teuchos::RCP problem_adapter = Teuchos::rcp(new inputAdapter_t(tpetra_vector)); + + // Create parameters for an RCB problem + + double tolerance = 1.05; + + Teuchos::ParameterList params("Node Partition Params"); + params.set("debug_level", "basic_status"); + params.set("debug_procs", "0"); + params.set("error_check_level", "debug_mode_assertions"); + + // params.set("algorithm", "rcb"); + params.set("algorithm", "multijagged"); + params.set("imbalance_tolerance", tolerance); + params.set("num_global_parts", nranks); + params.set("partitioning_objective", "minimize_cut_edge_count"); + + Teuchos::RCP> problem = + Teuchos::rcp(new Zoltan2::PartitioningProblem(&(*problem_adapter), ¶ms)); + + // Solve the problem + + problem->solve(); + + // create metric object where communicator is Teuchos default + + quality_t* metricObject1 = new quality_t(&(*problem_adapter), ¶ms, // problem1->getComm(), + &problem->getSolution()); + // // Check the solution. + + if (process_rank == 0) + { + metricObject1->printMetrics(std::cout); + } + + if (process_rank == 0) + { + real_t imb = metricObject1->getObjectCountImbalance(); + if (imb <= tolerance) + { + std::cout << "pass: " << imb << std::endl; + } + else + { + std::cout << "fail: " << imb << std::endl; + } + std::cout << std::endl; + } + delete metricObject1; + + // // migrate rows of the vector so they correspond to the partition recommended by Zoltan2 + + // Teuchos::RCP partitioned_node_coords_distributed = Teuchos::rcp(new MV(map, num_dim)); + + // Teuchos::RCP xpartitioned_node_coords_distributed = + // Teuchos::rcp(new Xpetra::TpetraMultiVector(partitioned_node_coords_distributed)); + + problem_adapter->applyPartitioningSolution(*tpetra_vector, tpetra_vector, problem->getSolution()); + + 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); + //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(); + length_ = (dims_[0] * dims_[1]); + // // migrate density vector if this is a restart file read + // if (simparam.restart_file&&repartition_node_densities) + // { + // Teuchos::RCP partitioned_node_densities_distributed = Teuchos::rcp(new MV(partitioned_map, 1)); + + // // create import object using local node indices map and all indices map + // Tpetra::Import importer(map, partitioned_map); + + // // comms to get ghosts + // partitioned_node_densities_distributed->doImport(*design_node_densities_distributed, importer, Tpetra::INSERT); + // design_node_densities_distributed = partitioned_node_densities_distributed; + // } + + // // update nlocal_nodes and node map + // map = partitioned_map; + // nlocal_nodes = map->getLocalNumElements(); +} + // Return size of the submap template size_t TpetraMVArray::submap_size() const {