Commit 5bb5aaf932b8ed30aebfbb0036e1532abfc6af46 | zig

author	Matthew Lugg <mlugg@mlugg.co.uk> 2025-06-12 09:56:37+01:00
committer	Matthew Lugg <mlugg@mlugg.co.uk> 2025-06-12 17:51:31+01:00
log	5bb5aaf932b8ed30aebfbb0036e1532abfc6af46
tree	059ef1f8ab00e70585895068e5561681d49514f5
parent	f9a670d46de3c62be16202f186eacfee6ec096d4
signature	Commit is signed but in an unrecognized format.

compiler: don't queue too much AIR/MIR

Without this cap, unlucky scheduling and/or details of what pipeline
stages perform best on the host machine could cause many gigabytes of
MIR to be stuck in the queue. At a certain point, pause the main thread
until some of the functions in flight have been processed.

3 files changed, 55 insertions(+), 0 deletions(-)

src
- Compilation.zig
- link.zig
- link
  - Queue.zig

src/Compilation.zig+6

 @@ -4607,12 +4607,17 @@ fn processOneJob(tid: usize, comp: *Compilation, job: Job) JobError!void {
             };
             assert(zcu.pending_codegen_jobs.rmw(.Add, 1, .monotonic) > 0); // the "Code Generation" node hasn't been ended
             zcu.codegen_prog_node.increaseEstimatedTotalItems(1);
