Rework AST memory layout for better memory usage and performance

[andrewrk]

This is a proof-of-concept of switching to a new memory layout for
tokens and AST nodes. The goal is threefold:

• smaller memory footprint
• faster performance for tokenization and parsing
• most importantly, a proof-of-concept that can be also applied to ZIR
  and TZIR to improve the entire compiler pipeline in this way.

I had a few key insights here:

• Underlying premise: using less memory will make things faster, because
  of fewer allocations and better cache utilization. Also using less
  memory is valuable in and of itself.
• Using a Struct-Of-Arrays for tokens and AST nodes, saves the bytes of
  padding between the enum tag (which kind of token is it; which kind
  of AST node is it) and the next fields in the struct. It also improves
  cache coherence, since one can peek ahead in the tokens array without
  having to load the source locations of tokens.
• Token memory can be conserved by only having the tag (1 byte) and byte
  offset (4 bytes) for a total of 5 bytes per token. It is not necessary
  to store the token ending byte offset because one can always re-tokenize
  later, but also most tokens the length can be trivially determined from
  the tag alone, and for ones where it doesn't, string literals for
  example, one must parse the string literal again later anyway in
  astgen, making it free to re-tokenize.
• AST nodes do not actually need to store more than 1 token index because
  one can poke left and right in the tokens array very cheaply.

So far we are left with one big problem though: how can we put AST nodes
into an array, since different AST nodes are different sizes?

This is where my key observation comes in: one can have a hash table for
the extra data for the less common AST nodes! But it gets even better than
that:

I defined this data that is always present for every AST Node:

• tag (1 byte)
  • which AST node is it
• main_token (4 bytes, index into tokens array)
  • the tag determines which token this points to
• struct{lhs: u32, rhs: u32}
  • enough to store 2 indexes to other AST nodes, the tag determines
    how to interpret this data

You can see how a binary operation, such as a * b would fit into this
structure perfectly. A unary operation, such as *a would also fit,
and leave rhs unused. So this is a total of 13 bytes per AST node.
And again, we don't have to pay for the padding to round up to 16 because
we store in struct-of-arrays format.

I made a further observation: the only kind of data AST nodes need to
store other than the main_token is indexes to sub-expressions. That's it.
The only purpose of an AST is to bring a tree structure to a list of tokens.
This observation means all the data that nodes store are only sets of u32
indexes to other nodes. The other tokens can be found later by the compiler,
by poking around in the tokens array, which again is super fast because it
is struct-of-arrays, so you often only need to look at the token tags array,
which is an array of bytes, very cache friendly.

So for nearly every kind of AST node, you can store it in 13 bytes. For the
rarer AST nodes that have 3 or more indexes to other nodes to store, either
the lhs or the rhs will be repurposed to be an index into an extra_data array
which contains the extra AST node indexes. In other words, no hash table needed,
it's just 1 big ArrayList with the extra data for AST Nodes.

Final observation, no need to have a canonical tag for a given AST. For example:
The expression foo(bar) is a function call. Function calls can have any
number of parameters. However in this example, we can encode the function
call into the AST with a tag called FunctionCallOnlyOneParam, and use lhs
for the function expr and rhs for the only parameter expr. Meanwhile if the
code was foo(bar, baz) then the AST node would have to be FunctionCall
with lhs still being the function expr, but rhs being the index into
extra_data. Then because the tag is FunctionCall it means
extra_data[rhs] is the "start" and extra_data[rhs+1] is the "end".
Now the range extra_data[start..end] describes the list of parameters
to the function.

Point being, you only have to pay for the extra bytes if the AST actually
requires it. There's no limit to the number of different AST tag encodings.

Results of parsing only:

• ✔️ 15% fewer cache-misses
• ✔️ 28% fewer total instructions executed
• ✔️ 26% fewer total CPU cycles
• ✔️ 22% faster wall clock time

Checklist before this can actually be merged:

• parser
• render (zig fmt)
• astgen
• translate-c

Also there are a couple common forms of AST that I did not implement yet:

• FnCall with 1 or 0 args
• Builtin call with 0, 1, or 2 args
• block with 0, 1, or 2 statements
• container decl with 0, 1, or 2 fields/decls
• extern fn foo() void; should be a FnProto, not FnDecl

[Vexu]

Could you make the Token and Node Tags snake_case now while you're breaking the API anyways?

[andrewrk]

Yep

[ifreund]

I like the MultiArrayList implementation with a single allocation. It seems the main disadvantage compared to a version with multiple allocations is that shrinking requires copying memory around and the toOwnedSlice() API is a little less useful because of that.

[ifreund]

Suggested change

	self.len = new_len;
	// TODO memset the invalidated items to undefined
	inline for (fields) |field_info, i| {
	const field = @intToEnum(Field, i);
	mem.set(field_info.field_type, self.slice().items(field)[new_len..], undefined);
	}
	self.len = new_len;

[ifreund]

