[inductor][cpp] GEMM template (infra and fp32) #124021

jgong5 · 2024-04-14T14:21:12Z

Stack from ghstack (oldest at bottom):

This PR adds the Cpp template infrastructure and the initial FP32 gemm template. See RFC #125683 for more background info.

Cpp template infrastructure
Similar template abstractions as the CUTLASS template, i.e., CppTemplate, CppTemplateKernel, CppTemplateBuffer. The MicroGemm micro-kernel abstraction that can be used by Cpp GEMM templates.
Initial FP32 gemm template
This involves a GEMM template implementation CppPackedGemmTemplate that supports GEMM with constant weight (B) requiring N to be a multiple of register blocking while allows the static or dynamic sizes for the M (batch dim) of A. The B matrix would be prepacked. This is a typical setting for inference workloads. The template handles the thread decomposition (via thread_blocking) and cache blocking (via cache_blocking). Then it invokes CppMicroGemm which handles register blocking, instruction selection, and other CPU architecture-specific optimizations. A CppMicroGemmFP32Vec micro-kernel implementation is provided for fp32 matmuls implemented with ATen vec abstraction.
Correctness and performance
The changes have been validated with fp32 inference on the three benchmark suites (torchbench, huggingface and timm_models) with both static shape and dynamic shapes. Since it is an initial implementation, we are still working on further performance improves with follow-up PRs including the optimizations in kernels as well as fusions. The perf gains are only observed from a selective number of models compared to the ATen kernels which are implemented with MKL. The perf gains are more obvious with dynamic shapes since MKL only supports packed gemm for static shapes. Below are details.

Static shapes

Benchmark	torchbench	huggingface	timm_models
Multi-threaded (baseline)	1.47x	1.36x	1.91x
Multi-threaded (max-autotune)	1.47x	1.36x	1.92x
Single-threaded (baseline)	1.56x	1.19x	1.51x
Single-threaded (max-autotune)	1.56x	1.19x	1.52x

Key models being sped up:
drq: 1.14x
soft_act: 1.12
cait_m36_384: 1.18x

Dynamic shapes

Benchmark	torchbench	huggingface	timm_models
Multi-threaded (baseline)	1.43x	1.28x	1.85x
Multi-threaded (max-autotune)	1.47x	1.28x	1.85x
Single-threaded (baseline)	1.55x	1.20x	1.51x
Single-threaded (max-autotune)	1.56x	1.19x	1.53x

Key models being sped up:
BERT_pytorch: 1.22x
pyhpc_turbulent: 1.13x
soft_actor_critic: 1.77x
BlenderbotForCausalLM: 1.09x
cait_m36_384: 1.17x

cc @voznesenskym @penguinwu @EikanWang @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler @amjames @desertfire @chauhang

Differential Revision: D57585365

[ghstack-poisoned]

pytorch-bot · 2024-04-14T14:21:15Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/124021

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (4 Unrelated Failures)

As of commit 6ad24d3 with merge base 7a506dd ():

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

inductor / cuda12.1-py3.10-gcc9-sm86 / test (dynamic_inductor_timm, 2, 2, linux.g5.4xlarge.nvidia.gpu) (gh) (similar failure)
sebotnet33ts_256
trunk / macos-13-py3-arm64 / test (default, 1, 3, macos-m1-stable) (gh) (similar failure)
Process completed with exit code 1.

UNSTABLE - The following jobs failed but were likely due to flakiness present on trunk and has been marked as unstable:

