Find a common home for decompositions, perhaps outside of the obliquely named _refs directory

[dan-garvey]

🚀 The feature, motivation and pitch

Downstream torch.compile user here.

I recently went through a bit of a frustrating process where I was under the naive and mistaken impression that the torch decompositions entirely lived here: https://github.com/pytorch/pytorch/blob/main/torch/_decomp/decompositions.py

So when hitting an issue with an op, my go to is to check and see if there is a decomposition. (ctrl +f in that file) In my case, I was looking for the decomposition of aten.linspace.Tensor_Tensor, which does not live in that file but instead lives under _refs/init.py. Unfortunately for me, I only discovered this a two working days later after trying to resolve the issue with that op at a different compiler level.

Now I could have used a more holistic search method like: repo:pytorch/pytorch @register_decomposition(aten.linspace

but I do think that the current set of decompositions are at least a little fragmented across the codebase: https://github.com/search?q=repo%3Apytorch%2Fpytorch%20%40register_decomposition(&type=code shows: torch/_inductor/decomposition.py torch/refs/__init_\.py torch/_decomp/decompositions.py torch/jit/_decompositions.py torch/_refs/fft.py torch/_refs/_conversions.py (and several more cases under _refs)

I'm by no means an expert on pytorch repo structure, and the way it is now could totally make sense and I just don't get it, but from my narrow, external perspective, the 6500 line refs/__init_\.py feels disorganized and unintuitive.

I realize I'm not offering much in terms of a solution here, but I expect that experienced contributors who are interested in keeping the codebase as organized as possible will be better equipped than I am to propose solutions. (if it is even a problem)

Alternatives

No response

Additional context

No response

cc @ezyang @msaroufim @bdhirsh @anijain2305 @chauhang

[ezyang]

The historic reason for this organization is torch/_refs was intended to be a pure Python reimplementation of the entirety of the PyTorch frontend (e.g., you could substitute torch with torch._refs) but the camp for that lost out to the camp that just wrote decomps for aten operators because this was easier to do and more expedient.

I do agree it's kind of disorganized but I'm not really sure what a more intuitive organization is.

[masnesral]

Discussed in the triage meeting. We'll take proposals for reorganization, but it's low priority internally.