Build global data on all processes.
(self, layout=None)
| 780 | self.change_scales(self.domain.dealias) |
| 781 | |
| 782 | def allgather_data(self, layout=None): |
| 783 | """Build global data on all processes.""" |
| 784 | # Change layout |
| 785 | if layout is not None: |
| 786 | self.change_layout(layout) |
| 787 | # Shortcut for serial execution |
| 788 | if self.dist.comm.size == 1: |
| 789 | return self.data.copy() |
| 790 | # Build global buffers |
| 791 | tensor_shape = tuple(cs.dim for cs in self.tensorsig) |
| 792 | global_shape = tensor_shape + self.layout.global_shape(self.domain, self.scales) |
| 793 | local_slices = tuple(slice(None) for cs in self.tensorsig) + self.layout.slices(self.domain, self.scales) |
| 794 | send_buff = np.zeros(shape=global_shape, dtype=self.dtype) |
| 795 | recv_buff = np.empty_like(send_buff) |
| 796 | # Combine data via allreduce -- easy but not communication-optimal |
| 797 | # Should be optimized using Allgatherv if this is used past startup |
| 798 | send_buff[local_slices] = self.data |
| 799 | self.dist.comm.Allreduce(send_buff, recv_buff, op=MPI.SUM) |
| 800 | return recv_buff |
| 801 | |
| 802 | def gather_data(self, root=0, layout=None): |
| 803 | # Change layout |
no test coverage detected