Implementing complex tensor operations like det()

[softmaximalist]

I am interested in implementing complex tensor operations like det() (determinant) but I am not really sure about their implementation process. I will mostly write my current understanding along with some questions I have. I am not really confident about knowledge of Burn backends and tensor operations so please let me know if I am misunderstanding something.

---

My current understanding is that there are two types of tensor operations: primitive (cannot be composed of other tensor ops) and non-primitive (more complex tensor ops that are computed using other primitive and non-primitive tensor ops). I assume most (or all) of the primitive tensor ops have already been implemented given that they are foundational operations. On the other hand, most non-primitive tensor operations can be implemented either using existing Burn tensor operations (a high level implementation) such as lingalg::diag() or implemented as an operation that does not rely on (most?) existing Burn tensor operations (a lower level implementation). When we need to implement a complex and computationally expensive non-primitive tensor operation like det(), lower-level implementation is prefered due to performance reasons.

When we want to implement these types of complex non-primitive tensor operations, we generally want to make the Backend perform the actual computation either by delegating the operation to a Backend's native operation or by writing a custom cubecl kernel.

I think the torch backend is special because it provides functions like det() which we can simply call when the torch backend is selected. On the other hand, for other backends, I think we need to write cubecl kernels. Am I right? However, there are already many high performance implementations of linear algebra operations in many different languages such as BLAS, LAPACK, cuBLAS, cuSOLVER, faes-rs, etc. If we do need to write a cubecl kernel for det(), my question is why can't we rely on these existing implementations (or at least use some code from these existing implementations to implement det() or other complex tensor ops)? Is it because of the language conflict (rust vs fortran, c++, etc)? Is it because kernel fusions, autodiff, etc will not work? Or is there some other reason?

[antimora]

Linking an existing ticket: #4609

[softmaximalist]

I am aware of this issue and I have implemented some tensor operations before. However, I am still not quite sure about the implementation process of complex tensor operations. Please note that my question above is about the general implementation process of complex tensor operations (not just det()). The det() operation is just one of the many operations I want to implement.

[laggui]

My current understanding is that there are two types of tensor operations: primitive (cannot be composed of other tensor ops) and non-primitive (more complex tensor ops that are computed using other primitive and non-primitive tensor ops).

It's a bit more nuanced than that.

Primitive ops must be implemented by backends and non-primitive ops have default implementations built from primitives.

Some ops have a default implementation that can be overridden by backends for performance.

For det(), I don't think you would write an actual kernel for that. It would likely be composed of other linalg ops to compute the determinant, such as LU decomposition. LU decomposition in tensor::linalg currently has no backend specialization, only a basic implementation using high-level tensor operations. This is not optimal, hence the performance note in the docs. It would need to be optimized. And then could be reused to compute the determinant.

I think the torch backend is special because it provides functions like det() which we can simply call when the torch backend is selected. On the other hand, for other backends, I think we need to write cubecl kernels. Am I right?

Yeah torch is a special case, because it already exposes an API to compute a lot of the tensor operations. So it could be specialized.

However, there are already many high performance implementations of linear algebra operations in many different languages such as BLAS, LAPACK, cuBLAS, cuSOLVER, faes-rs, etc. If we do need to write a cubecl kernel for det(), my question is why can't we rely on these existing implementations?

Using specialized libraries like cuBLAS or BLAS doesn't break the Burn / CubeCL abstractions per se, but it does reduce portability and composability. Those libraries are backend-specific, while CubeCL aims to provide a portable execution model across targets which can be heavily optimized. We want to leverage that as much as possible.

That said, not all “operations” are the same. Some problems, like determinants or LU decompositions, are not really single operations but multi-step algorithms involving synchronization, data-dependent control flow (e.g. pivoting), and multiple kernel launches. While these can be expressed in CubeCL, achieving high performance requires careful global scheduling and numerical considerations, not just fast kernels. It's very different from adding support for another element-wise operation that can be defined using an intrinsic.

There are also cases like NCCL, where the problem is fundamentally about inter-device communication and hardware topology rather than computation. They're more difficult to abstract and generalize. I think we're going to add some support in cubecl to leverage specialized libraries to bridge the gap.