Design of add-ons for mdarray

[grothesque]

In our research group, we have students who are implementing algorithms using mdarray and using a naive matrix multiplication routine for now. We will (soon) need to use efficient matrix multiplication and also other subroutines à la LAPACK.

I have to work at least on a stopgap solution, but would prefer to eventually come up with something proper. Here are some thoughts on the topic. I’d greatly appreciate comments and suggestions.

• In addition to classic BLAS and LAPACK, there are at least faer, and matrixmultiply. Faer contains both a high-level and a low-level interface.

• One way to supply dot-product and matrix multiplication is through a single crate where the backend(s) are selected through features. Here one problem that I see is that crate features must be additive. This seems to imply that there would have ot be some form of precedence for the case where there are more than two backends.

• Another way to supply dot-product and matrix multiplication is through multiple crates. One for LAPACK, say, one for faer, one for matrixmultiply. This is more flexible, and it also avoids complicated features (because all of these backends have their own features which are relevant as well). Since traits can be implemented only once for each trait-type pair in Rust, I presume that they way to go would be to have a separate API for each crates. (One could strive to have a common form for these APIs and ideally to keep them interchangeable.) An advantage of this approach would be that it keeps the individual crates less complex.

• Alternatively, one could try to introduce common traits, and require newtypes for the mdarray types. Or perhaps use some exotic approach like generic traits parametrized by marker types? But I fear complexity and boilerplate here.

• Here is an idea that I just had. Traits could be implemented not for arrays, but for particular libraries. For example there could be a MatMul trait that describes the multiplication of two mdarray slices. Then a crate might implement this trait and the usage would look something like

  use mdarray_blas::Blas;  // Implements the mdarray_traits::MatMul trait
  
  fn main() {
      let a: &mdarray::Slice<...> = ...;
      let b: &mdarray::Slice<...> = ...;
  
      let c = Blas.matmul(a, b).into_new();  // into_new is a method of mdarray_traits::MatMulResult
  }

Implementing traits not on the matrices themselves but rather on types that represent the library (or library contexts - they could also take parameters) seems to have several advantages:

• The use of the builder pattern for operations allows to have optional arguments and different kinds of dealing with the result (reusing allocations or creating new ones, say)

• There is a way to implement a common standard API (as specified by the trait), but each library can have own extensions, parameters, etc.

[fre-hu]

Great that mdarray is being used! It would be really good if you can work on BLAS and LACK support.

The idea about using a trait for the library makes sense, and gives a way to both change backend and as a state where you can configure for example parallelization. You could also have less coupling to the array implementation if you have an array trait that models the behaviour you need instead of direcly using the array types.

On the other hand it is also good to start simple, and try the existing blas and lapack bindings and look into how they are used for example in ndarray. They already give a way to change between different C backends.

[grothesque]

This is what I have so far: https://git.grothesque.org/mdarray-linalg-prototype/ (there’s also a GitHub mirror).

So far there’s only a simple gemm generic function, that however should cover all mappings that are supported by BLAS. There’s no operation-builder interface yet.

I’m grateful for any comments. In particular, I am considering the following questions:

• It seems to me that the gemm function in its current form is as transparent/efficient as possible, up to being zero-overhead for const-shaped arrays. (This is not a precise statement, since I do not completely grasp mdarray’s layout/mapping/shape system, but I hope that it’s clear what I mean.) In other words I hope that the compiler should be able to figure out statically that some (or even all) of the checks are no-ops, depending on the types of the arrays.

• I introduced the trait BlasScalar to abstract over the different BLAS scalar types. Perhaps one could try to find some way to directly use the num_complex::ComplexFloat trait for this with some trick, but this is not straightforward, and I’m not even sure whether it’s a good idea, since in theory ComplexFloat could be implemented for scalar types that are not supported by BLAS. If we continue with this approach, do we want to have one huge trait that contains all of BLAS and LAPACK? We can split it up, but then writing generic code that uses these functions will require complicated bounds (more complicated than T: ComplexFloat + BlasScalar).

• The current test function can be modified into eight different variants. It would be nice to always test all variants (there could be even more), but I don’t know how to do this without a lot of code duplication. Declarative macros don’t seem to work for this, because of problems with types and life times. At least I did not manage to come up with a nice solution.

