[Inductor][CPU] Add GEMM templates for _weight_int4pack_mm_for_cpu with AVX512 #146756
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/146756
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit cb28987 with merge base c644f4c ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
Xia-Weiwen
added
topic: not user facing
jgong5
changed the title
|
Hi @jansel Could you please review this PR? Thanks. |
| # define functions to generate example inputs for weight and group size | ||
| # otherwise, autotuner generates example inputs of all zeros for them | ||
| def get_example_weight(x: torch._inductor.ir.IRNode) -> torch.Tensor: | ||
| shape = x.get_size() | ||
| device = x.get_device() | ||
| return torch.randint(0, 255, shape, dtype=torch.uint8, device=device) |
There was a problem hiding this comment.
don't we need to account for strides here? This assumes contiguous.
Can we share some code with the other example input generation?
There was a problem hiding this comment.
Thanks for your comments.
- Now I have added a check: if weight is not contiguous, we won't take GEMM template as a choice for max-autotuning. And I have also added an assertion inside
get_example_weight. - As for the second question, I have searched in the PyTorch source code and didn't find any code that can be shared. This case is special because weight's dtype is uint8.
Thanks.
|
@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 |
Merge failedReason: 1 jobs have failed, first few of them are: trunk / macos-py3-arm64 / build Details for Dev Infra teamRaised by workflow job |
|
@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 |
…th AVX512 (pytorch#146756) **Summary** It's part of the task to enable max-autotune with GEMM template for WoQ INT4 GEMM on CPU. This PR adds GEMM templates for `torch.ops.aten_weight_int4pack_mm_for_cpu`. The micro kernel used for the templates is based on AVX512 and it's a copy of the ATen implementation of `torch.ops.aten_weight_int4pack_mm_for_cpu` with minor changes. Due to better blocking and loop schedule, the GEMM template based implementation outperforms the ATen implementation in all cases we tested. **Test plan** ``` python test/inductor/test_cpu_select_algorithm.py -k test_int4_woq_mm_avx512 ``` Pull Request resolved: pytorch#146756 Approved by: https://github.com/jgong5, https://github.com/leslie-fang-intel, https://github.com/jansel
Stack from ghstack (oldest at bottom):
Summary
It's part of the task to enable max-autotune with GEMM template for WoQ INT4 GEMM on CPU.
This PR adds GEMM templates for
torch.ops.aten_weight_int4pack_mm_for_cpu. The micro kernel used for the templates is based on AVX512 and it's a copy of the ATen implementation oftorch.ops.aten_weight_int4pack_mm_for_cpuwith minor changes.Due to better blocking and loop schedule, the GEMM template based implementation outperforms the ATen implementation in all cases we tested.
Test plan
cc @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10 @voznesenskym @penguinwu @EikanWang @Guobing-Chen @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @amjames @chauhang @aakhundov