Trinary Contractions

[SebastianSemper]

cuTensor 2.2 introduces a feature of trinary tensor contraction [1] which contracts three tensors in one function. It looks like cupy might be getting support for this [2]. Can opt_einsum be extended to support the possibility for trinary contractions?

[1] https://docs.nvidia.com/cuda/cutensor/latest/api/cutensor.html#cutensorcreatecontractiontrinary
[2] cupy/cupy#9041

[dgasmith]

When you say "support" what are you thinking of here?

The simplest thing we can do is not modify the optimization algorithms and then group the results into be "good" trinary computations. Modifying the algorithms to support trinary operations without intermediates would be more difficult and we would need heuristics when this appropriate (e.g., avoiding large intermediates).

My guess is letting this mature a bit with multiple implementations at higher levels than cutensor is appropriate as this gets outside of how we currently define path planning and into more subtle implementation details which will vary much more on the implementation. For example we always use BLAS where we can; however, depending on the shape and implementation einsum may be faster. This usually isn't a problem due to the shapes that people use opt_einsum, but I see trinary being a bit more of a foot gun in this way as there is more heuristics for these higher order operations.

If there are good short term gains I would be very interested in discussing!