54 #include <Teuchos_RCP.hpp>
55 #include <Teuchos_ArrayView.hpp>
57 #include <unordered_map>
59 #include <Tpetra_Map.hpp>
60 #include <Tpetra_Vector.hpp>
64 #include <zoltan_dd_cpp.h>
65 #include <unordered_set>
74 std::ostream &ostr = std::cout)
76 ostr << msg << std::endl;
77 Teuchos::FancyOStream fancy(Teuchos::rcpFromRef(ostr));
78 thing.describe(fancy, Teuchos::VERB_EXTREME);
81 template <
typename vector_t>
83 std::ostream &ostr = std::cout)
85 if (vec.getGlobalLength() >
TOOMANY)
return;
86 printTpetraThing<vector_t>(vec, msg, ostr);
89 template <
typename map_t>
91 std::ostream &ostr = std::cout)
93 if (map.getGlobalNumElements() >
TOOMANY)
return;
94 printTpetraThing<map_t>(map, msg, ostr);
99 template <
typename id_t>
103 IDs(
size_t nIds_,
float fracShared_, id_t idBase_,
int idStride_,
104 Teuchos::RCP<
const Teuchos::Comm<int> > &comm_) :
106 nShared(std::ceil(nIds * fracShared_)),
107 nUnique(nIds - nShared),
110 contiguous(idStride_ == 1),
111 ids(nIds, 0), comm(comm_)
113 int me = comm->getRank();
114 int np = comm->getSize();
118 for (
size_t i = 0; i < nUnique; i++)
119 ids[cnt++] = idBase + ((me * nUnique + i) * idStride);
122 size_t gnUnique = nUnique * np;
125 for (
size_t i = 0; i < nShared; i++) {
126 size_t r = rand() % nShared;
127 ids[cnt++] = idBase + ((gnUnique + r) * idStride);
137 int me = comm->getRank();
139 std::cout << me <<
" nIds = " << nIds <<
"; nUnique = " << nUnique
140 <<
"; nShared = " << nShared << std::endl;
142 std::cout << me <<
" Unique: ";
143 for (
size_t i = 0; i < nUnique; i++) std::cout << ids[i] <<
" ";
144 std::cout << std::endl;
146 std::cout << me <<
" Shared: ";
147 for (
size_t i = 0; i < nShared; i++) std::cout << ids[nUnique+i] <<
" ";
148 std::cout << std::endl;
157 typedef int scalar_t;
168 std::vector<id_t> ids;
170 Teuchos::RCP<const Teuchos::Comm<int> > comm;
188 template <
typename id_t>
191 typedef typename Tpetra::Map<int, id_t>
map_t;
192 typedef typename Teuchos::RCP<const map_t> rcpmap_t;
194 typedef typename Tpetra::Vector<scalar_t, int, id_t> vector_t;
195 typedef typename Teuchos::ArrayRCP<scalar_t> vectordata_t;
202 id_t minId = std::numeric_limits<id_t>::max();
204 std::unordered_map<id_t, size_t> uniqueIds;
205 uniqueIds.reserve(nIds);
207 for (
size_t i = 0; i < nIds; i++) {
209 if (
id < minId) minId = id;
210 if (uniqueIds.find(
id) == uniqueIds.end())
219 Teuchos::reduceAll(*comm, Teuchos::REDUCE_MIN, 1, &minId, &indexBase);
223 size_t nUniqueIds = uniqueIds.size();
224 Teuchos::Array<id_t> uniqueIdsList(nUniqueIds);
227 for (
auto it = uniqueIds.begin(); it != uniqueIds.end(); it++)
228 uniqueIdsList[cnt++] = it->first;
232 size_t dummy = Teuchos::OrdinalTraits<Tpetra::global_size_t>::invalid();
234 rcpmap_t idMap = Teuchos::rcp(
new map_t(dummy, uniqueIdsList(), indexBase,
240 vector_t idVec(idMap, 0.);
245 vectordata_t idData = idVec.getDataNonConst();
246 for (
auto it = uniqueIds.begin(); it != uniqueIds.end(); it++) {
247 id_t idx = idMap->getLocalElement(it->first);
248 idData[idx] = it->second;
263 id_t gnUnique = idMap->getMaxAllGlobalIndex() - indexBase + 1;
264 oto_idMap = Teuchos::rcp(
new map_t(gnUnique, indexBase, comm));
268 oto_idMap = Tpetra::createOneToOne(idMap);
271 vector_t oto_idVec(oto_idMap);
277 typedef Tpetra::Export<int, id_t> export_t;
279 export_t idExporter(idMap, oto_idMap);
290 oto_idVec.doExport(idVec, idExporter, Tpetra::ADD);
299 idVec.doImport(oto_idVec, idExporter, Tpetra::REPLACE);
304 size_t cntShared = 0;
306 auto idData = idVec.getDataNonConst();
307 for (
size_t i = 0; i < idVec.getLocalLength(); i++)
308 if (idData[i] > 1) cntShared++;
311 std::cout << comm->getRank() <<
" cntShared = " << cntShared
312 <<
"; nShared = " << nShared << std::endl;
314 return (cntShared == nShared);
324 template <
typename id_t>
328 if (nIds >
size_t(std::numeric_limits<int>::max()))
329 throw std::runtime_error(
"Problem too large for Zoltan_DD");
331 int inIds = int(nIds);
336 MPI_Comm mpicomm = Teuchos::getRawMpiComm(*comm);
340 MPI_Initialized(&flag);
344 MPI_Init(&narg, &argv);
347 MPI_Comm mpicomm = MPI_COMM_WORLD;
352 Zoltan_DD zz(mpicomm, nIdEnt, 0,
sizeof(scalar_t), inIds, 0);
360 std::vector<scalar_t> user(nIds, 1.);
365 ZOLTAN_ID_PTR zgids = NULL;
367 Teuchos::ArrayView<id_t> av(&(ids[0]), nIds);
374 zz.Update(zgids, NULL, (
char*)(nIds ? &(user[0]) : NULL), NULL, inIds);
377 zz.Find(zgids, NULL, (
char*)(nIds ? &(user[0]) : NULL), NULL, inIds, NULL);
381 size_t cntShared = 0;
382 std::unordered_set<id_t> alreadyCounted;
383 for (
size_t i = 0; i < nIds; i++) {
386 if (alreadyCounted.find(ids[i]) == alreadyCounted.end()) {
387 alreadyCounted.insert(ids[i]);
393 if ((nIds * comm->getSize()) <=
TOOMANY) zz.Print();
395 return (cntShared == nShared);
401 Tpetra::ScopeGuard tscope(&narg, &arg);
402 Teuchos::RCP<const Teuchos::Comm<int> > comm = Tpetra::getDefaultComm();
407 float fracShared = 0.2;
411 typedef Tpetra::Map<>::global_ordinal_type
gno_t;
412 IDs<gno_t> myIds(nIds, fracShared, idBase, idStride, comm);
422 float fracShared = 0.2;
426 typedef Tpetra::Map<>::global_ordinal_type
gno_t;
427 IDs<gno_t> myIds(nIds, fracShared, idBase, idStride, comm);
434 #ifdef HAVE_TPETRA_INT_INT
438 float fracShared = 0.2;
442 IDs<int> myIds(nIds, fracShared, idBase, idStride, comm);
452 float fracShared = 0.2;
456 IDs<int> myIds(nIds, fracShared, idBase, idStride, comm);
464 #ifdef HAVE_TPETRA_INT_LONG_LONG
468 float fracShared = 0.4;
void printTpetraThing(T &thing, const std::string &msg, std::ostream &ostr=std::cout)
map_t::global_ordinal_type gno_t
int main(int narg, char **arg)
void printVector(vector_t &vec, const std::string &msg, std::ostream &ostr=std::cout)
void printMap(map_t &map, const std::string &msg, std::ostream &ostr=std::cout)
static const size_t TOOMANY
IDs(size_t nIds_, float fracShared_, id_t idBase_, int idStride_, Teuchos::RCP< const Teuchos::Comm< int > > &comm_)
Traits class to handle conversions between gno_t/lno_t and TPL data types (e.g., ParMETIS's idx_t...
static void ASSIGN_ARRAY(first_t **a, ArrayView< second_t > &b)