• Currently, the gemm function will panic not only when the matrix shapes are not compatible, but also when the strides are too general for BLAS. I guess that this could and should be relaxed to accept any strides, using temporary arrays when necessary. I believe that conditional code for such cases will not even be emitted in common cases (for example when all three arrays are Dense it is clear at compile time that no temporary arrays will be necessary).

[fre-hu]

It looks fine as a start.

On the second question, maybe have separate Blas and Lapack traits and a trait Linalg: Blas + Lapack? It looks similar to the Lapack trait in ndarray-linalg, I guess you have seen it: https://docs.rs/ndarray-linalg/latest/ndarray_linalg/types/trait.Lapack.html.

[grothesque]

Indeed it looks like I can reuse that Lapack trait (which is an export from a separate crate called lax) as-is. It seems to be a very thin but well thought-out wrapper.

Or at least piggy-back on that trait for implementing an own one (since we might want to allow faer as another backend).

[grothesque]

I’ve made some progress: the builder-interface to gemm is now complete. For example one can write

Blas.matmul(&a, &b).scale(3.0).to_owned()

This will allocate a Tensor of unitialized memory and initialize it by the call to gemm. See the tests for complete usage examples.

However the builder interface is not completely generic yet. While there is a trait that provides matmul, the builder methods are not yet covered by a trait.

Comments are most welcome. Any naming suggestions?

[grothesque]

Concerning the names, I wonder for example whether to replace to_onwned by eval (shorter and in-line with mdarray), or scale by times (to better differentiate it from add_to_scaled). matmul itself could be also called gemm or even mm to make it shorter.

[grothesque]

I finished the prototype of a generic interface to gemm! What do you think about it?

Here are question that I ask myself:

• What is the right level of abstraction for the gemm operation? The current interface supports products of matrices with arbitrary layouts (beyond what BLAS itself allows), but the scalar types of the LHS and RHS must be the same, whereas products of a complex with a real matrix are useful in some contexts.

• What about the other BLAS operations? Level 1 & 2, are they even useful with modern compilers? What about the triangular and Hermitian products?

• An instance of MatMulBuilder is in some ways similar to an mdarray expression: it can be used to initialize an array. If there was a way to extend what mdarray’s assign and add_assign methods accept (i.e. some kind of “weak expressions”), one could write: array.assign(Blas.matmul(&a, &b)) or even array -= Blas.matmul(&a, &b). But of course a MatMulBuilder is not a full expression in that it doesn’t support all the operations that mdarray’s expressions do.

[grothesque]

I consider the “prototyping” phase finished, and as such I created the new repository mdarray-linalg. The old mdarray-linalg-prototype repos have been deleted (their history is preserved in the new repo).

So far there is no new functionality, but the prototype code has been reorganized into a workspace containing two crates (these are not on crates.io yet):

• mdarray-linalg: contains the trait definitions (so far only one trait for matrix multiplication)
• mdarray-linalg-blas: implements matrix multiplication for BLAS

We will be adding more functionality soon. We already have working SVD for LAPACK and faer.

[pwulles]

Hi,
I published today the first version of mdarray-linalg. It provides several backends (BLAS, LAPACK, and Faer) for linear algebra operations.
Currently, most of the basic operations and decompositions are covered by BLAS/LAPACK, and a good portion of the decompositions are also implemented through Faer. Some more “textbook” algorithms are also included in a simple, non-optimized Naive backend (Tensordot, PRRLU, ...).
The current priorities are:

• Completing Faer integration as a backend
• Improving documentation (for now, you can check the tests, but proper docs are on the way)
• Adding Nalgebra as a backend
  I hope this project will interest those of you working with linear algebra or multidimensional tensors in Rust.
  Feel free to share any feedback or suggestions you might have about the crate!
  There is probably still a lot of room for improvement, as this is my first project with Rust!

[grothesque]

We took down the crates again after realizing that we can and should absorb mdarray-linalg-tests and mdarray-linalg-naive into the main crate mdarray-linalg, thus reducing the total number of crates by two. We will republish the refactored set of crates in 24h, once the quarantine is over.

Sorry for any inconvenience. Meanwhile, it remains available at https://github.com/grothesque/mdarray-linalg.

[fre-hu]

It looks great! The structure with different backends works out really well. Good documentation also!

[fre-hu]

I have been thinking on the Array and Tensor types and whether it would be better to merge them.

Some drawbacks now is that when evaluating expressions with a generic shape, the result is an Owned implementation which has limited functionality. It is not possible to create new arrays with generic shape, for example with uninit() when writing generic code. It would also simplify the library to have fewer types, and align better to other libraries.

