CopyICFirst3dTensorToChannelsLast3dTensor Class — pytorch Architecture

Source Code

aten/src/ATen/native/quantized/cpu/fbgemm_utils.cpp lines 71–93

template <typename T>
void CopyICFirst3dTensorToChannelsLast3dTensor(
    int64_t G,
    int64_t IC_G,
    int64_t OC_G,
    int64_t D,
    int64_t H,
    int64_t W,
    const T* src,
    T* dst) {
  // IC OC/G THW -> G OC/G THW IC/G
  const int64_t inner_size = D * H * W;
  for (int64_t i = 0; i < G * OC_G; ++i) {
    for (const auto j : c10::irange(inner_size)) {
      for (const auto ic : c10::irange(IC_G)) {
        int g = static_cast<int>(i / OC_G);
        int oc = static_cast<int>(i % OC_G);
        dst[(i * inner_size + j) * IC_G + ic] =
            src[((g * IC_G + ic) * OC_G + oc) * inner_size + j];
      }
    }
  }
}