adaptive_avg_pool3d_out_frame Class — pytorch Architecture

Source Code

aten/src/ATen/native/AdaptiveAveragePooling3d.cpp lines 26–84

template <typename scalar_t>
void adaptive_avg_pool3d_out_frame(
    const scalar_t* input_p,
    scalar_t* output_p,
    int64_t sizeD,
    int64_t isizeT,
    int64_t isizeH,
    int64_t isizeW,
    int64_t osizeT,
    int64_t osizeH,
    int64_t osizeW,
    int64_t istrideD,
    int64_t istrideT,
    int64_t istrideH,
    int64_t istrideW) {
  at::parallel_for(0, sizeD, 1, [&](int64_t start, int64_t end) {
    for (const auto d : c10::irange(start, end)) {
      /* loop over output */
      for (const auto ot : c10::irange(osizeT)) {
        auto istartT = start_index(ot, osizeT, isizeT);
        auto iendT = end_index(ot, osizeT, isizeT);
        auto kT = iendT - istartT;

        for (const auto oh : c10::irange(osizeH)) {
          auto istartH = start_index(oh, osizeH, isizeH);
          auto iendH = end_index(oh, osizeH, isizeH);
          auto kH = iendH - istartH;

          for (const auto ow : c10::irange(osizeW)) {
            auto istartW = start_index(ow, osizeW, isizeW);
            auto iendW = end_index(ow, osizeW, isizeW);
            auto kW = iendW - istartW;

            /* local pointers */
            const scalar_t* ip = input_p + d * istrideD + istartT * istrideT +
                istartH * istrideH + istartW * istrideW;
            scalar_t* op = output_p + d * osizeT * osizeH * osizeW +
                ot * osizeH * osizeW + oh * osizeW + ow;

            /* compute local average: */
            scalar_t sum = 0;
            for (const auto it : c10::irange(kT)) {
              for (const auto ih : c10::irange(kH)) {
                for (const auto iw : c10::irange(kW)) {
                  scalar_t val =
                      *(ip + it * istrideT + ih * istrideH + iw * istrideW);
                  sum += val;
                }
              }
            }

            /* set output to local average */
            *op = sum / kT / kH / kW;
          }
        }
      }
    }
  });
}