[Help wanted] LLVM ERROR: out of memory on WSL

[kraudy]

Hi, i'm kinda stuck with this code because i'm not sure if the error comes from the code or directly from the environment.

Maybe some one can run it on a native Linux or point my unknown error

Here is the code. Is a simple back propagation

from collections import Optional, List, Dict, InlineList
from memory import UnsafePointer, memset_zero, ArcPointer
from memory import pointer

from utils import Variant
alias ValueOrFloat = Variant[Value, Float32]
alias ValueOrNone  = Variant[Value, NoneType]
alias SelfOrNone  = Variant[Value, NoneType]

fn otro_fun ():
    print("Hello from fun")


@value 
struct Value(AnyType):
    var data: ArcPointer[Float32]
    var grad : ArcPointer[Float32]

    # Maybe these two can be done in the same
    var _prev1 : List[ArcPointer[Value]]
    var _prev2 : List[ArcPointer[Value]]
    var _op : ArcPointer[String]

    fn __init__(out self, data: Float32):
        self.data = ArcPointer[Float32](data)
        self.grad = ArcPointer[Float32](0)


        self._prev1 = List[ArcPointer[Value]]() 
        self._prev2 = List[ArcPointer[Value]]() 

        self._op = ArcPointer[String]('') 
    

    fn __moveinit__(out self, owned existing: Self):
        self.data = existing.data
        self.grad = existing.grad
        self._prev1 = existing._prev1
        self._prev2 = existing._prev2
        self._op = existing._op
    
    fn __add__(self, other: Value) -> Value:
        var out = Value(data = (self.data[] + other.data[]))
        # Maybe i can just append this
        out._prev1 = List[ArcPointer[Value]](self) 
        out._prev2 = List[ArcPointer[Value]](other) 
        out._op = ArcPointer[String]('+')

        return out

    fn __add__(self, other: Float32) -> Value:
        # If the value passed is not Value
        # This isa can be useful to accept multiples types on a paramete
        var v = Value(other)
        # We are only making the conversion and reusing the value logic
        return self.__add__(v)
    
    fn __eq__(self, other: Self) -> Bool:
        #return self is other
        return self.data.__is__(other.data) and
               self.grad.__is__(other.data) and
               self._op.__is__(other._op)
    
    fn __pow__(self, other : Value) -> Value:
        var out = Value(data = (self.data[] ** other.data[])) 
         # We need to add the previous nodes
        out._prev1 = List[ArcPointer[Value]](self) 
        out._prev2 = List[ArcPointer[Value]](other) 
        out._op = ArcPointer[String]('**')

        return out
    
    fn __pow__(self, other: Float32) -> Value:
        var v = Value(other)
        return self.__pow__(v)

    fn backward_pow(mut v: Value):
        # This basically updates the grad
        var vv = v
        print("backward_pow")
        vv.__print()

        if len(v._prev1) == 1 and len(v._prev2) == 1:
            var l = v._prev1[0][].grad[]
            var r = v._prev2[0][].grad[]

            v._prev1[0][].grad = ArcPointer[Float32](l + 
                                      (r * l ** (r - 1) * vv.grad[]) )
    
    fn backward_add(mut v: Value):
        var vv = v
        print("backward_add")
        vv.__print()

        if len(v._prev1) == 1:
            var _children1 = v._prev1[0][]
            print("_children1.grad = ", _children1.grad[], "vv.grad = ",vv.grad[])
            #TODO: Valide which works
            #_children1.grad = ArcPointer[Float32](_children1.grad[] + vv.grad[])
            v._prev1[0][].grad = ArcPointer[Float32](_children1.grad[] + vv.grad[])
            #v.grad = _children1.grad
        
        if len(v._prev2) == 1:
            var _children2 = v._prev2[0][]
            print("_children2.grad = ", _children2.grad[], "vv.grad = ",vv.grad[])
            v._prev2[0][].grad = ArcPointer[Float32](_children2.grad[] + vv.grad[])
    
    fn _backward(mut v: Value):
        var op = String[](v._op[])
        print("op")
        print(op)

        print("_backward")
        v.__print()

        if op == '+':
            print("Option +")
            Value.backward_add(v)
        if op == '**':
            print("Option **")
            Value.backward_pow(v)
        else:
            print("OP not suported")
    
    fn build_topo(self, mut visited: List[ArcPointer[Value]], mut topo: List[ArcPointer[Value]]):
        var is_visited = Bool[](False)
        var size = Int[](len(visited))

        print("Build topo")

        for i in range(size):
            if self == visited[i][]:
                is_visited = True
        
        if not is_visited:
            #visited.append(ArcPointer[Value](self))
            print("Entering visited")
            visited.append(self)
            print(len(visited))
            if len(self._prev1) == 1:
                print("Entered _prev1 == 1")
                var _children1 = self._prev1[0][]
                if len(_children1._prev1) == 1:
                    Value.build_topo(_children1, visited, topo)
                #else:
                #    return

            if len(self._prev2) == 1:
                print("Entered _prev2 == 1")
                var _children2 = self._prev2[0][]
                if len(_children2._prev2) == 1:
                    Value.build_topo(_children2, visited, topo)
                #else:
                #    return
            
            topo.append(self)
            print(len(topo))

    
    fn backward(mut self):
        var visited = List[ArcPointer[Value]]()
        var topo = List[ArcPointer[Value]]()

        Value.build_topo(self, visited, topo)

        self.grad = Float32(1.0)
        var reversed_topo = reversed(topo) # this returns an iterator

        #for ref in reversed_topo:
        for i in range(len(topo), -1, -1):
            #var v = ref[]
            print("i: ", i)
            # If there is no elements, this gives error, maybe use try catch
            var v = topo[i][]
            print("Previous _backward: ")
            v.__print()
            Value._backward(v)
    
    fn __print(self):
        print("data: ", self.data[], "grad: ", self.grad[])
    
            
