q_maxpool_3d Class — pytorch Architecture

Source Code

aten/src/ATen/native/quantized/cpu/Pooling.cpp lines 343–457

template <typename Q>
Tensor q_maxpool_3d(
    Tensor qx, // Input Tensor (Quantized)
    int64_t kT,
    int64_t kH,
    int64_t kW, // kernel size
    int64_t sT,
    int64_t sH,
    int64_t sW, // strides
    int64_t pT,
    int64_t pH,
    int64_t pW, // padding
    int64_t dT,
    int64_t dH,
    int64_t dW,
    bool ceil_mode) { // dilation
  // Check input dimensions.
  TORCH_CHECK(kT > 0 && kH > 0 && kW > 0, "kernel_size should be greater than zero.");
  TORCH_CHECK(sT > 0 && sH > 0 && sW > 0, "strides should be greater than zero.");
  TORCH_CHECK(
      dT > 0 && dH > 0 && dW > 0,
      "dilation should be greater than zero. "
      "Got (",
      dT,
      ", ",
      dH,
      ", ",
      dW,
      ")");
  int ndim = qx.dim();
  // TODO leslie: Support non batch mode input when input is THWC which is 4-d tensor.
  TORCH_CHECK(ndim == 5, "Expecting the input tensor of rank 5.");

  // act: n, c, t, h, w
  int dimc = 1;
  int dimt = 2;
  int dimh = 3;
  int dimw = 4;
  int nbatch = qx.size(0);
  // Check if inputs are valid.
  int64_t iC = qx.size(dimc);
  int64_t iT = qx.size(dimt);
  int64_t iH = qx.size(dimh);
  int64_t iW = qx.size(dimw);
  TORCH_CHECK(iC > 0 && iT > 0 && iH > 0 && iW > 0, "input dimensions must be non-zero.");
  TORCH_CHECK(
      kT / 2 >= pT && kH / 2 >= pH && kW / 2 >= pW,
      "padding should be smaller than half of kernel_size.");

  // Check output dimensions.
  int64_t oC = iC;
  int64_t oT = pooling_output_shape(iT, kT, pT, sT, dT, ceil_mode);
  int64_t oH = pooling_output_shape(iH, kH, pH, sH, dH, ceil_mode);
  int64_t oW = pooling_output_shape(iW, kW, pW, sW, dW, ceil_mode);
  TORCH_CHECK(oT > 0 && oH > 0 && oW > 0,
              "Given input size: (",
              iC, "t", iT , "x", iH, "x", iW,
              "). Calculated output size: (",
              oC, "t", oT , "x", oH, "x", oW,
              "). Output size is too small.");

  std::vector<int64_t> oSizes = {nbatch, oC, oT, oH, oW};

  if (qx.is_contiguous(c10::MemoryFormat::ChannelsLast3d)) {
    // Fast path case for channels-last case.
    // In this case, we can preserve the data layout in memory
    // as well as use a loop nest that is more amenable to
    // vectorization.
    Tensor qy = at::_empty_affine_quantized(
        oSizes,
        qx.options()
          .dtype(toQIntType(qx.scalar_type()))
          .memory_format(qx.suggest_memory_format()),
        qx.q_scale(),
        qx.q_zero_point(),
        std::nullopt);
    qmaxpool_3d_nthwc_stub(qx.device().type(), qx, iC, iT, iH, iW, oT, oH, oW, kT, kH, kW, sT, sH, sW, pT, pH, pW, dT, dH, dW, qy);
    return qy;
  } else {
    Tensor qy = at::_empty_affine_quantized(
      oSizes,
      qx.options().dtype(toQIntType(qx.scalar_type())),
      qx.q_scale(),
      qx.q_zero_point());
    auto qx_contig = qx.contiguous();
    auto qxd = qx_contig.data_ptr<Q>();
    auto qyd = qy.data_ptr<Q>();

    spatial_dilated_max_pooling3d<Q>(
        qxd,
        nbatch,
        iC,
        iT,
        iH,
        iW,
        oT,
        oH,
        oW,
        kT,
        kH,
        kW,
        sT,
        sH,
        sW,
        pT,
        pH,
        pW,
        dT,
        dH,
        dW,
        qyd);

    return qy;
  }
}