[Quant][X86] add an op to compute uint8 batch norm 2d #152811
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/152811
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 6770c6e with merge base 7e637de ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
aten/src/ATen/native/quantized/cpu/kernels/QuantizedOpKernels.cpp
Outdated
Show resolved
Hide resolved
| @@ -407,4 +489,8 @@ TORCH_LIBRARY_IMPL(quantized, QuantizedCPU, m) { | |||
| m.impl(TORCH_SELECTIVE_NAME("quantized::batch_norm3d_relu"), TORCH_FN(q_batch_norm3d_impl<true>)); | |||
| } | |||
|
|
|||
| TORCH_LIBRARY_IMPL(onednn, CPU, m) { | |||
| m.impl(TORCH_SELECTIVE_NAME("onednn::qbatch_norm2d"), TORCH_FN(int8_batch_norm2d_cpu_impl)); | |||
There was a problem hiding this comment.
hmm, even though the op is put under "onednn" namespace but seems the implementation does not rely on onednn at all?
There was a problem hiding this comment.
Yes, that's true. It's because all these ops used for PT2E quantization are put under the onednn namespace, and they should only be used when using the X86Inductor backend. We decide to keep that practice. Do you have any concerns? Thanks.
There was a problem hiding this comment.
Perhaps we can separate the "op" namespace and the quantization backend such as "onednn". For op, "onednn" means that the implementation relies on the oneDNN library and don't have to put other quantized ops under "onednn" namespace. On the other hand, the quantization support with "onednn" quant backend can leverage ops implemented with oneDNN library and native ATen. What do you think?
There was a problem hiding this comment.
Thanks for the suggestion. It makes sense to me. How would you suggest we name the namespace? x86 or still quantized? cc @leslie-fang-intel @jerryzh168
Stack from ghstack (oldest at bottom):
Summary
This PR adds a new op,
onednn.qbatch_norm2d, which accepts uint8 inputs on CPU device (instead of QuantizedCPU).The new ops are implemented with AVX512 instructions and it provides similar performance as its counterpart for QuantizedCPU device
quantized.batch_norm2d.The new op supports output dtypes other than uint8 (fp32, fp16 and bf16 are supported).
Test plan
cc @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10 @jerryzh168