Declare a LayoutTensor in a struct?

Mojo
1

Hi I have a 2D array in an UnsafePointer.

For convenience I would like to use a LayoutTensor all wrapped in a struct:

Problem is that mutability is inferred only in LayoutTensor hence I cant declare the field.

Here is kind of what I am trying to do:

alias MyLayout = Layout.row_major(1024, 1024)

struct MyStruct:
    var data: UnsafePointer[UInt32]
    var tensor: LayoutTensor[DType.uint32, MyLayout]   <--- This guy xD

    fn __init__(out self):
        self.data = UnsafePointer[UInt32].alloc(1024 * 1024)
        self.tensor = LayoutTensor(self.data)

As someone how is trying to understand the Mojonic way to represent a 2D array (compile time size known)?

2

This works with specifying the origin

from layout import Layout, LayoutTensor
from memory import UnsafePointer

alias MyLayout = Layout.row_major(1024, 1024)


struct MyStruct:
    var _data: UnsafePointer[UInt32]
    var tensor: LayoutTensor[mut=True, DType.uint32, MyLayout, MutableAnyOrigin]

    fn __init__(out self):
        self._data = UnsafePointer[UInt32].alloc(1024 * 1024)
        self.tensor = LayoutTensor[
            mut=True, DType.uint32, MyLayout, MutableAnyOrigin
        ](self._data)

    fn __del__(owned self):
        self._data.free()


fn main():
    var my_struct = MyStruct()
    _ = my_struct.tensor.fill(1)
    print(my_struct.tensor[0, 0])
3

This is it. Thanks so much!!

1 Like
4

This topic was automatically closed 7 days after the last reply. New replies are no longer allowed.