Repository: pytorch
Issue: [cpp extension] making it possible to parallelize the building process
Labels: module: build,module: cpp-extensions,module: ci,triaged

## 🚀 Feature

Using ninja to build cpp extensions works great for extensions with many objects, but for extensions with very few objects it doesn't parallelize the build.

## Motivation

I have been trying to diagnose why https://github.com/NVIDIA/Megatron-LM takes forever to build its CUDA extensions.

Other than it building each extension sequentially I see that since it uses ninja to build JIT load

https://github.com/NVIDIA/Megatron-LM/blob/b31e1296354e979722627a6c4dedafe19b51fa97/megatron/fused_kernels/__init__.py#L47

ninja appears not to parallelize the build when there are just a few independent objects to build.

I artificially added copies of the object to build and at 6 objects it started to do 2 parallel builds. And I have 12 cores.

I read the ninja documentation and tried to search for answers yet I don't seem to be able to force a simple build like:

```
build scaled_upper_triang_masked_softmax.o: compile /mnt/nvme1/code/huggingface/Megatron-DeepSpeed-master-3/megatron/fused_kernels/scaled_upper_triang_masked_softmax.cpp

build scaled_upper_triang_masked_softmax_cuda.cuda.o: cuda_compile /mnt/nvme1/code/huggingface/Megatron-DeepSpeed-master-3/megatron/fused_kernels/scaled_upper_triang_masked_softmax_cuda.cu

build scaled_upper_triang_masked_softmax_cuda.so: link scaled_upper_triang_masked_softmax.o scaled_upper_triang_masked_softmax_cuda.cuda.o 

default scaled_upper_triang_masked_softmax_cuda.so
```
to not do it sequentially. 

I thought it had to do with having different rules, but it's not the case. It's not about build dependency either.

The first 2 targets are not dependent on each other. The 3rd one is dependent on the previous 2.

----

Additionally as you can see the code performs each extension build in a sequence:
https://github.com/NVIDIA/Megatron-LM/blob/b31e1296354e979722627a6c4dedafe19b51fa97/megatron/fused_kernels/__init__.py#L62-L89
overall this takes ~5minutes to build 3 .so objects which is insanely slow.

Is there a way to tell cpp_extension to build multiple extensions at once? (other than using some sort of pool of processes)

I suspect that if a single ninja.build file were to be written at once for all extensions ninja would parallelize the whole build. But since the build and load of cpp_extension are mushed together into one function, I'm not sure how one could approach that.

## Pitch

So I'm looking for a way to:

1. parallelize the build of a single extension that needs to build more than one but not many .o objects first.
2. parallelize the build of multiple extensions which currently have to be executed in a sequence

Other than wasting developer time, this also impacts the CI where one pays for the run-time and having dozens of idling cores is a huge waste. e.g. our tests take 5 min to build the extensions and the tests take less than that. 

It should be possible to build Megatron-LM extensions 3-5x faster with both parallelizations.

## Other common examples

https://github.com/NVIDIA/apex is another fine example of a project that takes really really long time to build, but I haven't tried to dig into its case. Hint: Megatron-LM requires `apex`.

## Additional info

The author of the ninja system writes in his notes:

> It is very difficult to communicate to users the many possible meanings of "fast". The Ninja manual tries to warn people away from using it on small programs. Literally the second paragraph after the introduction says "If your project is small, Ninja's speed impact is likely unnoticeable" and recommends using a different build system. Unfortunately "fast" sells, and the Ninja list often had users trying to use it for their miniature apps and getting frustrated by its lack of features.
http://neugierig.org/software/blog/2020/05/ninja.html

Thank you!

cc @malfet @seemethere @zou3519 @pytorch/pytorch-dev-infra