Parallel execution of multiple unrelated statements written sequentially

[Susmit-A]

🚀 The feature, motivation and pitch

Often, I write statements in a sequential manner that do not relate to each other - the order of their execution does not matter. Such statements can be parallelized to achieve better performance. One example: resize three different images using nearest, bilinear, and bicubic:

Old method:

img1 = F.interpolate(img1, mode="nearest", ...)
img2 = F.interpolate(img2, mode="bilinear", ...)
img3 = F.interpolate(img3, mode="bicubic", ...)

New method:

img1, img2, img3 = parallel_api_fn([
    {
        "function": F.interpolate, 
        "args": {"inputs": img1, "mode": "nearest"}
    },
    {
        "function": F.interpolate, 
        "args": {"inputs": img2, "mode": "bilinear"}
    },
    {
        "function": F.interpolate, 
        "args": {"inputs": img3, "mode": "bicubic"}
    }
])

Of course, the implementation mentioned above is just an example. The core idea is to take a list of functions (same or different), and their respective arguments. These functions are then executed in parallel using the new API function.

Alternatives

Currently, the functionality can be implemented using vmap or python multiprocessing. However, each of these requires a lot of additional code and has a lot of added complexity. If we assume that every line to be executed in parallel is independent and all arguments are immutable, I believe the lines can easily be run in parallel on GPU.

Additional context

No response

cc @ngimel @bdhirsh

[ngimel]

You can use cuda streams for that.