Repository: pytorch
Issue: Better fusion on triton templates
Labels: triaged,oncall: pt2,module: inductor

### 🚀 The feature, motivation and pitch

Hi PyTorch Team,

I'm optimizing a gather-compute-scatter pattern and would like to fuse gather and scatter operations into `matmul` Triton templates. My current implementation performs two gathers (for `input` and `weight`) and a scatter-add operation into `output`, with some intermediate channel computation.

In an einsum-like notation, this operation is expressed as:

```
out[ocrd[po, pi], m] += input[icrd[po, pi], c] * weight[wcrd[po], c, m]
```

Here’s the sample code:

```python
import torch

N = 87680
M = 27
Po = 4529
Pi = 128
out_channel = 64
in_channel = 64

wcrd = torch.randint(M, (Po,), device="cuda")
ocrd = torch.randint(N, (Po,Pi), device="cuda")
icrd = torch.randint(N, (Po,Pi), device="cuda")

input = torch.rand(*(N, in_channel), dtype=torch.float32, device="cuda")
weight = torch.rand(*(M, in_channel, out_channel), dtype=torch.float32, device="cuda")
output = torch.zeros((N, out_channel), dtype=torch.float32, device="cuda")

def test(wcrd, ocrd, icrd, input, weight, output):

  icrd_flat = icrd.view(-1)                                  # (Po*Pi,)
  input_selected = input.index_select(0, icrd_flat)          # (Po*Pi, in_channel)
  input_selected = input_selected.view(Po, Pi, in_channel)   # (Po, Pi, in_channel)
  weight_selected = weight.index_select(0, wcrd)             # (Po, in_channel, out_channel)
  output_values = torch.bmm(input_selected, weight_selected) # (Po, Pi, out_channel)
  ocrd_flat = ocrd.view(-1)                                  # (Po*Pi, )
  output_values_flat = output_values.view(-1, out_channel)   # (Po*Pi, out_channel)
  output.index_add_(dim=0, index=ocrd_flat,
      source=output_values_flat)                             # (N, out_channel)
  return output
compiled = torch.compile(test, mode="max-autotune-no-cudagraphs")
output = compiled(wcrd, ocrd, icrd, input, weight, output)
```

### Observations

When compiled with `torch.compile`, the generated `output_code.py` shows four separate Triton kernels: (1,2) two gathers, (3) bmm template, and (4) one scatter.

If I manually fuse these into a single `bmm` template that includes the gathers in the prologue and scatter-add in the epilogue, the performance improves significantly:

- **Without Fusion (current `torch.compile`):** 1.41 ms
- **With Manual Fusion:** 0.62 ms

### Questions

1. **Fusion of Pointwise Operations:**  
   Are `index_select` and `index_add` treated as pointwise operations in PyTorch's `torch/_inductor/kernel`? Can these operations be fused with `matmul`, `batchmatmul`, or `flexattention` in Triton templates?

2. **Prologue and Epilogue Fusion:**  
   If full fusion of multiple pointwise operations is challenging, can gather and scatter at least be fused in the prologue and epilogue of the `bmm` template?

3. **Prologue Fusion (Related to [#134532](https://github.com/pytorch/pytorch/issues/134532)):**  
   I came across [#134532](https://github.com/pytorch/pytorch/issues/134532) for prologue fusion. Will this resolve the issue in the example case above?

4. **Epilogue Fusion Support:**  
   Since `torch.compile` already supports epilogue fusion, why is `index_add_` not being fused into the epilogue of the `bmm` template? Is this a potential bug?


@jansel 

### Alternatives

_No response_

### Additional context

_No response_

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