Zig optimizer (my bachelor thesis)

Showcase
, ,
1

Hi :smiley:

As my bachelor thesis I decided to write a patch into the Zig compiler that introduces a new compilation phase dedicated to optimizations. It uses a Sea of Nodes representation adapted for the DoD style used in the compiler. https://codeberg.org/prokop/zig/src/branch/optimizer/src/opt. It was a lot of fun to write! I’m actually pretty proud of the way I store the program in memory (it went through many changes and i’m still not finished).

It is remarkable how unimpressive the code transformations are, given how much time i sank into it xD So instead I would like to highlight from the thesis how a simple optimization pass implementation looks like. Most of the code I wrote so far is support code that only allows writing efficient optimization passes and I didn’t have time left to actually implement many passes.

pub fn monotone(gpa: Allocator, oil: *Oil) error{OutOfMemory}!void {
	var queue: std.Deque(Index) = .empty;
	var qset: std.AutoHashMapUnmanaged(Index, void) = .empty;

	defer queue.deinit(gpa);
	defer qset.deinit(gpa);

	const len = oil.insts.len;

	try queue.ensureTotalCapacity(gpa, len);
	try qset.ensureTotalCapacity(gpa, len);

	for (0..len) |i| {
		const index: Index = @enumFromInt(i);
		queue.pushBackAssumeCapacity(index);
		qset.putAssumeCapacity(index, {});
	}

	// compute def->use chains
	const du, const du_buffer = try defuse(gpa, oil);
	defer gpa.free(du_buffer);
	defer gpa.free(du);

	while (queue.popBack()) |index| {
		assert(qset.remove(index));
		const inst = index.inst(oil);
		const oty = index.ty(oil);
		const new_oty = try compute.compute(
			gpa,
			oil,
			inst.tag,
			inst.data,
			oty,
		);

		if (oty.eq(new_oty)) continue;

		index.setType(oil, new_oty);

		// enqueue all uses
		for (du[@intFromEnum(index)]) |use| {
			if (qset.getOrPutAssumeCapacity(use).found_existing) continue;
			queue.pushBackAssumeCapacity(use);
		}
	}
}

This pass is able to propagate constants through non-trivial programs such as:

fn fun(x: u8) u8 {
	const a: u8 = 200;
	const b: u8 = if (x > 5) 100 else 100;
	const c: u8 = if (x > 10)
		a + b
	else
		a - b;
	return c / 2;
}

:down_arrow:

%0 : ⊥ ctrl = startcontrol
%1 : ⊥ mem = startmemory
%2 : ⊥ u8 = arg 0
%3 : ⊥ ctrl = region [%10, %9]
%4 : ⊥ mem = phi %3, [%1, %1]
%6 : 5 u8 = const
%7 : ⊥ u8 = cmp_lt %6, %2
%9 : ⊥ ctrl, %10 : ⊥ ctrl = if %0, %7
%14 : ⊥ ctrl = region [%21, %20]
%15 : ⊥ mem = phi %14, [%4, %4]
%17 : 10 u8 = const
%18 : ⊥ u8 = cmp_lt %17, %2
%20 : ⊥ ctrl, %21 : ⊥ ctrl = if %3, %18
%27 : 50 u8 = const
%28 : ⊤ ctrl = ret %14, %15, %27

It is unable to remove the if (yet) but it is able to infer that the function always returns 50 (the last two lines). It must always return 50 because it would be UB to execute the true branch of the if c:<

I didn’t have time to implement support for debug mode so the only compilation mode currently implemented is ReleaseFast.

The speed is pretty good. I measured that the speed of the optimization passes (excluding lowering) to be ~0.4 μs/inst (or ~751k lines of source code per second on a single core, if all you do are the optimization passes), which is alright for the first version but tbh I think it should be way faster! (Additionally my lowering is very slow, I have to look into that).

The bachelor thesis is available here. I hope now that I submitted it you all find a bunch of mistakes that will haunt me forever :smiley: I was lucky enough to have Jan Hubička as my supervisor, who has been maintaining GCC for the last 20 years or so, his input is very helpful.

I don’t have nearly enough street cred to get this merged any time soon (it would be +5k PR) but i’m 100% going to extend this into a masters thesis and optimize more code :star_struck:

AI / LLM usage disclosure

I tried to use some llm to generate some code a couple of months ago. I was meh so I didn’t continue but some of the code is still there somewhere probably.

20 Likes