Unify Quantization APIs for add, pool and relu

[supriyar]

Stack from ghstack:

• Unify Quantization APIs for add, pool and relu #26335 Unify Quantization APIs for add, pool and relu
• Changes to support int8 weight and fp32 bias in QNNPACK #26307 Changes to support int8 weight and fp32 bias in QNNPACK

Summary:
Use the backend engine flag to call QNNPACK for quantized ops.

Test Plan:
python test/test_quantized.py TestQNNPACKOps

Reviewers:

Subscribers:

Tasks:

Tags:

Differential Revision: D17504331

[raghuramank100]

Do you need a check for input layout here? How do you know that the input is nhwc contiguous?

[supriyar]

I'll make the change similar to #26242 once that lands for all qnnpack ops.

[t-vi]

It might well be me, but on raspbian (32bit armv7 hf) and USE_PYTORCH_QNNPACK=1, I get None from conv_prepack called like this:

r = torch.ones([2,2,2,2], dtype=torch.float)
scale, zero_point = 1, 2
q = torch.quantize_linear(r, scale, zero_point, torch.qint8)
torch.ops.quantized.conv_prepack(q, None, [1,1], [0,0], [1,1], 1)

I do get some torch.uint8 tensor when I do the same with a vanilla build on x86 (i.e. with FBGEMM).

[supriyar]

It might well be me, but on raspbian (32bit armv7 hf) and USE_PYTORCH_QNNPACK=1, I get None from conv_prepack called like this:

r = torch.ones([2,2,2,2], dtype=torch.float)
scale, zero_point = 1, 2
q = torch.quantize_linear(r, scale, zero_point, torch.qint8)
torch.ops.quantized.conv_prepack(q, None, [1,1], [0,0], [1,1], 1)

I do get some torch.uint8 tensor when I do the same with a vanilla build on x86 (i.e. with FBGEMM).

Did you try setting torch.backends.quantized.engine = torch.qnnpack? You can set the same on server side (x86) as well to confirm qnnpack works.

[t-vi]

Indeed, this fixes the example case and seems to return something in the quantized model. Thank you!

At some point, conv_prepack might raise a more descriptive exception except of just returning None... 🙂
Should the engine default to qnnpack when it available but fbgemm isn't?

[supriyar]

@dzhulgakov please review this. I changed it to permute input tensor to NHWC similar to your changes for qconv.

[dzhulgakov]

Yep, it should be fine

[dzhulgakov]

Looks good, I'll let Jerry to review in the details.

Would be really nice to unify the tests too and extend the benchmarks in benchmarks/op_benchmarks to go through both engines

[dzhulgakov]

Yep, it should be fine

[jerryzh168]

LGTM

[facebook-github-bot]

This pull request has been merged in f337459.