Auto-optimizing an algorithm

After I found that llms are useful, I wanted to get a better feel for what they are useful for. Clearly, they are good at generating boilerplate and writing bash one-liners, but what else?

I wanted to see how difficult it was to set up an agent loop that has some objective that it can evaluate autonomously, let it run and see what happens. Also, I wanted some idea what was happening in the process, but without human intervention. Code optimization is an example of something with a lot of interesting properties: it's easy to test for correctness, and has a (fairly) clear objective function¹.

So which code should we optimize? I figured I'd start at a blank slate, choose a problem that could have some different good solutions, that is easy to manually test or verify (if needed), and that I sort-of know from before.

After some brainstorming with the machine I chose the Levenshtein edit distance. Given two strings a and b, how many edits do you have to make to turn a into b? One edit is a single insertion, a single removal, or replacing a single character, and we want to find the minimal number of edits. We don't care what the edits actually are.

The Reference Solution

I guess you could brute force this and go through a search three trying all combinations, but that would be very inefficient. The standard "proper"² solution uses dynamic programming: we create a table dist[i][j] that computes the edit-distance between a[..i] and b[..j], and build this table bottom-up.

dist[0][k] == dist[k][0] == k because to go from the empty string to any other string you have to insert each character. This is the base-case.

For the other entries there are three cases:

1. remove a character: cost dist[i-1][j] + 1,
2. add a character: cost dist[i][j-1] + 1,
3. replace: cost dist[i-1][j-1], plus 1 if they are different and 0 otherwise.

Confused yet? Here's a half-filled out table:

Finding ? in the table is computing edit("pa", "ag"). We have some options:

1. Pay 1 to remove the 'a' in "pa" and then do the rest of the edits: edit("p", "ag") + 1.
2. Do the edit from "pa" to just "a" and then insert the 'g': edit("pa", "a") + 1
3. Replace the last charater and edit the rest: edit("p", "a") + 1. We get the +1 part since the chars doesn't match.

We only need to consider all three and pick the cheapest:

1. edit("p", "ag") + 1 = dist[1][2] + 1 == 3
2. edit("pa", "a") + 1 = dist[2][1] + 1 == 2
3. edit("p", "a") + 1 = dist[1][1] + 1 == 2

In this case is a tie: for the full edits, we can either go "ag" -> "aa" -> "pa" where we replace both times, or we can go "ag" -> "a" -> "pa" where we first delete and then insert.

I don't know, it's kinda confusing, and very tempting to frame as peeling off characters on both strings until nothing is left, dist[0][0], but then we have to show that this solves the same problem. Computing the table, though, is now super easy. Look at the three neighbors to the top left and add one to the smallest of those. If the letters match at your position, you get to choose the diagonal one for free.

Here's the filled out table with arrows showing a path we can take for the optimal cost of 3:

agent -> agen -> agan -> pagan

Setup

Look, this isn't actual science, and you're not my boss, so I can do what I want. Or rather, I don't have to do things I don't want to do. And I didn't want to write tests or benchmarks, or even think too hard about edge-cases or input distributions. So instead of designing a test suite and a benchmark system I first had agents generate these too.

Also, the agent generated the reference implementation:

// THIS IS PURE LLM GENERATED CODE
int ref_edit_distance(
    const unsigned char *s,
    int s_len, 
    const unsigned char *t,
    int t_len
) {
  if (s_len == 0) return t_len;
  if (t_len == 0) return s_len;

  int *prev = (int *)malloc((t_len + 1) * sizeof(int));
  int *curr = (int *)malloc((t_len + 1) * sizeof(int));

  for (int j = 0; j <= t_len; j++)
    prev[j] = j;

  for (int i = 1; i <= s_len; i++) {
    curr[0] = i;

    for (int j = 1; j <= t_len; j++) {
      int cost = (s[i - 1] == t[j - 1]) ? 0 : 1;

      int del = prev[j] + 1;
      int ins = curr[j - 1] + 1;
      int sub = prev[j - 1] + cost;

      int d = del;
      if (ins < d) d = ins;
      if (sub < d) d = sub;
      curr[j] = d;
    }

    int *tmp = prev;
    prev = curr;
    curr = tmp;
  }

  int r = prev[t_len];
  free(prev);
  free(curr);
  return r;
}

It has one funny trick, namely that instead of computing the whole table it computes one row at a time and switches which of the two rows it looks at and which is writes to. In iteration i, prev is row i-1 and curr is row i.

Setting up the harness was a little tricy because I tried to over-engineer it at first, with sandboxing and custom tools, trying to limit the exposure the agent had to the benchmarks and previous versions. In the end I used pi and wrote separate agents for separate stages of the process: one planner, one implementor, one reviewer, and the "top" model to orchestrate. A pretty standard setup, except that the implementor wasn't allowed to run the benchmarks.

The setup was such that the planner wrote an hypothesis, the implementor created a submission that passed the test suite, the orchestrator ran the benchmarks that spat out a json file with wall-clock time and perf numbers, and the reviewer looked at those and wrote a conclusion. Then it updated a global FINDINGS.md so that the planner in the next iteration wouldn't have to read too much to make a new plan.

I also made sure that the agents wrote plans to communicate so that I could store these, in case they would be interesting. pi can probably enforce this with an extension or something, but I didn't bother trying to figure out how.

I also found some third party code that looked like it would be a good reference and ran the benchmark on those.

Results

Short version: llms are capabable of optimizing code when given a harness like this. Here's some timings. Benchmark names have their size, alphabet size, and mutation rate, in the name. IDs are incremental and #99x are external implementations. It's sorted on the middle column.

Here's a table of three good versions, one of which is external.

Some of these differences are close to the benchmark noise, but the generated code is undeniably faster on certain patterns of inputs compared to the external #991. The flipside is also true.

The code is pretty bad. It's very typical "highly-optimized-and-butt-ugly" code. There's no abstractions pulled out, apart from some functions that are optimized for smaller string lengths and marked to go in certain sections of the elf. In other words, no care was taken into constructing a program that makes sense. But, the llm has no problem explaining what's going on, and with some external help it's very doable to sit down and work through it, just as you would have to do if you wanted to understand any other code that you didn't write.

As it was running, I was looking at the output every now and then and saw the agent run objdump on the target executable to confirm that certain regions got auto-vectorized properly. With the right tools, I think it could have done way better; what if it could easily see how instructions are scheduled on the execution ports to detect backend-pressure? These tools probably do exist, but neither I nor the agent seemed to know them.

What's useful for humans is also often useful for agents. I hope this wave of llms will both lower the bar and raise the leverage to create good tools that humans and agents can use in order to understand the systems we create.

Thanks for reading.

This work is licensed under a Creative Commons Attribution-ShareAlike 4.0 International License