mojo icon indicating copy to clipboard operation
mojo copied to clipboard

[stdlib] Add optional small buffer optimization in `List`

Open gabrieldemarmiesse opened this issue 9 months ago • 2 comments

Related to https://github.com/modularml/mojo/issues/2467

This is in the work for SSO. I'm trying things and I'd like to gather community feedback.

At first, I wanted to implement SSO using Variant[InlineList, List], while that would have worked, we would still need to look at the flag of the variant for pretty much any operation we need to do. For example reading the first element in the list of bytes.

If we instead store the pointer to the data being used, the size of the data and the capacity, then the pointer can point to the stack or the heap, we just follow the pointer and that's it. No if-else necessary. We can know if we are in the stack or the heap by looking at the address of the pointer. We don't even require a if-else to get the size of the data since it's always stored at the same spot.

I was planning to implement this using a custom structure, but in the end, I thought that I could use the powerful metaprogramming features of Mojo to add the small buffer optimisation (SBO) directly to the List[] struct. By default it would be disabled (=0) and we can use @parameter if at the few places needed, for example append() to check if we are using the SBO or not.

In theory, this should all be zero cost when unused, but I'm not sure right now. We also don't have a good benchmarking story. Also there might be some incompatibility if users make functions that don't take into account this new argument. e.g.

fn get_size(x: List[Float64]) -> Int:  # Only accepts List without SBO
    return len(x)

So I am trying to get some community feedback here, as well as feedback from the maintainers, to understand if this makes sense or not. If not, I can implement a whole new struct, but there might be a lot of duplication.

I'll continue to work on this a bit (it's a good way to have a quick prototype up and running) and see if I can beat the numbers I got in this PR:

  • https://github.com/modularml/mojo/pull/2507

EDIT: I run every List tests with 10 different buffer sizes (0 to 9), so I believe we can consider that it works. I'm currently working on using it inside a String. And it's having memory issues left and right. I wouldn't be surprised if String was misusing List in a few places. Most of the diff is the unit tests being places inside a for-loop. The changes for SBO itself are pretty small.

EDIT 2: Despite passing the full test suite of List, I found something that didn't crash with the previous list but crashes with this one:

fn foo():
    alias my_list: List[Int8, 10] = List[Int8, 10](0, 1, 2)
    print("Materializing my_list")
    var my_list_materialized = my_list
    print("all done, exiting function")


def main():
    foo()
    print("main exiting.")

EDIT 3: The whole report and benchmarks are here: https://github.com/modularml/mojo/issues/2467#issuecomment-2106263163

EDIT 4: The bug report is here: https://github.com/modularml/mojo/issues/2637

gabrieldemarmiesse avatar May 10 '24 20:05 gabrieldemarmiesse

This is interesting, my initial thoughts are I like this idea, as long as we can get all the same behavior. Carrying the extra member for the SBO doesn't add any size to the struct (when SBO=0) or anything like that, so maybe it's possible to get it practically seamless.

helehex avatar May 10 '24 20:05 helehex

If there's going to be a std Array type, we could even push it out to that

helehex avatar May 16 '24 00:05 helehex

If there's going to be a std Array type, we could even push it out to that

InlineArray is exactly that, no?

JoeLoser avatar May 22 '24 16:05 JoeLoser

If there's going to be a std Array type, we could even push it out to that

InlineArray is exactly that, no?

more like the current InlinedFixedVector, but thats only for deprecated AnyRegType's

helehex avatar May 22 '24 22:05 helehex

Closing in favor of https://github.com/modularml/mojo/pull/2827

gabrieldemarmiesse avatar May 26 '24 18:05 gabrieldemarmiesse