Precondition Based Optimization (Library Design Proposal)

Mojo Performance
1

Many functions in the standard library have to deal with many different cases and likely have many code branches due to this. A possible optimization is to allow users to create parameterized functions such that the programmer can allow the function to make some assumption in order to execute faster.

The most obvious example of this is to allow the list’s append function to have an assume_capacity parameter:

  fn append[assume_capacity: Bool = False](mut self, item: T):
      if assume_capacity:
          self.unsafe_set(self._len, item)
          self._len += 1
      else:
          self.append(item)

But the optimizations don’t have to stop there. Mojo strings have to support utf-8 encoding, but by adding an assume_ascii parameter, we can make many functions faster (.upper(), .lower(), .split(), and pretty much any other function that needs to iterate over the entire string). In the upper and lower cases, removing branches required by utf-8 support also allows us to reduce the entire function to 3 SIMD operations. Since mojo’s string type also supports cow, and sso, we can add additional parameters such as assume_static, assume_heap, assume_inline.

These performance improvements are not just hypothetical either: I implemented a very similar string type in rust, and allowed the users to pass preconditions to certain functions: sso_string_rs/src/lib.rs at master · akneni/sso_string_rs · GitHub. These preconditions often resulted in 10-20% performance improvements.

While this works in rust, mojo is already better prepared to support these design choices. In order to be able to pass an enum as a generic parameter in rust, I needed to convert it to an integer and back since regular enums can’t be used as generic parameters. Additionally, there is no way to set defaults for generic parameters, so I needed to split the functions into two, one for the normal operation and another where the user could pass assumptions (though splitting into 2 may be ideal since the assumption based function should be marked as unsafe in idiomatic rust).

4 Likes
2

Hi @akneni,
I think that unsafe but more performant functions in the standard lib have a unsafe_ prefix. See List.unsafe_get() and List.unsafe_set().

So a simple solution would be to propose a List.unsafe_append() function. This better indicates to the caller and reader of the code that this function is unsafe vs. list.append[true](item) or list.append[assume_capacity=true](item).

But for other use cases your proposal might be a good idea.

2 Likes
3

Agreed, we do need some way of marking these as unsafe. I’m a little opposed to the idea of having two versions of every function. Maybe we could just name the parameters with unsafe_? If this gets enough support here, I’ll create a proper proposal on GitHub.

1 Like
4

This is a really interesting idea, and I like how it would lead to “defragmenting” and reducing our API surface by adding unsafe knobs to existing algorithms, rather than adding new ones.

Perhaps a way to split the difference is to add the unsafe word to the parameter, e.g.:

fn append[unsafe_assume_capacity: Bool = False](mut self, item: T):

I wonder what the stdlib folks like @joe think about this?

2 Likes
5

I like this idea. We already have some level of that with the batch_size parameter on InlineArray, and I think more tunables that can be used by kernel authors or eventually be handed over to autotune is a good thing.

+1 on adding unsafe where relevant.

I think this would be especially helpful in String processing, since it would allow a developer to do a scan over, for instance, a JSON string and verify it’s ASCII or Latin1, then enter a “Latin1 fast path”.

2 Likes
6

I’ve posted a more formal proposal on github. I’ve also changed the name of this pattern to “Precondition Assisted Execution” as that seems to be more descriptive (optimization makes me think compiler optimization)

2 Likes