Repository: pytorch
Issue: Potential 7.5x compilation time improvement for foreach kernel
Labels: triaged,oncall: pt2,module: inductor,module: pt2 optimizer

### 🐛 Describe the bug

Foreach kernel (especially those generated for the optimizer) is very large and takes quite some time to compile. The reason is we duplicate the main block of code to process each tensor's gradient update. Those code need to be repeatedly lowered and transformed in triton which takes time.

Here is one initial idea which does not ultimately work but can be a good step stone to think about. Instead of codegen duplicate code for each tensor, we only codegen the code that's specific to each tensor
- xpid_offset
- xnumel
- in-ptrs
- out-ptrs

And then we move the common code to a function and call this function for each tensor.

Unfortunately this does not work, since triton by default will inline every function call.

But based on this idea, we can do something like this:

```
    if xpid >= 0 and xpid < 1:
        # setup code for this group
        # e.g. setup pid range and in/out ptrs and assign them to placeholder variables.
        # Here shows an example
        xpid_offset = xpid - 0
        xnumel = 23440896
        i0 = in_ptr0
        i1 = in_ptr1
        i2 = in_ptr2
        i3 = in_ptr3
        i4 = in_ptr4
        o0 = out_ptr0
        o1 = out_ptr2
        o2 = out_ptr3
    elif xpid >= 1 and xpid < 577:
        # setup code for this group
     ...
    else:
         tl.device_assert(False, "can not reach here")

    # Common body that uses the placeholder variables.
    # here is some simplified example. Full example will be provided below
    xoffset = xpid_offset * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(i0 + (x0), xmask)
    tmp1 = tl.load(i1 + (x0), xmask)
    tmp6 = tl.load(i2 + (x0), xmask)
    tmp13 = tl.load(i3 + (x0), xmask)
    tmp15 = tl.load(i4 + (0))
    ....
    tl.store(o0 + (x0), tmp5, xmask)
    tl.store(o1 + (x0), tmp34, xmask)
    tl.store(o2 + (x0), tmp12, xmask)
```

The trick here is to codegen the 'setup code' for each branch in a way that 'parameters' are assigned to placeholder variables. Those placeholder variables work like parameter and will be used in the 'common body'.


Here is an example optimizer wrapper generated for BertForMaskedLM: 
https://gist.github.com/shunting314/50343179d4cb9889526a34a2a00eb685#file-orig_wrapper-py-L1414

I manually apply the idea to triton_for_fused_2 kernel, and get https://gist.github.com/shunting314/e3396bd9753045313f90d829b42bc906#file-wrapper-py-L1414

The compilation time for triton_for_fused_2 improve from 4.574s to 0.613s (7.5x improvement).

cc @ezyang @msaroufim @bdhirsh @anijain2305 @chauhang @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @mlazos @janeyx99 @jansel @Chillee @eellison 

### Error logs

_No response_

### Minified repro

_No response_

### Versions

.

cc @ezyang @msaroufim @bdhirsh @anijain2305 @chauhang