+            // This value is used as a heuristic to avoid queueing too much AIR/MIR at once (hence
+            // using a lot of memory). If this would cause too many AIR bytes to be in-flight, we
+            // will block on the `dispatchZcuLinkTask` call below.
+            const air_bytes: u32 = @intCast(air.instructions.len * 5 + air.extra.items.len * 4);
             if (comp.separateCodegenThreadOk()) {
                 // `workerZcuCodegen` takes ownership of `air`.
                 comp.thread_pool.spawnWgId(&comp.link_task_wait_group, workerZcuCodegen, .{ comp, func.func, air, shared_mir });
                 comp.dispatchZcuLinkTask(tid, .{ .link_func = .{
                     .func = func.func,
                     .mir = shared_mir,
+                    .air_bytes = air_bytes,
                 } });
             } else {
                 {
 @@ -4624,6 +4629,7 @@ fn processOneJob(tid: usize, comp: *Compilation, job: Job) JobError!void {
                 comp.dispatchZcuLinkTask(tid, .{ .link_func = .{
                     .func = func.func,
                     .mir = shared_mir,
+                    .air_bytes = air_bytes,
                 } });
                 air.deinit(gpa);
             }

src/link.zig+5

 @@ -1267,6 +1267,11 @@ pub const ZcuTask = union(enum) {
         /// the codegen job to ensure that the linker receives functions in a deterministic order,
         /// allowing reproducible builds.
         mir: *SharedMir,
+        /// This is not actually used by `doZcuTask`. Instead, `Queue` uses this value as a heuristic
+        /// to avoid queueing too much AIR/MIR for codegen/link at a time. Essentially, we cap the
+        /// total number of AIR bytes which are being processed at once, preventing unbounded memory
+        /// usage when AIR is produced faster than it is processed.
+        air_bytes: u32,
         pub const SharedMir = struct {
             /// This is initially `.pending`. When `value` is populated, the codegen thread will set

src/link/Queue.zig+44

...	@@ -39,6 +39,21 @@ wip_zcu: std.ArrayListUnmanaged(ZcuTask),	...	@@ -39,6 +39,21 @@ wip_zcu: std.ArrayListUnmanaged(ZcuTask),
39	/// index into `wip_zcu` which we have reached.	39	/// index into `wip_zcu` which we have reached.
40	wip_zcu_idx: usize,	40	wip_zcu_idx: usize,
41		41
		42	/// The sum of all `air_bytes` for all currently-queued `ZcuTask.link_func` tasks. Because
		43	/// MIR bytes are approximately proportional to AIR bytes, this acts to limit the amount of
		44	/// AIR and MIR which is queued for codegen and link respectively, to prevent excessive
		45	/// memory usage if analysis produces AIR faster than it can be processed by codegen/link.
		46	/// The cap is `max_air_bytes_in_flight`.
		47	/// Guarded by `mutex`.
		48	air_bytes_in_flight: u32,
		49	/// If nonzero, then a call to `enqueueZcu` is blocked waiting to add a `link_func` task, but
		50	/// cannot until `air_bytes_in_flight` is no greater than this value.
		51	/// Guarded by `mutex`.
		52	air_bytes_waiting: u32,
		53	/// After setting `air_bytes_waiting`, `enqueueZcu` will wait on this condition (with `mutex`).
		54	/// When `air_bytes_waiting` many bytes can be queued, this condition should be signaled.
		55	air_bytes_cond: std.Thread.Condition,
		56
42	/// Guarded by `mutex`.	57	/// Guarded by `mutex`.
43	state: union(enum) {	58	state: union(enum) {
44	/// The link thread is currently running or queued to run.	59	/// The link thread is currently running or queued to run.
...	@@ -52,6 +67,11 @@ state: union(enum) {	...	@@ -52,6 +67,11 @@ state: union(enum) {
52	wait_for_mir: *ZcuTask.LinkFunc.SharedMir,	67	wait_for_mir: *ZcuTask.LinkFunc.SharedMir,
53	},	68	},
54		69
		70	/// In the worst observed case, MIR is around 50 times as large as AIR. More typically, the ratio is
		71	/// around 20. Going by that 50x multiplier, and assuming we want to consume no more than 500 MiB of
		72	/// memory on AIR/MIR, we see a limit of around 10 MiB of AIR in-flight.
		73	const max_air_bytes_in_flight = 10 * 1024 * 1024;
		74
55	/// The initial `Queue` state, containing no tasks, expecting no prelink tasks, and with no running worker thread.	75	/// The initial `Queue` state, containing no tasks, expecting no prelink tasks, and with no running worker thread.
56	/// The `pending_prelink_tasks` and `queued_prelink` fields may be modified as needed before calling `start`.	76	/// The `pending_prelink_tasks` and `queued_prelink` fields may be modified as needed before calling `start`.
57	pub const empty: Queue = .{	77	pub const empty: Queue = .{
...	@@ -64,6 +84,9 @@ pub const empty: Queue = .{	...	@@ -64,6 +84,9 @@ pub const empty: Queue = .{
64	.wip_zcu = .empty,	84	.wip_zcu = .empty,
65	.wip_zcu_idx = 0,	85	.wip_zcu_idx = 0,
66	.state = .finished,	86	.state = .finished,
		87	.air_bytes_in_flight = 0,
		88	.air_bytes_waiting = 0,
		89	.air_bytes_cond = .{},
67	};	90	};
68	/// `lf` is needed to correctly deinit any pending `ZcuTask`s.	91	/// `lf` is needed to correctly deinit any pending `ZcuTask`s.
69	pub fn deinit(q: Queue, comp: Compilation) void {	92	pub fn deinit(q: Queue, comp: Compilation) void {
...	@@ -131,6 +154,16 @@ pub fn enqueueZcu(q: Queue, comp: Compilation, task: ZcuTask) Allocator.Error!	...	@@ -131,6 +154,16 @@ pub fn enqueueZcu(q: Queue, comp: Compilation, task: ZcuTask) Allocator.Error!
131	{	154	{
132	q.mutex.lock();	155	q.mutex.lock();
133	defer q.mutex.unlock();	156	defer q.mutex.unlock();
		157	// If this is a `link_func` task, we might need to wait for `air_bytes_in_flight` to fall.
		158	if (task == .link_func) {
		159	const max_in_flight = max_air_bytes_in_flight -\| task.link_func.air_bytes;
		160	while (q.air_bytes_in_flight > max_in_flight) {
		161	q.air_bytes_waiting = task.link_func.air_bytes;
		162	q.air_bytes_cond.wait(&q.mutex);
		163	q.air_bytes_waiting = 0;
		164	}
		165	q.air_bytes_in_flight += task.link_func.air_bytes;
		166	}
134	try q.queued_zcu.append(comp.gpa, task);	167	try q.queued_zcu.append(comp.gpa, task);
135	switch (q.state) {	168	switch (q.state) {
136	.running, .wait_for_mir => return,	169	.running, .wait_for_mir => return,
...	@@ -221,6 +254,17 @@ fn flushTaskQueue(tid: usize, q: Queue, comp: Compilation) void {	...	@@ -221,6 +254,17 @@ fn flushTaskQueue(tid: usize, q: Queue, comp: Compilation) void {
221	}	254	}
222	link.doZcuTask(comp, tid, task);	255	link.doZcuTask(comp, tid, task);
223	task.deinit(comp.zcu.?);	256	task.deinit(comp.zcu.?);
		257	if (task == .link_func) {
		258	// Decrease `air_bytes_in_flight`, since we've finished processing this MIR.
		259	q.mutex.lock();
		260	defer q.mutex.unlock();
		261	q.air_bytes_in_flight -= task.link_func.air_bytes;
		262	if (q.air_bytes_waiting != 0 and
		263	q.air_bytes_in_flight <= max_air_bytes_in_flight -\| q.air_bytes_waiting)
		264	{
		265	q.air_bytes_cond.signal();
		266	}
		267	}
224	q.wip_zcu_idx += 1;	268	q.wip_zcu_idx += 1;
225	}	269	}
226	}	270	}