My experience with using LayoutTensors

Mojo Standard Library
1

At the community meeting I was asked to further talk about my experience and bumps I found with trying to use LayoutTensor exclusively for MojoSplat.

WARNING 1: This will probably be outdated information super fast as Mojo evolves

WARNING 2: A lot of the points here can be pointed as my own “Skill Issues”, but nevertheless I think it can be interesting to learn also with my own troubles learning to use them. I am fully open to learning why I am wrong or what I am not understanding here

Off my head here are a few more comprehensive bumps I found in using LayoutTensors for MojoSplat:

  1. Passing them around as arguments can be very tricky. I guess this is mostly a problem of a compiled language with strict typing but I think it is one of the major things breaking the “magic” of the LayoutTensors. For example, it is near impossible to pass them to a function from inside a GPU kernel. This for example is a function i wanted to use but ended up not being able to find a way that worked and instead made all code inline: https://github.com/bertaveira/mojosplat/blob/dfe20e3c7bd13e20ffc06eb04e70a67c76458501/mojosplat/kernels/projection.mojo#L13

    1. This gets exponentially harder if one wants to manipulate the tensor. For example if one slices one dimension this variable becomes a very complex type which is very hard to create a function to accept it, if even possible

  2. LayoutTensor manipulation misses many possibilities. For example slicing a 2D Layout Tensor from a 3D LayoutTensor is very hard and something I really wanted to do. In the end one ends up doing exactly the same code as if we had a regular pointer but with LayoutTensor. Indexing one value at a time.
    Example here I’d love to be able to slice 2D matrix from this 3D LayoutTensore but instead needed to make 2 nested for loops: mojosplat/mojosplat/kernels/projection.mojo at dfe20e3c7bd13e20ffc06eb04e70a67c76458501 · bertaveira/mojosplat · GitHub

  3. Indexing can output complicated types. This likely is my own lack of knowledge but I cannot understand why we get a SIMD when indexing one value. It can cause many annoyances. For example if we try to multiply this SIMD indexed value with another value of another type it just breaks. It is probably my lack of knowledge fo the typing system but the fact that most but not all types are SIMD and the DType is another thing confuses me quite a bit (coming from C++ background).
    E.g. of tiny annoyance is the fact that we need [0] on every index because we need the raw type and not SIMD[1] type mojosplat/mojosplat/kernels/projection.mojo at dfe20e3c7bd13e20ffc06eb04e70a67c76458501 · bertaveira/mojosplat · GitHub

  4. The UNKNOWN_VALUE usage is confusing and very verbose. Not sure it is working as of now but there was a few bugs I reported in the past with it so I am not using it. Even when using it it still compiles the kernel on different shapes since the compiler is not size agnostic

  5. Would be nice to have some bound checks. Not sure what the side effects of doing it, but finding layout issues is very hard

Again most of these are annoyances and maybe lack of my own knowledge of the typing system but hope this thread gives some insight into my own struggles and helps others :upside_down_face:

3 Likes
2

Pingging @denis as this post came from your request