Bug 1869365 - wasm: Add shell function to dump MIR/LIR of wasm functi…

…on. r=jseward

Usage:
  wasmDumpIon(bytecode, funcIndex, "mir" | "lir" | "opt-mir");

Differential Revision: https://phabricator.services.mozilla.com/D196078

js/src/builtin/TestingFunctions.cpp

@@ -2045,6 +2045,95 @@ static bool WasmDisassemble(JSContext* cx, unsigned argc, Value* vp) {
   return false;
 }
 
+static bool ToIonDumpContents(JSContext* cx, HandleValue value,
+                              wasm::IonDumpContents* contents) {
+  RootedString option(cx, JS::ToString(cx, value));
+
+  if (!option) {
+    return false;
+  }
+
+  bool isEqual = false;
+  if (!JS_StringEqualsLiteral(cx, option, "mir", &isEqual) || isEqual) {
+    *contents = wasm::IonDumpContents::UnoptimizedMIR;
+    return isEqual;
+  } else if (!JS_StringEqualsLiteral(cx, option, "unopt-mir", &isEqual) ||
+             isEqual) {
+    *contents = wasm::IonDumpContents::UnoptimizedMIR;
+    return isEqual;
+  } else if (!JS_StringEqualsLiteral(cx, option, "opt-mir", &isEqual) ||
+             isEqual) {
+    *contents = wasm::IonDumpContents::OptimizedMIR;
+    return isEqual;
+  } else if (!JS_StringEqualsLiteral(cx, option, "lir", &isEqual) || isEqual) {
+    *contents = wasm::IonDumpContents::LIR;
+    return isEqual;
+  } else {
+    return false;
+  }
+}
+
+static bool WasmDumpIon(JSContext* cx, unsigned argc, Value* vp) {
+  if (!wasm::HasSupport(cx)) {
+    JS_ReportErrorASCII(cx, "wasm support unavailable");
+    return false;
+  }
+
+  CallArgs args = CallArgsFromVp(argc, vp);
+
+  args.rval().set(UndefinedValue());
+
+  SharedMem<uint8_t*> dataPointer;
+  size_t byteLength;
+  if (!args.get(0).isObject() || !IsBufferSource(args.get(0).toObjectOrNull(),
+                                                 &dataPointer, &byteLength)) {
+    JS_ReportErrorASCII(cx, "argument is not a buffer source");
+    return false;
+  }
+
+  uint32_t targetFuncIndex;
+  if (!ToUint32(cx, args.get(1), &targetFuncIndex)) {
+    JS_ReportErrorASCII(cx, "argument is not a func index");
+    return false;
+  }
+
+  wasm::IonDumpContents contents = wasm::IonDumpContents::Default;
+  if (args.length() > 2 && !ToIonDumpContents(cx, args.get(2), &contents)) {
+    JS_ReportErrorASCII(cx, "argument is not a valid dump contents");
+    return false;
+  }
+
+  wasm::MutableBytes bytecode = cx->new_<wasm::ShareableBytes>();
+  if (!bytecode) {
+    return false;
+  }
+  if (!bytecode->append(dataPointer.unwrap(), byteLength)) {
+    ReportOutOfMemory(cx);
+    return false;
+  }
+
+  UniqueChars error;
+  JSSprinter out(cx);
+  if (!out.init()) {
+    ReportOutOfMemory(cx);
+    return false;
+  }
+
+  if (!wasm::DumpIonFunctionInModule(*bytecode, targetFuncIndex, contents, out,
+                                     &error)) {
+    if (error) {
+      JS_ReportErrorNumberUTF8(cx, GetErrorMessage, nullptr,
+                               JSMSG_WASM_COMPILE_ERROR, error.get());
+      return false;
+    }
+    ReportOutOfMemory(cx);
+    return false;
+  }
+
+  args.rval().set(StringValue(out.release(cx)));
+  return true;
+}
+
 enum class Flag { Tier2Complete, Deserialized };
 
 static bool WasmReturnFlag(JSContext* cx, unsigned argc, Value* vp, Flag flag) {
@@ -9758,6 +9847,15 @@ JS_FOR_WASM_FEATURES(WASM_FEATURE)
 "      ImportJitExit    - wasm-to-jitted-JS stubs\n"
 "      all              - all kinds, including obscure ones\n"),
 
