[Feature Request] Unify SSO between InlinedString and String type

[gabrieldemarmiesse]

Review Mojo's priorities

• I have read the roadmap and priorities and I believe this request falls within the priorities.

What is your request?

We currently have https://docs.modular.com/mojo/stdlib/utils/inlined_string#inlinedstring and String. Does it make sense to have SSO enabled in String too? Is it something that's currently possible or would that be too much trouble for the compiler?

What is your motivation for this change?

I guess that having SSO enabled on String would provide speed improvements. Furthermore, as we are writing more and more functions that work with Strings, I believe that we'll be avoiding a big refactor if we want SSO afterwards.

Any other details?

Maybe that would require to make the internal buffer of String actually private? I wonder if some functions use the underlying pointer or List directly.

[JoeLoser]

Thanks for filing this! We actually discussed this exact topic earlier in the week in the team's design meeting. Here are my notes:

• One first step is to remove _ArrayMem in utils/inlined_string and use InlineArray directly. We may run into trouble with our buffer formatting things in builtin/int.mojo and builtin/hex.mojo since they use the default constructor, which previously left undef values in the backing pop.array whereas we don’t allow undef values in default constructor in the new InlineArray for safety purposes, so you’ll have to pick a default fill value (such as 0) I suspect (FYI @lsh).

With SSO specifically, we talked about a few things like ABI concerns and the like, but decided on the following:

• Add SSO parameter to the existing String type. Make it a keyword-only parameter such as _Size that defaults to the fixed SSO size (such as the existing value of 24). This will make it so it's clearly a "use at your own risk" if you explicitly set this parameter in caller code (by needing to specify the keyword parameter). This is because different parameters for the SSO value would generate a distinct String type in our elaborator and have the same consequences as "template explosion" in C++ and make it difficult to use different types across APIs, e.g. String and String[_Size=32] for example. One thing to note is nothing is ABI stable in the library, and so we have freedom to change the details of the SSO which can affect the layout of the type. This is also one reason why the keyword was chosen to have an underscore prefix.
• Remove InlinedString.
• Keep FixedString as is. It has value still, but we can safely remove InlinedString which was being that initial stepping stone as we were thinking about SSO.

There are some future enhancements we can make for SSO implementation such as:

• Use tagged pointers instead of tagging the variant. This will save one byte here.

How does this sound to you and @lsh (who was asking me about this offline the other day)?

[gabrieldemarmiesse]

Thanks for the detailed plan. Looks good overall, I'm wondering about one thing:

One first step is to remove _ArrayMem in utils/inlined_strings and use InlineArray directly.

If we do this step mentionned afterwards:

Remove InlinedString.

Wouldn't that be work that we'll quickly throw away? If I understand correctly the goals here, removing _ArrayMem from InlinedString is not useful as InlinedString will be removed too.

[JoeLoser]

Thanks for the detailed plan. Looks good overall, I'm wondering about one thing:

One first step is to remove _ArrayMem in utils/inlined_strings and use InlineArray directly.

If we do this step mentionned afterwards:

Remove InlinedString.

Wouldn't that be work that we'll quickly throw away? If I understand correctly the goals here, removing _ArrayMem from InlinedString is not useful as InlinedString will be removed too.

Ah yes, good point. Sorry about that: feel free to skip the step about removing _ArrayMem. Was still drinking coffee and waking up when replying earlier 😄

[gabrieldemarmiesse]

I'm going through the code right now and I'm wondering if it makes sense to have a struct like "SmallSizeOptimizedListOfBytes" (name TBD) which has SSO but the interface of a list. I think the number of methods to implement would be smaller and the String stuct would use it as a list, it would be transparent and not require many changes compared to now. Does that make sense or am I missing something?

[gabrieldemarmiesse]

I'm posting here the benchmarks results as it's better to have them here than in a PR.

SBO: Small buffer optimization
SSO: Short string optimization (SBO applied to strings)

Pull requests (draft only, they're proof of concept)

Here are the relevant PR and highlights, note that all stdlib tests are passing for each PR, with the exception of the materialization bug:

• [stdlib] Add small string optimization to the String struct #2507 Use a Variant. Used to produce two columns on the next table.
• [stdlib] Implement SSO using small buffer optimization on List #2632 Use a pointer that points to the stack or the heap. ⚠️ List and String not materializable from alias when SBO is in use: possible compiler bug. It's the last column on the benchmark table.

This PR was here only to show what is needed to implement SBO in List:

• [stdlib] Add optional small buffer optimization in List #2613 ⚠️ List not materializable from alias when SBO is in use: possible compiler bug

The materialization compiler bug

It can be observed by doing a checkout on this PR:

• [stdlib] Add optional small buffer optimization in List #2613

and running the following code:

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


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

This produces a src/tcmalloc.cc:302] Attempt to free invalid pointer 0x7fff4f1bcf60. Note that using the list at compile time is okay, and at runtime is fine too. The bug appear when converting an alias to a var. I suspect this is because the compiler doesn't like that a pointer can point to something on the stack or the heap dynamically.

The full bug report with minimal reproducible example on nightly is here: #2637

Benchmarks

System: Debian bookworm/sid in Docker in WSL2
Processor Intel(R) Core(TM) i7-10700KF CPU @ 3.80GHz 3.80 GHz
Installed RAM 16,0 GB
System type 64-bit operating system, x64-based processor
Edition Windows 11 Pro
Version 21H2
Installed on ‎01/‎07/‎2021
OS build 22000.2538

Results

