-
Notifications
You must be signed in to change notification settings - Fork 240
/
default_transfer.py
267 lines (220 loc) · 9.87 KB
/
default_transfer.py
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
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
"""Define the default Transfer class."""
from collections import defaultdict
import numpy as np
from openmdao.core.constants import INT_DTYPE
from openmdao.vectors.transfer import Transfer
from openmdao.utils.array_utils import _global2local_offsets
from openmdao.utils.mpi import MPI
_empty_idx_array = np.array([], dtype=INT_DTYPE)
def _merge(indices_list):
if len(indices_list) > 0:
return np.concatenate(indices_list)
else:
return _empty_idx_array
class DefaultTransfer(Transfer):
"""
Default NumPy transfer.
Parameters
----------
in_vec : <Vector>
Pointer to the input vector.
out_vec : <Vector>
Pointer to the output vector.
in_inds : int ndarray
Input indices for the transfer.
out_inds : int ndarray
Output indices for the transfer.
comm : MPI.Comm or <FakeComm>
Communicator of the system that owns this transfer.
"""
@staticmethod
def _setup_transfers(group):
"""
Compute all transfers that are owned by our parent group.
Parameters
----------
group : <Group>
Parent group.
"""
iproc = group.comm.rank
rev = group._mode == 'rev' or group._mode == 'auto'
for subsys in group._subgroups_myproc:
subsys._setup_transfers()
abs2meta = group._var_abs2meta
group._transfers = transfers = {}
vectors = group._vectors
offsets = _global2local_offsets(group._get_var_offsets())
mypathlen = len(group.pathname + '.' if group.pathname else '')
# Initialize empty lists for the transfer indices
xfer_in = []
xfer_out = []
fwd_xfer_in = defaultdict(list)
fwd_xfer_out = defaultdict(list)
if rev:
rev_xfer_in = defaultdict(list)
rev_xfer_out = defaultdict(list)
allprocs_abs2idx = group._var_allprocs_abs2idx
sizes_in = group._var_sizes['input']
offsets_in = offsets['input']
if sizes_in.size > 0:
sizes_in = sizes_in[iproc]
offsets_in = offsets_in[iproc]
offsets_out = offsets['output']
if offsets_out.size > 0:
offsets_out = offsets_out[iproc]
# Loop through all connections owned by this group
for abs_in, abs_out in group._conn_abs_in2out.items():
# This weeds out discrete vars (all vars are local if using this Transfer)
if abs_in in abs2meta['input']:
indices = None
# Get meta
meta_in = abs2meta['input'][abs_in]
idx_in = allprocs_abs2idx[abs_in]
idx_out = allprocs_abs2idx[abs_out]
# Read in and process src_indices
src_indices = meta_in['src_indices']
if src_indices is not None:
src_indices = src_indices.shaped_array()
# 1. Compute the output indices
offset = offsets_out[idx_out]
if src_indices is None:
output_inds = np.arange(offset, offset + meta_in['size'], dtype=INT_DTYPE)
else:
output_inds = src_indices + offset
# 2. Compute the input indices
input_inds = np.arange(offsets_in[idx_in],
offsets_in[idx_in] + sizes_in[idx_in], dtype=INT_DTYPE)
if indices is not None:
input_inds = input_inds.reshape(indices.shape)
# Now the indices are ready - input_inds, output_inds
sub_in = abs_in[mypathlen:].split('.', 1)[0]
fwd_xfer_in[sub_in].append(input_inds)
fwd_xfer_out[sub_in].append(output_inds)
if rev and abs_out in abs2meta['output']:
sub_out = abs_out[mypathlen:].split('.', 1)[0]
rev_xfer_in[sub_out].append(input_inds)
rev_xfer_out[sub_out].append(output_inds)
tot_size = 0
for sname, inds in fwd_xfer_in.items():
fwd_xfer_in[sname] = arr = _merge(inds)
fwd_xfer_out[sname] = _merge(fwd_xfer_out[sname])
tot_size += arr.size
if rev:
for sname, inds in rev_xfer_in.items():
rev_xfer_in[sname] = _merge(inds)
rev_xfer_out[sname] = _merge(rev_xfer_out[sname])
if tot_size > 0:
try:
xfer_in = np.concatenate(list(fwd_xfer_in.values()))
xfer_out = np.concatenate(list(fwd_xfer_out.values()))
except ValueError:
xfer_in = xfer_out = np.zeros(0, dtype=INT_DTYPE)
xfer_all = DefaultTransfer(vectors['input']['nonlinear'],
vectors['output']['nonlinear'], xfer_in, xfer_out,
group.comm)
else:
xfer_all = None
transfers['fwd'] = xfwd = {}
xfwd[None] = xfer_all
if rev:
transfers['rev'] = xrev = {}
xrev[None] = xfer_all
for sname, inds in fwd_xfer_in.items():
if inds.size > 0:
xfwd[sname] = DefaultTransfer(vectors['input']['nonlinear'],
vectors['output']['nonlinear'],
inds, fwd_xfer_out[sname], group.comm)
else:
xfwd[sname] = None
if rev:
for sname, inds in rev_xfer_out.items():
if inds.size > 0:
xrev[sname] = DefaultTransfer(vectors['input']['nonlinear'],
vectors['output']['nonlinear'],
rev_xfer_in[sname], inds, group.comm)
else:
xrev[sname] = None
@staticmethod
def _setup_discrete_transfers(group):
"""
Compute all transfers that are owned by our parent group.
Parameters
----------
group : <Group>
Parent group.
"""
group._discrete_transfers = transfers = defaultdict(list)
name_offset = len(group.pathname) + 1 if group.pathname else 0
iproc = group.comm.rank
owns = group._owning_rank
for tgt, src in group._conn_discrete_in2out.items():
src_sys, src_var = src[name_offset:].split('.', 1)
tgt_sys, tgt_var = tgt[name_offset:].split('.', 1)
xfer = (src_sys, src_var, tgt_sys, tgt_var)
transfers[tgt_sys].append(xfer)
if group.comm.size > 1:
# collect all xfers for each tgt system
for tgt, src in group._conn_discrete_in2out.items():
src_sys, src_var = src[name_offset:].split('.', 1)
tgt_sys, tgt_var = tgt[name_offset:].split('.', 1)
xfer = (src_sys, src_var, tgt_sys, tgt_var)
transfers[tgt_sys].append(xfer)
total_send = set()
total_recv = []
total_xfers = []
for tgt_sys, xfers in transfers.items():
send = set()
recv = []
for src_sys, src_var, tgt_sys, tgt_var in xfers:
if group.pathname:
src_abs = '.'.join([group.pathname, src_sys, src_var])
else:
src_abs = '.'.join([src_sys, src_var])
tgt_rel = '.'.join((tgt_sys, tgt_var))
src_rel = '.'.join((src_sys, src_var))
if iproc == owns[src_abs]:
# we own this var, so we'll send it out to others
send.add(src_rel)
if (tgt_rel in group._var_discrete['input'] and
src_rel not in group._var_discrete['output']):
# we have the target locally, but not the source, so we need someone
# to send it to us.
recv.append(src_rel)
transfers[tgt_sys] = (xfers, send, recv)
total_xfers.extend(xfers)
total_send.update(send)
total_recv.extend(recv)
transfers[None] = (total_xfers, total_send, total_recv)
# find out all ranks that need to receive each discrete source var
allproc_xfers = group.comm.allgather(transfers)
allprocs_recv = defaultdict(lambda: defaultdict(list))
for rank, rank_transfers in enumerate(allproc_xfers):
for tgt_sys, (_, _, recvs) in rank_transfers.items():
for recv in recvs:
allprocs_recv[tgt_sys][recv].append(rank)
group._allprocs_discrete_recv = allprocs_recv
# if we own a src var but it's local for every rank, we don't need to send it to anyone.
total_send = total_send.intersection(allprocs_recv)
for tgt_sys in transfers:
xfers, send, _ = transfers[tgt_sys]
# update send list to remove any vars that don't have a remote receiver,
# and get rid of recv list because allprocs_recv has the necessary info.
transfers[tgt_sys] = (xfers, send.intersection(allprocs_recv[tgt_sys]))
def _transfer(self, in_vec, out_vec, mode='fwd'):
"""
Perform transfer.
Parameters
----------
in_vec : <Vector>
pointer to the input vector.
out_vec : <Vector>
pointer to the output vector.
mode : str
'fwd' or 'rev'.
"""
if mode == 'fwd':
# this works whether the vecs have multi columns or not due to broadcasting
in_vec.set_val(out_vec.asarray()[self._out_inds.flat], self._in_inds)
else: # rev
out_vec.iadd(np.bincount(self._out_inds, in_vec._get_data()[self._in_inds],
minlength=out_vec._data.size))