def halo_exchange_forward(self, x):
halo_region_send = x[:, :, :, -self.halo_size:]
if self.comm.rank < 3 :
chainermn.functions.send(halo_region_send, self.comm, rank=self.comm.rank + 1, tag=(self.comm.rank +1)* self.index)
if self.comm.rank > 0 :
received_halo_region = chainermn.functions.recv(self.comm, rank=self.comm.rank -1, tag=self.comm.rank*self.index)
x = concat((x, received_halo_region), axis=-1)
return x
def halo_exchange_backward(self, x):
halo_region_send = x[:, :, :, :self.halo_size]
if self.comm.rank > 0:
sent = chainermn.functions.send(halo_region_send,
self.comm,
rank=self.comm.rank - 1,
tag=(self.comm.rank - 1) *
self.index * 2)
if self.comm.rank < 3:
received_halo_region = chainermn.functions.recv(
self.comm,
rank=self.comm.rank + 1,
tag=self.comm.rank * self.index * 2)
x = concat((received_halo_region, x), axis=-1)
return x