linalg-rs / sandbox

An experimental repository to test Rust linear algebra traits
Apache License 2.0
0 stars 0 forks source link

Associated types for Operator traits #2

Open tbetcke opened 1 year ago

tbetcke commented 1 year ago

We require a way to assoicate the type of the underlying Scalar field with Operators and Vectors. My suggestion is to introduce for the operator the associated types:

type In: GeneralScalar;
type Out: General Scalar;

corresponding to the input scalar field and output scalar field (input scalar field and output scalar field may not be identical, e.g. when it passes through an offload device that only accepts f32).

For Vector traits I suggest the associated type

type T: GeneralScalar;

Happy to create a pull request for this.

jedbrown commented 1 year ago

Would this put us on a path toward VectorF32 andVectorC64structs that implementVector` with corresponding associated type? Do you have thoughts on what a mixed precision algorithm might look like? Like suppose you had a Galerkin operation $y \gets P A^{-1} P^T x$ where $x$ is f64 and $A^{-1}$ should be applied using f32. What would the types be in something like this?

p.matvec_t(x, &mut u)?;
a_factored.matvec(&u, &mut v)?; // factors behaving as a matrix? a_factored.solve(...)
p.matvec(&v, &mut y)?;
tbetcke commented 1 year ago

We have the following mappings:

P: f32->f64 P^T: f64 -> f32 A: f32-f32 x: f64 v: f32 y: f64

This will be straight forward if we define for each operator an input and an output type.

matvec would have a signature of the form matvec(x: &Vec<T=Self::In>, y: &mut Vec<T=Self::Out>). Correspondingly, matvec_t would have matvec(x: &Vec<T=Self::Out>, y: &mut Vec<T=Self::In>)

I will create a corresponding pull request. Then we can discuss.

jedbrown commented 1 year ago

Could we have a .with_output::<f32> to return lightweight adaptors that use different precisions? BLIS (which makes no attempt at compile-time type safety) remembers only the representation type of each object (matrix/vector), and allows a non-default "compute precision" associated with the output type. https://github.com/flame/blis/blob/master/docs/MixedDatatypes.md

The implementation handles packing each input to the compute precision (for the microkernel) and unpacking to the output precision. Since (large size) dense linear algebra always has packing, there's no need to match precisions in the type system for user-visible objects. In Rust, I'd imagine this could be done with traits like PackableTo<ComputeScalar>.

This gets messier for sparse matrices because there is normally no packing (it's unclear what to pack when because sparsity isn't structured and it adds nontrivial overhead).

tbetcke commented 1 year ago

Having leightweight adaptors is a good idea. Trying to imagine how non-explicit data types look like though and not fully convinced yet.

In practice, with non-explicit types, I would assume that for example if we want an iterator on a vector we would write:

vec.with_output::<f32>::iter()

to get an iterator to an f32 representation, and correspondingly

vec.with_output::<f64>::iter()

We would always have to use this notation since we do not necessarily know on the user side what the native type of the vector is.

Correspondingly, we would have get and get_mut traits that take an input of a given type and translate it to the native type.

It would be slightly more complicated for complex types. Trying to access those with get for example would have to result in an error.

In principle this would all work. My biggest concern is verbosity about when data conversions take place. As a user of the library I would like to know explicitly when conversions are done. But by calling .with_output::<f32> on a vector or matrix whose type is not explicit I don't know whether a data conversion is internally taking place or not. This is a design choice of expressiveness vs runtime flexibility.

jedbrown commented 1 year ago

Yeah, I don't know how to balance this. For high performance work at larger sizes, we'll generally need access to tiles (even with specified alignment and strides) so iterators won't be over scalar entries and packing would have to be done even if the types aren't converted. It's messier for small sizes and sparse things because conversion has overhead and it's hard to vectorize. With an SpMV $y \gets A x$ where $A$ is f32 and $x$ is f64, it's likely better to convert $x$ up-front unless there are very few nonzeros per row.

It's true that explicit conversions are idiomatic Rust. Some interfaces make their function arguments impl Into<T> so the caller doesn't have to think about it. I wonder if it would make sense to create homogeneous type traits for now, and we could make another layer of type-fluid traits when working on mixed precision algorithms. We could implement both traits for the same types and a caller would normally only use one at a time. The type-fluid traits would mostly be implemented by packing to a compute type (at the appropriate granularity) and calling homogeneous-type functions. I think that would help prevent us from getting bogged down on mixed precision algorithms right now, and users who want explicit types could still do that.

tbetcke commented 1 year ago

I am currently experimenting with type-fluid traits. Assume we have an Operator op of type dyn OperatorBase and a vector of type dyn Vector. The question is how do we access the data?

The operator has an internal data representation and inside its matvec expects a certain type. So we could try to create a method for .with_output<T> for vector that returns some type of VectorTypedView with getters, iterators, etc.

The problem is how to deal with the generic type parameter T. If we add a generic method to vector we are losing object safety. But somewhere we need to specify the type.

If we add associated types to OperatorBase, Vector, etc. we don't have these problems. But we would then specify operators and vectors always with a given native type (maybe we have to do that anyway), somehow going counter full type fluidity.

Is there any way to achieve generic access with a trait object? I have played around with Boxes, etc. but always run into this issue.

tbetcke commented 1 year ago

I have put a proposed implementation using associated types into pull request #4 .