When quantized::max_pool2d is used?

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?

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

cc @raghuramank100 @jerryzh168

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

1 Like