feat: Tracing JIT

- Profile loops: super simple right now, a simple counter - When hits threshold, compile the node into a native function with no arguments - Patch the bytecode to replace the loop with the function call closes #134
buzz-language · Apr 30, 2024 · 4123bfa · 4123bfa
1 parent f26067d
commit 4123bfa
Show file tree

Hide file tree

Showing 8 changed files with 532 additions and 106 deletions.
diff --git a/build.zig b/build.zig
@@ -23,12 +23,14 @@ const BuzzDebugOptions = struct {
 const BuzzJITOptions = struct {
     on: bool,
     always_on: bool,
+    hotspot_always_on: bool,
     debug: bool,
     prof_threshold: f128 = 0.05,
 
     pub fn step(self: BuzzJITOptions, options: *Build.Step.Options) void {
         options.addOption(@TypeOf(self.debug), "jit_debug", self.debug);
         options.addOption(@TypeOf(self.always_on), "jit_always_on", self.always_on);
+        options.addOption(@TypeOf(self.hotspot_always_on), "jit_hotspot_always_on", self.hotspot_always_on);
         options.addOption(@TypeOf(self.on), "jit", self.on);
         options.addOption(@TypeOf(self.prof_threshold), "jit_prof_threshold", self.prof_threshold);
     }
@@ -252,6 +254,11 @@ pub fn build(b: *Build) !void {
                 "jit_always_on",
                 "JIT engine will compile any function encountered",
             ) orelse false,
+            .hotspot_always_on = !is_wasm and b.option(
+                bool,
+                "jit_hotspot_always_on",
+                "JIT engine will compile any hotspot encountered",
+            ) orelse false,
             .on = !is_wasm and b.option(
                 bool,
                 "jit",

diff --git a/src/Ast.zig b/src/Ast.zig
@@ -137,6 +137,16 @@ pub const Node = struct {
         While,
         Yield,
         Zdef,
+
+        pub fn isHotspot(self: Tag) bool {
+            return switch (self) {
+                .While,
+                .For,
+                .ForEach,
+                => true,
+                else => false,
+            };
+        }
     };
 
     pub const Components = union(Tag) {

diff --git a/src/Chunk.zig b/src/Chunk.zig
@@ -118,14 +118,17 @@ pub const OpCode = enum(u8) {
     OP_LIST_APPEND,
 
     OP_MAP,
-    // FIXMEL delete and only use OP_SET_MAP_SUBSCRIPT
+    // FIXME: delete and only use OP_SET_MAP_SUBSCRIPT
     OP_SET_MAP,
 
     OP_EXPORT,
     OP_IMPORT,
 
     OP_TO_STRING,
     OP_TYPEOF,
+
+    OP_HOTSPOT,
+    OP_HOTSPOT_CALL,
 };
 
 /// A chunk of code to execute

diff --git a/src/Codegen.zig b/src/Codegen.zig
@@ -57,8 +57,8 @@ jit: ?*JIT,
 
 reporter: Reporter,
 
-const generators = [_]NodeGen{
-    noGen, // AnonymousObjectType,
+const generators = [_]?NodeGen{
+    null, // AnonymousObjectType,
     generateAs, // As,
     generateAsyncCall, // AsyncCall,
     generateBinary, // Binary,
@@ -73,27 +73,27 @@ const generators = [_]NodeGen{
     generateEnum, // Enum,
     generateExport, // Export,
     generateExpression, // Expression,
-    noGen, // FiberType,
+    null, // FiberType,
     generateFloat, // Float,
     generateFor, // For,
     generateForceUnwrap, // ForceUnwrap,
     generateForEach, // ForEach,
     generateFunction, // Function,
-    noGen, // FunctionType,
+    null, // FunctionType,
     generateFunDeclaration, // FunDeclaration,
     generateGenericResolve, // GenericResolve,
-    noGen, // GenericResolveType,
-    noGen, // GenericType,
+    null, // GenericResolveType,
+    null, // GenericType,
     generateGrouping, // Grouping,
     generateIf, // If,
     generateImport, // Import,
     generateInteger, // Integer,
     generateIs, // Is,
     generateList, // List,
-    noGen, // ListType,
+    null, // ListType,
     generateMap, // Map,
-    noGen, // MapType,
-    noGen, // Namespace,
+    null, // MapType,
+    null, // Namespace,
     generateNamedVariable, // NamedVariable,
     generateNull, // Null,
     generateObjectDeclaration, // ObjectDeclaration,
@@ -105,7 +105,7 @@ const generators = [_]NodeGen{
     generateResolve, // Resolve,
     generateResume, // Resume,
     generateReturn, // Return,
-    noGen, // SimpleType,
+    null, // SimpleType,
     generateString, // String,
     generateStringLiteral, // StringLiteral,
     generateSubscript, // Subscript,
@@ -115,7 +115,7 @@ const generators = [_]NodeGen{
     generateTypeOfExpression, // TypeOfExpression,
     generateUnary, // Unary,
     generateUnwrap, // Unwrap,
-    noGen, // UserType,
+    null, // UserType,
     generateVarDeclaration, // VarDeclaration,
     generateVoid, // Void,
     generateWhile, // While,
@@ -245,12 +245,12 @@ pub fn patchJump(self: *Self, offset: usize) void {
         (@as(u32, @intCast(instruction)) << 24) | @as(u32, @intCast(jump));
 }
 
-pub fn patchTry(self: *Self, offset: usize) void {
+pub fn patchTryOrJit(self: *Self, offset: usize) void {
     std.debug.assert(offset < self.currentCode());
 
     const jump: usize = self.currentCode();
 
-    if (jump > 16777215) {
+    if (jump > std.math.maxInt(u24)) {
         self.reportError(
             .block_too_large,
             "Try block too large.",
@@ -408,7 +408,11 @@ inline fn generateNode(self: *Self, node: Ast.Node.Index, breaks: ?*std.ArrayLis
         return null;
     }
 
-    return try Self.generators[@intFromEnum(self.ast.nodes.items(.tag)[node])](self, node, breaks);
+    if (Self.generators[@intFromEnum(self.ast.nodes.items(.tag)[node])]) |generator| {
+        return generator(self, node, breaks);
+    }
+
+    return null;
 }
 
 fn nodeValue(self: *Self, node: Ast.Node.Index) Error!?Value {
@@ -1768,6 +1772,9 @@ fn generateFor(self: *Self, node: Ast.Node.Index, breaks: ?*std.ArrayList(usize)
         return null;
     }
 
+    const jit_jump = try self.emitJump(locations[node], .OP_HOTSPOT);
+    try self.emit(locations[node], node);
+
     for (components.init_declarations) |decl| {
         _ = try self.generateNode(decl, breaks);
     }
@@ -1829,6 +1836,8 @@ fn generateFor(self: *Self, node: Ast.Node.Index, breaks: ?*std.ArrayList(usize)
     try self.patchOptJumps(node);
     try self.endScope(node);
 
+    self.patchTryOrJit(jit_jump);
+
     return null;
 }
 
@@ -2034,6 +2043,8 @@ fn generateForEach(self: *Self, node: Ast.Node.Index, breaks: ?*std.ArrayList(us
     _ = try self.generateNode(components.iterable, breaks);
 
     const loop_start: usize = self.currentCode();
+    const jit_jump = try self.emitJump(locations[node], .OP_HOTSPOT);
+    try self.emit(locations[node], node);
 
     // Calls `next` and update key and value locals
     try self.emitOpCode(
@@ -2096,6 +2107,8 @@ fn generateForEach(self: *Self, node: Ast.Node.Index, breaks: ?*std.ArrayList(us
     );
     try self.endScope(node);
 
+    self.patchTryOrJit(jit_jump);
+
     return null;
 }
 
@@ -2272,7 +2285,7 @@ fn generateFunction(self: *Self, node: Ast.Node.Index, breaks: ?*std.ArrayList(u
     current_function.upvalue_count = @intCast(components.upvalue_binding.count());
 
     if (BuildOptions.debug) {
-        try disassembler.disassembleChunk(&current_function.chunk, current_function.name.string);
+        disassembler.disassembleChunk(&current_function.chunk, current_function.name.string);
         std.debug.print("\n\n", .{});
     }
 
@@ -3484,7 +3497,7 @@ fn generateTry(self: *Self, node: Ast.Node.Index, breaks: ?*std.ArrayList(usize)
     var exit_jumps = std.ArrayList(usize).init(self.gc.allocator);
     defer exit_jumps.deinit();
 
-    self.patchTry(try_jump);
+    self.patchTryOrJit(try_jump);
     var has_unconditional = components.unconditional_clause != null;
     for (components.clauses) |clause| {
         const error_type = type_defs[clause.type_def].?;
@@ -3833,6 +3846,9 @@ fn generateWhile(self: *Self, node: Ast.Node.Index, breaks: ?*std.ArrayList(usiz
 
     const loop_start: usize = self.currentCode();
 
+    const jit_jump = try self.emitJump(locations[node], .OP_HOTSPOT);
+    try self.emit(locations[node], node);
+
     if (condition_type_def.def_type == .Placeholder) {
         self.reporter.reportPlaceholder(self.ast, condition_type_def.resolved_type.?.Placeholder);
     }
@@ -3868,6 +3884,8 @@ fn generateWhile(self: *Self, node: Ast.Node.Index, breaks: ?*std.ArrayList(usiz
     try self.patchOptJumps(node);
     try self.endScope(node);
 
+    self.patchTryOrJit(jit_jump);
+
     return null;
 }