GlobalStructInference: Un-nest struct.news in globals when that is helpful (#6688)

If we have

(global $g (struct.new $S
  (i32.const 1)
  (struct.new $T ..)
  (ref.func $f)
))

then before this PR if we wanted to read the middle field we'd stop, as it is non-constant.
However, we can un-nest it, making it constant:

(global $g.unnested (struct.new $T ..))
(global $g (struct.new $S
  (i32.const 1)
  (global.get $g.unnested)
  (ref.func $f)
))

Now the field is a global.get of an immutable global, which is constant. Using this
technique we can handle anything in a struct field, constant or not. The cost of adding
a global is likely offset by the benefit of being able to refer to it directly, as that opens
up more opportunities later.

Concretely, this replaces the constant values we look for in GSI with a variant over
constants or expressions (we do still want to group constants, as multiple globals
with the same constant field can be treated as a whole). And we note cases where we
need to un-nest, and handle those at the end.

1 files changed, 220 insertions, 62 deletions

diff --git a/src/passes/GlobalStructInference.cpp b/src/passes/GlobalStructInference.cpp
index c0b92e3bb..0ddbe9b0f 100644
--- a/src/passes/GlobalStructInference.cpp
+++ b/src/passes/GlobalStructInference.cpp
@@ -48,8 +48,11 @@
 // TODO: Only do the case with a select when shrinkLevel == 0?
 //
 
+#include <variant>
+
 #include "ir/find_all.h"
 #include "ir/module-utils.h"
+#include "ir/names.h"
 #include "ir/possible-constant.h"
 #include "ir/subtypes.h"
 #include "ir/utils.h"
@@ -211,16 +214,98 @@ struct GlobalStructInference : public Pass {
       std::sort(globals.begin(), globals.end());
     }
 
-    // Optimize based on the above.
-    struct FunctionOptimizer
-      : public WalkerPass<PostWalker<FunctionOptimizer>> {
-      bool isFunctionParallel() override { return true; }
+    // We are looking for the case where we can pick between two values using a
+    // single comparison. More than two values, or more than a single
+    // comparison, lead to tradeoffs that may not be worth it.
+    //
+    // Note that situation may involve more than two globals. For example we may
+    // have three relevant globals, but two may have the same value. In that
+    // case we can compare against the third:
+    //
+    //  $global0: (struct.new $Type (i32.const 42))
+    //  $global1: (struct.new $Type (i32.const 42))
+    //  $global2: (struct.new $Type (i32.const 1337))
+    //
+    // (struct.get $Type (ref))
+    //   =>
+    // (select
+    //   (i32.const 1337)
+    //   (i32.const 42)
+    //   (ref.eq (ref) $global2))
+    //
+    // To discover these situations, we compute and group the possible values
+    // that can be read from a particular struct.get, using the following data
+    // structure.
+    struct Value {
+      // A value is either a constant, or if not, then we point to whatever
+      // expression it is.
+      std::variant<PossibleConstantValues, Expression*> content;
+      // The list of globals that have this Value. In the example from above,
+      // the Value for 42 would list globals = [$global0, $global1].
+      // TODO: SmallVector?
+      std::vector<Name> globals;
+
+      bool isConstant() const {
+        return std::get_if<PossibleConstantValues>(&content);
+      }
 
-      std::unique_ptr<Pass> create() override {
-        return std::make_unique<FunctionOptimizer>(parent);
+      const PossibleConstantValues& getConstant() const {
+        assert(isConstant());
+        return std::get<PossibleConstantValues>(content);
       }
 
-      FunctionOptimizer(GlobalStructInference& parent) : parent(parent) {}
+      Expression* getExpression() const {
+        assert(!isConstant());
+        return std::get<Expression*>(content);
+      }
+    };
+
+    // Constant expressions are easy to handle, and we can emit a select as in
+    // the last example. But we can handle non-constant ones too, by un-nesting
+    // the relevant global. Imagine we have this:
+    //
+    //  (global $g (struct.new $S
+    //    (struct.new $T ..)
+    //
+    // We have a nested struct.new here. That is not a constant value, but we
+    // can turn it into a global.get:
+    //
+    //  (global $g.nested (struct.new $T ..)
+    //  (global $g (struct.new $S
+    //    (global.get $g.nested)
+    //
+    // After this un-nesting we end up with a global.get of an immutable global,
+    // which is constant. Note that this adds a global and may increase code
+    // size slightly, but if it lets us infer constant values that may lead to
+    // devirtualization and other large benefits. Later passes can also re-nest.
+    //
+    // We do most of our optimization work in parallel, but we cannot add
+    // globals in parallel, so instead we note the places we need to un-nest in
+    // this data structure and process them at the end.
+    struct GlobalToUnnest {
+      // The global we want to refer to a nested part of, by un-nesting it. The
+      // global contains a struct.new, and we want to refer to one of the
+      // operands of the struct.new directly, which we can do by moving it out
+      // to its own new global.
+      Name global;
+      // The index of the struct.new in the global named |global|.
+      Index index;
+      // The global.get that should refer to the new global. At the end, after
+      // we create a new global and have a name for it, we update this get to
+      // point to it.
+      GlobalGet* get;
+    };
+    using GlobalsToUnnest = std::vector<GlobalToUnnest>;
+
+    struct FunctionOptimizer : PostWalker<FunctionOptimizer> {
+    private:
+      GlobalStructInference& parent;
+      GlobalsToUnnest& globalsToUnnest;
+
+    public:
+      FunctionOptimizer(GlobalStructInference& parent,
+                        GlobalsToUnnest& globalsToUnnest)
+        : parent(parent), globalsToUnnest(globalsToUnnest) {}
 
       bool refinalize = false;
 
@@ -278,65 +363,85 @@ struct GlobalStructInference : public Pass {
           return;
         }
 
-        // We are looking for the case where we can pick between two values
-        // using a single comparison. More than two values, or more than a
-        // single comparison, add tradeoffs that may not be worth it, and a
-        // single value (or no value) is already handled by other passes.
-        //
-        // That situation may involve more than two globals. For example we may
-        // have three relevant globals, but two may have the same value. In that
-        // case we can compare against the third:
-        //
-        //  $global0: (struct.new $Type (i32.const 42))
-        //  $global1: (struct.new $Type (i32.const 42))
-        //  $global2: (struct.new $Type (i32.const 1337))
-        //
-        // (struct.get $Type (ref))
-        //   =>
-        // (select
-        //   (i32.const 1337)
-        //   (i32.const 42)
-        //   (ref.eq (ref) $global2))
-
-        // Find the constant values and which globals correspond to them.
-        // TODO: SmallVectors?
-        std::vector<PossibleConstantValues> values;
-        std::vector<std::vector<Name>> globalsForValue;
-
-        // Check if the relevant fields contain constants.
+        // TODO: SmallVector?
+        std::vector<Value> values;
+
+        // Scan the relevant struct.new operands.
         auto fieldType = field.type;
         for (Index i = 0; i < globals.size(); i++) {
           Name global = globals[i];
           auto* structNew = wasm.getGlobal(global)->init->cast<StructNew>();
-          PossibleConstantValues value;
+          // The value that is read from this struct.new.
+          Value value;
+
+          // Find the value read from the struct and represent it as a Value.
+          PossibleConstantValues constant;
           if (structNew->isWithDefault()) {
-            value.note(Literal::makeZero(fieldType));
+            constant.note(Literal::makeZero(fieldType));
+            value.content = constant;
           } else {
-            value.note(structNew->operands[fieldIndex], wasm);
-            if (!value.isConstant()) {
-              // Give up entirely.
-              return;
+            Expression* operand = structNew->operands[fieldIndex];
+            constant.note(operand, wasm);
+            if (constant.isConstant()) {
+              value.content = constant;
+            } else {
+              value.content = operand;
             }
           }
 
-          // Process the current value, comparing it against the previous.
-          auto found = std::find(values.begin(), values.end(), value);
-          if (found == values.end()) {
-            // This is a new value.
-            assert(values.size() <= 2);
+          // If the value is constant, it may be grouped as mentioned before.
+          // See if it matches anything we've seen before.
+          bool grouped = false;
+          if (value.isConstant()) {
+            for (auto& oldValue : values) {
+              if (oldValue.isConstant() &&
+                  oldValue.getConstant() == value.getConstant()) {
+                // Add us to this group.
+                oldValue.globals.push_back(global);
+                grouped = true;
+                break;
+              }
+            }
+          }
+          if (!grouped) {
+            // This is a new value, so create a new group, unless we've seen too
+            // many unique values. In that case, give up.
             if (values.size() == 2) {
-              // Adding this value would mean we have too many, so give up.
               return;
             }
+            value.globals.push_back(global);
             values.push_back(value);
-            globalsForValue.push_back({global});
-          } else {
-            // This is an existing value.
-            Index index = found - values.begin();
-            globalsForValue[index].push_back(global);
           }
         }
 
+        // Helper for optimization: Given a Value, returns what we should read
+        // for it.
+        auto getReadValue = [&](const Value& value) -> Expression* {
+          if (value.isConstant()) {
+            // This is known to be a constant, so simply emit an expression for
+            // that constant.
+            return value.getConstant().makeExpression(wasm);
+          }
+
+          // Otherwise, this is non-constant, so we are in the situation where
+          // we want to un-nest the value out of the struct.new it is in. Note
+          // that for later work, as we cannot add a global in parallel.
+
+          // There can only be one global in a value that is not constant, which
+          // is the global we want to read from.
+          assert(value.globals.size() == 1);
+
+          // Create a global.get with temporary name, leaving only the updating
+          // of the name to later work.
+          auto* get = builder.makeGlobalGet(value.globals[0],
+                                            value.getExpression()->type);
+
+          globalsToUnnest.emplace_back(
+            GlobalToUnnest{value.globals[0], fieldIndex, get});
+
+          return get;
+        };
+
         // We have some globals (at least 2), and so must have at least one
         // value. And we have already exited if we have more than 2 values (see
         // the early return above) so that only leaves 1 and 2.
@@ -345,7 +450,7 @@ struct GlobalStructInference : public Pass {
           // otherwise return the value.
           replaceCurrent(builder.makeSequence(
             builder.makeDrop(builder.makeRefAs(RefAsNonNull, curr->ref)),
-            values[0].makeExpression(wasm)));
+            getReadValue(values[0])));
           return;
         }
         assert(values.size() == 2);
@@ -353,11 +458,11 @@ struct GlobalStructInference : public Pass {
         // We have two values. Check that we can pick between them using a
         // single comparison. While doing so, ensure that the index we can check
         // on is 0, that is, the first value has a single global.
-        if (globalsForValue[0].size() == 1) {
+        if (values[0].globals.size() == 1) {
           // The checked global is already in index 0.
-        } else if (globalsForValue[1].size() == 1) {
+        } else if (values[1].globals.size() == 1) {
+          // Flip so the value to check is in index 0.
           std::swap(values[0], values[1]);
-          std::swap(globalsForValue[0], globalsForValue[1]);
         } else {
           // Both indexes have more than one option, so we'd need more than one
           // comparison. Give up.
@@ -365,15 +470,19 @@ struct GlobalStructInference : public Pass {
         }
 
         // Excellent, we can optimize here! Emit a select.
-        //
+
+        auto checkGlobal = values[0].globals[0];
+        // Compute the left and right values before the next line, as the order
+        // of their execution matters (they may note globals for un-nesting).
+        auto* left = getReadValue(values[0]);
+        auto* right = getReadValue(values[1]);
         // Note that we must trap on null, so add a ref.as_non_null here.
-        auto checkGlobal = globalsForValue[0][0];
         replaceCurrent(builder.makeSelect(
           builder.makeRefEq(builder.makeRefAs(RefAsNonNull, curr->ref),
                             builder.makeGlobalGet(
                               checkGlobal, wasm.getGlobal(checkGlobal)->type)),
-          values[0].makeExpression(wasm),
-          values[1].makeExpression(wasm)));
+          left,
+          right));
       }
 
       void visitFunction(Function* func) {
@@ -381,12 +490,61 @@ struct GlobalStructInference : public Pass {
           ReFinalize().walkFunctionInModule(func, getModule());
         }
       }
-
-    private:
-      GlobalStructInference& parent;
     };
 
-    FunctionOptimizer(*this).run(getPassRunner(), module);
+    // Find the optimization opportunitites in parallel.
+    ModuleUtils::ParallelFunctionAnalysis<GlobalsToUnnest> optimization(
+      *module, [&](Function* func, GlobalsToUnnest& globalsToUnnest) {
+        if (func->imported()) {
+          return;
+        }
+
+        FunctionOptimizer optimizer(*this, globalsToUnnest);
+        optimizer.walkFunctionInModule(func, module);
+      });
+
+    // Un-nest any globals as needed, using the deterministic order of the
+    // functions in the module.
+    Builder builder(*module);
+    auto addedGlobals = false;
+    for (auto& func : module->functions) {
+      // Each work item here is a global with a struct.new, from which we want
+      // to read a particular index, from a particular global.get.
+      for (auto& [globalName, index, get] : optimization.map[func.get()]) {
+        auto* global = module->getGlobal(globalName);
+        auto* structNew = global->init->cast<StructNew>();
+        assert(index < structNew->operands.size());
+        auto*& operand = structNew->operands[index];
+
+        // If we already un-nested this then we don't need to repeat that work.
+        if (auto* nestedGet = operand->dynCast<GlobalGet>()) {
+          // We already un-nested, and this global.get refers to the new global.
+          // Simply copy the target.
+          get->name = nestedGet->name;
+          assert(get->type == nestedGet->type);
+        } else {
+          // Add a new global, initialized to the operand.
+          auto newName = Names::getValidGlobalName(
+            *module,
+            global->name.toString() + ".unnested." + std::to_string(index));
+          module->addGlobal(builder.makeGlobal(
+            newName, get->type, operand, Builder::Immutable));
+          // Replace the operand with a get of that new global, and update the
+          // original get to read the same.
+          operand = builder.makeGlobalGet(newName, get->type);
+          get->name = newName;
+          addedGlobals = true;
+        }
+      }
+    }
+
+    if (addedGlobals) {
+      // Sort the globals so that added ones appear before their uses.
+      PassRunner runner(module);
+      runner.add("reorder-globals-always");
+      runner.setIsNested(true);
+      runner.run();
+    }
   }
 };