Modular•11mo ago

List of memory-only structs

I am trying to create a struct for a neural network layer with memory-only list of neurons:

struct Neuron:
    '''Implementation of a neuron.'''
    var weights: Tensor[DType.float64]
    var bias: Float64
    var nin: Int

    fn __init__(inout self, nin: Int):
        self.weights = rand[DType.float64](nin)
        self.bias = random_float64(-1.0, 1.0)
        self.nin = nin
    
    fn __call__(self, x: Tensor[DType.float64]) raises -> Float64:
        if x.num_elements() != self.nin:
            raise Error('Input values do not match number of weights')
        
        var sum: Float64 = self.bias
        for i in range(self.nin):
            sum += x[i] * self.weights[i]
        return math.tanh(sum)

struct Neuron:
    '''Implementation of a neuron.'''
    var weights: Tensor[DType.float64]
    var bias: Float64
    var nin: Int

    fn __init__(inout self, nin: Int):
        self.weights = rand[DType.float64](nin)
        self.bias = random_float64(-1.0, 1.0)
        self.nin = nin
    
    fn __call__(self, x: Tensor[DType.float64]) raises -> Float64:
        if x.num_elements() != self.nin:
            raise Error('Input values do not match number of weights')
        
        var sum: Float64 = self.bias
        for i in range(self.nin):
            sum += x[i] * self.weights[i]
        return math.tanh(sum)

But I'm struggling to compose these into a list of any sort, whether with a heap array, vector, or variadic list, getting this error: candidate not viable: method argument #1 cannot bind generic !mlirtype to memory-only type 'Neuron'. Is there a way to compose structs in this way?

6 Replies

Stole•11mo ago

Yeah, Pointers to any type that isn't a "simple" struct (not register_passable) don't work. I wrote a more detailed explanation here: https://discord.com/channels/1087530497313357884/1098713601386233997/1164304975816556655, albeit for a slightly different use case. The main point is the same, if you want such a thing you're essentially just going to be managing the memory yourself. Here, you can mark Neuron with @register_passable, then maybe add a copy method that then copies the struct's Tensor over so we don't duplicate pointers, and other similar operations.

Discord

Discord - A New Way to Chat with Friends & Communities

Discord is the easiest way to communicate over voice, video, and text. Chat, hang out, and stay close with your friends and communities.

Stijn•11mo ago

@Stole @seb Currently running into the same error. I tried to store the pointers of your class in a DynamicVector and that works. But when you are trying to dereference again it produces the same error. Any of you, or anyone else, that have found a workaround while we are waiting for traits? I'm here now (trying to work around it)

struct MyClass[dtype: DType]:
    var tensor: Tensor[dtype]
    fn init(inout self, tensor: Tensor[dtype]):
        self.tensor = tensor


fn main():
    alias dtype = DType.float32
    alias nelts: Int = simdwidthof[dtype]()

    let tensor: Tensor[dtype] = rand[dtype](2, 5)
    var myelement = MyClass[dtype](tensor)


    let xPtr = Pointer[MyClass[dtype]].address_of(myelement)  # get address memory of x
    print(xPtr.as_index())  # print address memory of x

    var vec = DynamicVector[Pointer[MyClass[dtype]]]()
    vec.push_back(xPtr)
    for i in range(vec.size):
        print(vec[i].as_index())

    let x : MyClass[dtype] = xPtr.load()
    print(x.tensor)

struct MyClass[dtype: DType]:
    var tensor: Tensor[dtype]
    fn init(inout self, tensor: Tensor[dtype]):
        self.tensor = tensor


fn main():
    alias dtype = DType.float32
    alias nelts: Int = simdwidthof[dtype]()

    let tensor: Tensor[dtype] = rand[dtype](2, 5)
    var myelement = MyClass[dtype](tensor)


    let xPtr = Pointer[MyClass[dtype]].address_of(myelement)  # get address memory of x
    print(xPtr.as_index())  # print address memory of x

    var vec = DynamicVector[Pointer[MyClass[dtype]]]()
    vec.push_back(xPtr)
    for i in range(vec.size):
        print(vec[i].as_index())

    let x : MyClass[dtype] = xPtr.load()
    print(x.tensor)

But you can't do the pointer.load() as well. I know you could probably do it with Low level IR. But i would rather not go down that rabithole

seb•11mo ago

@Stijn I ended up abandoning a struct-based approach to neural nets; couldnt figure out any other workarounds without using register_passable("trivial"). However, if that class were only to contain things like pointers, numbers, and dynamic vectors, its size might not be so large as to make full copying prohibitive. One thing: If your register_passable("trivial") struct does contain pointers, you might want to implement a manual delete() method of some sort on the class which frees the pointer, and call that method in the __del__() method of whatever class owns the passable class, so that you don't end up with memory leaks

Stijn•11mo ago

@seb This looks like a working workaround. Will have a better look at it later: cfr. @rd4com : new small tutorial: 🤹 making lists of structs with magic operators (until lifetimes) https://github.com/rd4com/mojo-learning/blob/main/tutorials/lists-of-structs-magic-operators-pre-lifetimes.md

Michael K•11mo ago

NDBuffer is register_passable and seems suitable for neural nets. I am not sure if it is 'trivial' but it is efficiently copied. You can have a Network struct with weights stored as NDBuffers or DynamicVectors of NDBuffers. The network needs to allocate the pointers when the weights are loaded or created for training and then free them in its del function. It is manually managing the memory but it is not that cumbersome. I have also built my own simple struct using `@register_passable('trivial') but I don't see any performance difference from using NDBuffer so it is probably easier to use that.

Modular Docs - buffer

Module

Stijn•11mo ago

Just found this: https://docs.modular.com/mojo/stdlib/builtin/builtin_list.html#variadiclistmem. Just handling pointers without having to deal with register_passable("trivial")

Gaming

Programming

List of memory-only structs