[inductor] Lower aten.cumsum

[desertfire]

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305 @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @ngimel @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @soumith

[jansel]

I don't think it is currently possible to write this one in Triton, so this is blocked on Triton language features.

[lezcano]

What about doing a expand + tril + sum? It would do some extra computation but it unblock this + it would be parallelizable. The only issue is that we would need play around a bit with the dimensions as tril just works on matrices, and cumsum takes an optional dim.

[jansel]

That would be correct, but I would expect it to be slow.

[lezcano]

And you would be right. I will implement the multi-kernel approach from JAX, see if it gets a more reasonable speed.

[lezcano]

I'm leaving this open, as the current decomposition is very slow. See #86229 (comment)
See #86372 for a different approach (that currently breaks TorchInductor).