[Wasm GC] Local-Subtyping pass (#3765)

If a local is say anyref, but all values assigned to it are something
more specific like funcref, then we can make the type of the local
more specific. In principle that might allow further optimizations, as
the local.gets of that local will have a more specific type that their
users can see, like this:

(import .. (func $get-funcref (result funcref)))

(func $foo (result i32)
  (local $x anyref)
  (local.set $x (call $get-funcref))
  (ref.is_func (local.get $x))
)
=>
(func $foo (result i32)
  (local $x funcref) ;; updated to a subtype of the original
  (local.set $x (call $get-funcref))
  (ref.is_func (local.get $x)) ;; this can now be optimized to "1"
)

A possible downside is that using more specific types may not end
up allowing optimizations but may end up increasing the size of
the binary (say, replacing lots of anyref with various specific
types that compress more poorly; also, for recursive types the LUB
may be a unique type appearing nowhere else in the wasm). We
should investigate the code size factors more later.

1 files changed, 142 insertions, 0 deletions

diff --git a/src/passes/LocalSubtyping.cpp b/src/passes/LocalSubtyping.cpp
new file mode 100644
index 000000000..f435e53d4
--- /dev/null
+++ b/src/passes/LocalSubtyping.cpp
@@ -0,0 +1,142 @@
+/*
+ * Copyright 2021 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.
+ */
+
+//
+// Refines the types of locals where possible. That is, if a local is assigned
+// types that are more specific than the local's declared type, refine the
+// declared type. This can then potentially unlock optimizations later when the
+// local is used, as we have more type info. (However, it may also increase code
+// size in theory, if we end up declaring more types - TODO investigate.)
+//
+
+#include <ir/find_all.h>
+#include <ir/linear-execution.h>
+#include <ir/utils.h>
+#include <pass.h>
+#include <wasm.h>
+
+namespace wasm {
+
+struct LocalSubtyping : public WalkerPass<PostWalker<LocalSubtyping>> {
+  bool isFunctionParallel() override { return true; }
+
+  Pass* create() override { return new LocalSubtyping(); }
+
+  // Shared code to find all sets or gets for each local index. Returns a vector
+  // that maps local indexes => vector of LocalGet*|LocalSet* expressions.
+  template<typename T>
+  std::vector<std::vector<T*>> getLocalOperations(Function* func) {
+    std::vector<std::vector<T*>> ret;
+    ret.resize(func->getNumLocals());
+    FindAll<T> operations(func->body);
+    for (auto* operation : operations.list) {
+      ret[operation->index].push_back(operation);
+    }
+    return ret;
+  }
+
+  void doWalkFunction(Function* func) {
+    if (!getModule()->features.hasGC()) {
+      return;
+    }
+
+    auto varBase = func->getVarIndexBase();
+    auto numLocals = func->getNumLocals();
+
+    auto setsForLocal = getLocalOperations<LocalSet>(func);
+    auto getsForLocal = getLocalOperations<LocalGet>(func);
+
+    // Keep iterating while we find things to change. There can be chains like
+    // X -> Y -> Z where one change enables more. Note that we are O(N^2) on
+    // that atm, but it is a rare pattern as general optimizations
+    // (SimplifyLocals and CoalesceLocals) break up such things; also, even if
+    // we tracked changes more carefully we'd have the case of nested tees
+    // where we could still be O(N^2), so we'd need something more complex here
+    // involving topological sorting. Leave that for if the need arises.
+
+    // TODO: handle cycles of X -> Y -> X etc.
+
+    bool more;
+
+    do {
+      more = false;
+
+      // First, refinalize which will recompute least upper bounds on ifs and
+      // blocks, etc., potentially finding a more specific type. Note that
+      // that utility does not tell us if it changed anything, so we depend on
+      // the next step for knowing if there is more work to do.
+      ReFinalize().walkFunctionInModule(func, getModule());
+
+      // Second, find vars whose actual applied values allow a more specific
+      // type.
+
+      for (Index i = varBase; i < numLocals; i++) {
+        // Find all the types assigned to the var, and compute the optimal LUB.
+        // Note that we do not need to take into account the initial value of
+        // zero or null that locals have: that value has the type of the local,
+        // which is a supertype of all the assigned values anyhow. It will never
+        // be able to tell us of a more specific subtype that is possible.
+        std::unordered_set<Type> types;
+        for (auto* set : setsForLocal[i]) {
+          types.insert(set->value->type);
+        }
+        if (types.empty()) {
+          // Nothing is assigned to this local (other opts will remove it).
+          continue;
+        }
+
+        auto oldType = func->getLocalType(i);
+        auto newType = Type::getLeastUpperBound(types);
+        assert(newType != Type::none); // in valid wasm there must be a LUB
+
+        // Remove non-nullability if we disallow that in locals.
+        if (!getModule()->features.hasGCNNLocals() && newType.isNonNullable()) {
+          newType = Type(newType.getHeapType(), Nullable);
+          // Note that the old type must have been nullable as well, as non-
+          // nullable types cannot be locals without that feature being enabled,
+          // which means that we will not have to do any extra work to handle
+          // non-nullability if we update the type: we are just updating the
+          // heap type, and leaving the type nullable as it was.
+          assert(oldType.isNullable());
+        }
+
+        if (newType != oldType) {
+          // We found a more specific type!
+          assert(Type::isSubType(newType, oldType));
+          func->vars[i - varBase] = newType;
+          more = true;
+
+          // Update gets and tees.
+          for (auto* get : getsForLocal[i]) {
+            get->type = newType;
+          }
+
+          // NB: These tee updates will not be needed if the type of tees
+          //     becomes that of their value, in the spec.
+          for (auto* set : setsForLocal[i]) {
+            if (set->isTee() && set->type != Type::unreachable) {
+              set->type = newType;
+            }
+          }
+        }
+      }
+    } while (more);
+  }
+};
+
+Pass* createLocalSubtypingPass() { return new LocalSubtyping(); }
+
+} // namespace wasm