code: def single_test(
      peer_rank,
      peer_group_size,
      halo_ex,
      C,
      H,
      W,
      half_halo,
      dtype,
      memory_format,
      H_split,
      num_steps,
      numSM=1,
  ):
      if memory_format == 1:
          # 1 -> explicit nhwc
          explicit_nhwc = True
          if H_split:
              y = torch.randn(
                  [1, H + 2 * half_halo, W, C],
                  dtype=dtype,
                  device="cuda",
              )
              ym = y[:, half_halo : H + half_halo, :, :]
          else:
              y = torch.randn(
                  [1, H, W + 2 * half_halo, C],
                  dtype=dtype,
                  device="cuda",
              )
              ym = y[:, :, half_halo : W + half_halo, :]
      else:
          # 2 -> native nhwc
          # 3 -> nchw
          explicit_nhwc = False
          if H_split:
              y = torch.randn(
                  [1, C, H + 2 * half_halo, W],
                  dtype=dtype,
                  device="cuda",
              )
              if memory_format == 2:
                  y = y.to(memory_format=torch.channels_last)
              ym = y[:, :, half_halo : H + half_halo, :]
          else:
              y = torch.randn(
                  [1, C, H, W + 2 * half_halo],
                  dtype=dtype,
                  device="cuda",
              )
              if memory_format == 2:
                  y = y.to(memory_format=torch.channels_last)
              ym = y[:, :, :, half_halo : W + half_halo]
      y3 = y.clone()
      list_y = []
      for step in range(num_steps):
          halo_ex(y, H_split, explicit_nhwc, numSM)
          list_y.append(y.clone())
          y.copy_(y3)
          halo_ex.peer_pool.reset()
          torch.distributed.barrier()
      y2 = y3.clone()
      list_y2 = []
      for step in range(num_steps):
          nccl_halo_ex(
              peer_rank,
              peer_group_size,
              y2,
              half_halo,
              explicit_nhwc,
              H_split,
          )
          list_y2.append(y2.clone())
          y2.copy_(y3)
      if memory_format == 1:
          memory_format_str = "explicit_nhwc"
      elif memory_format == 2:
          memory_format_str = "native nhwc"
      elif memory_format == 3:
          memory_format_str = "nchw"
      else:
          memory_format_str = "???"
      torch.testing.assert_close(list_y, list_y2, msg=memory_format_str)
  
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 2 for calling a global function with example usage: single_test(peer_rank, peer_group_size, halo_ex, C, H, W, half_halo, dtype, memory_format, H_split, num_steps, numSM) without return types