From 745518be4ede1cfa525e0740ca998dda95840be6 Mon Sep 17 00:00:00 2001 From: Benoit Giannangeli Date: Tue, 23 Jun 2026 12:56:49 +0200 Subject: [PATCH] fix(jit): Compute complexity scores while generating bytecode --- src/Ast.zig | 98 ++++++++++++------------------------------------- src/Chunk.zig | 98 ------------------------------------------------- src/Codegen.zig | 84 ++++++++++++++++++++++++++++++++++++++++-- src/Jit.zig | 25 +++++-------- src/vm.zig | 24 +++++++++++- 5 files changed, 136 insertions(+), 193 deletions(-) diff --git a/src/Ast.zig b/src/Ast.zig index 6fb7ba76..6413facc 100644 --- a/src/Ast.zig +++ b/src/Ast.zig @@ -618,32 +618,6 @@ pub const Slice = struct { return ctx.result orelse &.{}; } - const UsesFiberContext = struct { - result: bool = false, - - pub fn processNode(self: *UsesFiberContext, _: std.mem.Allocator, ast: Self.Slice, node: Self.Node.Index) (std.mem.Allocator.Error || std.fmt.BufPrintError)!bool { - switch (ast.nodes.items(.tag)[node]) { - .AsyncCall, - .Resolve, - .Resume, - .Yield, - => { - self.result = true; - return true; - }, - else => return false, - } - } - }; - - pub fn usesFiber(self: Self.Slice, allocator: std.mem.Allocator, node: Node.Index) !bool { - var ctx = UsesFiberContext{}; - - try self.walk(allocator, &ctx, node, .breadthFirst); - - return ctx.result; - } - const IsConstantContext = struct { result: ?bool = null, @@ -827,56 +801,26 @@ pub const Slice = struct { return ctx.result orelse false; } - /// Mirrors Chunk.score (even though Chunk.score and Node.score won't be comparable) - /// Is used to compute complexity of a hotspot node (which doesn't have a Chunk available to evaluate) - const ComplexityContext = struct { + /// JIT complexity metadata stored only on function and hotspot candidate components. + pub const JitComplexity = struct { + /// Complexity score computed during codegen to help evaluate if the node is worth JIT compiling. score: usize = 0, - - pub fn processNode( - ctx: *ComplexityContext, - _: std.mem.Allocator, - ast: Self.Slice, - node: Self.Node.Index, - ) (std.mem.Allocator.Error || std.fmt.BufPrintError)!bool { - if (ast.nodes.items(.complexity_score)[node]) |sc| { - ctx.score += sc; - return true; // Don't go deeper we already computed this node score - } - - ctx.score += switch (ast.nodes.items(.tag)[node]) { - .AsyncCall, - .Resolve, - .Resume, - => { // Blacklist because of fiber use - ctx.score = 0; - return true; - }, - .Call, - .DoUntil, - .For, - .ForEach, - .Throw, - .Try, - .While, - => @as(usize, @intCast(1)), - else => @as(usize, @intCast(0)), - } + 1; // At least 1 per node - - return false; - } + /// Parent in the codegen-time JIT complexity tree, used to update ancestor scores after a hotspot is compiled. + parent: ?Node.Index = null, }; - pub fn score(self: Self.Slice, allocator: std.mem.Allocator, node: Node.Index) !usize { - const complexity_score = &self.nodes.items(.complexity_score)[node]; - if (complexity_score.* == null) { - var ctx = ComplexityContext{}; - - try self.walk(allocator, &ctx, node, .breadthFirst); - - complexity_score.* = ctx.score; - } + /// Returns mutable JIT complexity metadata for nodes that can be compiled directly by the JIT. + pub fn jitComplexity(self: Self.Slice, node: Node.Index) ?*JitComplexity { + const tags = self.nodes.items(.tag); + const components = self.nodes.items(.components); - return complexity_score.* orelse 0; + return switch (tags[node]) { + .Function => &components[node].Function.jit, + .For => &components[node].For.jit, + .ForEach => &components[node].ForEach.jit, + .While => &components[node].While.jit, + else => null, + }; } const NamespaceContext = struct { @@ -1944,8 +1888,6 @@ pub const Node = struct { /// How many time it was visited at runtime (used to decide wether its a hotspot that needs to be compiled) count: usize = 0, - /// Complexity score computed once to help evaluate if the node is worth JIT compiling - complexity_score: ?usize = null, /// Node status: blacklisted, queued/generated/compiled by the JIT, compilable jit_status: JitStatus = .compilable, /// Once compiled @@ -2363,6 +2305,8 @@ pub const For = struct { post_loop: []const Node.Index, body: Node.Index, label: ?TokenIndex, + /// JIT complexity metadata for this hotspot candidate. + jit: Slice.JitComplexity = .{}, }; pub const ForEach = struct { @@ -2372,6 +2316,8 @@ pub const ForEach = struct { body: Node.Index, key_omitted: bool, label: ?TokenIndex, + /// JIT complexity metadata for this hotspot candidate. + jit: Slice.JitComplexity = .{}, }; pub const Function = struct { @@ -2391,6 +2337,8 @@ pub const Function = struct { // Should be .FunctionType // Only function without a function_signature is a script function_signature: ?Node.Index, + /// JIT complexity metadata for this function candidate. + jit: Slice.JitComplexity = .{}, upvalue_binding: std.AutoArrayHashMapUnmanaged(u8, bool), @@ -2628,6 +2576,8 @@ pub const WhileDoUntil = struct { condition: Node.Index, body: Node.Index, label: ?TokenIndex, + /// JIT complexity metadata used when this component represents a hotspot candidate. + jit: Slice.JitComplexity = .{}, }; pub const Zdef = struct { diff --git a/src/Chunk.zig b/src/Chunk.zig index a322db6c..95e555ec 100644 --- a/src/Chunk.zig +++ b/src/Chunk.zig @@ -13,10 +13,6 @@ code: std.ArrayList(u32) = .empty, locations: std.ArrayList(Ast.TokenIndex) = .empty, /// List of constants defined in this chunk constants: std.ArrayList(Value) = .empty, -/// Ranges of bytecode skipped by compiled hotspots -compiled_hotspot_ranges: std.ArrayList(InstructionRange) = .empty, -/// Complexity score computed once to help evaluate if the chunk is worth JIT compiling -complexity_score: ?u32 = null, pub fn init(allocator: std.mem.Allocator, ast: Ast.Slice) Self { return Self{ @@ -29,7 +25,6 @@ pub fn deinit(self: *Self) void { self.code.deinit(self.allocator); self.constants.deinit(self.allocator); self.locations.deinit(self.allocator); - self.compiled_hotspot_ranges.deinit(self.allocator); } pub fn write(self: *Self, code: u32, where: Ast.TokenIndex) !void { @@ -45,94 +40,6 @@ pub fn addConstant(self: *Self, vm: ?*VM, value: Value) !u24 { return @intCast(self.constants.items.len - 1); } -/// Compute a basic complexity score based on size and presence "costly" opcodes -pub fn score(self: *Self) u32 { - if (self.complexity_score) |sc| return sc; - - var complexity_score: u32 = 0; - - for (self.code.items, 0..) |op, index| { - if (self.isInCompiledHotspotRange(index)) { - continue; - } - - complexity_score += 1; - - switch (VM.getCode(op)) { - .OP_HOTSPOT, // Those cover any loop - .OP_CALL, - .OP_TAIL_CALL, - .OP_CALL_INSTANCE_PROPERTY, - .OP_TAIL_CALL_INSTANCE_PROPERTY, - .OP_INSTANCE_INVOKE, - .OP_INSTANCE_TAIL_INVOKE, - .OP_PROTOCOL_INVOKE, - .OP_PROTOCOL_TAIL_INVOKE, - .OP_TRY, - .OP_TRY_END, - .OP_THROW, - => complexity_score += 1, - .OP_FIBER_FOREACH, - .OP_RESUME, - .OP_RESOLVE, - => return 0, // A chunk with fiber op codes will not be compiled so the score is 0 - else => {}, - } - } - - self.complexity_score = complexity_score; - - return complexity_score; -} - -pub fn addCompiledHotspotRange(self: *Self, start: usize, end: usize) !void { - if (start >= end) { - return; - } - - var merged = InstructionRange{ - .start = start, - .end = end, - }; - - var index: usize = 0; - while (index < self.compiled_hotspot_ranges.items.len) { - const range = self.compiled_hotspot_ranges.items[index]; - - if (merged.end < range.start) { - try self.compiled_hotspot_ranges.insert(self.allocator, index, merged); - self.complexity_score = null; - return; - } - - if (merged.start > range.end) { - index += 1; - continue; - } - - merged.start = @min(merged.start, range.start); - merged.end = @max(merged.end, range.end); - _ = self.compiled_hotspot_ranges.orderedRemove(index); - } - - try self.compiled_hotspot_ranges.append(self.allocator, merged); - self.complexity_score = null; -} - -fn isInCompiledHotspotRange(self: *const Self, index: usize) bool { - for (self.compiled_hotspot_ranges.items) |range| { - if (index < range.start) { - return false; - } - - if (index >= range.start and index < range.end) { - return true; - } - } - - return false; -} - pub const OpCode = enum(u8) { OP_CONSTANT, OP_NULL, @@ -278,11 +185,6 @@ const Self = @This(); pub const max_constants = std.math.maxInt(u24); -const InstructionRange = struct { - start: usize, - end: usize, -}; - const RegistryContext = struct { pub fn hash(_: RegistryContext, key: Self) u64 { return std.hash.Wyhash.hash( diff --git a/src/Codegen.zig b/src/Codegen.zig index 8320ec30..4ac5790a 100644 --- a/src/Codegen.zig +++ b/src/Codegen.zig @@ -66,6 +66,16 @@ const GeneratedNode = struct { flow: TerminalFlow = .{}, }; +/// Tracks the JIT complexity accumulated for a node while its code is generated. +const ComplexityFrame = struct { + /// Node currently being generated. + node: Ast.Node.Index, + /// Accumulated complexity for the node and generated children. + score: usize, + /// Nearest generated ancestor with stored JIT complexity metadata. + nearest_candidate: ?Ast.Node.Index, +}; + const Break = struct { ip: usize, // The op code will tell us if this is a continue or a break statement label_node: ?Ast.Node.Index = null, @@ -87,6 +97,8 @@ flavor: RunFlavor, opt_jumps: std.ArrayList(std.ArrayList(usize)) = .empty, /// Jumps emitted by `out` statements to the end of the current block expression. block_expression_out_jumps: std.ArrayList(std.ArrayList(usize)) = .empty, +/// Stack of nodes currently being generated, used to accumulate JIT complexity without another AST walk. +complexity_stack: std.ArrayList(ComplexityFrame) = .empty, /// Used to generate error messages parser: *Parser, jit: ?*JIT, @@ -188,6 +200,7 @@ pub fn init( } pub fn deinit(self: *Self) void { + self.complexity_stack.deinit(self.gc.allocator); self.reporter.deinit(); } @@ -539,6 +552,26 @@ fn generateStatements( return result; } +/// Returns the intrinsic JIT complexity for a generated node, before generated children contribute. +fn jitComplexityBase(tag: Ast.Node.Tag) usize { + return switch (tag) { + .AsyncCall, + .Resolve, + .Resume, + .Yield, + => 0, + .Call, + .DoUntil, + .For, + .ForEach, + .Throw, + .Try, + .While, + => 2, + else => 1, + }; +} + fn generateNode(self: *Self, node: Ast.Node.Index, breaks: ?*Breaks) Error!GeneratedNode { if (self.synchronize(node)) { return .{}; @@ -552,11 +585,50 @@ fn generateNode(self: *Self, node: Ast.Node.Index, breaks: ?*Breaks) Error!Gener node, ); - if (Self.generators[@intFromEnum(self.ast.nodes.items(.tag)[node])]) |generator| { - return generator(self, node, breaks); + const tag = self.ast.nodes.items(.tag)[node]; + const generator = Self.generators[@intFromEnum(tag)] orelse return .{}; + + const nearest_candidate = if (self.complexity_stack.items.len > 0) + self.complexity_stack.items[self.complexity_stack.items.len - 1].nearest_candidate + else + null; + + try self.complexity_stack.append(self.gc.allocator, .{ + .node = node, + .score = jitComplexityBase(tag), + .nearest_candidate = if (self.ast.jitComplexity(node) != null) node else nearest_candidate, + }); + errdefer { + const popped = self.complexity_stack.pop().?; + std.debug.assert(popped.node == node); + } + + const generated = try generator(self, node, breaks); + const frame = self.complexity_stack.pop().?; + std.debug.assert(frame.node == node); + + if (self.ast.jitComplexity(node)) |jit| { + jit.* = .{ + .score = frame.score, + .parent = nearest_candidate, + }; } - return .{}; + if (self.complexity_stack.items.len > 0) { + // Function bodies are complexity boundaries; the parent only pays for creating the function. + const contribution = switch (tag) { + .Function => jitComplexityBase(.Function), + else => frame.score, + }; + const current = &self.complexity_stack.items[self.complexity_stack.items.len - 1]; + if (current.score == 0 or contribution == 0) { + current.score = 0; + } else { + current.score += contribution; + } + } + + return generated; } fn generateAs(self: *Self, node: Ast.Node.Index, breaks: ?*Breaks) Error!GeneratedNode { @@ -1824,6 +1896,12 @@ fn generateForEach(self: *Self, node: Ast.Node.Index, breaks: ?*Breaks) Error!Ge const components = node_components[node].ForEach; const iterable_type_def = type_defs[components.iterable].?; + if (iterable_type_def.def_type == .Fiber) { + // Fiber foreach cannot be JIT compiled; a zero score propagates that blacklist upward. + const current = &self.complexity_stack.items[self.complexity_stack.items.len - 1]; + std.debug.assert(current.node == node); + current.score = 0; + } // If iterable constant and empty, skip the node if (try self.ast.isConstant(self.gc.allocator, components.iterable)) { diff --git a/src/Jit.zig b/src/Jit.zig index a6cd742b..863ebe87 100644 --- a/src/Jit.zig +++ b/src/Jit.zig @@ -218,7 +218,7 @@ pub fn compileFunctionIfNeeded(self: *Self, closure: *o.ObjClosure) StartError!b } if (BuildOptions.jit_always_on or closure.function.call_count > BuildOptions.jit_call_threshold) { - const score = closure.function.call_count * closure.function.chunk.score(); + const score = closure.function.call_count * function_ast.jitComplexity(closure.function.node).?.score; if (score == 0) { function_ast.nodes.items(.jit_status)[closure.function.node] = .blacklisted; @@ -260,7 +260,7 @@ pub fn compileHotspotIfNeeded(self: *Self, ast: Ast.Slice, frame_closure: *o.Obj } if (BuildOptions.jit_hotspot_always_on or ast.nodes.items(.count)[node] > BuildOptions.jit_hotspot_threshold) { - const score = ast.nodes.items(.count)[node] * try ast.score(self.gc.allocator, node); + const score = ast.nodes.items(.count)[node] * ast.jitComplexity(node).?.score; if (score == 0) { ast.nodes.items(.jit_status)[node] = .blacklisted; @@ -299,13 +299,10 @@ pub fn compile(self: *Self, ast: Ast.Slice, closure: *o.ObjClosure, hotspot_node .compilable => {}, } - if (try ast.usesFiber( - self.gc.allocator, - ast_node, - )) { + if (ast.jitComplexity(ast_node).?.score == 0) { if (BuildOptions.jit_debug) { log.debug( - "Not compiling node {s}#{}, likely because it uses a fiber\n", + "Not compiling node {s}#{}, it has zero JIT complexity\n", .{ @tagName(ast.nodes.items(.tag)[ast_node]), ast_node, @@ -365,7 +362,6 @@ pub fn compileFunctionSynchronously(self: *Self, closure: *o.ObjClosure) Error!v const function_ast = closure.function.chunk.ast; const function_node = closure.function.node; - _ = closure.function.chunk.score(); switch (function_ast.nodes.items(.jit_status)[function_node]) { .blacklisted, .generated, .queued => return error.CantCompile, @@ -373,10 +369,7 @@ pub fn compileFunctionSynchronously(self: *Self, closure: *o.ObjClosure) Error!v .compilable => {}, } - if (try function_ast.usesFiber( - self.gc.allocator, - function_node, - )) { + if (function_ast.jitComplexity(function_node).?.score == 0) { function_ast.nodes.items(.jit_status)[function_node] = .blacklisted; return error.CantCompile; @@ -436,7 +429,7 @@ fn work(self: *Self) Error!void { "Worker starting for compiling function `{s}` with score {}", .{ job.closure.function.type_def.resolved_type.?.Function.name.string, - job.closure.function.call_count * job.closure.function.chunk.complexity_score.?, + job.closure.function.call_count * job.ast.jitComplexity(job.node).?.score, }, ) else @@ -445,7 +438,7 @@ fn work(self: *Self) Error!void { .{ job.node, @tagName(job.ast.nodes.items(.tag)[job.node]), - job.ast.nodes.items(.count)[job.node] * job.ast.nodes.items(.complexity_score)[job.node].?, + job.ast.nodes.items(.count)[job.node] * job.ast.jitComplexity(job.node).?.score, job.closure.function.type_def.resolved_type.?.Function.name.string, }, ); @@ -570,7 +563,7 @@ fn doJob(self: *Self, job: *const Job) Error!CompletedJob { "Finished job function `{s}` with score {} in {}ms", .{ job.closure.function.type_def.resolved_type.?.Function.name.string, - job.closure.function.call_count * job.closure.function.chunk.complexity_score.?, + job.closure.function.call_count * job.ast.jitComplexity(job.node).?.score, time, }, ); @@ -580,7 +573,7 @@ fn doJob(self: *Self, job: *const Job) Error!CompletedJob { .{ job.node, @tagName(job.ast.nodes.items(.tag)[job.node]), - job.ast.nodes.items(.count)[job.node] * job.ast.nodes.items(.complexity_score)[job.node].?, + job.ast.nodes.items(.count)[job.node] * job.ast.jitComplexity(job.node).?.score, job.closure.function.type_def.resolved_type.?.Function.name.string, time, }, diff --git a/src/vm.zig b/src/vm.zig index 3433fc44..09cc7521 100644 --- a/src/vm.zig +++ b/src/vm.zig @@ -4578,7 +4578,7 @@ pub const VM = struct { }; obj_native.mark(self.gc); - // The now compile hotspot must be a new constant for the current function + // The newly compiled hotspot must be a new constant for the current function. frame.closure.function.chunk.constants.append( frame.closure.function.chunk.allocator, obj_native.toValue(), @@ -4596,6 +4596,27 @@ pub const VM = struct { self.panic("Out of memory"); unreachable; }; + + const compiled_hotspot_score = function_ast.jitComplexity(node).?.score; + if (compiled_hotspot_score > 0) { + const function_node = frame.closure.function.node; + var current_node: ?Ast.Node.Index = node; + + // Once the hotspot bytecode is patched out, remove its cost from candidate ancestors. + while (current_node) |score_node| { + const jit = function_ast.jitComplexity(score_node).?; + jit.score = if (jit.score > compiled_hotspot_score) + jit.score - compiled_hotspot_score + else + 0; + + if (score_node == function_node) { + break; + } + + current_node = jit.parent; + } + } } frame = self.currentFrame().?; @@ -5563,7 +5584,6 @@ pub const VM = struct { ); const hotspot_call_start = to - hotspot_call.len; - try chunk.addCompiledHotspotRange(frame.ip - 1, hotspot_call_start); // In the event that we are in a nested loop, we put a jump instruction in place of OP_HOTSPOT chunk.code.items[frame.ip - 2] = (@as(u32, @intCast(@intFromEnum(Chunk.OpCode.OP_JUMP))) << 24) | @as(