[wasm] Disable jiterpreter null check optimization for mid-method traces (#82585)

* Disable null check optimization when starting mid-method since we don't know which locals have had their address taken
* Improve null check optimization diagnostics

diff --git a/src/mono/wasm/runtime/jiterpreter-support.ts b/src/mono/wasm/runtime/jiterpreter-support.ts
index 3e91b3d..bd28587 100644 (file)
--- a/src/mono/wasm/runtime/jiterpreter-support.ts
+++ b/src/mono/wasm/runtime/jiterpreter-support.ts
@@ -35,6 +35,7 @@ export class WasmBuilder {
     stackSize!: number;
     inSection!: boolean;
     inFunction!: boolean;
+    allowNullCheckOptimization!: boolean;
     locals = new Map<string, [WasmValtype, number]>();
 
     permanentFunctionTypeCount = 0;
@@ -90,6 +91,8 @@ 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.allowNullCheckOptimization = this.options.eliminateNullChecks;
     }
 
     push () {

diff --git a/src/mono/wasm/runtime/jiterpreter-trace-generator.ts b/src/mono/wasm/runtime/jiterpreter-trace-generator.ts
index b62ab62..bffbd59 100644 (file)
--- a/src/mono/wasm/runtime/jiterpreter-trace-generator.ts
+++ b/src/mono/wasm/runtime/jiterpreter-trace-generator.ts
@@ -1165,7 +1165,7 @@ function append_local_null_check (builder: WasmBuilder, localOffset: number, ip:
 
 // Loads the specified i32 value and then bails out if it is null, leaving it in the cknull_ptr local.
 function append_ldloc_cknull (builder: WasmBuilder, localOffset: number, ip: MintOpcodePtr, leaveOnStack: boolean) {
-    if (knownNotNull.has(localOffset)) {
+    if (builder.allowNullCheckOptimization && knownNotNull.has(localOffset)) {
         counters.nullChecksEliminated++;
         if (cknullOffset === localOffset) {
             // console.log(`cknull_ptr already contains ${localOffset}`);
@@ -1180,7 +1180,9 @@ function append_ldloc_cknull (builder: WasmBuilder, localOffset: number, ip: Min
 
         if (nullCheckValidation) {
             builder.local("cknull_ptr");
+            append_ldloc(builder, localOffset, WasmOpcode.i32_load);
             builder.i32_const(builder.base);
+            builder.i32_const(ip);
             builder.callImport("notnull");
         }
         return;
@@ -1196,8 +1198,12 @@ function append_ldloc_cknull (builder: WasmBuilder, localOffset: number, ip: Min
     if (leaveOnStack)
         builder.local("cknull_ptr");
 
-    if (!addressTakenLocals.has(localOffset) && builder.options.eliminateNullChecks) {
+    if (
+        !addressTakenLocals.has(localOffset) &&
+        builder.allowNullCheckOptimization
+    ) {
         knownNotNull.add(localOffset);
+        // FIXME: This breaks the ldelema1 implementation somehow
         cknullOffset = localOffset;
     }
 }

diff --git a/src/mono/wasm/runtime/jiterpreter.ts b/src/mono/wasm/runtime/jiterpreter.ts
index be6dc01..15a12d2 100644 (file)
--- a/src/mono/wasm/runtime/jiterpreter.ts
+++ b/src/mono/wasm/runtime/jiterpreter.ts
@@ -46,7 +46,8 @@ export const
     // Very expensive!!!!
     trapTraceErrors = false,
     // When eliminating a null check, replace it with a runtime 'not null' assertion
-    //  that will print a diagnostic message if the value is actually null
+    //  that will print a diagnostic message if the value is actually null or if
+    //  the value does not match the value on the native interpreter stack in memory
     nullCheckValidation = false,
     // 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
@@ -77,6 +78,8 @@ export const disabledOpcodes : Array<MintOpcode> = [
 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"
+    // "ResolveMethodArguments"
+    // "HashCode"
 ];
 
 export class InstrumentedTraceState {
@@ -534,7 +537,9 @@ function initialize_builder (builder: WasmBuilder) {
     builder.defineType(
         "notnull", {
             "ptr": WasmValtype.i32,
+            "expected": WasmValtype.i32,
             "traceIp": WasmValtype.i32,
+            "ip": WasmValtype.i32,
         }, WasmValtype.void, true
     );
     builder.defineType(
@@ -570,12 +575,12 @@ function initialize_builder (builder: WasmBuilder) {
 }
 
 function assert_not_null (
-    value: number, traceIp: MintOpcodePtr
+    value: number, expectedValue: number, traceIp: MintOpcodePtr, ip: MintOpcodePtr
 ) {
-    if (value)
+    if (value && (value === expectedValue))
         return;
     const info = traceInfo[<any>traceIp];
-    throw new Error(`expected non-null value in trace ${info.name} but found null`);
+    throw new Error(`expected non-null value ${expectedValue} but found ${value} in trace ${info.name} @ 0x${(<any>ip).toString(16)}`);
 }
 
 // returns function id
@@ -605,6 +610,15 @@ function generate_wasm (
     const endOfBody = <any>startOfBody + <any>sizeOfBody;
     const traceName = `${methodName}:${(traceOffset).toString(16)}`;
 
+    // HACK: If we aren't starting at the beginning of the method, we don't know which
+    //  locals may have already had their address taken, so the null check optimization
+    //  is potentially invalid since there could be addresses on the stack
+    // FIXME: The interpreter maintains information on which locals have had their address
+    //  taken, so if we flow that information through it will allow us to make this optimization
+    //  robust in all scenarios and remove this hack
+    if (traceOffset > 0)
+        builder.allowNullCheckOptimization = false;
+
     if (useDebugCount) {
         if (cwraps.mono_jiterp_debug_count() === 0) {
             if (countLimitedPrintCounter-- >= 0)
@@ -775,7 +789,7 @@ function generate_wasm (
     } catch (exc: any) {
         threw = true;
         rejected = false;
-        console.error(`MONO_WASM: ${traceName} code generation failed: ${exc} ${exc.stack}`);
+        console.error(`MONO_WASM: ${methodFullName || traceName} code generation failed: ${exc} ${exc.stack}`);
         recordFailure();
         return 0;
     } finally {
@@ -793,7 +807,7 @@ function generate_wasm (
                     console.log(builder.traceBuf[i]);
             }
 
-            console.log(`// MONO_WASM: ${traceName} generated, blob follows //`);
+            console.log(`// MONO_WASM: ${methodFullName || traceName} generated, blob follows //`);
             let s = "", j = 0;
             try {
                 // We may have thrown an uncaught exception while inside a block,