code: def nccl_halo_ex(
      peer_rank, peer_group_size, y, half_halo, explicit_nhwc, H_split
  ):
      if explicit_nhwc:
          if H_split:
              _, Hp, _, _ = list(y.shape)
              H = Hp - 2 * half_halo
              top_out_halo = y[:, half_halo : 2 * half_halo, :, :]
              top_inp_halo = y[:, :half_halo, :, :]
              btm_out_halo = y[:, H : H + half_halo, :, :]
              btm_inp_halo = y[
                  :, H + half_halo : H + 2 * half_halo, :, :
              ]
          else:
              _, _, Wp, _ = list(y.shape)
              W = Wp - 2 * half_halo
              top_out_halo = y[:, :, half_halo : 2 * half_halo, :]
              top_inp_halo = y[:, :, :half_halo, :]
              btm_out_halo = y[:, :, W : W + half_halo, :]
              btm_inp_halo = y[
                  :, :, W + half_halo : W + 2 * half_halo, :
              ]
      else:
          if H_split:
              _, _, Hp, _ = list(y.shape)
              H = Hp - 2 * half_halo
              top_out_halo = y[:, :, half_halo : 2 * half_halo, :]
              top_inp_halo = y[:, :, :half_halo, :]
              btm_out_halo = y[:, :, H : H + half_halo, :]
              btm_inp_halo = y[
                  :, :, H + half_halo : H + 2 * half_halo, :
              ]
          else:
              _, _, _, Wp = list(y.shape)
              W = Wp - 2 * half_halo
              top_out_halo = y[:, :, :, half_halo : 2 * half_halo]
              top_inp_halo = y[:, :, :, :half_halo]
              btm_out_halo = y[:, :, :, W : W + half_halo]
              btm_inp_halo = y[
                  :, :, :, W + half_halo : W + 2 * half_halo
              ]
  
      mf = (
          torch.channels_last
          if y.is_contiguous(memory_format=torch.channels_last)
          else torch.contiguous_format
      )
      top_out_halo = top_out_halo.contiguous()
      btm_out_halo = btm_out_halo.contiguous()
  
      top_inp_halos = [
          torch.empty_like(top_out_halo) for _ in range(peer_group_size)
      ]
      torch.distributed.all_gather(top_inp_halos, top_out_halo)
      btm_inp_halos = [
          torch.empty_like(btm_out_halo) for _ in range(peer_group_size)
      ]
      torch.distributed.all_gather(btm_inp_halos, btm_out_halo)
      top_rank = (peer_rank + peer_group_size - 1) % peer_group_size
      btm_rank = (peer_rank + 1) % peer_group_size
      if peer_rank == 0:
          top_inp_halo.zero_()
      else:
          top_inp_halo.copy_(
              btm_inp_halos[top_rank].to(memory_format=mf)
          )
      if peer_rank == peer_group_size - 1:
          btm_inp_halo.zero_()
      else:
          btm_inp_halo.copy_(
              top_inp_halos[btm_rank].to(memory_format=mf)
          )
  
file_path: apex/apex/contrib/test/peer_memory/test_peer_halo_exchange_module.py
imports: import unittest
  import torch
  from torch.testing._internal import common_utils
  from apex.transformer.testing.distributed_test_base import NcclDistributedTestBase
label_desc: Use image node_id 1 for calling a global function with example usage: nccl_halo_ex(peer_rank, peer_group_size, y, half_halo, explicit_nhwc, H_split) without return types