+    JS_FN_HELP("wasmDumpIon", WasmDumpIon, 2, 0,
+"wasmDumpIon(bytecode, funcIndex, [, contents])\n",
+"wasmDumpIon(bytecode, funcIndex, [, contents])"
+"  Returns a dump of compiling a function in the specified module with Ion."
+"  The `contents` flag controls what is dumped. one of:"
+"    `mir` | `unopt-mir`: Unoptimized MIR (the default)"
+"    `opt-mir`: Optimized MIR"
+"    `lir`: LIR"),
+
     JS_FN_HELP("wasmHasTier2CompilationCompleted", WasmHasTier2CompilationCompleted, 1, 0,
 "wasmHasTier2CompilationCompleted(module)",
 "  Returns a boolean indicating whether a given module has finished compiled code for tier2. \n"

js/src/wasm/WasmBaselineCompile.cpp

@@ -11972,10 +11972,7 @@ bool js::wasm::BaselineCompileFunctions(const ModuleEnvironment& moduleEnv,
     // Build the local types vector.
 
     ValTypeVector locals;
-    if (!locals.appendAll(moduleEnv.funcs[func.index].type->args())) {
-      return false;
-    }
-    if (!DecodeLocalEntries(d, *moduleEnv.types, moduleEnv.features, &locals)) {
+    if (!DecodeLocalEntriesWithParams(d, moduleEnv, func.index, &locals)) {
       return false;
     }
 

js/src/wasm/WasmCompile.cpp

@@ -726,8 +726,8 @@ void CompilerEnvironment::computeParameters(Decoder& d) {
   state_ = Computed;
 }
 
-template <class DecoderT>
-static bool DecodeFunctionBody(DecoderT& d, ModuleGenerator& mg,
+template <class DecoderT, class ModuleGeneratorT>
+static bool DecodeFunctionBody(DecoderT& d, ModuleGeneratorT& mg,
                                uint32_t funcIndex) {
   uint32_t bodySize;
   if (!d.readVarU32(&bodySize)) {
@@ -751,9 +751,9 @@ static bool DecodeFunctionBody(DecoderT& d, ModuleGenerator& mg,
                            bodyBegin + bodySize);
 }
 
-template <class DecoderT>
+template <class DecoderT, class ModuleGeneratorT>
 static bool DecodeCodeSection(const ModuleEnvironment& env, DecoderT& d,
-                              ModuleGenerator& mg) {
+                              ModuleGeneratorT& mg) {
   if (!env.codeSection) {
     if (env.numFuncDefs() != 0) {
       return d.fail("expected code section");
@@ -996,3 +996,46 @@ SharedModule wasm::CompileStreaming(
 
   return mg.finishModule(*bytecode, streamEnd.tier2Listener);
 }
+
+class DumpIonModuleGenerator {
+ private:
+  ModuleEnvironment& moduleEnv_;
+  uint32_t targetFuncIndex_;
+  IonDumpContents contents_;
+  GenericPrinter& out_;
+  UniqueChars* error_;
+
+ public:
+  DumpIonModuleGenerator(ModuleEnvironment& moduleEnv, uint32_t targetFuncIndex,
+                         IonDumpContents contents, GenericPrinter& out,
+                         UniqueChars* error)
+      : moduleEnv_(moduleEnv),
+        targetFuncIndex_(targetFuncIndex),
+        contents_(contents),
+        out_(out),
+        error_(error) {}
+
+  bool finishFuncDefs() { return true; }
+  bool compileFuncDef(uint32_t funcIndex, uint32_t lineOrBytecode,
+                      const uint8_t* begin, const uint8_t* end) {
+    if (funcIndex != targetFuncIndex_) {
+      return true;
+    }
+
+    FuncCompileInput input(funcIndex, lineOrBytecode, begin, end,
+                           Uint32Vector());
+    return IonDumpFunction(moduleEnv_, input, contents_, out_, error_);
+  }
+};
+
+bool wasm::DumpIonFunctionInModule(const ShareableBytes& bytecode,
+                                   uint32_t targetFuncIndex,
+                                   IonDumpContents contents,
+                                   GenericPrinter& out, UniqueChars* error) {
+  UniqueCharsVector warnings;
+  Decoder d(bytecode.bytes, 0, error, &warnings);
+  ModuleEnvironment moduleEnv(FeatureArgs::allEnabled());
+  DumpIonModuleGenerator mg(moduleEnv, targetFuncIndex, contents, out, error);
+  return moduleEnv.init() && DecodeModuleEnvironment(d, &moduleEnv) &&
+         DecodeCodeSection(moduleEnv, d, mg);
+}

js/src/wasm/WasmCompile.h

@@ -93,6 +93,19 @@ SharedModule CompileStreaming(const CompileArgs& args, const Bytes& envBytes,
                               const Atomic<bool>& cancelled, UniqueChars* error,
                               UniqueCharsVector* warnings);
 
+// What to print out from dumping a function from Ion.
+enum class IonDumpContents {
+  UnoptimizedMIR,
+  OptimizedMIR,
+  LIR,
+
+  Default = UnoptimizedMIR,
+};
+
+bool DumpIonFunctionInModule(const ShareableBytes& bytecode,
+                             uint32_t targetFuncIndex, IonDumpContents contents,
+                             GenericPrinter& out, UniqueChars* error);
+
 }  // namespace wasm
 }  // namespace js
 

js/src/wasm/WasmCompileArgs.h

@@ -103,6 +103,15 @@ struct FeatureArgs {
   FeatureArgs(FeatureArgs&&) = default;
 
   static FeatureArgs build(JSContext* cx, const FeatureOptions& options);
+  static FeatureArgs allEnabled() {
+    FeatureArgs args;
+#define WASM_FEATURE(NAME, LOWER_NAME, ...) args.LOWER_NAME = true;
+    JS_FOR_WASM_FEATURES(WASM_FEATURE)
+#undef WASM_FEATURE
+    args.sharedMemory = Shareable::True;
+    args.simd = true;
+    return args;
+  }
 
 #define WASM_FEATURE(NAME, LOWER_NAME, ...) bool LOWER_NAME;
   JS_FOR_WASM_FEATURES(WASM_FEATURE)