GlobalStructInference: Un-nest struct.news in globals when that is helpful

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,48 +450,101 @@ 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);
 
         // 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.
           return;
         }
 
         // 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) {
         if (refinalize) {
           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();
+    }
   }
 };