1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
|
// SPDX-FileCopyrightText: Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
// SPDX-License-Identifier: BSD-3-Clause
#include "vtkDIYExplicitAssigner.h"
#include "vtkMath.h"
#include "vtkObjectFactory.h"
#include <algorithm>
#include <cassert>
#include <iterator>
#include <numeric>
//------------------------------------------------------------------------------
VTK_ABI_NAMESPACE_BEGIN
vtkDIYExplicitAssigner::vtkDIYExplicitAssigner(
diy::mpi::communicator comm, int local_blocks, bool force_power_of_two /*=false*/)
: diy::StaticAssigner(comm.size(), local_blocks)
{
std::vector<int> block_counts;
if (comm.size() > 1)
{
block_counts.resize(comm.size());
diy::mpi::all_gather(comm, local_blocks, block_counts);
}
else
{
block_counts.push_back(local_blocks);
}
assert(block_counts.size() >= 1);
if (force_power_of_two)
{
const int global_num_blocks = std::accumulate(block_counts.begin(), block_counts.end(), 0);
const int global_block_counts_pow_2 = vtkMath::NearestPowerOfTwo(global_num_blocks);
// we pad each rank with extra blocks
auto extra_blocks = global_block_counts_pow_2 - global_num_blocks;
const auto extra_blocks_per_rank =
static_cast<int>(std::ceil(extra_blocks / static_cast<double>(block_counts.size())));
for (auto& count : block_counts)
{
if (extra_blocks > 0)
{
const auto padding = std::min(extra_blocks_per_rank, extra_blocks);
count += padding;
extra_blocks -= padding;
}
}
assert(
std::accumulate(block_counts.begin(), block_counts.end(), 0) == global_block_counts_pow_2);
}
// convert to inclusive-scan
this->IScanBlockCounts = std::move(block_counts);
for (size_t cc = 1; cc < this->IScanBlockCounts.size(); ++cc)
{
this->IScanBlockCounts[cc] += this->IScanBlockCounts[cc - 1];
}
this->set_nblocks(this->IScanBlockCounts.back());
assert(
force_power_of_two == false || vtkMath::NearestPowerOfTwo(this->nblocks()) == this->nblocks());
}
//------------------------------------------------------------------------------
int vtkDIYExplicitAssigner::rank(int gid) const
{
auto iter =
std::lower_bound(this->IScanBlockCounts.begin(), this->IScanBlockCounts.end(), gid + 1);
assert(iter != this->IScanBlockCounts.end());
return static_cast<int>(std::distance(this->IScanBlockCounts.begin(), iter));
}
//------------------------------------------------------------------------------
void vtkDIYExplicitAssigner::local_gids(int rank, std::vector<int>& gids) const
{
const auto min = rank == 0 ? 0 : this->IScanBlockCounts[rank - 1];
const auto max = this->IScanBlockCounts[rank];
gids.resize(max - min);
std::iota(gids.begin(), gids.end(), min);
}
VTK_ABI_NAMESPACE_END
|
