Feat: WIP - Enable ND-Distributed arrays on operators

[mrava87]

Just toying around with the idea of allowing ND-Distributed arrays to be consumed by operators (like in PyLops)

A few thoughts from this initial experiment:

• self.dims/self.dimsd must be added to MPILinearOperator like we do for PyLops LinearOperator and we must ensure that they are defined for every operator (without them it is impossible to know the kind of ND input the operator would accept) - Added in Add dims and dimsd to MPI Operators #199.
• in pylops we follow this convention:
  • matvec/rmatvec always take 1d arrays
  • dot can take ndarray (ravels them, runs matvec, reshapes them back)
  • in each operator the _matvec/_rmatvec implementations assume 1d inputs but with @reshape can reshape it for processing and then flatten it back
  • solvers can taken Nd arrays but internally the flatten, run and reshape
  • everything is a view so back and forth between 1d and nd is rather cheap
• in pylops-mpi, any ravel/reshape does inherently create a new Distributed array and does copy local arrays so it is kind of expensive. We may need to consider to change a bit the places where ravel/reshape is done, perhaps:
  • in matvec/rmatvec of MPILinearOperator instead of in dot, so that also in solvers we can just call those methods on NDarrays
  • avoid ravel before calling _matvec/_rmatvec and instead change @reshape to just pass blank if the input is already ND (maybe with a check that its global_shape is aligned with dims/dimsd... and similar avoid flattening it back

@rohanbabbar04 FYI, this is as far as I got, some tentative code changes, an example, and some thoughts that came to my mind whilst toying around 😄

[rohanbabbar04]

@mrava87, thanks this makes sense, looks like a good start.

I will add some sub-task pointers in a couple of days here itself, which we would need to cover 🙂 .

[mrava87]

@rohanbabbar04 sounds good, feel free to take it from here 😉

[rohanbabbar04]

@mrava87, I made some changes to make our dot work with ND arrays matching self.dims. Have added some code for 4 operators until now - MPIFirstDerivative, MPISecondDerivative, MPIBlockDiag, MPIVStack.

• I agree that numpy.reshape and numpy.ravel return views, which is why they can be safely used in PyLops. However, in pylops_mpi they usually create a copy. To work around this, I added a local_array parameter to the DistributedArray class. If provided, it is directly assigned to local_array. I don't expect users to use this parameter directly in most cases, but we can leverage it internally in our reshape and ravel methods, and potentially in other places in the future.

• I actually think that having _matvec and _rmatvec return ND DistributedArrays could introduce some issues. For example:

# Running with 2 ranks
second_deriv = MPISecondDerivative(dims=200)
b_diag = MPIBlockDiag(ops=[MatrixMult(A=np.ones((100, 100)))])
summa = b_diag + second_deriv

For summa to compute matvec correctly, it expects the local_shape to be the same on every rank. This may no longer hold if _matvec returns ND arrays, since they would be reshaped according to dimsd.

• I introduced _local_dims and _local_dimsd as SimpleNamespaces containing two fields: dim and axis.(We can have a tuple as well). These specify the shape of the final local array and the axis along which it should be reshaped. I think storing the axis is important for other operators as well, such as Halo and MatrixMult.

[mrava87]

@mrava87, I made some changes to make our dot work with ND arrays matching self.dims. Have added some code for 4 operators until now - MPIFirstDerivative, MPISecondDerivative, MPIBlockDiag, MPIVStack.

• I agree that numpy.reshape and numpy.ravel return views, which is why they can be safely used in PyLops. However, in pylops_mpi they usually create a copy. To work around this, I added a local_array parameter to the DistributedArray class. If provided, it is directly assigned to local_array. I don't expect users to use this parameter directly in most cases, but we can leverage it internally in our reshape and ravel methods, and potentially in other places in the future.
• I actually think that having _matvec and _rmatvec return ND DistributedArrays could introduce some issues. For example:

# Running with 2 ranks
second_deriv = MPISecondDerivative(dims=200)
b_diag = MPIBlockDiag(ops=[MatrixMult(A=np.ones((100, 100)))])
summa = b_diag + second_deriv

For summa to compute matvec correctly, it expects the local_shape to be the same on every rank. This may no longer hold if _matvec returns ND arrays, since they would be reshaped according to dimsd.

• I introduced _local_dims and _local_dimsd as SimpleNamespaces containing two fields: dim and axis. These specify the shape of the final local array and the axis along which it should be reshaped. I think storing the axis is important for other operators as well, such as Halo and MatrixMult.

@rohanbabbar04 thanks for continuing on this 😄

Let me have a look in the next couple of days so I can provide a more informed feedback

[mrava87]

I follow this but I still have a question 😉

Before we always assumed that the input of a _matvec/_rmatvec was a 1D Distributed Array so when we decorated them with @Reshaped, inside here we will always have to create a new ND Distributed array and fill it (for a moment I am ignoring the ghost cell path).

Now, we want to enable the scenario when a ND Distributed array is passed directly to the _matvec/_rmatvec of an operator, so when we get into this the array may already by ND. At that stage if cells_front=cells_back=0, would you not say that this method could just pass x to f?

Also, again along similar lines

if len(y.global_shape) > 1:
   # Make sure y is distributed along axis=0 before applying ravel
    y = y.redistribute(axis=0).ravel()

does this not mean that the output is still again always flattened, so we can't yet have an operator that takes a ND Distributed arrays and returns a ND Distributed array? Maybe somehow related to you second point ('I actually...) But then, is this not defeating a bit the real purpose of this PR?