I have a prototype with a merged array, where I have a "storage" type to select the array type. Any allocator type can be used as storage for dynamic arrays, and there is a special type for stack arrays. The storage type can be used to get an internal buffer type and to create new arrays.

Do you think this makes sense, and would it fit mdarray-linalg?

[grothesque]

I have been thinking on the Array and Tensor types and whether it would be better to merge them.

Some drawbacks now is that when evaluating expressions with a generic shape, the result is an Owned implementation which has limited functionality. It is not possible to create new arrays with generic shape, for example with uninit() when writing generic code. It would also simplify the library to have fewer types, and align better to other libraries.

I was not even aware of the Owned trait in mdarray so far.

I see that this is what allows

    let a = array![[1, 2, 3], [4, 5, 6]];
    let r = (&a + &view![1, 2, 3]).eval();

instead of

    let a = array![[1, 2, 3], [4, 5, 6]];
    let r = (&a + &view![1, 2, 3]).eval::<Array<_, _>>();

By the way, while playing with the above, I managed to modify the old variant to eval to a Tensor despite a being an Array (by writing let r = (&a + &view![1, 2, 3]).eval::<Tensor<_, _>>();), I did not manage to find a way to do this with the new variant. (I’m not stating that this is a problem. Just being curious.)

I have a prototype with a merged array, where I have a "storage" type to select the array type. Any allocator type can be used as storage for dynamic arrays, and there is a special type for stack arrays. The storage type can be used to get an internal buffer type and to create new arrays.

Do you think this makes sense, and would it fit mdarray-linalg?

As far as I understand, we would not loose anything because of this unification, or would we?

As far as mdarray-linalg is concerned, so far it does not do anything with expressions or even views. Most functions take &DSlice<_, 2, _> as input and newly computed results are either written into newly allocated Tensor instances, or into a &mut DSlice<_, 2, _> space provided by the user.

This brings me to a point that I wanted to raise:

Currently, most of mdarray-linalg takes arrays as arguments of type &DSlice<_, 2, _>. This implies that the shape is of type (usize, usize), which is not true for Array (for example), where a valid shape could be of type (Const<3>, Const<2>), say. The consequence of this is that our mdarray-linalg’s matrix multiplication facitilies only work with arrays that have an in-memory-stored-shape.

This, however, goes against one, as I see it, central advantages of mdarray: it can impersonate both ndarray (where, at least currently, all dimensions are dynamic), and nalgebra (which is efficient for small fixed-size matrices).

In mdarray-linalg, we have a Naive implementation of our MatMul trait that could in principle be very efficient even for the tiniest matrices, if the matrix size could be passed on as a generic argument up to the point where the actual matrix multiplication happens. But currently fixed-size arrays cannot be used with MatMul without first converting them to views with a (usize, usize) shape...

In order to resolve this, I am considering to modify mdarray-linalg traits to take arguments of type &Slice<_, (_, _), _> instead of &DSlice<_, 2, _>. This would have the double advantage of working for types without a (usize, usize) shape (like the current Array), and of being able to take advantage of type-encoded shape information.

Does the above change sound like a good idea? Or is it perhaps that already today the compiler will be able to understand simple cases like the example program of mdarray-linalg and take advantage of the fact that while the shape dimensions in that example are dynamic, effectively they are compile-time constants. (I probably should have a look at the assembly.)

Is there any conflict with your plan of unifying Array and Tensor? I believe not, but I think it’s crucial that mdarray keeps the possibility to have shapes that feature an arbitrary mix of in-memory and in-type encoding.

---

Here is another aspect that is somewhat related to the above:
In mdarray-linalg, we have exposed BLAS-like functionality such that no heap allocations are performed.
However, LAPACK-like functionality (e.g.) SVD often requires multiple temporary buffers. Currently, we allocate these on the heap by creating temporary uninitialized tensors.

If we had the shape of the input matrix to SVD as a generic mdarray shape, we could in principle allocate the temporary buffers on the stack when the shape is known at compile time. Is there some functionality in mdarray that given generic shape (i.e. a type that implements Shape) allocates an array of element type std::mem::MaybeUninit<_> either on the stack (if the shape is known at compile time), or on the heap (otherwise)?

I believe that this is not possible currently, but perhaps the unification of Array and Tensor would allow for something like that?