When quantized::max_pool2d is used?

masahi · January 26, 2020, 11:10pm

In my IR that came from the quantized resnet model from torchvision, I have one max pool op which operates on quantized tensors:

  %input.3 : QUInt8(1, 64, 56, 56) = aten::max_pool2d(%input.2, %1279, %1282, %1285, %1288, %1289), scope: __module.maxpool # /home/masa/anaconda3/lib/python3.7/site-packages/torch/nn/functional.py:488:0

Does this dispatch into quantized::max_pool2d below?

github.com

pytorch/pytorch/blob/master/aten/src/ATen/native/quantized/cpu/qpool.cpp#L414


    #ifdef USE_PYTORCH_QNNPACK
    if (at::globalContext().qEngine() == at::QEngine::QNNPACK && qx.scalar_type() == kQUInt8) {
      return qnnpack_maxpool(qx, kernel_size, stride, padding, dilation, ceil_mode);
    }
    #endif
    return at::max_pool2d(qx, kernel_size, stride, padding, dilation, ceil_mode);
  }
};


static auto registry = torch::RegisterOperators().op(
    "quantized::max_pool2d(Tensor qx, "
    "int[] kernel_size, "
    "int[] stride, "
    "int[] padding, "
    "int[] dilation,"
    "bool ceil_mode) -> Tensor",
    torch::RegisterOperators::options().kernel<QMaxPool2D_arr_args>(
        DispatchKey::QuantizedCPUTensorId));


} // namespace
} // namespace native

Or more generally, I want to know about aten dispatch mechanism. Any pointer is appreciated.

cc @raghuramank100 @jerryzh168

jerryzh168 · February 14, 2020, 6:48pm

Yes, that is correct. it is dispatch here: https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/Pooling.cpp#L128
We have multiple ways to do dispatch right now in PyTorch, one common place is in native_functions.yaml, you can take a look at: https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/README.md