Over-engineered? Generics over item type and item ownership

help
5

Is the set of tensor types (dense, sparse...) known in one central place at compile time? If so, you could also try an enum-based approach as an alternative to a trait/generics approach:

enum Tensor {
    Dense(DenseTensor),
    Sparse(SparseTensor),
    ...
}

As I wrote here a long time ago, the tradeoff between trait objects and enums has many sides. To paraphrase my old comment with some extra considerations applicable to your specific problem...

  • With enums, you don't need to Box and Arc everything (although you will likely want a Box or Arc somewhere for the actual floating-point data, at least for larger tensors. For smaller tensors, something like the smallvec approach may allow you to avoid those tiny allocations that are more expensive than actually manipulating the tensor behind them.)
  • In exchange for this avoidance of heap allocation, you will pay the price of having every individual Tensor whose subtype the compiler does not infer use a slot on the stack which is as large as the largest Tensor subtype. This can become a problem if the subtypes have a widely different stack size.
  • With enums, dispatching over the right contraction implementation becomes as simple as a match (tensor1, tensor2), which can be compiled into a simple jump table when the compiler's optimizer is in a good mood. Code for this will also be a lot easier to read than an emulation of multiple dispatch over Rust traits, which only provide single dispatch.
  • Empirically, the compiler optimizer seems better at "de-enumifying" than devirtualizing traits in contexts where it actually knows at compile time which type of tensor you're dealing with.
  • In terms of code quality and cache locality, the trait approach tends to break down when you have too many methods, whereas the enum approach tends to break down when you have too many tensor types.

Concerning f32 vs f64 genericity, as we discussed previously I think that's extremely expensive in terms of code complexity with the current state of num-traits, which is why I try to get as much mileage as I can out of the feature-based approach until it breaks down and I am absolutely forced to implement such genericity.

1 Like