pytorch / pytorch/pytorch

[RFC] Adding CPU Intrinsics into codegen

Open
#191,257 11 comments 1 reaction 0 assignees View on GitHub
bot-triaged feature module: cpu module: flex attention module: inductor needs design oncall: pt2 triaged
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.