PS: whether we can or not enable ND Arrays flowing through ops and solvers, one by-product of this PR, which we should definitely merge is the improvement in efficiency when it comes to the setting of local_array in various places 😄

[mrava87]

@rohanbabbar04 I left some comments (mainly questions and thoughts).

In general:

• I do agree with your new local_array approach, I think it already saves quite a few numpy array copies we used to make 😄
• _local_dims and _local_dimsd, I do see their value (and fine with SimpleNamespaces approach)... I was just not sure if they are always added or just in some cases (ie for some operators); mostly unsure if this will lead to scenarios where an operator expects them but they may not be there?
• Not sure about your example, but just to be clear, we do not want/can to allow every possible combination of operators to lead to a ndarray compatible case.. I tried to mimic something like this with PyLops

`second_deriv = SecondDerivative(dims=200)
b_diag = BlockDiag(ops=[MatrixMult(A=np.ones((100, 100))), MatrixMult(A=np.ones((100, 100)))])
summa = b_diag + second_deriv

print(b_diag.dims, second_deriv.dims)

x = np.ones(200)
y = summa @ x
print(y.shape)

x = np.ones((2, 100))
y = summa @ x
print(y.shape)

and get (2, 100) (np.int64(200),) (200,) (2, 100), so basically the dims of the 2 ops differ but I think because we flatten in dot -> reshape in each op -> reflatten in each op -> reshape in dot, both case work... so I guess the same applies to the pylops-mpi, hence you sticking to this flatten in dot -> reshape in each op -> reflatten in each op -> reshape in dot pattern?

One last suggestion: try to have a go with one of the operators you have now enabled ND arrays to go end-to-end with the solution of an inverse problem. Can we go all the way ND? One of the initial use case that led me to think about this was your FFT, which may benefit from avoid to ever go back to 1d with redistribution but sticking all the time to ND (other than any code we have fro example in tutorials where it would be nicer for a user not to have to think in flattened local arrays of a global nd array but just work with nd local arrays

[mrava87]

Why this change?

[mrava87]

Nice idea about the local_array

[mrava87]

Does this assume that if the array is scattered all dims but the one of axis are the same? Do you think its worth a check+raise or we shall just assume the user does it correctly (at least maybe mention this in the docstring?)

[mrava87]

Since most of the inputs are positional, I would move this a bit earlier maybe close to local_shapes... and a few things to consider when one passes local_array:

• dtype can be conflicting, so maybe we should say if local_array is passed dtype is ignored and local_array.dtype used instead?
• same for engine?
• similarly local_shapes may be contrasting, maybe here I would raise an error if both are passed to avoid confusion?

I understand users may not find this useful (aka it is more for developers) but since it surface as an input of the class constructor, I think we should treat it as is and make sure it does not lead to confusing scenarios...

[mrava87]

So this _local_dims is not set in MPILinearOperator constructor, so if one does MPILinearOperator(Op) with Op being a PyLops operator, this will not be present... thinking if this could lead to any scenario where it is always expected but it may not always be there?

[mrava87]

For example here you don't set _local_dimsd ?

[mrava87]

Mmh back to one of my original comments: why do we need ravel here, and is with the new local_array now a much cheaper (view-only) operation?

[mrava87]

But since redistributed is for sure not cheap, why you see this as needed here? And why ravel, other than maybe to be as close as possible to what we do in pylops?

[rohanbabbar04]

@mrava87 My main thought was to try and mimic what we do in pylops, which was flattening and then apply the _matvec and then reshape back to dimsd like we do in pylops. With the introduction of local_array, copy would be cheap, but yeah with redistribution it will be expensive. So would you suggest we keep NDArray only for specific ops like FFTs(and maybe other....), since we used to redistribute(expensive operation) it to axis 0 to apply the ravel at the end.

[mrava87]

@rohanbabbar04 I am not sure, my initial thought was 'let's follow the PyLops steps' but then already when I started toying around I saw the issue with us recreating and reassigning all of these numpy arrays (as our local arrays) so I had the idea that maybe for the purpose of pylops-mpi (solve one very large problem in a distributed manner*) we could just avoid doing any ravel-(reshape-apply-flattened)-reshape with (reshape-apply-flattened) done in the reshape decorator if an input was not a 1d array.

Now if this process was cheap as it seems now (up to redistribute) this would be purely stylistical but given the fact we do that, maybe we need to rethink the entire process instead of fitting the PyLops pattern in it?

• for pylops the idea of solving one or more flattened problems (with multiple RHS) makes somehow sense as in some applications one has a small problem, but many of them... this kind of goes away in pylops-mpi as if you reach for it the likely cause is your problem is so big it cant even fit in one single machine VRAM (or is very slow when solved in single-cpu mode), so my thought that maybe we can loosely allow a solver to operate with either flat or unflattened arrays just considering them all the time as if they were flattened (where basically only for dot and norm I see we needed some modifications, all +/-/* are already transparent to that)

. So would you suggest we keep NDArray only for specific ops like FFTs(and maybe other....), since we used to redistribute(expensive operation) it to axis 0 to apply the ravel at the end.

but what I pointed here is applied to all operators where one passes a ND-array and self.dims is not 1d, so I am not sure what you say would work? Perhaps I would suggest we take a tutorial, say poststack and see if what we have now (plus likely some small changes to the solvers) would allow going from doing

# Create flattened model data
m3d_dist = pylops_mpi.DistributedArray(global_shape=ny * nx * nz)
m3d_dist[:] = m3d_i.flatten()

to

# Create distributed model data
m3d_dist = pylops_mpi.DistributedArray(global_shape=(ny, nx, nz), axis=0)
m3d_dist[:] = m3d_i

where note that here this is kind of just stylistic and for ease of use (I always thought it would be nice for a user not to think in 1d distributed flattened arrays, but for other scenarios where an operator like FFT is part of the final Op it would become more important.

Right now I have the feeling that even doing d_dist = BDiag @ m3d_dist to the unflatted m3d_dist would hit the heavy redistribute?

[rohanbabbar04]

Hmm... Yeah, I was also thinking about the best way to tackle ND arrays, especially for operators like FFTs and derivatives where users would naturally want to preserve the multidimensional shape.

I was thinking of changing matvec slightly:

def matvec(self, x: DistributedArray) -> DistributedArray:
    M, N = self.shape
    if x.ndim == 1:
        if x.global_shape != (N,):
            raise ValueError("dimension mismatch")
        return self._matvec(x)

    if np.prod(x.global_shape) != N:
        raise ValueError("dimension mismatch")

    return self._matvec_nd(x)

The idea would be that _matvec_nd is optional. Operators that make more sense/reduce communication work with ND arrays (such as FFTs or derivatives) can implement it and preserve the input shape. Operators that don't have a meaningful ND implementation can simply leave _matvec_nd unimplemented (e.g., raising NotImplementedError) and continue to support only 1D distributed arrays through _matvec.

For FFTs, for example, we could move the implementation entirely into _matvec_nd, since the FFT is expensive when we are creating a 1-d array and the users would typically want to work with the ND dimensions rather than flattening the array.

Again, this is just an idea for now, and we can think it through some more. 🙂

[mrava87]

Right, the idea seems appealing but l would need to know a bit more to see if this makes sense 😉

First question: as we want to mantain backward compatibility, I guess for operators where it makes sense to implement _matvec_nd their _matvec would still leverage the @reshape decorator and then just call _matvec_nd?

Second question: this is something I am always very careful we honor, as I think it's one of the great features of pylops. Users should always be able to mix any combination of operators (provided shapes make sense) in the most transparent possible way (i.e., not having to know all the internals of the actual implementations of the operators). Now let' say we go for this _matvec_nd implementation, do you think we can end up in a situation where:
- a user can always create a combined operator and pass a 1D distributed array, even if somehow performance-wise suboptimal, like when you have an FFT operatr
- a user can always create a combined operator and pass a ND distributed array provided all of the operators in the combined operator support this (implement _matvec_nd)?

Again, I would maybe start by taking a tutorial like I suggested above and see if any of what you are thinking would enable/not-enable solving that problem with ND distributed arrays, and then maybe work backwards adding support to more operators and checking against other tutorials/examples

[rohanbabbar04]

Yup, sounds good, we can start with a tutorial and see the best way possible 🙂.

[mrava87]

Mmh I have not tried if it is possible but say we define the input DistributedArray with custom local_shapes (x) and then Op=MPIFirstDerivative and then do Op @ x... this is legit right, but now the local_split dims would not match the dims of the input array?

[mrava87]

What is this testing (the name is cryptic)?

You do global_shape=BDiag_MPI.dims but Op as is defined is supposed (even in numpy) to take a 1d array (you would need to specify otherdims for this operator to allow one to pass a Numpy ndarray... or is this having global_dims[0] equal to the number of ranks so it is really a (1,par['nx']) local shape?