Benchmark id	nightly 11/05/2024 (836626b)	Using Variant with custom inline bytes list (3271101)	Using Variant with InlineList (59f19bc)	Using SBO in List directly (f27b43a)
1	212750 us	70254 us (x3.02)	72997 us (x2.91)	71117 us (x2.98)
2	69441 us	85479 us (x0.81)	84488 us (x0.82)	64898 us (x1.07)
3	35758 us	8291 us (x4.36)	8423 us (x4.29)	6066 us (x5.97)
4	189726 us	225251 us (x0.85)	221874 us (x0.86)	195324 us (x0.98)
size of buffer (in bytes)	0	24	24	16
size of String (in bytes)	24	40	40	40

The benchmark code

import time


fn get_dummy_int_list(n: Int) -> List[Int]:
    var result = List[Int](capacity=n)
    for i in range(n):
        result.append(i)
    return result

fn get_json_looking_string() -> String:
    var int_list = get_dummy_int_list(50000)
    # dang this is slow
    return __type_of(int_list).__str__(int_list)

fn benchmark_1(list_of_ints: List[Int]):
    var result_list = List[String](capacity=len(list_of_ints))
    for i in range(len(list_of_ints)):
        result_list.append(str(list_of_ints[i]))

fn benchmark_2():
    var result = String()
    for _ in range(10000):
        result += "a"
        result += "b"
        result += "c"
        result = result[:-2]

fn benchmark_3(json_looking_string: String):
    # we want to get a list of ints from a string looking like [1, 2,3 ,4,5,6, 7,8 ,9,10]
    var stripped = json_looking_string.removeprefix("[").removesuffix("]")
    var split: List[String]
    try:
        split = stripped.split(",")
    except:
        abort("this should never happen, split didn't work")
        split = List[String]("hello")
    var result = List[Int](capacity=len(split))
    for x in split:
        try:
            result.append(int(x[].strip()))
        except e:
            abort("this should never happen, error" + str(e))
            result.append(0)

fn benchmark_4(input_list: List[Int]):
    var result = __type_of(input_list).__str__(input_list)


def main():
    print("starting benchmarks")
    
    var list_of_ints = get_dummy_int_list(1000000)
    for i in range(10):
        benchmark_1(list_of_ints)
    t1 = time.now()
    benchmark_1(list_of_ints)
    t2 = time.now()
    total_time = (t2-t1)//1_000
    speedup  = 212614 / total_time
    print("benchmark 1:", (t2-t1)//1_000, "us (x" + str(speedup)[:4] + ")")

    for i in range(10):
        benchmark_2()
    t1 = time.now()
    benchmark_2()
    t2 = time.now()
    total_time = (t2-t1)//1_000
    speedup  = 69444  / total_time
    print("benchmark 2:", (t2-t1)//1_000, "us (x" + str(speedup)[:4] + ")")

    json_looking_string = get_json_looking_string()
    for i in range(10):
        benchmark_3(json_looking_string)
    t1 = time.now()
    benchmark_3(json_looking_string)
    t2 = time.now()
    total_time = (t2-t1)//1_000
    speedup  = 36218  / total_time
    print("benchmark 3:", (t2-t1)//1_000, "us (x" + str(speedup)[:4] + ")")

    list_of_ints = get_dummy_int_list(10000)
    for i in range(10):
        benchmark_4(list_of_ints)
    t1 = time.now()
    benchmark_4(list_of_ints)
    t2 = time.now()
    total_time = (t2-t1)//1_000
    speedup  = 191769  / total_time
    print("benchmark 4:", (t2-t1)//1_000, "us (x" + str(speedup)[:4] + ")")

[gabrieldemarmiesse]

Hello there, I haven't had much feedback on this issue. I had some feedback on the implementation details in #2632 but no confirmation that this is the direction we want to take. There is also the materialization bug at play #2637 .

I'd like to ask a couple question to unblock the situation and avoid doing work which will be thrown away:

• Do we have confirmation that [BUG] Incorrect pointer behavior when materializing a type #2637 is really a compiler bug and will be fixed and is not a misuse of the language?
• Looking at the benchmarks do we have confirmation from the maintainers that we'll go with the 4th option (Using SBO in List directly) as this is the one that has the largest speedups?
• Do we want to have SBO optionally present in List? There wasn't a clear "Yes" from the maintainers.

If the answer to all three questions is "Yes" then the plan will be:

1. Create a new PR for the SBO with all the comments of [stdlib] Implement SSO using small buffer optimization on List #2632 addressed and get it merged. At the same time, the compiler devs will work to fix [BUG] Incorrect pointer behavior when materializing a type #2637.
2. Make a pull request where I add a parameter to the String struct, something like sso_size. It will default to 0 to avoid any disruption and users will be able to play with it.
3. After enough time has passed and we are confident all bugs are flushed out, we choose a new default sso_size which will be non-null.

This will open the door for those future works:

• Decide how we handle the null terminator (two different strings, à la rust? All strings are null terminated à la C++? A parameter on the struct? etc...)
• Decide how we handle UTF-8 and the support for fast indexing, slicing, len().
• All kind of complicated optimizations can be added to the String struct after the previous two steps are done.

Do I have a go of the maintainers on this plan? @JoeLoser @rparolin @ConnorGray

[gabrieldemarmiesse]

I went forward with the plan and here are the PRs, to review in this order:

• [stdlib] Add optional small buffer optimization to List, take 2 #2825
• [stdlib] Work around the materialization bug present in #2825 #2826
• [stdlib] Add SSO to String by using SBO in List with materialization workaround #2827
  The new benchmarks are in [stdlib] Add SSO to String by using SBO in List with materialization workaround #2827 's description