From bdc3d745bd6fffcefb76b415c51b9ad33db7859f Mon Sep 17 00:00:00 2001 From: Adrian-Diaz Date: Fri, 1 Nov 2024 00:28:02 -0600 Subject: [PATCH] WIP: tpetra wrappers --- src/include/tpetra_wrapper_types.h | 140 ++++++++++++++++++----------- 1 file changed, 86 insertions(+), 54 deletions(-) diff --git a/src/include/tpetra_wrapper_types.h b/src/include/tpetra_wrapper_types.h index bfa1c912..f427a45d 100644 --- a/src/include/tpetra_wrapper_types.h +++ b/src/include/tpetra_wrapper_types.h @@ -1032,22 +1032,26 @@ TpetraMVArray::~TpetraMVArray() {} //////////////////////////////////////////////////////////////////////////////// ///////////////////////// -// TpetraCArray: CArray Tpetra wrapper. +// TpetraCArray: CArray Tpetra wrapper around Tpetra CRSMatrix; it was the only distributed Layout Right type we knew about in Trilinos. ///////////////////////// template class TpetraCArray { - // this is manage + // this is managed using TArray1D = CArrayKokkos; using TArray1D_Host = RaggedRightArrayKokkos; + using row_map_type = Kokkos::View; using input_row_map_type = DCArrayKokkos; using values_array = Kokkos::View; using global_indices_array = Kokkos::View; using indices_array = Kokkos::View; - size_t dim1_; - size_t global_dim1_; + size_t dims_[2]; + size_t global_dim0_; + size_t submap_size_; + size_t length_; + size_t order_; // tensor order (rank) size_t column_map_size_; size_t length_; MPI_Comm mpi_comm_; @@ -1094,8 +1098,11 @@ class TpetraCArray { *this = temp; } + TpetraCArray(size_t global_dim0, + const std::string& tag_string = DEFAULTSTRINGARRAY, MPI_Comm mpi_comm = MPI_COMM_WORLD); + //CRS matrix constructor for banded matrix case - TpetraCArray(size_t dim1, size_t dim2, + TpetraCArray(size_t global_dim0, size_t dim2, const std::string& tag_string = DEFAULTSTRINGARRAY, MPI_Comm mpi_comm = MPI_COMM_WORLD); //CRS matrix constructor with arbitrary row graph and column map supplied @@ -1104,6 +1111,9 @@ class TpetraCArray { //CRS matric constructor with arbitrary row graph; builds column map for you and thus one less arg TpetraCArray(Teuchos::RCP> input_pmap, size_t dim1, const std::string& tag_string = DEFAULTSTRINGARRAY); + KOKKOS_INLINE_FUNCTION + T& operator()(size_t i) const; + KOKKOS_INLINE_FUNCTION T& operator()(size_t i, size_t j) const; @@ -1135,7 +1145,7 @@ class TpetraCArray { size_t extent() const; KOKKOS_INLINE_FUNCTION - size_t dim1() const; + size_t dims(size_t i) const; size_t global_dim() const; @@ -1175,31 +1185,65 @@ TpetraCArray::TpetraCArray(): tpetra_pmap(NULL) } } +// 1D Constructor +template +TpetraCArray::TpetraCArray(size_t global_dim0, const std::string& tag_string, MPI_Comm mpi_comm) { +} + // Constructor that takes local data in a matar ragged type template -TpetraCArray::TpetraCArray(size_t global_dim1, size_t dim2, const std::string& tag_string, MPI_Comm mpi_comm) { +TpetraCArray::TpetraCArray(size_t global_dim0, size_t dim1, const std::string& tag_string, MPI_Comm mpi_comm) { mpi_comm_ = mpi_comm; - global_dim1_ = global_dim1; + global_dim0_ = global_dim0; Teuchos::RCP> teuchos_comm = Teuchos::rcp(new Teuchos::MpiComm(mpi_comm_)); - tpetra_pmap = Teuchos::rcp(new Tpetra::Map((long long int) global_dim1, 0, teuchos_comm)); + tpetra_pmap = Teuchos::rcp(new Tpetra::Map((long long int) global_dim0, 0, teuchos_comm)); pmap = TpetraPartitionMap(tpetra_pmap); - dim1_ = tpetra_pmap->getLocalNumElements(); + dims_[0] = tpetra_pmap->getLocalNumElements(); + dims_[1] = dim1; + order = 2; //construct strides that are constant mystrides_ = row_map_type("mystrides_",dim1_); - for(int irow = 0; irow < dim1_; irow++){ - mystrides_(irow) = dim2; + FOR_ALL(irow, 0, dims_[0], { + mystrides_(irow) = dims_[1]; + }); + + this_array_ = TArray1D(dims_[0], dims_[1], "tag_string"); + size_t nnz = dims_[0]*dims_[1]; + + //assign indices 0 to dim1-1 since tpetra constructor still needs a graph + row_map_type row_offsets_pass("row_offsets", dims_[0] + 1); + + Kokkos::parallel_for("StartValuesInit", dim1_+1, KOKKOS_CLASS_LAMBDA(const int i) { + row_offsets_pass(i) = 0; + }); + + Kokkos::parallel_scan("StartValuesSetup", dim1_, KOKKOS_CLASS_LAMBDA(const int i, int& update, const bool final) { + // Load old value in case we update it before accumulating + const size_t count = mystrides_(i); + update += count; + if (final) { + row_offsets_pass((i+1)) = update; + } + + }); + + row_offsets_pass(0) = 0; + for(int ipass = 1; ipass < dim1_ + 1; ipass++){ + row_offsets_pass(ipass) = row_offsets_pass(ipass-1) + dims_[1]; } - this_array_ = input_values; - global_indices_array input_crs_graph = crs_graph.get_kokkos_dual_view().d_view; - + global_indices_array input_crs_graph = global_indices_array("dummy_graph", nnz); + FOR_ALL(i,0,dims_[0], { + for(int igraph = 0; igraph < dim1_; igraph++){ + input_crs_graph(row_offsets_pass(i) + igraph) = igraph; + } + }); //build column map for the global conductivity matrix Teuchos::RCP > colmap; const Teuchos::RCP > dommap = tpetra_pmap; - Tpetra::Details::makeColMap(colmap, tpetra_pmap, input_crs_graph.get_kokkos_dual_view().d_view, nullptr); + Tpetra::Details::makeColMap(colmap, tpetra_pmap, input_crs_graph, nullptr); tpetra_column_pmap = colmap; - size_t nnz = input_crs_graph.size(); //debug print //std::cout << "DOF GRAPH SIZE ON RANK " << myrank << " IS " << nnz << std::endl; @@ -1208,23 +1252,20 @@ TpetraCArray::TpetraCArray(size_t global_dim1, crs_local_indices_ = indices_array("crs_local_indices", nnz); //row offsets with compatible template arguments - row_map_type row_offsets_pass("row_offsets", dim1_ + 1); - for(int ipass = 0; ipass < dim1_ + 1; ipass++){ - row_offsets_pass(ipass) = input_values.start_index_(ipass); - } + size_t entrycount = 0; for(int irow = 0; irow < dim1_; irow++){ for(int istride = 0; istride < mystrides_(irow); istride++){ - crs_local_indices_(entrycount) = tpetra_column_pmap->getLocalElement(crs_graph(entrycount)); + crs_local_indices_(entrycount) = tpetra_column_pmap->getLocalElement(input_crs_graph(entrycount)); entrycount++; } } //sort values and indices - Tpetra::Import_Util::sortCrsEntries(row_offsets_pass, crs_local_indices_.d_view, this_array_.get_kokkos_view()); + Tpetra::Import_Util::sortCrsEntries(row_offsets_pass, crs_local_indices_, this_array_.get_kokkos_view()); - tpetra_crs_matrix = Teuchos::rcp(new MAT(tpetra_pmap, tpetra_column_pmap, start_index_.d_view, crs_local_indices_.d_view, this_array_.get_kokkos_view())); + tpetra_crs_matrix = Teuchos::rcp(new MAT(tpetra_pmap, tpetra_column_pmap, row_offsets_pass, crs_local_indices_, this_array_.get_kokkos_view())); tpetra_crs_matrix->fillComplete(); } @@ -1232,18 +1273,6 @@ TpetraCArray::TpetraCArray(size_t global_dim1, template TpetraCArray::TpetraCArray(TpetraPartitionMap &input_pmap, size_t dim1, const std::string& tag_string) { - // mpi_comm_ = input_pmap.mpi_comm_; - // global_dim1_ = input_pmap.num_global_; - // 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(tpetra_pmap, this_array_)); } // Overloaded 2D constructor taking an RPC pointer to a Tpetra Map @@ -1251,17 +1280,6 @@ template ::TpetraCArray(Teuchos::RCP> input_pmap, size_t dim1, const std::string& tag_string) { - // global_dim1_ = input_pmap->getGlobalNumElements(); - // dims_[0] = input_pmap->getLocalNumElements(); - // dims_[1] = dim1; - // 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(tpetra_pmap, this_array_)); } template @@ -1290,11 +1308,20 @@ void TpetraCArray::set_mpi_type() { } } +template +KOKKOS_INLINE_FUNCTION +T& TpetraCArray::operator()(size_t i) const { + assert(order_ == 1 && "Tensor order (rank) does not match constructor in TpetraMVArray 1D!"); + assert(i >= 0 && i < dims_[0] && "i is out of bounds in TpetraCArray 1D!"); + return this_array_(i); +} + template KOKKOS_INLINE_FUNCTION T& TpetraCArray::operator()(size_t i, size_t j) const { - assert(i >= 0 && i < dim1_ && "i is out of bounds in TpetraCArray!"); - assert(j >= 0 && j < mystrides_(i) && "j is out of bounds in TpetraCArray!"); + assert(order_ == 2 && "Tensor order (rank) does not match constructor in TpetraMVArray 2D!"); + assert(i >= 0 && i < dims_[0] && "i is out of bounds in TpetraCArray 2D!"); + assert(j >= 0 && j < dims_[1] && "j is out of bounds in TpetraCArray 2D!"); return this_array_(i,j); } @@ -1343,11 +1370,14 @@ TpetraCArray& TpetraCArray::extent() const { template KOKKOS_INLINE_FUNCTION -size_t TpetraCArray::dim1() const { - return dim1_; +size_t TpetraCArray::dims(size_t i) const { + assert(i < order_ && "TpetraMVArray order (rank) does not match constructor, dim[i] does not exist!"); + assert(i >= 0 && dims_[i]>0 && "Access to TpetraMVArray dims is out of bounds!"); + return dims_[i]; } template size_t TpetraCArray::global_dim() const { - return global_dim1_; + return global_dim0_; } template