Need control over where allocations go (registers vs. stack vs. heap)

[dsharlet]

Hi all, I started playing with mojo in the last few days. I wrote a more comprehensive set of notes from my experience here: https://github.com/dsharlet/mojo_comments

One of the issues mentioned there is it would be really helpful to be able to declare a chunk of intermediate storage and control where it goes (registers, stack, heap). Right now, the only way I can find to allocate memory looks like a heap allocation, which is not very appropriate for use cases like the one I documented above.

As far as I can tell, it's not possible to have some kind of array "native" to the language itself (much like Python), which seems like it will be necessary in order to have allocations be placed on the stack.

#400 (sort of) mentions this as should be deferred to library authors, but I'm not sure that makes sense for allocations intended to be on the stack. This seems like it needs compiler support and understanding, maybe even some higher level conceptual support (it's not necessarily the case that a stack is conceptually a thing at all, and just an implementation detail...). This is not currently mentioned in the roadmap (unless I missed it), and I think it might require some deep thought with how this interacts with the language design, so doing it early could be important.

Maybe just adding some alloca-like thing that works with the pointer types used for heap allocations currently would work? It would be a headache to use though, because presumably it couldn't be used inside a helper function, like all the existing heap allocation usages I've seen (such as to make a Matrix object).

[rarebreed]

Perhaps I am mistaken, but I assumed the following:

• heap : only through the Pointer[T].alloc or DTypePointer[T].alloc method currently
  • are there other ways?
• register
  • if decorated with @register_passable("trivial") , it goes to some register (general purpose?)
  • but would an SIMD variable, load to an available SIMD register? and which one? (avx, sse, etc?)
    • eg not stored in a Buffer, but just a let simd = SIMD[f32, 4] declaration
• stack:
  • local vars in fn body and probably top-level variables once that is enabled in mojo
  • through the memory.stack_allocation method

I agree that it would be nice to make this more clear in the documentation.

As for creating an array on the stack, I believe you should be able to create that through the aforementioned stack_allocate method. It'd be nice to have syntax sugar to create them with a literal (and differentiate dynamically growable vector types from statically sized array types that can all live on the stack).

But mojo is very young, and I think a LOT is going to change. For example, they want to make this a memory safe language (unlike C++). But without lifetime tracking (and mutable reference aliasing tracking) mojo will be (and currently is) a massively unsafe language. Since they are still working on lifetimes, I am curious how similar it will be to rust and what changes it will bring about.

In fact, if modular wants to entice away rustaceans to mojo, they will need to make sure it's a memory-safe language (I was surprised...though I should not have been...to see segmentation faults again after being mercifully free of them for so long using rust). I am actually curious how mojo is going to handle the equivalent of rust's unsafe I wonder if they will make the Pointer type only usable behind some kind of unsafe declaration (like rust does), and only allow references (once they have their lifetime system fleshed out) as part of the "safe" language.

In rust, there are only a handful of ways to heap allocate, but more importantly, rust's ownership rules track who is responsible for freeing the memory (and rust's drop method is called implicitly by the compiler when the variable goes out of scope).

[dsharlet]

Thanks for the stack_allocation link, that does seem like what I was looking for! (Seems like alloca.) Unfortunately, the performance characteristics of this function are not what I expected: I'm still seeing a huge performance boost by lifting the call to memory.stack_allocation to be outside all of my loops, and in general performs almost the same as heap allocation. Other tests I tried indicate it really is a stack allocation (e.g. using free() on it fails like one would expect, and allocating 1 GB with stack_allocation causes a stack overflow-ish looking failure). So... I am not sure what to think.

[dsharlet]

I found the problem here. The Matrix object in matrix.mojo initializes its data with rand in its constructor :) So, now stack allocations perform as expected.

[jackos]

Great answer @rarebreed thanks, for the question on using unsafe features we're likely to go with the Swift approach of naming things as unsafe, rather than Rust's approach of coloring with unsafe which also has different behavior to normal code. But this is subject to change as the language matures.