adaptive_avg_pool_single_out_frame Class — pytorch Architecture

Source Code

aten/src/ATen/native/quantized/cpu/AdaptiveAveragePooling.cpp lines 56–130

template <typename scalar_t>
void adaptive_avg_pool_single_out_frame(
    scalar_t* input_p,
    scalar_t* output_p,
    int64_t sizeC,
    int64_t isizeD, // Set to 1 for 2D
    int64_t isizeH,
    int64_t isizeW,
    int64_t osizeD, // Set to 1 for 2D
    int64_t osizeH,
    int64_t osizeW,
    int64_t istrideC,
    int64_t istrideD,  // Set to 1 for 2D
    int64_t istrideH,
    int64_t istrideW) {
  at::parallel_for(0, sizeC, 0, [&](int64_t start, int64_t end) {
    for (const auto c : c10::irange(start, end)) {
      /* loop over output */
      for (int64_t od = 0; od < osizeD; od++) {
        int istartD = start_index(od, osizeD, isizeD);
        int iendD = end_index(od, osizeD, isizeD);
        int kD = iendD - istartD;
        // NOLINTNEXTLINE(cppcoreguidelines-narrowing-conversions,bugprone-narrowing-conversions)
        float kDr = 1.0 / kD;
        for (int64_t oh = 0; oh < osizeH; oh++) {
          int istartH = start_index(oh, osizeH, isizeH);
          int iendH = end_index(oh, osizeH, isizeH);
          int kH = iendH - istartH;
          // NOLINTNEXTLINE(cppcoreguidelines-narrowing-conversions,bugprone-narrowing-conversions)
          float kDHr = kDr / kH;

          for (int64_t ow = 0; ow < osizeW; ow++) {
            int istartW = start_index(ow, osizeW, isizeW);
            int iendW = end_index(ow, osizeW, isizeW);
            int kW = iendW - istartW;
            // NOLINTNEXTLINE(cppcoreguidelines-narrowing-conversions,bugprone-narrowing-conversions)
            float kDHWr = kDHr / kW;

            /* local pointers */
            scalar_t* ip = input_p +
                           c * istrideC +
                           istartD * istrideD +
                           istartH * istrideH +
                           istartW * istrideW;
            scalar_t* op = output_p +
                           c * osizeD * osizeH * osizeW +
                           od * osizeH * osizeW +
                           oh * osizeW +
                           ow;

            /* compute local average: */
            int64_t sum = 0;
            for (int id = 0; id < kD; id++) {
              for (int ih = 0; ih < kH; ih++) {
                for (int iw = 0; iw < kW; iw++) {
                  // NOLINTNEXTLINE(bugprone-signed-char-misuse)
                  int64_t val = (ip +
                                 id * istrideD +
                                 ih * istrideH +
                                 iw * istrideW)->val_;
                  sum += val;
                }
              }
            }

            /* set output to local average */
            // TODO: add the max/min clip
            op->val_ = static_cast<typename scalar_t::underlying>(
                std::nearbyint(sum * kDHWr));
          } // ow
        } // oh
      } // od
    }
  });
}