This comment seems to be incorrect, the ptrs array is indeed in field order.

[ifreund]

sizes.indexes is never used, we should probably get rid of it.

[andrewrk]

Heads up to @alexnask I think this is going to be a major breaking change to ZLS that may require quite some effort to update. Sorry about that. The good news, however, is that it should in theory improve perf & memory usage of ZLS.

[trishume]

One thing you might want to consider is writing a quick benchmark of a simple traversal over the whole AST. I would expect this change to make producing the AST faster but it might make traversing slower due to potentially taking misses on 3 cache lines when visiting a node instead of one. Similarly if you implement the two types of function calls you might add more branch mispredictions.

It would also maybe help double-check the ergonomics of traversal. Although maybe you can just wait to port zig fmt to do all of this since presumably it's a tree traversal.

Edit: The extra cache misses may only be an issue if in practice the order you generate the tree differs substantially from the order you traverse it. This may depend on details of your parser, like if you (pre)allocate the + node before parsing either argument you might end up with the same depth-first order you traverse in and thus minimize cache misses, but if you allocate it after or in-between the two children you won't (although it may still be quite similar and thus cache-efficient).

[Avokadoen]

Shouldn't the Allocator arguments in multi_array_list.zig simply be called "allocator" and not "gpa"?

[andrewrk]

Shouldn't the Allocator arguments in multi_array_list.zig simply be called "allocator" and not "gpa"?

I've been experimenting with a new convention: if the parameter is named gpa it indicates the allocator will be used in a way that would be better suited for a general purpose allocator - e.g. one that reclaims memory on free(). It's not a requirement, but it's a hint to callsites that, if they have the choice, a gpa would be preferable to an arena.

[andrewrk]

Nope.
Wanna see a magic trick? Apply this patch:

Wow, thanks for pointing this out. That's a footgun if I've ever seen one. Wait so this has nothing to do with SIMD vectors, instead it has to do with using std.mem.copy on a u7? I do not yet understand why this patch would change anything. The ABI size of u7 is 1 byte, so writing a u7 should clear the extra bit.

[LemonBoy]

Long story short, the optimizations are kicking in and are vectorizing (exceptionally well, I must add) the mem.copy loop by turning the single i7 loads into <4 x i7> ones (and coalescing more than a single read/write when possible, the aarch64 asm accelerates the copy by reading/writing 8 bytes at once).

Now, the problem with LLVM's vector is that they are tight, there's no padding in between the elements, unlike in a array...
You probably understand the problem now, the generated aarch64 code is slicing and dicing the bits in the wrong way.

[andrewrk]

the generated aarch64 code is slicing and dicing the bits in the wrong way.

Ah, I see. So it is an LLVM bug in vectorization, but not with @reduce. Is that your conclusion? Or have we made some horrific language design flaw?

[LemonBoy]

Ah, I see. So it is an LLVM bug in vectorization, but not with @reduce. Is that your conclusion?

Yes, the @reduce thing was a red herring, all it did was shift the code around to hide the problem.
It remains to be seen if the problem is limited to the loop vectorizer or is extended to other regular loads/stores.

[andrewrk]

Benchmark for running zig fmt on the whole std lib:

master branch

• peak memory usage: 84.7 MiB
• zig executable binary size (without -Denable-llvm): 3.37 MiB

[nix-shell:~/dev/zig/build-release]$ perf stat -e cache-misses -e instructions -e cycles ./zig fmt ~/tmp/std

 Performance counter stats for './zig fmt /home/andy/tmp/std':

         2,803,440      cache-misses                                                
     4,592,565,483      instructions              #    2.70  insn per cycle         
     1,702,969,462      cycles                                                      

       0.389872088 seconds time elapsed

       0.370892000 seconds user
       0.018994000 seconds sys

ast-memory-layout branch

• peak memory usage: 68.4 MiB
• zig executable binary size (without -Denable-llvm): 3.17 MiB

[nix-shell:~/Downloads/zig/build-release]$ perf stat -e cache-misses -e instructions -e cycles ./zig fmt ~/tmp/std2

 Performance counter stats for './zig fmt /home/andy/tmp/std2':

         3,771,657      cache-misses                                                
     3,941,606,705      instructions              #    2.62  insn per cycle         
     1,502,290,493      cycles                                                      

       0.345329438 seconds time elapsed

       0.332314000 seconds user
       0.013012000 seconds sys

Deltas:

• ✔️ 19.3% less peak memory usage
• ✔️ 10.8% reduction in binary size
• ❌ 34.5% more cache-misses
• ✔️ 14.2% fewer instructions
• ✔️ 11.8% fewer CPU cycles
• ✔️ 11.4% faster wall clock time

Keep in mind, zig fmt does a lot of extra work that astgen from the main pipeline does not have to do. So I think these are very promising results!

[LemonBoy]

Link to the LLVM patch, keep all your fingers crossed and hopefully it'll make into LLVM 12