cpu_channel_shuffle_cl Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/ChannelShuffleKernel.cpp lines 60–85

template <typename scalar_t>
void cpu_channel_shuffle_cl(
    TensorBase& output,
    const TensorBase& input,
    int64_t groups) {
  auto input_data = input.data_ptr<scalar_t>();
  auto output_data = output.data_ptr<scalar_t>();

  int64_t nbatch = input.size(0);
  int64_t channels = input.size(1);
  int64_t channels_per_group = channels / groups;
  int64_t image_size = input.numel() / nbatch / channels;

  // 4d: parallel on dimension of n, h, w
  // 5d: parallel on dimension of n, d, h, w
  at::parallel_for(0, nbatch * image_size, 0, [&](int64_t begin, int64_t end) {
    for (const auto i : c10::irange(begin, end)) {
      scalar_t* output_ptr = output_data + i * channels;
      scalar_t* input_ptr = input_data + i * channels;

      // transpose each channel lane:
      // from [groups, channels_per_group] to [channels_per_group, groups]
      utils::transpose(groups, channels_per_group, input_ptr, channels_per_group, output_ptr, groups);
    }
  });
}