linux-binary-manywheel / manywheel-py3_8-cuda12_1-test / test (gh) (#127288)
ImportError: libcudnn.so.9: cannot open shared object file: No such file or directory
linux-binary-manywheel / manywheel-py3_8-cuda12_4-test / test (gh) (#127289)
ImportError: libcudnn.so.9: cannot open shared object file: No such file or directory

This comment was automatically generated by Dr. CI and updates every 15 minutes.

ghstack-source-id: 16b145b Pull Request resolved: #124021

[ghstack-poisoned]

ghstack-source-id: f918cd8 Pull Request resolved: #124021

[ghstack-poisoned]

ghstack-source-id: 4a203e5 Pull Request resolved: #124021

[ghstack-poisoned]

jgong5 · 2024-05-28T06:52:58Z

Hi @jgong5! Sorry, I need to cancel your PRs because they broke a bunch of inductor tests. @jansel could you please help @jgong5 find out the cause of the problem - D57776425, D57776411.

@DanilBaibak @Chillee helped to share the error log to me (https://pastebin.com/LPxPf6nM). It seems the error is not caused by this PR but from the PR on the top of the stack: #126545. Can you confirm if this is the case? If so, I can land the other PRs and focus on this one.
Even though I cannot repro the problem locally, but I tried a fix in #126545 according to the error due to the lack of -mf16c option with clang17 inductor AVX2 build.

This PR adds the Cpp template infrastructure and the initial FP32 gemm template. See RFC pytorch#125683 for more background info. 1. Cpp template infrastructure Similar template abstractions as the CUTLASS template, i.e., `CppTemplate`, `CppTemplateKernel`, `CppTemplateBuffer`. The MicroGemm micro-kernel abstraction that can be used by Cpp GEMM templates. 2. Initial FP32 gemm template This involves a GEMM template implementation `CppPackedGemmTemplate` that supports GEMM with constant weight (`B`) requiring `N` to be a multiple of register blocking while allows the static or dynamic sizes for the `M` (batch dim) of `A`. The `B` matrix would be prepacked. This is a typical setting for inference workloads. The template handles the thread decomposition (via `thread_blocking`) and cache blocking (via `cache_blocking`). Then it invokes `CppMicroGemm` which handles register blocking, instruction selection, and other CPU architecture-specific optimizations. A `CppMicroGemmFP32Vec` micro-kernel implementation is provided for fp32 matmuls implemented with ATen vec abstraction. 3. Correctness and performance The changes have been validated with fp32 inference on the three benchmark suites (torchbench, huggingface and timm_models) with both static shape and dynamic shapes. Since it is an initial implementation, we are still working on further performance improves with follow-up PRs including the optimizations in kernels as well as fusions. The perf gains are only observed from a selective number of models compared to the ATen kernels which are implemented with MKL. The perf gains are more obvious with dynamic shapes since MKL only supports packed gemm for static shapes. Below are details. Static shapes | Benchmark | torchbench | huggingface | timm_models | |------------|-------------|--------------|--------------| | Multi-threaded (baseline) | 1.47x | 1.36x | 1.91x | | Multi-threaded (max-autotune) | 1.47x | 1.36x | 1.92x | | Single-threaded (baseline) | 1.56x | 1.19x | 1.51x | | Single-threaded (max-autotune) | 1.56x | 1.19x | 1.52x | Key models being sped up: drq: 1.14x soft_act: 1.12 cait_m36_384: 1.18x Dynamic shapes | Benchmark | torchbench | huggingface | timm_models | | --- | --- | --- | --- | | Multi-threaded (baseline) | 1.43x | 1.28x | 1.85x | | Multi-threaded (max-autotune) | 1.47x | 1.28x | 1.85x | | Single-threaded (baseline) | 1.55x | 1.20x | 1.51x | | Single-threaded (max-autotune) | 1.56x | 1.19x | 1.53x | Key models being sped up: BERT_pytorch: 1.22x pyhpc_turbulent: 1.13x soft_actor_critic: 1.77x BlenderbotForCausalLM: 1.09x cait_m36_384: 1.17x Differential Revision: [D57585365](https://our.internmc.facebook.com/intern/diff/D57585365) Pull Request resolved: pytorch#124021 Approved by: https://github.com/jansel

As part of pytorch#125683, this PR adds the epilogue support for c++ gemm template by reusing the c++ vector codegen on sub-slices of tensors. This is implemented by retracing the epilogue IR nodes with new ranges and offsets. The new `codegen_loop_bodies` and `codegen_functions` methods are added to c++ vector codegen for this purpose. This is leveraged by the `store_output` method of the template kernel for epilogue codegen and store to the final result. Pull Request resolved: pytorch#126019 Approved by: https://github.com/jansel ghstack dependencies: pytorch#124021

…ue fusion (pytorch#126068) As part of pytorch#125683, this PR adds the initial bf16/fp16 gemm template support with micro-gemm implemented with fused type casting and fp32 computation. It doesn't provide epilogue fusion support yet which will be added in the next PR. Pull Request resolved: pytorch#126068 Approved by: https://github.com/jansel ghstack dependencies: pytorch#124021, pytorch#126019

…rch#126545) As part of pytorch#125683, this PR adds epilogue fusion support for bf16/fp16 gemms. The key changes are as follows: 1. bf16 linear w/ epilogue fusion of some ops was originally supported via ATen oneDNN linear pointwise ops. In order to match the ATen op semantics, in-template epilogue support is added to the cpp gemm template so that we would have: "gemm + in-template epilogues -> template buffer". If the template is chosen for codegen, the in-template epilogues will be concatenated with the out-of-template epilogues that are appended during the scheduling. 2. Support bf16/fp16 legalization for `codegen_loop_bodies` which is used to generate the epilogue loops. 3. We used to leverage the in-place buffer mechanism to handle the in-place buffers in the epilogue codegen, in particular, for the reuses for output buffers of GEMM, template and epilogues. This is not correct since the output buffer is an "output" not an "in-place" buffer of the template kernel itself. Now, we use a dedicated "aliases" dict to manage such buffer reuses and the intermediate aliasing buffers are removed after codegen. 4. Add `localize_buffer` method to `LocalBufferScope` to allow the replacement of a global buffer with a local one in the given inductor IR nodes. This helps the fused loops to work on smaller-sized local buffers for better data locality. Pull Request resolved: pytorch#126545 Approved by: https://github.com/jansel ghstack dependencies: pytorch#124021, pytorch#126019, pytorch#126068

…on (pytorch#126545)" This reverts commit 43baabe. Reverted pytorch#126545 on behalf of https://github.com/DanilBaibak due to Break internal build ([comment](pytorch#124021 (comment)))

…o epilogue fusion (pytorch#126068)" This reverts commit 4aa43d1. Reverted pytorch#126068 on behalf of https://github.com/DanilBaibak due to Break internal build ([comment](pytorch#124021 (comment)))

…26019)" This reverts commit 56c412d. Reverted pytorch#126019 on behalf of https://github.com/DanilBaibak due to Break internal build ([comment](pytorch#124021 (comment)))

This reverts commit 0d1e228. Reverted pytorch#124021 on behalf of https://github.com/DanilBaibak due to Break internal build ([comment](pytorch#124021 (comment)))

[ghstack-poisoned]

jgong5 · 2024-05-29T01:08:13Z

Hi @jgong5! Sorry, I need to cancel your PRs because they broke a bunch of inductor tests. @jansel could you please help @jgong5 find out the cause of the problem - D57776425, D57776411.

@DanilBaibak @Chillee helped to share the error log to me (https://pastebin.com/LPxPf6nM). It seems the error is not caused by this PR but from the PR on the top of the stack: #126545. Can you confirm if this is the case? If so, I can land the other PRs and focus on this one. Even though I cannot repro the problem locally, but I tried a fix in #126545 according to the error due to the lack of -mf16c option with clang17 inductor AVX2 build.

@DanilBaibak So I will land #126068, #126019, #124021 first since I don't see errors from them so far.

[ghstack-poisoned]

jgong5 · 2024-05-29T03:38:58Z

@pytorchbot merge

pytorchmergebot · 2024-05-29T03:41:06Z

Merge started

Your 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

Advanced Debugging

Check the merge workflow status
here

As part of #125683, this PR adds the epilogue support for c++ gemm template by reusing the c++ vector codegen on sub-slices of tensors. This is implemented by retracing the epilogue IR nodes with new ranges and offsets. The new `codegen_loop_bodies` and `codegen_functions` methods are added to c++ vector codegen for this purpose. This is leveraged by the `store_output` method of the template kernel for epilogue codegen and store to the final result. Pull Request resolved: #126019 Approved by: https://github.com/jansel ghstack dependencies: #124021

…ue fusion (#126068) As part of #125683, this PR adds the initial bf16/fp16 gemm template support with micro-gemm implemented with fused type casting and fp32 computation. It doesn't provide epilogue fusion support yet which will be added in the next PR. Pull Request resolved: #126068 Approved by: https://github.com/jansel ghstack dependencies: #124021, #126019

…on (pytorch#126545)" This reverts commit 43baabe. Reverted pytorch#126545 on behalf of https://github.com/DanilBaibak due to Break internal build ([comment](pytorch#124021 (comment)))

…o epilogue fusion (pytorch#126068)" This reverts commit 4aa43d1. Reverted pytorch#126068 on behalf of https://github.com/DanilBaibak due to Break internal build ([comment](pytorch#124021 (comment)))

…26019)" This reverts commit 56c412d. Reverted pytorch#126019 on behalf of https://github.com/DanilBaibak due to Break internal build ([comment](pytorch#124021 (comment)))

This reverts commit 0d1e228. Reverted pytorch#124021 on behalf of https://github.com/DanilBaibak due to Break internal build ([comment](pytorch#124021 (comment)))

This PR adds the Cpp template infrastructure and the initial FP32 gemm template. See RFC pytorch#125683 for more background info. 1. Cpp template infrastructure Similar template abstractions as the CUTLASS template, i.e., `CppTemplate`, `CppTemplateKernel`, `CppTemplateBuffer`. The MicroGemm micro-kernel abstraction that can be used by Cpp GEMM templates. 2. Initial FP32 gemm template This involves a GEMM template implementation `CppPackedGemmTemplate` that supports GEMM with constant weight (`B`) requiring `N` to be a multiple of register blocking while allows the static or dynamic sizes for the `M` (batch dim) of `A`. The `B` matrix would be prepacked. This is a typical setting for inference workloads. The template handles the thread decomposition (via `thread_blocking`) and cache blocking (via `cache_blocking`). Then it invokes `CppMicroGemm` which handles register blocking, instruction selection, and other CPU architecture-specific optimizations. A `CppMicroGemmFP32Vec` micro-kernel implementation is provided for fp32 matmuls implemented with ATen vec abstraction. 3. Correctness and performance The changes have been validated with fp32 inference on the three benchmark suites (torchbench, huggingface and timm_models) with both static shape and dynamic shapes. Since it is an initial implementation, we are still working on further performance improves with follow-up PRs including the optimizations in kernels as well as fusions. The perf gains are only observed from a selective number of models compared to the ATen kernels which are implemented with MKL. The perf gains are more obvious with dynamic shapes since MKL only supports packed gemm for static shapes. Below are details. Static shapes | Benchmark | torchbench | huggingface | timm_models | |------------|-------------|--------------|--------------| | Multi-threaded (baseline) | 1.47x | 1.36x | 1.91x | | Multi-threaded (max-autotune) | 1.47x | 1.36x | 1.92x | | Single-threaded (baseline) | 1.56x | 1.19x | 1.51x | | Single-threaded (max-autotune) | 1.56x | 1.19x | 1.52x | Key models being sped up: drq: 1.14x soft_act: 1.12 cait_m36_384: 1.18x Dynamic shapes | Benchmark | torchbench | huggingface | timm_models | | --- | --- | --- | --- | | Multi-threaded (baseline) | 1.43x | 1.28x | 1.85x | | Multi-threaded (max-autotune) | 1.47x | 1.28x | 1.85x | | Single-threaded (baseline) | 1.55x | 1.20x | 1.51x | | Single-threaded (max-autotune) | 1.56x | 1.19x | 1.53x | Key models being sped up: BERT_pytorch: 1.22x pyhpc_turbulent: 1.13x soft_actor_critic: 1.77x BlenderbotForCausalLM: 1.09x cait_m36_384: 1.17x Differential Revision: [D57585365](https://our.internmc.facebook.com/intern/diff/D57585365) Pull Request resolved: pytorch#124021 Approved by: https://github.com/jansel

As part of pytorch#125683, this PR adds the epilogue support for c++ gemm template by reusing the c++ vector codegen on sub-slices of tensors. This is implemented by retracing the epilogue IR nodes with new ranges and offsets. The new `codegen_loop_bodies` and `codegen_functions` methods are added to c++ vector codegen for this purpose. This is leveraged by the `store_output` method of the template kernel for epilogue codegen and store to the final result. Pull Request resolved: pytorch#126019 Approved by: https://github.com/jansel ghstack dependencies: pytorch#124021

…ue fusion (pytorch#126068) As part of pytorch#125683, this PR adds the initial bf16/fp16 gemm template support with micro-gemm implemented with fused type casting and fp32 computation. It doesn't provide epilogue fusion support yet which will be added in the next PR. Pull Request resolved: pytorch#126068 Approved by: https://github.com/jansel ghstack dependencies: pytorch#124021, pytorch#126019

Update

00eb31a

[ghstack-poisoned]

pytorch-bot bot added ciflow/inductor module: inductor labels Apr 14, 2024

jgong5 pushed a commit that referenced this pull request Apr 14, 2024

[inductor][cpp] GEMM template

9fde2fd

ghstack-source-id: 16b145b Pull Request resolved: #124021

jgong5 marked this pull request as draft April 14, 2024 14:21

pytorchbot added the open source label Apr 14, 2024

Update

b6ff5fe

[ghstack-poisoned]

Update

0355c46

[ghstack-poisoned]

Update

ba94cdf

[ghstack-poisoned]

Update

5ad7899

[ghstack-poisoned]

Update

1c4edcd

[ghstack-poisoned]

Update

a56957d

[ghstack-poisoned]

Update

5bf33c4

[ghstack-poisoned]

Update

f780f9c

[ghstack-poisoned]

Update

0580a46

[ghstack-poisoned]

jgong5 pushed a commit that referenced this pull request Apr 18, 2024

[inductor][cpp] GEMM template

028df20

ghstack-source-id: f918cd8 Pull Request resolved: #124021

Update

d795f31

[ghstack-poisoned]

jgong5 pushed a commit that referenced this pull request Apr 26, 2024

[inductor][cpp] GEMM template

3eb23b1

ghstack-source-id: 4a203e5 Pull Request resolved: #124021

Update

002bedb

[ghstack-poisoned]

Update

2bfc603

[ghstack-poisoned]

Update

c872b7c

[ghstack-poisoned]

jgong5 mentioned this pull request May 28, 2024

[inductor][cpp] BF16 AMX micro-gemm support #127195

Closed

Update

bdb239e

[ghstack-poisoned]

Update

6ad24d3

[ghstack-poisoned]

pytorchmergebot added the merging label May 29, 2024

pytorchmergebot closed this in cef776b May 29, 2024

pytorchmergebot removed the merging label May 29, 2024

github-actions bot deleted the gh/jgong5/40/head branch June 29, 2024 01:54

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[inductor][cpp] GEMM template (infra and fp32) #124021

[inductor][cpp] GEMM template (infra and fp32) #124021

Uh oh!

jgong5 commented Apr 14, 2024 •

edited

Loading

Uh oh!

pytorch-bot bot commented Apr 14, 2024 •

edited

Loading

Uh oh!

jgong5 commented May 28, 2024 •

edited

Loading

Uh oh!

jgong5 commented May 29, 2024

Uh oh!

jgong5 commented May 29, 2024

Uh oh!

pytorchmergebot commented May 29, 2024

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

10 participants

[inductor][cpp] GEMM template (infra and fp32) #124021

[inductor][cpp] GEMM template (infra and fp32) #124021

Uh oh!

Conversation

jgong5 commented Apr 14, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Apr 14, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/124021

✅ You can merge normally! (4 Unrelated Failures)

Uh oh!

jgong5 commented May 28, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jgong5 commented May 29, 2024

Uh oh!

jgong5 commented May 29, 2024

Uh oh!

pytorchmergebot commented May 29, 2024

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

10 participants

jgong5 commented Apr 14, 2024 •

edited

Loading

pytorch-bot bot commented Apr 14, 2024 •

edited

Loading

jgong5 commented May 28, 2024 •

edited

Loading