[wasm] Make jiterpreter module size limit configurable + clean up technical debt (dotnet#107318)

Fix spurious "block stack not empty" errors when an error occurs during jiterp codegen
Add missing overflow checks to blob builder appends
Make maximum trace size a runtime option
Move blob builder capacity to a named constant near the top of the file

Author: kg

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

@@ -449,7 +449,7 @@ function flush_wasm_entry_trampoline_jit_queue () {
                 ;
             }
 
-            const buf = builder.getArrayView();
+            const buf = builder.getArrayView(false, true);
             for (let i = 0; i < buf.length; i++) {
                 const b = buf[i];
                 if (b < 0x10)

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

@@ -519,7 +519,7 @@ export function mono_interp_flush_jitcall_queue (): void {
                 ;
             }
 
-            const buf = builder.getArrayView();
+            const buf = builder.getArrayView(false, true);
             for (let i = 0; i < buf.length; i++) {
                 const b = buf[i];
                 if (b < 0x10)

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

@@ -17,7 +17,10 @@ import {
 export const maxFailures = 2,
     maxMemsetSize = 64,
     maxMemmoveSize = 64,
-    shortNameBase = 36;
+    shortNameBase = 36,
+    // NOTE: This needs to be big enough to hold the maximum module size since there's no auto-growth
+    //  support yet. If that becomes a problem, we should just make it growable
+    blobBuilderCapacity = 16 * 1024;
 
 // uint16
 export declare interface MintOpcodePtr extends NativePointer {
@@ -120,6 +123,8 @@ export class WasmBuilder {
 
     clear (constantSlotCount: number) {
         this.options = getOptions();
+        if (this.options.maxModuleSize >= blobBuilderCapacity)
+            throw new Error(`blobBuilderCapacity ${blobBuilderCapacity} is not large enough for jiterpreter-max-module-size of ${this.options.maxModuleSize}`);
         this.stackSize = 1;
         this.inSection = false;
         this.inFunction = false;
@@ -920,8 +925,8 @@ export class WasmBuilder {
         this.appendU8(WasmOpcode.i32_add);
     }
 
-    getArrayView (fullCapacity?: boolean) {
-        if (this.stackSize > 1)
+    getArrayView (fullCapacity?: boolean, suppressDeepStackError?: boolean) {
+        if ((suppressDeepStackError !== true) && this.stackSize > 1)
             throw new Error("Jiterpreter block stack not empty");
         return this.stack[0].getArrayView(fullCapacity);
     }
@@ -942,8 +947,9 @@ export class BlobBuilder {
     textBuf = new Uint8Array(1024);
 
     constructor () {
-        this.capacity = 16 * 1024;
+        this.capacity = blobBuilderCapacity;
         this.buffer = <any>Module._malloc(this.capacity);
+        mono_assert(this.buffer, () => `Failed to allocate ${blobBuilderCapacity}b buffer for BlobBuilder`);
         localHeapViewU8().fill(0, this.buffer, this.buffer + this.capacity);
         this.size = 0;
         this.clear();
@@ -1051,18 +1057,26 @@ export class BlobBuilder {
         if (typeof (count) !== "number")
             count = this.size;
 
+        if ((destination.size + count) >= destination.capacity)
+            throw new Error("Destination buffer full");
+
         localHeapViewU8().copyWithin(destination.buffer + destination.size, this.buffer, this.buffer + count);
         destination.size += count;
     }
 
     appendBytes (bytes: Uint8Array, count?: number) {
         const result = this.size;
         const heapU8 = localHeapViewU8();
+        const actualCount = (typeof (count) !== "number")
+            ? bytes.length
+            : count;
+
+        if ((this.size + actualCount) >= this.capacity)
+            throw new Error("Buffer full");
+
         if (bytes.buffer === heapU8.buffer) {
-            if (typeof (count) !== "number")
-                count = bytes.length;
-            heapU8.copyWithin(this.buffer + result, bytes.byteOffset, bytes.byteOffset + count);
-            this.size += count;
+            heapU8.copyWithin(this.buffer + result, bytes.byteOffset, bytes.byteOffset + actualCount);
+            this.size += actualCount;
         } else {
             if (typeof (count) === "number")
                 bytes = new Uint8Array(bytes.buffer, bytes.byteOffset, count);
@@ -1950,6 +1964,7 @@ export type JiterpreterOptions = {
     wasmBytesLimit: number;
     tableSize: number;
     aotTableSize: number;
+    maxModuleSize: number;
 }
 
 const optionNames: { [jsName: string]: string } = {
@@ -1986,6 +2001,7 @@ const optionNames: { [jsName: string]: string } = {
     "wasmBytesLimit": "jiterpreter-wasm-bytes-limit",
     "tableSize": "jiterpreter-table-size",
     "aotTableSize": "jiterpreter-aot-table-size",
+    "maxModuleSize": "jiterpreter-max-module-size",
 };
 
 let optionsVersion = -1;

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

@@ -38,7 +38,7 @@ import {
     traceEip, nullCheckValidation,
     traceNullCheckOptimizations,
     nullCheckCaching, defaultTraceBackBranches,
-    maxCallHandlerReturnAddresses,
+    maxCallHandlerReturnAddresses, moduleHeaderSizeMargin,
 
     mostRecentOptions,
 
@@ -275,10 +275,7 @@ export function generateWasmBody (
             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 maxBytesGenerated = 3840,
+        const maxBytesGenerated = builder.options.maxModuleSize - moduleHeaderSizeMargin,
             spaceLeft = maxBytesGenerated - builder.bytesGeneratedSoFar - builder.cfg.overheadBytes;
         if (builder.size >= spaceLeft) {
             // mono_log_info(`trace too big, estimated size is ${builder.size + builder.bytesGeneratedSoFar}`);

src/mono/browser/runtime/jiterpreter.ts

@@ -72,8 +72,8 @@ export const
     useFullNames = false,
     // Use the mono_debug_count() API (set the COUNT=n env var) to limit the number of traces to compile
     useDebugCount = false,
-    // Web browsers limit synchronous module compiles to 4KB
-    maxModuleSize = 4080;
+    // Subtracted from the maxModuleSize option value to make space for a typical header
+    moduleHeaderSizeMargin = 300;
 
 export const callTargetCounts: { [method: number]: number } = {};
 
@@ -828,7 +828,7 @@ function generate_wasm (
             mono_log_info(`${(<any>(builder.base)).toString(16)} ${methodFullName || traceName} generated ${buffer.length} byte(s) of wasm`);
         modifyCounter(JiterpCounter.BytesGenerated, buffer.length);
 
-        if (buffer.length >= maxModuleSize) {
+        if (buffer.length >= builder.options.maxModuleSize) {
             mono_log_warn(`Jiterpreter generated too much code (${buffer.length} bytes) for trace ${traceName}. Please report this issue.`);
             return 0;
         }
@@ -913,7 +913,7 @@ function generate_wasm (
                 ;
             }
 
-            const buf = builder.getArrayView();
+            const buf = builder.getArrayView(false, true);
             for (let i = 0; i < buf.length; i++) {
                 const b = buf[i];
                 if (b < 0x10)

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

@@ -140,7 +140,8 @@ DEFINE_BOOL(jiterpreter_constant_propagation, "jiterpreter-constant-propagation"
 // 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")
-// any trace that doesn't have at least this many meaningful (non-nop) opcodes in it will be rejected
+// when deciding whether to generate a trace, we sum the value of sequential opcodes that will fit into it
+//  and reject any trace entry point where the score is below this value
 DEFINE_INT(jiterpreter_minimum_trace_value, "jiterpreter-minimum-trace-value", 18, "Reject traces that perform less than this amount of (approximate) work")
 // ensure that we don't create trace entry points too close together
 DEFINE_INT(jiterpreter_minimum_distance_between_traces, "jiterpreter-minimum-distance-between-traces", 4, "Don't insert entry points closer together than this")
@@ -175,6 +176,7 @@ DEFINE_INT(jiterpreter_table_size, "jiterpreter-table-size", 6 * 1024, "Size of
 //  to bloat to an unacceptable degree. In practice this is still better than nothing.
 // FIXME: In the future if we find a way to reduce the number of unique tables we can raise this constant
 DEFINE_INT(jiterpreter_aot_table_size, "jiterpreter-aot-table-size", 3 * 1024, "Size of the jiterpreter AOT trampoline function tables")
+DEFINE_INT(jiterpreter_max_module_size, "jiterpreter-max-module-size", 4080, "Size limit for jiterpreter generated WASM modules")
 #endif // HOST_BROWSER
 
 #if defined(TARGET_WASM) || defined(TARGET_IOS)  || defined(TARGET_TVOS) || defined (TARGET_MACCAT)