authorgravatar for abbix@riseup.netlg <abbix@riseup.net> 2026-08-31 11:54:22+02:00
committergravatar for mlugg@mlugg.co.ukMatthew Lugg <mlugg@mlugg.co.uk> 2026-08-31 11:54:22+02:00
logc14a630041a196c6af4959c789d90178e0fe7ae1
tree00932ac3c10efa411d3d76bc8556ddb13ceccbc0
parentfd3048c5a91220102420ef51ba5193d6febca521

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 <mlugg@mlugg.co.uk>

2 files changed, 17 insertions(+), 1 deletions(-)

src/InternPool.zig+4
...@@ -2596,6 +2596,7 @@ pub const Key = union(enum) {...@@ -2596,6 +2596,7 @@ pub const Key = union(enum) {
2596 arg_values: []const Index,2596 arg_values: []const Index,
2597 result: Index,2597 result: Index,
2598 branch_count: u32,2598 branch_count: u32,
2599 branch_quota: u32,
2599 };2600 };
26002601
2601 pub fn hash32(key: Key, ip: *const InternPool) u32 {2602 pub fn hash32(key: Key, ip: *const InternPool) u32 {
...@@ -6272,6 +6273,7 @@ pub const MemoizedCall = struct {...@@ -6272,6 +6273,7 @@ pub const MemoizedCall = struct {
6272 args_len: u32,6273 args_len: u32,
6273 result: Index,6274 result: Index,
6274 branch_count: u32,6275 branch_count: u32,
6276 branch_quota: u32,
6275};6277};
62766278
6277pub fn init(ip: *InternPool, gpa: Allocator, io: Io, available_threads: usize) !void {6279pub 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 {...@@ -6913,6 +6915,7 @@ pub fn indexToKey(ip: *const InternPool, index: Index) Key {
6913 .arg_values = @ptrCast(extra_list.view().items(.@"0")[extra.end..][0..extra.data.args_len]),6915 .arg_values = @ptrCast(extra_list.view().items(.@"0")[extra.end..][0..extra.data.args_len]),
6914 .result = extra.data.result,6916 .result = extra.data.result,
6915 .branch_count = extra.data.branch_count,6917 .branch_count = extra.data.branch_count,
6918 .branch_quota = extra.data.branch_quota,
6916 } };6919 } };
6917 },6920 },
6918 };6921 };
...@@ -7992,6 +7995,7 @@ pub fn get(ip: *InternPool, gpa: Allocator, io: Io, tid: Zcu.PerThread.Id, key:...@@ -7992,6 +7995,7 @@ pub fn get(ip: *InternPool, gpa: Allocator, io: Io, tid: Zcu.PerThread.Id, key:
7992 .args_len = @intCast(memoized_call.arg_values.len),7995 .args_len = @intCast(memoized_call.arg_values.len),
7993 .result = memoized_call.result,7996 .result = memoized_call.result,
7994 .branch_count = memoized_call.branch_count,7997 .branch_count = memoized_call.branch_count,
7998 .branch_quota = memoized_call.branch_quota,
7995 }),7999 }),
7996 });8000 });
7997 extra.appendSliceAssumeCapacity(.{@ptrCast(memoized_call.arg_values)});8001 extra.appendSliceAssumeCapacity(.{@ptrCast(memoized_call.arg_values)});
src/Sema.zig+13-1
...@@ -132,6 +132,10 @@ dependencies: std.array_hash_map.Auto(InternPool.Dependee, void) = .empty,...@@ -132,6 +132,10 @@ dependencies: std.array_hash_map.Auto(InternPool.Dependee, void) = .empty,
132/// by `analyzeCall`.132/// by `analyzeCall`.
133allow_memoize: bool = true,133allow_memoize: bool = true,
134134
135/// The largest quota requested by `@setEvalBranchQuota` within the comptime call
136/// currently being analyzed.
137quota_request: u32 = 0,
138
135/// The `BranchHint` for the current branch of runtime control flow.139/// The `BranchHint` for the current branch of runtime control flow.
136/// This state is on `Sema` so that `cold` hints can be propagated up through blocks with less special handling.140/// This state is on `Sema` so that `cold` hints can be propagated up through blocks with less special handling.
137branch_hint: ?std.lang.BranchHint = null,141branch_hint: ?std.lang.BranchHint = null,
...@@ -4914,7 +4918,7 @@ fn zirSetEvalBranchQuota(sema: *Sema, block: *Block, inst: Zir.Inst.Index) Compi...@@ -4914,7 +4918,7 @@ fn zirSetEvalBranchQuota(sema: *Sema, block: *Block, inst: Zir.Inst.Index) Compi
4914 const src = block.nodeOffset(inst_data.src_node);4918 const src = block.nodeOffset(inst_data.src_node);
4915 const quota: u32 = @intCast(try sema.resolveInt(block, src, inst_data.operand, .u32, .{ .simple = .operand_setEvalBranchQuota }));4919 const quota: u32 = @intCast(try sema.resolveInt(block, src, inst_data.operand, .u32, .{ .simple = .operand_setEvalBranchQuota }));
4916 sema.branch_quota = @max(sema.branch_quota, quota);4920 sema.branch_quota = @max(sema.branch_quota, quota);
4917 sema.allow_memoize = false;4921 sema.quota_request = @max(sema.quota_request, quota);
4918}4922}
49194923
4920fn zirStoreNode(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!void {4924fn zirStoreNode(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!void {
...@@ -7218,6 +7222,7 @@ fn analyzeCall(...@@ -7218,6 +7222,7 @@ fn analyzeCall(
7218 .arg_values = memoized_arg_values,7222 .arg_values = memoized_arg_values,
7219 .result = undefined, // ignored by hash+eql7223 .result = undefined, // ignored by hash+eql
7220 .branch_count = undefined, // ignored by hash+eql7224 .branch_count = undefined, // ignored by hash+eql
7225 .branch_quota = undefined, // ignored by hash+eql
7221 },7226 },
7222 }) orelse break :memoize;7227 }) orelse break :memoize;
7223 const memoized_call = ip.indexToKey(memoized_call_index).memoized_call;7228 const memoized_call = ip.indexToKey(memoized_call_index).memoized_call;
...@@ -7227,6 +7232,8 @@ fn analyzeCall(...@@ -7227,6 +7232,8 @@ fn analyzeCall(
7227 break :memoize;7232 break :memoize;
7228 }7233 }
7229 sema.branch_count += memoized_call.branch_count;7234 sema.branch_count += memoized_call.branch_count;
7235 sema.branch_quota = @max(sema.branch_quota, memoized_call.branch_quota);
7236 sema.quota_request = @max(sema.quota_request, memoized_call.branch_quota);
7230 const result = Air.internedToRef(memoized_call.result);7237 const result = Air.internedToRef(memoized_call.result);
7231 if (ensure_result_used) {7238 if (ensure_result_used) {
7232 try sema.ensureResultUsed(block, sema.typeOf(result), call_src);7239 try sema.ensureResultUsed(block, sema.typeOf(result), call_src);
...@@ -7353,6 +7360,10 @@ fn analyzeCall(...@@ -7353,6 +7360,10 @@ fn analyzeCall(
7353 defer sema.allow_memoize = old_allow_memoize and sema.allow_memoize;7360 defer sema.allow_memoize = old_allow_memoize and sema.allow_memoize;
7354 sema.allow_memoize = true;7361 sema.allow_memoize = true;
73557362
7363 const old_quota_request = sema.quota_request;
7364 defer sema.quota_request = @max(old_quota_request, sema.quota_request);
7365 sema.quota_request = 0;
7366
7356 // Store the current eval branch count so we can find out how many eval branches7367 // Store the current eval branch count so we can find out how many eval branches
7357 // the comptime call caused.7368 // the comptime call caused.
7358 const old_branch_count = sema.branch_count;7369 const old_branch_count = sema.branch_count;
...@@ -7388,6 +7399,7 @@ fn analyzeCall(...@@ -7388,6 +7399,7 @@ fn analyzeCall(
7388 .arg_values = memoized_arg_values,7399 .arg_values = memoized_arg_values,
7389 .result = result_val.toIntern(),7400 .result = result_val.toIntern(),
7390 .branch_count = sema.branch_count - old_branch_count,7401 .branch_count = sema.branch_count - old_branch_count,
7402 .branch_quota = sema.quota_request,
7391 } });7403 } });
7392 }7404 }
7393 }7405 }