How is it possible for a function to return a data type that is different from the data type expected by the calling function?

[redbov]

I have a question related to a code snippet from the mojo docs: https://docs.modular.com/mojo/notebooks/Matmul. The code snippet is the one exactly after the text: "Adding types to the Python implementation" in the docs.

The code:

alias type = DType.float32

struct Matrix[rows: Int, cols: Int]:
    var data: DTypePointer[type]

    # Initialize zeroeing all values
    fn __init__(inout self):
        self.data = DTypePointer[type].alloc(rows * cols)
        memset_zero(self.data, rows * cols)

    # Initialize taking a pointer, don't set any elements
    fn __init__(inout self, data: DTypePointer[type]):
        self.data = data

    # Initialize with random values
    @staticmethod
    fn rand() -> Self:
        var data = DTypePointer[type].alloc(rows * cols)
        rand(data, rows * cols)
        return Self(data)

    fn __getitem__(self, y: Int, x: Int) -> Scalar[type]:
        return self.load[1](y, x)

    fn __setitem__(self, y: Int, x: Int, val: Scalar[type]):
        self.store[1](y, x, val)

    fn load[nelts: Int](self, y: Int, x: Int) -> SIMD[type, nelts]:
        return self.data.load[width=nelts](y * self.cols + x)

    fn store[nelts: Int](self, y: Int, x: Int, val: SIMD[type, nelts]):
        return self.data.store[width=nelts](y * self.cols + x, val)

Here, getitem() returns Scalar[type] and fn load() returns SIMD[type, nelts]. These functions are connected, because getitem() calls fn load().

My question is the following: The data type, that gets returned from fn load() to getitem is a SIMD, which is I think a vector, containing 1 value that is DType.float32 . But after this, getitem() has to return a Scalar[type], which in this case is DType.float32. The same happens for setitem() and fn store().

So how can getitem() receive one data type and return another one? Is there an implicit conversion happening from a SIMD vector of a single element to a scalar? If so, could you please explain how this works?

Also is this for all the functions in mojo, or is it a dunder method or just a regular struct method thing? Or is it because in the current way of implementation, fn load() only returns a SIMD vector with only 1 element? So what would happen if the load function is used within getitem with nelts set to more than 1? It seems like it would return a SIMD vector, which would not be compatible with the expected return type of getitem.

[Benny-Nottonson]

There is one level of implicit casting allowed in Mojo returns. I.E. if I have a struct that has an initializer with the signature

struct Example:
    fn __init__(inout self, data: Scalar[DType.float32]):
        ...

If I then have a function that returns a Float32, it can be implicitly converted by the compiler, i.e.

fn make_example(val: Float32) -> Example:
    return val + 1

Will run as intended and return an Example struct.

[Benny-Nottonson]

This is known formally as implicit type conversion (coercion) between types, common in many other languages.

[masters3d]

Does this implicit conversion incur a runtime cost? (I imagine some conversions could be expensive)
I wonder if this implicit conversion is part of the mojo "spec" or required when dealing with python.

[Benny-Nottonson]

Does this implicit conversion incur a runtime cost? (I imagine some conversions could be expensive)

I wonder if this implicit conversion is part of the mojo "spec" or required when dealing with python.

It would obviously depend on the use case, in the example I gave, pretty much no, if there is any cost it is trivially small. You are correct however that there is some runtime cost for more complex structures, so the full answer is yes. (Although I suspect Mojo optimizes away a lot of this when it knows it will need an implicit cast, but that is out of the scope of my knowledge)