[wasm] Don't generate unused imports in jiterpreter traces (#83061)

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

@@ -7733,9 +7733,11 @@ MINT_IN_CASE(MINT_BRTRUE_I8_SP) ZEROP_SP(gint64, !=); MINT_IN_BREAK;
 						break;
 					case JITERPRETER_NOT_JITTED:
 						// Patch opcode to disable it because this trace failed to JIT.
-						mono_memory_barrier ();
-						*mutable_ip = MINT_TIER_NOP_JITERPRETER;
-						mono_memory_barrier ();
+						if (!mono_opt_jiterpreter_estimate_heat) {
+							mono_memory_barrier ();
+							*mutable_ip = MINT_TIER_NOP_JITERPRETER;
+							mono_memory_barrier ();
+						}
 						ip += 3;
 						break;
 					default:

src/mono/wasm/runtime/jiterpreter-interp-entry.ts

@@ -283,10 +283,10 @@ function flush_wasm_entry_trampoline_jit_queue () {
         for (let i = 0; i < trampImports.length; i++) {
             mono_assert(trampImports[i], () => `trace #${i} missing`);
             const wasmName = compress ? i.toString(shortNameBase) : undefined;
-            builder.defineImportedFunction("i", trampImports[i][0], trampImports[i][1], wasmName);
+            builder.defineImportedFunction("i", trampImports[i][0], trampImports[i][1], true, wasmName);
         }
 
-        builder.generateImportSection();
+        builder._generateImportSection();
 
         // Function section
         builder.beginSection(3);
@@ -326,6 +326,7 @@ function flush_wasm_entry_trampoline_jit_queue () {
                 throw new Error(`Failed to generate ${info.traceName}`);
 
             builder.appendU8(WasmOpcode.end);
+            builder.endFunction(true);
         }
 
         builder.endSection();

src/mono/wasm/runtime/jiterpreter-jit-call.ts

@@ -394,9 +394,9 @@ export function mono_interp_flush_jitcall_queue () : void {
         // Emit function imports
         for (let i = 0; i < trampImports.length; i++) {
             const wasmName = compress ? i.toString(shortNameBase) : undefined;
-            builder.defineImportedFunction("i", trampImports[i][0], trampImports[i][1], wasmName);
+            builder.defineImportedFunction("i", trampImports[i][0], trampImports[i][1], true, wasmName);
         }
-        builder.generateImportSection();
+        builder._generateImportSection();
 
         // Function section
         builder.beginSection(3);
@@ -432,6 +432,7 @@ export function mono_interp_flush_jitcall_queue () : void {
             if (!ok)
                 throw new Error(`Failed to generate ${info.name}`);
             builder.appendU8(WasmOpcode.end);
+            builder.endFunction(true);
         }
 
         builder.endSection();

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

@@ -1,6 +1,7 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+import { mono_assert } from "./types";
 import { NativePointer, ManagedPointer, VoidPtr } from "./types/emscripten";
 import { Module } from "./imports";
 import { WasmOpcode } from "./jiterpreter-opcodes";
@@ -30,6 +31,26 @@ type FunctionTypeByIndex = [
     returnType: WasmValtype,
 ];
 
+type FunctionInfo = {
+    index: number;
+    name: string;
+    typeName: string;
+    typeIndex: number;
+    locals: { [name: string]: WasmValtype };
+    export: boolean;
+    generator: Function;
+    error: Error | null;
+    blob: Uint8Array | null;
+}
+
+type ImportedFunctionInfo = {
+    index?: number;
+    typeIndex: number;
+    module: string;
+    name: string;
+    friendlyName: string;
+}
+
 export class WasmBuilder {
     stack: Array<BlobBuilder>;
     stackSize!: number;
@@ -49,10 +70,12 @@ export class WasmBuilder {
     functionTypesByIndex: { [index: number] : FunctionTypeByIndex } = {};
 
     importedFunctionCount!: number;
-    importedFunctions!: { [name: string] : [
-        index: number, typeIndex: number, unknown: string
-    ] };
-    importsToEmit!: Array<[string, string, number, number]>;
+    importedFunctions!: { [name: string] : ImportedFunctionInfo };
+    nextImportIndex = 0;
+
+    functions: Array<FunctionInfo> = [];
+    estimatedExportBytes = 0;
+
     argumentCount!: number;
     activeBlocks!: number;
     base!: MintOpcodePtr;
@@ -80,9 +103,13 @@ export class WasmBuilder {
         this.functionTypesByShape = Object.create(this.permanentFunctionTypesByShape);
         this.functionTypesByIndex = Object.create(this.permanentFunctionTypesByIndex);
 
+        this.nextImportIndex = 0;
         this.importedFunctionCount = 0;
         this.importedFunctions = {};
-        this.importsToEmit = [];
+
+        this.functions.length = 0;
+        this.estimatedExportBytes = 0;
+
         this.argumentCount = 0;
         this.current.clear();
         this.traceBuf.length = 0;
@@ -97,27 +124,42 @@ export class WasmBuilder {
         this.allowNullCheckOptimization = this.options.eliminateNullChecks;
     }
 
-    push () {
+    _push () {
         this.stackSize++;
         if (this.stackSize >= this.stack.length)
             this.stack.push(new BlobBuilder());
         this.current.clear();
     }
 
-    pop () {
+    _pop (writeToOutput: boolean) {
         if (this.stackSize <= 1)
             throw new Error("Stack empty");
 
         const current = this.current;
         this.stackSize--;
 
-        this.appendULeb(current.size);
         const av = current.getArrayView();
-        this.appendBytes(av);
+        if (writeToOutput) {
+            this.appendULeb(current.size);
+            this.appendBytes(av);
+            return null;
+        } else
+            return av.slice();
     }
 
+    // HACK: Approximate amount of space we need to generate the full module at present
+    // FIXME: This does not take into account any other functions already generated if they weren't
+    //  emitted into the module immediately
     get bytesGeneratedSoFar () {
-        return this.stack[0].size;
+        return this.stack[0].size +
+            // HACK: A random constant for section headers and padding
+            32 +
+            // mod (2 bytes) name (2-3 bytes) type (1 byte) typeidx (1-2 bytes)
+            (this.importedFunctionCount * 8) +
+            // type index for each function
+            (this.functions.length * 2) +
+            // export entry for each export
+            this.estimatedExportBytes;
     }
 
     get current() {
@@ -274,17 +316,23 @@ export class WasmBuilder {
         this.endSection();
     }
 
-    generateImportSection () {
+    _generateImportSection () {
+        const allImports = Object.values(this.importedFunctions);
+        const importsToEmit = allImports.filter(f => f.index !== undefined);
+        importsToEmit.sort((lhs, rhs) => lhs.index! - rhs.index!);
+
         // Import section
         this.beginSection(2);
-        this.appendULeb(1 + this.importsToEmit.length + this.constantSlots.length);
-
-        for (let i = 0; i < this.importsToEmit.length; i++) {
-            const tup = this.importsToEmit[i];
-            this.appendName(tup[0]);
-            this.appendName(tup[1]);
-            this.appendU8(tup[2]);
-            this.appendULeb(tup[3]);
+        this.appendULeb(1 + importsToEmit.length + this.constantSlots.length);
+
+        // console.log(`referenced ${importsToEmit.length}/${allImports.length} import(s)`);
+        for (let i = 0; i < importsToEmit.length; i++) {
+            const ifi = importsToEmit[i];
+            // console.log(`  #${ifi.index} ${ifi.module}.${ifi.name} = ${ifi.friendlyName}`);
+            this.appendName(ifi.module);
+            this.appendName(ifi.name);
+            this.appendU8(0x0); // function
+            this.appendU8(ifi.typeIndex);
         }
 
         for (let i = 0; i < this.constantSlots.length; i++) {
@@ -306,42 +354,122 @@ export class WasmBuilder {
 
     defineImportedFunction (
         module: string, name: string, functionTypeName: string,
-        wasmName?: string
-    ) {
-        const index = this.importedFunctionCount++;
+        assumeUsed: boolean, wasmName?: string
+    ) : ImportedFunctionInfo {
         const type = this.functionTypes[functionTypeName];
         if (!type)
             throw new Error("No function type named " + functionTypeName);
         const typeIndex = type[0];
-        this.importedFunctions[name] = [
-            index, typeIndex, type[3]
-        ];
-        this.importsToEmit.push([module, wasmName || name, 0, typeIndex]);
-        return index;
+        const result = this.importedFunctions[name] = {
+            index: assumeUsed ? this.importedFunctionCount++ : undefined,
+            typeIndex,
+            module,
+            name: wasmName || name,
+            friendlyName: name,
+        };
+        return result;
+    }
+
+    defineFunction (
+        options: {
+            type: string,
+            name: string,
+            export: boolean,
+            locals: { [name: string]: WasmValtype }
+        }, generator: Function
+    ) {
+        const rec : FunctionInfo = {
+            index: this.functions.length,
+            name: options.name,
+            typeName: options.type,
+            typeIndex: this.functionTypes[options.type][0],
+            export: options.export,
+            locals: options.locals,
+            generator,
+            error: null,
+            blob: null
+        };
+        this.functions.push(rec);
+        if (rec.export)
+            this.estimatedExportBytes += rec.name.length + 8;
+        return rec;
+    }
+
+    emitImportsAndFunctions () {
+        let exportCount = 0;
+        for (let i = 0; i < this.functions.length; i++) {
+            const func = this.functions[i];
+            if (func.export)
+                exportCount++;
+
+            this.beginFunction(func.typeName, func.locals);
+            try {
+                func.generator();
+            /*
+            } catch (exc) {
+                func.error = <any>exc;
+            */
+            } finally {
+                func.blob = this.endFunction(false);
+            }
+        }
+
+        this._generateImportSection();
+
+        // Function section
+        this.beginSection(3);
+        this.appendULeb(this.functions.length);
+        for (let i = 0; i < this.functions.length; i++)
+            this.appendULeb(this.functions[i].typeIndex);
+
+        // Export section
+        this.beginSection(7);
+        this.appendULeb(exportCount);
+        for (let i = 0; i < this.functions.length; i++) {
+            const func = this.functions[i];
+            if (!func.export)
+                continue;
+            this.appendName(func.name);
+            this.appendU8(0); // func export
+            this.appendULeb(this.importedFunctionCount + i);
+        }
+
+        // Code section
+        this.beginSection(10);
+        this.appendULeb(this.functions.length);
+        for (let i = 0; i < this.functions.length; i++) {
+            const func = this.functions[i];
+            mono_assert(func.blob, () => `expected function ${func.name} to have a body`);
+            this.appendULeb(func.blob.length);
+            this.appendBytes(func.blob);
+        }
+        this.endSection();
     }
 
     callImport (name: string) {
         const func = this.importedFunctions[name];
         if (!func)
             throw new Error("No imported function named " + name);
+        if (func.index === undefined)
+            func.index = this.importedFunctionCount++;
         this.appendU8(WasmOpcode.call);
-        this.appendULeb(func[0]);
+        this.appendULeb(func.index);
     }
 
     beginSection (type: number) {
         if (this.inSection)
-            this.pop();
+            this._pop(true);
         this.appendU8(type);
-        this.push();
+        this._push();
         this.inSection = true;
     }
 
     endSection () {
         if (!this.inSection)
             throw new Error("Not in section");
         if (this.inFunction)
-            this.endFunction();
-        this.pop();
+            this.endFunction(true);
+        this._pop(true);
         this.inSection = false;
     }
 
@@ -350,8 +478,8 @@ export class WasmBuilder {
         locals?: {[name: string]: WasmValtype}
     ) {
         if (this.inFunction)
-            this.endFunction();
-        this.push();
+            throw new Error("Already in function");
+        this._push();
 
         const signature = this.functionTypes[type];
         this.locals.clear();
@@ -449,13 +577,14 @@ export class WasmBuilder {
         this.inFunction = true;
     }
 
-    endFunction () {
+    endFunction (writeToOutput: boolean) {
         if (!this.inFunction)
             throw new Error("Not in function");
         if (this.activeBlocks > 0)
             throw new Error(`${this.activeBlocks} unclosed block(s) at end of function`);
-        this.pop();
+        const result = this._pop(writeToOutput);
         this.inFunction = false;
+        return result;
     }
 
     block (type?: WasmValtype, opcode?: WasmOpcode) {

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

@@ -18,7 +18,7 @@ import {
     getMemberOffset, JiterpMember
 } from "./jiterpreter-support";
 import {
-    sizeOfDataItem, maxModuleSize,
+    sizeOfDataItem,
 
     disabledOpcodes, countCallTargets,
     callTargetCounts, trapTraceErrors,
@@ -178,7 +178,13 @@ export function generate_wasm_body (
             record_abort(traceIp, ip, traceName, "end-of-body");
             break;
         }
+
+        // HACK: Browsers set a limit of 4KB, we lower it slightly since a single opcode
+        //  might generate a ton of code and we generate a bit of an epilogue after
+        //  we finish
+        const maxModuleSize = 3850;
         if (builder.size >= maxModuleSize - builder.bytesGeneratedSoFar) {
+            // console.log(`trace too big, estimated size is ${builder.size + builder.bytesGeneratedSoFar}`);
             record_abort(traceIp, ip, traceName, "trace-too-big");
             break;
         }
@@ -1238,6 +1244,7 @@ export function generate_wasm_body (
     //  from there.
     builder.local("eip");
     builder.appendU8(WasmOpcode.return_);
+    builder.appendU8(WasmOpcode.end);
 
     return result;
 }