Device management and pmap

[jessegrabowski]

With the implementation of Blockwise and pt.vectorize, one of the remaining major features pytensor doesn't have vis-a-vis the big boy tensor libraries is device management and pmap.

I'm not suggesting we commit to developing and maintaining our own cuda ops (though it would be nice in an ideal world). But it would be nice to at least be able to support CPU parallelism out of the box. How far away are we from something like that? I don't even have a sense of how to think about it.

Even if we couldn't do it ourselves at first, it would still be nice be able to send information about which Ops should go to which devices when compiling to JAX or torch.

Tagging @Ch0ronomato because he got me thinking about this.

[ricardoV94]

To get some perspective, I suggest looking at what we already have.

Elemwise uses/can use OpenMP in the C backend.

Fusion rewrite is important here. We don't currently fuse Blockwise.

Scan could be implemented with prange in Numba. Not sure anything needs to be done with JAX.

We used to have Ops to transfer data to and from GPU (check Theano).

If thinking about users like PyMC, note that we work with multiple graphs and automatically translate from forward to logp graphs, and Ops like send to GPU don't necessarily make sense in the same places.

This raises the question how much stuff can we do automatically so it doesn't require user input.

Also, how do we avoid clobbering graph analysis with low level device management operations?

Just some things to muse on

[Ch0ronomato]

I don't know if pytensor is or is not intended to be able to do that type of lowering ala cuda ops. Torch and Jax are all pretty intensely maintained so they can manage to keep the whole device management to be done well ™️ .

Does it make sense to think about the different kinds of possible parallelism we would want to accomplish? For example, I could see a feature where we take the fgraph, and break it into pieces to better accomplish a form a pipeline parallelism / force gpu sync points.