[RFC] Adding CPU Intrinsics into codegen
- Dominant language
- Python
- Stars
- 103k
- Forks
- 29.5k
- PR merge metrics
- PR metrics pending
Description
### 🚀 The feature, motivation and pitch
With the rise of agentic AI and edge computing, CPU inference performance has become increasingly critical. To unlock the full hardware potential of modern CPUs, it is often necessary to utilize low-level intrinsics. One example is [this](https://github.com/cfgfung/pytorch/blob/raymond/overlap_amx_avx_only/torch/_inductor/codegen/cpu_intrinsics/cpp_flex_attention_amx.py), which leverages the CPU's out-of-order execution engine to overlap GEMM and Softmax operations within Flex Attention 3.
### Alternatives
We are considering two approaches for placing this implementation:
**Approach 1: Place it directly under torch/_inductor/codegen**
Pros: Easy to track and review.
Cons: Tightly couples the implementation to a specific function (flex_attention).
**Approach 2: Create a new torch/_inductor/codegen/cpu_intrinsics/ directory**
Pros: Extensible for future intrinsic implementations.
Cons: Potentially over-engineered, given that flex_attention is currently the sole use case.
Recommendation:
I favor Approach 2. With the rise of agentic AI workloads, we can anticipate a growing need for CPU-specific optimizations.
### Additional context
_No response_
cc @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10 @aditew01 @chauhang @penguinwu @voznesenskym @EikanWang @Guobing-Chen @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @aakhundov @coconutruben @jataylo @Chillee @drisspg @yanboliang @BoyuanFeng @liangel-02 @howardzhang-cv @jas @jansel
Contributor guide
Research direction
Start by reading torch/_inductor/codegen and the linked example at torch/_inductor/codegen/cpu_intrinsics/cpp_flex_attention_amx.py. Compare the two proposed locations for CPU intrinsic implementations and review the discussion to determine the agreed design. Done means the placement and scope are settled for future CPU-specific optimizations.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Needs clarification
- Newbie friendliness
- 30/100