Repository: pytorch
Issue: Inductor pattern-matching seems to be sensitive to the presence/absence of clone nodes
Labels: triaged,oncall: pt2,module: dynamic shapes,module: inductor,inductor_pattern_match

### 🐛 Describe the bug

### Summary
 This issue had been discussed in https://github.com/pytorch/pytorch/pull/113004#discussion_r1383666731.
The pattern-matching process seems to be sensitive to the presence or absence of `clone` nodes.
A workaround may be possible. Creating this issue for tracking it, and for documenting this behavior (will add this issue's link in comments in the code when I'd add a new pattern).

### Problem
As in the script below, if I add its corresponding pattern in `torch/_inductor/fx_passes/fuse_attention.py`, it's able to _match itself_, but not the `BERT_pytorch` SDPA pattern. 

```python
def BERT_pytorch_mha_pattern(query, key, value, attn_mask, dropout_p):
    # From https://github.com/codertimo/BERT-pytorch/blob/master/bert_pytorch/model/attention/single.py#L14
    scores = torch.matmul(query, key.transpose(-2, -1)).div(math.sqrt(query.size(-1)))
    fill_value = torch.full((), -float(1e9), dtype=query.dtype, device=query.device)
    scores = scores.masked_fill(attn_mask == 0, fill_value)   
    return torch.nn.functional.dropout(
        scores.softmax(dim=-1), dropout_p
    ).matmul(value)

with torch.no_grad():
    q = torch.randn(1, 12, 256, 64)
    k = torch.randn(1, 12, 256, 64)
    v = torch.randn(1, 12, 256, 64)
    attn_mask = torch.rand(1, 12, 256, 256).to(torch.bool)
    BERT_pytorch_mha_model = torch.compile(BERT_pytorch_mha_pattern, fullgraph=True, options={"freezing": True})
    BERT_pytorch_mha_model(q, k, v, attn_mask, 0.0)
```
### Workaround

However, if I `transpose`  `q`, `k` & `v` in a [new SDPA pattern corresponding to the script below](https://github.com/pytorch/pytorch/compare/main...sanchitintel:pytorch:sanchitj/sdpa_pattern_example?expand=1), then Inductor is unable to match it _with itself_ because the serialized pattern (generated by adding a pattern in `torch/_inductor/fx_passes/fuse_attention.py` corresponding to the script below has `clone` nodes after `permute` and `expand` of `query`, `key`, and `value`, but the `torch.fx.GraphModule` graph does not when `joint_graph_passes` are run.
But I discovered that this new pattern is actually able to match the SDPA pattern in `BERT_pytorch` because the generated graph for the model also has the aforementioned `clone` nodes.

```python
import torch

def BERT_pytorch_mha_pattern(q, k, v, attn_mask, dropout_p):
    # attn_mask is already of bool dtype in this model 
    q = q.permute(0, 2, 1, 3)
    k = k.permute(0, 2, 1, 3)
    v = v.permute(0, 2, 1, 3)
    # From https://github.com/codertimo/BERT-pytorch/blob/master/bert_pytorch/model/attention/single.py#L14
    scores = torch.matmul(q, k.transpose(-2, -1)).div(math.sqrt(q.size(-1)))
    fill_value = torch.full((), -float(1e9), dtype=q.dtype, device=q.device)
    scores = scores.masked_fill(attn_mask == 0, fill_value)   
    return torch.nn.functional.dropout(
        scores.softmax(dim=-1), dropout_p
    ).matmul(v)
    
with torch.no_grad():
    q = torch.randn(1, 256, 12, 64)
    k = torch.randn(1, 256, 12, 64)
    v = torch.randn(1, 256, 12, 64)
    attn_mask = torch.rand(1, 12, 256, 256).to(torch.bool)
    BERT_pytorch_mha_model = torch.compile(BERT_pytorch_mha_pattern, fullgraph=True, options={"freezing": True})
    BERT_pytorch_mha_model(q, k, v, attn_mask, 0.0)
```

So, the problem I had been facing has been resolved with this workaround, but I thought we should clearly document this behavior. Thanks!


 


### Versions

Current main branch

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