Minimal JS legalization (#1824)

Even when we don't want to fully legalize code for JS, we should still legalize things that only JS cares about. In particular, dynCall_* methods are used from JS to call into the wasm table, and if they exist they are only for JS, so we should only legalize them. The use case motivating this is that in dynamic linking you may want to disable legalization, so that wasm=>wasm module calls are fast even with i64s, but you do still need dynCalls to be legalized even in that case, otherwise an invoke with an i64 parameter would fail.
WebAssembly · Dec 15, 2018 · 6cc5984 · 6cc5984
1 parent e8f5842
commit 6cc5984
Show file tree

Hide file tree

Showing 9 changed files with 175 additions and 52 deletions.
diff --git a/src/abi/js.h b/src/abi/js.h
@@ -0,0 +1,43 @@
+/*
+ * Copyright 2018 WebAssembly Community Group participants
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef wasm_abi_abi_h
+#define wasm_abi_abi_h
+
+#include "wasm.h"
+
+namespace wasm {
+
+namespace ABI {
+
+enum LegalizationLevel {
+  Full = 0,
+  Minimal = 1
+};
+
+inline std::string getLegalizationPass(LegalizationLevel level) {
+  if (level == Full) {
+    return "legalize-js-interface";
+  } else {
+    return "legalize-js-interface-minimally";
+  }
+}
+
+} // namespace ABI
+
+} // namespace wasm
+
+#endif // wasm_abi_abi_h
diff --git a/src/asm2wasm.h b/src/asm2wasm.h
@@ -41,6 +41,7 @@
 #include "wasm-builder.h"
 #include "wasm-emscripten.h"
 #include "wasm-module-building.h"
+#include "abi/js.h"
 
 namespace wasm {
 
@@ -1452,9 +1453,9 @@ void Asm2WasmBuilder::processAsm(Ref ast) {
   // finalizeCalls also does autoDrop, which is crucial for the non-optimizing case,
   // so that the output of the first pass is valid
   passRunner.add<FinalizeCalls>(this);
-  if (legalizeJavaScriptFFI) {
-    passRunner.add("legalize-js-interface");
-  }
+  passRunner.add(ABI::getLegalizationPass(
+    legalizeJavaScriptFFI ? ABI::Full : ABI::Minimal
+  ));
   if (runOptimizationPasses) {
     // autodrop can add some garbage
     passRunner.add("vacuum");

diff --git a/src/passes/LegalizeJSInterface.cpp b/src/passes/LegalizeJSInterface.cpp
@@ -22,6 +22,13 @@
 // stub methods added in this pass, that thunk i64s into i32, i32 and
 // vice versa as necessary.
 //
+// We can also legalize in a "minimal" way, that is, only JS-specific
+// components, that only JS will care about, such as dynCall methods
+// (wasm will never call them, as it can share the table directly). E.g.
+// is dynamic linking, where we can avoid legalizing wasm=>wasm calls
+// across modules, we still want to legalize dynCalls so JS can call into the
+// table even to a signature that is not legal.
+//
 // This pass also legalizes according to asm.js FFI rules, which
 // disallow f32s. TODO: an option to not do that, if it matters?
 //
@@ -40,68 +47,74 @@
 namespace wasm {
 
 struct LegalizeJSInterface : public Pass {
+  bool full;
+
+  LegalizeJSInterface(bool full) : full(full) {}
+
   void run(PassRunner* runner, Module* module) override {
     // for each illegal export, we must export a legalized stub instead
     for (auto& ex : module->exports) {
       if (ex->kind == ExternalKind::Function) {
         // if it's an import, ignore it
         auto* func = module->getFunction(ex->value);
-        if (isIllegal(func)) {
+        if (isIllegal(func) && isRelevant(ex.get(), func)) {
           auto legalName = makeLegalStub(func, module);
           ex->value = legalName;
         }
       }
     }
-    // Avoid iterator invalidation later.
-    std::vector<Function*> originalFunctions;
-    for (auto& func : module->functions) {
-      originalFunctions.push_back(func.get());
-    }
-    // for each illegal import, we must call a legalized stub instead
-    for (auto* im : originalFunctions) {
-      if (im->imported() && isIllegal(module->getFunctionType(im->type))) {
-        auto funcName = makeLegalStubForCalledImport(im, module);
-        illegalImportsToLegal[im->name] = funcName;
-        // we need to use the legalized version in the table, as the import from JS
-        // is legal for JS. Our stub makes it look like a native wasm function.
-        for (auto& segment : module->table.segments) {
-          for (auto& name : segment.data) {
-            if (name == im->name) {
-              name = funcName;
+    if (full) {
+      // Avoid iterator invalidation later.
+      std::vector<Function*> originalFunctions;
+      for (auto& func : module->functions) {
+        originalFunctions.push_back(func.get());
+      }
+      // for each illegal import, we must call a legalized stub instead
+      for (auto* im : originalFunctions) {
+        if (im->imported() && isIllegal(module->getFunctionType(im->type))) {
+          auto funcName = makeLegalStubForCalledImport(im, module);
+          illegalImportsToLegal[im->name] = funcName;
+          // we need to use the legalized version in the table, as the import from JS
+          // is legal for JS. Our stub makes it look like a native wasm function.
+          for (auto& segment : module->table.segments) {
+            for (auto& name : segment.data) {
+              if (name == im->name) {
+                name = funcName;
+              }
             }
           }
         }
       }
-    }
-    if (illegalImportsToLegal.size() > 0) {
-      for (auto& pair : illegalImportsToLegal) {
-        module->removeFunction(pair.first);
-      }
+      if (illegalImportsToLegal.size() > 0) {
+        for (auto& pair : illegalImportsToLegal) {
+          module->removeFunction(pair.first);
+        }
 
-      // fix up imports: call_import of an illegal must be turned to a call of a legal
+        // fix up imports: call_import of an illegal must be turned to a call of a legal
 
-      struct FixImports : public WalkerPass<PostWalker<FixImports>> {
-        bool isFunctionParallel() override { return true; }
+        struct FixImports : public WalkerPass<PostWalker<FixImports>> {
+          bool isFunctionParallel() override { return true; }
 
-        Pass* create() override { return new FixImports(illegalImportsToLegal); }
+          Pass* create() override { return new FixImports(illegalImportsToLegal); }
 
-        std::map<Name, Name>* illegalImportsToLegal;
+          std::map<Name, Name>* illegalImportsToLegal;
 
-        FixImports(std::map<Name, Name>* illegalImportsToLegal) : illegalImportsToLegal(illegalImportsToLegal) {}
+          FixImports(std::map<Name, Name>* illegalImportsToLegal) : illegalImportsToLegal(illegalImportsToLegal) {}
 
-        void visitCall(Call* curr) {
-          auto iter = illegalImportsToLegal->find(curr->target);
-          if (iter == illegalImportsToLegal->end()) return;
+          void visitCall(Call* curr) {
+            auto iter = illegalImportsToLegal->find(curr->target);
+            if (iter == illegalImportsToLegal->end()) return;
 
-          if (iter->second == getFunction()->name) return; // inside the stub function itself, is the one safe place to do the call
-          replaceCurrent(Builder(*getModule()).makeCall(iter->second, curr->operands, curr->type));
-        }
-      };
+            if (iter->second == getFunction()->name) return; // inside the stub function itself, is the one safe place to do the call
+            replaceCurrent(Builder(*getModule()).makeCall(iter->second, curr->operands, curr->type));
+          }
+        };
 
-      PassRunner passRunner(module);
-      passRunner.setIsNested(true);
-      passRunner.add<FixImports>(&illegalImportsToLegal);
-      passRunner.run();
+        PassRunner passRunner(module);
+        passRunner.setIsNested(true);
+        passRunner.add<FixImports>(&illegalImportsToLegal);
+        passRunner.run();
+      }
     }
   }
 
@@ -118,6 +131,11 @@ struct LegalizeJSInterface : public Pass {
     return false;
   }
 
+  bool isRelevant(Export* ex, Function* func) {
+    if (full) return true;
+    // We are doing minimal legalization - just what JS needs.
+    return ex->name.startsWith("dynCall_");
+  }
 
   // JS calls the export, so it must call a legal stub that calls the actual wasm function
   Name makeLegalStub(Function* func, Module* module) {
@@ -256,7 +274,11 @@ struct LegalizeJSInterface : public Pass {
 };
 
 Pass *createLegalizeJSInterfacePass() {
-  return new LegalizeJSInterface();
+  return new LegalizeJSInterface(true);
+}
+
+Pass *createLegalizeJSInterfaceMinimallyPass() {
+  return new LegalizeJSInterface(false);
 }
 
 } // namespace wasm

diff --git a/src/passes/pass.cpp b/src/passes/pass.cpp
@@ -85,6 +85,7 @@ void PassRegistry::registerPasses() {
   registerPass("inlining", "inline functions (you probably want inlining-optimizing)", createInliningPass);
   registerPass("inlining-optimizing", "inline functions and optimizes where we inlined", createInliningOptimizingPass);
   registerPass("legalize-js-interface", "legalizes i64 types on the import/export boundary", createLegalizeJSInterfacePass);
+  registerPass("legalize-js-interface-minimally", "legalizes i64 types on the import/export boundary in a minimal manner, only on things only JS will call", createLegalizeJSInterfaceMinimallyPass);
   registerPass("local-cse", "common subexpression elimination inside basic blocks", createLocalCSEPass);
   registerPass("log-execution", "instrument the build with logging of where execution goes", createLogExecutionPass);
   registerPass("i64-to-i32-lowering", "lower all uses of i64s to use i32s instead", createI64ToI32LoweringPass);

diff --git a/src/passes/passes.h b/src/passes/passes.h
@@ -42,6 +42,7 @@ Pass* createI64ToI32LoweringPass();
 Pass* createInliningPass();
 Pass* createInliningOptimizingPass();
 Pass* createLegalizeJSInterfacePass();
+Pass* createLegalizeJSInterfaceMinimallyPass();
 Pass* createLocalCSEPass();
 Pass* createLogExecutionPass();
 Pass* createInstrumentLocalsPass();

diff --git a/src/tools/asm2wasm.cpp b/src/tools/asm2wasm.cpp
@@ -96,7 +96,7 @@ int main(int argc, const char *argv[]) {
            [&wasmOnly](Options *o, const std::string& ) {
              wasmOnly = true;
            })
-      .add("--no-legalize-javascript-ffi", "-nj", "Do not legalize (i64->i32, f32->f64) the imports and exports for interfacing with JS", Options::Arguments::Zero,
+      .add("--no-legalize-javascript-ffi", "-nj", "Do not fully legalize (i64->i32, f32->f64) the imports and exports for interfacing with JS", Options::Arguments::Zero,
            [&legalizeJavaScriptFFI](Options *o, const std::string& ) {
              legalizeJavaScriptFFI = false;
            })

diff --git a/src/tools/wasm-emscripten-finalize.cpp b/src/tools/wasm-emscripten-finalize.cpp
@@ -30,6 +30,7 @@
 #include "wasm-io.h"
 #include "wasm-printing.h"
 #include "wasm-validator.h"
+#include "abi/js.h"
 
 using namespace cashew;
 using namespace wasm;
@@ -83,7 +84,7 @@ int main(int argc, const char *argv[]) {
       .add("--input-source-map", "-ism", "Consume source map from the specified file",
            Options::Arguments::One,
            [&inputSourceMapFilename](Options *o, const std::string& argument) { inputSourceMapFilename = argument; })
-      .add("--no-legalize-javascript-ffi", "-nj", "Do not legalize (i64->i32, "
+      .add("--no-legalize-javascript-ffi", "-nj", "Do not fully legalize (i64->i32, "
            "f32->f64) the imports and exports for interfacing with JS",
            Options::Arguments::Zero,
            [&legalizeJavaScriptFFI](Options *o, const std::string& ) {
@@ -158,13 +159,13 @@ int main(int argc, const char *argv[]) {
   EmscriptenGlueGenerator generator(wasm);
   generator.fixInvokeFunctionNames();
 
-  if (legalizeJavaScriptFFI) {
-    PassRunner passRunner(&wasm);
-    passRunner.setDebug(options.debug);
-    passRunner.setDebugInfo(debugInfo);
-    passRunner.add("legalize-js-interface");
-    passRunner.run();
-  }
+  PassRunner passRunner(&wasm);
+  passRunner.setDebug(options.debug);
+  passRunner.setDebugInfo(debugInfo);
+  passRunner.add(ABI::getLegalizationPass(
+    legalizeJavaScriptFFI ? ABI::Full : ABI::Minimal
+  ));
+  passRunner.run();
 
   std::vector<Name> initializerFunctions;
 

diff --git a/test/passes/legalize-js-interface-minimally.txt b/test/passes/legalize-js-interface-minimally.txt
@@ -0,0 +1,39 @@
+(module
+ (type $FUNCSIG$j (func (result i64)))
+ (type $FUNCSIG$vi (func (param i32)))
+ (import "env" "imported" (func $imported (result i64)))
+ (import "env" "setTempRet0" (func $setTempRet0 (param i32)))
+ (export "func" (func $func))
+ (export "dynCall_foo" (func $legalstub$dyn))
+ (func $func (; 2 ;) (type $FUNCSIG$j) (result i64)
+  (drop
+   (call $imported)
+  )
+  (unreachable)
+ )
+ (func $dyn (; 3 ;) (type $FUNCSIG$j) (result i64)
+  (drop
+   (call $imported)
+  )
+  (unreachable)
+ )
+ (func $legalstub$dyn (; 4 ;) (result i32)
+  (local $0 i64)
+  (set_local $0
+   (call $dyn)
+  )
+  (call $setTempRet0
+   (i32.wrap/i64
+    (i64.shr_u
+     (get_local $0)
+     (i64.const 32)
+    )
+   )
+  )
+  (i32.wrap/i64
+   (get_local $0)
+  )
+ )
+)
+(module
+)
diff --git a/test/passes/legalize-js-interface-minimally.wast b/test/passes/legalize-js-interface-minimally.wast
@@ -0,0 +1,15 @@
+(module
+  (import "env" "imported" (func $imported (result i64)))
+  (export "func" (func $func))
+  (export "dynCall_foo" (func $dyn))
+  (func $func (result i64)
+    (drop (call $imported))
+    (unreachable)
+  )
+  (func $dyn (result i64)
+    (drop (call $imported))
+    (unreachable)
+  )
+)
+(module)
+