[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 ab2a1c8 with merge base f5e0806 ( 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
There was a problem hiding this comment.
Hi @jerryzh168 @leslie-fang-intel Do you think it's ok to use onednn namespace even if the we don't use oneDNN library for implementation? If not, do you have any suggestions on name of the namespace? Thanks.
There was a problem hiding this comment.
btw, it could be a long-term refactoring work, not necessarily blocking this PR landing.
There was a problem hiding this comment.
it can be x86 I think, since it's used in x86 inductor backend, we are deprecating the quantized namespace as we deprecate eager and fx graph mode quant in pytorch
There was a problem hiding this comment.
@jgong5 @jerryzh168 Thanks for your suggestions. I will keep this PR as is and refactor this part in a future PR.
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
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