code: def backward(ctx, grad_o):
      if ctx.spatial_group_size > 1:
          out1_pad = ctx.saved_tensors[-1]
          outputs = ctx.saved_tensors[-4:-1]
      else:
          outputs = ctx.saved_tensors[-3:]
  
      if ctx.downsample:
          grad_conv3, grad_conv4 = drelu_dscale2(
              grad_o,
              outputs[2],
              ctx.saved_tensors[6],
              ctx.saved_tensors[11],
          )
      else:
          grad_conv3, grad_conv4 = drelu_dscale1(
              grad_o, outputs[2], ctx.saved_tensors[6]
          )
  
      # create input vector for backward
      t_list = [*ctx.saved_tensors[0:10]]
      t_list.append(grad_conv3)
      t_list.append(grad_conv4)
  
      # outputs used for wgrad and generating drelu mask
      t_list.append(outputs[0])
      t_list.append(outputs[1])
  
      # in case there is downsample
      if ctx.downsample:
          t_list.append(ctx.saved_tensors[10])
  
      grads = fast_bottleneck.backward_init(
          ctx.explicit_nhwc, ctx.stride_1x1, t_list
      )
      wgrad3_stream = torch.cuda.Stream()
      wgrad3_stream.wait_stream(torch.cuda.current_stream())
      grad_out2 = fast_bottleneck.backward_grad_out2(
          ctx.explicit_nhwc, ctx.stride_1x1, t_list, grads
      )
      wgrad2_stream = torch.cuda.Stream()
      wgrad2_stream.wait_stream(torch.cuda.current_stream())
      # do halo exchange of grad_out2 here
      # compute halo cells for grad_out1
      if ctx.spatial_group_size > 1:
          if ctx.explicit_nhwc:
              N, Hs, W, C = list(grad_out2.shape)
          else:
              N, C, Hs, W = list(grad_out2.shape)
          relu1 = t_list[12]
          ctx.stream1.wait_stream(torch.cuda.current_stream())
          with torch.cuda.stream(ctx.stream1):
              (
                  top_halo,
                  btm_halo,
              ) = ctx.spatial_halo_exchanger.left_right_halo_exchange(
                  grad_out2[:, :1, :, :], grad_out2[:, Hs - 1 :, :, :]
              )
              # copy halos to send buffer
          if ctx.spatial_method == 1 or ctx.spatial_method == 2:
              # 1 -> halo recompute approach
              # 2 -> wait for concatenated halos, then do single conv on full input (not implemented yet for bprop)
              if ctx.spatial_group_rank < ctx.spatial_group_size - 1:
                  ctx.stream2.wait_stream(ctx.stream1)
                  with torch.cuda.stream(ctx.stream2):
                      if ctx.explicit_nhwc:
                          btm_fat_halo = torch.empty(
                              (N, 3, W, C),
                              dtype=grad_out2.dtype,
                              device=grad_out2.device,
                          )
                          btm_fat_halo[:, :2, :, :].copy_(
                              grad_out2[:, Hs - 2 :, :, :]
                          )
                          btm_fat_halo[:, 2:, :, :].copy_(btm_halo)
                          btm_fat_relu_halo = torch.empty(
                              (N, 3, W, C),
                              dtype=grad_out2.dtype,
                              device=grad_out2.device,
                          )
                          btm_fat_relu_halo[:, :2, :, :].copy_(
                              relu1[:, Hs - 2 :, :, :]
                          )
                          btm_fat_relu_halo[:, 2:, :, :].zero_()
                      else:
                          btm_fat_halo = torch.empty(
                              (N, C, 3, W),
                              dtype=grad_out2.dtype,
                              device=grad_out2.device,
                          )
                          btm_fat_halo[:, :, :2, :].copy_(
                              grad_out2[:, :, Hs - 2 :, :]
                          )
                          btm_fat_halo[:, :, 2:, :].copy_(btm_halo)
                          btm_fat_relu_halo = torch.empty(
                              (N, C, 3, W),
                              dtype=grad_out2.dtype,
                              device=grad_out2.device,
                          )
                          btm_fat_relu_halo[:, :, :2, :].copy_(
                              relu1[:, :, Hs - 2 :, :]
                          )
                          btm_fat_relu_halo[:, :, 2:, :].zero_()
                      btm_grad_out1_halo = (
                          fast_bottleneck.backward_grad_out1_halo(
                              ctx.explicit_nhwc,
                              ctx.stride_1x1,
                              t_list,
                              grads,
                              btm_fat_halo,
                              btm_fat_relu_halo,
                          )
                      )
                      if ctx.explicit_nhwc:
                          btm_grad_out1_halo = btm_grad_out1_halo[
                              :, 1:2, :, :
                          ]
                      else:
                          btm_grad_out1_halo = btm_grad_out1_halo[
                              :, :, 1:2, :
                          ]
              if ctx.spatial_group_rank > 0:
                  with torch.cuda.stream(ctx.stream1):
                      if ctx.explicit_nhwc:
                          top_fat_halo = torch.empty(
                              (N, 3, W, C),
                              dtype=grad_out2.dtype,
                              device=grad_out2.device,
                          )
                          top_fat_halo[:, :1, :, :].copy_(top_halo)
                          top_fat_halo[:, 1:, :, :].copy_(
                              grad_out2[:, :2, :, :]
                          )
                          top_fat_relu_halo = torch.empty(
                              (N, 3, W, C),
                              dtype=grad_out2.dtype,
                              device=grad_out2.device,
                          )
                          top_fat_relu_halo[:, :1, :, :].zero_()
                          top_fat_relu_halo[:, 1:, :, :].copy_(
                              relu1[:, :2, :, :]
                          )
                      else:
                          top_fat_halo = torch.empty(
                              (N, C, 3, W),
                              dtype=grad_out2.dtype,
                              device=grad_out2.device,
                          )
                          top_fat_halo[:, :, :1, :].copy_(top_halo)
                          top_fat_halo[:, :, 1:, :].copy_(
                              grad_out2[:, :, :2, :]
                          )
                          top_fat_relu_halo = torch.empty(
                              (N, C, 3, W),
                              dtype=grad_out2.dtype,
                              device=grad_out2.device,
                          )
                          top_fat_relu_halo[:, :, :1, :].zero_()
                          top_fat_relu_halo[:, :, 1:, :].copy_(
                              relu1[:, :, :2, :]
                          )
                      top_grad_out1_halo = (
                          fast_bottleneck.backward_grad_out1_halo(
                              ctx.explicit_nhwc,
                              ctx.stride_1x1,
                              t_list,
                              grads,
                              top_fat_halo,
                              top_fat_relu_halo,
                          )
                      )
                      if ctx.explicit_nhwc:
                          top_grad_out1_halo = top_grad_out1_halo[
                              :, 1:2, :, :
                          ]
                      else:
                          top_grad_out1_halo = top_grad_out1_halo[
                              :, :, 1:2, :
                          ]
              if ctx.use_delay_kernel:
                  inc.add_delay(10)
          elif ctx.spatial_method != 3:
              assert False, "spatial_method must be 1, 2 or 3"
  
      # compute grad_out1 for internal cells
      if (
          ctx.spatial_group_size <= 1
          or ctx.spatial_method == 1
          or ctx.spatial_method == 2
      ):
          grad_out1 = fast_bottleneck.backward_grad_out1(
              ctx.explicit_nhwc,
              ctx.stride_1x1,
              t_list,
              grads,
              grad_out2,
          )
      elif ctx.spatial_group_size > 1 and ctx.spatial_method == 3:
          grad_out1 = fast_bottleneck.backward_grad_out1_mask(
              ctx.explicit_nhwc,
              ctx.stride_1x1,
              t_list,
              grads,
              grad_out2,
              ctx.thresholdTop,
              ctx.thresholdBottom,
          )
  
      # apply halo cells to grad_out1
      if ctx.spatial_group_size > 1:
          w = t_list[2]
          z = t_list[4]
          relu1 = t_list[12]
          # print("w.shape = %s, z.shape = %s, relu1.shape = %s" % (str(list(w.shape)), str(list(z.shape)), str(list(relu1.shape))))
          if ctx.spatial_method == 1 or ctx.spatial_method == 2:
              if ctx.spatial_group_rank < ctx.spatial_group_size - 1:
                  torch.cuda.current_stream().wait_stream(ctx.stream2)
                  if ctx.explicit_nhwc:
                      grad_out1[:, Hs - 1 :, :, :].copy_(
                          btm_grad_out1_halo
                      )
                  else:
                      grad_out1[:, :, Hs - 1 :, :].copy_(
                          btm_grad_out1_halo
                      )
                  # print("ctx.spatial_group_rank = %d, apply grad_out1 btm halo (grad_out1.shape = %s)" % (ctx.spatial_group_rank, str(list(grad_out1.shape))))
              if ctx.spatial_group_rank > 0:
                  torch.cuda.current_stream().wait_stream(ctx.stream1)
                  if ctx.explicit_nhwc:
                      grad_out1[:, :1, :, :].copy_(top_grad_out1_halo)
                  else:
                      grad_out1[:, :, :1, :].copy_(top_grad_out1_halo)
                  # print("ctx.spatial_group_rank = %d, apply grad_out1 top halo (grad_out1.shape = %s)" % (ctx.spatial_group_rank, str(list(grad_out1.shape))))
          elif ctx.spatial_method == 3:
              if ctx.spatial_group_rank < ctx.spatial_group_size - 1:
                  if ctx.explicit_nhwc:
                      btm_relu_halo = relu1[:, Hs - 1 :, :, :].clone()
                      btm_grad_out1 = grad_out1[:, Hs - 1 :, :, :]
                  else:
                      btm_relu_halo = relu1[:, :, Hs - 1 :, :].clone()
                      btm_grad_out1 = grad_out1[:, :, Hs - 1 :, :]
                  w1by3 = w[:, :1, :, :].clone()
                  ctx.stream2.wait_stream(
                      ctx.stream1
                  )  # wait for halo transfers to finish
                  ctx.stream2.wait_stream(
                      torch.cuda.current_stream()
                  )  # wait for backward_grad_out1_mask to finish before launching halo correction kernel
                  with torch.cuda.stream(ctx.stream2):
                      btm_grad_out1_halo = (
                          fast_bottleneck.backward_grad_out1_halo_corr(
                              ctx.explicit_nhwc,
                              ctx.stride_1x1,
                              t_list,
                              w1by3,
                              grads,
                              btm_halo,
                              btm_relu_halo,
                              btm_grad_out1.clone(),
                          )
                      )
                      btm_grad_out1.copy_(btm_grad_out1_halo)
              if ctx.spatial_group_rank > 0:
                  if ctx.explicit_nhwc:
                      top_relu_halo = relu1[:, :1, :, :].clone()
                      top_grad_out1 = grad_out1[:, :1, :, :]
                  else:
                      top_relu_halo = relu1[:, :, :1, :].clone()
                      top_grad_out1 = grad_out1[:, :, :1, :]
                  w1by3 = w[:, 2:, :, :].clone()
                  ctx.stream1.wait_stream(
                      torch.cuda.current_stream()
                  )  # wait for backward_grad_out1_mask to finish before launching halo correction kernel
                  with torch.cuda.stream(ctx.stream1):
                      top_grad_out1_halo = (
                          fast_bottleneck.backward_grad_out1_halo_corr(
                              ctx.explicit_nhwc,
                              ctx.stride_1x1,
                              t_list,
                              w1by3,
                              grads,
                              top_halo,
                              top_relu_halo,
                              top_grad_out1.clone(),
                          )
                      )
                      top_grad_out1.copy_(top_grad_out1_halo)
              if ctx.spatial_group_rank < ctx.spatial_group_size - 1:
                  torch.cuda.current_stream().wait_stream(
                      ctx.stream2
                  )  # wait for halo correction to finish
              if ctx.spatial_group_rank > 0:
                  torch.cuda.current_stream().wait_stream(ctx.stream1)
  
      wgrad1_stream = torch.cuda.Stream()
      wgrad1_stream.wait_stream(torch.cuda.current_stream())
      fast_bottleneck.backward_rest(
          ctx.explicit_nhwc,
          ctx.stride_1x1,
          t_list,
          grads,
          grad_out2,
          grad_out1,
      )
      with torch.cuda.stream(wgrad3_stream):
          fast_bottleneck.backward_wgrad3(
              ctx.explicit_nhwc, ctx.stride_1x1, t_list, grads
          )
      with torch.cuda.stream(wgrad2_stream):
          if ctx.spatial_group_size > 1:
              fast_bottleneck.backward_wgrad2_pad(
                  ctx.explicit_nhwc,
                  ctx.stride_1x1,
                  t_list,
                  grads,
                  out1_pad,
                  grad_out2,
              )
          else:
              fast_bottleneck.backward_wgrad2(
                  ctx.explicit_nhwc,
                  ctx.stride_1x1,
                  t_list,
                  grads,
                  grad_out2,
              )
      with torch.cuda.stream(wgrad1_stream):
          fast_bottleneck.backward_wgrad1(
              ctx.explicit_nhwc,
              ctx.stride_1x1,
              t_list,
              grads,
              grad_out1,
          )
      torch.cuda.current_stream().wait_stream(wgrad3_stream)
      torch.cuda.current_stream().wait_stream(wgrad2_stream)
      torch.cuda.current_stream().wait_stream(wgrad1_stream)
  
      return (
          None,
          None,
          None,
          None,
          None,
          None,
          None,
          None,
          None,
          None,
          None,
          None,
          *grads,
      )
  
imports: import functools
  import torch
  import torch.distributed
  from torch import nn
  from apex import check_cudnn_version_and_warn
  import fast_bottleneck
  import nccl_p2p_cuda
label_desc: Use image node_id 2 for calling the SpatialBottleneckFunction obj's underlying member method code with example usage: obj.backward(ctx, grad_o) and returns: None,  None,  None,  None,  None,  None,  None,  None,  None,  None,  None,  None,  unknown