Add an optimized C implementation of STFLD_O that uses the correct ki…

…nd of write barrier (#81806)

This PR moves most of the jiterpreter's STFLD_O implementation into a C function that is responsible for also doing the null check. As a bonus that function is able to use the correct kind of write barrier (though it's not clear to me whether the previous one was broken in any way).

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

@@ -1137,6 +1137,23 @@ mono_jiterp_get_array_rank (gint32 *dest, MonoObject **src)
 	return 1;
 }
 
+// Returns 1 on success so that the trace can do br_if to bypass its bailout
+EMSCRIPTEN_KEEPALIVE int
+mono_jiterp_set_object_field (
+	uint8_t *locals, guint32 fieldOffsetBytes,
+	guint32 targetLocalOffsetBytes, guint32 sourceLocalOffsetBytes
+) {
+	MonoObject * targetObject = *(MonoObject **)(locals + targetLocalOffsetBytes);
+	if (!targetObject)
+		return 0;
+	MonoObject ** target = (MonoObject **)(((uint8_t *)targetObject) + fieldOffsetBytes);
+	mono_gc_wbarrier_set_field_internal (
+		targetObject, target,
+		*(MonoObject **)(locals + sourceLocalOffsetBytes)
+	);
+	return 1;
+}
+
 EMSCRIPTEN_KEEPALIVE int
 mono_jiterp_debug_count ()
 {

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

@@ -23,6 +23,7 @@ import {
     callTargetCounts, trapTraceErrors,
     trace, traceOnError, traceOnRuntimeError,
     emitPadding, traceBranchDisplacements,
+    traceEip,
 
     mostRecentOptions,
 
@@ -146,7 +147,7 @@ export function generate_wasm_body (
             break;
         }
 
-        if (instrumentedTraceId) {
+        if (instrumentedTraceId && traceEip) {
             builder.i32_const(instrumentedTraceId);
             builder.ip_const(ip);
             builder.callImport("trace_eip");
@@ -1067,7 +1068,10 @@ function emit_fieldop (
         fieldOffset = getArgU16(ip, 3),
         localOffset = getArgU16(ip, isLoad ? 1 : 2);
 
-    if (opcode !== MintOpcode.MINT_LDFLDA_UNSAFE)
+    if (
+        (opcode !== MintOpcode.MINT_LDFLDA_UNSAFE) &&
+        (opcode !== MintOpcode.MINT_STFLD_O)
+    )
         append_ldloc_cknull(builder, objectOffset, ip, false);
 
     let setter = WasmOpcode.i32_store,
@@ -1114,16 +1118,29 @@ function emit_fieldop (
             getter = WasmOpcode.i64_load;
             setter = WasmOpcode.i64_store;
             break;
-        case MintOpcode.MINT_STFLD_O:
-            // dest
-            builder.local("cknull_ptr");
+        case MintOpcode.MINT_STFLD_O: {
+            /*
+             * Writing a ref-type field has to call an import to perform the write barrier anyway,
+             *  and technically it should use a different kind of barrier from copy_pointer. So
+             *  we define a special import that is responsible for performing the whole stfld_o
+             *  operation with as little trace-side overhead as possible
+             * Previously the pseudocode looked like:
+             *  cknull_ptr = *(MonoObject *)&locals[objectOffset];
+             *  if (!cknull_ptr) bailout;
+             *  copy_pointer(cknull_ptr + fieldOffset, *(MonoObject *)&locals[localOffset])
+             */
+            builder.block();
+            builder.local("pLocals");
             builder.i32_const(fieldOffset);
-            builder.appendU8(WasmOpcode.i32_add);
-            // src
-            append_ldloca(builder, localOffset);
-            // FIXME: Use mono_gc_wbarrier_set_field_internal
-            builder.callImport("copy_pointer");
+            builder.i32_const(objectOffset);
+            builder.i32_const(localOffset);
+            builder.callImport("stfld_o");
+            builder.appendU8(WasmOpcode.br_if);
+            builder.appendLeb(0);
+            append_bailout(builder, ip, BailoutReason.NullCheck);
+            builder.endBlock();
             return true;
+        }
         case MintOpcode.MINT_LDFLD_VT: {
             const sizeBytes = getArgU16(ip, 4);
             // dest

src/mono/wasm/runtime/jiterpreter.ts

@@ -40,6 +40,8 @@ export const
     traceBranchDisplacements = false,
     // Trace when we reject something for being too small
     traceTooSmall = false,
+    // For instrumented methods, trace their exact IP during execution
+    traceEip = false,
     // Wraps traces in a JS function that will trap errors and log the trace responsible.
     // Very expensive!!!!
     trapTraceErrors = false,
@@ -57,11 +59,18 @@ export const callTargetCounts : { [method: number] : number } = {};
 export let mostRecentTrace : InstrumentedTraceState | undefined;
 export let mostRecentOptions : JiterpreterOptions | undefined = undefined;
 
+// You can disable an opcode for debugging purposes by adding it to this list,
+//  instead of aborting the trace it will insert a bailout instead. This means that you will
+//  have trace code generated as if the opcode were otherwise enabled
 export const disabledOpcodes : Array<MintOpcode> = [
 ];
 
+// Detailed output and/or instrumentation will happen when a trace is jitted if the method fullname has a match
+// Having any items in this list will add some overhead to the jitting of *all* traces
+// These names can be substrings and instrumentation will happen if the substring is found in the full name
 export const instrumentedMethodNames : Array<string> = [
     // "System.Collections.Generic.Stack`1<System.Reflection.Emit.LocalBuilder>& System.Collections.Generic.Dictionary`2<System.Type, System.Collections.Generic.Stack`1<System.Reflection.Emit.LocalBuilder>>:FindValue (System.Type)"
+    // "InternalInsertNode"
 ];
 
 export class InstrumentedTraceState {
@@ -247,6 +256,7 @@ function getTraceImports () {
         ["hascsize", "hascsize", getRawCwrap("mono_jiterp_object_has_component_size")],
         ["hasflag", "hasflag", getRawCwrap("mono_jiterp_enum_hasflag")],
         ["array_rank", "array_rank", getRawCwrap("mono_jiterp_get_array_rank")],
+        ["stfld_o", "stfld_o", getRawCwrap("mono_jiterp_set_object_field")],
     ];
 
     if (instrumentedMethodNames.length > 0) {
@@ -493,6 +503,14 @@ function initialize_builder (builder: WasmBuilder) {
             "source": WasmValtype.i32,
         }, WasmValtype.i32, true
     );
+    builder.defineType(
+        "stfld_o", {
+            "locals": WasmValtype.i32,
+            "fieldOffsetBytes": WasmValtype.i32,
+            "targetLocalOffsetBytes": WasmValtype.i32,
+            "sourceLocalOffsetBytes": WasmValtype.i32,
+        }, WasmValtype.i32, true
+    );
 }
 
 // returns function id
@@ -534,7 +552,11 @@ function generate_wasm (
     let compileStarted = 0;
     let rejected = true, threw = false;
 
-    const instrument = methodFullName && (instrumentedMethodNames.indexOf(methodFullName) >= 0);
+    const instrument = methodFullName && (
+        instrumentedMethodNames.findIndex(
+            (filter) => methodFullName.indexOf(filter) >= 0
+        ) >= 0
+    );
     const instrumentedTraceId = instrument ? nextInstrumentedTraceId++ : 0;
     if (instrument) {
         console.log(`instrumenting: ${methodFullName}`);