From c14a630041a196c6af4959c789d90178e0fe7ae1 Mon Sep 17 00:00:00 2001 From: lg Date: Mon, 31 Aug 2026 11:54:22 +0200 Subject: [PATCH] Sema: allow comptime call memoization when callee uses `@setEvalBranchQuota` (#36494) Previously, memoization was disabled as soon as setEvalBranchQuota was called. This commit keeps memoization enabled when quota is being set by keeping track of the requested quota in the call itself and externally. This has shown significant improvements when calling functions where the callee calls `@setEvalBranchQuota` For example: ```zig fn fib(n: u64) u64 { @setEvalBranchQuota(1_000_000_000); // before, this disabled memoization if (n < 2) return n; return fib(n - 1) + fib(n - 2); } pub export fn entry() u64 { var acc: u64 = 0; inline for (0..40) |_| { acc +%= comptime fib(25); } return acc; } ``` This went from ~30s to near instant on my machine, admittedly this is the perfect case for memoization though. A more "real-world" case like this: ```zig pub export fn entry() usize { @setEvalBranchQuota(10_000_000); var n: usize = 0; inline for (0..200) |_| { // comptimePrint calls @setEvalBranchQuota! const s = std.fmt.comptimePrint("value {d} of {s} at {x}", .{ 42, "abcdefgh", 7 }); n += s.len; } return n; } ``` Showed around ~2.2x improvement. Reviewed-on: https://codeberg.org/ziglang/zig/pulls/36494 Reviewed-by: mlugg --- src/InternPool.zig | 4 ++++ src/Sema.zig | 14 +++++++++++++- 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/src/InternPool.zig b/src/InternPool.zig index b477ff3b75eb235cab65bbef70de40bf7c01cad6..fd475c6391cab19462e5240b9cc70215b1ac3484 100644 --- a/src/InternPool.zig +++ b/src/InternPool.zig @@ -2596,6 +2596,7 @@ pub const Key = union(enum) { arg_values: []const Index, result: Index, branch_count: u32, + branch_quota: u32, }; pub fn hash32(key: Key, ip: *const InternPool) u32 { @@ -6272,6 +6273,7 @@ pub const MemoizedCall = struct { args_len: u32, result: Index, branch_count: u32, + branch_quota: u32, }; pub fn init(ip: *InternPool, gpa: Allocator, io: Io, available_threads: usize) !void { @@ -6913,6 +6915,7 @@ pub fn indexToKey(ip: *const InternPool, index: Index) Key { .arg_values = @ptrCast(extra_list.view().items(.@"0")[extra.end..][0..extra.data.args_len]), .result = extra.data.result, .branch_count = extra.data.branch_count, + .branch_quota = extra.data.branch_quota, } }; }, }; @@ -7992,6 +7995,7 @@ pub fn get(ip: *InternPool, gpa: Allocator, io: Io, tid: Zcu.PerThread.Id, key: .args_len = @intCast(memoized_call.arg_values.len), .result = memoized_call.result, .branch_count = memoized_call.branch_count, + .branch_quota = memoized_call.branch_quota, }), }); extra.appendSliceAssumeCapacity(.{@ptrCast(memoized_call.arg_values)}); diff --git a/src/Sema.zig b/src/Sema.zig index d828999ddf32555ef39dd3fe936d88389fa8e9ea..7eda20bf930243883aa8f87d0b454da0b74eb6c6 100644 --- a/src/Sema.zig +++ b/src/Sema.zig @@ -132,6 +132,10 @@ dependencies: std.array_hash_map.Auto(InternPool.Dependee, void) = .empty, /// by `analyzeCall`. allow_memoize: bool = true, +/// The largest quota requested by `@setEvalBranchQuota` within the comptime call +/// currently being analyzed. +quota_request: u32 = 0, + /// The `BranchHint` for the current branch of runtime control flow. /// This state is on `Sema` so that `cold` hints can be propagated up through blocks with less special handling. branch_hint: ?std.lang.BranchHint = null, @@ -4914,7 +4918,7 @@ fn zirSetEvalBranchQuota(sema: *Sema, block: *Block, inst: Zir.Inst.Index) Compi const src = block.nodeOffset(inst_data.src_node); const quota: u32 = @intCast(try sema.resolveInt(block, src, inst_data.operand, .u32, .{ .simple = .operand_setEvalBranchQuota })); sema.branch_quota = @max(sema.branch_quota, quota); - sema.allow_memoize = false; + sema.quota_request = @max(sema.quota_request, quota); } fn zirStoreNode(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!void { @@ -7218,6 +7222,7 @@ fn analyzeCall( .arg_values = memoized_arg_values, .result = undefined, // ignored by hash+eql .branch_count = undefined, // ignored by hash+eql + .branch_quota = undefined, // ignored by hash+eql }, }) orelse break :memoize; const memoized_call = ip.indexToKey(memoized_call_index).memoized_call; @@ -7227,6 +7232,8 @@ fn analyzeCall( break :memoize; } sema.branch_count += memoized_call.branch_count; + sema.branch_quota = @max(sema.branch_quota, memoized_call.branch_quota); + sema.quota_request = @max(sema.quota_request, memoized_call.branch_quota); const result = Air.internedToRef(memoized_call.result); if (ensure_result_used) { try sema.ensureResultUsed(block, sema.typeOf(result), call_src); @@ -7353,6 +7360,10 @@ fn analyzeCall( defer sema.allow_memoize = old_allow_memoize and sema.allow_memoize; sema.allow_memoize = true; + const old_quota_request = sema.quota_request; + defer sema.quota_request = @max(old_quota_request, sema.quota_request); + sema.quota_request = 0; + // Store the current eval branch count so we can find out how many eval branches // the comptime call caused. const old_branch_count = sema.branch_count; @@ -7388,6 +7399,7 @@ fn analyzeCall( .arg_values = memoized_arg_values, .result = result_val.toIntern(), .branch_count = sema.branch_count - old_branch_count, + .branch_quota = sema.quota_request, } }); } } -- 2.54.0