search

facebookresearch / fmmax

Fourier modal method with Jax
MIT License
95 stars 9 forks source link

Enable vectorized eigensolve #86

Closed mfschubert closed 7 months ago

smartalecH commented 7 months ago

Do you know in what sense this vectorizes things?

Is it vectorizing the underlying eig decomposition? Or is allowing one to vectorize the operation of many eig decompositions? If so, does this impact the performance of the underlying Lapack routine, which expects it to already have access to all the threads?

mfschubert commented 7 months ago

My understanding is that this simply determines whether a batch of matrices is passed to the function, or if the function is called once for each matrix in the batch. If the single-matrix code is poorly parallelized, then this may make a difference. If it is already well-parallelized, the difference should be minor.

I am not observing a major difference, but since the eig is in fact vectorized, this seems like the right thing to do. That way we are not inserting any artificial serialization bottlenecks.