def main():
    var a = Value(data = 1.0)
    var b = Value(data = 2.0)
    a.__print()

    var c = a + b
    c.__print()

    var d = c + Float32(3.0)
    d.__print()

    c._prev1[0][].__print()
    c._prev2[0][].__print()

    c.backward()
    print("Resultado ====")
    c.__print()
    c._prev1[0][].__print()
    c._prev2[0][].__print()


    var f = Value(data = 1.0)
    var g = Value(data = 2.0)
    var h = f + g
    #this breaks
    h.backward()

The system

Linux Ubuntu 5.15.146.1-microsoft-standard-WSL2 #1 SMP Thu Jan 11 04:09:03 UTC 2024 x86_64 x86_64 x86_64 GNU/Linux

mojo version

mojo 24.6.0 (4487cd6e)

Error

LLVM ERROR: out of memory
Allocation failed
[187884:187884:20250118,172648.512482:ERROR file_io_posix.cc:144] open /sys/devices/system/cpu/cpu0/cpufreq/scaling_cur_freq: No such file or directory (2)
[187884:187884:20250118,172648.512585:ERROR file_io_posix.cc:144] open /sys/devices/system/cpu/cpu0/cpufreq/scaling_max_freq: No such file or directory (2)
Please submit a bug report to https://github.com/modularml/mojo/issues and include the crash backtrace along with all the relevant source codes.
Stack dump:
0.      Program arguments: mojo engine.mojo --debug-level = full
 #0 0x000055add47e26bb llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) Signals.cpp:0:0
 #1 0x000055add47e0619 llvm::sys::RunSignalHandlers() Signals.cpp:0:0
 #2 0x000055add47e2d3d SignalHandler(int) Signals.cpp:0:0
 #3 0x00007f961c717520 (/lib/x86_64-linux-gnu/libc.so.6+0x42520)
 #4 0x00007f961c76b9fc __pthread_kill_implementation ./nptl/./nptl/pthread_kill.c:44:76
 #5 0x00007f961c76b9fc __pthread_kill_internal ./nptl/./nptl/pthread_kill.c:78:10
 #6 0x00007f961c76b9fc pthread_kill ./nptl/./nptl/pthread_kill.c:89:10
 #7 0x00007f961c717476 gsignal ./signal/../sysdeps/posix/raise.c:27:6
 #8 0x00007f961c6fd7f3 abort ./stdlib/./stdlib/abort.c:81:7
 #9 0x000055add47b2da9 llvm::report_bad_alloc_error(char const*, bool) ErrorHandling.cpp:0:0
#10 0x000055add47b2df5 (/home/kraudy/Mojo/mojo-grad/.magic/envs/default/bin/mojo+0x5715df5)
#11 0x00007f961cbff3e3 (anonymous namespace)::handle_oom(void* (*)(void*), void*, bool, bool) tcmalloc.cc:0:0
#12 0x00007f961cc0e1a4 tcmalloc::allocate_full_cpp_throw_oom(unsigned long) tcmalloc.cc:0:0
#13 0x00007f95a80063b5 
mojo crashed!