[wasm] New jiterpreter backbranch scan pass (#98530)

Remove the C code responsible for generating the backward_branch_offsets table stored in InterpMethod that was used by the jiterpreter.
Add a new pre-pass to the jiterpreter that does a quick opcode scan to find back-branch targets on the fly when a trace is actually being compiled.
Reimplement interpreter branch decoding in a central place in the jiterpreter, and generalize it based on opcode info tables instead of a bunch of hand-written decode logic.
Remove some dead options and dead code from the jiterpreter typescript.
Fix an enum that got out of sync.
Minor debug logging improvements.

Author: kg

src/mono/browser/runtime/jiterpreter-opcodes.ts

@@ -1,8 +1,6 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-// Keep this file in sync with mintops.def. The order and values need to match exactly.
-
 import cwraps from "./cwraps";
 import { utf8ToString } from "./strings";
 import { OpcodeInfoType } from "./jiterpreter-enums";
@@ -28,27 +26,29 @@ export type SimdInfoTable = {
     [argument_count: number]: SimdInfoSubtable
 }
 
+// Keep in sync with mintops.h
 export const enum MintOpArgType {
-    MintOpNoArgs = 0,
-    MintOpShortInt,
-    MintOpUShortInt,
-    MintOpInt,
-    MintOpLongInt,
-    MintOpFloat,
-    MintOpDouble,
-    MintOpBranch,
-    MintOpShortBranch,
-    MintOpSwitch,
-    MintOpMethodToken,
-    MintOpFieldToken,
-    MintOpClassToken,
-    MintOpTwoShorts,
-    MintOpTwoInts,
-    MintOpShortAndInt,
-    MintOpShortAndShortBranch,
-    MintOpPair2,
-    MintOpPair3,
-    MintOpPair4
+	MintOpNoArgs = 0,
+	MintOpShortInt,
+	MintOpUShortInt,
+	MintOpInt,
+	MintOpLongInt,
+	MintOpFloat,
+	MintOpDouble,
+	MintOpBranch,
+	MintOpShortBranch,
+	MintOpSwitch,
+	MintOpMethodToken,
+	MintOpFieldToken,
+	MintOpClassToken,
+	MintOpVTableToken,
+	MintOpTwoShorts,
+	MintOpTwoInts,
+	MintOpShortAndInt,
+	MintOpShortAndShortBranch,
+	MintOpPair2,
+	MintOpPair3,
+	MintOpPair4
 }
 
 // keep in sync with jiterpreter.c, see mono_jiterp_relop_fp

src/mono/browser/runtime/jiterpreter-support.ts

@@ -1321,6 +1321,9 @@ class Cfg {
                 this.builder.appendU8(WasmOpcode.br_if);
                 this.builder.appendULeb(this.blockStack.indexOf(this.backDispatchOffsets[0]));
             } else {
+                if (this.trace > 0)
+                    mono_log_info(`${this.backDispatchOffsets.length} back branch offsets after filtering.`);
+
                 // the loop needs to start with a br_table that performs dispatch based on the current value
                 //  of the dispatch index local
                 // br_table has to be surrounded by a block in order for a depth of 0 to be fallthrough

src/mono/browser/runtime/jiterpreter-trace-generator.ts

@@ -9,7 +9,7 @@ import {
 } from "./memory";
 import {
     WasmOpcode, WasmSimdOpcode, WasmValtype,
-    getOpcodeName,
+    getOpcodeName, MintOpArgType
 } from "./jiterpreter-opcodes";
 import {
     MintOpcode, SimdInfo,
@@ -33,7 +33,7 @@ import {
 
     disabledOpcodes, countCallTargets,
     callTargetCounts,
-    trace, traceOnError, traceOnRuntimeError,
+    trace, traceOnError,
     emitPadding, traceBranchDisplacements,
     traceEip, nullCheckValidation,
     traceNullCheckOptimizations,
@@ -109,6 +109,7 @@ function is_backward_branch_target(
     if (!backwardBranchTable)
         return false;
 
+    // TODO: sort the table and exploit that for faster scan. Not important yet
     for (let i = 0; i < backwardBranchTable.length; i++) {
         const actualOffset = (backwardBranchTable[i] * 2) + <any>startOfBody;
         if (actualOffset === ip)
@@ -128,6 +129,74 @@ function get_known_constant_value(builder: WasmBuilder, localOffset: number): Kn
     return knownConstantValues.get(localOffset);
 }
 
+// Perform a quick scan through the opcodes potentially in this trace to build a table of
+//  backwards branch targets, compatible with the layout of the old one that was generated in C.
+// We do this here to match the exact way that the jiterp calculates branch targets, since
+//  there were previously corner cases where jiterp and interp disagreed.
+export function generateBackwardBranchTable(
+    ip: MintOpcodePtr, startOfBody: MintOpcodePtr, sizeOfBody: MintOpcodePtr,
+): Uint16Array | null {
+    const endOfBody = <any>startOfBody + <any>sizeOfBody;
+    // TODO: Cache this table object instance and reuse it to reduce gc pressure?
+    const table : number[] = [];
+    // IP of the start of the trace in U16s, relative to startOfBody.
+    const rbase16 = (<any>ip - <any>startOfBody) / 2;
+
+    while (ip < endOfBody) {
+        // IP of the current opcode in U16s, relative to startOfBody. This is what the back branch table uses
+        const rip16 = (<any>ip - <any>startOfBody) / 2;
+        const opcode = <MintOpcode>getU16(ip);
+        // HACK
+        if (opcode === MintOpcode.MINT_SWITCH)
+            break;
+
+        const opLengthU16 = cwraps.mono_jiterp_get_opcode_info(opcode, OpcodeInfoType.Length);
+        // Any opcode with a branch argtype will have a decoded displacement, even if we don't
+        //  implement the opcode. Everything else will return undefined here and be skipped
+        const displacement = getBranchDisplacement(ip, opcode);
+        if (typeof (displacement) !== "number") {
+            ip += <any>(opLengthU16 * 2);
+            continue;
+        }
+
+        // These checks shouldn't fail unless memory is corrupted or something is wrong with the decoder.
+        // We don't want to cause decoder bugs to make the application exit, though - graceful degradation.
+        if (displacement === 0) {
+            mono_log_info(`opcode @${ip} branch target is self. aborting backbranch table generation`);
+            break;
+        }
+
+        const rtarget16 = rip16 + (displacement);
+        if (rtarget16 < 0) {
+            mono_log_info(`opcode @${ip}'s displacement of ${displacement} goes before body: ${rtarget16}. aborting backbranch table generation`);
+            break;
+        }
+
+        // If the relative target is before the start of the trace, don't record it.
+        // The trace will be unable to successfully branch to it so it would just make the table bigger.
+        if (rtarget16 >= rbase16)
+            table.push(rtarget16);
+
+        switch (opcode) {
+            case MintOpcode.MINT_CALL_HANDLER:
+            case MintOpcode.MINT_CALL_HANDLER_S:
+                // While this formally isn't a backward branch target, we want to record
+                //  the offset of its following instruction so that the jiterpreter knows
+                //  to generate the necessary dispatch code to enable branching back to it.
+                table.push(rip16 + opLengthU16);
+                break;
+        }
+
+        ip += <any>(opLengthU16 * 2);
+    }
+
+    if (table.length <= 0)
+        return null;
+    // Not important yet, so not doing it
+    // table.sort((a, b) => a - b);
+    return new Uint16Array(table);
+}
+
 export function generateWasmBody(
     frame: NativePointer, traceName: string, ip: MintOpcodePtr,
     startOfBody: MintOpcodePtr, endOfBody: MintOpcodePtr,
@@ -1546,7 +1615,7 @@ export function generateWasmBody(
                 }
             }
 
-            if ((trace > 1) || traceOnError || traceOnRuntimeError || mostRecentOptions!.dumpTraces || instrumentedTraceId) {
+            if ((trace > 1) || traceOnError || mostRecentOptions!.dumpTraces || instrumentedTraceId) {
                 let stmtText = `${(<any>ip).toString(16)} ${opname} `;
                 const firstDreg = <any>ip + 2;
                 const firstSreg = firstDreg + (numDregs * 2);
@@ -2572,13 +2641,45 @@ function append_call_handler_store_ret_ip(
     builder.callHandlerReturnAddresses.push(retIp);
 }
 
+function getBranchDisplacement(
+    ip: MintOpcodePtr, opcode: MintOpcode
+) : number | undefined {
+    const opArgType = cwraps.mono_jiterp_get_opcode_info(opcode, OpcodeInfoType.OpArgType),
+        payloadOffset = cwraps.mono_jiterp_get_opcode_info(opcode, OpcodeInfoType.Sregs),
+        payloadAddress = <any>ip + 2 + (payloadOffset * 2);
+
+    let result : number;
+    switch (opArgType) {
+        case MintOpArgType.MintOpBranch:
+            result = getI32_unaligned(payloadAddress);
+            break;
+        case MintOpArgType.MintOpShortBranch:
+            result = getI16(payloadAddress);
+            break;
+        case MintOpArgType.MintOpShortAndShortBranch:
+            result = getI16(payloadAddress + 2);
+            break;
+        default:
+            return undefined;
+    }
+
+    if (traceBranchDisplacements)
+        mono_log_info(`${getOpcodeName(opcode)} @${ip} displacement=${result}`);
+
+    return result;
+}
+
 function emit_branch(
     builder: WasmBuilder, ip: MintOpcodePtr,
-    frame: NativePointer, opcode: MintOpcode, displacement?: number
+    frame: NativePointer, opcode: MintOpcode
 ): boolean {
     const isSafepoint = (opcode >= MintOpcode.MINT_BRFALSE_I4_SP) &&
         (opcode <= MintOpcode.MINT_BLT_UN_I8_IMM_SP);
 
+    const displacement = getBranchDisplacement(ip, opcode);
+    if (typeof (displacement) !== "number")
+        return false;
+
     // If the branch is taken we bail out to allow the interpreter to do it.
     // So for brtrue, we want to do 'cond == 0' to produce a bailout only
     //  when the branch will be taken (by skipping the bailout in this block)
@@ -2592,15 +2693,7 @@ function emit_branch(
         case MintOpcode.MINT_BR_S: {
             const isCallHandler = (opcode === MintOpcode.MINT_CALL_HANDLER) ||
                 (opcode === MintOpcode.MINT_CALL_HANDLER_S);
-            displacement = (
-                (opcode === MintOpcode.MINT_BR) ||
-                (opcode === MintOpcode.MINT_CALL_HANDLER)
-            )
-                ? getArgI32(ip, 1)
-                : getArgI16(ip, 1);
 
-            if (traceBranchDisplacements)
-                mono_log_info(`br.s @${ip} displacement=${displacement}`);
             const destination = <any>ip + (displacement * 2);
 
             if (displacement <= 0) {
@@ -2653,7 +2746,6 @@ function emit_branch(
 
             // Load the condition
 
-            displacement = getArgI16(ip, 2);
             append_ldloc(builder, getArgU16(ip, 1), is64 ? WasmOpcode.i64_load : WasmOpcode.i32_load);
             if (
                 (opcode === MintOpcode.MINT_BRFALSE_I4_S) ||
@@ -2684,11 +2776,6 @@ function emit_branch(
         }
     }
 
-    if (!displacement)
-        throw new Error("Branch had no displacement");
-    else if (traceBranchDisplacements)
-        mono_log_info(`${getOpcodeName(opcode)} @${ip} displacement=${displacement}`);
-
     const destination = <any>ip + (displacement * 2);
 
     if (displacement < 0) {
@@ -2741,10 +2828,6 @@ function emit_relop_branch(
     if (!relopInfo && !intrinsicFpBinop)
         return false;
 
-    const displacement = getArgI16(ip, 3);
-    if (traceBranchDisplacements)
-        mono_log_info(`relop @${ip} displacement=${displacement}`);
-
     const operandLoadOp = relopInfo
         ? relopInfo[1]
         : (
@@ -2779,7 +2862,7 @@ function emit_relop_branch(
         builder.callImport("relop_fp");
     }
 
-    return emit_branch(builder, ip, frame, opcode, displacement);
+    return emit_branch(builder, ip, frame, opcode);
 }
 
 function emit_math_intrinsic(builder: WasmBuilder, ip: MintOpcodePtr, opcode: MintOpcode): boolean {

src/mono/browser/runtime/jiterpreter.ts

@@ -4,23 +4,23 @@
 import { MonoMethod } from "./types/internal";
 import { NativePointer } from "./types/emscripten";
 import { Module, mono_assert, runtimeHelpers } from "./globals";
-import { getU16, getU32_unaligned, localHeapViewU8 } from "./memory";
+import { getU16 } from "./memory";
 import { WasmValtype, WasmOpcode, getOpcodeName } from "./jiterpreter-opcodes";
 import { MintOpcode } from "./mintops";
 import cwraps from "./cwraps";
 import {
     MintOpcodePtr, WasmBuilder, addWasmFunctionPointer,
     _now, isZeroPageReserved,
     getRawCwrap, importDef, JiterpreterOptions, getOptions, recordFailure,
-    getMemberOffset, getCounter, modifyCounter,
+    getCounter, modifyCounter,
     simdFallbackCounters, getWasmFunctionTable
 } from "./jiterpreter-support";
 import {
-    JiterpMember, BailoutReasonNames, BailoutReason,
+    BailoutReasonNames, BailoutReason,
     JiterpreterTable, JiterpCounter,
 } from "./jiterpreter-enums";
 import {
-    generateWasmBody
+    generateWasmBody, generateBackwardBranchTable
 } from "./jiterpreter-trace-generator";
 import { mono_jiterp_free_method_data_interp_entry } from "./jiterpreter-interp-entry";
 import { mono_jiterp_free_method_data_jit_call } from "./jiterpreter-jit-call";
@@ -34,10 +34,6 @@ export const
     // Record a trace of all managed interpreter opcodes then dump it to console
     //  if an error occurs while compiling the output wasm
     traceOnError = false,
-    // Record trace but dump it when the trace has a runtime error instead
-    //  requires trapTraceErrors to work and will slow trace compilation +
-    //  increase memory usage
-    traceOnRuntimeError = false,
     // Trace the method name, location and reason for each abort
     traceAbortLocations = false,
     // Count the number of times a given method is seen as a call target, then
@@ -61,11 +57,6 @@ export const
     // Print diagnostic information when generating backward branches
     // 1 = failures only, 2 = full detail
     traceBackBranches = 0,
-    // If we encounter an enter opcode that looks like a loop body and it was already
-    //  jitted, we should abort the current trace since it's not worth continuing
-    // Unproductive if we have backward branches enabled because it can stop us from jitting
-    //  nested loops
-    abortAtJittedLoopBodies = true,
     // Enable generating conditional backward branches for ENDFINALLY opcodes if we saw some CALL_HANDLER
     //  opcodes previously, up to this many potential return addresses. If a trace contains more potential
     //  return addresses than this we will not emit code for the ENDFINALLY opcode
@@ -1030,11 +1021,8 @@ export function mono_interp_tier_prepare_jiterpreter(
     const methodName = utf8ToString(cwraps.mono_wasm_method_get_name(method));
     info.name = methodFullName || methodName;
 
-    const imethod = getU32_unaligned(getMemberOffset(JiterpMember.Imethod) + <any>frame);
-    const backBranchCount = getU32_unaligned(getMemberOffset(JiterpMember.BackwardBranchOffsetsCount) + imethod);
-    const pBackBranches = getU32_unaligned(getMemberOffset(JiterpMember.BackwardBranchOffsets) + imethod);
-    let backwardBranchTable = backBranchCount
-        ? new Uint16Array(localHeapViewU8().buffer, pBackBranches, backBranchCount)
+    let backwardBranchTable = mostRecentOptions.noExitBackwardBranches
+        ? generateBackwardBranchTable(ip, startOfBody, sizeOfBody)
         : null;
 
     // If we're compiling a trace that doesn't start at the beginning of a method,

src/mono/mono/mini/interp/interp-internals.h

@@ -181,8 +181,6 @@ struct InterpMethod {
 	unsigned int is_verbose : 1;
 #if HOST_BROWSER
 	unsigned int contains_traces : 1;
-	guint16 *backward_branch_offsets;
-	unsigned int backward_branch_offsets_count;
 	MonoBitSet *address_taken_bits;
 #endif
 #if PROFILE_INTERP

src/mono/mono/mini/interp/jiterpreter.c

@@ -1165,7 +1165,9 @@ enum {
 	JITERP_MEMBER_SPAN_LENGTH,
 	JITERP_MEMBER_SPAN_DATA,
 	JITERP_MEMBER_ARRAY_LENGTH,
+	// Kept as-is but no longer implemented
 	JITERP_MEMBER_BACKWARD_BRANCH_OFFSETS,
+	// Ditto
 	JITERP_MEMBER_BACKWARD_BRANCH_OFFSETS_COUNT,
 	JITERP_MEMBER_CLAUSE_DATA_OFFSETS,
 	JITERP_MEMBER_PARAMS_COUNT,
@@ -1195,10 +1197,6 @@ mono_jiterp_get_member_offset (int member) {
 			return offsetof (InterpFrame, imethod);
 		case JITERP_MEMBER_DATA_ITEMS:
 			return offsetof (InterpMethod, data_items);
-		case JITERP_MEMBER_BACKWARD_BRANCH_OFFSETS:
-			return offsetof (InterpMethod, backward_branch_offsets);
-		case JITERP_MEMBER_BACKWARD_BRANCH_OFFSETS_COUNT:
-			return offsetof (InterpMethod, backward_branch_offsets_count);
 		case JITERP_MEMBER_CLAUSE_DATA_OFFSETS:
 			return offsetof (InterpMethod, clause_data_offsets);
 		case JITERP_MEMBER_RMETHOD:

src/mono/mono/mini/interp/mintops.h

@@ -8,6 +8,7 @@
 #include <config.h>
 #include <glib.h>
 
+// If you change this, update jiterpreter-opcodes.ts.
 typedef enum
 {
 	MintOpNoArgs,