[wasm] Implement partial backward branch support in the Jiterpreter

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

@@ -173,6 +173,11 @@ struct InterpMethod {
 	unsigned int needs_thread_attach : 1;
 	// If set, this method is MulticastDelegate.Invoke
 	unsigned int is_invoke : 1;
+#if HOST_BROWSER
+	unsigned int contains_traces : 1;
+	guint16 *backward_branch_offsets;
+	unsigned int backward_branch_offsets_count;
+#endif
 #if PROFILE_INTERP
 	long calls;
 	long opcounts;

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

@@ -462,46 +462,6 @@ mono_jiterp_relop_fp (double lhs, double rhs, int opcode) {
 
 #undef JITERP_RELOP
 
-#define JITERP_MEMBER_VT_INITIALIZED 0
-#define JITERP_MEMBER_ARRAY_DATA 1
-#define JITERP_MEMBER_STRING_LENGTH 2
-#define JITERP_MEMBER_STRING_DATA 3
-#define JITERP_MEMBER_IMETHOD 4
-#define JITERP_MEMBER_DATA_ITEMS 5
-#define JITERP_MEMBER_RMETHOD 6
-#define JITERP_MEMBER_SPAN_LENGTH 7
-#define JITERP_MEMBER_SPAN_DATA 8
-#define JITERP_MEMBER_ARRAY_LENGTH 9
-
-// we use these helpers at JIT time to figure out where to do memory loads and stores
-EMSCRIPTEN_KEEPALIVE size_t
-mono_jiterp_get_member_offset (int member) {
-	switch (member) {
-		case JITERP_MEMBER_VT_INITIALIZED:
-			return MONO_STRUCT_OFFSET (MonoVTable, initialized);
-		case JITERP_MEMBER_ARRAY_DATA:
-			return MONO_STRUCT_OFFSET (MonoArray, vector);
-		case JITERP_MEMBER_ARRAY_LENGTH:
-			return MONO_STRUCT_OFFSET (MonoArray, max_length);
-		case JITERP_MEMBER_STRING_LENGTH:
-			return MONO_STRUCT_OFFSET (MonoString, length);
-		case JITERP_MEMBER_STRING_DATA:
-			return MONO_STRUCT_OFFSET (MonoString, chars);
-		case JITERP_MEMBER_IMETHOD:
-			return offsetof (InterpFrame, imethod);
-		case JITERP_MEMBER_DATA_ITEMS:
-			return offsetof (InterpMethod, data_items);
-		case JITERP_MEMBER_RMETHOD:
-			return offsetof (JiterpEntryDataHeader, rmethod);
-		case JITERP_MEMBER_SPAN_LENGTH:
-			return offsetof (MonoSpanOfVoid, _length);
-		case JITERP_MEMBER_SPAN_DATA:
-			return offsetof (MonoSpanOfVoid, _reference);
-		default:
-			g_assert_not_reached();
-	}
-}
-
 EMSCRIPTEN_KEEPALIVE size_t
 mono_jiterp_get_size_of_stackval () {
 	return sizeof(stackval);
@@ -733,7 +693,7 @@ jiterp_should_abort_trace (InterpInst *ins, gboolean *inside_branch_block)
 				if (*inside_branch_block)
 					return TRACE_CONTINUE;
 				else
-					return TRACE_ABORT;
+					return mono_opt_jiterpreter_backward_branches_enabled ? TRACE_CONTINUE : TRACE_ABORT;
 			}
 
 			*inside_branch_block = TRUE;
@@ -771,6 +731,7 @@ jiterp_should_abort_trace (InterpInst *ins, gboolean *inside_branch_block)
 			(opcode >= MINT_BRFALSE_I4) &&
 			(opcode <= MINT_BLT_UN_I8_IMM_SP)
 		) {
+			// FIXME: Detect negative displacement and abort appropriately
 			*inside_branch_block = TRUE;
 			return TRACE_CONTINUE;
 		}
@@ -941,16 +902,17 @@ trace_info_alloc () {
  *  instruction that would end up there instead and not waste any resources trying to compile it.
  */
 void
-jiterp_insert_entry_points (void *_td)
+jiterp_insert_entry_points (void *_imethod, void *_td)
 {
-	if (!mono_opt_jiterpreter_traces_enabled)
-		return;
+	InterpMethod *imethod = (InterpMethod *)_imethod;
 	TransformData *td = (TransformData *)_td;
-
 	// Insert an entry opcode for the next basic block (call resume and first bb)
 	// FIXME: Should we do this based on relationships between BBs instead of insn sequence?
 	gboolean enter_at_next = TRUE;
 
+	if (!mono_opt_jiterpreter_traces_enabled)
+		return;
+
 	for (InterpBasicBlock *bb = td->entry_bb; bb != NULL; bb = bb->next_bb) {
 		// Enter trace at top of functions
 		gboolean is_backwards_branch = FALSE,
@@ -979,6 +941,7 @@ jiterp_insert_entry_points (void *_td)
 			gint32 trace_index = trace_info_alloc ();
 
 			td->cbb = bb;
+			imethod->contains_traces = TRUE;
 			InterpInst *ins = mono_jiterp_insert_ins (td, NULL, MINT_TIER_PREPARE_JITERPRETER);
 			memcpy(ins->data, &trace_index, sizeof (trace_index));
 
@@ -1261,6 +1224,52 @@ mono_jiterp_trace_transfer (
 	return displacement + relative_displacement;
 }
 
+#define JITERP_MEMBER_VT_INITIALIZED 0
+#define JITERP_MEMBER_ARRAY_DATA 1
+#define JITERP_MEMBER_STRING_LENGTH 2
+#define JITERP_MEMBER_STRING_DATA 3
+#define JITERP_MEMBER_IMETHOD 4
+#define JITERP_MEMBER_DATA_ITEMS 5
+#define JITERP_MEMBER_RMETHOD 6
+#define JITERP_MEMBER_SPAN_LENGTH 7
+#define JITERP_MEMBER_SPAN_DATA 8
+#define JITERP_MEMBER_ARRAY_LENGTH 9
+#define JITERP_MEMBER_BACKWARD_BRANCH_OFFSETS 10
+#define JITERP_MEMBER_BACKWARD_BRANCH_OFFSETS_COUNT 11
+
+// we use these helpers at JIT time to figure out where to do memory loads and stores
+EMSCRIPTEN_KEEPALIVE size_t
+mono_jiterp_get_member_offset (int member) {
+	switch (member) {
+		case JITERP_MEMBER_VT_INITIALIZED:
+			return MONO_STRUCT_OFFSET (MonoVTable, initialized);
+		case JITERP_MEMBER_ARRAY_DATA:
+			return MONO_STRUCT_OFFSET (MonoArray, vector);
+		case JITERP_MEMBER_ARRAY_LENGTH:
+			return MONO_STRUCT_OFFSET (MonoArray, max_length);
+		case JITERP_MEMBER_STRING_LENGTH:
+			return MONO_STRUCT_OFFSET (MonoString, length);
+		case JITERP_MEMBER_STRING_DATA:
+			return MONO_STRUCT_OFFSET (MonoString, chars);
+		case JITERP_MEMBER_IMETHOD:
+			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_RMETHOD:
+			return offsetof (JiterpEntryDataHeader, rmethod);
+		case JITERP_MEMBER_SPAN_LENGTH:
+			return offsetof (MonoSpanOfVoid, _length);
+		case JITERP_MEMBER_SPAN_DATA:
+			return offsetof (MonoSpanOfVoid, _reference);
+		default:
+			g_assert_not_reached();
+	}
+}
+
 // HACK: fix C4206
 EMSCRIPTEN_KEEPALIVE
 #endif // HOST_BROWSER

src/mono/mono/mini/interp/jiterpreter.h

@@ -34,7 +34,7 @@ jiterp_preserve_module ();
 
 // HACK: Pass void* so that this header can include safely in files without definition for TransformData
 void
-jiterp_insert_entry_points (void *td);
+jiterp_insert_entry_points (void *imethod, void *td);
 
 // used by the typescript JIT implementation to notify the runtime that it has finished jitting a thunk
 //  for a specific callsite, since it can take a while before it happens

src/mono/mono/mini/interp/transform.c

@@ -8442,13 +8442,18 @@ add_patchpoint_data (TransformData *td, int patchpoint_data_index, int native_of
 
 // Generates the final code, after we are done with all the passes
 static void
-generate_compacted_code (TransformData *td)
+generate_compacted_code (InterpMethod *rtm, TransformData *td)
 {
 	guint16 *ip;
 	int size;
 	int patchpoint_data_index = 0;
 	td->relocs = g_ptr_array_new ();
 	InterpBasicBlock *bb;
+#if HOST_BROWSER
+	#define BACKWARD_BRANCH_OFFSETS_SIZE 64
+	unsigned int backward_branch_offsets_count = 0;
+	guint16 backward_branch_offsets[BACKWARD_BRANCH_OFFSETS_SIZE] = { 0 };
+#endif
 
 	// This iteration could be avoided at the cost of less precise size result, following
 	// super instruction pass
@@ -8467,6 +8472,14 @@ generate_compacted_code (TransformData *td)
 		InterpInst *ins = bb->first_ins;
 		bb->native_offset = GPTRDIFF_TO_INT (ip - td->new_code);
 		td->cbb = bb;
+
+#if HOST_BROWSER
+		if (bb->backwards_branch_target && rtm->contains_traces) {
+			if (backward_branch_offsets_count < BACKWARD_BRANCH_OFFSETS_SIZE)
+				backward_branch_offsets[backward_branch_offsets_count++] = ip - td->new_code;
+		}
+#endif
+
 		if (bb->patchpoint_data)
 			patchpoint_data_index = add_patchpoint_data (td, patchpoint_data_index, bb->native_offset, bb->index);
 		if (bb->emit_patchpoint) {
@@ -8493,6 +8506,16 @@ generate_compacted_code (TransformData *td)
 	handle_relocations (td);
 
 	g_ptr_array_free (td->relocs, TRUE);
+
+#if HOST_BROWSER
+	if (backward_branch_offsets_count > 0) {
+		rtm->backward_branch_offsets = mono_mem_manager_alloc(td->mem_manager, backward_branch_offsets_count * sizeof(guint16));
+		rtm->backward_branch_offsets_count = backward_branch_offsets_count;
+		memcpy(rtm->backward_branch_offsets, backward_branch_offsets, backward_branch_offsets_count * sizeof(guint16));
+	}
+
+	#undef BACKWARD_BRANCH_OFFSETS_SIZE
+#endif
 }
 
 static void
@@ -10759,11 +10782,11 @@ generate (MonoMethod *method, MonoMethodHeader *header, InterpMethod *rtm, MonoG
 		interp_optimize_code (td);
 		interp_alloc_offsets (td);
 #if HOST_BROWSER
-		jiterp_insert_entry_points (td);
+		jiterp_insert_entry_points (rtm, td);
 #endif
 	}
 
-	generate_compacted_code (td);
+	generate_compacted_code (rtm, td);
 
 	if (td->total_locals_size >= G_MAXUINT16) {
 		if (td->disable_inlining) {

src/mono/mono/utils/options-def.h

@@ -108,6 +108,8 @@ DEFINE_BOOL(jiterpreter_dump_traces, "jiterpreter-dump-traces", FALSE, "Dump the
 DEFINE_BOOL(jiterpreter_use_constants, "jiterpreter-use-constants", FALSE, "Use runtime imports for pointer constants")
 // Attempt to eliminate redundant null checks in compiled traces
 DEFINE_BOOL(jiterpreter_eliminate_null_checks, "jiterpreter-eliminate-null-checks", TRUE, "Attempt to eliminate redundant null checks in traces")
+// enables performing backward branches without exiting traces
+DEFINE_BOOL(jiterpreter_backward_branches_enabled, "jiterpreter-backward-branches-enabled", TRUE, "Enable performing backward branches without exiting traces")
 // When compiling a jit_call wrapper, bypass sharedvt wrappers if possible by inlining their
 //  logic into the compiled wrapper and calling the target AOTed function with native call convention
 DEFINE_BOOL(jiterpreter_direct_jit_call, "jiterpreter-direct-jit-calls", TRUE, "Bypass gsharedvt wrappers when compiling JIT call wrappers")

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

@@ -60,6 +60,7 @@ export class WasmBuilder {
     branchTargets = new Set<MintOpcodePtr>();
     options!: JiterpreterOptions;
     constantSlots: Array<number> = [];
+    backBranchOffsets: Array<MintOpcodePtr> = [];
     nextConstantSlot = 0;
 
     constructor (constantSlotCount: number) {
@@ -91,6 +92,7 @@ export class WasmBuilder {
         this.constantSlots.length = this.options.useConstants ? constantSlotCount : 0;
         for (let i = 0; i < this.constantSlots.length; i++)
             this.constantSlots[i] = 0;
+        this.backBranchOffsets.length = 0;
 
         this.allowNullCheckOptimization = this.options.eliminateNullChecks;
     }
@@ -723,6 +725,8 @@ export const counters = {
     failures: 0,
     bytesGenerated: 0,
     nullChecksEliminated: 0,
+    backBranchesEmitted: 0,
+    backBranchesNotEmitted: 0,
 };
 
 export const _now = (globalThis.performance && globalThis.performance.now)
@@ -950,6 +954,8 @@ export const enum JiterpMember {
     SpanLength = 7,
     SpanData = 8,
     ArrayLength = 9,
+    BackwardBranchOffsets = 10,
+    BackwardBranchOffsetsCount = 11,
 }
 
 const memberOffsets : { [index: number] : number } = {};
@@ -995,6 +1001,8 @@ export type JiterpreterOptions = {
     dumpTraces: boolean;
     // Use runtime imports for pointer constants
     useConstants: boolean;
+    // Enable performing backward branches without exiting traces
+    noExitBackwardBranches: boolean;
     // Unwrap gsharedvt wrappers when compiling jitcalls if possible
     directJitCalls: boolean;
     eliminateNullChecks: boolean;
@@ -1022,6 +1030,7 @@ const optionNames : { [jsName: string] : string } = {
     "dumpTraces": "jiterpreter-dump-traces",
     "useConstants": "jiterpreter-use-constants",
     "eliminateNullChecks": "jiterpreter-eliminate-null-checks",
+    "noExitBackwardBranches": "jiterpreter-backward-branches-enabled",
     "directJitCalls": "jiterpreter-direct-jit-calls",
     "minimumTraceLength": "jiterpreter-minimum-trace-length",
     "minimumTraceHitCount": "jiterpreter-minimum-trace-hit-count",